[micah-wil]

AITER MegaMoEV2 (initially integrated in #59685) has two dispatch schemes:

1. Compact: a separate megamoe_prepare_compact kernel packs tokens before dispatch.
2. Fixed-slot: each token writes straight into a per-rank slot, with no prepare kernel.

AITER takes fixed-slot only when the runtime's max_tok_per_rank is at most FIXED_SLOT_MAX_MTPR. That limit is 255 by default and can be raised with AITER_MEGA_FIXED_SLOT_MAX_MTPR. In the initial integration, we built a single runtime with the compact path, so decode always pays for the compact prepare step.

This PR builds a second runtime and picks between the two paths at each decode step.

The decode runtime is sized max_tok_per_rank = prev_power_of_2(FIXED_SLOT_MAX_MTPR), which ends up being 128 by default. It's only built when AITER's bundle plan says that size actually uses fixed-slot. Otherwise behavior is unchanged. The fused shared expert makes 49 experts per rank, so no decode runtime is built with FSE enabled.

Perf and Accuracy

Without MTP

Concurrency	Tokens/rank/step	main TPOT p50	This PR TPOT p50	main tok/s	This PR tok/s
32	4 (fixed-slot)	40.1 ms	34.5 ms (−14%)	775	894 (+15%)
128	16 (fixed-slot)	49.1 ms	43.3 ms (−12%)	2490	2827 (+14%)
2048	256 (compact)	121.3 ms	121.9 ms	14456	13934

With MTP (DSpark, 3 draft tokens)

Concurrency	Tokens/rank/step	main TPOT p50	This PR TPOT p50	main ITL p50	This PR ITL p50
32	16 (fixed-slot)	15.4 ms	14.0 ms (−9%)	50.6 ms	44.9 ms (−11%)
128	64 (fixed-slot)	23.0 ms	22.0 ms (−4%)	67.6 ms	58.1 ms (−14%)
512	256 (compact)	45.8 ms	45.5 ms	113.9 ms	113.7 ms

At 2048 without MTP and 512 with MTP, both builds use the compact runtime, so they run the same kernels. So it is expected that there is no diff from the current main.

Accuracy (gsm8k 5-shot, flexible / strict, ±0.006)

	main	This PR
Without MTP	0.9530 / 0.9538	0.9545 / 0.9538
With MTP	0.9538 / 0.9538	0.9515 / 0.9515

MTP acceptance rate is 91–92% on both.

Edit: A useful note from @benenzhu

One note from that setup: with K6 drafts each decode request carries 7 tokens, so at concurrency 128 a rank holds 112+ tokens per step, about 143 on average in our runs because of uneven load. With the default FIXED_SLOT_MAX_MTPR=255 the decode runtime holds 128 tokens, and those steps fall back to the compact path. We ran with AITER_MEGA_FIXED_SLOT_MAX_MTPR=511, which gives a 256-token decode runtime. It may be worth mentioning this knob in the description.

[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.

[micah-wil]

/ci run

[github-actions]

✅ Triggered Buildkite CI #94957 for commit 4090b90b7b16.

[benenzhu]

LGTM. Built the same thing while tuning DeepSeek-V4-Pro on 8x MI355X (DP8 attention + EP8 AITER MegaMoE, AgentX agentic workload, DSpark K6): a second fixed-slot MegaMoEV2 runtime for decode, picked by the max token count across DP. This PR matches that design, including the DP-max selection, so it looks good to us.

One note from that setup: with K6 drafts each decode request carries 7 tokens, so at concurrency 128 a rank holds 112+ tokens per step, about 143 on average in our runs because of uneven load. With the default FIXED_SLOT_MAX_MTPR=255 the decode runtime holds 128 tokens, and those steps fall back to the compact path. We ran with AITER_MEGA_FIXED_SLOT_MAX_MTPR=511, which gives a 256-token decode runtime. It may be worth mentioning this knob in the description.