feat(moe): route before the rows move; per-pass real row counts; V4.1 FP8 KV by valarLip · Pull Request #2725 · ROCm/ATOM

valarLip · 2026-10-26T14:57:51Z

Motivation

MoE routing used to run after the DP / PCP gather, on every rank's rows, and DP padding rows were masked afterwards per consumer (step_pad_rows reading cu_seqlens_q, which a token-graph bucket records at its own request count). This PR routes on the rank that owns the rows, before they move, and gives every pass its real row count on the device, so padding is routed to expert -1 at the source.

Two commits:

feat(v41): MXFP8 KV cache on the original path and in mono — --kv_cache_dtype fp8 for DeepSeek-V4.1 (1.84x the rows of BF16), attention through AITER's mla_sparse_prefill_v41 ([ASM] [HIP] [OPUS] Add gfx950 MLA v4.1 and split-K for the opus MLA sparse prefill and move it to ctypes aiter#6675, merged).
feat(moe): route before the rows move; per-pass real row counts — described below.

Changes

Routing before the rows move

A MoE method splits into select_experts (logical routes), place_experts (EPLB remap and load record, the shared column in the layer's expert layout) and apply(topk_weights, topk_ids, ..., x_scale).
The DP gather, the PCP gather (ATOM_PCP_MOE_MERGE, now inside FusedMoE._run_experts) and the comm-fused DPA gather move hidden states with their top-k weights and ids in one grouped all_gatherv, in the method's RoutedRowsWire (MXFP8 where the experts take MXFP8, MXFP4 under ATOM_MORI_FP4_DISPATCH); the scales reach the experts as x_scale, so nothing is quantized twice.
Rows are placed after every gather: a gathered row picks the same EPLB replica on every rank, and the legacy layout's prefilled shared ids are re-marked for the rows' gathered positions (reassign_gathered_shared_ids).
EPLB remap, load record and shared-column append are one Triton kernel; the forced-remote replica hash mixes in the DP rank for a rank's own rows.
The TP Triton path sorts the routed ids through routing_from_dispatched.

Per-pass real row counts

forward_vars["batch_info"] publishes each step's scheduled requests / tokens / draft tokens; Context.scheduled_tokens_gpu points at the pass's entry, read on device so a replayed graph follows its own step. Captures stage their synthetic step's counts.
Only a consumer of the pass's whole rows gets the count (pass_scheduled_tokens_gpu); a pass over other rows declares its own (ForwardContext.use_pass_rows) — V4.1's decoder over a prefill's tails (ReplayPlan.filled_gpu / finishing_gpu). TBO micro-batches take theirs, clamped to their rows.
The count reaches every routing kernel (num_valid_rows), fused_moe (num_local_tokens) and the Triton experts (gate_valid).
Removes the DP pad-row mask (step_pad_rows, ATOM_MEGA_MASK_PAD_ROWS, zero_pad_rows_); MegaMoE asks its combine to skip -1 slots.

Also

MoRI v1 decode launch geometry by token count.
Dead code removed: triton_kernel_moe_forward, pcp_allgather_rankmajor, rocm_aiter_topk_softmax_fake.
DeepSeek-V4.1 CI (accuracy, benchmark, agentic test/nightly) and docs on the FP8 KV cache; the vLLM-plugin cases stay on auto (that path refuses fp8).
Kimi-K2.7-Code-MXFP4 CI cases removed.

Dependencies

[ASM] [HIP] [FlyDSL] feat: pad-row routing count, custom all_gatherv, MoE quant-sort race fix aiter#6916 (merged): num_valid_rows on the routing ops, custom all_gatherv.
[ASM] [HIP] [OPUS] Add gfx950 MLA v4.1 and split-K for the opus MLA sparse prefill and move it to ctypes aiter#6675 (merged): mla_sparse_prefill_v41.
[Triton/Gluon] [HIP] [OPUS] feat: pad-row count for HERD, gfx1250 sigmoid top-k and MXFP4 sort+quant aiter#6938 (open, performance only): HERD / gfx1250 top-k with num_valid_rows and the fused MXFP4 sort+quant with num_local_tokens. Without it those fast paths fall back; results are unchanged.

Test plan

Unit tests (affected directories: model_ops, deepseek_v41 attention/models, eplb, topk, mega, mxfp4, tbo, envs, benchmark catalog): 2102 passed; the 2 failures are test_compilation's known torch.compile issue in the dev container.
DeepSeek-V4.1-Flash GSM8K 3-shot (1319), DSpark 5, FP4 index, --level 0, MI355X:
- TP4, BF16 KV, mono off: 0.9181 / 0.9181 (flexible / strict), acceptance 49.36%
- DP4 + EP, MoRI, BF16 KV: 0.9174 / 0.9181, acceptance 49.42%
- DP4 + EP, DP gather (MoRI off), BF16 KV: 0.9158 / 0.9158 and 0.9227 / 0.9227, acceptance 46.7% / 47.8%
- TP4, FP8 KV, mono on: 0.9204 / 0.9181, acceptance 49.08%
- CI baseline 0.9224 (threshold 0.91).
1k/1k concurrency 64, FP8 KV, DSpark 5: TP2 6165.92 tok/s total (TPOT 19.99 ms, TTFT 200.66 ms); DEP2 MoRI 6749.39 tok/s (TPOT 17.28 ms, TTFT 898.73 ms). Single runs.
DP-gather DSpark acceptance is ~2pp under TP / MoRI; a paired comparison against main is pending.
Not yet run end to end: gpt-oss (TP Triton), DeepSeek-R1 on the legacy shared layout with DP gather, TBO, PCP, EPLB, MegaMoE.

- Pool geometry takes kv_dtype bf16 / fp4 / fp8. An fp8 row is 512 e4m3 (RoPE dims included) then 16 E8M0 codes, 528 B, row-addressed; the window ring is written and read by byte address. - Original path: sparse_attn_v41_paged_2src over aiter mla_sparse_prefill_v41, q quantized to MXFP8 by rope_quant_window. - Mono: K1 writes the ring row as MXFP8. K2a's QK runs scaled FP8 MFMAs with q quantized once a CTA by the KV row's rule; its PV dequantizes the rows to bf16 and keeps the bf16 MFMAs. - Mono PV reads its V tile back with ds_read_tr16_b64 instead of per-element LDS reads and shifts. - Mono MoE sums the routed contributions in f32 and rounds once, as the reference MoE.forward does, instead of a bf16 rounding at every add. - Check-mode reference, hand-off contract and tests cover the fp8 pool. Requires aiter #6675 (mla_sparse_prefill_v41).

MoE routing now runs on the rank that owns the rows, before any of them move, and the routing kernels route the rows past the pass's real row count to expert -1, so padding is neither dispatched nor computed. Routing: - A MoE method splits into select_experts (logical routes), place_experts (the EPLB remap and load record, the shared column in the layer's expert layout) and apply(topk_weights, topk_ids, ..., x_scale). Rows are placed after every gather: a gathered row picks the same EPLB replica on every rank, and the legacy layout's prefilled shared ids are re-marked for the rows' gathered positions (reassign_gathered_shared_ids). - The DP gather, the PCP gather (ATOM_PCP_MOE_MERGE, now inside FusedMoE._run_experts) and the comm-fused DPA gather move hidden states with their top-k weights and ids in one grouped all_gatherv, in the method's RoutedRowsWire (MXFP8 where the experts take it, MXFP4 under ATOM_MORI_FP4_DISPATCH), whose scales reach the experts as x_scale. - The EPLB remap, load record and shared-column append are one Triton kernel; the forced-remote replica hash mixes in the DP rank for a rank's own rows. - The TP Triton path sorts the routed ids through routing_from_dispatched. Real row counts: - forward_vars["batch_info"] publishes each step's scheduled requests, tokens and draft tokens; Context.scheduled_tokens_gpu points at the pass's entry, read on device so a replayed graph follows its own step. Captures stage their synthetic step's counts. - Only a consumer of the pass's whole rows gets the count (pass_scheduled_tokens_gpu); a pass over other rows declares its own (ForwardContext.use_pass_rows): V4.1's decoder over a prefill's tails (ReplayPlan.filled_gpu / finishing_gpu). TBO micro-batches take theirs, clamped to their rows. - The count reaches every routing kernel (num_valid_rows), fused_moe (num_local_tokens) and the Triton experts (gate_valid). - Removes the DP pad-row mask (step_pad_rows, ATOM_MEGA_MASK_PAD_ROWS, zero_pad_rows_); MegaMoE asks its combine to skip -1 slots. Also: MoRI v1 decode launch geometry by token count; a PCP + DP-attention refusal reworded; dead code removed (triton_kernel_moe_forward, pcp_allgather_rankmajor); V4.1 CI and docs on the FP8 KV cache; the Kimi-K2.7-Code-MXFP4 CI cases removed.

