[DeepEP V2] Bound num_max_tokens_per_rank in do_expand=False

[WoosukKwon]

Purpose

In do_expand=False (decode / cudagraph) mode, DeepEPV2PrepareAndFinalize._do_dispatch left num_max_tokens_per_rank unset, so ElasticBuffer.dispatch fell back to the buffer's init capacity (= max_num_batched_tokens). The recv buffer was therefore sized to the worst case R * max_num_batched_tokens, and the expert kernels processed ~R * 8192 rows even when only a handful of decode tokens were present — which dominated decode step time.

This PR bounds num_max_tokens_per_rank to the DP-padded batch size (max(num_tokens_across_dp), uniform across ranks), rounded up to a power of two.

The pow2 bucketing is important: DeepEP JIT-compiles a separate dispatch kernel per distinct num_max_tokens_per_rank. Feeding the raw per-step size makes it recompile for every batch size — a cicc recompile storm that starves the GPU at high concurrency. Rounding up to a power of two bounds the compiled set to ~log2(max_num_batched_tokens) values (compiled once, then cached), while staying small for decode (e.g. 1 token → 1) and capped at the buffer's init capacity for prefill.

Prefill (do_expand=True) is unchanged — it keeps the existing num_max_tokens_per_rank=None / CPU-sync path.

Not a duplicate

Searched open PRs for num_max_tokens_per_rank, DeepEPV2, DeepEP v2, and deepep_v2. The related open DeepEP v2 PRs — #45282 (NVFP4 dispatch), #40718 (combine_v2), #45193 (topk_ids optional for do_expand), #45321 (Dockerfile version bump) — address different areas. None touch decode-mode recv-buffer sizing.

Test plan

• Lint/type: ruff check, ruff format, and mypy (via pre-commit) pass on the touched file.
• Benchmarked DeepSeek-V4-Flash on GB200 ×4, dp=4 -ep, DeepEP v2 + Triton MoE:
  • concurrency=1 decode step time: 100.7 ms → 16.4 ms (6.1×)
  • concurrency=64: no cicc recompile storm (previously ~0% GPU utilization during the storm)
  • concurrency=1024: ~+6% throughput
  • GSM8K accuracy unchanged (~95%)

Note

AI assistance (Claude) was used in preparing this change.

Signed-off-by: Woosuk Kwon woosuk@inferact.ai