Repository navigation
Conversation
JCNTH
force-pushed
the
webgpu-decode-attention
branch
from
September 24, 2026 16:58
0d4fd4a to
ceb29e2
Compare
…exts The paged decode attention kernel spreads threads across the head dimension (bdx = head_dim / vec), so each q.k dot product ends in a cross-thread allreduce and one trip of the serial loop over the context covers only bdz * bdy positions. On WebGPU, with at most 256 invocations per workgroup and every workgroupBarrier a real sync, that loop dominates decode time at long context. Add a second thread nest for WebGPU that cuts the context into shards owned by a few lanes each (64 dims of q and O per lane in registers), exchanges partial scores through a double-buffered workgroup array so one sync per position suffices, and merges the per-shard online-softmax states in two phases that fit the default 16 KiB of workgroup storage. Wide heads (128, 256) split across threadIdx.y lanes so the workgroup stays at 256 invocations. The existing tiled nest is wrapped in a T.macro unchanged and remains in charge of short contexts; the choice is made on the host from the symbolic page count via gpu_thread_invariant. Other targets are unaffected. Attention time per decode step at 3060 tokens of context, Apple M4 Max / Chrome, sharded vs tiled nest: 2.9x / 4.0x / 3.4x faster for head_dim 64 / 128 / 256 (Llama-3.1-8B shape: 36.6 -> 9.0 ms). Output matches an f64 reference to f16 rounding (0.02% relative L2) at 1..3060 positions.
JCNTH
force-pushed
the
webgpu-decode-attention
branch
from
October 13, 2026 23:05
ceb29e2 to
0abb605
Compare
Preserve upstream tensor layouts and standalone TIRx APIs.
This branch has not been deployed
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
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.
Long-context WebGPU decode attention spends much of its time reducing across the head dimension. Add a sequence-sharded kernel selected above 256 tokens, while retaining the tiled kernel for short contexts and existing behavior on other targets. The layouts fit 256 invocations and 16 KiB shared memory.
Earlier Apple M4 Max/Chrome Canary measurements show 2.91–4.05× lower attention time at 3,060 context, with 0.02% relative L2 against an f64 reference and identical greedy text. The merged implementation passes 27 kernel-configuration tests, including WebGPU layouts, sliding windows and unchanged non-WebGPU behavior. End-to-end results in the companion work also include GEMV/runtime changes and are not attributed solely to this kernel.