perf(webgpu): subgroup cooperative shaders for MatMul, BatchMatMul, Softmax (#141) - #151
Merged
Merged
Conversation
…#141) Add matmulSubgroupShader and batchMatMulSubgroupShader using subgroupAdd for cooperative K-reduction across lanes. Workgroup_size(32) maps to GPU warp/wavefront. Runtime feature gate: subgroup shaders only when device supports FeatureSubgroupOperations, scalar fallback otherwise. Backend requests subgroup feature from adapter when available. Software renderer (CI) always uses scalar path — zero regression risk. 4 new tests: correctness (8 edge cases), scalar fallback, batch, WGSL syntax.
Three-phase subgroup Softmax: subgroupMax for max, subgroupAdd for exp-sum, per-lane normalize. One workgroup(32) per row instead of one thread per row. Runtime feature gate + scalar fallback. Flash Attention subgroup deferred: requires full algorithm redesign (per-query cooperative dot product), not an incremental change. 5 tests: correctness (10 cases), scalar fallback, explicit subgroup path, numerical stability (large values), WGSL syntax validation.
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
This was referenced Aug 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements P0+P1 items from #141 — subgroup cooperative reductions for GPU compute shaders.
MatMul + BatchMatMul (P0)
matmulSubgroupShader: workgroup_size(32), lanes stride over K dimension,subgroupAdd(partial)reduces across lanes, thread 0 writes resultbatchMatMulSubgroupShader: same pattern with batch index viawid.zadapter.Features().Contains(FeatureSubgroupOperations)→ request feature on device →subgroupsEnabled=trueSoftmax (P1)
softmaxSubgroupShader: three-phase cooperative —subgroupMaxfor max,subgroupAddfor exp-sum, per-lane normalizeFlash Attention (P1 — deferred)
Current Flash Attention uses per-thread independent accumulators (64 threads × 1 query each). Subgroup optimization would require restructuring the Q@K dot product loop for per-query cooperative reduction — a full algorithm rewrite, not an incremental change. Tracked as follow-up.
Tests
Test plan
go build ./...GOOS=js GOARCH=wasm go build ./...golangci-lint run— 0 issuesCloses #141