Support mnnvl all2allv from Flashinfer by wenscarl · Pull Request #21003 · vllm-project/vllm

wenscarl · 2025-07-15T17:26:53Z

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.

Purpose

Test Plan

VLLM_ALL2ALL_BACKEND="flashinfer_all2allv" \
CUDA_VISIBLE_DEVICES=0,1,2,3 \
VLLM_USE_FLASHINFER_MOE_FP4=1 \
VLLM_FLASHINFER_MOE_BACKEND="throughput" \
  /home/shuw/.local/bin/vllm serve nvidia/DeepSeek-R1-FP4 \
    --quantization="modelopt_fp4" \
    --trust-remote-code \
    --max-model-len=2048 \
    --block-size=128 \
    --max-num-seqs=256 \
    --enable-expert-parallel \
    --gpu_memory_utilization=0.8 \
    --tensor-parallel-size 1 \
    --data-parallel-size 4
    
python benchmarks/benchmark_serving.py \
  --model nvidia/DeepSeek-R1-FP4 \
  --dataset-name random \
  --ignore-eos \
  --num-prompts 256 \
  --max-concurrency 256 \
  --random-input-len 128 \
  --random-output-len 1024

vs.
VLLM_ALL2ALL_BACKEND="naive" \

...

Test Result

accuracy:

VLLM_USE_FLASHINFER_MOE_FP4=1 VLLM_FLASHINFER_MOE_BACKEND="throughput" VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_ALL2ALL_BACKEND="flashinfer_all2allv" \
lm_eval --model vllm --model_args pretrained=nvidia/DeepSeek-R1-FP4,quantization=modelopt_fp4,data_parallel_size=4,enable_expert_parallel=True,tensor_parallel_size=1,enforce_eager=True,max_model_len=2048 --trust_remote_code --tasks gsm8k --num_fewshot 5 --batch_size auto

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9371|±  |0.0067|
|     |       |strict-match    |     5|exact_match|↑  |0.9340|±  |0.0068|

perf:
Alltoallv(this PR):

Successful requests:                     256       
Benchmark duration (s):                  111.43    
Total input tokens:                      32512     
Total generated tokens:                  262144    
Request throughput (req/s):              2.30      
Output token throughput (tok/s):         2352.63   
Total Token throughput (tok/s):          2644.42   
---------------Time to First Token----------------
Mean TTFT (ms):                          9734.68   
Median TTFT (ms):                        9769.39   
P99 TTFT (ms):                           9794.07   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          99.37     
Median TPOT (ms):                        99.34     
P99 TPOT (ms):                           99.35     
---------------Inter-token Latency----------------
Mean ITL (ms):                           99.37     
Median ITL (ms):                         90.51     
P99 ITL (ms):                            143.30

allgather-reducescatter

Successful requests:                     256       
Benchmark duration (s):                  134.83    
Total input tokens:                      32512     
Total generated tokens:                  262144    
Request throughput (req/s):              1.90      
Output token throughput (tok/s):         1944.28   
Total Token throughput (tok/s):          2185.42   
---------------Time to First Token----------------
Mean TTFT (ms):                          9614.91   
Median TTFT (ms):                        9648.26   
P99 TTFT (ms):                           9673.86   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          122.36    
Median TPOT (ms):                        122.33    
P99 TPOT (ms):                           122.34    
---------------Inter-token Latency----------------
Mean ITL (ms):                           122.36    
Median ITL (ms):                         122.26    
P99 ITL (ms):                            129.50

(Optional) Documentation Update

github-actions · 2025-07-15T17:27:02Z

👋 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.

🚀

mergify · 2025-07-15T17:27:28Z

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

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

gemini-code-assist

Code Review

This pull request introduces support for Flashinfer's mnnvl all2allv for Mixture-of-Experts (MoE) layers, which is a significant performance enhancement for distributed inference. The changes are comprehensive, touching custom ops, distributed communicators, the MoE layer implementation, and quantization methods.

The core of the change is the new FlashInferAllToAllManager and its integration into the MoE forward pass. The review focuses on potential issues like hardcoded values, code duplication, and correctness of the communication logic to ensure the new feature is robust and maintainable.

