[Perf][MoE] Optimize deepep_v2 receiver CPU Overhead - #51114
Merged
robertgshaw2-redhat merged 4 commits intoAug 17, 2026
Merged
Conversation
LucasWilkinson
requested review from
mgoin,
pavanimajety and
zyongye
as code owners
August 5, 2026 04:19
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
…epeat_interleave
In DeepEPV2PrepareAndFinalize._receiver, the do_expand=True (prefill) branch
built recv_topk_idx with a Python loop issuing one slice+fill_ per local
expert. Rows already arrive grouped by expert, so this is a repeat_interleave.
The loop cost one kernel launch per local expert, per MoE layer, per step. On
Kimi-K3-pruned75 (224 experts) at DP=2/TP=4/EP=8 that measured 22 aten::fill_
and 146us of self CPU per layer-step, i.e. ~13ms per engine step across 92 MoE
layers, against 234ms of dispatch+combine GPU time.
Two details keep the replacement off the critical path:
* The repeats tensor is borrowed from expert_tokens_meta, whose counts are
already staged with .to(non_blocking=True); its construction moves above
recv_topk_idx so it can be reused. Building a separate one with
torch.tensor(list, device=cuda) is a pageable copy -- synchronous with
respect to the host -- and blocks until the enqueued GPU work drains:
52us on an idle GPU but 5249us behind a deep queue, worse than the loop.
* output_size is passed, so repeat_interleave does not sum the counts on
device and sync to learn the output length. It is known on the host
because do_expand=True implies do_cpu_sync=True.
arange(num_local_experts) + rank_expert_offset is rank-constant and cached.
Signed-off-by: Lucas Wilkinson <wilkinson.lucas@gmail.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
LucasWilkinson
force-pushed
the
lwilkinson/deepep-v2-receiver-repeat-interleave
branch
from
August 5, 2026 04:20
909949f to
770f8af
Compare
robertgshaw2-redhat
approved these changes
Aug 5, 2026
Collaborator
Author
|
/run ci |
tlrmchlsmth
approved these changes
Aug 6, 2026
mgoin
added
performance
Member
|
/ci run |
|
✅ Triggered Buildkite CI #82695 for commit |
robertgshaw2-redhat
enabled auto-merge (squash)
August 7, 2026 18:48
Collaborator
|
/ci run |
|
✅ Triggered Buildkite CI #82900 for commit |
Collaborator
|
/ci run |
|
✅ Triggered Buildkite CI #84222 for commit |
Contributor
|
Hi @LucasWilkinson, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, |
robertgshaw2-redhat
merged commit Aug 17, 2026
75dde08
into
vllm-project:main
107 of 110 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Optimize away excessive fills in deepep_v2's receiver
BEFORE
AFTER