github-actions · 2026-10-26T14:58:13Z

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

✅ Run after the PR is approved and Pre Checkin passes
✅ Run immediately when an approval review is submitted
✅ Can be requested before approval with labels

Label	Tests
`ci:full`	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
`ci:atom`	Run native ATOM model accuracy tests
`ci:vllm`	Run ATOM vLLM OOT model accuracy tests
`ci:sglang`	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2725 --add-label <label>

…-gather

…count

Route through FusedMoE.select_experts, which reads the pass's real row count from the forward context since #2725: the captured Triton route matches the eager AITER one bit for bit on every replay as the count changes, and with a fused shared expert writes the topk_meta_rows slice at the TBO offset. Co-authored-by: Cursor <cursoragent@cursor.com>

Main's #2725 moved MoE's real row count onto the device (`Context.scheduled_tokens_gpu`, read from `forward_vars["batch_info"]`) and added `batch_info` to the `input_ids` publication group. The mixed deferred branch builds its own counts, and the rebase onto #2725 adds `counts[batch_info] = 3` there; without it a fresh runner's first mixed step uploads nothing and MoE reads 0 rows (measured: `[0, 0]` on the device), routing every row to no expert. The end-to-end mixed input test now checks `batch_info` reaches the device as `[bs, total_tokens]`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

