JCNTH opened a new pull request, #20426: URL: https://github.com/apache/tvm/pull/20426
The paged-KV decode kernel spends its threads across the head dimension (`bdx = D / 4`), so every q·k dot product is finished by a cross-thread allreduce, K/V are staged through workgroup memory, and one iteration of the serial loop over the context covers only `bdz * bdy` positions: 16 for a GQA model on WebGPU's 256 invocations. Decode attention is bound by memory latency per thread, so that loop dominates at long context. In a per-kernel GPU-timestamp profile of live WebLLM (Llama-3.2-1B, Apple M4 Max) attention is 10% of GPU time per token at 84 tokens of context, 50% at 1k and 75% at 3k, moving KV at ~13 GB/s while the q4 GEMVs reach 430. ## Change A second thread nest for WebGPU (`_make_seq_sharded_decode_nest`), selected on the host from the page count (`nnz_pages > B * 256 / page_size`); the existing tiled nest becomes a `T.macro` and stays in charge of short contexts, where it is faster. Other targets are untouched (the generated PrimFunc is structurally identical for cuda/metal/vulkan/rocm/opencl). - One workgroup per query head; the context is cut into up to 128 shards, each owned by `head_dim / 64` lanes (`threadIdx.y`), so a thread keeps 64 dims of q and of the running output in registers (2 × head_dim floats spills beyond 64 dims). Layouts: 128×1 for D=64, 128×2 for D=128, 64×4 for D=256, always 256 invocations or fewer. - Per position, the lanes of a shard each form a partial dot product and exchange it through workgroup memory, double-buffered by position parity, so one `workgroupBarrier` per position suffices; the online-softmax state is updated per lane. - Shard states are merged in two phases, (max, denom) first, then the weighted outputs a `CHUNK` of dims at a time summed one dim per thread, so everything fits the default 16 KiB of workgroup storage. - A first version with fewer threads and no barriers was *slower* (0.63× at 3k): per-thread cost is memory latency, so throughput follows the number of threads in flight. Hence one workgroup per head. ## Measurements (Apple M4 Max, Chrome Canary, GPU timestamps) Attention time per decode step at 3060 tokens of context, tiled → sharded: **2.95× / 3.89× / 3.49×** for head_dim 64 / 128 / 256 (Hq32/Hkv8 shapes, 16–32 layers). Break-even is at ~100–300 tokens for every head_dim, so the 256-token host switch holds. End to end vs what WebLLM ships (npm bundle + prebuilt lib; this PR's libs also carry the GEMV config of the companion PR), decode tok/s: | model | short prompt | after a ~2750-token prompt | |---|---|---| | Llama-3.1-8B-Instruct-q4f16_1 | 56.5 → 72.1 (1.29×) | 20.3 → 47.9 (**2.35×**) | | Gemma-2-2b-it-q4f16_1 | 91.9 → 126.0 (1.37×) | 22.2 → 59.0 (**2.68×**) | | Llama-3.2-1B-Instruct-q4f16_1 | 200.7 → 279.4 (1.39×) | 83.8 → 167.4 (**2.00×**) | | Qwen3-0.6B-q4f16_1 | 124.3 → 253.1 (2.02×) | 40.8 → 79.5 (1.95×) | Generated text identical in every case (temperature 0). The short-prompt rows are the runtime + GEMV work; the long-prompt rows are where this kernel is selected. ## Correctness Output and LSE match an f64 reference as well as the tiled kernel does (0.02% relative L2 = f16 output rounding, LSE within 2e-6) at nine context lengths from 1 to 3060 including odd shard lengths, with a shuffled page table; with `rotary_mode = 1`, sliding window and batch > 1 the two nests agree to ≤ 0.003% on identical inputs. Through the WebLLM API at ~2.7k tokens, top-5 logprobs shift by a mean 0.078 nat, the same as the floor between the prebuilt lib and a self-built upstream lib (0.070), with identical greedy text. ## Tests `test_webgpu_decode_attention_dispatches_on_context_length` (two kernels, host-side switch, no `blockIdx.z`), `test_webgpu_decode_attention_lanes_follow_head_dim` (layouts for D = 64/128/256/96, storage ≤ 16 KiB), `test_decode_attention_unchanged_off_webgpu`. `tests/python/relax/test_frontend_nn_llm_kernel_config.py`: 18 passed. pre-commit clean. ## Open for reviewers - Switch point (256 tokens) and shard count (128) tuned on one adapter. - The same layout may help other latency-bound backends (Metal, Vulkan); not tried. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
