feat(deepep_v2): expanded (do_expand=True) prefill dispatch for DeepEP-V2

[cyhdmjzzy]

Motivation

Add a faster prefill path for the DeepEP-V2 (ElasticBuffer) MoE A2A backend
(--moe-a2a-backend deepep_v2) that uses DeepEP's expanded dispatch layout, gated by
SGLANG_DEEPEP_V2_PREFILL_DO_EXPAND (default True). DeepEP-V2 uses ElasticBuffer:
one buffer serves both prefill and decode, switching behavior at runtime via
do_cpu_sync / do_expand.

The existing deepep_v2 backend already had an "expanded non-masked" prefill branch, but
it was locked off at dispatch and used a slower fallback. This change opens and
optimizes that branch so prefill runs on the expanded layout, while leaving the decode
(expand + masked slab) path unchanged. The goal is a faster prefill MoE all-to-all
with identical model outputs. In production PD-disaggregated (prefill/decode split)
serving, decode runs on the V1 low_latency backend in a separate process (rationale in
Modifications → Why decode stays on V1); this PR's optimization deliberately targets
prefill only.

Modifications

Core

All changes live inside the expand branch; the non-expand and decode paths are
untouched.

• token_dispatcher/deepep_v2.py — unlock expand at dispatch. The backend bound
  use_masked = use_expand_layout = not is_extend, pinning prefill to non-expand; split
  them into use_masked = is_decode and
  use_expand_layout = is_decode or _prefill_expand_enabled so prefill can run expanded
  but never masked. Quantize the expanded recv scale column-major only under ue8m0
  (non-ue8m0 stays row-major for the GEMM's own tma_align). Derive num_sms from
  get_theoretical_num_sms when SGLANG_DEEPEP_V2_NUM_SMS is 0 (dispatch has no
  zero-SM fallback and would hang the NVLink barrier otherwise); use ElasticBuffer
  hybrid mode automatically when nnodes > 1 (direct is NVLink-only), still overridable
  via --deepep-v2-mode.
• kernels/ops/moe/ep_moe_kernels.py — add fill_m_indices_from_psum (label rows
  from the device psum, no cumsum / H2D / sentinel) and scale_expanded_rows_
  (in-place row weighting, stride-aware so it also handles the column-major fp8 input
  scale).
• moe_runner/deep_gemm.py — the expanded non-masked pre/post permute use the two
  new kernels; leave hidden_states_scale_tma_aligned False so the contiguous GEMM still
  runs tma_align_input_scale on the expanded recv scale. The expand path applies the
  top-k weight before combine (which ignores topk_weights in expand mode) by folding
  the weight into down_proj's fp32 input scale — down_proj is linear and fp8 dequant is
  fp8_q * scale, so a per-row scale factor is numerically exact, shrinking a ~112 µs
  pass over the [expanded, H] bf16 tensor to ~4 µs over the [expanded, H/128] scale.
  (Under ue8m0 the scale is a power of two and multiplying by the weight would need
  re-rounding, so that case instead weights down_output with scale_expanded_rows_
  before combine.)
• model_runner.py + model_runner_components/moe_ep_setup.py — prebuild the
  ElasticBuffer at deployment time (in init_cuda_graphs, after capture and before
  serving), skipping it when decode graph capture already built it, so the first request
  does not pay the ~2 GB symmetric alloc + cross-rank NCCL barrier.

Wiring

• models/mimo_v2.py — include is_deepep_v2() in the EP / _enable_a2a_moe
  branches. This is required: without it _enable_a2a_moe stays False, MoE runs
  forward_normal whose extra tensor_model_parallel_all_reduce (which forward_deepep
  skips, since dispatcher.combine already reduces across ranks) hangs the first
  cross-node request.
• arg_groups/moe_hook.py — add MiMoV2ForCausalLM / MiMoV2FlashForCausalLM to
  the validated deepep_v2 architectures; divide the prefill token-budget check by
  attn_tp_size (a per-DP chunk scatters across the DP group's attn-TP ranks before
  dispatch, so a per-DP chunk of N needs N / attn_tp_size per EP rank, not N).
• environ.py — add SGLANG_DEEPEP_V2_PREFILL_DO_EXPAND and
  SGLANG_DEEPEP_V2_GPU_TIMEOUT_SECS. The per-rank dispatch-token cap keeps the existing
  SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK unchanged.

SGLANG_DEEPEP_V2_PREFILL_DO_EXPAND (default: on)

Prefill dispatch has two layouts; the only differing dispatch argument is
do_expand (do_cpu_sync stays True in both, so shapes are exact and neither uses
CUDA graph):

• do_expand=False — dispatch returns recv_x=[num_recv, H], one row per
  deduplicated token, not grouped by expert. A bridge kernel ep_scatter
  (inherited from V1) physically expands it into per-(token, expert) rows grouped by
  expert for the contiguous grouped GEMM. ep_scatter is two triton kernels: a light
  layout kernel (ep_scatter_1, ~1.7 µs) and a heavy full-hidden copy (ep_scatter_2,
  ~58 µs), plus one H2D of the per-expert counts. This path's operator chain is
  identical to the community's original non-expand design (ep_scatter → contiguous
  GEMM → ep_gather), for easy A/B comparison against do_expand=True (here weighting
  is folded into ep_gather, with no down_proj scale folding).
