[yhl-amd]

Summary

DeepSeek-V4.1-Flash does not start on gfx1250 (MI455X) because of two AITER gaps (1, 2). Once those are fixed it runs, but GSM8K is about 3 points below MI355X because the gfx1250 grouped MoE GEMM drops part of K (3). This PR fixes all three. The ATOM side is ROCm/ATOM#2617, which depends on this PR.

1. deepgemm_fp8_paged_mqa_logits: 8-token pages on gfx1250

Why it failed. V4.1's FP8 index plane stores 8 rows per page (candidate_block_size: 8 in the checkpoint config). This is model semantics, so the page size cannot be changed. On gfx1250 the preshuffle path TDM-loads one page per iteration and uses it directly as the WMMA B tile, which is 16 tokens wide. So any page below 16 was rejected:

aiter/ops/triton/attention/pa_mqa_logits.py:660
AssertionError: gfx1250 preshuffle (TDM block-load) needs KVBlockSize>=16; got 8.

gfx950 already supports 8-token pages (two pages per MFMA tile). This brings gfx1250 to the same contract.

What changed (aiter/ops/triton/gluon/pa_mqa_logits.py, gfx1250 branch of _gluon_deepgemm_fp8_paged_mqa_logits_preshuffle only):

• A 16-token tile is built from PagesPerTile = 16 // KVBlockSize pages. Each page is its own TDM load, for both values and scales, into adjacent shared-memory slots.
• The tile is read through reinterpret(shape=[NUM_BUFFERS, 1, TileRows * HiddenDim]). Each double buffer's pages are therefore one contiguous view, which is reshaped in the page's (8, 16) shuffle layout (shuffle_weight(kv, layout=(min(page, 16), 16))). Scales are read the same way.
• The async_wait counts scale with the 2 * PagesPerTile loads each tile issues: values first, then scales.
• Page ids are still read one tile ahead, as before. Without that, pages of 16 and 64 regressed by 14-20% at B=64.
• For pages of 16 and up, PagesPerTile == 1, so the loop issues the same loads in the same order as before.
• The KVBlockSize >= 16 assert in the wrapper is removed. The compile path's own gfx1250 check (KVBlockSize > 1 and ChunkK % KVBlockSize == 0) still applies.

Two Gluon details shaped this design, in case they matter for other gfx1250 kernels:

• A TDM async_load into a smem.slice(...) view loses the slice's offset: both pages landed in slot 0. Loading into smem.index(i) is correct.
• slice needs a constexpr start, and slice followed by reshape on this allocation crashed LLVM. Hence the reinterpret view.

2. gemm_a8w8_blockscale_group32: allow gfx1250

The entry point asserted get_gfx_runtime() == "gfx950". The kernel it calls, the shared AFP8WFP8 Triton kernel (tl.dot_scaled), is correct on gfx1250, so the gate now admits gfx1250 as well. V4.1's dense FP8 projections (32x32 / 1x32 E8M0 scales) all go through this entry point.

3. Grouped MoE (gfx1250 TDM): fit tile_k to each stage's K

Why it was wrong. The gfx1250 grouped MoE GEMM (grouped_gemm_gfx1250_a8w4, which is the default MoE path on gfx1250) walks K in whole tile_k tiles and silently drops any remainder. gemm2's K is inter_dim. DeepSeek-V4.1-Flash at TP4 has inter_dim = 2560 / 4 = 640 per rank, and with the default tile_k = 256 every routed expert lost 128 of its 640 K columns. Nothing errored and the server answered sensibly, but:

• Each layer's MoE output on se3 was 10-25% off MI355X, and its norm 3-7% low, even with bit-identical inputs and identical expert selection. Routing matched to 0.3% on logits and 0.1% on weights.
• GSM8K 3-shot (full set) scored 0.892 on se3, against 0.924 for the same ATOM/AITER source, serve flags and lm_eval command on MI355X.

What changed (aiter/ops/flydsl/grouped_moe_gfx1250.py): _grouped_a8w4_tdm_moe now halves each stage's tile_k until it divides that stage's K into at least two tiles. gemm1 is checked against model_dim, gemm2 against inter_dim. Below 128 (one 4-scale WMMA step) it raises instead of computing a wrong result. The e8m0 scale preshuffle keeps k order for any 128-multiple tile, so weights prepared for 256 are read correctly at 128. Shapes that already divide keep their tile, so no other model changes.

Results on se3.

• New test: test_grouped_inter_dim_not_tile_k_multiple (a4w4 Silu, inter_dim 384 and 640; a4w4 runs the same TDM K loop as a8w4) gives logits_diff 0.19 / 0.08 and output norm 0.83 / 0.93 of the fp32 reference before the fix, and 0.0000 / 1.000 after.
• DeepSeek-V4.1-Flash TP4, GSM8K 3-shot full set: 0.8916 → 0.9280 (MI355X: 0.9242). Decode TPOT is 8.8 ms (it was 9.2 ms before).

Accuracy

On gfx1250 (MI455X, se3), with this container: triton 3.8.0.dev20260903+git9551a57b, flydsl 0.3.3.dev886.

