Repository navigation
[HIP] [Bugfix] Honor num_local_tokens on the MXFP4 inline-sort path - #7083
ZLkanyo009 wants to merge 2 commits into
Conversation
Co-authored-by: Cursor <cursoragent@cursor.com>
test_moe_sorting_mxfp4.py checks the fused BM16 aux sort against a sort over the leading rows only, including a CUDA graph replayed with new counts. test_moe_inline_sort_num_local_tokens.py runs fused_moe on the tuned Qwen3.8-Flash-Next rows with padded rows (NaN activations, -1 routes) and checks that the tuned inline-sort config stays selected and that the leading rows match an all-real call. Co-authored-by: Cursor <cursoragent@cursor.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
|
@valarLip @yzhou103 Update: this is now covered by On aiter
So the low-concurrency regression this PR targeted (ROCm/ATOM#2725 × #6708) is fixed on |
|
Closing this one: We confirmed it end to end with Qwen3.8-Flash-Next on MI355X, on a tree with aiter
|
Summary
The MXFP4 inline-sort path of
fused_moenow acceptsnum_local_tokens. The fused BM16 aux sort (mxfp4_moe_sort) learns to read the real row count on device, and_mxfp4_inline_sort_unsupportedstops rejecting calls that pass one. ATOM main always passesnum_local_tokens, so the Qwen3.8-Flash-Next rows #6708 tuned had stopped being used there. With this PR they are used again, and ATOM still skips the padded CUDA-graph rows.Motivation
Two changes interact here:
num_local_tokenstofused_moeon every prefill and decode forward. The value is the real row count for that pass, so the MoE sort skips padded CUDA-graph rows.qwen3_8_flash_fp4_tuned_fmoe.csv). Every row it added for 1–512 tokens is an inline-sort row (flydsl_mxmoe_g1_a4w4_16x*with the layout-v2 atomic stage 2)._mxfp4_inline_sort_unsupportedrejects every call that carriesnum_local_tokens. On ATOM main, all of those tuned rows are therefore discarded:Each hour of serving logs tens of thousands of these lines. Falling back to the default heuristics costs Qwen3.8-Flash-Next 3–11% p90 ITL interactivity at concurrency 1–8 on AgentX (MI355X, TP1, MTP 3). If ATOM stops passing
num_local_tokens, the tuned rows come back, but the padded-row skipping from #2725 is lost.Changes
The rest of the inline-sort path already handles
num_local_tokens:3 * M * topk >= 4 * E. It reads the count on device, routes only the leading rows, prefillsreverse_sortedwith -1 and writes the count tocumsum_tensor[1].use_valid_token_countwhennum_local_tokensis given, so it only writes rows belowcumsum_tensor[1]. The AOT build already covers both variants.m_indices. Padding slots point at rowM, past the end ofA, and the hardware drops those loads.The one exception was the fused BM16 aux sort
mxfp4_moe_sort(aux_sortzi_*/aux_sortonly_*), which handles smallM. It sorted allMhost rows, so the -1 ids of padded rows would have been sorted as if they were experts. This PR changes:moe_sort_quant.cuh:sort_subkerneltakes an optional devicenum_local_tokensand sorts only the leadingmin(max(*num_local_tokens, 0), M)rows. Past those rows,reverse_sortedis -1, andcumsum_tensor[1]holds the count, matchingmoe_sorting_opus_fwd. Padding slots still point at physical rowM. The zero-init variant still clears allMoutput rows, because stage 2 accumulates into them.gen_instances.py,mxfp4_moe_aux_dispatch.h),mxfp4_moe_sort_kernel, the pybind (num_local_tokens=None) and_adaptive_moe_sort/_moe_sorting_impl.mxfp4_moe_sort_kernelchecks that the tensor is int32 with at least one element on the same device, and accepts it only with prologue 0. That is the only prologue the inline-sort path uses, since the path is BM16-only.num_local_tokensis removed from the options that_mxfp4_inline_sort_unsupportedrejects.Calls without
num_local_tokensare unchanged: they passnullptrand run the same kernels.Related: #7062. For
M >= 63, the inline-sort path now also reaches the Opus multi-phase sort withnum_local_tokens, since E = 513 on gfx950 always takes multi-phase. That kernel's tail over-read is fixed separately in #7062. The over-read values are never used, ATOM main already hits the same read on its non-inline-sort rows, and this change does not depend on the fix. All runs below include #7062.Performance
AgentX, Qwen3.8-Flash-Next MXFP4, MI355X, TP1, MTP 3, no offload, 1 h per point, ATOM main
fbb82673e. "Before" is aiter main and "after" is this PR, both on the same node.Three groups were run: "before", "after", and "no count", which is ATOM main patched to stop passing
num_local_tokens(pre-#2725 behaviour) on aiter main. "No count" is the upper bound this PR should reach. In each group, the four concurrencies run at the same time on four GPUs, and the groups run back to back. There were two rounds: round 1 ran after → before → no count, and round 2 ran before → no count → after with the GPU assignment reversed. Averages over the two rounds:Per round, p90 ITL interactivity before → after:
All 24 points had 0 errors. The "after" group is at or above "no count" except for c8, where the two are within 0.3%. *The c2 throughput gap comes from round 1's replay alone. There, the "before" and "no count" runs replayed nearly the same requests (605 requests, 59.8M prompt tokens), while the "after" run replayed 574 requests and 53.3M prompt tokens. In round 2 the three groups replayed 586–588 requests each, and c2 throughput is −0.1%.
Single-layer
fused_moeon the same shape (all rows real, i.e. the c ≤ 16 decode case), replayed in a CUDA graph, in µs:num_local_tokensWith this PR, passing the count costs at most 0.6 µs per layer, and each MoE layer is 11–39 µs faster than main whenever a caller passes
num_local_tokens.Testing
op_tests/test_moe_sorting_mxfp4.pygets new_check_inline_sort_*cases (E/topk 513/11 and 896/16; rows/count (1,1), (4,0), (4,3), (16,5), (48,47), (48,48)). The sort with a count must match a sort of only the leading rows, including padding slots, thereverse_sortedtail (-1),cumsum_tensor[1](= count) and the zeroedmoe_buf. A CUDA graph replayed with new counts is also checked. All cases pass.New
op_tests/test_moe_inline_sort_num_local_tokens.pyrunsfused_moeon the Qwen3.8-Flash-Next TP1 shape (E=513, 2560/640, topk 11, random MXFP4 weights) with padded rows: NaN activations, -1 ids and NaN weights. It covers rows/count 1/1, 4/1, 4/3, 16/5, 16/16, 44/1, 44/43, 128/64, 128/127 and 256/200, which exercises both the fused BM16 sort and the Opus branch. It checks that:Two identical calls already differ by 0.32–0.43% because of the stage 2 bf16 atomic order, so these errors are within run-to-run noise. The test skips off gfx950 and when cu_num != 256, because the tuned rows are pinned there. On main it fails as expected (
M=1: tuned inline-sort row not used).fused_moetest also passes withAITER_MOE_SORT_BACKEND=opus, where smallMtakes the Opus aux sort too.004245be3(plus [Test] Regression test for topk_ids tail reads in multi-phase MoE sorting #7062). The branch has since been rebased ontobbef5a6f2without conflicts, and both op tests pass again on the rebased head (plus [Test] Regression test for topk_ids tail reads in multi-phase MoE sorting #7062).Dependencies
Breaking Changes
None.
mxfp4_moe_sortgains a trailing optional argument.cc @valarLip (ROCm/ATOM#2725) @yzhou103 (#6708)