perf(linear_attention): halve the prefill chunk and split the scan over dv (-25% op, -1.07 s TTFT, 129 -> 49 MiB scratch) - #627
Draft
BoarQing wants to merge 6 commits into
Draft
Conversation
Rewrite the gated_delta linear-attention prefill from a 587-deep serial chunk recurrence (grid=B*Hkv=32 blocks, 25% occupancy, latency-bound) into a three-pass chunk-parallel form that exploits the affine chunk-state map S_out = L_c*S_in + b_c, with the dominant [Q;W]@s0 GEMM on the RDNA3 wave-matrix cores: Pass 1 (grid = B*Hkv*nchunks): S_in-independent locals Uloc=(I-trilT)^-1(beta.*v), W=(I-trilT)^-1(diag(beta.*A)K), Ktilde, a_last. Pass 2 (grid = B*Hkv, light serial scan): U = Uloc - W*S0, S1 = a_last*S0 + Ktilde^T U (fp16-rounded per chunk). Pass 3 (grid = B*Hkv*nchunks): O = scale*(A.*(Q*S0) + P_lower*U), with the stacked [Q;W]@s0 GEMM via __builtin_amdgcn_wmma_f32_16x16x16_f16_w32. Numerically equivalent to the serial kernel (same fp32 math, same per-chunk fp16 rounding); validated cosine=1.0000000 vs an fp64 per-token gated_delta reference. Measured (Qwen3.6-35B, gfx1151, seq 18778): kernel 604.8 ms -> 110.0 ms (5.5x); RGP peak-clock occupancy 25% -> 25-100%, wavefronts 256 -> 150272; E2E 16k prefill TTFT 33476 ms -> 17580 ms (1.90x). No regression at shorter sequences (in128 1.51x, in2048 1.77x E2E). Co-authored-by: Cursor <cursoragent@cursor.com>
… path
Promote the optimized path to the only implementation and drop the bring-up
scaffolding now that it is validated:
- Delete the original serial chunked kernel (linear_attention_prefill_-
chunked_kernel) and the scalar Pass-3 variant (la_pf_pass3_out); they are
superseded by the chunk-parallel + WMMA path.
- Remove all debug/tuning env gates: HIPDNN_EP_LA_PARALLEL, _WMMA, _PROF,
_SCANBLK, _P3BLK, and the per-pass hipEvent timing instrumentation.
- The launcher now always dispatches the three-pass parallel path; on any
launch/alloc failure it returns nonzero so the runtime falls back to the
per-token loop (unchanged higher-level safety net). WMMA needs dk,dv
multiples of 16, so decline otherwise (-> per-token loop).
Behavior for the supported gated_delta config is unchanged: cosine=1.0000000
vs the fp64 reference, kernel 109.6 ms at seq 18778 (5.5x vs the removed
serial kernel).
Note: HIPDNN_EP_LA_PARALLEL still exists for the separate decode-path G-split
(la_par_G) and is untouched here.
Co-authored-by: Cursor <cursoragent@cursor.com>
…buffer Reduce the chunk-parallel gated_delta prefill's device scratch and fix its lifetime, on top of the chunk-parallel + WMMA prefill (#589). - Move the prefill scratch out of a process-static, never-freed g_la_scratch into the per-session RuntimeState::la_scratch pool (grow-on-demand, freed in hipdnn_ep_state_cleanup), mirroring qmoe_scratch/conv_scratch. The kernel now takes a caller-owned scratch + size; the wrapper sizes it via the new hip_linear_attention_prefill_scratch_bytes(). Footprint is now bounded to the session and released on teardown instead of held for the whole process. - Store only the per-token rlast scalar (C floats/chunk) instead of the full Ktilde = rlast*K tile (C*dk floats/chunk); the scan rebuilds Ktilde from the original key tensor. Removes the Kt_g buffer (~1/dk the size) at bit-identical values. Both changes are numerically exact (same fp32 math / per-chunk fp16 rounding); only op-order-invariant scratch and buffer sizing change. Validated end-to-end on Qwen3.5-9B (gfx1151): chunk-parallel prefill path exercised (dk=dv=128), coherent output, no regression. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Bound the chunk-parallel gated_delta prefill's device scratch to a fixed window instead of the whole sequence, on top of the session-scoped scratch change (#596) and the chunk-parallel + WMMA prefill (#589). The per-(head,chunk) buffers (Uloc, W, rlast, a_last and the chunk-start states) were sized by nchunks = seq_len/LA_CHUNK, so the arena grew without bound with prompt length. The three passes now run over LA_WINDOW_CHUNKS = 64 chunks (2048 tokens) at a time, so the arena is sized per window and stops scaling with the sequence (measured via the exported hip_linear_attention_prefill_scratch_bytes, B=1, Hkv=32, dk=dv=128): seq 4096: 256.5 MiB -> 129.3 MiB (1.98x) seq 13529: 847.7 MiB -> 129.3 MiB (6.56x) seq 16384: 1026.1 MiB -> 129.3 MiB (7.94x) seq 32768: 2052.1 MiB -> 129.3 MiB (15.9x) A window still launches B*Hkv*64 = 2048 blocks at Hkv=32, far more than the 40 CUs need to stay saturated, so the chunk-parallel decomposition keeps its parallelism. Ordering needs no extra synchronisation: the three passes are issued back-to-back on one stream, so pass 3 of a window has finished reading Uloc/W/S0 before pass 1 of the next window overwrites them. The scan is already sequential over chunks, so windowing only changes where its running state lives between launches. Across windows it hands off through a small fp16 carry buffer instead of re-reading `state`: the running state is fp16 in LDS, so a round trip through a narrower T (bf16 keeps 8 mantissa bits vs fp16's 10) would perturb it at every boundary. Single-window prefills are bit-identical to before -- is_first == is_last, the carry is never read or written, and it is not allocated at all. Co-authored-by: Cursor <cursoragent@cursor.com>
…er dv The windowed prefill from #606 is limited by LDS, not by arithmetic. Pass 3 asks for 61,952 B of LDS at LA_CHUNK=32, over half of a CU's 64 KB, so only one block is resident per CU and a stall has nothing to hide behind. Two changes address that. LA_CHUNK 32 -> 16 halves every chunk-sized LDS tile. Pass 3 drops to 29,824 B, which fits two blocks per CU. LA_WINDOW_CHUNKS follows 64 -> 32 to hold the window at 512 tokens, which also shrinks the scratch arena. The extra launches are not a real cost: holding LA_CHUNK at 16, a 32-chunk window beats a 128-chunk one by 9.6% despite issuing 3.7x as many. The scan then becomes the tail. It is serial in the chunk index and ran one block per (batch, kv head) -- 32 blocks for this model, fewer than the 40 CUs. But both of its inner products touch only column j of S0, so the recurrence factors completely over dv and splits across blocks. At NSPLIT=4 that is 128 blocks, and the narrower 256-thread block that a dv/4 slice wants is what keeps them resident. Each split re-reads the dv-independent W tile; that duplication is what is traded for the parallelism, and it only pays once LA_CHUNK=16 has shrunk the tile. NSPLIT=1 restores the previous scan exactly. Measured on gfx1151 (Radeon 8060S, 40 CUs), Qwen3-Next-80B-A3B fp16, 30 layers, 18,778-token prompt, paired interleaved runs against #606: isolated op -25.2% (10/10 paired wins) TTFT -1,069 ms (-5.4%, 5/6 paired rounds) scratch arena 129.3 -> 49.1 MiB per pass pass1 -54.3%, scan -5.8%, pass3 -42.2% Correctness is unchanged: an fp64 per-token reference at 18,778 tokens, which crosses 36 window boundaries, gives cosine 0.999999960 and a trailing-half relative L2 of 2.749e-4, matching #606 to every reported digit. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> Made-with: Cursor Co-authored-by: Cursor <cursoragent@cursor.com>
|
Thanks for opening a PR! This project follows LLVM's incremental-development and AI-tool-use Before requesting review, please check that:
Reviewers are assigned through |
L2 Accuracy Results (EP vs CPU)
Threshold: 0.01 | Run: 3532 - Commit: |
MorphiZen EP Performance Results
EPContext Export Performance
EPContext Import Performance
OGA Benchmark Results
OGA Wheel Smoke (Python benchmark_e2e.py)
Run: 3532 - Commit: |
Contributor
Author
Contributor
Author
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
Halves the linear-attention prefill chunk (
LA_CHUNK32 -> 16,LA_WINDOW_CHUNKS64 -> 32) and splits the prefill scan over thedvdimension (LA_SCAN_NSPLIT4, 256-thread scan blocks). On an 18,778-token prompt this takes the linear-attention op down 25.2% and TTFT down 1,069 ms (-5.4%), and shrinks the prefill scratch arena from 129.3 MiB to 49.1 MiB.Numerics are unchanged; this is an occupancy and blocking change only.
Related issue or design
Builds directly on #606. See "Notes for reviewers" for the branch ordering.
Why
#606 capped the prefill scratch but left the kernels' shape alone, and profiling that shape shows the limiter is LDS, not arithmetic.
At
LA_CHUNK=32pass 3 asks for 61,952 B of LDS. A CU has 64 KB, so exactly one block is resident, and with one block per CU a memory stall has nothing to hide behind. The scan has the opposite problem: it is serial in the chunk index and ran one block per(batch, kv head), which is 32 blocks for this model against 40 CUs, so the machine is not even filled once.Both are addressed by changing how work is divided, not what work is done.
What
Chunk 32 -> 16. Every chunk-sized LDS tile halves; pass 3 drops to 29,824 B, which fits two blocks per CU.
LA_WINDOW_CHUNKSfollows 64 -> 32 so the window stays at 512 tokens, which is also where most of the arena reduction comes from. The extra launches this costs are not a real cost: holdingLA_CHUNKat 16 and varying only the window, 32 chunks (37 windows, 49.1 MiB) beats 128 chunks (10 windows, 193.3 MiB) by 9.6%.Scan split over
dv. The scan's two inner products areWS[t][j] = sum_i W[t][i] S0[i][j]andS1[i][j] = a S0[i][j] + sum_s Ktilde[s][i] U[s][j]. Both touch only columnjofS0, so the recurrence factors completely overdvand the columns can be split across blocks without changing the serial structure in the chunk index. AtNSPLIT=4the scan runs 128 blocks instead of 32, each owningdv/4columns, and the narrower 256-thread block is what keeps them resident.Two things worth knowing about the split:
dv-independentWtile. That duplicated read is the cost traded for the parallelism, and it only pays onceLA_CHUNK=16has shrunk the tile: applied on top ofLA_CHUNK=32the split is a 12.8% regression.NSPLIT=1restores the previous scan exactly (dvs == dv,j0 == 0). The launcher declines the config ifdv % LA_SCAN_NSPLIT != 0, so the runtime falls back rather than computing a wrong answer on an unsupported shape.Test plan
Correctness. A standalone harness calls the exported launcher on synthetic fp16 inputs and compares against an fp64 per-token reference. The existing gate only ran the reference at
T <= 512, which never crosses a window boundary and therefore could not see any of this; since every change here alters the multi-window carry hand-off, the reference was run at the full 18,778 tokens, crossing 36 boundaries:Identical to every reported digit.
Cosine alone was not trusted as the gate: a deliberately corrupted build (inter-window carry scaled by 0.99) still passes
cosine > 0.9999atT=4096, because a single perturbed boundary barely moves a whole-tensor cosine. The trailing-half relative L2 moves 2.5x on that same corruption, so it is what the gate keys on.Performance. All figures are paired interleaved runs, arm order rotated per round. This is necessary rather than fussy: the machine has at least two performance states, and the same DLL measured anywhere from 75 to 165 ms per call across this session. Absolute numbers are therefore not comparable across rounds and only within-round paired deltas are reported.
Notes for reviewers
Branch ordering. This targets
mainper CONTRIBUTING, but only the last commit is mine. The other five belong to #606 (feat/linattn-window-tiling) and its basefeat/linattn-scratch-mem, which has no PR of its own. Review just072ce177, or the compare against #606's head:That is 1 file, +67/-32. This should land after #606, or be rebased onto
mainonce #606 does.A third change was measured and rejected. Storing
S0transposed makes pass 3's WMMA B fragment one 32-byte vector load instead of 16 two-byte loads. The ISA confirms the mechanism exactly as predicted,global_load_u16per WMMA falling from 20.7 to 4.7 and the inner loop from 78 instructions to 13, and it is still not worth taking: the transposing store has to read LDS down a column, and the scan loses more than pass 3 gains. The conclusion is that the B-fragment load was never the limiter. It is not in this PR and no dead code for it is left behind.Coverage limits.
NSPLIT=4was verified atB=1, Hkv=32, where the scan was starved of blocks. On a shape whereB*Hkvalready exceeds the CU count the split has nothing to win and the duplicatedWread may make it a small loss; it is worth a second look before this is relied on for large-batch prefill.AI assistance. The profiling, the measurement harnesses, the implementation and this description were produced with Cursor (Claude Opus 4.5), and the commit carries the trailers. The claims above are all from local runs on the hardware listed below, and the correctness gate was validated against a known-bad build before being trusted. I have reviewed the change and own it.
Performance
Hardware, workload, build configuration, and measurement method:
B=1,Hkv=32,dk=dv=128hipEventtiming around each launch inside the window loop, first 60 forwards skipped