[Bugfix][Wide EP] Fix redundant work when using DeepEP, TP Attn, and EP MoE

[tlrmchlsmth]

Purpose

Currently, when running attention with TP and using --enable-expert-parallel, the MoE layers will do duplicate work when using DeepEP. In this case, the output of attention will be replicated across TP ranks and each token copy will be dispatched to the EP ranks it gets routed to, multiplying the amount of work by tp_size.

This PR avoids this duplicate work by ensuring the input to the MoE layer is sequence parallel instead of replicated.

Notes:

• The performance bug applies to any MoE model but this PR only fixes it for DeepSeekV3
• We can do better by replacing the all_reduce at the end of attention with a reduce_scatter. This reduces the amount of computation but is a little more invasive to the model definition since we need to handle the sharding of the residuals. This has an extra effect of de-duplicating the work done during layer norms (minor improvement).

Test Plan

lm_eval --model local-completions --tasks gsm8k --model_args model=deepseek-ai/DeepSeek-R1-0528,base_url=http://infra-wide-ep-inference-gateway-istio.llm-d-wide-ep.svc.cluster.local/v1/completions,num_concurrent=1000,max_retries=3,tokenized_requests=False

Test Result

This PR

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9583|±  |0.0055|
|     |       |strict-match    |     5|exact_match|↑  |0.9560|±  |0.0056|

5b31cb1 (last good commit before #24119 landed):

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9591|±  |0.0055|
|     |       |strict-match    |     5|exact_match|↑  |0.9568|±  |0.0056|

[gemini-code-assist]

Code Review

This pull request correctly introduces sequence parallelism to the MoE layer in the DeepseekV2 model to prevent redundant computations when using both Tensor Parallelism and Expert Parallelism. The approach of chunking the input before the MoE layer and gathering the output afterward is sound. I've found one critical issue that could lead to a runtime error, which I've detailed in a specific comment.

[tlrmchlsmth]

@abcdabcd987 and @nandor could you double-check me here: Can pplx handle replicated input tokens in the TP attn + EP MoE case?

[robertgshaw2-redhat]

I think we should call this use_sequence_parallel_mlp since we use seq parallelism for just the mlp layer here

[tlrmchlsmth]

Not sure I like this because the MLP being sequence parallel is kind of a side effect. And we need to pass it into the fused_moe layer for the chunking.

I'm not a fan of the sequence_parallel name though so definitely open to suggestions

[robertgshaw2-redhat]

This looks good to me. Just left some comments on explaining the parallelism setup.

[robertgshaw2-redhat]

there are genuine failures in the CI related to DeepSeek MTP

[mergify]

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

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

[tlrmchlsmth]

Seeing some issues with CUDA graphs on:

ms-wide-ep-llm-d-modelservice-decode-0-1 vllm-worker-decode (EngineCore_7 pid=290) (VllmWorker TP0 pid=314) ERROR 09-03 19:20:29 [multiproc_executor.py:611]   File "/workspace/vllm/vllm/model_executor/layers/fused_moe/layer.py", line 1663, in process_chunk
ms-wide-ep-llm-d-modelservice-decode-0-1 vllm-worker-decode (EngineCore_7 pid=290) (VllmWorker TP0 pid=314) ERROR 09-03 19:20:29 [multiproc_executor.py:611]     staged_hidden_states.copy_(hidden_states, non_blocking=True)
ms-wide-ep-llm-d-modelservice-decode-0-1 vllm-worker-decode (EngineCore_7 pid=290) (VllmWorker TP0 pid=314) ERROR 09-03 19:20:29 [multiproc_executor.py:611] RuntimeError: output with shape [1, 7168] doesn't match the broadcast shape [0, 7168]

[tlrmchlsmth]

Seeing some issues with CUDA graphs on:

ms-wide-ep-llm-d-modelservice-decode-0-1 vllm-worker-decode (EngineCore_7 pid=290) (VllmWorker TP0 pid=314) ERROR 09-03 19:20:29 [multiproc_executor.py:611]   File "/workspace/vllm/vllm/model_executor/layers/fused_moe/layer.py", line 1663, in process_chunk
ms-wide-ep-llm-d-modelservice-decode-0-1 vllm-worker-decode (EngineCore_7 pid=290) (VllmWorker TP0 pid=314) ERROR 09-03 19:20:29 [multiproc_executor.py:611]     staged_hidden_states.copy_(hidden_states, non_blocking=True)
ms-wide-ep-llm-d-modelservice-decode-0-1 vllm-worker-decode (EngineCore_7 pid=290) (VllmWorker TP0 pid=314) ERROR 09-03 19:20:29 [multiproc_executor.py:611] RuntimeError: output with shape [1, 7168] doesn't match the broadcast shape [0, 7168]

It was a torch.compile issue. I ended up having to work around it by wrapping the sp chunking code in a custom op

[nvpohanh]

cc @weireweire

[tlrmchlsmth]

cc @zou3519 @ProExpertProg in case you have better ideas than this wrap-it-in-a-custom-op hack

[zou3519]

@tlrmchlsmth do you have the original error message and/or a stack trace?

[zou3519]

Also, this custom operator is technically incorrect. The output is a view of the input, which means bad things can happen in the presence of mutation. I don't know if vLLM specifically will hit any of those issues, it depends on how it's being used.