[MoE][Bugfix] Skip SP padded rows in grouped MoE routing by elvircrn · Pull Request #56079 · vllm-project/vllm

elvircrn · 2026-09-09T12:55:24Z

The gist of the issue is that sequence parallelism needs to pad token counts divisible by TP size and that vllm incorrectly doesn't mark these padded tokens as ignored.

Purpose

While running Kimi k3 e2e experiments, an imbalance in expert routing was observed:

making rank 0 a straggler rank:

tanking prefill performance.

With this patch, we get the following accumulated expert token counts during a subset of the 15-min aiperf benchmark run with claude code traces:

and get the following e2e speed-ups:

with the following gsm8k scores:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9727|±  |0.0045|
|     |       |strict-match    |     5|exact_match|↑  |0.9735|±  |0.0044|

Test Plan

Test Result

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

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.

tlrmchlsmth

Looks like there is an issue with DBO -- details via codex

Details

Review comment:

[P2] Propagate the mask into DBO's microbatch contexts — /Users/tms/code/vllm/vllm/v1/worker/gpu_model_runner.py:4555-4555
When the legacy runner enables DBO, _model_forward enters UBatchWrapper, whose _make_ubatch_metadata() replaces this outer context with per-microbatch contexts created without is_padding. Kimi's
sp_padding_mask therefore treats existing DP/CUDA-graph padding in each microbatch as real tokens, and the grouped router still dispatches those rows under DeepEP v2. Pass the appropriately sliced mask
into each DBO context as the v2 runner does.

DBO needs to be handled in the legacy vllm/v1/worker/gpu_ubatch_wrapper.py:252, because it replaces the outer forward context.

The flow should be:

gpu_model_runner.py puts the full-batch is_padding mask on the outer context.
UBatchWrapper.call() (vllm/v1/worker/gpu_ubatch_wrapper.py:375) reads forward_context.is_padding.
Both _make_ubatch_metadata() calls pass that mask onward.
Each per-microbatch create_forward_context() gets:

is_padding=(
is_padding[ubatch_slice.token_slice]
if is_padding is not None
else None
)

The two paths to cover are:

CUDA-graph capture at line 406 (vllm/v1/worker/gpu_ubatch_wrapper.py:406)
Ordinary DBO execution at line 432 (vllm/v1/worker/gpu_ubatch_wrapper.py:432)

That mirrors the v2 implementation in vllm/v1/worker/gpu/ubatch_utils.py:351, where each microbatch context receives its sliced is_padding tensor. Merely passing the mask into the outer context in
gpu_model_runner.py is insufficient because DBO overrides it.

Otherwise LGTM

Propagate persistent padding masks through both GPU model runners. Opt grouped-router invalidation in only for Kimi-K3 and require the resolved MoE backend to be DeepEP v2. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

github-actions · 2026-09-14T12:37:09Z

✅ @elvircrn, CI is now available for this PR.

/ci run starts upstream CI; /amd-ci run starts AMD CI only.
/ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
/amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
/ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

tlrmchlsmth · 2026-09-14T12:43:41Z

/ci run

github-actions · 2026-09-14T12:43:58Z

✅ Triggered Buildkite CI #88818 for commit 28e965fdb8fc.

tlrmchlsmth · 2026-09-16T18:42:14Z

/ci run

github-actions · 2026-09-16T18:42:56Z

✅ Triggered Buildkite CI #89406 for commit 6796c4e32dd0.

tlrmchlsmth · 2026-09-17T19:03:23Z

/ci retry

github-actions · 2026-09-17T19:03:41Z

✅ Triggered Buildkite CI #89695 for commit 1b01ae9b7dc1, running 1 failed step(s) from Buildkite CI #89406.

Root cause: upstream vllm#56079 added a `skip_padding: bool = False` kwarg to create_fused_moe_router()/GroupedTopKRouter to invalidate sequence-parallel padding rows in grouped MoE routing. The vllm-gaudi override mirrors this factory's signature and did not accept the new kwarg, so every call site now raises: TypeError: create_fused_moe_router() got an unexpected keyword argument 'skip_padding' Upstream: vllm-project/vllm#56079 Fix: accept skip_padding on the override (signature parity) and forward it to GroupedTopKRouter; the guard it enables only fires under VLLM_MOE_SKIP_PADDING + DeepEP-v2-kernel forward context, which HPU never exercises, so this is a no-op on HPU beyond restoring the call signature. Signed-off-by: Paweł Olejniczak <pawelx.olejniczak@intel.com>

