PR #36838: enable flashinfer moe kernel for DP + EP

• Repository: vllm-project/vllm
• Author: czhu-cohere
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/36838

Description (problem / solution / changelog)

Purpose

Previously the BF16 flashinfer moe kernel is disabled when dp > 1. I think the kernel itself should be able to support it, we just need to enable on the vLLM side. Also add test to verify the kernel selection logic works as intended.

Test Plan

pytest tests/kernels/moe/test_unquantized_backend_selection.py run gsm8k with bf16 qwen 3a30b on 2xB200 DP2 EP2 and compare the result with different moe backend.

server command

# triton/default backend
vllm serve Qwen/Qwen3-30B-A3B \
  --data-parallel-size 2 \
  --enable-expert-parallel \
  --trust-remote-code \
  --port 8000

# flashinfer cutlass
VLLM_USE_FLASHINFER_MOE_FP16=1 VLLM_FLASHINFER_MOE_BACKEND=throughput vllm serve Qwen/Qwen3-30B-A3B \
  --data-parallel-size 2 \
  --enable-expert-parallel \
  --trust-remote-code \
  --port 8000

# flashinfer trtllm
VLLM_USE_FLASHINFER_MOE_FP16=1 VLLM_FLASHINFER_MOE_BACKEND=latency vllm serve Qwen/Qwen3-30B-A3B \
  --data-parallel-size 2 \
  --enable-expert-parallel \
  --trust-remote-code \
  --port 8000

test command

python -m lm_eval \
  --model local-completions \
  --model_args "model=Qwen/Qwen3-30B-A3B,base_url=http://localhost:8000/v1/completions,num_concurrent=128,max_retries=5,tokenized_requests=False,tokenizer=Qwen/Qwen3-30B-A3B" \
  --tasks gsm8k_cot \
  --batch_size auto \
  --log_samples \
  --output_path /tmp/lm_eval_qwen_dp2_ep

Test Result

pytest tests/kernels/moe/test_unquantized_backend_selection.py pass

• 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.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

Changed files

• tests/kernels/moe/test_unquantized_backend_selection.py (modified, +88/-2)
• vllm/model_executor/layers/fused_moe/oracle/unquantized.py (modified, +1/-7)
• vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py (modified, +8/-0)

PR #39593: [Bug]: Fix FlashInfer CUTLASS BF16 + CUDA graphs IMA

• Repository: vllm-project/vllm
• Author: yzong-rh
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/39593

Description (problem / solution / changelog)

Purpose

Fix illegal memory access crash in FlashInfer CUTLASS BF16 MoE with DP/DEP + CUDA graphs.

Closes https://github.com/vllm-project/vllm/issues/37758

Root cause: When CUDA graphs are enabled under data parallelism, all ranks pad to the same batch size. Attention often produces NaN for padding tokens. In BF16, NaN propagates through RMSNorm, projections, and residual adds until it reaches the MoE router. The topK kernel picks the same expert index K times when encountering NaN values. FlashInfer Cutlass MoE kernels do not expect duplicate expert indices per token — the resulting out-of-bounds access causes a hard CUDA error.

Instrumentation of attention output (Qwen3-30B-A3B BF16, piecewise CUDA graphs, FlashInfer CUTLASS MoE, GSM8K eval) confirms NaN is present in padding rows on virtually every padded step:

DP/DEP produce ~13× more padding events with much larger pad regions than single-GPU bucket rounding, which is why the bug manifests under DP but not single GPU.

cutlass_fp8_bf16_simple.py shows that calling FlashInfer's cutlass_fused_moe in eager mode with all tokens routed to the same expert (simulating the duplicate-ID condition) crashes with an illegal memory access on both BF16 and FP8 paths.

cutlass_fp8_bf16_simple.py

FP8 was immune because the FP8 quantization kernel (scaled_fp8_conversion) uses fminf/fmaxf, which silently clamp NaN to ±448.0, acting as an implicit NaN firewall at every linear layer.

Fix: Sanitize NaN in the fused softmax/sigmoid+topk CUDA kernel (csrc/moe/topk_softmax_kernels.cu). This is applied in three places:

1. moeSoftmax (fallback path for non-power-of-2 expert counts): NaN is replaced with -FLT_MAX in the max-reduce, exp-sum, and final softmax loops.
2. moeSigmoid (fallback sigmoid path): NaN is replaced before the sigmoid computation.
3. topkGating (fused warp-level path for power-of-2 / multiple-of-64 expert counts): NaN is replaced in-register immediately after loading, before softmax/sigmoid and argmax.

topK now returns K distinct indices, never duplicates.

Test Plan

1. tests/kernels/moe/test_routing.py

   • Includes additional tests test_topk_nan_row_distinct_experts and test_grouped_topk_nan_row_distinct_experts, which inject all-NaN rows into top-k. They verify that NaN rows return distinct values and that non NaN rows match the baseline.

