[Triton/Gluon] fp8_mqa_logits: gate buffer ops on the int32 offset, not tensor bytes

[xiaobochen-amd]

Problem

fp8_mqa_logits crashes the process on prefill shapes whose fp32 logits exceed
2 GiB. BLOCK_M=2 is selected for seq_len > 4096 but only compiles on the
buffer-store path, while the buffer-store gate switches off at 2 GiB — the two
combine into an AMDGCN backend abort at JIT time (Sequence.h:275 "Begin must be less or equal to End"). For GLM-5.x this is any chunked prefill of 8192 tokens
against a context past 65,536.

Fix

The gate had the wrong unit. Buffer ops address through a 32-bit offset, but the
kernel re-bases the pointer per row and per KV tile, so that offset never has to
span the tensor. What must fit in int32 is the largest element offset the
kernel forms. Counting bytes made the limit 4x too tight for an fp32 output.

Results

MI355X, 32 heads x 128 head_dim, against the plain-store path. Bit-identical
output in every case (max rel err 0.00e+00):

s_q	s_k	logits	plain store	buffer store	speedup
8192	65536	2.00 GiB	3.557 ms	2.941 ms	1.21x
8192	131072	4.00 GiB	9.812 ms	6.489 ms	1.51x
16384	131072	8.00 GiB	21.738 ms	12.879 ms	1.69x
8192	262144	8.00 GiB	25.710 ms	14.184 ms	1.81x

The new boundary is exact: 8192x262144 and 16384x131072 both place the largest
offset at exactly INT32_MAX and are bit-correct, while 16384x139264 and
32768x131072 fault.

End to end, a 21-layer GLM-5.x indexer prefill of 9,695 new tokens against a
96,960-token context at chunk 8192 drops from 150.2 ms to 118.4 ms.

Test

Existing cases top out at s_q=1024, s_k=1560, four orders of magnitude below
the gate, which is why nothing caught this. Added 8192x65664 and 8192x98304,
which core-dump on current main. The reference is computed per row so s_k can
be large, and the test skips when free VRAM is short.

op_tests/triton_tests/attention/test_fp8_mqa_logits.py: 146 passed.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5121 --add-label <label>

PR title tags:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf] and op tags like [MLA] are left untouched. Add the no-auto-title label to opt this PR out of title tagging.

[Copilot]

Pull request overview

Fixes a crash in the fp8_mqa_logits Gluon path for very large prefill shapes by correcting the buffer-op gating logic to use the maximum int32 element offset formed by the kernel (rather than tensor byte size), and adds regression tests that exercise logits tensors > 2 GiB.

Changes:

• Update Gluon buffer load/store gating to check max int32 element offsets derived from strides and extents.
• Gate BLOCK_M=2 selection on use_buffer_store to avoid AMDGCN backend JIT aborts on the plain-store path.
• Add large-shape regression tests with a per-row reference implementation to avoid materializing huge [num_heads, s_q, s_k] tensors.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
aiter/ops/triton/attention/fp8_mqa_logits.py	Corrects buffer-op gating units (element offsets) and prevents BLOCK_M=2 from being selected when buffer stores are disabled.
op_tests/triton_tests/attention/test_fp8_mqa_logits.py	Adds regression coverage for logits tensors exceeding 2 GiB using a per-row reference to keep reference memory bounded.

Suppressed comments (1)

op_tests/triton_tests/attention/test_fp8_mqa_logits.py:206

• Precompute kv in float32 (and transpose once) before looping over sampled rows, then pass it into the row reference helper to avoid repeated large casts/transposes.

    # Sample rows across the grid: first, last, and the BLOCK_M=2 block seam.
    for i in (0, 1, s_q // 2, s_q // 2 + 1, s_q - 1):
        ref_row = ref_fp8_mqa_logits_row(q[i], kv, weights[i], int(ks[i]), int(ke[i]))
        got_row = logits[i]

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.