[Triton/Gluon] [MI350] Optimize FP8 MQA logits kernel - #5216
Merged
Merged
Conversation
The gfx950 gluon path disabled buffer stores whenever the whole fp32 logits tensor exceeded 2 GiB. That bound does not describe what buffer ops actually require here: the kernel bakes the row and the tile into the base pointer (logits_ptr + row_id * stride_logits_s, then += BLOCK_KV * stride_logits_k per tile) and rebuilds the descriptor from it, so the i32 buffer offset only ever spans one tile. What does bind is that the row base is computed in i32 *elements*, so the largest index the kernel forms must fit in an int32. Measured on gfx950: 16384 x 131072 (max index 2**31 - 1, 8 GiB) is correct, 32768 x 131072 (2**32) faults. Falling off buffer stores was not just slower, it did not build. A masked global store lowers to a branch, while the masked buffer store is branch-free (select mask ? offset : -2**31). With BLOCK_M=2 the extra join blocks land inside the double-buffered async-copy region and trip an assert in LLVM's SIInsertWaitcnts, WaitcntBrackets::mergeAsyncMarks(): Assertion `Begin <= End && "Begin must be less or equal to End."' failed. llvm/ADT/Sequence.h when one predecessor has pending llvm.amdgcn.asyncmark and the other has none, so MergeCount == 0 and seq_inclusive<unsigned>(1, 0) fires. Fixed upstream by llvm/llvm-project#193499 (81d618b6bc1e); the LLVM Triton 3.8 pins carries it, 3.7's does not. BLOCK_M=2 is therefore gated on buffer stores being available OR Triton >= 3.8, which post-fix is only reachable past 2**31 logits elements. Verified on MI355X / gfx950, Triton 3.7.0+amd.rocm7.2.0 and 3.8.0 ToT: all eight GLM-5.2 / DSV4 prefill configs compile (GLM-5.2 2x8kx32k aborted before), cosine diff vs a torch reference ~1e-10 at 4, 8 and 16 GiB, and the >2 GiB shapes keep buffer stores: DSV4 2x8kx32k 879 -> 1359 TFLOP/s. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Profiling the shipped config on MI355X (rocprofv3 counters + ISA) shows the kernel is not MFMA bound. For GLM-5.2 1x8kx8k, MFMA is 2.7% of issued instructions while VALU is 63.5% and SALU 25.5%, with 35% of wave cycles stalled. The compiled inner loop is 16 v_mfma against 140 v_accvgpr_read/write (AGPR<->VGPR shuffles), 128 v_maximum3_f32 (the relu in the head fold) and 16 scratch_store spills. The MFMA count matches the arithmetic exactly, so the matrix work is waste-free; the cost is the per-head reduction epilogue and the register pressure it creates. waves_per_eu = 4 pins the kernel to 128 VGPRs (gfx950's 512-entry unified file divided by 4) and the reduction does not fit, which is where both the AGPR traffic and the spills come from. A 48-combo sweep per model-shape over waves_per_eu x NUM_CHAINS x (BLOCK_KV, NUM_WARPS) picks the lowest-pressure point instead of the highest occupancy: waves_per_eu 4 -> 3 for <=32 heads, 4 otherwise num_chains 4 -> 2 num_warps 2 if heads<=32 else 1 -> 2 block_kv 64 if heads<=32 else 32 -> 64 NUM_CHAINS=4 was a net loss everywhere: the parallel FMA chains cost more in live registers than they save in dependency depth. DSV4's narrow 32-wide, 1-warp tile was the larger miss and accounts for most of its gain. Kernel-only GPU time (torch.profiler, kernel-name filtered), shipped vs tuned, interleaved in one session on MI355X / Triton 3.7: glm5.2 1x8kx8k 1253 -> 1562 TFLOP/s 1.25x glm5.2 2x8kx8k 1334 -> 1584 TFLOP/s 1.19x glm5.2 1x8kx32k 1354 -> 1548 TFLOP/s 1.14x glm5.2 2x8kx32k (did not build) -> 1790 TFLOP/s dsv4 1x8kx8k 1193 -> 1721 TFLOP/s 1.44x dsv4 2x8kx8k 1168 -> 1697 TFLOP/s 1.45x dsv4 1x8kx32k 1315 -> 1758 TFLOP/s 1.34x dsv4 2x8kx32k 883 -> 1837 TFLOP/s 2.08x All eight configs verified against a torch reference, worst cosine diff 7e-10. Tuned on prefill shapes (1x8kx8k, 4x8kx8k) at 32 and 64 heads, Triton 3.7, gfx950 only. Decode shapes, other head counts and gfx942 are not covered by this sweep and the space is not smooth -- BLOCK_KV=128 was bimodal, 0.39x in one combo -- so a wider validation pass is warranted before relying on these defaults elsewhere. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…x950 The kernel is VALU bound, not MFMA bound. Measured on MI355X: the matrix pipe runs ~40% busy against a 4.4 PFLOP/s ceiling, and a probe that adds 16 packed-FMA slots per KV tile costs a flat ~14% of runtime, so roughly two thirds of the time tracks the 76 VALU-class ops in the inner loop. Of those only 48 are irreducible (32 relu + 16 packed FMA); the rest is the per-head reduction epilogue and the register pressure it creates. Three changes, worth 10-15% on 64-head shapes: * MFMA_NONK_DIM 16 -> 32 wherever num_heads >= 32. The 16x16x128 output layout splits the head axis into 4 register bits and 2 lane bits, so the head sum needs two cross-lane butterfly steps (13 VALU ops, 6 of them v_mov because permlane*_swap is destructive). The 32x32x64 layout leaves one lane bit, so one step. Inner-loop VALU drops 70 -> 55. * M_CHUNK: fold each head chunk as its MFMA retires instead of materialising the whole score tile first. Accumulator liveness drops from 32 registers to 16, which is what keeps the wider 32x32 tile off the spill path -- without it it spills 89 VGPRs and loses 15%. The partial sums stay in the pre-cross-lane layout so the lane-crossing step still runs once per KV tile. * NUM_CHAINS 2 -> 1 when chunked, waves_per_eu 4 -> 3 (170 VGPRs, no spills; 4 spills and 5+ collapses), and a 2x unroll of the KV walk. Kernel-only TFLOP/s, no -inf prefill, best of 2 interleaved reps: dsv4 1x4kx4k 1639 -> 1879 1.15x dsv4 1x8kx8k 1589 -> 1750 1.10x dsv4 2x8kx8k 1650 -> 1856 1.13x dsv4 4x8kx8k 1733 -> 1949 1.13x dsv4 1x8kx32k 1837 -> 2062 1.12x glm5.2 (32 heads, already 32x32 and single-M-tile) +1-4% from the unroll All 8 compile-matrix configs still build, aiter's 144 unit tests pass, and both models check out against the torch reference at ~1e-10 cosine diff. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Two independent launcher-gated changes to the gfx950 fp8_mqa_logits kernel,
both found by profiling it with an ATT capture and hardware counters. The
kernel is matrix-core bound -- the MFMA runs at the full CDNA4 FP8 rate and
wastes nothing on tiling -- but the matrix core sits idle 37% of the time on
64 q-heads and 52% on 32, behind s_barrier, s_waitcnt and MFMA-hazard s_nop.
Both changes buy back some of that idle.
RELAXED_STORE (32 q-heads). With BLOCK_M=2 the loop walks the union of both
rows' KV ranges, so every store is masked to the part its row owns: 14 VALU
and 8 SALU per two KV tiles, on every tile, when the mask is all-true on all
but the last one or two. clean_logits=False already tells the caller that
out-of-window positions are theirs to fill or ignore, so there the mask
protects nothing and the multi-row store can collapse to the single-row one.
The union bound stays -- that is what keeps the store inside the row.
Docstring change: clean_logits=False now says those positions are
*unspecified* rather than untouched. A caller that wants -inf there must
fill it in after the call, not before. Both existing test suites that use
clean_logits=False already rehydrate that region themselves.
num_warps=1 + BLOCK_KV=32 (64 q-heads). A one-wave workgroup emits no
s_barrier at all. At num_warps=2 the async copy hands warp w the odd/even KV
columns while the MFMA layout has it consume a contiguous half, so the two
genuinely alias and Triton has to synchronise every tile -- 8.9% of wave
time, 100% stall. Halving BLOCK_KV keeps the per-wave work identical. The
cost is half as many waves, so it is gated on seq_len > 4096: measured 0.98x
at 4096 and 1.02-1.06x from 8192 up. BLOCK_M=2 halves the grid again, so
32 q-heads stays out.
Measured, warmup 300 then the median of five 500-iteration windows, arms
interleaved (the part is power limited, so short best-of-N windows are not
comparable):
64 q-heads 32 q-heads
1x4kx4k 1.000x (gates off) 1.001x (gates off)
1x8kx8k 1.026x 1.045x
2x8kx8k 1.046x 1.044x
4x8kx8k 1.021x 1.034x
1x8kx32k 1.058x 1.044x
2x8kx32k 1.032x 1.043x
Peak 2314 TF/s and 1967 TF/s, from 2242 and 1887. Against FlyDSL PR #4538 the
kernel goes from behind at two cells and level at three more to ahead
everywhere the gates fire: 1.04-1.15x at 64 q-heads, 1.01-1.13x at 32.
Verification: clean_logits=True compiles to a byte-identical instruction
stream to before (1287 instructions, only DWARF line numbers move);
clean_logits=False is bitwise identical in-window to the masked kernel and
deterministic run to run at six varlen layouts, including two where a
BLOCK_M=2 block straddles a batch boundary; 149/149 unit tests pass.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Conflicts were all in fp8_mqa_logits, where both sides had moved: - launcher: main replaced the byte-based buffer-op gate with an int32 element offset bound (#5121) and gated BLOCK_M=2 on use_buffer_store. Took main's bound, since it is the better general rule and the branch only overrode the result for gfx950; kept the branch's gfx950 tuning block, which supersedes main's. The BLOCK_M=2 guard is vacuous on gfx950 -- the rebasing makes buffer stores unconditional there -- so it is noted rather than applied. Note the auto-merge left both gate computations in place, with the old byte rule silently overwriting the new one; the old block is removed. - gluon kernel: kept main's config-aware repr (#5097) and added the branch's M_CHUNK, UNROLL and RELAXED_STORE to its field list. - tests: took main's file as-is. Its test_fp8_mqa_logits_logits_past_2gib covers the logits-past-2-GiB case; the branch's large_kv, long_query and kv_over_2gib tests are dropped. 146 tests pass; black and ruff clean; the gfx950 launcher still picks num_warps=1/BLOCK_KV=32 only at block_m == 1 and seq_len > 4096, and RELAXED_STORE only when clean_logits is False. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
cagrikymk
force-pushed
the
cagri/opt_mqa_logits
branch
from
September 2, 2026 21:36
dfe89ad to
15602d3
Compare
cagrikymk
marked this pull request as ready for review
September 3, 2026 00:24
github-actions
Bot
changed the title
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Hardcoded tuning must move to shared configuration, and the relaxed multi-row store path needs targeted coverage.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Optimizes the gfx950 FP8 MQA logits path for long-context workloads.
Changes:
- Rebases pointers to support large tensors with buffer operations.
- Adds chunked head reduction, MFMA selection, loop unrolling, and relaxed stores.
- Retunes launch parameters for Triton 3.8.
File summaries
| File | Description |
|---|---|
aiter/ops/triton/attention/fp8_mqa_logits.py |
Updates gfx950 dispatch and tuning heuristics. |
aiter/ops/triton/_gluon_kernels/gfx950/attention/fp8_mqa_logits.py |
Implements optimized loading, reduction, looping, and storage. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 3
- Review effort level: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
vgokhale
approved these changes
Sep 9, 2026
5 tasks
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.
Motivation
This PR optimizes FP8 MQA logits gluon kernel on MI355, both DSV4 and GLM5.2. Also, it resolves buffer/global store/load issues by changing how base pointer calculated so that the offsets are within the limit of buffer ops.
Other aspects of the kernels also improved, reduction, mfma op selection logic, loop unrolling to better compensate loop overhead, relaxing store masking when clean logits past boundaries are not needed.
The data is collected with Triton 3.7.0+amd.rocm7.2.0.git89002410.
Technical Details
dsv4 — 64 q-heads x 128
glm5.2 — 32 q-heads x 128
1x8kx8k, TF/s, with some retuning for newly released 3.8:Test Plan
Existing tests pass, also stressed tested with 16k queries and 1M kv length to verify buffer op solution wors.
Test Result
All tests pass.
Includes certain ideas from #5048