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

[jiejingzhangamd]

What breaks today

Any prefill whose fp32 logits tensor passes 2 GiB aborts the AMDGCN backend at JIT time:

llvm/include/llvm/ADT/Sequence.h:275: llvm::iota_range<unsigned int>::iota_range(T, T, bool)
  [T = unsigned int]: Assertion `Begin <= End && "Begin must be less or equal to End."' failed.

Two independent decisions collide. BLOCK_M=2 is selected whenever seq_len > 4096, but it only compiles on the buffer-store path. The buffer-store gate switches off at 2 GiB. Past that size we therefore ask for BLOCK_M=2 with plain stores, which is exactly the combination that does not compile.

For GLM-5.x this is not a corner case: it is every chunked prefill of 8192 tokens against a context beyond 65,536.

Why the gate was wrong

The gate counted bytes:

BUFFER_LIMIT_BYTES = 2 * 1024 * 1024 * 1024
use_buffer_store = logits.numel() * logits.element_size() < BUFFER_LIMIT_BYTES

Buffer ops do address through a 32-bit offset, but the kernel re-bases logits_ptr per row and again per KV tile before every store, so that offset never has to span the tensor. What actually has to fit in int32 is the largest element offset the kernel forms, because the row strides stay 32-bit on the buffer path and are only widened to int64 on the fallback (_gluon_fp8_mqa_logits_kernel, the if not USE_BUFFER_STORE: stride_logits_s.to(gl.int64) branch).

Counting bytes where the constraint is elements made the limit 4x too tight for an fp32 output.

The new bound is exact, not conservative

Measured on MI355X, 32 heads x 128 head_dim, against the plain-store path. Output is bit-identical in every passing case (max relative error 0.00e+00, full-tensor comparison, not sampled):

s_q	s_k	logits	plain store	buffer store	speedup
8192	65536	2.00 GiB	3.557 ms	2.941 ms	1.21x
8192	95457	2.91 GiB	6.105 ms	4.577 ms	1.33x
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 boundary predicted by the new expression matches the hardware to the element. 8192x262144 and 16384x131072 both place the largest offset at exactly INT32_MAX and are bit-correct; 16384x139264 (1.06 x 2^31) and 32768x131072 (2^32) both fault with a memory access fault. So the fix is not "raise the limit and hope" -- the limit is now the real one.

End to end

A 21-layer GLM-5.x indexer prefill of 9,695 new tokens against a 96,960-token context (90% prefix cache hit) at chunk 8192, 3 repeats, spread under 1%:

	mean
before	150.23 ms
after	118.36 ms

21.2% off the indexer stage, with no serving-config change. For reference, working around the cliff by dropping --chunked-prefill-size to 5120 measured 119.08 ms -- the same win, but it has to be retuned per context length. At ISL 65,536 the default 8192 is already optimal, and at 131,072 the best chunk moves again; the kernel-side fix removes the whole tuning problem.

Tests

The existing cases top out at s_q=1024, s_k=1560, four orders of magnitude below the gate, which is why this was never caught. Added 8192x65664 and 8192x98304, which crash the process on the current code and pass here. They skip themselves if the device lacks the free memory, and they compare row-wise against a per-row reference because the existing ref_fp8_mqa_logits materializes [num_heads, s_q, s_k] -- hundreds of GB at these shapes.

146 passed on gfx950 (144 existing + 2 new). ruff 0.16.0 format and check both clean.

Note on use_buffer_load

The KV gate is corrected in the same way for consistency. It never binds in practice -- KV would have to reach 2 GiB, i.e. ~16M tokens -- so this half is reasoning, not measurement.

---

Note on base. Retargeted from upstream-sync to main and rebased onto it. The fix has no meaning on main by itself: main predates BLOCK_M in this kernel, so the crash it prevents cannot happen there and there is no block_m condition to gate. The PR therefore also carries upstream's ROCm#4563, the commit that introduces BLOCK_M, plus the two-line block_m = 1 for the non-gfx950 branch from ROCm#4246 — ROCm#4563 routes the launch grid through block_m but only defines it under arch == "gfx950", so on its own it leaves every other arch raising NameError.

Not carried: ROCm#4508's strip_annotate (a Python 3.14 compatibility change to the same gluon kernel), which is unrelated to BLOCK_M.

The dispatcher and test files end up byte-identical to what was reviewed on upstream-sync. The previous head is kept at backup/mqa-logits-buffer-store-gate-20260827.