• do_expand=True (default) — DeepEP's dispatch_copy_epilogue produces the
  per-(token, expert) rows grouped by expert within dispatch's own TMA copy, so
  ep_scatter is not needed. The heavy ep_scatter_2 (~58 µs) and that H2D both
  disappear; the light layout kernel is replaced by fill_m_indices_from_psum (~1.5 µs,
  reusing DeepEP's device-side psum instead of recomputing a cumsum). It removes one
  full HBM round-trip per MoE layer at no accuracy cost.

dispatch_copy_epilogue_impl produces different outputs based on the do_expand
template parameter (nsys traces):

V1 dispatch:

V2 dispatch (do_expand=False):

V2 dispatch (do_expand=True):

V2 combine (do_expand=True, topk weighted before combine):

V2 combine (do_expand=True, topk weight folded into down_proj's input scale) —
_fwd_kernel_scale_expanded_rows moved from before combine to before down_proj, acting
only on the small fp8 input scale:

In expand mode ElasticBuffer's combine refuses topk_weights (it does the un-expand and
reduction itself inside the kernel), so the top-k weighting must be done by SGLang before
combine. Instead of scaling down_proj's bf16 output, this design folds each row's
routing weight w into down_proj's fp32 input scale down_input_scale=[M, I/128]
([~18432, 16]) before the GEMM: down_proj is linear and fp8 dequant is
value = q × scale, and w is one positive scalar per row, so it merges into the input
scale, the output carries the weight automatically, and the result is numerically exact.
The weighted tensor shrinks from [M, H] bf16 ([~18432, 6144], ~453 MB) to
[M, I/128] fp32 (~1.2 MB), and the multiply folds into the pre-GEMM handling that
already runs — measured ~112 µs → ~4 µs (ue8m0-packed scales fall back to weighting
the output).

Default is True because it is strictly less work than False with identical results;
set 0 to fall back to the validated non-expanded path for A/B comparison. "Less work"
here means do_expand=True versus do_expand=False:

• do_expand=False needs ep_scatter to physically expand the deduplicated recv_x
  into the per-expert layout; its heavy ep_scatter_2 reads and writes all of hidden
  once (a ~58 µs full HBM round-trip), plus one H2D copying the per-expert counts
  back to the device.
• do_expand=True lets DeepEP's dispatch_copy_epilogue do the expansion within
  dispatch's own TMA copy, so ep_scatter_2 (~58 µs) and that H2D both disappear,
  leaving only a ~1.5 µs fill_m_indices_from_psum to label rows. That removes one full
  hidden HBM round-trip per MoE layer (fewer kernels, one fewer large-tensor pass), with
  the GEMM and expert compute unchanged and results identical.

Why decode stays on V1 low_latency (not DeepEP-V2)

This PR's optimization targets prefill only. The deepep_v2 decode (expand + masked slab)
path is left untouched; in production PD-disaggregated serving, decode runs on V1's
low_latency mode in a separate process.

Decode requires a masked GEMM (fixed-height slab [num_local_experts, max_m, H]) to be
CUDA-graph-capturable:

	V1 low_latency	V2 direct
Output layout	packed_recv_x = [num_local_experts, num_ranks×cap, H] — fixed-height grid	compact per-expert buckets — dynamic bucket lengths
Feeds masked GEMM directly?	Yes — the fixed-height layout is the masked-GEMM input	No — buckets don't align with e × max_m
Bridge kernels needed	None	Two: expand_to_masked_slab (before GEMM) + masked_slab_to_expand (after)

V1's low_latency output is natively the masked-GEMM input, so decode needs no bridge
kernels. V2's direct output must be reshaped into a slab and back — two extra HBM
round-trips every layer, every step, with no offsetting benefit in SGLang (unlike vLLM,
whose V2 decode gain comes from a decode-only activation kernel SGLang does not use).

Accuracy Tests

Verified with --disaggregation-mode null (single process; deepep_v2 drives both
prefill and decode MoE dispatch — a stricter check than the production split), GSM8K
5-shot, 1319 questions, greedy decoding, MiMoV2, tp16/dp2:

python3 -m sglang.test.few_shot_gsm8k --host 127.0.0.1 --port 9001 \
    --num-questions 1319 --num-shots 5 --parallel 128

MoE a2a backend	Accuracy	Invalid
deepep_v1 (baseline)	0.958	0.002
deepep_v2	0.964	0.002

The 0.006 gap (~8 of 1319 questions) is within greedy-decoding noise, so DeepEP-V2 has
no accuracy regression. (Production uses V2 for prefill + V1 low_latency for
decode; both paths are independently validated.)

Speed Tests

End-to-end prefill throughput, MiMoV2, TP8/EP16/DP2, chunk 32k per DP, bench_serving's
Input token throughput (tok/s). "single DP throughput" = all-throughput / 2. A
three-way comparison: V1 baseline, the community's existing deepep_v2 (do_expand=False),
and this PR's final design (do_expand=True with the topk weight folded into down_proj's
input scale, weighted-before-down_proj). Δ is the gain over V1.