Paged MQA logits.

• test_paged_mqa_logits_preshuffle is now enabled on gfx1250: 44 passed, 4 skipped. The skipped cases are page 256 with ChunkK 128, which the gfx1250 compile path does not admit. Coverage:
  • pages 8 / 16 / 32 / 64 / 128 / 256;
  • ChunkK 128 and 256;
  • compact and padded tables;
  • decode, and MTP with next_n=3 and an empty context.
• A separate sweep against a torch reference with a randomly permuted page table: pages 8 / 16 / 64, next_n 1 and 2, nine context lengths from 2 to 8192. Max relative error 1.25e-4.
• The same sweep was run on the unmodified kernel at pages 16 / 64 to validate the reference itself; it agrees there too.

group32 GEMM. The group32 entry-point tests in test_gemm_afp8wfp8.py now run on gfx1250, and all 64 pass. They cover panel scales and tails, extreme scale codes, CUDA-graph replay with live inputs, row scales with split-K tails, the public dispatch, raw views, preallocated output, and the invalid contract. test_tuned_variants_use_production_config stays gfx950-only, because no tuned group32 configs exist for gfx1250 yet.

End to end. With this PR and the ATOM PR, DeepSeek-V4.1-Flash serves on gfx1250: TP4, --level 0, BF16 KV, FP8 index.

• Smoke tests answer correctly, and TPOT is about 9.2 ms at concurrency 1.
• GSM8K 3-shot on the first 200 questions (lm_eval local-completions) is 0.87 at concurrency 32 and 0.855 at concurrency 1 (+/- 0.025).
• The ATOM recipe lists 0.92-0.945 on MI355X. That figure uses DSpark, and its exact command is not recorded, so it is not a like-for-like baseline. A same-command MI355X run is still to do. The failing answers are coherent reasoning errors, not garbage.

Performance

deepgemm_fp8_paged_mqa_logits(Preshuffle=True) on gfx1250, decode, 64 heads, D=128, mean of 50 calls:

page	B	context	before	after
8	1 / 16 / 64	8k / 8k / 32k	assert	19.2 / 18.0 / 45.1 us
16	1 / 16 / 64	8k / 8k / 32k	19.0 / 18.7 / 37.2 us	18.0 / 18.7 / 37.6 us
64	1 / 16 / 64	8k / 8k / 32k	17.6 / 17.6 / 31.6 us	17.6 / 18.0* / 31.3 us

* This cell reads 18-31 us depending on what ran earlier in the same process. Repeated isolated runs give 17.8-18.9 us patched and 18.2-18.6 us original.

Test plan

• op_tests/test_pa_mqa_logits_offset.py on gfx1250: 61 passed, 64 skipped. The skips are pre-existing gfx1250 skips plus the 4 above.
• op_tests/triton_tests/gemm/basic/test_gemm_afp8wfp8.py on gfx1250: 714 passed, 125 skipped (650 / 189 before this PR).
  • Note: the existing test_gemm_afp8wfp8_preshuffle[...-gluon-...] cases sometimes abort the process on gfx1250. This also happens with the unmodified test file, and those tests are not touched here.
• black and ruff check are clean on the changed files.
• op_tests/flydsl_tests/test_flydsl_grouped_gemm.py on gfx1250: 24 passed, 3 failed (11 failed, 16 passed before). The 3 fail before and after this change: missing AOT cache in run-only mode, and the a8w4 situv2 reference.
• DeepSeek-V4.1-Flash TP4 on gfx1250, GSM8K 3-shot full set: 0.9280 (MI355X with the same code and command: 0.9242).
• gfx950: the gfx950 code paths are not modified. CI should confirm.

🤖 Generated with Claude Code

[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 (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:windows	Windows JIT op smoke test on the gfx1201 ROCm runner
ci:multi-gpu	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: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
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 6626 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[Copilot]

🟡 Changes recommended

An accepted small-chunk configuration permits overlapping writes that can produce incorrect logits.

1 open finding

• Prevent overlapping split-KV writes when ChunkK is misaligned · New

What changed in this PR

This PR enables gfx1250 support needed by DeepSeek-V4.1-Flash for paged MQA logits and group32 FP8 GEMM.

Changes:

• Assemble 8-token pages into 16-token WMMA tiles.
• Allow group32 FP8 GEMM on gfx1250.
• Enable related tests on gfx1250.

File	Description
op_tests/​triton_tests/​gemm/​basic/​test_gemm_afp8wfp8.py	Enables gfx1250 group32 tests.
op_tests/​test_pa_mqa_logits_offset.py	Enables gfx1250 preshuffle tests.
aiter/​ops/​triton/​gluon/​pa_mqa_logits.py	Adds multi-page tile assembly.
aiter/​ops/​triton/​gemm/​basic/​gemm_a8w8_blockscale_group32.py	Allows gfx1250 dispatch.
aiter/​ops/​triton/​attention/​pa_mqa_logits.py	Removes the short-page restriction.

🧠 Review effort: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[azaidy]

Please create different PRs for different backends and different kernels.