Skip to content

[Kernels] Overlap shared experts with send/recv #23273

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
robertgshaw2-redhat merged 28 commits into from
Sep 3, 2025
Merged

[Kernels] Overlap shared experts with send/recv #23273

robertgshaw2-redhat merged 28 commits into from
Sep 3, 2025

Conversation

@bnellnm
Copy link
Collaborator

@bnellnm bnellnm commented Aug 20, 2025 •
edited by github-actions bot
Loading

Purpose

Overlap shared experts computation with send and receive operations in all2all dispatcher.

TODO

  • Implement prepare_async for flashinfer prepare/finalize class?
  • Move sum of shared/fused experts and final reduce of into the FusedMoE layer.

Test Plan

Tried it with llama4 and deepseek_v2
Added unit test for shared experts in test_pplx_moe.py

Test Result

This trace shows the MLP (nvjet...) overlapped with DeepEP dispatch send/receive (green).

Screenshot 2025-08-29 at 3 19 55 PM

Documentation Update

Updated modular kernel docs

cc @SageMoore , @LucasWilkinson

@mergify mergify bot added the llama Related to Llama models label Aug 20, 2025
@mergify
Copy link

mergify bot commented Aug 20, 2025

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @bnellnm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Aug 20, 2025
@github-actions
Copy link

github-actions bot commented Aug 20, 2025

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

@bnellnm bnellnm changed the title Interleave shared experts with send/recv WIP [Kernels] Interleave shared experts with send/recv Aug 20, 2025
@bnellnm bnellnm marked this pull request as ready for review August 21, 2025 13:45
@mergify mergify bot added the documentation Improvements or additions to documentation label Aug 21, 2025
@bnellnm bnellnm changed the title [Kernels] Interleave shared experts with send/recv [Kernels] Overlap shared experts with send/recv Aug 21, 2025
@mergify mergify bot added the deepseek Related to DeepSeek models label Aug 22, 2025
@bnellnm bnellnm force-pushed the shared-fused-experts branch from d19e631 to 679ff7b Compare August 23, 2025 00:53
@mergify mergify bot removed the needs-rebase label Aug 23, 2025
@bnellnm bnellnm requested a review from WoosukKwon as a code owner August 27, 2025 03:50
LucasWilkinson
LucasWilkinson reviewed Aug 27, 2025
View reviewed changes
LucasWilkinson
LucasWilkinson reviewed Aug 28, 2025
View reviewed changes
SageMoore
SageMoore reviewed Aug 28, 2025
View reviewed changes
Copy link
Contributor

@SageMoore SageMoore left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice work, @bnellnm. I only have minor nits. Otherwise looks good.

LucasWilkinson
LucasWilkinson approved these changes Aug 29, 2025
View reviewed changes
Copy link
Collaborator

@LucasWilkinson LucasWilkinson left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Would be good to get a trace though to show the overlap

@bnellnm bnellnm force-pushed the shared-fused-experts branch from 1e81a2c to 5939843 Compare August 29, 2025 18:16
@mergify
Copy link

mergify bot commented Aug 30, 2025

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @bnellnm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Aug 30, 2025
@bnellnm bnellnm force-pushed the shared-fused-experts branch from 17eeab8 to e2a93cc Compare September 1, 2025 19:33
@mergify mergify bot removed the needs-rebase label Sep 1, 2025
hmellor
hmellor reviewed Sep 2, 2025
View reviewed changes
docs/design/fused_moe_modular_kernel.md Outdated
Copy link
Member

@hmellor hmellor Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thunk?

Copy link
Collaborator Author

@bnellnm bnellnm Sep 2, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A parameterless function. I can change the wording if it's not clear.