2. End-to-end GSM8K correctness (Qwen3-30B-A3B BF16 and Qwen3.5-35B-A3B BF16, FlashInfer CUTLASS MoE, piecewise CUDA graphs):

   Model	Measured	Expected
   Qwen3-30B-A3B BF16 (DP=2)	0.8870	0.8800
   Qwen3-30B-A3B BF16 (DEP=2)	0.8901	0.8800
   Qwen3.5-35B-A3B BF16 (DP=2)	0.8408	0.8400
   Qwen3.5-35B-A3B BF16 (DEP=2)	0.8560	0.8400

3. Micro-benchmark (benchmark_router_select_experts.py, 128 tokens × 128 experts, top_k=8, bfloat16, 10% NaNs, B200, mirrors Qwen3-30B):

   The NaN check is fused into the existing kernel rather than added as a separate torch.nan_to_num(-float('inf')) call before routing. A separate nan_to_num launch adds ~30% overhead:

   Variant	Kernel	p20	p50	p80
   Before fix	topk_softmax	10.21 µs	10.27 µs	10.37 µs
   Before fix	topk_sigmoid	10.24 µs	12.13 µs	12.32 µs
   After fix (fused)	topk_softmax	10.11 µs	10.24 µs	10.30 µs
   After fix (fused)	topk_sigmoid	10.24 µs	12.16 µs	10.32 µs
   torch.nan_to_num	topk_softmax	11.42 µs	13.28 µs	13.44 µs
   torch.nan_to_num	topk_sigmoid	13.15 µs	13.31 µs	13.44 µs

benchmark_router_select_experts.py

1. E2E Latency:

vllm serve Qwen/Qwen3-30B-A3B \
  --data-parallel-size 2 \
  --max-model-len 8192

vllm bench serve \
  --backend vllm \
  --model Qwen/Qwen3-30B-A3B \
  --endpoint /v1/completions \
  --num-prompts 128 \
  --random-input-len 2 \
  --random-output-len 512 \
  --num-warmups 128 \
  --request-rate inf \
  --temperature 0

Results within run to run variance

============ Serving Benchmark Result ============
Successful requests:                     128       
Failed requests:                         0         
Benchmark duration (s):                  4.74      
Total input tokens:                      256       
Total generated tokens:                  65536     
Request throughput (req/s):              27.01     
Output token throughput (tok/s):         13826.70  
Peak output token throughput (tok/s):    14336.00  
Peak concurrent requests:                128.00    
Total token throughput (tok/s):          13880.71  
---------------Time to First Token----------------
Mean TTFT (ms):                          97.14     
Median TTFT (ms):                        99.54     
P99 TTFT (ms):                           116.43    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          9.06      
Median TPOT (ms):                        9.06      
P99 TPOT (ms):                           9.06      
---------------Inter-token Latency----------------
Mean ITL (ms):                           9.06      
Median ITL (ms):                         9.06      
P99 ITL (ms):                            10.59     
==================================================

============ Serving Benchmark Result ============
Successful requests:                     128       
Failed requests:                         0         
Benchmark duration (s):                  4.72      
Total input tokens:                      256       
Total generated tokens:                  65536     
Request throughput (req/s):              27.11     
Output token throughput (tok/s):         13881.52  
Peak output token throughput (tok/s):    14336.00  
Peak concurrent requests:                128.00    
Total token throughput (tok/s):          13935.74  
---------------Time to First Token----------------
Mean TTFT (ms):                          95.84     
Median TTFT (ms):                        97.08     
P99 TTFT (ms):                           109.43    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          9.03      
Median TPOT (ms):                        9.03      
P99 TPOT (ms):                           9.03      
---------------Inter-token Latency----------------
Mean ITL (ms):                           9.03      
Median ITL (ms):                         9.00      
P99 ITL (ms):                            11.15     
==================================================

============ Serving Benchmark Result ============
Successful requests:                     128       
Failed requests:                         0         
Benchmark duration (s):                  4.83      
Total input tokens:                      256       
Total generated tokens:                  65536     
Request throughput (req/s):              26.48     
Output token throughput (tok/s):         13556.04  
Peak output token throughput (tok/s):    14031.00  
Peak concurrent requests:                128.00    
Total token throughput (tok/s):          13608.99  
---------------Time to First Token----------------
Mean TTFT (ms):                          97.80     
Median TTFT (ms):                        104.58    
P99 TTFT (ms):                           115.65    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          9.23      
Median TPOT (ms):                        9.23      
P99 TPOT (ms):                           9.24      
---------------Inter-token Latency----------------
Mean ITL (ms):                           9.23      
Median ITL (ms):                         9.22      
P99 ITL (ms):                            11.01     
==================================================

Test Result

See above.

• 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.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)

Changed files

• csrc/moe/topk_softmax_kernels.cu (modified, +8/-3)
• tests/kernels/moe/test_routing.py (modified, +168/-3)
• vllm/model_executor/layers/fused_moe/oracle/unquantized.py (modified, +0/-6)