Repository navigation
fix(fused_moe): pad MXFP4 A4W4 MoE sort extent to a block_size multiple - #5573
Conversation
The flydsl MXFP4 A4W4 atomic MoE gemm2 kernel hard-faults with a HIP
"illegal memory access" under load. Root cause is a broken block-alignment
invariant between the sort stage and the gemm kernel, not a bug in the
arithmetic:
* gemm2s atomic epilog runs the NON-persistent grid and issues
_issue_all_a_loads() unconditionally for every grid block (incl. the
trailing padding block) BEFORE the `if bx_i32 < bound` guard, relying on
the A buffer descriptor to clamp. The descriptor is sized to
max_m_blocks*BM rows (max_m_blocks = ceil(max_sorted/BM)).
* The A buffer (inter_sorted_quant) is allocated with exactly max_sorted
rows. When max_sorted is not a multiple of BM the descriptor over-reads
ceil(max_sorted/BM)*BM - max_sorted rows and faults on an unmapped page.
_adaptive_moe_sort always rounds max_sorted up to a BM multiple, so the
unconditional load is safe there. PR ROCm#4526 added _aux_uses_opus(), which
diverts the A4W4 atomic path to the Opus sort whose
max_num_tokens_padded = topk_ids.numel() + num_experts*block_size - topk
is NOT a block multiple, breaking the invariant.
Fix: round max_num_tokens_padded up to max_num_m_blocks*block_size in both
Opus/flydsl sort sizing sites, matching what _adaptive_moe_sort maintains.
max_num_m_blocks is unchanged; only the row extent is enlarged.
Validated on Kimi-K2.7-Code-MXFP4 (-tp 4): full gsm8k sweep 1319/1319 clean
(0 illegal-memory-access), accuracy 0.9484 flexible / 0.9477 strict.
Adds op_tests/flydsl_tests/test_mxfp4_a4w4_moe_oob.py, a standalone,
model-free repro with three tests:
* test_moe_sorting_opus_row_extent_block_aligned - drives the production
sort (output_aux=AUX_SORT_OPUS) and asserts the sorted extent is a
block multiple (deterministic; fails pre-fix, passes post-fix).
* test_mxfp4_gemm2_descriptor_fits_allocation - the same invariant as pure
arithmetic, naming the exact 21 rows over-read at the K2.7 fault shape.
* test_mxfp4_a4w4_gemm2_illegal_access - runs the real production kernel
flydsl_mxfp4_gemm2(atomic=True) at the fault shape with the A buffer
against a VMM guard page, reproducing the actual HIP illegal memory
access in a subprocess (SIGABRT pre-fix, clean exit post-fix).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
Run the allocation invariant from the top-level MoE test so standard Aiter CI collects it. Keep the VMM guard-page check behind gfx950, CuPy, and HIP VMM capability checks, and retain a focused CLI entry for local reproduction.
Enforce the block-aligned allocation contract in the existing MoE sorting reference and output comparisons. Add a focused AUX_SORT_OPUS case across supported block sizes, replacing the bespoke VMM and subprocess machinery in test_moe_2stage.
|
Advisory review (static + hand-run; not a merge gate). Validation/Perf ran no GPU stage — reasons are on their lines below. Findings tagged *Rounds the Opus/FlyDSL MoE sort's padded row extent in aiter/fused_moe.py up to a block_size multiple so the MXFP4 A4W4 atomic GEMM2's A-buffer descriptor (sized ceil(rows/BM)BM rows) never spans more rows than the allocation, plus two regression tests in op_tests/test_moe_2stage.py. Review (advisory): 🔴 HIGH RISK 🔴 [verified] aiter/fused_moe.py sits on the serving path of every downstream MoE model and this diff changes its sort output extent for all of them, yet the PR carries no ci:* label, so the Atom/SGLang/vLLM jobs — the only CI that runs Kimi-K2.7-Code-MXFP4 and DeepSeek-R1-MXFP4 over this exact sort path — skip by default and a downstream break would surface only after merge. Author must add ci:all (or at least ci:sglang + ci:vllm) and require it green before merge. |
|
Advisory review (static + hand-run; not a merge gate). Validation/Perf ran no GPU stage — reasons are on their lines below. Findings tagged *Rounds the Opus/FlyDSL MoE sort's padded row extent in aiter/fused_moe.py up to a block_size multiple so the MXFP4 A4W4 atomic GEMM2's A-buffer descriptor (sized ceil(rows/BM)BM rows) never spans more rows than the allocation, plus two regression tests in op_tests/test_moe_2stage.py. Review (advisory): 🔴 HIGH RISK 🔴 [verified] aiter/fused_moe.py sits on the serving path of every downstream MoE model and this diff changes its sort output extent for all of them, yet the PR carries no ci:* label, so the Atom/SGLang/vLLM jobs — the only CI that runs Kimi-K2.7-Code-MXFP4 and DeepSeek-R1-MXFP4 over this exact sort path — skip by default and a downstream break would surface only after merge. Author must add ci:all (or at least ci:sglang + ci:vllm) and require it green before merge. |
…ing fixes, fp8 MQA logits split-k and DSv4 tunings (#5573, #5295, #5558, #5603, #5627, #5485) (#5638) Cherry-picks six already-merged `main` PRs onto `release/v0.1.22` for the `v0.1.22.post1` post release. | PR | `main` commit | Backport commit | What | |---|---|---|---| | #5485 | `9252f4672` | `155534984` | Extend the DeepSeek-V4 a8w8 blockscale GEMM tunings for gfx950 (tuning CSV only) | | #5573 | `22d2c7c91` | `6b23ba866` | Pad the MXFP4 A4W4 MoE sort extent to a block_size multiple (fixes a HIP illegal memory access) | | #5295 | `972c8e1fd` | `7d68b0edb` | Skip invalid expert IDs in MoE sorting | | #5558 | `3fdfca11e` | `dd83a9d17` | Fix MoE routing kernel compile failure | | #5603 | `a84bd368c` | `a96461997` | Add split-k support for fp8 MQA logits on gfx950 | | #5627 | `f5ed7dc54` | `a41214712` | Follow-up to #5603: drop chunking when summation folding is unavailable (fixes Triton 3.6 compile) | Original PRs: - #5485: #5485 - #5573: #5573 - #5295: #5295 - #5558: #5558 - #5603: #5603 - #5627: #5627 To be published as `v0.1.22.post1` once merged (tag on the merge commit, release automation builds the wheel set).
4 participants
The flydsl MXFP4 A4W4 atomic MoE gemm2 kernel hard-faults with a HIP "illegal memory access" under load. Root cause is a broken block-alignment invariant between the sort stage and the gemm kernel, not a bug in the arithmetic:
if bx_i32 < boundguard, relying on the A buffer descriptor to clamp. The descriptor is sized to max_m_blocks*BM rows (max_m_blocks = ceil(max_sorted/BM))._adaptive_moe_sort always rounds max_sorted up to a BM multiple, so the unconditional load is safe there. PR #4526 added _aux_uses_opus(), which diverts the A4W4 atomic path to the Opus sort whose max_num_tokens_padded = topk_ids.numel() + num_experts*block_size - topk is NOT a block multiple, breaking the invariant.
Fix: round max_num_tokens_padded up to max_num_m_blocks*block_size in both Opus/flydsl sort sizing sites, matching what _adaptive_moe_sort maintains. max_num_m_blocks is unchanged; only the row extent is enlarged.
Validated on Kimi-K2.7-Code-MXFP4 (-tp 4): full gsm8k sweep 1319/1319 clean (0 illegal-memory-access), accuracy 0.9484 flexible / 0.9477 strict.
Adds op_tests/flydsl_tests/test_mxfp4_a4w4_moe_oob.py, a standalone, model-free repro with three tests:
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist