[chennyiiis]

Overview

AITER biased_grouped_topk runs an LDS group scan even when the grouping cannot eliminate any expert — that is, when topk_group >= num_expert_group, so every group is selected. This adds a register-resident Triton sigmoid+bias top-k and substitutes it at rocm_aiter_ops.biased_grouped_topk, the single call site both ROCm routers reach.

Previously this PR specialised rocm_aiter_grouped_topk only. Moving it to the wrapper came out of review (thanks @simondanielsson) and covers both paths without touching routing metadata:

caller	group arguments it passes
GroupedTopKRouter	the model's own 1/1 for a degenerate router
FusedTopKBiasRouter	a count synthesised to fit an AITER kernel — measured 2/2 for a 64-expert router — with topk_group == num_expert_group

vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py is no longer modified; the diff is three files, additions only.

Claims

• 1.67×–2.15× on the routing op on gfx942 under CUDA graph replay, across the configurations listed below.
• Equivalent to the call it replaces whenever no group is eliminated, including num_expert_group == topk_group > 1, which is the shape FusedTopKBiasRouter produces.
• Eager execution is not faster, so it is excluded by the substitution rather than claimed as a win.
• Selected expert ids match AITER on 3323 of 3328 routing calls in a live Moonlight-16B-A3B forward and on 1248 of 1248 in Kimi-VL-A3B.
• No end-to-end speedup is claimed. Interleaved A/B on two models is +0.4%–0.8%, comparable to the observed same-arm variation.
• No new env var, and no change to router selection or to the advertised routing method.

Dispatch scope: why this is restricted to CUDA graph steps

The kernel's advantage is device time. Timed from a plain Python loop it is slower — it is a Triton launch from Python against a single C++ dispatch, and that overhead is not amortised.

Because the gate is a token-count bound rather than an execution-phase check, eager steps can reach it, so this was measured end to end with --enforce-eager, interleaved B-A-B-A on Moonlight:

in / out	concurrency	AITER tok/s	ungated PR tok/s	Δ	same-arm drift
64 / 256	8	488.7	471.8	−3.5%	0.6% / 1.2%
64 / 256	32	1910.4	1876.9	−1.8%	0.8% / 0.1%
1024 / 1024	8	780.8	758.1	−2.9%	0.8% / 0.2%
1024 / 1024	32	3008.8	2915.3	−3.1%	0.0% / 0.9%

Consistent sign and a magnitude several times the drift, so this is a real regression, not noise. The dispatch therefore checks get_forward_context().cudagraph_runtime_mode, mirroring the check in vllm/models/deepseek_v41/amd/ops/sparse_indexer.py, and eager steps keep AITER. Re-running the same eager A/B with the check in place:

in / out	concurrency	AITER tok/s	this PR tok/s	Δ	same-arm drift
64 / 256	8	490.2	492.5	+0.5%	0.7% / 0.1%
64 / 256	32	1934.9	1933.3	−0.1%	1.2% / 0.0%
1024 / 1024	8	787.1	786.8	−0.0%	0.1% / 1.0%
1024 / 1024	32	3018.7	3009.5	−0.3%	1.9% / 0.8%

Instrumentation observed no Triton dispatches with enforce_eager=True, while runs with CUDA graphs enabled exercised the new path.

Validation

All runs on one MI300X (gfx942, 304 CUs), torch 2.10.0, triton 3.6.0.

Three checkpoints are used, all with num_expert_group=1, topk_group=1, a sigmoid router and a noaux_tc correction bias. Which one carries which measurement, since they differ:

model	experts / top-k	context	used for
Moonlight-16B-A3B	64 / 6	8192	routing comparison, e2e ≤6k, NIAH 4096, GPQA
Kimi-VL-A3B-Instruct	64 / 6	131072	routing comparison, e2e 1k
Kimi-Linear-48B-A3B	256 / 8	1M	e2e 1k, GPQA, and the 8k/64k attempt

NIAH ran at 4096 and GPQA at roughly 1k prompts, both well under the 8192 input length at which the engine fault described below appears, so those results and that fault are not in conflict.

Which model carries which measurement, since they differ:

measurement	model	context used
unit, micro	— (kernel in isolation)	n/a
routing comparison	Moonlight-16B-A3B and Kimi-VL-A3B	short decode rows
e2e 1k/1k, 6k/1k	Moonlight-16B-A3B	≤ 7k
e2e 1k/1k	Kimi-VL-A3B	≤ 2k
e2e 8k/1k, 64k/1k	not obtained — see below	—
NIAH	Moonlight-16B-A3B	4096
GPQA	Moonlight-16B-A3B	~1k prompts

NIAH and GPQA ran on Moonlight at 4096 and ~1k respectively, both well under the 8192 input length at which the engine crash below appears, so there is no contradiction between those results and that crash.

Unit

Against upstream db3b0f0fa:

pytest tests/kernels/moe/test_rocm_aiter_topk.py -k "flat_topk or substitutes"
# 72 passed

Each correctness test compares against rocm_aiter_ops.biased_grouped_topk(..., 1, 1, ...) — the exact call being replaced — through the same wrapper production uses.

test	coverage
test_flat_topk_matches_aiter	40 cases: {bf16, fp32} × {renorm on, off} × 10 shapes, including primes/odd sizes that exercise the masked tail of BLOCK_E = next_power_of_2(num_experts): (3,257,8), (7,131,5), (13,1021,31), (1,2,1), (5,65,3)
test_flat_topk_casts_bias_like_aiter	bf16 gating with an fp32 bias — see Bias dtype
test_flat_topk_matches_aiter_with_uneliminating_groups	num_expert_group == topk_group ∈ {1, 2, 4, 8} × renorm on/off — the FusedTopKBiasRouter shape
test_flat_topk_honours_routed_scaling	routed_scaling_factor 1.0 / 2.446 / 2.827
test_flat_topk_gate	9 parametrized shape cases the gate must accept or refuse
test_flat_topk_dispatches_only_under_cudagraph	the NONE / PIECEWISE / FULL mode check
test_flat_topk_falls_back_without_forward_context	direct calls outside a forward pass keep AITER
test_biased_grouped_topk_substitutes_only_when_expected	the wrapper substitutes for 1/1 and 2/2 under a graph, and not for 8/4 or in eager

The rest of that file reports 34 failures in my container which reproduce identically on unmodified db3b0f0fa (34 failed / 5 passed before this change, 34 failed / 77 passed after); this PR does not touch them.

pre-commit run --files vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py \
  vllm/model_executor/layers/fused_moe/router/rocm_sigmoid_bias_flat_topk.py \
  tests/kernels/moe/test_rocm_aiter_topk.py

Micro

Per-call device time under CUDA graph replay, 50 calls captured per graph so the replay overhead is amortised ~50× instead of landing on both arms.

gating	experts	tokens	AITER (µs)	PR (µs)	speedup
bf16	256	1 / 8 / 32 / 128	5.47 / 5.67 / 5.72 / 5.76	3.00 / 3.18 / 3.20 / 3.27	1.82× / 1.78× / 1.79× / 1.76×
bf16	384	1 / 8 / 32 / 128	6.80 / 7.03 / 7.07 / 7.14	4.01 / 4.14 / 4.16 / 4.26	1.69× / 1.70× / 1.70× / 1.67×
fp32	256	1 / 8 / 32 / 128	6.05 / 6.12 / 6.24 / 6.26	2.82 / 2.99 / 3.00 / 3.16	2.15× / 2.05× / 2.08× / 1.98×
fp32	384	1 / 8 / 32 / 128	7.33 / 7.54 / 7.61 / 7.62	3.56 / 3.70 / 3.73 / 3.79	2.06× / 2.04× / 2.04× / 2.01×

The Triton arm moves little with expert count (3.0 µs at 256, 3.9 µs at 384) while AITER grows (5.9 → 7.2 µs). This is the measured range; expert counts outside 256/384 and architectures other than gfx942 are not covered here.

Routing comparison on live models

Comparing generated text between arms is not a usable instrument on this setup: it is not deterministic run to run — two baseline-only runs of the same four greedy prompts already diverge. Instead both kernels were run on identical inputs inside a live forward, AITER into the real buffers and the new kernel into scratch, with generation continuing on the AITER result so the measurement cannot perturb the run.

model	calls compared	calls where all selected ids match	max abs weight delta
Moonlight-16B-A3B	3328	3323	4.77e-07
Kimi-VL-A3B-Instruct	1248	1248	4.77e-07

Both have n_group=1, topk_group=1, scoring_func="sigmoid", noaux_tc, norm_topk_prob=true, 64 experts, top-6. The five differing Moonlight calls each carry a weight delta of 1.19e-07 or 4.77e-07; that is consistent with a near-tie in the biased score resolving to a different expert, but the weight delta alone does not prove the scores were exactly equal.

End to end

vllm serve + vllm bench serve, interleaved, one fresh server per block.

An initial sequential A/B (all of arm 0, then all of arm 1) showed +16–24% throughput. That did not reproduce under interleaved testing and has been withdrawn — the same AITER baseline measured 1789 tok/s sequentially and 2250 tok/s interleaved on a shared host.

model	in / out	concurrency	Δ throughput	Δ TPOT	same-arm drift
Moonlight	1024 / 1024	8	+0.7%	−0.8%	0.4% / 1.4%
Moonlight	1024 / 1024	32	+0.8%	−0.9%	1.1% / 0.6%
Kimi-VL	1024 / 1024	8	+0.4%	−0.6%	1.0% / 0.9%
Kimi-VL	1024 / 1024	32	+0.7%	−0.9%	3.6% / 0.7%
Kimi-Linear	1024 / 1024	8	+1.1%	−1.1%	0.4% / 0.4%
Kimi-Linear	1024 / 1024	32	+0.1%	−0.1%	1.1% / 0.7%

Comparable to the same-arm variation, so no end-to-end speedup is claimed. The Kimi-Linear rows discard the first repetition of each block: the very first benchmark run of the campaign returned 1066 tok/s against 1510-1543 for the other seven runs of that point, a warm-up artefact rather than a property of either arm.

The 8k/1k and 64k/1k points are now measured. The earlier report that they could not be served was wrong, and the reason is worth stating because it is a bug in the benchmark harness rather than in serving.

vllm bench serve --dataset-name random samples token ids from tokenizer.vocab_size. For this checkpoint that is 163842, while the model's embedding has config.vocab_size = 163840 rows, so ids 163840 and 163841 index past the table and the embedding gather faults. The dataset emits a contiguous run of ids, so a 65536-token prompt covers 40% of the vocabulary and nearly always includes one of the two; an 8192-token prompt covers 5%, which is why 8k completed under some seeds and not others. On stock vLLM, with no part of this PR installed, the identical 8192 configuration completes 3/3 with --seed 7 and faults 3/3 with --seed 1234:

seed	completed	Memory access fault
7	8, 8, 8	0
1234	0, 0, 0	1

That also explains the fault seen earlier on a different machine with Kimi-VL, where the kernel named in the HSA exception was at::native::vectorized_gather_kernel — the embedding gather.

Replacing the synthesised prompts with real text (a custom dataset whose maximum token id is 138177, verified below the embedding size) removes the fault entirely: 8 of 8 runs complete at 64k with zero faults.

in / out	concurrency	AITER tok/s	this PR tok/s	Δ throughput	Δ TPOT	same-arm drift
8192 / 1024	4	4305.4	4466.3	+3.7%	−4.0%	4.1% / 1.7%
8192 / 1024	8	6154.7	6162.5	+0.1%	−1.2%	1.4% / 1.8%
65536 / 1024	1	13528.3	14305.7	+5.7%	−7.0%	8.6% / 8.4%
65536 / 1024	2	18045.6	17927.5	−0.7%	−1.6%	1.7% / 0.9%

Every point has a same-arm drift of the same size as the difference between arms, so these do not change the conclusion: no end-to-end speedup is claimed. The 64k rows carry only 2 and 4 requests per run, which is why their drift is the largest in the table.

Accuracy

A/B on Moonlight through the served endpoint, lm-eval 0.4.13.

NIAH (RULER, 4096): the three reported task scores match between arms — niah_multikey_1 1.00/1.00, niah_single_2 0.94/0.94, niah_single_1 0.00/0.00. The subtask scoring zero on both arms is inconclusive; this is task-level agreement, not per-sample output agreement.

GPQA (gpqa_diamond_zeroshot, 198 questions, loglikelihood). A first single-run A/B looked like +2.0 points for this PR. Three interleaved runs per arm show the difference changes sign:

run	AITER	this PR
1	0.3131	0.3333
2	0.3384	0.3182
3	0.3485	0.3283

AITER alone spans 0.3131–0.3485 across its own three runs.

Repeated on Kimi-Linear-48B-A3B, three interleaved runs per arm:

run	AITER	this PR
1	0.3889	0.4040
2	0.3990	0.4192
3	0.3990	0.4091
median	0.3990	0.4091

Here this PR is nominally higher in all three runs, by 1.0 point. That is not claimed as an accuracy improvement: 1.0 point is 2 questions out of 198, well inside the ±3.4 point standard error, and a routing difference confined to 0.15% of calls has no mechanism by which it could make a model more accurate. Taken with the Moonlight runs, these results show no regression, but are not sufficient to establish accuracy equivalence.

Bias dtype

rocm_aiter_ops.biased_grouped_topk casts correction_bias to the gating dtype before dispatching, because AITER raises on a mismatch. A bf16 gating output next to an fp32 correction bias is a common router layout, so without the same cast the new path would score experts at a precision the replaced path never uses, and could select different experts. sigmoid_bias_flat_topk therefore makes the same cast at its own entry; Removing the cast fails 4 of the 6 cases in test_flat_topk_casts_bias_like_aiter; the other two shapes happen not to be sensitive to the bias precision.

Known behavioural differences

input	result
well-formed gating	weights agree to 2e-3 in unit tests; selected ids match on 4571 of 4576 live routing calls across two models
near-ties in the biased score	weights agree to 1.19e-07; the index chosen among near-equal experts may differ
NaN in gating output	materially different from AITER (6/10 synthetic configs, max weight delta 0.93)

If maintainers would rather the gate reject non-finite input, that is a small addition to can_use_sigmoid_bias_flat_topk.

Details

Why the substitution sits at the wrapper

Review asked whether a model with no grouping should simply not be routed to GroupedTopKRouter. Two things make the wrapper the better place.

First, the advertised routing method is not a label. layer.py:378 takes it from router.kernel_routing_method_type and stores it on the MoE config, where the AITER mxfp4 expert kernels consume it: _supports_routing_method lists [Renormalize, RenormalizeNaive, DeepseekV4, DeepSeekV3], self.renormalize is derived from the same tuple, and the apply path branches on DeepSeekV3 to pick score_mode="sigmoid". MiniMax2 is in none of them, so degenerating num_expert_group 1 → None would make those backends refuse these models and would flip renormalize to False for a checkpoint whose config says norm_topk_prob: true.

Second, and more decisively, changing the router would not remove the op. Calling fused_topk_bias directly — the routine FusedTopKBiasRouter uses — with a 64-expert, top-6 sigmoid+bias router on ROCm with AITER enabled:

synthesised group count for 64 experts: 2
biased_grouped_topk CALLED 1x with (num_expert_group, topk_group) = (2, 2)

_aiter_get_num_expert_group picks a divisor AITER has a kernel for, with topk_group == num_expert_group so the grouping is a no-op. A model moved off GroupedTopKRouter would still reach the grouped kernel, with a fabricated two-group configuration in place of its real single group.

Specialising at rocm_aiter_ops.biased_grouped_topk covers both callers, needs no routing-metadata change, and generalises the condition from "one group" to "no group is eliminated".

Relation to #49870

#49870 short-circuits the same degenerate shape for the case where AITER fused MoE is not enabled, which this PR does not touch.

Scope and bounds

can_use_sigmoid_bias_flat_topk requires topk_group >= num_expert_group, a contiguous bias of exactly num_experts, and bounded token/expert/topk counts. The scoring function is not checked: at this call site the op is AITER's biased (sigmoid) top-k by construction, so matching that op is the contract. _MAX_EXPERTS/_MAX_TOPK bound the register file the kernel holds a row in; _MAX_TOKENS=128 selects small batches, which mostly means decode and speculative verification but is a token-count bound, not an execution-phase check — a short prefill can satisfy it. These are the currently accepted range, not a demonstrated optimum: the measurements above cover 256/384 experts on gfx942, so the thresholds and the single-wave configuration are not shown to be best across the whole accepted range.

Limitations

• Measured on gfx942 only.
• The e2e and accuracy comparisons used the container's own older vLLM build for both arms, because today's main cannot start against this image's AITER (aiter.ops.flydsl.moe_common missing, fused_moe() without swiglu_limit). Both arms are that same build differing only in this kernel. Unit and micro results are against db3b0f0fa.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: One branch in rocm_aiter_grouped_topk behind a shape gate and a graph-mode check, no new env var.

• Testing and Validation: 62 tests, including a cast-removal regression caught by 4 of 6 cases, and tests for the graph-mode check and the missing-context fallback; routing compared on 4576 live calls across two models.

• Code Quality and Style: pre-commit clean on the three changed files.

• Pull Request Contents: Claims, evidence, reproduction commands, a withdrawn measurement and the limitations are stated above.

---

Found by Hyperloom; reviewed and measured by hand.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[simondanielsson]

Thanks for the fix!

Question before any other changes are made: is it rather possible to make a change so that we never route to this grouped topk custom op at all in this case, rather than adding a new "pure topk" fallback here? There's already some handling of degenerate_grouping in the router factory, please check if we just need to make a change there instead

Otherwise I'm also thinking if we should add this fallback in aiter instead

Suggestion: Can we please add

1. e2e perf numbers, i.e. vllm bench serve at 1k/1k, 8k/1k, 64k/1k at a few different concurrencies?
2. gpqa and niah results to ensure accuracy is on point

[simondanielsson]

Suggestion: Let's find a way to auto-enable this feature whenever it's beneficial. We don't want to give users another optional knob :)

[simondanielsson]

Suggestion: can we also add some weird or prime shapes here?

[simondanielsson]

Suggestion: Can we get rid of this thin helper to limit the diff? We can put the scoring_func check into can_use_sigmoid_bias_topk itself

[simondanielsson]

Suggestion: Can we make this name a bit more specific? I.e. it's a triton sigmoid single group topk

[Fangzhou-Ai]

Are these parameters derived from kimi k3 only?

[chennyiiis]

Thanks @simondanielsson and @Fangzhou-Ai for the review, and @shen-shanshan for the verified label. Implementation and description are updated.

@simondanielsson — router selection and implementation placement

I checked the existing factory logic. For biased sigmoid routing with renormalize=True, changing num_expert_group from 1 to None changes the advertised routing method from DeepSeekV3 to MiniMax2, so the current routing_method_preserved check keeps GroupedTopKRouter. With renormalize=False, the factory already falls through.

Separately, FusedTopKBiasRouter still calls AITER's biased_grouped_topk for some supported shapes (fused_topk_bias_router.py:258-272), so changing the selected router alone does not necessarily eliminate the grouped op.

This explains the current dispatch, but does not rule out a refactor that preserves the routing metadata. Moving the specialization into AITER is also a valid alternative; I'm open to that placement if it is preferred.

Auto-enabling only where it is beneficial

You were right to push on this, and the first version did not satisfy it.

The previous dispatch still allowed the fast path in eager execution whenever its shape gate passed. Interleaved --enforce-eager measurements on Moonlight showed a 1.8–3.5% throughput regression.

The dispatch now checks get_forward_context().cudagraph_runtime_mode, following the existing pattern in deepseek_v41/amd/ops/sparse_indexer.py. It retains AITER when the runtime mode is NONE or no forward context is available.

Repeating the same eager A/B comparisons after this change gave throughput differences of −0.3% to +0.5%, within the observed same-arm variation of up to 1.9%. Instrumentation observed no Triton dispatches with --enforce-eager, while runs with CUDA graphs enabled exercised the new path. Tests now cover the NONE, PIECEWISE and FULL mode checks, plus the missing-context fallback.

Requested cleanups

The new environment variable is removed, the original _use_sigmoid_bias_topk helper is removed, the scoring-function check is part of the gate, the kernel is renamed to sigmoid_single_group_topk, and the tests add prime/odd shapes and a parametrized gate test — 62 passing cases in the targeted suite.

I also aligned the bias dtype conversion with the existing AITER wrapper and added regression coverage: rocm_aiter_ops.biased_grouped_topk casts correction_bias to the gating dtype, and without the same cast a bf16 gating output next to an fp32 bias is scored at a precision the replaced path never uses. Removing the cast fails 4 of the 6 cases in that test; the other two shapes happen not to be sensitive to the bias precision.

Performance and accuracy

The routing microbenchmark shows 1.67–2.15× on gfx942 under CUDA graph replay, over 256/384 experts, bf16/fp32, 1–128 tokens. That is the measured range, not a claim about the whole accepted range.

I am not claiming an end-to-end speedup. The initial sequential A/B result did not reproduce under interleaved testing and has been withdrawn. Interleaved results on Moonlight and Kimi-VL are +0.4%–0.8%, comparable to the observed same-arm variation.

The requested 8k/1k and 64k/1k points remain incomplete. Kimi-VL has a 131072 context, so the window is not the obstacle, but both the AITER baseline and this PR encounter the same HSA exception in the current test environment; the root cause has not been isolated.

The three reported NIAH task scores match between arms, although the subtask scoring zero on both arms is inconclusive. Across three interleaved GPQA runs per arm, the difference changes sign. These tests do not show a consistent accuracy change, but are insufficient to establish accuracy equivalence.

For the routing comparison, all selected expert ids match in 3323/3328 calls on Moonlight and 1248/1248 calls on Kimi-VL. The five differing Moonlight calls are reported in the PR description.

One validation limitation: the e2e and accuracy comparisons used the same older container vLLM build for both arms because of AITER compatibility issues. Unit and microbenchmark results were obtained against db3b0f0fa.

@Fangzhou-Ai — parameter bounds

The gate uses shape limits rather than model names, and correctness coverage now includes Moonlight, Kimi-K2 and synthetic prime/odd shapes. However, the current measurements do not establish that these thresholds or the single-wave configuration are optimal across the entire accepted range.

_MAX_TOKENS=128 selects small batches, primarily covering decode and speculative verification; it does not distinguish execution phases, and a short prefill can satisfy it. The larger prefill batches observed in the live run fell back to AITER. Performance evidence is currently limited to the listed configurations on gfx942.

[simondanielsson]

router selection and implementation placement

I checked the existing factory logic. For biased sigmoid routing with renormalize=True, changing num_expert_group from 1 to None changes the advertised routing method from DeepSeekV3 to MiniMax2, so the current routing_method_preserved check keeps GroupedTopKRouter. With renormalize=False, the factory already falls through.

Separately, FusedTopKBiasRouter still calls AITER's biased_grouped_topk for some supported shapes (fused_topk_bias_router.py:258-272), so changing the selected router alone does not necessarily eliminate the grouped op.

This explains the current dispatch, but does not rule out a refactor that preserves the routing metadata. Moving the specialization into AITER is also a valid alternative; I'm open to that placement if it is preferred.

Can you please check what changes we would need to make to change this existing dispatch logic?

Just thinking intuitively: if there is no grouping to be done by the model, I think the model should not be routed to a GroupedTopKRouter. But I might be missing some details here

[chennyiiis]

You were right to push on this. I checked, and it changed where the patch belongs — the PR is respun accordingly.

What blocks the factory change today

The advertised routing method is not a label. layer.py:378 takes it from router.kernel_routing_method_type and stores it on the MoE config, where the AITER mxfp4 expert kernels consume it three ways: _supports_routing_method lists [Renormalize, RenormalizeNaive, DeepseekV4, DeepSeekV3], self.renormalize is derived from the same tuple, and the apply path branches on DeepSeekV3 to pick score_mode="sigmoid". MiniMax2 is in none of them, so degenerating num_expert_group 1 → None would make those backends refuse these models and would flip renormalize to False for a checkpoint whose config says norm_topk_prob: true. The specific blocker is fused_topk_bias_router.py:422, where routing_method_type hardcodes num_expert_group=None.

But changing the router would not remove the op

I measured this rather than reading it. Calling fused_topk_bias directly — the routine FusedTopKBiasRouter uses — with a 64-expert, top-6 sigmoid+bias router on ROCm with AITER enabled:

synthesised group count for 64 experts: 2
aiter fused moe enabled: True
biased_grouped_topk CALLED 1x with (num_expert_group, topk_group) = (2, 2)

_aiter_get_num_expert_group picks a divisor AITER has a kernel for, with topk_group == num_expert_group so the grouping is a no-op ("Grouping is a no-op here, so any divisor", lines 108-124). So a model moved off GroupedTopKRouter would still reach the grouped kernel, and would reach it with a fabricated two-group configuration in place of its real single group.

What I changed

Your intuition — no grouping, no grouped work — is right; it just needs to act one level down. The specialization now sits in rocm_aiter_ops.biased_grouped_topk, the single call site both routers reach, and the condition is generalised from "one group" to "no group is eliminated" (topk_group >= num_expert_group). That covers GroupedTopKRouter passing the model's real 1/1 and FusedTopKBiasRouter passing a synthesised 2/2, needs no routing-metadata change, and leaves router selection alone.

rocm_aiter_moe.py is no longer modified at all. The diff is three files, additions only.

New coverage for the generalised condition:

• test_flat_topk_matches_aiter_with_uneliminating_groups — num_expert_group == topk_group ∈ {1, 2, 4, 8} × renormalize on/off, each compared against AITER.
• test_biased_grouped_topk_substitutes_only_when_expected — the wrapper substitutes for 1/1 and 2/2 under a graph, and does not for 8/4 (a group is eliminated) or in eager.

72 passing cases in the targeted suite, pre-commit clean.

One caveat on scope: the gate no longer checks scoring_func. At this call site the op is AITER's biased (sigmoid) top-k by construction, so matching that op rather than the router's declared scoring function is the contract. Say the word if you would rather keep an explicit check.

Everything else from the previous round is unchanged: no env var, the CUDA-graph restriction with the eager regression numbers behind it, and no end-to-end speedup claimed.

Measurements are on the container's vLLM build, which is older than main; the same fused_topk_bias_router code is on main at lines 108-124 and 258-282.

[simondanielsson]

Thanks for making the changes!

On this:

The requested 8k/1k and 64k/1k points are still incomplete. Kimi-VL has a 131072 context, so the window is not the obstacle. Every request at 8192 and 65536 input tokens kills the engine in this container:

Kernel Name: _ZN2at6native24vectorized_gather_kernelILi16ElEEvPcS2_PT0_illllb
rocdevice.cpp:3586: HSA_STATUS_ERROR_EXCEPTION: hardware exception. code: 0x1016

Both the AITER baseline and this PR hit it: all four interleaved blocks logged one HSA exception and the same 392× HTTP 200 / 24× HTTP 500 split. The root cause has not been isolated. If someone with a long-context ROCm environment wants to run those two points I will take the numbers.

Suggestions:

1. Can you find another model that dispatched to this kernel and that doesn't crash at 8k+ context to validate the perf of this on?
2. For that model, please also re-run gpqa. I would have expected us to be unable to run GPQA on Kimi VL if it indeed crashes at 8k ISL?

[chennyiiis]

Both points are resolved, and the 8k/1k and 64k/1k numbers you asked for are now in the description. No code changed in this round — the previous revision stands.

On GPQA and Kimi-VL

Different models, so there was no contradiction: NIAH and GPQA ran on Moonlight-16B-A3B at 4096 and ~1k prompts, both far below the input length where the fault appeared. The description did not make that obvious, so Validation now opens with a table of which checkpoint carries which measurement.

Another model that reaches this kernel

Kimi-Linear-48B-A3B-Instruct: num_expert_group=1, topk_group=1, sigmoid router with a noaux_tc correction bias, moe_renormalize, 256 experts / top-8, context well beyond 64k. A probe confirms it reaches this kernel — 988 Triton dispatches at experts=256 topk=8 grp=1/1. 256 experts is the widest point in the microbenchmark, so it is the strongest case for the kernel rather than the weakest.

Why the long-context runs were failing

This turned out to be a bug in the benchmark harness, not in serving, and not in this PR.

vllm bench serve --dataset-name random samples token ids from tokenizer.vocab_size. For this checkpoint that is 163842, while the model's embedding has config.vocab_size = 163840 rows. Ids 163840 and 163841 therefore index past the embedding table. The dataset emits a contiguous run of ids, so the probability of hitting one scales with prompt length: a 65536-token prompt covers 40% of the vocabulary and nearly always includes one, an 8192-token prompt covers 5%.

That is exactly the behaviour we saw. On stock vLLM, with no part of this PR installed, the identical 8192-token configuration gives:

seed	completed (3 reps)	Memory access fault
7	8, 8, 8	0
1234	0, 0, 0	1

It also explains the earlier fault on a different machine with Kimi-VL, where the HSA exception named at::native::vectorized_gather_kernel — the embedding gather.

Replacing the synthesised prompts with real text (a custom dataset whose maximum token id is 138177, asserted below the embedding size) removes it completely: 8 of 8 runs complete at 64k with zero faults.

Happy to open a separate issue for the harness side of this if it is useful — the fix would be to clamp the sampled ids to the model's config.vocab_size rather than the tokenizer's.

The numbers

Interleaved B-A-B-A, one server per block, on Kimi-Linear:

in / out	concurrency	AITER tok/s	this PR tok/s	Δ throughput	Δ TPOT	same-arm drift
8192 / 1024	4	4305.4	4466.3	+3.7%	−4.0%	4.1% / 1.7%
8192 / 1024	8	6154.7	6162.5	+0.1%	−1.2%	1.4% / 1.8%
65536 / 1024	1	13528.3	14305.7	+5.7%	−7.0%	8.6% / 8.4%
65536 / 1024	2	18045.6	17927.5	−0.7%	−1.6%	1.7% / 0.9%

Every point has a same-arm drift the size of the difference between arms, so this does not change the conclusion: no end-to-end speedup is claimed. The 64k rows carry only 2 and 4 requests per run, which is why their drift is the largest here.

GPQA on the same model, three interleaved runs per arm: AITER 0.3889 / 0.3990 / 0.3990, this PR 0.4040 / 0.4192 / 0.4091. This PR is nominally higher in all three, by 1.0 point — I am not claiming that as an improvement, since 1.0 point is 2 questions out of 198, inside the ±3.4 point standard error, and a routing difference confined to 0.15% of calls has no mechanism to make a model more accurate. No regression; equivalence not established either.

Where that leaves the PR

The measurable effect remains what it was: the routing op is 1.67×–2.15× faster on gfx942 under CUDA graph replay, eager is excluded by the dispatch because it measured slower, and end to end is inside the noise on three models and four input lengths. That is the same basis on which #58354 landed, with its own "no end-to-end speedup is claimed", so I would rather state it plainly than stretch it.

[simondanielsson]

Let's discuss if we want to add this if it has no perf benefits. Happy to hear @tjtanaa's or @shen-shanshan's thoughts

[simondanielsson]

Suggestion: IMO we can drop this and test_flat_topk_falls_back_without_forward_context as it's not testing much

[simondanielsson]

Minor: let's use set_random_seed

[simondanielsson]

Suggestion: IMO we can drop this as well and just keep the correctness tests