[yhl-amd]

Motivation

Since ROCm/ATOM#2725, ATOM passes num_local_tokens (an int32 [1] device tensor holding the pass's real row count) on every MXFP4 fused_moe call. Two kernels on the gfx950 a8w4 path then touch memory outside their tensors. Both are intermittent: whether one faults depends on what the caching allocator placed after the tensor. On DeepSeek-V4.1-Flash plain TP2 (MI355X, --kv_cache_dtype bf16, mono decode on), InferenceX AgentX c16 / c128 / c256 died with HIP error: an illegal memory access was encountered. With the stock aiter, each c64 run on aac (4 servers per node) died after about 200 requests.

1. Out-of-bounds row write in fused_mx_quant_moe_sort_kernel (csrc/kernels/quant_kernels.cu).
A counted sort fills a block's padding slots with the sentinel topk << 24 | count, which is a real token id paired with an out-of-range slot. The kernel only checks the token (token_idx < num_tokens), so the sentinel passes. Its weight is 0, so it takes the zero-weight branch. #6916 removed the topk_id < topk guard from store_zero_mx_quant_moe_sort_row, so that branch now zeroes row count * topk + topk. When count is the last real row, that is the row after out, and the stage-2 call writes zeros into whatever tensor follows. Before ATOM#2725 the sentinel token was tokens, which is rejected by the token check, so this path was never reached.

2. Out-of-bounds topk_ids loads in the multi-phase sort (csrc/include/moe_sorting_opus.h).
With LocalToken, the P0_v1, P0_v2 and P01 kernels run their vector loop up to the count rounded up to the vector width. During prefill the count is usually the tensor height. When the height is not a multiple of the width, the last vector load goes past the end of [tokens, topk]. The extra ids are discarded by the token check, so the results are correct, but the load faults if the next page is unmapped. (#7062 targets the same loads with a different approach. This change was written independently while bisecting the crash above. Happy to drop this half in favour of #7062.)

Technical Details

• fused_mx_quant_moe_sort_kernel: a route is valid when topk == 1 || topk_id < topk. An invalid route zeroes only its scale slot, which is indexed by its sorted row and so always in range. It no longer reads its input row or writes an out row. store_zero_mx_quant_moe_sort_row gets the guard back as route_valid. Valid routes behave exactly as before.
• moe_sorting_opus.h: new impl::moe_sorting_load_topk_vec. A vector that would cross tokens * topk is read id by id, and ids past the end are set to -1 (no expert). All other vectors use the same single load as before. Here kargs.tokens is the host-side tensor height, which is the maximum in LocalToken mode.

Test Plan

• op_tests/test_moe_sorting_mxfp4.py: new section "Padding sentinels under a num_local_tokens count". It sorts with num_local_tokens = tokens - 1 and the padding row set to -1, places out in front of sentinel bytes, and asserts that nothing past out is written and that the real routes match an uncounted sort byte for byte. It covers fp8 / fp4x2, tokens 2/5/7/33, and both the bounded and unbounded launchers.
• Counted vs uncounted moe_sorting on the same real rows, compared as (expert, token, slot) multisets, for M in {1, 3, 7, 33, 257, 1023, 2090, 4097, 8191, 16383, 16384} and counts M, M-1 and M/2.
• op_tests/test_moe_sorting.py (including CUDA graph capture/replay).
• End to end: ATOM e893780f4 (ATOM#2725), DeepSeek-V4.1-Flash TP2, AgentX trace replay with the InferenceX recipe, 30 min per run, four servers per node.

Test Result

• New sentinel section: stock aiter fails the first case (wrote past out; exactly one 1152-byte row after out is zeroed). With this PR, 16/16 cases pass.
• Full test_moe_sorting_mxfp4.py passes (exit 0), including the 12 zero-weight cases from [ASM] [HIP] [FlyDSL] feat: pad-row routing count, custom all_gatherv, MoE quant-sort race fix #6916. test_moe_sorting.py passes. Counted vs uncounted multisets match for every M and count, before and after this change, so routing results are unchanged.
• End to end (illegal-access crashes):

aiter	runs	requests	crashes
stock	c64 ×2	~500	2
fix 1 only	c64 ×4	~17,000	1 (large-M prefill)
fix 1 + 2 (this PR)	c64 ×4, c16, c64, c128, c256	~36,000	0

The request error rate with this PR is 0.02–0.09%, the same as a run that does not pass num_local_tokens at all.

The end-to-end and kernel tests ran on the aiter in rocm/atom-dev:nightly_202610261558 (MI355X) with these same hunks applied. This branch is that patch rebased onto current main, which still contains both issues.

🤖 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 7085 --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.