Perf/pr596 linear attention - #603
Draft
donghanbj wants to merge 5 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>
Vectorize adjacent scan columns and use all pass-3 wavefronts, cutting seq-2048 prefill latency by about 20% with bit-identical output. 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 |
Contributor
Author
|
/run-test models=all |
Contributor
|
ROCm OGA CI finished with result: UNSTABLE.
|
L2 Accuracy Results (EP vs CPU)
Threshold: 0.01 | Run: 3450 - Commit: |
MorphiZen EP Performance Results
EPContext Export Performance
EPContext Import Performance
OGA Benchmark Results
OGA Wheel Smoke (Python benchmark_e2e.py)
Run: 3450 - Commit: |
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
3 participants
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
Related issue or design
Why
What
Test plan
Notes for reviewers
Checklist