main (#2725) moved MoE routing out of apply() into select_experts / place_experts, so the FlyDSL HERD routing moves with it: - Mxfp4MoEMethod.select_experts returns the Top-(K+1) candidates on a HERD step; apply() recognises them by their width and passes herd_topk. - HERD steps keep real routes on pad rows and pass no num_local_tokens: aiter's HERD routing takes neither -1 ids nor a pad-row count. - _flydsl_herd_step reads the routing settings off the layer and now also declines steps the Triton experts or the comm-fused backend run, since routing is decided before apply().

Conflicts, all from main's V4.1 FP8 KV (#2725), which replaced the pool geometry's `packed` flag with `kv_dtype`: - `build_v41_pool_geometry` takes `kv_dtype` and passes it on; the native builder hands it `kv_cache_dtype` and the plugin `v41_kv_cache_dtype`, so both still derive the pool from the one function. - `tests/test_v41_geometry_parity.py` stays deleted: with one function there is no second copy to compare. - DSpark's target-pool check takes main's BF16-or-FP8 rule; the note on the draft's own `fp8_ds_mla` cache is kept. - A plugin test built the geometry with `packed=False`; it now names `kv_dtype="bf16"`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…2125) * [PCP] Use TP for shared_experts * [PCP] Sync sampled ids outside the speculative branch, and reject PCP+TBO+spec * [Fix] Fix PCP + DSpark part I * [PCP] Step the shared-expert fold down to a block-aligned grid `pcp_merged_tp_grid` cut unconditionally on the full pcp*tp grid, which faults in the weight loader whenever `moe_inter_dim / (pcp*tp)` is not a 128 multiple -- any tp*pcp=16 deployment, where 3072/16 = 192 splits a block-quant scale block in half. It now takes the largest divisor that stays aligned and reports a `prescale` of g/(pcp*tp). The two cases that were hand-written fall out of that formula -- no scale on the full grid, the familiar 1/pcp_size when the module stays replicated -- so the pre-scale branch that had become unreachable carries weight again. 8-card configs are unchanged: still cut 8 ways, still prescale 1.0. Also from review: - Refuse ATOM_PCP_MOE_MERGE together with MoE's flatten-TP-across-DP. That fold shards routed across pcp*dp*tp, while a module summed only by the tp all_reduce and the pcp collective never sees a DP reduce. - Refuse online re-quantisation when the effective TP grid is wider than the real TP group: the gather runs on the real group and the shard-back on the override, so it would narrow 1/pcp of a weight as if it were whole. DCP's override narrows instead and is unaffected. - One `_needs_id_sync` predicate behind both the PCP and the TP id broadcast, which have to stay the same condition. - Tests: the two folds asserted against each other on every rank (a reversed flatten goes red on tp4pcp2/tp2pcp4 and, as expected, not on pcp8 where tp=1 cannot tell them apart), the aligned step-down, and that shards tile the ranks evenly -- which is what makes prescale a scalar rather than a per-rank factor. - The import guard uses skip_if_dependency_missing rather than `except Exception`, which would read a SyntaxError as a skip. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [PCP] Fold the padding rows into the last request's FP4 indexer group Since #2688 the FP4 prefill scorer is the row-group kernel, which reads group b's keys off block_tables[b] and asserts the table covers every group. pcp_owned_rows_by_request still gave the PCP padding rows a group of their own, one more than the requests the table has rows for, so any PCP prefill whose token count is not a multiple of pcp asserted (block_tables.shape[0] >= batch). The old prefill kernel looked rows up by batch_id_per_q_token and skipped bid < 0, so the extra group was harmless there. Padding rows sit at position 0, so their bound is 0 and they read no key in whichever group they land in. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Sum the DSpark draft FFN across PCP like Block.forward Under PCP the MoE weights are folded into the TP grid for every model, the draft included, so each rank's FFN output is a partial that needs the PCP all_reduce. Block.forward (target and MTP) gets it from moe_pcp_merge_forward; DSparkLayer.forward_block rewrote the FFN step as a plain self.ffn() call and never summed, which halved the DSpark acceptance rate under PCP. The draft graph is recorded with is_prefill=False and is_dummy_run=False, so the decode branch's all_reduce is what goes into it. V4-Pro-DSpark, gsm8k 200 x 5-shot, k=5: pcp8 32.96% -> 63.14% acceptance against a same-source tp8 at 63.48%. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * style: sort the pcp_utils imports in deepseek_v2 Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Address review: drop DSpark's second aux gather, settle pcp_split_rows in every forward - DSparkProposer.propose called aux_for() only to check it was not None, which under a PCP prefill all_gathered the aux buffers a second time (compute_draft_kv already gathered them) and threw the result away. Check _captures_aux instead. - pcp_split_rows was set only in ModelRunner.forward; RapidServe's decode forward overrides it and never did, so aux was not gathered there. Both forwards now call _settle_pcp_split_rows() after run_model. - Tests: the FP4 owned-rows case now has a real padding row, plus a last-request-owns-only-padding case and a brute-force comparison; a structural test pins DSpark's forward_block to the same moe_pcp_merge_forward branch as Block.forward. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Keep an online-requantized shared expert on the TP grid Online re-quantization gathers the full weight on the real TP group before quantizing (per_Token on the row-parallel w2 always needs it). Folded onto the wider pcp*tp grid, that gather rebuilds only 1/pcp of the weight, which linear.py now refuses at load -- turning a config that ran on main (the shared expert unsharded under PCP) into a load failure. When any of the shared expert's linears will be re-quantized online, do not fold it: it stays on TP exactly as on main. Offline-quantized checkpoints are unaffected. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Put forward's decorators back on forward f324bd9 inserted _settle_pcp_split_rows between ModelRunner.forward and its decorators, so the helper picked up @_profile_runner_stage, @torch.inference_mode() and @with_eplb_forward_monitor -- the last one binds a batch argument, so every forward raised TypeError -- and forward lost all three. A test now calls the helper bare and checks that forward is still wrapped. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Warn when a PCP model lacks pcp_split_active; trim comments - _settle_pcp_split_rows warns once when PCP is on, a drafter is present and the model has no pcp_split_active(): it then assumes no row-split, and a future PCP model that forgot the method would silently leave the drafter's aux ungathered. - Shorten the comments and docstrings this branch added to what/why; the derivations live in the PR description. No code change beyond the warning (AST-compared with docstrings stripped). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Drop shared-expert sharding and the DSpark FFN branch ahead of #2725 main's #2725 moves the PCP gather/reduction into FusedMoE._run_experts, around the routed experts only, and computes the shared expert on each rank's own 1/pcp rows. That already takes the prefill gain the shared-expert sharding was after, and a shared expert sharded across PCP can no longer be summed in prefill (the ranks hold different rows). It also deletes moe_pcp_merge_forward, which the DSpark draft FFN fix called; FusedMoE now all-reduces the draft's MoE itself. Removed: pcp_merged_tp_grid / PcpSharedGrid, Expert(tp_grid), the shared prescale, the online re-quant check, linear.py's override forwarding and gather guard, DSparkLayer.forward_block's merge branch, and their tests. Kept: sampled-id sync and the PCP+TBO+spec guard, pcp_split_active and the DSpark aux gather, the FP4 indexer padding fix. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * test(pcp): give the draft-split tests a forward context #2725 makes _pcp_split_draft_inputs flag the pass as split on the forward context for FusedMoE's PCP gather, so the CPU tests have to supply one. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * test(pcp): pin the DSpark aux gather and the draft-split decision aux_for must gather exactly when the target forward row-split, and the draft prefill split must follow that answer (pcp_split_rows), not the current pass's pcp_split, which drafting clears. Reverting either fix now fails a test instead of silently halving DSpark's acceptance. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Rename pcp_split_rows to target_pcp_split Next to #2725's per-pass pcp_split the old name read as a synonym. It is the target forward's answer, fixed right after it and kept through drafting, which clears pcp_split per pass; it is now also a Context __init__ argument like its sibling. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Snapshot pcp_split for the drafter instead of re-deriving it Since #2725 both deepseek_v2 and deepseek_v4 write the target forward's split decision to context.pcp_split, and FusedMoE gathers on it. target_pcp_split now copies that value right after run_model, before drafting clears it, instead of re-deriving it through a separate pcp_split_active(). The aux gather and FusedMoE thus read one decision. Drops pcp_split_active() from pcp_utils and both models, the _pcp_active delegation, and the missing-method warning; this PR no longer touches the model files. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

The first EPLB rebalance hangs the DEP8 server on rocm/atom-dev:nightly_202610261558 (ROCm/ATOM#2725 routing rewrite). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… route kernels Since ROCm/ATOM#2725 the TP path hands the grouped MoE its per-pass real row count, and topk_softmax writes expert -1 for every slot of a row past it (a graph bucket's pad rows). The two TP route kernels indexed their per-expert counters with the raw id: the fused small-batch kernel (moe_route_psum_fused) wrote LDS at counter[-1] and produced a garbage row for the route, and the plain moe_route kernel did the same in global memory. The gather that followed then read a row that was never allocated: HSA_STATUS_ERROR_MEMORY_FAULT in moe_route_psum_fused within seconds of TP4 + MTP3 serving on real traffic (Qwen3.5-35B-A3B on gfx1250; the MLPerf endpoints recipe's known issue 3). The EP route kernels already drop such routes; the TP ones never saw one before. A route with a negative id now claims no slot and gets DROPPED_ROUTE_ROW, as the g2l kernels do: it adds 0 to counter 0 and keeps the sentinel, so it costs no GEMM row, and the fused kernel's remap phase leaves the sentinel alone (its expert would index past `starts`). The TP branch also left `num_local_tokens` unused, so the kernels after the route had no bound: the multidest quant gather scattered each pad token to its -1 rows (the next fault, moe_token_multidest_quant), and the gather/reduce read them back. The TP branch now derives the same `num_valid_tokens` / `num_valid_routes` the EP branch does, which routes the quant to its bounded kernel and lets the remap and the gather/reduce skip the tail. Without a count nothing changes. Verified against a torch reference for both kernels at E=128/256, 3-512 tokens, 0-100% pad rows: counts exclude pad routes, pad routes map to -1, kept routes land uniquely in their expert's aligned segment, fused and plain psum agree. TP4 + MTP3 on the MLPerf endpoints agentic recipe (c=64, 30 minutes) then runs clean where it faulted within 30 s before.

Add `--sequence-parallel-size`, riding the PCP process group as an axis independent of TP: non-attention layers run on T/W tokens, and an exchange around attention trades the sequence for a head slice (GQA KV heads and an indexer tail are replicated to the ranks whose queries use them). MoE gathers the SP shards, runs this rank's expert slice and reduce-scatters the partial sums; large MXFP4 batches route and quantize before the gather. ATOM_SP_MOE_MEGA (default on) runs eligible SP MXFP4 prefill MoE blocks through AITER's MegaMoeTP, with the native norm, gate and top-k in front of it. Ineligible configs, layers or a failed runtime setup keep the native MoE. ATOM_TP_REPLICATED_O_PROJ (experimental, default off) keeps o_proj replicated under TP and exchanges heads instead of all-reducing. Both are generic and opted into per model; MiniMax-M3 is the first user. Ported onto the request/token graph runner (#2688) and the route-before-gather MoE (#2725): SP graphs capture this rank's token shard at a fixed address, gathered SP rows place their experts after the gather and skip the per-rank row count, and SP keeps its own reduce-scatter instead of the PCP-merge all-reduce. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…under num_local_tokens (#7085) * fix(quant): never write a sorting sentinel's row in the MX quant + MoE sort Sorted with a real-row count (num_local_tokens), a block's padding slots hold the sentinel `topk << 24 | count`: a real token id with an out-of-range slot. fused_mx_quant_moe_sort_kernel only checks the token, so the sentinel passes; its weight is 0, so it takes the zero-weight branch, and #6916 dropped the `topk_id < topk` guard from store_zero_mx_quant_moe_sort_row. With count the last real row, `count * topk + topk` is the row past `out`: a stage-2 call zeroes one row of whatever follows `out` in device memory. Since ROCm/ATOM#2725 passes num_local_tokens on every MXFP4 fused_moe call, this fires on DeepSeek-V4.1-Flash TP2 (a8w4, topk 6): AgentX c16+ dies with an intermittent illegal memory access, which surfaces in later kernels. - A route is valid when topk == 1 or topk_id < topk. An invalid one zeroes its scale slot (indexed by its sorted row, always in range) and neither reads its input row nor writes an out row. - op_tests/test_moe_sorting_mxfp4.py: padding sentinels under a count must leave rows past `out` untouched and the real routes byte-identical to an uncounted sort (fp8 / fp4x2, bounded and unbounded launchers). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(moe_sorting): never load topk_ids past the tensor under num_local_tokens With a device-side token count (LocalToken), the multi-phase sort kernels run their vector loops to the count rounded up to the vector width. When the count is the tensor's height and not a multiple of that width, the last vector loads past the end of the [tokens, topk] ids. The ids are dropped by the token check, so results are right, but the load faults when the tensor ends right before an unmapped page -- intermittently, by allocator layout. ROCm/ATOM#2725 passes num_local_tokens on every fused_moe call, so plain prefill (count == tokens) hits it: DeepSeek-V4.1-Flash TP2 AgentX c64 died about once per 17k requests after the quant-sort fix. - The P0 v1 / v2 and P01 kernels read a vector that would pass the end id by id, ids past the end -1 (no expert); every other vector loads as before. - Counted and uncounted sorts of the same real rows give the same routes for M in 1..16384 (checked as (expert, token, slot) multisets). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

#2725 writes get_forward_context().context.pcp_split on every forward, but the vLLM plugin installs no ATOM context, so every DeepSeek-V2-family model on the plugin (GLM-5.2 in CI) dies on its first forward with 'NoneType' object has no attribute 'pcp_split'. The plugin runs no PCP and the readers short-circuit on PCP being off, so the flag is only needed where the context exists.

Brings the stacked Ulysses-SP + gemm_fusion series (70 commits; 68 files vs main) onto current main, resolving its collisions with #2688 "one graph mechanism per region", #2725 "MoE route before rows move", and the other commits merged since 4438322. Conflicts resolved (7 files): - atom/model_engine/model_runner.py (9 hunks): SP hooks ported onto main's req_graph / ForwardMode.token_graph mechanism (padded-input staging, local width capture, SP gather on replay, tr-graph record path); the gemm_fusion launch-gate check stays immediately before the single req_graph.replay(); the SP graph-capture context manager and the gemm_fusion before/after_graph_capture hooks are kept. - atom/model_engine/engine_core.py (1 hunk): kept the KV-scrub RPC next to main's has_pending_kv_work gate. - atom/model_ops/attention_mha.py (1 hunk): kept the M3 sparse fused-prologue methods and the dense-prologue eligibility change. - atom/model_ops/moe.py (10 hunks): SP gather/prequantize/local-topk hooks re-expressed on main's select_experts + _run_experts flow; the x_scale and scheduled_tokens_gpu handoff is kept; sp_moe_reduce_scatter restored before the reduce_results all-reduce. Cross-hunk adaptations: `_run_experts` gained keyword-only `pcp_handled=False` (default False is main's behavior exactly; SP/token-gather callers pass True so the PCP fold does not move the rows a second time), and `FusedMoE.select_experts` keeps the union's `shared_experts_fused=None` argument (its default reproduces main's fuse_shared logic). - atom/utils/forward_context.py (2 hunks): imports unioned; main's DPMetadata(running_tokens_per_dp_rank) is kept, with the SP expansion of context.running_tokens_across_dp when Ulysses SP is enabled. - atom/utils/backends.py (1 hunk): took main's typing import; this branch's change there was a stale rebase artifact, so the file is now identical to main. - tests/test_model_runner_decode_padding.py (2 hunks): main's token_graph fixture plus the SP assertions (split input, gathered output, padded height). Re-baselined onto the new APIs (test-only): - tests/model_ops/test_gemm_fusion_rebase_compat.py: the compute_hash stand-in gained main's new env/geometry factors; the prepare_inputs pin follows the running_tokens_per_dp_rank API. - tests/model_ops/test_gemm_fusion_r12_abort_path.py: the replay pin follows self.req_graph.replay() and still requires the launch-gate check right before it. - tests/model_ops/test_sp_moe_dispatch.py: ported to main's route-before-gather API (26 tests). - tests/model_ops/test_shared_expert_dispatch.py: the branch's select_experts_with_record test re-expressed on FusedMoE.select_experts' shared_experts_fused argument. The 5 shipped files (atom/model_ops/gemm_fusion/resources.py, hooks.py, site.py, tests/model_ops/test_gemm_fusion.py, tests/model_ops/test_gemm_fusion_aiter_contract.py) are blob-identical to efb0252. Verified without a GPU: 474 CPU tests pass (8 PR-scoped files + 2 ported MoE test files; test_shared_expert_dispatch.py skips locally without aiter), Black clean, Ruff has no findings on lines this merge adds. GPU e2e for SP graph replay, gemm_fusion gate capture/replay placement, and the r18 band is still pending.

CI non-GPU unit tests failed on four test-side drifts introduced when main's #2688/#2725 refactors merged into the SP series; product code is unchanged: - test_moe_layouts_cannot_reuse_incompatible_sp_compilation_artifacts: Config.compute_hash now reads max_model_len / max_num_seqs directly. - test_ulysses_cpu_metadata prepare_inputs test: supply running_tokens_per_dp_rank, forward_vars['batch_info'], the sync ub fields, ubatch_split_bs and token_graph; assert the SP table the merged set_forward_context publishes. - test_forward_context_without_aiter_or_triton: subprocess snippet now uses the running_tokens_per_dp_rank signature and SP table semantics. Local: 34 passed across the two files; E0-equivalent selection 562 passed / 3 skipped; Black and Ruff clean.

…2 on Model Runner V2 (#2818) * fix(plugin): build the reused TP communicator with aiter's own __init__ ReusedPyNcclCudaCommunicator copied CudaCommunicator.__init__ by hand to avoid a second ncclCommInitRank, so every field aiter adds is missing on it. aiter now reads self.gluon_ar_plain in all_reduce, and every vLLM-plugin run dies at engine init with AttributeError. Build a plain CudaCommunicator instead, handing it vLLM's PyNccl communicator in place of the one its __init__ would create. aiter then sets every field it reads (Gluon all-reduce, ownership, all2all state) and RCCL still inits no duplicate communicator; a check fails loudly if aiter ever builds its own. * fix(deepseek_v2): set pcp_split only on an installed ATOM context #2725 writes get_forward_context().context.pcp_split on every forward, but the vLLM plugin installs no ATOM context, so every DeepSeek-V2-family model on the plugin (GLM-5.2 in CI) dies on its first forward with 'NoneType' object has no attribute 'pcp_split'. The plugin runs no PCP and the readers short-circuit on PCP being off, so the flag is only needed where the context exists. * fix(plugin): accept reduce_results in the vLLM GDN subclasses #2520 made the Qwen3.5 decoder layer pass reduce_results to its GatedDeltaNet, so the norm after it can fuse the TP all-reduce. The plugin swaps that class for Qwen3_5GatedDeltaNetVllm, whose __init__ spelled out the old signature, and model construction failed: TypeError: Qwen3_5GatedDeltaNetVllm.__init__() got an unexpected keyword argument 'reduce_results' Accept it and pass it through. allreduce_norm_fusion_enabled() keeps the fusion off in plugin mode, so the plugin layer passes True and computes what it did before #2520. Qwen3NextGatedDeltaNetVllm narrowed its parent's signature the same way; the Qwen3-Next layer does not pass the argument today, so it is aligned before it does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(plugin): carry V4.1's MX scale across the capture break as bytes vLLM's eager_break_during_capture weak-refs every argument of the call it breaks out through torch.ops._C.weak_ref_tensor, a stable-ABI op that accepts the fp8 e4m3/e5m2 family but not float8_e8m0fnu. Since the HIP Mega-MHC norm became the default (#2568) it hands QKV an e8m0 MX scale, and graph capture of the V4.1 plugin cell failed: Not yet supported ScalarType 44, please file an issue ... _forward now passes an e8m0 scale as a uint8 view of the same storage, with its dtype as a non-tensor keyword, and _attend_block views it back. The weak ref still points at the caller's buffer on every replay. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(docker): backport vLLM's Kimi-K3 KDA autotune key fix The base image moved from Triton 3.7.0 to 3.8.0 between the 10-13 and 10-25 nightlies. 3.8 checks autotune keys against the kernel's arguments at decoration time, and the pinned vLLM (db9527a46) keys chunk_kda_fwd_kernel_intra_sub_chunk on NC, which is not one. Importing vllm.models.kimi_k3 raised, and model inspection failed for every Kimi-K3 deployment: ValueError: key contains names that are not kernel arguments: NC. vllm-project/vllm#60446 (fc3931739) drops NC from the key, upstream on 10-09 but after the pinned commit. Apply it as a patch on top of the pinned commit. A scan of every @triton.autotune in the image's vLLM finds no other key naming a non-argument. The PR job reuses the nightly vLLM image unless the vLLM commit or the ATOM base recipe changes, so a change to the vLLM recipe or its patches was never built by its own PR. Rebuild the vLLM image when either changes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(plugin): serve DeepSeek-V4 under vLLM's Model Runner V2 vLLM runs DeepseekV4ForCausalLM on V1 on ROCm; with VLLM_USE_V2_MODEL_RUNNER=1 the V4 cell died in V2's kernel warmup: RuntimeError: DeepSeek-V4 request(s) resumed at chunk_start [2, 2, ...] into a fresh slot, but no state checkpoint path is bound The state-slot hooks were V1's alone. V2 has no _update_states; it drops finished and preempted requests in finish_requests, which execute_model calls first. Hook that method on V2 (the same body as V1's _update_states hook) and initialize_kv_cache on every runner class, so slots come back and the checkpoint groups get their ids. Without this, slots would run out after max_num_seqs requests. V2's kernel warmup also schedules real-looking requests. Their prefill was built before the forward bound the proxy cache, so it took no slots and the following decode found them fresh at chunk_start 2. The metadata builder now binds from the proxy, as V4.1's does, with the arguments the wrapper's forward would pass. V1 is unchanged: its dummy runs carry no request ids, so they still take throwaway arange slots after the bind. Prefix caching (state-checkpoint stores run from V1's _model_forward) and the MTP draft (a V1 get_kv_cache_spec hook, and V2 sizing the drafter's hidden buffer from the class) still hook V1 alone; under V2 they are refused at model construction instead of served wrong. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * ci(plugin): run the V4-Flash and GLM-5.2 cells on Model Runner V2 vLLM keeps DeepseekV4ForCausalLM and GlmMoeDsaForCausalLM on V1 on ROCm (ROCM_DEFAULT_MRV1_ARCHITECTURES) unless VLLM_USE_V2_MODEL_RUNNER is set, so no PR cell ran either model on V2; the other four cells already select V2 by default. Set it on both cells. GLM-5.2 needs no code change for V2; DeepSeek-V4 needs the previous commit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: XiaobingZhang <xiaobingzhangupc@gmail.com> Co-authored-by: perzhang <perzhang@amd.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Main's #2725 moved MoE's real row count onto the device (`Context.scheduled_tokens_gpu`, read from `forward_vars["batch_info"]`) and added `batch_info` to the `input_ids` publication group. The mixed deferred branch builds its own counts, and the rebase onto #2725 adds `counts[batch_info] = 3` there; without it a fresh runner's first mixed step uploads nothing and MoE reads 0 rows (measured: `[0, 0]` on the device), routing every row to no expert. The end-to-end mixed input test now checks `batch_info` reaches the device as `[bs, total_tokens]`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Add `--sequence-parallel-size`, riding the PCP process group as an axis independent of TP: non-attention layers run on T/W tokens, and an exchange around attention trades the sequence for a head slice (GQA KV heads and an indexer tail are replicated to the ranks whose queries use them). MoE gathers the SP shards, runs this rank's expert slice and reduce-scatters the partial sums; large MXFP4 batches route and quantize before the gather. ATOM_SP_MOE_MEGA (default on) runs eligible SP MXFP4 prefill MoE blocks through AITER's MegaMoeTP, with the native norm, gate and top-k in front of it. Ineligible configs, layers or a failed runtime setup keep the native MoE. ATOM_TP_REPLICATED_O_PROJ (experimental, default off) keeps o_proj replicated under TP and exchanges heads instead of all-reducing. Both are generic and opted into per model; MiniMax-M3 is the first user. Ported onto the request/token graph runner (#2688) and the route-before-gather MoE (#2725): SP graphs capture this rank's token shard at a fixed address, gathered SP rows place their experts after the gather and skip the per-rank row count, and SP keeps its own reduce-scatter instead of the PCP-merge all-reduce. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

#2776 gated pad-row masking on EP (`mask_pad_rows=layer.use_ep`) after a suspected "-1 route faults the TP a16w4 MoE". The fault does not reproduce: the TP regression from #2725 was the sort's pad sentinel (real token count) sitting inside the a16w4 gemm2's `token_id < M` bound (padded M), which piled every pad slot into one row. Routing pad rows to -1 is skipped by the sort, so restore it for TP and drop the switch. Kimi-K3 TP8, aiter bd4ff1c70: bs100 decode device time 44.56 -> 42.86 ms/step (-3.8%), c128 TPOT unchanged (102.46 -> 102.41 ms, full batches have no pad rows), GSM8K 0.9666, DSpark GSM8K 0.9651 / acceptance 51.57%, no fault. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…under num_local_tokens (#7085) * fix(quant): never write a sorting sentinel's row in the MX quant + MoE sort Sorted with a real-row count (num_local_tokens), a block's padding slots hold the sentinel `topk << 24 | count`: a real token id with an out-of-range slot. fused_mx_quant_moe_sort_kernel only checks the token, so the sentinel passes; its weight is 0, so it takes the zero-weight branch, and #6916 dropped the `topk_id < topk` guard from store_zero_mx_quant_moe_sort_row. With count the last real row, `count * topk + topk` is the row past `out`: a stage-2 call zeroes one row of whatever follows `out` in device memory. Since ROCm/ATOM#2725 passes num_local_tokens on every MXFP4 fused_moe call, this fires on DeepSeek-V4.1-Flash TP2 (a8w4, topk 6): AgentX c16+ dies with an intermittent illegal memory access, which surfaces in later kernels. - A route is valid when topk == 1 or topk_id < topk. An invalid one zeroes its scale slot (indexed by its sorted row, always in range) and neither reads its input row nor writes an out row. - op_tests/test_moe_sorting_mxfp4.py: padding sentinels under a count must leave rows past `out` untouched and the real routes byte-identical to an uncounted sort (fp8 / fp4x2, bounded and unbounded launchers). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(moe_sorting): never load topk_ids past the tensor under num_local_tokens With a device-side token count (LocalToken), the multi-phase sort kernels run their vector loops to the count rounded up to the vector width. When the count is the tensor's height and not a multiple of that width, the last vector loads past the end of the [tokens, topk] ids. The ids are dropped by the token check, so results are right, but the load faults when the tensor ends right before an unmapped page -- intermittently, by allocator layout. ROCm/ATOM#2725 passes num_local_tokens on every fused_moe call, so plain prefill (count == tokens) hits it: DeepSeek-V4.1-Flash TP2 AgentX c64 died about once per 17k requests after the quant-sort fix. - The P0 v1 / v2 and P01 kernels read a vector that would pass the end id by id, ids past the end -1 (no expert); every other vector loads as before. - Counted and uncounted sorts of the same real rows give the same routes for M in 1..16384 (checked as (expert, token, slot) multisets). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Main's #2725 moved MoE's real row count onto the device (`Context.scheduled_tokens_gpu`, read from `forward_vars["batch_info"]`) and added `batch_info` to the `input_ids` publication group. The mixed deferred branch builds its own counts, and the rebase onto #2725 adds `counts[batch_info] = 3` there; without it a fresh runner's first mixed step uploads nothing and MoE reads 0 rows (measured: `[0, 0]` on the device), routing every row to no expert. The end-to-end mixed input test now checks `batch_info` reaches the device as `[bs, total_tokens]`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

#2776 gated pad-row masking on EP (`mask_pad_rows=layer.use_ep`) after a suspected "-1 route faults the TP a16w4 MoE". The fault does not reproduce: the TP regression from #2725 was the sort's pad sentinel (real token count) sitting inside the a16w4 gemm2's `token_id < M` bound (padded M), which piled every pad slot into one row. Routing pad rows to -1 is skipped by the sort, so restore it for TP and drop the switch. Kimi-K3 TP8, aiter bd4ff1c70: bs100 decode device time 44.56 -> 42.86 ms/step (-3.8%), c128 TPOT unchanged (102.46 -> 102.41 ms, full batches have no pad rows), GSM8K 0.9666, DSpark GSM8K 0.9651 / acceptance 51.57%, no fault. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: JiaoliangYu <Jiaoliang.Yu@amd.com>

Add `--sequence-parallel-size`, riding the PCP process group as an axis independent of TP: non-attention layers run on T/W tokens, and an exchange around attention trades the sequence for a head slice (GQA KV heads and an indexer tail are replicated to the ranks whose queries use them). MoE gathers the SP shards, runs this rank's expert slice and reduce-scatters the partial sums; large MXFP4 batches route and quantize before the gather. ATOM_SP_MOE_MEGA (default on) runs eligible SP MXFP4 prefill MoE blocks through AITER's MegaMoeTP, with the native norm, gate and top-k in front of it. Ineligible configs, layers or a failed runtime setup keep the native MoE. ATOM_TP_REPLICATED_O_PROJ (experimental, default off) keeps o_proj replicated under TP and exchanges heads instead of all-reducing. Both are generic and opted into per model; MiniMax-M3 is the first user. Ported onto the request/token graph runner (#2688) and the route-before-gather MoE (#2725): SP graphs capture this rank's token shard at a fixed address, gathered SP rows place their experts after the gather and skip the per-rank row count, and SP keeps its own reduce-scatter instead of the PCP-merge all-reduce. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* perf(moe): Triton softmax top-k routing for captured graphs Co-authored-by: Cursor <cursoragent@cursor.com> * test(moe): Triton softmax top-k matches AITER bit for bit Co-authored-by: Cursor <cursoragent@cursor.com> * perf(moe): route num_valid_rows pad rows in the Triton top-k softmax Add the optional num_valid_rows count of ROCm/aiter#6916 to triton_topk_softmax and rocm_aiter_topk_softmax_impl: an int32 device tensor whose first element is the count of real rows. Rows at or past it get routed ids -1 and weights 0; shared columns are written as for any row. The count is read on device, so a replayed graph follows each step's count. Without it the kernel is compiled as before (constexpr switch). The AITER topk_softmax path gets the count only when it is set, so an AITER without #6916 keeps working for callers that do not pass it. Co-authored-by: Cursor <cursoragent@cursor.com> * test(moe): Triton top-k follows select_experts' pass row count Route through FusedMoE.select_experts, which reads the pass's real row count from the forward context since #2725: the captured Triton route matches the eager AITER one bit for bit on every replay as the count changes, and with a fused shared expert writes the topk_meta_rows slice at the TBO offset. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(moe): cap the Triton top-k at AITER's 512-expert fused kernel AITER's fused topkGatingSoftmax dispatch stops at 512 experts; 1024 takes moeSoftmax + moeTopK, which softmaxes the whole row before the top-k and rounds the weights differently. Accepting 1024 made eager (AITER) and captured (Triton) routing differ in the last bits. The GPU test now covers every expert count / k / shared-expert count triton_topk_softmax_supported() accepts, bit for bit against AITER, and checks that 1024 experts route through AITER in eager and in a graph. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(moe): Triton top-k clears picked experts to -inf like AITER AITER clears each picked slot to -inf and picks with `>`, so once a row's finite logits run out it picks expert 0 again (weight 0). The Triton kernel cleared with a key below every value, so on rows with fewer than k finite logits it moved on to the next index instead. Clear to -inf's key (index kept, so ties still go to the lowest id). The row max is now the first pick's value, as AITER's is. NaN used to sort first. It now sorts below -inf and is never picked, which AITER agrees with unless the NaN is the first element of one of its per-thread vector loads; there AITER may pick it (every weight NaN) or misorder the row, so NaN rows carry no bit-identity guarantee (stated in the docstring). Tests: rows with 1..k+1 finite logits (and none), and rows with NaNs off AITER's vector-load first slots, bit for bit against AITER. * fix(moe): Triton top-k requires unit last strides on its outputs The kernel stores at w_ptr + row * stride_w + kk (and likewise for the ids), but only the logits' last stride was checked. AITER honours any stride, so a transposed output view would silently be written to the wrong addresses. triton_topk_softmax_supported() now takes the outputs too (same argument order as the launcher) and refuses a non-unit last stride on either, so such calls stay on AITER; the launcher asserts it. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(moe): one shared tail for both top-k softmax kernels The Triton branch returned early with its own copy of the function's tail. Pick the kernel with if/else and share the tail, so a later change to the returned layout cannot land in only one copy. Co-authored-by: Cursor <cursoragent@cursor.com> * test(moe): check the top-k softmax kernel choice on the GPU-less CI The choice between the Triton kernel and AITER's topk_softmax was only tested by the GPU test, which skips on a GPU-less runner, and lived in modules that import Triton and AITER. Move it (triton_topk_softmax_supported plus the switch / scoring-function checks, as use_triton_topk_softmax) to topk_softmax_dispatch.py, which imports neither, and test it there: every shape AITER's fused kernel takes is accepted, 1024 experts and other shapes are not, nor non-unit output strides. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(moe): run the Triton top-k softmax in eager passes too Choosing the kernel on is_current_stream_capturing() routed a token through AITER when its batch ran eagerly (prefill) and through Triton when it was replayed from a graph (decode), so batch independence rested on the two staying bit-identical -- expf, the renormalization's summation order, tie rules, which expert counts take which AITER dispatch branch -- guarded only by a GPU test the GPU-less CI skips. Use the Triton kernel for every pass whose shape it supports; other shapes stay on AITER in both. ATOM_TRITON_TOPK_SOFTMAX=0 now puts every pass on AITER. The eager warmup forward thus also compiles the kernel the graphs capture. The impl / select_experts tests run with the switch on and off and check that eager and captured calls launch the same kernel. Co-authored-by: Cursor <cursoragent@cursor.com> * docs: ATOM_TRITON_TOPK_SOFTMAX and the Triton top-k routing kernel Document the switch (default, what it changes, what 0 does) in environment_variables.md; in model_ops_guide.md say which kernel rocm_aiter_topk_softmax runs and list triton_topk_softmax.py and topk_softmax_dispatch.py. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(moe): send long top-k softmax launches to AITER The Triton routing kernel runs one single-wave program per row and falls behind AITER's topk_softmax past about 20 rows per CU: on MI355X (256 CUs) the crossover is 20-28 rows per CU across the softmax top-k shapes ATOM's models use (5120-7168 rows), with AITER up to 2x faster at 16K rows. use_triton_topk_softmax now also requires gating_output.shape[0] <= 20 * get_cu_num() (triton_topk_softmax_max_rows). The choice still depends only on the launch, never on the capture state, so a pass of a given size runs the same kernel eager or replayed; every default CUDA graph (2048 decode rows with MTP 3, 3072-token prefill graphs) stays below the cutoff and on Triton. The CU count comes from AITER's get_cu_num, imported lazily so the dispatch module stays importable on the GPU-less CI. Tests: the CPU test pins the cutoff (5120 rows at 256 CUs, 6080 at 304) for each model shape on both sides, with the switch and for eager vs captured passes; the GPU test runs the impl at the cutoff (Triton) and one row past it (AITER), eager and replayed, with and without num_valid_rows, bit for bit against both kernels. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(moe): decide the top-k softmax kernel once per layer Each routing call used to read ATOM_TRITON_TOPK_SOFTMAX (an os.getenv per call) and rerun the shape, stride and expert-count checks, although all of them but the row count are fixed once the layer is built. FusedMoE now decides at construction: triton_topk_softmax_rows() takes the layer's expert count, top-k, renormalize, fused shared experts and their scoring func, reads the switch, and returns the row cutoff (20 rows per CU) or None. The layer passes it down select_experts -> rocm_aiter_topk_softmax as triton_max_rows, and the impl only compares the launch's row count. Changing the switch after the layer is built no longer moves it. The output buffers are the impl's own (unit last stride), and router logits are a gate's output or a column slice of it, so no per-launch layout check is kept; the Triton launcher still asserts a unit last stride. Under MoRI all2all fusion the MXFP4 routing drops a scored shared column whose logit stays in the router logits, so it passes None (AITER, as before). _cu_num() loses its functools.cache: aiter's get_cu_num is already cached, and the cutoff is now computed once per layer. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(moe): take the Triton top-k softmax cutoff from a per-arch table The 20 rows per CU crossover was measured on MI355X (gfx950) only, yet gfx942 and gfx1250 inherited it. TRITON_TOPK_SOFTMAX_MAX_ROWS_PER_CU is now a table keyed by AITER's get_gfx(), with the one measured entry, gfx950: 20. On an architecture without an entry, triton_topk_softmax_max_rows() is None, so no layer gets a cutoff and every launch stays on AITER, as before the Triton kernel existed. The architecture is read once per layer, with the switch and the CU count. Tests (CPU): the gfx950 cutoff still follows the CU count; gfx942 (304 and 80 CUs) and gfx1250 give every model shape no cutoff; a decision reads the switch, the architecture and the CU count once each; an unnormalized layer gets no cutoff for any model shape. The GPU tests that need a cutoff skip where there is none, and the hot-path test also counts _gfx calls. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(moe): fail loud on Triton top-k softmax launches the layer did not decide The Triton branch of rocm_aiter_topk_softmax_impl ignored renormalize and num_routing_experts: triton_topk_softmax always renormalizes and takes the routed width from the logits. A layer only gets a cutoff when it renormalizes, so the branch now asserts renormalize, and that num_routing_experts is the logits' width less the scored shared columns, so a caller that pairs a layer's cutoff with another routing fails instead of being routed differently. Two checks on the Triton branch only. For that check to hold, FusedMoEMethodBase.select_experts and the Qwen3.8 single-token decode path now pass the layer's routed expert count (expert_layout.num_routed, as Mxfp4MoEMethod already did) rather than global_num_experts, the dispatch space, which grows with EPLB's redundant experts. num_routing_experts reaches nothing else on these paths: FusedMoE.select_experts only hands it to fused_topk, whose AITER branch does not take it. Test (GPU): an unnormalized launch without a cutoff matches aiter.topk_softmax(need_renorm=False) bit for bit; with a cutoff, an unnormalized launch and launches whose num_routing_experts is 0, the logits' full width or half the experts all raise, and a matching one runs Triton. Also correct the triton_topk_softmax docstring: a row with no finite logit gets NaN for every routed weight (as in AITER and as the test asserts), not expert 0 with weight 0. Co-authored-by: Cursor <cursoragent@cursor.com> * test(moe): pin the routed expert count select_experts passes down Both FusedMoEMethodBase.select_experts and Mxfp4MoEMethod.select_experts must pass expert_layout.num_routed, the width the Triton launch checks the logits against, not global_num_experts (the dispatch space, 10 in the test layout). Co-authored-by: Cursor <cursoragent@cursor.com> * test(moe): use a routed width that is wrong with and without a shared expert Co-authored-by: Cursor <cursoragent@cursor.com> * perf(moe): count only a CU-masked stream's CUs in the Triton top-k softmax cutoff On RapidServe's --disagg-constrained streams a launch gets half of MI355X's CUs. There the Triton kernel is slower than AITER's from 26 rows per CU, by up to 36% at the 40 rows per CU the full-GPU cutoff allows. A layer that may run on such a stream now counts the fewest CUs one gets (RapidServeModelRunner._CU_POOL_FRACTIONS), so its cutoff is 2560 rows. Co-authored-by: Cursor <cursoragent@cursor.com> * docs(moe): note the CU-masked Triton top-k softmax cutoff Co-authored-by: Cursor <cursoragent@cursor.com> * docs(moe): state the CU-masked crossover per measured row count Co-authored-by: Cursor <cursoragent@cursor.com> * fix(moe): decide the Triton top-k softmax cutoff on the GPU running get_gfx() reports GPU_ARCHS' build target, so an MI300 running a gfx942;gfx950 build took MI355X's cutoff; get_gfx_runtime() reads the GPU itself. The kernel's docstring now states which expert a row out of finite logits repeats, and that a NaN is picked only from a row of NaNs, with GPU tests for rows without finite logits and rows of NaNs. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(moe): iq2r routes with the routed width; link the fraction-move issue Iq2rMoEMethod.apply passed global_num_experts (the dispatch space) as num_routing_experts, unlike the other select_experts callers. Pass layer.expert_layout.num_routed and pin it in the iq2r apply test, whose fake layer now has a routed width that differs from global_num_experts. Also link #2863 (move _CU_POOL_FRACTIONS to atom/config.py) next to the lazy model-runner import, and name the AITER commit the NaN vector-load behaviour was probed against. Co-authored-by: Cursor <cursoragent@cursor.com> * docs(moe): note the external CU mask and the NaN exception The dispatch comment says get_cu_num() does not count an HSA_CU_MASK set outside ATOM, and that AITER's CU_NUM covers that case. The kernel's module docstring now names rows with NaN logits as the exception to bit-identity, as the function docstring already does. Co-authored-by: Cursor <cursoragent@cursor.com> * docs(moe): note the external CU mask next to the routing table Co-authored-by: Cursor <cursoragent@cursor.com> * docs(moe): CU_NUM must be set before AITER caches the CU count Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(moe): fold the top-k softmax kernel choice into topK.py - Drop ATOM_TRITON_TOPK_SOFTMAX: the Triton kernel is bit-identical to AITER's, so it needs no off switch, and ATOM adds no env vars. - Move triton_topk_softmax_rows into topK.py, which already imports AITER; the supported / max-rows helpers and the lazy AITER wrappers go. The RapidServe import stays local: model_runner imports topK. - Shorten docstrings and comments; drop the num_valid_rows assert and the empty-launch guard (a zero grid launches nothing). - Tests: the dispatch tests join the GPU test file, the NaN / -inf cases are one parametrized test, the full shape sweep covers fp16 and fp32 logits too, and the switch / call-count tests go. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(moe): keep top-1 / top-2 softmax routing on AITER On MI355X, Triton is slower than AITER at top-1 past ~3K rows and at top-2 near the 5120-row cutoff (results are still bit-identical), so the Triton kernel now takes top-4 and up only; every model shape it served before is top-4 or more. A shared column with scoring func "" is refused too: AITER rejects that call, while the Triton kernel would sigmoid-score it. No caller passes "". Co-authored-by: Cursor <cursoragent@cursor.com> * revert(moe): leave the existing top-k softmax lines as they were - rocm_aiter_topk_softmax_impl allocates AITER's scratch where it did. - Mxfp4MoEMethod.select_experts keeps its existing line; the new argument reads layout.uses_all2all_fusion directly. - _select_logical's new argument defaults to None. - docs/model_ops_guide.md: the routing table is unchanged; the Triton kernel is described in a paragraph below it. Co-authored-by: Cursor <cursoragent@cursor.com> * Test the Triton top-k softmax gate on CPU; drop FusedMoE class attribute Move the shape, architecture and per-CU cutoff cases of triton_topk_softmax_rows to tests/model_ops/test_triton_topk_softmax_rows.py, which imports topK under test_comm_fused_moe's AITER / Triton stubs so Pre Checkin runs them. The CU-masked cases stay in the GPU file. Remove the triton_topk_softmax_max_rows class attribute; __init__ sets it. Co-authored-by: Cursor <cursoragent@cursor.com> * Decide the Triton top-k softmax kernel inside topK.py rocm_aiter_topk_softmax_impl now derives the routed width, top-k and shared columns from the launch and asks a cached triton_topk_softmax_rows, which reads RapidServe's CU masking from the current config. moe.py, qwen4_exp.py and the RTP-LLM topk_compat go back to main: no layer attribute, no triton_max_rows argument, no num_routed change. With the width taken from the launch, the renormalize and width asserts go too; a MoRI all2all launch that keeps a scored shared logit is 2^n + 1 wide and stays on AITER. The docstring says 32 experts is the fewest measured. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Cursor <cursoragent@cursor.com>

Add `--sequence-parallel-size`, riding the PCP process group as an axis independent of TP: non-attention layers run on T/W tokens, and an exchange around attention trades the sequence for a head slice (GQA KV heads and an indexer tail are replicated to the ranks whose queries use them). MoE gathers the SP shards, runs this rank's expert slice and reduce-scatters the partial sums; large MXFP4 batches route and quantize before the gather. ATOM_SP_MOE_MEGA (default on) runs eligible SP MXFP4 prefill MoE blocks through AITER's MegaMoeTP, with the native norm, gate and top-k in front of it. Ineligible configs, layers or a failed runtime setup keep the native MoE. ATOM_TP_REPLICATED_O_PROJ (experimental, default off) keeps o_proj replicated under TP and exchanges heads instead of all-reducing. Both are generic and opted into per model; MiniMax-M3 is the first user. Ported onto the request/token graph runner (#2688) and the route-before-gather MoE (#2725): SP graphs capture this rank's token shard at a fixed address, gathered SP rows place their experts after the gather and skip the per-rank row count, and SP keeps its own reduce-scatter instead of the PCP-merge all-reduce. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Main's #2725 moved MoE's real row count onto the device (`Context.scheduled_tokens_gpu`, read from `forward_vars["batch_info"]`) and added `batch_info` to the `input_ids` publication group. The mixed deferred branch builds its own counts, and the rebase onto #2725 adds `counts[batch_info] = 3` there; without it a fresh runner's first mixed step uploads nothing and MoE reads 0 rows (measured: `[0, 0]` on the device), routing every row to no expert. The end-to-end mixed input test now checks `batch_info` reaches the device as `[bs, total_tokens]`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Add `--sequence-parallel-size`, riding the PCP process group as an axis independent of TP: non-attention layers run on T/W tokens, and an exchange around attention trades the sequence for a head slice (GQA KV heads and an indexer tail are replicated to the ranks whose queries use them). MoE gathers the SP shards, runs this rank's expert slice and reduce-scatters the partial sums; large MXFP4 batches route and quantize before the gather. ATOM_SP_MOE_MEGA (default on) runs eligible SP MXFP4 prefill MoE blocks through AITER's MegaMoeTP, with the native norm, gate and top-k in front of it. Ineligible configs, layers or a failed runtime setup keep the native MoE. ATOM_TP_REPLICATED_O_PROJ (experimental, default off) keeps o_proj replicated under TP and exchanges heads instead of all-reducing. Both are generic and opted into per model; MiniMax-M3 is the first user. Ported onto the request/token graph runner (#2688) and the route-before-gather MoE (#2725): SP graphs capture this rank's token shard at a fixed address, gathered SP rows place their experts after the gather and skip the per-rank row count, and SP keeps its own reduce-scatter instead of the PCP-merge all-reduce. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Main's #2725 moved MoE's real row count onto the device (`Context.scheduled_tokens_gpu`, read from `forward_vars["batch_info"]`) and added `batch_info` to the `input_ids` publication group. The mixed deferred branch builds its own counts, and the rebase onto #2725 adds `counts[batch_info] = 3` there; without it a fresh runner's first mixed step uploads nothing and MoE reads 0 rows (measured: `[0, 0]` on the device), routing every row to no expert. The end-to-end mixed input test now checks `batch_info` reaches the device as `[bs, total_tokens]`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

valarLip added 2 commits October 26, 2026 14:34

valarLip added the ci:atom label Oct 26, 2026

valarLip added 4 commits October 26, 2026 15:04

test: skip the routed-rows gather tests without aiter

e838a0b

fix(vllm-plugin): port V4.1's FP8-pool attention into the plugin's copy

bbc275b

Merge remote-tracking branch 'origin/main' into feat/moe-route-before…

22013ed

…-gather

fix(dsv4-mono): give the mono DSpark proposer's passes the draft row …

79cb029

…count

valarLip merged commit e893780 into main Oct 26, 2026
32 of 35 checks passed

valarLip deleted the feat/moe-route-before-gather branch October 26, 2026 15:48

yhl-amd mentioned this pull request Oct 26, 2026

fix(dp): preserve prefill alignment with capped pending probes #2727

Merged

8 of 9 tasks

yitingw1 mentioned this pull request Oct 27, 2026

[PCP] Fix sampled-id sync, DSpark aux gather and FP4 indexer padding #2125

Merged

1 task

ZLkanyo009 mentioned this pull request Oct 27, 2026

Qwen3.8-Flash-Next on MI355X: review / merge order for the open ATOM and aiter PRs #2712

Open

28 tasks

ZhiweiYan-96 mentioned this pull request Oct 27, 2026

fix(moe): slice fused-shared top-k weights to the batch #2706

Closed

ZLkanyo009 mentioned this pull request Oct 27, 2026

perf(moe): Triton softmax top-k routing, bit-identical to AITER #2665

Merged

JiaoliangYu mentioned this pull request Oct 27, 2026

fix(gfx1250): honor real token counts in TP grouped MoE routing ROCm/aiter#6989

Closed

PerryZhang01 mentioned this pull request Oct 27, 2026

DeepSeek-V4.1-Flash: LMCache tier offload through the MP server (vLLM plugin) #2598

Open

charlieguo1106 mentioned this pull request Oct 27, 2026

[FlyDSL] [MoE] Handle -1 routes (pad rows) in the gfx1250 TP grouped MoE ROCm/aiter#7004

Merged

gbyu-amd mentioned this pull request Oct 27, 2026

fix(moe): keep aiter's tuned inline-sort MXFP4 rows on the TP path (K3 decode MoE +35-45%) #2756

Closed

1 of 2 tasks

JiaoliangYu mentioned this pull request Oct 27, 2026

[HIP] fix(fused_moe): give non-persistent per_1x128 1-stage kernels the full-height scale under num_local_tokens ROCm/aiter#7010

Open

Klaud-Cold mentioned this pull request Oct 27, 2026

[Klaud Cold] Update DeepSeek-V4-Pro FP4 MI355X ATOM AgentX DSpark image to nightly_202610261558 / [Klaud Cold] 将 DeepSeek-V4-Pro FP4 MI355X ATOM AgentX DSpark 镜像更新至 nightly_202610261558 SemiAnalysisAI/InferenceX#4577

Closed

gbyu-amd mentioned this pull request Oct 27, 2026

[HIP] [MoE] MXFP4 inline-sort: accept num_local_tokens, skip ids outside [0, E) ROCm/aiter#7020

Merged

yitingw1 mentioned this pull request Oct 27, 2026

perf(v4): Split the prefill sparse indexer by query row across TP #2722

Merged

1 task

ZhangLirong-amd mentioned this pull request Oct 28, 2026

ci(agentic): run the DSv4.1 DTP2/DEP2 c128 cells with InferenceX's DPA recipe #2784

Merged

sudhu2k mentioned this pull request Oct 28, 2026

[Lumen-RL] Enable Rollout Router Replay (R3) #2274

Open

1 task

XiaobingSuper mentioned this pull request Oct 29, 2026

fix(plugin): vLLM plugin engine init and first forward (reused TP communicator, deepseek_v2 pcp_split) #2816

Closed

ganyi1996ppo mentioned this pull request Nov 2, 2026

fix(moe): mask cudagraph pad rows under pure TP again #2891

Closed

gbyu-amd mentioned this pull request Nov 2, 2026

fix(moe): mask cudagraph pad rows under pure TP again #2893

Merged

JiaoliangYu mentioned this pull request Nov 2, 2026

fix(moe): restore TP MoE num_local_tokens #2905

Merged

yelkhamra mentioned this pull request Nov 2, 2026

[Bugfix] Clarify fused all-reduce limit units with MAX_BYTES ROCm/aiter#7357

Merged

2 tasks done

Navigation Menu

Repository navigation

feat(moe): route before the rows move; per-pass real row counts; V4.1 FP8 KV - #2725

feat(moe): route before the rows move; per-pass real row counts; V4.1 FP8 KV#2725
valarLip merged 6 commits into
mainfrom
feat/moe-route-before-gather

valarLip commented Oct 26, 2026

Uh oh!

github-actions Bot commented Oct 26, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

valarLip commented Oct 26, 2026

Motivation

Changes

Routing before the rows move

Per-pass real row counts

Also

Dependencies

Test plan

Uh oh!

github-actions Bot commented Oct 26, 2026

🏷️ CI Guide

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant