[nsakkine]

Summary

Fixes the MI350X 8-GPU all-to-all (A2A) CI failure that appeared after the mha_v4-sparse-geometry merge:

MX Q descale needs 38912 mapped bytes … but only 37888 are backed

The merge put the Q-scale backing check in the validator that dense and sparse launches share. The dense objects bound their Q-scale gather with a buffer descriptor, so they accept a Q scale at its logical size. The A2A op returns exactly that, since its Q scales are unpadded since #6841. This PR moves the check into fmha_v4_fwd_sparse, so only the gfx950 sorted-sparse path enforces it.

This is a workaround, not the full fix. The gfx950 sorted-sparse objects still fetch MX scales with unguarded global loads, so they still need Q scales backed to a whole 256-row query tile. The 64x64 MX-V rows also need V-scale slack. As a result, the sparse path is not compatible with A2A output. Passing A2A packed Q/K/V and scales to mha_v4_packed with a block-sparse LUT fails with "MX Q descale needs …". It is a clear error, not a fault or a wrong result. The planned follow-up rebuilds the sparse objects with bounded scale gathers, as the dense ones already have. That makes dense and sparse accept the same scales, removes the sparse-only checks, and removes the gfx950 sparse exclusion in test_mha_v4_dispatchable_objects_bound_their_scale_gathers.

PR Self-Review

Reviewed revision

a77e2fed5b4812e7ecb3d7998d4f9535b83f9b98

Behavior and intent

• csrc/py_itfs_cu/asm_mha_v4_fwd.cu:
  • The Q check_scale_backing_storage call moves out of validate_packed_mha_v4, which both paths share, into fmha_v4_fwd_sparse after find_config. It runs only for E8M0 per-1x32 Q.
  • The MX-V slack check is unchanged.
• Dense: logical-size Q scales are accepted again, as before the sparse-geometry merge.
• Sparse: unchanged. It still requires Q backed to MHA_V4_QUERY_TILE_ROWS (256) rows, which the aiter.ops.mha_v4_quant producers already reserve.
• Not changed: the A2A op, the producers and the kernels.

Validation and unverified scope

All runs were on MI355X (gfx950).

• op_tests/multigpu_tests/test_attention_a2a.py -w 2 4 8 --heads 8 40 (the CI command):
  • Before this change, it reproduces CI's exact error in the dense consumer call at world size 4. With 8 ranks, every default sequence is already a multiple of 256 rows, so 8-rank-only runs pass.
  • With this change: rows=273, checks=14686, failures=0.
• Script mode, as in CI: op_tests/test_mha_v4.py ran 328 checks and op_tests/test_mha_v4_sparse.py ran 290, with no failures.
• A dense call with logical-size MXFP6 Q scales gives output bitwise equal to a call with padded scales.
• Not verified: gfx942, which has no hardware here. The change is in arch-independent launcher code, and the gfx942 objects already bound their gathers.
• Known limitation, by design: sparse launches with A2A output (see above).

PR-specific question

Is a clear error for "A2A output fed into a sparse launch" acceptable until the sparse objects are rebuilt with bounded scale gathers? The alternative is to pad Q and mxfp6_p V scales inside the A2A op now. That works, but the kernel rebuild would make it unnecessary.

Changed op tests

• op_tests/test_mha_v4.py:
  • test_mha_v4_packed_accepts_mxfp8_scale_recipe: back to a logical-size Q scale. Dense must accept it; that is the regression this PR fixes.
  • New test_mha_v4_packed_sparse_mx_q_scales_need_the_whole_query_tile: the sparse launch must reject a logical-size MX Q scale with "MX Q descale needs", and must run once the storage backs the 256-row tile. This pins down the sparse requirement that the follow-up rebuild removes.

[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 Radeon R9700 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 7651 --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.

Reviews:
A reviewer is usually assigned within one working day after the PR is ready for review (see Code Review Process). If nobody has responded for about a week, mention the area contact listed in CONTACTS.md. Want to help review? See REVIEWERS.md.

[aiter-bot]

READY FOR HUMAN REVIEW

PR Self-Review

Revision a77e2fed5b48 · mode shadow

Author self-review accepted. Readiness report:

Author answers

• PASS — current-SHA answers are present and structurally specific

Show the author's self-review answers

Behavior and intent

- `csrc/py_itfs_cu/asm_mha_v4_fwd.cu`:
  - The Q `check_scale_backing_storage` call moves out of `validate_packed_mha_v4`, which both paths share, into `fmha_v4_fwd_sparse` after `find_config`. It runs only for E8M0 per-1x32 Q.
  - The MX-V slack check is unchanged.
- Dense: logical-size Q scales are accepted again, as before the sparse-geometry merge.
- Sparse: unchanged. It still requires Q backed to `MHA_V4_QUERY_TILE_ROWS` (256) rows, which the `aiter.ops.mha_v4_quant` producers already reserve.
- Not changed: the A2A op, the producers and the kernels.

Validation and unverified scope

All runs were on MI355X (gfx950).
- `op_tests/multigpu_tests/test_attention_a2a.py -w 2 4 8 --heads 8 40` (the CI command):
  - Before this change, it reproduces CI's exact error in the dense consumer call at world size 4. With 8 ranks, every default sequence is already a multiple of 256 rows, so 8-rank-only runs pass.
  - With this change: rows=273, checks=14686, failures=0.
- Script mode, as in CI: `op_tests/test_mha_v4.py` ran 328 checks and `op_tests/test_mha_v4_sparse.py` ran 290, with no failures.
- A dense call with logical-size MXFP6 Q scales gives output bitwise equal to a call with padded scales.
- Not verified: gfx942, which has no hardware here. The change is in arch-independent launcher code, and the gfx942 objects already bound their gathers.
- Known limitation, by design: sparse launches with A2A output (see above).

PR-specific question

Is a clear error for "A2A output fed into a sparse launch" acceptable until the sparse objects are rebuilt with bounded scale gathers? The alternative is to pad Q and `mxfp6_p` V scales inside the A2A op now. That works, but the kernel rebuild would make it unnecessary.

Changed op tests

- `op_tests/test_mha_v4.py`:
  - `test_mha_v4_packed_accepts_mxfp8_scale_recipe`: back to a logical-size Q scale. Dense must accept it; that is the regression this PR fixes.
  - New `test_mha_v4_packed_sparse_mx_q_scales_need_the_whole_query_tile`: the sparse launch must reject a logical-size MX Q scale with "MX Q descale needs", and must run once the storage backs the 256-row tile. This pins down the sparse requirement that the follow-up rebuild removes.

PR state

• PASS — pull request is marked ready for review

Copilot findings

• PASS — no active Copilot finding lacks a substantive author reply
• ADVISORY — Copilot did not review the current head SHA

Code and op-test hygiene

• PASS — no committed artifact or conflict marker detected

Required CI

• PASS — Black: passed
• PASS — Ruff: passed
• PASS — Repository dependency: passed
• PASS — Aiter Test Gate: passed

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}