[Bugfix][R3] Size monolithic routing replay buffer for DP by TomerBN-Nvidia · Pull Request #50874 · vllm-project/vllm

TomerBN-Nvidia · 2026-08-03T13:20:04Z

Purpose

Fix routing-replay capture for the FlashInfer monolithic MoE kernel under naive
data parallelism, including padded sequence-parallel shards when expert
parallelism is enabled.

Two related assumptions fail in a TP2/DP2 deployment:

Replay buffer capacity. max_num_tokens is a per-rank scheduler limit,
while the naive dispatch path all-gathers rank-local batches before invoking
the monolithic kernel. This produced an 8192-row replay buffer for a
16384-row kernel input and failed during warmup. The buffer is now sized for
the larger of the DP and EP dispatch groups, since under EP the gathered
batch is the flattened EP group rather than the DP group.
Gathered shard layout. Under DP+EP, topk_ids can contain gathered
sequence-parallel shards whose dp_metadata.local_sizes include
CUDA-graph/SP padding. Those rows do not match the unpadded per-DP token
counts used by the existing capture paths, so the batch-dimension assertion
fired. When the gathered tensor matches sum(local_sizes), we locate the
current DP rank's contiguous shard group using that exact all-gatherv layout
and copy only its real token count, trimming trailing padding.

Existing naive-DP, modular-local, and modular-SP layouts are unchanged, and no
routing or model-output semantics change.

Duplicate check

I searched open upstream PRs for routing replay buffer, routed experts capture, routing replay all-gather, and local_sizes routing; none fixes
this allocation or the padded-shard capture:

[Bugfix] Share FlashInfer B12x MoE workspace across layers #48698 scales a separate FlashInfer B12x MoE workspace and does not touch the
routing replay output added by [RL Infra][FlashInfer] Enable router replay output from FlashInfer monolithic MoE kernel #44214.
[R3] Unify routed expert shape configuration #50940 also edits routed_experts_capturer.py, but only the shape
configuration; it leaves the capture() batch-layout branches alone.
[Bugfix] Skip MoE padding across router paths #50759 skips MoE padding inside the router implementations, upstream of the
gathered-batch layout handled here.

Validation

git diff --check: passed
vLLM pre-commit on the changed files, run in the repo's Python 3.12 venv:
Ruff check/format, mypy-3.10, typos, SPDX, root-lazy-import,
forbidden-import, torch-CUDA-API, config-validation, and sign-off hooks all
passed. (Under a host Python 3.9 two of these hooks crash on 3.10+ syntax;
that is an environment artifact, not a finding.)
Runtime validation of the exact TP2/DP2+EP default-backend configuration was
run on the corresponding branch of my fork on GB200: mixed and decode
CUDA-graph capture completed, both API servers started, and a
/v1/completions request was served with routed-experts capture enabled. The
original routing batch-dimension assertion did not fire.
Model evaluation: not applicable; routing decisions and model outputs are
unchanged by this fix.

AI assistance was used. As the human submitter I have reviewed every changed
line and the recorded runtime validation.

Naive DP all-gathers per-rank token batches before the monolithic MoE kernel, so allocate routing replay output for the gathered maximum. Assisted-by: OpenAI Codex <noreply@openai.com> Signed-off-by: tbarnatan <tbarnatan@nvidia.com>

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.

TomerBN-Nvidia · 2026-08-03T13:22:00Z

@aoshen02 Can you please take a look?

aoshen02 · 2026-08-04T00:19:10Z

@aoshen02 Can you please take a look?

Will take a look today.

Under DP+EP the gathered routing tensor carries sequence-parallel shards whose sizes include CUDA-graph/SP padding, so size the monolithic replay buffer for the larger of the DP and EP dispatch groups and locate this DP rank's contiguous shard group from the exact all-gatherv layout. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: tbarnatan <tbarnatan@nvidia.com>

TomerBN-Nvidia · 2026-08-04T07:38:43Z

@aoshen02 I added some more logic that I needed. Can you please take another look? It passes my tests that include DP.

aoshen02 · 2026-08-04T07:40:48Z

@aoshen02 I added some more logic that I needed. Can you please take another look? It passes my tests that include DP.
sure.

aoshen02 · 2026-08-04T13:26:48Z

            return
+        # Naive DP and DP+EP prepare paths all-gather per-rank batches before
+        # invoking the monolithic kernel. ``max_num_tokens`` is the scheduler
+        # limit for one rank, so the replay output must cover the largest