request_lens (k)	request count	V1 (tok/s)	V2 community do_expand=False (tok/s)	Δ vs V1	V2 this PR weighted-before-down_proj (tok/s)	Δ vs V1
8	512	25255.66	27028.59	+7.0%	29106.99	+15.2%
16	512	25313.40	27189.06	+7.4%	28690.07	+13.3%
32	256	24334.59	26097.61	+7.2%	27480.15	+12.9%
64	256	22680.25	24432.20	+7.7%	25692.17	+13.3%
128	128	20537.96	21686.12	+5.6%	22272.49	+8.4%
256	64	16922.24	17753.09	+4.9%	18226.55	+7.7%
512	24	12291.39	12766.78	+3.9%	13027.60	+6.0%
768	16	9707.58	10036.62	+3.4%	10121.12	+4.3%
960	8	8175.53	8345.79	+2.1%	8585.99	+5.0%

Two takeaways:

• Versus V1: the community deepep_v2 (do_expand=False) is already ~2–8% faster than
  V1; this PR pushes the gain to ~4–15% (~13–15% at short/medium sequence lengths,
  narrowing to ~4–6% at very long ones as attention's O(L²) cost dominates and the MoE
  all-to-all share shrinks).
• Versus the community deepep_v2: this PR is further ahead at every sequence length
  (e.g. 8k: 27028 → 29107, +7.7%; 512k: 12767 → 13028, +2.0%), from the ep_scatter_2
  full-hidden copy that do_expand=True removes plus folding the topk weight into
  down_proj's input scale (shrinking a ~112 µs large-tensor weighting to ~4 µs). The
  rationale for each is in the SGLANG_DEEPEP_V2_PREFILL_DO_EXPAND section above and in
  weighted-before-down_proj.md.

Checklist

• Format your code according to the Format code with pre-commit.
• Add tests according to the Run and add unit tests. Validated end-to-end on a multi-node multi-GPU setup, and the GSM8K accuracy check above passed.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Notes for reviewers

Runtime environment

Developed and validated on (2-node × 8-GPU, Hopper-class, ~140 GB HBM per GPU):

Component	Version / detail
CUDA	12.9 (toolkit 12.9.2; base image cu12)
NCCL	2.30.7+cuda12.9 — DeepEP-V2's GIN API (ncclGinRequest_t / ginTrafficClass) needs 2.30.x; the bundled 2.29.x lacks these symbols, so the image overrides every libnccl.so (including torch's) with 2.30.7
DeepEP	2.0.0 (the ElasticBuffer build exposing deep_ep.ElasticBuffer), built against torch 2.13 / NCCL 2.30.7, cp312
PyTorch	2.13
Python	3.12
sgl-kernel	0.4.5

Communication backend: DeepEP-V2 (ElasticBuffer) communicates over NCCL GIN
(GPU-Initiated Networking, reusing PyTorch's NCCL communicator), not NVSHMEM/IBGDA —
the V2 path has no NVSHMEM dependency (NVSHMEM is the DeepEP-V1 transport).
NCCL_CUMEM_ENABLE=1 is required by ElasticBuffer and is set automatically. Single node
uses --deepep-v2-mode direct (NVLink); multi-node needs hybrid (auto-selected when
nnodes > 1, or set explicitly).

---

CI States

Latest PR Test (Base): ❌ Run #34617090627
Latest PR Test (Extra): ❌ Run #34617089426
Latest PR Test (AMD ROCm 10): ❌ Run #34617090277

[ch-wan]

Summary

This PR unlocks DeepEP-V2's expanded (do_expand=True) prefill dispatch, replaces the ep_scatter bridge with fill_m_indices_from_psum, and folds top-k weights into down_proj's fp32 input scale. The expand/non-expand split, MiMoV2 A2A wiring, and ElasticBuffer prebuild look internally consistent with the existing decode masked path.

The main correctness risk is the new prefill token-budget math, which divides by attn_tp_size even when tokens are not scattered.

Two of the PR's stated DeepEP premises do not hold against upstream: ElasticBuffer.dispatch has mapped num_sms=0 to get_theoretical_num_sms since 2.0.0, so the "zero-SM hang" derivation reproduces DeepEP's own default; and DeepEP 2.1.0 (#674) accepts 1-D expanded weights in combine, which would remove the pre-combine weighting (fold, ue8m0 fallback, scale_expanded_rows_) entirely. The version dependency should be recorded now and a gated follow-up planned.

The two extra positionals passed to ep_scatter_from_psum are a standalone fix: main's only deepep_v2 prefill path has raised a missing-argument error since #35758. That deserves disclosure in the PR body and a unit test, since CI never reaches the wrapper.

Remaining items are style and conventions: domain logic inlined into the frozen model_runner.py, env-var naming/sentinel conventions, a duplicate startup check, and a few comments that misdescribe the contracts they sit next to.

Issue counts by severity

• bugs: 1
• suggestions: 10
• nits: 1

[ch-wan]

[bug] The new prefill budget divides max_prefill_buffer_tokens by attn_tp_size unconditionally (attn_dp_size = dp_size if enable_dp_attention else 1, so with DP-attention off attn_tp_size == tp_size). Tokens are only sharded across attn-TP ranks when DP-attention is on; without it each EP rank still dispatches the full chunk. Example: --tp 8 without --enable-dp-attention, chunked_prefill_size=2048, cap=512 now validates (ceil(2048/8)=256) and then dies in _validate_common on a 2048-token dispatch. The formula also omits attn_cp_size (derive_attention_widths uses tp / dp / cp), so it disagrees with the live get_parallel().attn_tp_size used in prebuild_deepep_v2_buffers. The new unit test only covers enable_dp_attention=True.

Suggestion: Divide only when tokens are actually scattered (DP-attention / CP). Reuse derive_attention_widths (or get_parallel().attn_tp_size) instead of recomputing tp_size // attn_dp_size. Add a case with enable_dp_attention=False, tp_size>1 that still requires the undivided chunk.

[cyhdmjzzy]

Fixed. The division now only applies when tokens are actually scattered:

if view.enable_dp_attention or view.attn_cp_size > 1:
    attn_dp_size = view.dp_size if view.enable_dp_attention else 1
    attn_tp_size = max(1, view.tp_size // attn_dp_size // view.attn_cp_size)
    prefill_tokens = -(-prefill_tokens // attn_tp_size)

So with DP-attention off each EP rank keeps the full chunk (no division), and the width now mirrors derive_attention_widths (tp / attn_dp / cp) including attn_cp_size, matching get_parallel().attn_tp_size used in prebuild. Added a regression case (enable_dp_attention=False, tp_size=16) that keeps the undivided 2048-token chunk and still trips the cap.

[ch-wan]

[suggestion] Prefill-node prebuild uses chunked_prefill_size // attn_tp_size (floor) and ignores dynamic-chunking's 1.25x probe, while validate_deepep_v2_dispatch_token_budget uses ceil(max_prefill_buffer_tokens / attn_tp_size). Remainder tokens (2049 / 8 -> 257 vs 256) can pass this check and still exceed the ElasticBuffer cap at dispatch. The check is also redundant with the server-args validator when the two formulas agree.

Suggestion: Drop the duplicate check, or share one helper with the validator (ceil of max_prefill_buffer_tokens, same attn-TP width).

[cyhdmjzzy]

Dropped the duplicate check. The server-args validator (validate_deepep_v2_dispatch_token_budget) is now the single source of truth for the prefill budget — ceil of max_prefill_buffer_tokens over the same attn-TP width, so the floor/1.25x-probe discrepancy is gone. prebuild no longer recomputes min_tokens_per_rank.

[ch-wan]

[suggestion] _get_allow_hybrid_mode is mode == "hybrid" or nnodes > 1, so --deepep-v2-mode direct is silently ignored on multi-node. Forcing hybrid is the right hang-avoidance default, but the CLI then does not mean what it says (the PR text also claims it remains overridable).

Suggestion: Treat nnodes > 1 as the default when mode is unset/auto, and either honor an explicit direct or reject it with a clear error that NVLink-only mode cannot run across nodes.

[cyhdmjzzy]

Done. _get_allow_hybrid_mode now treats unset/auto as "hybrid when nnodes > 1", honors an explicit "hybrid", and rejects an explicit "direct" across nodes with a clear error instead of silently overriding it:

if mode == "direct" and nnodes > 1:
    raise ValueError(
        "--deepep-v2-mode direct is NVLink-only and cannot run across "
        f"nodes (nnodes={nnodes}); use hybrid or leave it unset (auto)."
    )

I also removed the stale "remains overridable" wording from the PR text.

[ch-wan]

[suggestion] The three-line block restates why folding w into down_input_scale is exact and what ue8m0 does instead. That design rationale already lives in the PR; next to the if it is history, not a line-local constraint.

Suggestion: Keep at most one trap (ue8m0 is a power of two, so do not fold into the scale) attached to the not DEEPGEMM_SCALE_UE8M0 guard.

[cyhdmjzzy]

Trimmed to a single line-local trap on the not-ue8m0 guard (ue8m0 scales are powers of two, so folding into the scale would break them). The design rationale stays in the PR body.

[ch-wan]

[suggestion] The comment narrates an omitted DeepGemmRunnerInput argument (hidden_states_scale_tma_aligned left at default False) and the GEMM's follow-on tma_align_input_scale. The default is already False; the load-bearing fact is that expanded DeepEP recv scales are not in the contiguous-GEMM TMA layout.

Suggestion: One line on the recv-scale layout constraint, or drop it if _run_contiguous_gemm already always aligns when the flag is False.

[cyhdmjzzy]

Reduced to one line stating the load-bearing constraint: expanded DeepEP recv scales are not in the contiguous-GEMM TMA layout, so the flag stays False and _run_contiguous_gemm aligns them. Dropped the narration of the default value.

[ch-wan]

[suggestion] "combine ignores topk_weights in expand mode" is true for DeepEP 2.0.0 (its docstring: "Not used in expand mode"), but DeepEP 2.1.0 (#674, 2026-07-03) changed the contract: combine(topk_weights=...) accepts a 1-D [num_expanded_tokens] tensor in expand mode when the buffer was built with allow_multiple_reduction=True, which is the constructor default DeepEPv2Buffer.get_buffer already inherits (csrc/elastic/buffer.hpp takes the get_shape<1> path; tests/elastic/test_ep.py:227-229). On 2.1.0 the whole pre-combine weighting (the down-scale fold, the ue8m0 fallback and scale_expanded_rows_) collapses to buffer.combine(hidden_states, handle=..., topk_weights=recv_topk_weights) with no extra kernel.

Suggestion: State the version dependency in this comment and in the PR body now, and plan a follow-up that gates on the DeepEP version and hands recv_topk_weights straight to combine. That needs a bench, since allow_multiple_reduction changes the reduction epilogue (DeepEP documents a precision / transfer-size trade-off for it).

[cyhdmjzzy]

Recorded the version dependency in the comment and will note it in the PR body:

# DeepEP 2.0.0 combine ignores topk_weights in expand mode; 2.1.0+ (#674)
# accepts a 1-D weight there, so a follow-up can gate on version and pass
# recv_topk_weights straight to combine.

Agreed this collapses the down-scale fold + ue8m0 fallback + scale_expanded_rows_ on 2.1.0. I'll keep it as a gated follow-up rather than in this PR, since allow_multiple_reduction changes the reduction epilogue and DeepEP documents a precision/transfer-size trade-off, so it needs its own bench.

[ch-wan]

[suggestion] model_runner.py is a frozen orchestration-only file (.claude/skills/large-class-style/SKILL.md §1.3-1.5). decode_runner_captured is domain knowledge, not coordination: it encodes that a captured decode graph already built the ElasticBuffer and that EagerRunner means no capture happened. That reasoning belongs with the DeepEP-v2 setup code, and inlining it here also removes an override point for downstream forks.

Suggestion: Move the gate into the collaborator, e.g. maybe_prebuild_deepep_v2_buffers(model=..., decode_cuda_graph_runner=self.decode_cuda_graph_runner, disaggregation_mode=..., chunked_prefill_size=..., attn_tp_size=...) doing the EagerRunner check itself, and leave init_cuda_graphs with the single delegate call (or a maybe_init_deepep_v2_buffers helper method).

[cyhdmjzzy]

Moved. The EagerRunner gate now lives in the collaborator as maybe_prebuild_deepep_v2_buffers(model=..., decode_cuda_graph_runner=...), which does the capture check itself. model_runner.py keeps only the single delegate call, so the frozen file stays orchestration-only and downstream forks get their override point back.

[ch-wan]

[suggestion] Three points against .claude/skills/env-var-conventions/SKILL.md:

• SGLANG_DEEPEP_V2_PREFILL_DO_EXPAND is a boolean without a verb category (Rule 4: ENABLE_ / DISABLE_ / USE_ ...). SGLANG_DEEPEP_V2_ENABLE_PREFILL_EXPAND (default True) reads at the call site as "if enabled". Renaming later costs an alias entry, so it is cheapest now.
• SGLANG_DEEPEP_V2_GPU_TIMEOUT_SECS = EnvInt(0) uses 0 as an "unset" sentinel; Rule 2 asks for EnvInt(None) when unset must be distinguishable from a real value.
• The knob only raises num_gpu_timeout_secs. Prefill dispatch runs with do_cpu_sync=True, so the wait that actually times out while a peer rank is JIT-compiling is DeepEP's CPU-side num_cpu_timeout_secs (300 s default). If the goal is surviving first-request stalls, that timeout matters at least as much, or the comment should say why only the GPU barrier needs raising.

Also, this file's convention is bare declarations; the rationale for both flags is already in the PR body, so the two comment blocks can go.

[cyhdmjzzy]

All three addressed:

(1) Renamed to SGLANG_DEEPEP_V2_ENABLE_PREFILL_EXPAND (default True), reads as "if enabled" at the call site.

(2)+(3) Rather than fix the sentinel, I removed SGLANG_DEEPEP_V2_GPU_TIMEOUT_SECS entirely and no longer pass num_gpu_timeout_secs to ElasticBuffer (it uses DeepEP's default). You're right that the GPU barrier wasn't the right lever. The stall it papered over comes from ranks finishing their lazy DeepGEMM warmup at different times, so the first-layer dispatch barrier waits past the default. The correct fix is to remove that skew via DeepGEMM precompilation (pre-run sglang.compile_deep_gemm, or set SGLANG_JIT_DEEPGEMM_PRECOMPILE=0 at deploy time) so every rank hits a warm cache and the default timeout is ample — no knob needed.

Also dropped both comment blocks per the bare-declaration convention.

[ch-wan]

[nit] Two inaccuracies that will mislead the next editor. (1) The psum is not an exclusive prefix sum: DeepEP's handle docstring defines expand-mode psum[i] = align(psum[i-1]) + unaligned_count_i (inclusive, only the earlier experts aligned), which is exactly why both start and seg_end are rounded up here. (2) DeepEP sizes the expanded recv_x as the sum of aligned per-expert counts under do_cpu_sync=True (num_recv_tokens_per_expert_list holds aligned counts and buffer.hpp accumulates them into num_expanded_tokens), so total_rows == align(psum[-1]) and the last-expert seg_end = total_rows branch is a no-op. If it is meant to guard a capacity-sized buffer (do_cpu_sync=False), labelling that tail with a real expert id would make DeepGEMM compute the whole buffer instead of skipping -1 rows.

Suggestion: Fix the comment, and replace the branch with the total_rows % ALIGN == 0 assert the old kernel had.

[cyhdmjzzy]

Both fixed. The comment now says the psum is DeepEP's inclusive per-expert count (psum[i] = align(psum[i-1]) + count_i, only earlier experts aligned), which is why both start and seg_end round up. And I replaced the no-op last-expert branch with the total_rows % expert_alignment == 0 assert the old kernel had, since do_cpu_sync=True sizes recv_x as the sum of aligned counts (total_rows == align(psum[-1])).

[cyhdmjzzy]

Thanks for the very thorough review. I've addressed all 12 points; a per-item summary:

• [bug] prefill budget: now divides by the attn-TP width only when tokens are actually scattered (DP-attention or CP), and folds attn_cp_size in. Added a pure-TP regression case.
• The two DeepEP premises you flagged are fixed: the num_sms=0 derivation is removed (it reproduced DeepEP's own default), and the combine version dependency (2.0.0 vs 2.1.0 Fix sampling #674) is now recorded with a planned gated follow-up.
• ep_scatter_from_psum missing-argument fix: called out below and split into its own PR (fix(moe): pass expert_start/num_experts in ep_scatter_from_psum #38518) with a dedicated unit test, since CI never reaches this wrapper.
• Remaining style/convention items (frozen model_runner.py gate, env naming/sentinel, duplicate prebuild check, misleading comments) are all applied.

Details inline. The GPU-side timeout knob is removed entirely rather than fixed — see the env-var thread for the reasoning.

[cyhdmjzzy]

https://aie-shared-objects.cnbj1.mi-fds.com/aie-shared-objects/prefill-DeepEP-V2-weighted-before-down_proj.nsys-rep
This is the nsys trace captured when using DeepEP-V2 on the prefill side.

[ch-wan]

[suggestion] Prior finding 3 is only half-fixed. Explicit --deepep-v2-mode direct now raises on nnodes > 1 (good), but the comment, the error text ("leave it unset (auto)"), and the final return nnodes > 1 still describe an auto/unset hybrid default that does not exist. deepep_v2_mode is Literal["direct", "hybrid"] and defaults to "direct", so the fallback branch is dead, and following the error's "unset (auto)" advice still raises.

Suggestion: Drop the auto/unset wording and the dead return nnodes > 1. The error should tell the user to pass --deepep-v2-mode hybrid. If multi-node should auto-select hybrid, that belongs in resolution (change the default / add "auto"), not in this helper.

[cyhdmjzzy]

You're right — deepep_v2_mode is Literal["direct","hybrid"] with default "direct", so the fallback and the "unset (auto)" wording described a state that can't occur. Dropped the dead return nnodes > 1; explicit direct on nnodes>1 now raises asking for --deepep-v2-mode hybrid, otherwise it returns mode == "hybrid". Auto-selection belongs in resolution (default/"auto"), which I can do as a follow-up if wanted.

[ch-wan]

[suggestion] Prior finding 5 is still present after the field revert. The comment narrates that hidden_states_scale_tma_aligned is omitted and stays False. _run_contiguous_gemm already TMA-aligns when DEEPGEMM_NEED_TMA_ALIGNED_SCALES is set (and that flag is False under UE8M0), so the omitted kwarg is the dataclass default, not a local invariant that needs restating.

Suggestion: Delete the comment. If the default-False choice is the trap, a one-liner on the DeepGemmRunnerInput field is the right home, not a narration of a dropped kwarg.

[cyhdmjzzy]

Deleted. You're right — with the field reverted, hidden_states_scale_tma_aligned=False is just the dataclass default, and _run_contiguous_gemm already TMA-aligns when DEEPGEMM_NEED_TMA_ALIGNED_SCALES is set, so the comment only restated the default. Removed it rather than moving it, since the field name is self-describing and the alignment is generic, not a per-field trap.

[ch-wan]

[suggestion] Prior finding 4 is trimmed around the scale-fold site, but post-permute still embeds design history and an unowned follow-up: DeepEP 2.0.0 vs 2.1.0+ (#674), and "a follow-up can gate on version". That belongs in the PR body, not next to the call.

Suggestion: Keep at most the live constraint ("combine ignores topk_weights in expand mode, so rows must be weighted here unless already folded into the down-proj scale"). Drop the version changelog and the unowned follow-up, or make it TODO(<owner>): if it is planned work.

[cyhdmjzzy]

Trimmed to just the live constraint (combine ignores topk_weights in expand mode, so rows are weighted here unless already folded into the down-proj scale). Dropped the 2.0.0-vs-2.1.0 changelog and the unowned follow-up from the call site — the version dependency stays in the PR body, and I'll track the 2.1.0 gating there rather than as an inline note.