Repository navigation
[ROCm][Bugfix] Fix Triton paged decode leaving output heads unwritten on gfx942 - #63268
Conversation
… on gfx942 With 64-bit store offsets, Triton's AMD buffer_store lowering drops the head stride for the dot-result layout of kernel_paged_attention_2d, so a decode row's heads land on top of each other and most stay unwritten. Offset the output pointer to the token row as a 64-bit scalar and keep the per-head offsets 32-bit. Fill the test output with NaN so an unwritten row fails deterministically. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
/ci run |
|
✅ Triggered Buildkite CI #96825 for commit ℹ️ This PR is 14 commits behind upstream |
| cur_batch_in_all_start_index * output_stride_0 | ||
| + query_head_idx * output_stride_1 | ||
| ) | ||
| # The row base stays a 64-bit scalar and the per-head offsets 32-bit: with |
There was a problem hiding this comment.
Could this also be an issue in triton_unified_attention.py?
There was a problem hiding this comment.
I checked triton_unified_attention.py on MI300X with our pinned Triton (3.8.0+git669b31ac), and it doesn't hit this. I pre-filled the output with NaN and swept head size 64/128/256, GQA 1/4/8, page size 16 to 256, decode and mixed batches, 2D and 3D (so reduce_segments too), fp8 out, TD on/off, forced tiles, and softcap/sinks/sliding window/ALiBi. Every output value was written, and the results match AMDGCN_USE_BUFFER_OPS=0 bit for bit. The same harness reproduces the drops in paged decode on main.
The LLVM IR for the i64 head offset is correct; the gfx942 ISA ignores output_stride_1 and multiplies the head index by an unrelated spilled SGPR in high-spill variants. The 32-bit per-head offset avoids it. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com>
|
/ci run |
|
✅ Triggered Buildkite CI #96953 for commit |
simondanielsson
left a comment
simondanielsson
left a comment
There was a problem hiding this comment.
LGTM
I think generally in the future when we review triton code we have to be careful about proper handling of 32bit offsets and lowering to buffer load/stores, seems a bit brittle
Both accuracy bugs in ROCM_ATTN's Triton fallback that this avoided are fixed on main: the sliding-window NaN (vllm-project#62828) no longer occurs since vllm-project#60767 nulls below-window block-table entries, and the gfx942 unwritten output heads were fixed by vllm-project#63268. What remains is a performance change, which belongs in vllm-project#61433. Keep only the gfx950 rule for vllm-project#62606. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Chen, Todd <zhenchen@amd.com>
… on gfx942 (vllm-project#63268) Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: vllm-agent <claw@inferact.ai>
chunked_prefill_paged_decode.py: offset the output pointer to the token's row as a 64-bit scalar and keep the per-head store offsets 32-bit.test_prefix_prefill.py: fill the output with NaN, so a row the kernel skips fails every time instead of passing on leftover memory.test_qwen3_nonstandard_block_size[True-chunked_prefill_paged_decode-...-8-4-256-1056], added with the interleaved V cache in #42755, fails on MI300 intermittently (build 96693). Locally it fails 6 of 20 runs on main. The Triton decode kernel only writes 263 of the decode token's 2048 output values, so the result depends on whatever memory the output buffer reuses. The kernel's logic is correct: underTRITON_INTERPRET=1, or withAMDGCN_USE_BUFFER_OPS=0, every value is written. The LLVM IR computes the 64-bit head offset correctly, but in the gfx942 ISA theoutput_stride_1argument is never read and the head index is multiplied by an unrelated spilled SGPR instead (the block-table row offset, in the variant traced). Rows land at the wrong offsets, so most values are left unwritten and some heads overwrite others. This only happens in variants with hundreds of SGPR spills; computing the per-head offset in 32 bits avoids it, so the change is a workaround for a backend miscompile. The interleaved V path is not the only one affected: the normal layout loses outputs too, on the Triton fallback with 128 and 256 token pages at head size 128 and 256, which hybrid models with large pages can hit. On a decode-only sweep of 60 head-size, GQA, page-size and layout combinations, 18 dropped stores before this change and none after; kernel time changes by -0.4% to +1.4%.Tests:
pytest kernels/attention/test_prefix_prefill.py10/10 (205 passed each; the updated test fails 3/3 without the kernel change), andpytest v1/attention/test_attention_backends.py v1/attention/test_rocm_attention_backends_selection.py(201 passed).AI assistance (Claude Code) was used to draft this change. Every line was reviewed by the submitter.