bnellnm added 11 commits September 3, 2025 02:35
@bnellnm
use more realistic MLP for test
a9548a1 
Signed-off-by: Bill Nell <bnell@redhat.com>
@bnellnm
fix cudagraphs
31f2d05 
Signed-off-by: Bill Nell <bnell@redhat.com>
@bnellnm
undo debugging code in activation.py
79df8e7 
Signed-off-by: Bill Nell <bnell@redhat.com>
@bnellnm
review comments + cleanup deepseek changes
544b62a 
Signed-off-by: Bill Nell <bnell@redhat.com>
@bnellnm
review comments
d617c83 
Signed-off-by: Bill Nell <bnell@redhat.com>
@bnellnm
remove debugging env var
46b2c3b 
Signed-off-by: Bill Nell <bnell@redhat.com>
@bnellnm
add all_reduce to fallback code
cc11426 
Signed-off-by: Bill Nell <bnell@redhat.com>
@bnellnm
fix data_parallel.py compilation config default
8173209 
Signed-off-by: Bill Nell <bnell@redhat.com>
@bnellnm
add check for SharedFusedMoE layer
9dc9205 
Signed-off-by: Bill Nell <bnell@redhat.com>
@bnellnm
ignore shared expert weights
2d5561f 
Signed-off-by: Bill Nell <bnell@redhat.com>
@bnellnm
ignore shared_expert parameters
068ad88 
Signed-off-by: Bill Nell <bnell@redhat.com>
auto-merge was automatically disabled September 3, 2025 02:40

Head branch was pushed to by a user without write access

@bnellnm bnellnm force-pushed the shared-fused-experts branch from 4f87a94 to 068ad88 Compare September 3, 2025 02:40
@robertgshaw2-redhat robertgshaw2-redhat enabled auto-merge (squash) September 3, 2025 02:43
@robertgshaw2-redhat
Copy link
Collaborator

robertgshaw2-redhat commented Sep 3, 2025

/ready

@bnellnm
make doc clearer
b94296f 
Signed-off-by: Bill Nell <bnell@redhat.com>
auto-merge was automatically disabled September 3, 2025 02:44

Head branch was pushed to by a user without write access

@bnellnm
Copy link
Collaborator Author

bnellnm commented Sep 3, 2025

@SageMoore , @LucasWilkinson could you guys take a final look? Should be ready to merge. @robertgshaw2-redhat verified that it works on multi-node systems now.

@robertgshaw2-redhat robertgshaw2-redhat merged commit e9b92dc into vllm-project:main Sep 3, 2025
54 checks passed
@bnellnm bnellnm mentioned this pull request Sep 3, 2025
Merged
eicherseiji pushed a commit to eicherseiji/vllm that referenced this pull request Sep 9, 2025
@bnellnm @eicherseiji
[Kernels] Overlap shared experts with send/recv (vllm-project#23273)
270c4d7 
Signed-off-by: Bill Nell <bnell@redhat.com>
@yewentao256
Copy link
Member

yewentao256 commented Sep 9, 2025

This PR introduces the R1 accuracy issue #24530

@whx-sjtu whx-sjtu mentioned this pull request Sep 15, 2025
Merged
5 tasks
FeiDaLI pushed a commit to FeiDaLI/vllm that referenced this pull request Sep 25, 2025
@bnellnm @FeiDaLI
[Kernels] Overlap shared experts with send/recv (vllm-project#23273)
d41bbb7 
Signed-off-by: Bill Nell <bnell@redhat.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@hmellor hmellor hmellor left review comments

@LucasWilkinson LucasWilkinson LucasWilkinson approved these changes

@mgoin mgoin Awaiting requested review from mgoin mgoin is a code owner

@robertgshaw2-redhat robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat robertgshaw2-redhat is a code owner

@tlrmchlsmth tlrmchlsmth Awaiting requested review from tlrmchlsmth tlrmchlsmth is a code owner

@yewentao256 yewentao256 Awaiting requested review from yewentao256 yewentao256 is a code owner

@WoosukKwon WoosukKwon Awaiting requested review from WoosukKwon WoosukKwon is a code owner

@njhill njhill Awaiting requested review from njhill

@ywang96 ywang96 Awaiting requested review from ywang96

@comaniac comaniac Awaiting requested review from comaniac

@alexm-redhat alexm-redhat Awaiting requested review from alexm-redhat

+1 more reviewer

@SageMoore SageMoore SageMoore approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

deepseek Related to DeepSeek models documentation Improvements or additions to documentation llama Related to Llama models ready ONLY add when PR is ready to merge/full CI is needed v1

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@bnellnm @robertgshaw2-redhat @yewentao256 @SageMoore @LucasWilkinson @hmellor