gemini-code-assist · 2025-07-15T17:28:32Z

vllm/distributed/device_communicators/all2all.py

+        gpus_per_node: int = 4, #TODO(shuw): remove hardcode
+    ):


The gpus_per_node parameter is hardcoded. This limits the flexibility of the implementation for different hardware configurations. Consider making this configurable.

vllm/distributed/device_communicators/all2all.py

gemini-code-assist · 2025-07-15T17:28:32Z

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

+                print("xxxx"*100)
+                print(all2all_manager)
+                print(f"ep_size:{self.ep_size}, {self.ep_rank}")


These print statements appear to be for debugging purposes and should be removed before merging.

gemini-code-assist · 2025-07-15T17:28:32Z

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

+                print("xxxx"*100)
+                print(all2all_manager)
+                print(f"ep_size:{self.ep_size}, {self.ep_rank}")
+                assert all2all_manager is not None
+                # TODO(shuw): need to consider chunking for global_num_tokens_cpu
+                x1, topk_ids1, topk_weights1, alltoall_info = all2all_manager.dispatch(
+                    get_dp_group().device_communicator,
+                    global_num_tokens_cpu,
+                    a1,
+                    topk_ids,
+                    topk_weights,
+                    top_k,
+                    num_experts,
+                    self.ep_rank,
+                    self.ep_size,
+                )
+                self.alltoall_info = alltoall_info


It appears that the all2all_info variable is assigned a value but never used. If this is the case, the assignment should be removed to avoid confusion.

gemini-code-assist · 2025-07-15T17:28:32Z

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

+            if enable_flashinfer_fp4_allgather:
+                topk_weights, topk_ids, a1q, a1q_scale = \
+                    get_dp_group().all_gatherv([topk_weights, topk_ids, a1q, a1q_scale],
+                                            dim=0,
+                                            sizes=get_local_sizes(local_tokens))
+


The code block if enable_flashinfer_fp4_allgather: seems to perform a redundant communication. The all2all_manager.dispatch call on line 122 already performs a gather operation. This is then followed by another get_dp_group().all_gatherv here. This appears to be redundant and could impact performance. Please verify if both are necessary. If not, the redundant call should be removed.

gemini-code-assist · 2025-07-15T17:28:32Z

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

+            # if enable_flashinfer_alltoall:
+            #     print("all2allcalling"*100)
+            #     a1q = MnnvlMoe.mnnvl_moe_alltoallv(a1q, self.alltoall_info,
+            #                                        self.alltoall_workspace,
+            #                                        self.ep_rank, self.ep_size)
+            #     a1q_scale = MnnvlMoe.mnnvl_moe_alltoallv(
+            #         a1q_scale, alltoall_info, self.alltoall_workspace,
+            #         self.ep_rank, self.ep_size)


There are large blocks of commented-out code related to mnnvl_moe_alltoallv. This code should be either implemented or removed to avoid confusion and keep the codebase clean.

mergify · 2025-08-29T21:07:25Z

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

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

mergify · 2025-09-04T09:57:38Z

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

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

vllm/envs.py

vllm/distributed/device_communicators/all2all.py

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

vllm/model_executor/layers/quantization/modelopt.py

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

Signed-off-by: Shu Wang <shuw@nvidia.com>

vllm/distributed/device_communicators/mnnvl_compat.py

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

mergify · 2025-09-23T16:06:44Z

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

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

mgoin · 2025-09-23T19:59:05Z

@wenscarl can you fix the latest merge conflict?

Signed-off-by: Shu Wang <shuw@nvidia.com>

Signed-off-by: Shu Wang <shuw@nvidia.com> Signed-off-by: Shu Wang. <shuw@nvidia.com> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Signed-off-by: yewentao256 <zhyanwentao@126.com>

Signed-off-by: Shu Wang <shuw@nvidia.com> Signed-off-by: Shu Wang. <shuw@nvidia.com> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

mergify bot added the needs-rebase label Jul 15, 2025

gemini-code-assist bot reviewed Jul 15, 2025

View reviewed changes

wenscarl force-pushed the dev_a2a branch from 5cb01c7 to 2527891 Compare August 25, 2025 21:11