…d skip_padding kwarg, drop stale manager.touch() (+3 more) (#1806) This PR consolidates 5 hourly-CI fixes against vllm@`9a70c233cd5d3ef3b14d79c582a6f9ddfa03cf80`. ## Bug 1: forward skip_padding kwarg through HPU MoE router factory override - **State machine id**: create_fused_moe_router_unexpected_kwarg - **Commit**: 65cc65e ### Root cause upstream vllm#56079 added a `skip_padding: bool = False` kwarg to create_fused_moe_router()/GroupedTopKRouter to invalidate sequence-parallel padding rows in grouped MoE routing. The vllm-gaudi override mirrors this factory's signature and did not accept the new kwarg, so every call site now raises: TypeError: create_fused_moe_router() got an unexpected keyword argument 'skip_padding' ### Culprit Regression introduced by [PR #56079](vllm-project/vllm#56079). ### Fix accept skip_padding on the override (signature parity) and forward it to GroupedTopKRouter; the guard it enables only fires under VLLM_MOE_SKIP_PADDING + DeepEP-v2-kernel forward context, which HPU never exercises, so this is a no-op on HPU beyond restoring the call signature. ## Bug 2: drop stale manager.touch() assertions in offloading connector test - **State machine id**: offloading_connector_touch_assertion - **Commit**: c5c5c55 ### Root cause upstream vllm#51787 ("Track cache recency once per request") removed OffloadingConnectorScheduler._touch and every manager.touch() call site, replacing per-call cache-recency tracking with req_context.set_offload_key_position() tracked once per request. vllm-gaudi's vendored copy of the unit test still asserted the old manager.touch() calls, so it now fails: AssertionError: Expected 'touch' to have been called. ### Culprit Probable culprit: [PR #51787](vllm-project/vllm#51787) - pinned by symbol archaeology, bisect not run. Candidate range: [compare](vllm-project/vllm@4be3dcf...9a70c23). ### Fix mirror upstream's own post-#51787 test update (its analogous tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py) by dropping the removed-call assertions and asserting touch is no longer called for the same scenario. ## Bug 3: add missing release_kv_cache_memory to MultiModelEngineClient - **State machine id**: online_model_swap_multimodelengineclient_abstract - **Commit**: 7ba9ac9 ### Root cause upstream vllm#44890 added a new abstract method release_kv_cache_memory to EngineClient; MultiModelEngineClient never implemented it, so instantiation raised TypeError. ### Culprit Probable culprit: [PR #44890](vllm-project/vllm#44890) - pinned by symbol archaeology, bisect not run. Candidate range: [compare](vllm-project/vllm@4be3dcf...9a70c23). ### Fix delegate to the wrapped AsyncLLM instance, mirroring the existing sleep/wake_up passthrough pattern. ## Bug 4: skip is_sleeping assignment on read-only Executor property - **State machine id**: online_model_swap_multimodelengineclient_abstract - **Commit**: 198d7e9 ### Root cause upstream vllm#44890 (same PR as the MultiModelEngineClient abstract-method break) turned Executor.is_sleeping into a read-only property computed from sleeping_tags. gaudi_reconfigure_engine's _reset_executor_sleep_state still assigned model_executor.is_sleeping directly, raising AttributeError: property 'is_sleeping' of 'UniProcExecutor' object has no setter during online model swap. ### Culprit Probable culprit: [PR #44890](vllm-project/vllm#44890) - pinned by symbol archaeology, bisect not run. Candidate range: [compare](vllm-project/vllm@4be3dcf...9a70c23). ### Fix feature-detect whether is_sleeping is a class-level property before assigning it; clearing sleeping_tags already makes the property read False on the new executor. ## Bug 5: implement apply_lora_full_linear on PunicaWrapperHPU - **State machine id**: hpu_unit_tests_engine_core_init_failed - **Commit**: 18d4414 ### Root cause upstream vllm PR 53555 ("[LoRA] Support modules_to_save for sequence classification") added a new abstract method apply_lora_full_linear to PunicaWrapperABC and implemented it for CPU, GPU and XPU. PunicaWrapperHPU never implemented it, so it is no longer a concrete class and get_punica_wrapper() raises TypeError: Can't instantiate abstract class PunicaWrapperHPU without an implementation for abstract method 'apply_lora_full_linear' while loading any LoRA model on HPU, which the engine surfaces as the generic "Engine core initialization failed". ### Culprit Probable culprit: [PR #53555](vllm-project/vllm#53555) - pinned by symbol archaeology, bisect not run. Candidate range: [compare](vllm-project/vllm@4be3dcf...9a70c23). ### Fix implement the hook with plain torch ops, mirroring the CPU and GPU versions: route rows through the prompt-level sampler_indices mapping and hand off to the inherited _select_full_linear_output helper for the bias and module-enabled selection. Multiplying against every adapter and gathering the routed row afterwards keeps the weight operand statically shaped for HPU graph compilation. The new method only references _select_full_linear_output at call time, and its only caller (ClassificationHeadWithLoRA) postdates PR 53555, so the change stays compatible with older VLLM_STABLE_COMMIT builds. --------- Signed-off-by: Paweł Olejniczak <pawelx.olejniczak@intel.com>

elvircrn requested review from AndreasKaratzas, WoosukKwon, mgoin, njhill, pavanimajety, tlrmchlsmth, yewentao256 and zyongye as code owners September 9, 2026 12:55

claude Bot reviewed Sep 9, 2026

View reviewed changes

elvircrn force-pushed the fix/grouped-topk-padding-routing branch from 0d3e2c9 to 169a56b Compare September 9, 2026 13:25

mergify Bot added kimi k3 mrv2 Model Runner V2 specific labels Sep 9, 2026

tlrmchlsmth reviewed Sep 11, 2026

View reviewed changes

elvircrn and others added 2 commits September 13, 2026 18:16

Fix DBO padding mask propagation

28e965f

Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

elvircrn force-pushed the fix/grouped-topk-padding-routing branch from 169a56b to 28e965f Compare September 13, 2026 16:17

elvircrn mentioned this pull request Sep 13, 2026

Fix moe scratch buffer sizing for the TP>1 and DP>1 case #55105

Closed

tlrmchlsmth added the ready ONLY add when PR is ready to merge/full CI is needed label Sep 14, 2026

tlrmchlsmth approved these changes Sep 14, 2026

View reviewed changes

tlrmchlsmth changed the title ~~Skip padded rows in grouped MoE routing~~ [MoE][Bugfix] Skip SP padded rows in grouped MoE routing Sep 14, 2026

mergify Bot added the bug Something isn't working label Sep 14, 2026

mergify Bot and others added 2 commits September 14, 2026 13:06

Merge branch 'main' into fix/grouped-topk-padding-routing

22b09c3

Merge branch 'main' into fix/grouped-topk-padding-routing

6796c4e

tlrmchlsmth enabled auto-merge (squash) September 16, 2026 18:42

Merge branch 'main' into fix/grouped-topk-padding-routing

1b01ae9

tlrmchlsmth merged commit ac2f0ea into vllm-project:main Sep 17, 2026
13 checks passed

vllm-agent mentioned this pull request Sep 18, 2026

Revert "[MoE][Bugfix] Skip SP padded rows in grouped MoE routing" (#56079) #57471

Closed

pawel-olejniczak mentioned this pull request Sep 21, 2026

[FIX_FOR_VLLM_CUSTOM=9a70c233cd5d3ef3b14d79c582a6f9ddfa03cf80] forward skip_padding kwarg, drop stale manager.touch() (+3 more) vllm-project/vllm-gaudi#1806

Merged

Navigation Menu

Repository navigation

[MoE][Bugfix] Skip SP padded rows in grouped MoE routing - #56079

[MoE][Bugfix] Skip SP padded rows in grouped MoE routing#56079
tlrmchlsmth merged 5 commits into
vllm-project:mainfrom
elvircrn:fix/grouped-topk-padding-routing

elvircrn commented Sep 9, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

tlrmchlsmth left a comment •

edited

Loading

Uh oh!

github-actions Bot commented Sep 14, 2026

Uh oh!

tlrmchlsmth commented Sep 14, 2026

Uh oh!

github-actions Bot commented Sep 14, 2026

Uh oh!

tlrmchlsmth commented Sep 16, 2026

Uh oh!

github-actions Bot commented Sep 16, 2026

Uh oh!

tlrmchlsmth commented Sep 17, 2026

Uh oh!

github-actions Bot commented Sep 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

elvircrn commented Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

tlrmchlsmth left a comment • 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!

github-actions Bot commented Sep 14, 2026

Uh oh!

tlrmchlsmth commented Sep 14, 2026

Uh oh!

github-actions Bot commented Sep 14, 2026

Uh oh!

tlrmchlsmth commented Sep 16, 2026

Uh oh!

github-actions Bot commented Sep 16, 2026

Uh oh!

tlrmchlsmth commented Sep 17, 2026

Uh oh!

github-actions Bot commented Sep 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

elvircrn commented Sep 9, 2026 •

edited

Loading

tlrmchlsmth left a comment •

edited

Loading