[MoE] Optimize fused MoE dispatch and workspace allocation with TP leader restriction

[skyloevil]

Purpose

Optimize MoE (Mixture of Experts) performance by improving dispatch efficiency and workspace memory allocation. The changes address under-allocation issues during tensor parallel operations.

Changes

1. Restrict all2all dispatch to TP leaders: Limit expensive all-to-all communication operations to tensor parallel leaders only, reducing redundant communication overhead.
2. Improved workspace sizing:
   - Base capacity calculation on aq.size(1) (backend T dimension)
   - Guard against under-allocation when TP > 1 by using observed M values
   - Prevent _resize_cache under-allocation during distributed operations
3. Debug logging: Added one-time debug logging for total tokens and workspace comments to aid in performance analysis and troubleshooting.

Test Plan

Command:

VLLM_MOE_DP_CHUNK_SIZE=32 VLLM_ALL2ALL_BACKEND="deepep_low_latency" VLLM_USE_DEEP_GEMM=1  vllm serve Qwen/Qwen3-30B-A3B-FP8  --trust-remote-code  --data-parallel-size 2 --tensor-parallel-size 2 --enable-expert-parallel --port 9010  --no-enable-prefix-caching

lm_eval --model local-completions --tasks gsm8k --model_args model=Qwen/Qwen3-30B-A3B-FP8,base_url=http://127.0.0.1:9010/v1/completions,num_concurrent=30,max_retries=3,tokenized_requests=False --limit 100

Results:

lm_eval result

4 GPU H100 80GB HBM3

moe log

server log

Files Modified

• vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
• vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py
• vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
• vllm/model_executor/layers/fused_moe/layer.py

Related Work

This PR builds upon and is related to:

• Optimize MoE Token Dispatch for Tensor Parallel Configurations #22993: Implements core TP leader dispatch optimization strategy, reducing cross-rank communication overhead by 2x-8x
• [BugFix] : Fix Batched DeepGemm Experts #19515: Provides foundation for MoE workspace allocation fixes and debugging improvements

[mergify]

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

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 several key optimizations for Fused MoE layers, primarily by restricting all-to-all dispatch to Tensor Parallel leaders and improving workspace memory allocation. The logic for these optimizations appears correct and well-implemented. The added debug logging is also a welcome addition for better observability.

However, I've identified a significant maintainability concern with code duplication for the TP leader dispatch logic. This logic is copied in both the high-throughput and low-latency DeepEP finalization files. I've left a comment with a suggestion to refactor this into a shared utility to prevent potential future bugs. Addressing this will make the codebase more robust.

[gemini-code-assist]

This block of code for restricting dispatch to the Tensor Parallel leader is duplicated in vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py (lines 152-160).

Duplicating this logic creates a maintainability issue. If this logic needs to be updated in the future, it's easy to forget to update it in both places, which could lead to subtle and hard-to-debug inconsistencies between the high-throughput and low-latency MoE paths.

To improve maintainability and reduce the risk of future bugs, I recommend refactoring this logic into a shared utility function. For example:

# In a shared utility module
def restrict_dispatch_to_tp_leader(*tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:
    """If the current rank is not a TP leader, returns empty tensors."""
    tp_world_size = get_tensor_model_parallel_world_size()
    if tp_world_size <= 1:
        return tensors

    tp_rank_in_group = get_tp_group().rank_in_group
    if tp_rank_in_group != 0:
        return tuple(t[:0] for t in tensors)
    
    return tensors

# In prepare_async methods
a1, topk_ids, topk_weights = restrict_dispatch_to_tp_leader(
    a1, topk_ids, topk_weights
)

This would centralize the logic, making the code cleaner and safer to maintain.

[robertgshaw2-redhat]

cc @tlrmchlsmth @bnellnm - I believe we are already using SP in this case, but correct me if Im wrong.

[tlrmchlsmth]

cc @tlrmchlsmth @bnellnm - I believe we are already using SP in this case, but correct me if Im wrong.

We use SP for DeepSeekV3 (ntroduced in #24134), but we need to do it for Qwen and other MoE models. I'll put up a PR today

[bnellnm]

I think @varun-sundar-rabindranath should check out this PR as well.

[bnellnm]

I think these changes might be redundant with another PR that makes the workspaces effectively global. See #23693

[mergify]

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

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

[skyloevil]

@varun-sundar-rabindranath PTAL..Thanks.

[mergify]

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

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

[tlrmchlsmth]

@skyloevil this should be covered by #24982, thank you!