mergify bot added v1 and removed needs-rebase labels Aug 25, 2025

wenscarl marked this pull request as ready for review August 25, 2025 21:14

wenscarl requested review from WoosukKwon, alexm-redhat, comaniac, mgoin, njhill, robertgshaw2-redhat, tlrmchlsmth, yewentao256 and ywang96 as code owners August 25, 2025 21:14

wenscarl force-pushed the dev_a2a branch 2 times, most recently from 888fe50 to 1ead2de Compare August 28, 2025 19:47

mergify bot added the needs-rebase label Aug 29, 2025

wenscarl force-pushed the dev_a2a branch from 89e20a0 to 0a20453 Compare August 31, 2025 00:14

mergify bot removed the needs-rebase label Aug 31, 2025

mergify bot added the needs-rebase label Sep 4, 2025

tlrmchlsmth self-assigned this Sep 9, 2025

mergify bot removed the needs-rebase label Sep 9, 2025

tlrmchlsmth reviewed Sep 9, 2025

View reviewed changes

vllm/envs.py Outdated Show resolved Hide resolved

tlrmchlsmth reviewed Sep 9, 2025

View reviewed changes

vllm/distributed/device_communicators/all2all.py Outdated Show resolved Hide resolved

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py Outdated Show resolved Hide resolved

vllm/model_executor/layers/quantization/modelopt.py Show resolved Hide resolved

bnellnm reviewed Sep 9, 2025

View reviewed changes

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py Outdated Show resolved Hide resolved

Add back naive_multicast

92ee54a

Signed-off-by: Shu Wang <shuw@nvidia.com>

wenscarl force-pushed the dev_a2a branch from dbb9b18 to 92ee54a Compare September 19, 2025 19:00

wenscarl requested a review from tlrmchlsmth September 19, 2025 19:30

wenscarl added 2 commits September 19, 2025 21:28

Merge branch 'main' into dev_a2a

23d280a

Merge branch 'main' into dev_a2a

1bc14eb

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

tlrmchlsmth reviewed Sep 23, 2025

View reviewed changes

vllm/distributed/device_communicators/mnnvl_compat.py Outdated Show resolved Hide resolved

remove allgather_bytes

c08aeec

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

tlrmchlsmth approved these changes Sep 23, 2025

View reviewed changes

pre-commit

d2cd8cc

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

tlrmchlsmth enabled auto-merge (squash) September 23, 2025 16:00

mergify bot added the needs-rebase label Sep 23, 2025

Merge branch 'main' into dev_a2a

e9fcf44

Signed-off-by: Shu Wang <shuw@nvidia.com>

auto-merge was automatically disabled September 24, 2025 16:05
Head branch was pushed to by a user without write access

mergify bot removed the needs-rebase label Sep 24, 2025

wenscarl requested review from bnellnm and tlrmchlsmth September 24, 2025 16:15

mgoin merged commit 54e42b7 into vllm-project:main Sep 24, 2025
56 checks passed

pavanimajety mentioned this pull request Sep 24, 2025

[Bug]: [B200] Enabling FlashInfer FP8 MoE results in an exception #25506

Closed

1 task

hjjq mentioned this pull request Jan 12, 2026

[Draft][Kernel] Add new flashinfer A2A kernel #32217

Closed

5 tasks

This was referenced Mar 4, 2026

[Kernel] Add FlashInfer MoE A2A Kernel #36022

Merged

Create tests/distributed/test_mnnvl_alltoall.py #35241

Open

Uh oh!

Conversation

wenscarl commented Jul 15, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Essential Elements of an Effective PR Description Checklist

Purpose

Test Plan

Test Result

(Optional) Documentation Update

Uh oh!

github-actions bot commented Jul 15, 2025

Uh oh!

mergify bot commented Jul 15, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 15, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

gemini-code-assist bot Jul 15, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 15, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 15, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 15, 2025

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Aug 29, 2025

Uh oh!

mergify bot commented Sep 4, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mergify bot commented Sep 23, 2025

Uh oh!

mgoin commented Sep 23, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

wenscarl commented Jul 15, 2025 •

edited by github-actions bot

Loading