A bit verbose.

aoshen02 · 2026-08-04T13:27:34Z

+        # dispatch group rather than only the local batch. Under EP that group
+        # can also contain sequence-parallel shards flattened into the EP
+        # group, including padding.
+        dispatch_group_size = max(


I think it would be more clearer to change to

if self.moe_config.use_ep: self.moe_config.ep_size else: self.moe_config.dp_size

aoshen02 · 2026-08-04T13:42:30Z

LGTM in general.

Assisted-by: OpenAI Codex Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

aoshen02 · 2026-08-13T03:47:47Z

/ci run

github-actions · 2026-08-13T03:48:04Z

✅ Triggered Buildkite CI #83664 for commit a4f624a76984.

ZJY0516

approved by @aoshen02

…ct#50874) Signed-off-by: tbarnatan <tbarnatan@nvidia.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>

…ct#50874) Signed-off-by: tbarnatan <tbarnatan@nvidia.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>

…ct#50874) Signed-off-by: tbarnatan <tbarnatan@nvidia.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: khushali9 <khushali.desai9@gmail.com>

[Bugfix] Size monolithic routing replay buffer for DP

34ecf50

Naive DP all-gathers per-rank token batches before the monolithic MoE kernel, so allocate routing replay output for the gathered maximum. Assisted-by: OpenAI Codex <noreply@openai.com> Signed-off-by: tbarnatan <tbarnatan@nvidia.com>

mergify Bot added the bug Something isn't working label Aug 3, 2026

TomerBN-Nvidia marked this pull request as ready for review August 3, 2026 13:21

TomerBN-Nvidia requested review from mgoin, pavanimajety and zyongye as code owners August 3, 2026 13:21

claude Bot reviewed Aug 3, 2026

View reviewed changes

aoshen02 added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 4, 2026

aoshen02 mentioned this pull request Aug 4, 2026

[R3] Unify routed expert shape configuration #50940

Merged

TomerBN-Nvidia changed the title ~~[Bugfix] Size monolithic routing replay buffer for DP~~ [Bugfix][R3] Size monolithic routing replay buffer for DP Aug 4, 2026

aoshen02 reviewed Aug 4, 2026

View reviewed changes

Address routing replay review feedback

a4f624a

Assisted-by: OpenAI Codex Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

ZJY0516 approved these changes Aug 13, 2026

View reviewed changes

ywang96 approved these changes Aug 13, 2026

View reviewed changes

ywang96 merged commit 399f974 into vllm-project:main Aug 13, 2026
105 of 112 checks passed

Navigation Menu

Uh oh!

Uh oh!

[Bugfix][R3] Size monolithic routing replay buffer for DP - #50874

[Bugfix][R3] Size monolithic routing replay buffer for DP#50874
ywang96 merged 3 commits into
vllm-project:mainfrom
TomerBN-Nvidia:tbn/fix-routing-replay-buffer-dp-main

TomerBN-Nvidia commented Aug 3, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

TomerBN-Nvidia commented Aug 3, 2026

Uh oh!

aoshen02 commented Aug 4, 2026

Uh oh!

TomerBN-Nvidia commented Aug 4, 2026

Uh oh!

aoshen02 commented Aug 4, 2026

Uh oh!

aoshen02 Aug 4, 2026

Uh oh!

aoshen02 Aug 4, 2026 •

edited

Loading

Uh oh!

aoshen02 commented Aug 4, 2026

Uh oh!

aoshen02 commented Aug 13, 2026

Uh oh!

github-actions Bot commented Aug 13, 2026

Uh oh!

ZJY0516 left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Uh oh!

Conversation

TomerBN-Nvidia commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Duplicate check

Validation

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

TomerBN-Nvidia commented Aug 3, 2026

Uh oh!

aoshen02 commented Aug 4, 2026

Uh oh!

TomerBN-Nvidia commented Aug 4, 2026

Uh oh!

aoshen02 commented Aug 4, 2026

Uh oh!

aoshen02 Aug 4, 2026

Choose a reason for hiding this comment

Uh oh!

aoshen02 Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aoshen02 commented Aug 4, 2026

Uh oh!

aoshen02 commented Aug 13, 2026

Uh oh!

github-actions Bot commented Aug 13, 2026

Uh oh!

ZJY0516 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

TomerBN-Nvidia commented Aug 3, 2026 •

edited

Loading

aoshen02 Aug 4, 2026 •

edited

Loading