[wangxingran222]

Purpose

When deploying multimodal models on the current vLLM main branch, or deploying text models with the --enable-prompt-embeds option enabled, Sequence Parallelism (SP) cannot be enabled due to torch compile errors. Specifically, two issues occur:

1. Shape mismatch in aten.copy_ operations: When torch compile detects that the forward function modifies input parameters (e.g., inputs_embeds), the compiler inserts aten.copy_ nodes to write modified values back to input tensors (functionalization and copy_ epilogue). In SP mode, the computation is sharded while input parameters remain full-length, causing shape mismatch errors in these copy operations.

(Worker_TP0 pid=1489154) ERROR 01-27 20:18:16 [multiproc_executor.py:852]   File "/workspace/fixmultimodalsp/lib/python3.12/site-packages/torch/__init__.py", line 1677, in _check_with
(Worker_TP0 pid=1489154) ERROR 01-27 20:18:16 [multiproc_executor.py:852]     raise error_type(message_evaluated)
(Worker_TP0 pid=1489154) ERROR 01-27 20:18:16 [multiproc_executor.py:852] torch._inductor.exc.InductorError: RuntimeError: expand: attempting to expand a dimension of length (s59//4) -> s59!

2. Incorrect residual slicing: The middle RMSNorm patterns incorrectly slice the residual tensor using residual[0:local_len, ...], which always takes the first shard regardless of TP rank. This causes incorrect results when tp_rank > 0.

This PR fixes both issues by:

• Adding an InputMutationEliminationPass to remove shape-mismatched copy_ operations on graph inputs
• Correcting the residual slicing to use residual[tp_rank * local_len : (tp_rank + 1) * local_len, ...]

Supplementary Material

Details

Issue 1: Computation Graph Error Details

The tail end of the original computation graph shows the shape mismatch issue:

# No stacktrace found for following nodes
reduce_scatter_default: "bf16[(s59//4), 5120]" = torch.ops.vllm.reduce_scatter.default(mm_255, 0, 4, 'tp:0');  mm_255 = None
sym_size_int_2: "Sym((s59//4))" = torch.ops.aten.sym_size.int(reduce_scatter_default, 0);  sym_size_int_2 = None
auto_functionalized_321 = torch.ops.higher_order.auto_functionalized(torch.ops._C.fused_add_rms_norm.default, input = reduce_scatter_default, residual = getitem_1925, weight = arg516_1, epsilon = 1e-06);  reduce_scatter_default = getitem_1925 = arg516_1 = None
getitem_1922: "bf16[(s59//4), 5120]" = auto_functionalized_321[1]
getitem_1923: "bf16[(s59//4), 5120]" = auto_functionalized_321[2];  auto_functionalized_321 = None
all_gather_default: "bf16[4*((s59//4)), 5120]" = torch.ops.vllm.all_gather.default(getitem_1922, 0, 4, 'tp:0');  getitem_1922 = None

 # File: /workspace/mydev/fixmultimodalsp/vllm/_custom_ops.py:341 in fused_add_rms_norm, code: torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon)
copy_: "bf16[s59, 5120]" = torch.ops.aten.copy_.default(arg1_1, getitem_1923);  arg1_1 = getitem_1923 = copy_ = None
return (all_gather_default,)

What's happening here:

• arg1_1 is the input_embeds with shape bf16[s59, 5120] (full length)
• getitem_1923 has shape bf16[(s59//4), 5120] (sharded length after Sequence Parallelism)
• The copy_ operation tries to copy the sharded tensor back into the full-length input tensor, causing a shape mismatch
• This triggers the Inductor RuntimeError: expand: attempting to expand a dimension of length (s59//4) -> s59!

Issue 2: Incorrect Residual Slicing Details

For multimodal models or text models with --enable-prompt-embeds, the vocabulary embedding computation doesn't happen inside the model's forward function—it's done in gpu_model_runner instead. This means the FirstAllReduceRMSNormPattern never gets matched, so its output residual never gets sharded by SP. As a result, the first MiddleAllReduceRMSNormPattern receives the full-length residual instead of a sharded one.

The problem is in the MiddleAllReduceRMSNormPattern replacement, which uses residual = residual[0 : reduce_scatter.size(0), ...]. This causes all TP ranks to slice the same first chunk of the residual [0, sym_size_int_129], leading to incorrect computation results.

All TP ranks end up with the same computation graph:

reduce_scatter_default_127: "bf16[(s59//4), 5120]" = torch.ops.vllm.reduce_scatter.default(mm_1, 0, 4, 'tp:0');  mm_1 = None
sym_size_int_129: "Sym((s59//4))" = torch.ops.aten.sym_size.int(reduce_scatter_default_127, 0)
slice_tensor_127: "bf16[(s59//4), 5120]" = torch.ops.aten.slice.Tensor(arg1_1, 0, 0, sym_size_int_129);  arg1_1 = sym_size_int_129 = None
auto_functionalized_448 = torch.ops.higher_order.auto_functionalized(torch.ops._C.fused_add_rms_norm.default, input = reduce_scatter_default_127, residual = slice_tensor_127, weight = arg9_1, epsilon = 1e-06);  reduce_scatter_default_127 = slice_tensor_127 = arg9_1 = None

Notice how slice_tensor_127 always slices from index 0 to sym_size_int_129, regardless of TP rank. This is wrong—each rank should be slicing its own chunk of the residual.

Test Plan

We ran both performance and accuracy tests to validate the fix.

Performance Testing:

• Tested SP + Async TP on H800 with TP=4
• Measured TTFT improvements on:
  • Text models with --enable-prompt-embeds enabled
  • Multimodal models
• Test workload: single concurrent request with input length of 16384 tokens

Accuracy Testing:

• Verified correctness of SP + Async TP on:
  • Text models with --enable-prompt-embeds: tested on gsm8k
  • Multimodal models: tested on mmmu_val
• Compared results against baseline (SP + Async TP disabled)
• No significant accuracy degradation observed in either case
• Note: Accuracy tests were conducted with [Bugfix] Use 'sum' reduction instead of 'avg' in Async TP reduce-scatter #33088 applied (not yet merged)

Test Result

Performance Testing

Text model(Qwen3-32B) with --enable-prompt-embeds:

Details

This PR, Qwen3-32B, enable-prompt-embeds, SP + Async TP on

============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  37.58     
Total input tokens:                      524288    
Total generated tokens:                  512       
Request throughput (req/s):              0.85      
Output token throughput (tok/s):         13.62     
Peak output token throughput (tok/s):    16.00     
Peak concurrent requests:                2.00      
Total token throughput (tok/s):          13963.15  
---------------Time to First Token----------------
Mean TTFT (ms):                          644.54    
Median TTFT (ms):                        643.55    
P99 TTFT (ms):                           668.19    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          35.31     
Median TPOT (ms):                        35.30     
P99 TPOT (ms):                           35.51     
---------------Inter-token Latency----------------
Mean ITL (ms):                           35.31     
Median ITL (ms):                         35.26     
P99 ITL (ms):                            38.46     
==================================================

Main, Qwen3-32B, enable-prompt-embeds, SP + Async TP off

============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  28.89     
Total input tokens:                      524288    
Total generated tokens:                  512       
Request throughput (req/s):              1.11      
Output token throughput (tok/s):         17.72     
Peak output token throughput (tok/s):    27.00     
Peak concurrent requests:                3.00      
Total token throughput (tok/s):          18166.43  
---------------Time to First Token----------------
Mean TTFT (ms):                          757.27    
Median TTFT (ms):                        755.71    
P99 TTFT (ms):                           781.50    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          9.68      
Median TPOT (ms):                        9.68      
P99 TPOT (ms):                           9.75      
---------------Inter-token Latency----------------
Mean ITL (ms):                           9.68      
Median ITL (ms):                         9.64      
P99 ITL (ms):                            10.30     
==================================================

Main, Qwen3-32B, Token IDs, SP + Async TP on

============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  37.69     
Total input tokens:                      524288    
Total generated tokens:                  512       
Request throughput (req/s):              0.85      
Output token throughput (tok/s):         13.58     
Peak output token throughput (tok/s):    16.00     
Peak concurrent requests:                2.00      
Total token throughput (tok/s):          13923.24  
---------------Time to First Token----------------
Mean TTFT (ms):                          645.47    
Median TTFT (ms):                        647.78    
P99 TTFT (ms):                           668.92    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          35.47     
Median TPOT (ms):                        35.47     
P99 TPOT (ms):                           35.52     
---------------Inter-token Latency----------------
Mean ITL (ms):                           35.47     
Median ITL (ms):                         35.45     
P99 ITL (ms):                            36.05     
==================================================

Multimodal model(Qwen2.5-VL-32B):

Details

This PR, Qwen2.5-VL-32B, SP + Async TP on

============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  35.47     
Total input tokens:                      524288    
Total generated tokens:                  512       
Request throughput (req/s):              0.90      
Output token throughput (tok/s):         14.44     
Peak output token throughput (tok/s):    16.00     
Peak concurrent requests:                2.00      
Total token throughput (tok/s):          14797.33  
---------------Time to First Token----------------
Mean TTFT (ms):                          630.27    
Median TTFT (ms):                        630.89    
P99 TTFT (ms):                           653.25    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          31.85     
Median TPOT (ms):                        31.85     
P99 TPOT (ms):                           31.89     
---------------Inter-token Latency----------------
Mean ITL (ms):                           31.85     
Median ITL (ms):                         31.84     
P99 ITL (ms):                            32.57     
==================================================

Main, Qwen2.5-VL-32B, SP + Async TP off

============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  27.37     
Total input tokens:                      524288    
Total generated tokens:                  512       
Request throughput (req/s):              1.17      
Output token throughput (tok/s):         18.71     
Peak output token throughput (tok/s):    30.00     
Peak concurrent requests:                3.00      
Total token throughput (tok/s):          19174.74  
---------------Time to First Token----------------
Mean TTFT (ms):                          718.52    
Median TTFT (ms):                        717.72    
P99 TTFT (ms):                           758.97    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          9.10      
Median TPOT (ms):                        9.11      
P99 TPOT (ms):                           9.14      
---------------Inter-token Latency----------------
Mean ITL (ms):                           9.10      
Median ITL (ms):                         9.12      
P99 ITL (ms):                            9.74      
==================================================

Accuracy Testing

Text model(Qwen3-32B) with --enable-prompt-embeds:

Details

This PR, Qwen3-32B, enable-prompt-embeds, SP + Async TP on

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.6232	±	0.0133
		strict-match	5	exact_match	↑	0.7392	±	0.0121

Main, Qwen3-32B, enable-prompt-embeds, SP + Async TP off

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.6240	±	0.0133
		strict-match	5	exact_match	↑	0.7392	±	0.0121

Multimodal model(Qwen2.5-VL-32B):

Details

This PR, Qwen2.5-VL-32B, SP + Async TP on

Tasks	Version	Filter	n-shot	Metric		Value		Stderr	Stderr_CLT	Stderr_Clustered
mmmu_val	0	none	0	mmmu_acc	↑	0.5922	±	N/A	N/A	N/A

Main, Qwen2.5-VL-32B, SP + Async TP off

Tasks	Version	Filter	n-shot	Metric		Value		Stderr	Stderr_CLT	Stderr_Clustered
mmmu_val	0	none	0	mmmu_acc	↑	0.5922	±	N/A	N/A	N/A

---

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

[gemini-code-assist]

Code Review

This pull request introduces two key fixes to enable Sequence Parallelism (SP) for multimodal models and models using prompt embeddings. First, it adds a new InputMutationEliminationPass to remove erroneous aten.copy_ operations inserted by AOT Autograd, which cause shape mismatches in SP mode. This is a clean solution to a tricky compiler-induced problem. Second, it corrects the residual tensor slicing logic in MiddleAllReduceRMSNormPattern and MiddleAllReduceRMSNormStaticFP8Pattern to correctly handle sharded computation across different tensor parallel ranks. The changes are well-explained, logical, and supported by performance and accuracy tests. The implementation is solid and effectively resolves the described bugs. Overall, this is an excellent contribution that improves vLLM's support for advanced parallelism features.

[wangxingran222]

cc @jasonlizhengjian @angelayi @cascade812 @ProExpertProg

[ProExpertProg]

Sorry why do prompt embeds and MM models trigger different behavior here? Can you share the resulting FX graph that has this problem? Also please share the full commands that trigger this issue.

Maybe we should just only use SP+AsyncTP with fullgraph compilation (inductor partition or splitting_ops=[]).

[wangxingran222]

Sorry why do prompt embeds and MM models trigger different behavior here?

The distinct behavior stems from how embeddings are computed for these cases in gpu_model_runner:

vllm/vllm/v1/worker/gpu_model_runner.py

Lines 2801 to 2853 in 07bdabe

	if self.supports_mm_inputs and is_first_rank and not is_encoder_decoder:
	# Run the multimodal encoder if any.
	with self.maybe_get_ec_connector_output(
	scheduler_output,
	encoder_cache=self.encoder_cache,
	) as ec_connector_output:
	self._execute_mm_encoder(scheduler_output)
	mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output)
	# NOTE(woosuk): To unify token ids and soft tokens (vision
	# embeddings), we always use embeddings (rather than token ids)
	# as input to the multimodal model, even when the input is text.
	inputs_embeds_scheduled = self.model.embed_input_ids(
	self.input_ids.gpu[:num_scheduled_tokens],
	multimodal_embeddings=mm_embeds,
	is_multimodal=is_mm_embed,
	)
	# TODO(woosuk): Avoid the copy. Optimize.
	self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
	input_ids, inputs_embeds = self._prepare_mm_inputs(num_input_tokens)
	model_kwargs = {
	**self._init_model_kwargs(),
	**self._extract_mm_kwargs(scheduler_output),
	}
	elif self.enable_prompt_embeds and is_first_rank:
	# Get the input embeddings for the tokens that are not input embeds,
	# then put them into the appropriate positions.
	# TODO(qthequartermasterman): Since even when prompt embeds are
	# enabled, (a) not all requests will use prompt embeds, and (b)
	# after the initial prompt is processed, the rest of the generated
	# tokens will be token ids, it is not desirable to have the
	# embedding layer outside of the CUDA graph all the time. The v0
	# engine avoids this by "double compiling" the CUDA graph, once
	# with input_ids and again with inputs_embeds, for all num_tokens.
	# If a batch only has token ids, then including the embedding layer
	# in the CUDA graph will be more performant (like in the else case
	# below).
	token_ids_idx = (
	self.is_token_ids.gpu[:num_scheduled_tokens]
	.nonzero(as_tuple=False)
	.squeeze(1)
	)
	# Some tokens ids may need to become embeds
	if token_ids_idx.numel() > 0:
	token_ids = self.input_ids.gpu[token_ids_idx]
	tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids)
	self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
	inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
	model_kwargs = self._init_model_kwargs()
	input_ids = None

For prompt embeddings and multimodal models, the embedding lookup happens outside the model's forward function. Consequently, inputs_embeds (rather than input_ids) are passed directly to the model, meaning the VocabParallelEmbedding layer is excluded from the FX Graph. This leads to two specific issues:

1. Shape Mismatch in copy_ Ops: Since inputs_embeds enters the graph directly, it acts as the residual for the fused_add_rms_norm operator (an in-place op). During AOT Autograd functionalization, the compiler inserts a copy_ epilogue to write results back to the input tensor. However, with Sequence Parallelism (SP) enabled, the computed output is sharded (seq_len // tp_size), while the original input tensor remains full-length (seq_len), causing a shape mismatch error.
2. Incorrect Residual Slicing: Because VocabParallelEmbedding is missing from the graph, the FirstAllReduceRMSNormPattern is never matched. As a result, the residual tensor exiting the first norm layer remains unsharded. When the subsequent MiddleAllReduceRMSNormPattern applies slicing (residual[0 : local_len]), every TP rank incorrectly slices the first chunk of the global residual (intended only for Rank 0), leading to computation errors.

share the full commands that trigger this issue.

For prompt embeds:

vllm serve \
    Qwen/Qwen2.5-72B-Instruct \
    --enable-prompt-embeds \
    --tensor-parallel-size 4 \
    --compilation-config '{"pass_config": {"enable_sp": true, "fuse_gemm_comms": true, "fuse_allreduce_rms": false}, "splitting_ops":[], "cudagraph_mode": "FULL"}' \
    --load-format dummy

For MM models:

vllm serve \
    Qwen/Qwen2.5-VL-72B-Instruct \
    --tensor-parallel-size 4 \
    --compilation-config '{"pass_config": {"enable_sp": true, "fuse_gemm_comms": true, "fuse_allreduce_rms": false}, "splitting_ops":[], "cudagraph_mode": "FULL"}' \
    --load-format dummy

Can you share the resulting FX graph that has this problem?

Sure. Below are the FX graphs for Qwen2.5-VL-72B with TP=4 on vLLM main branch, with serve command:

vllm serve \
    Qwen/Qwen2.5-VL-72B-Instruct \
    --tensor-parallel-size 4 \
    --compilation-config '{"pass_config": {"enable_sp": true, "fuse_gemm_comms": true, "fuse_allreduce_rms": false}, "splitting_ops":[], "cudagraph_mode": "FULL"}' \
    --load-format dummy

The first file shows the graph before the SP Pass applied.

The second file shows the graph after the SP Pass applied, where the shape mismatch and slicing issues described above occur.

qwen25_vl72b_fx_graph_before.py

qwen25_vl72b_fx_graph_after.py

[ProExpertProg]

Thank you so much for the detailed response! This makes sense to me. Proper slicing seems right in this case. Perhaps include some of the details above in a comment somewhere? And can you add a relevant case (prompt embeds enabled) to the E2E correctness tests?

[ProExpertProg]

can you rewrite this to if not ... continue to avoid deep if nesting?

[wangxingran222]

can you rewrite this to if not ... continue to avoid deep if nesting?

Sure, fixed.

[ProExpertProg]

@wangxingran222 can you also rewrite this comment because your approach now guarantees semantic correctness during the pass as well!

[ProExpertProg]

(if I understand correctly)

[wangxingran222]

@wangxingran222 can you also rewrite this comment because your approach now guarantees semantic correctness during the pass as well!

Yes, fixed.

[wangxingran222]

Thank you so much for the detailed response! This makes sense to me. Proper slicing seems right in this case. Perhaps include some of the details above in a comment somewhere? And can you add a relevant case (prompt embeds enabled) to the E2E correctness tests?

Thanks for the review! I have integrated the requested E2E test case. It is currently failing, and I am investigating the root cause. I will push the fix once resolved.

[ProExpertProg]

@wangxingran222 I think I might know why it's failing. When compiling piecewise (no inductor partition and non-empty splitting ops), the residual passed between graphs gets sliced, which I'm now remembering was the original reason my slicing was not tp-aware.i think we should just remove the piecewise support for SP, let me put up an RFC.

EDIT: RFC here: #35771

[wangxingran222]

I think I might know why it's failing. When compiling piecewise (no inductor partition and non-empty splitting ops), the residual passed between graphs gets sliced, which I'm now remembering was the original reason my slicing was not tp-aware.

Yes. I observed failures in all use_inductor_graph_partition=False cases, as well as all pp=2 cases. The proper slicing implementation currently breaks SP + PP compatibility.

[ProExpertProg]

Do you mind chiming in on the RFC?

[wangxingran222]

There are 4 CI failures, but they all look unrelated to SP or PP.

[ProExpertProg]

Do all the tests fail on main as well? If yes, can you post in #pr-merge-requests in the slack?

[wangxingran222]

Do all the tests fail on main as well? If yes, can you post in #pr-merge-requests in the slack?

Did a quick scan on recent PRs: the 3 AMD failures (amd-v1-spec-decode-mi325-1, amd-multi-modal-models-standard-3, amd-multi-modal-models-standard-4) also fail on #41486. elastic-ep-scaling-test seems to be consistently failing since ~05-05, e.g. #40731