A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 00:00 UTC
[Bugfix] Fix SP pass for multimodal models and PP+SP residual handling (#33322)
Signed-off-by: Xingran Wang <wangxingran123456@outlook.com> Signed-off-by: Hongjian Zhang <hirokenovo@gmail.com> Co-authored-by: Hongjian Zhang <hirokenovo@gmail.com>
Wang Xingran committed
May 10, 2026 at 10:44 UTC
0b272a6e01ab49274ce392dc1684d52143d5fd8e
3 files changed
+116
-34
tests/compile/correctness_e2e/test_sequence_parallel.py
+48
@@ -170,6 +170,7 @@ def _compare_sp(
170
num_gpus_available: int,
171
use_inductor_graph_partition: bool,
172
fuse_gemm_comms: bool,
173
+ enable_prompt_embeds: bool,
174
*,
175
method: Literal["generate", "encode"],
176
is_multimodal: bool,
@@ -253,6 +254,8 @@ def _compare_sp(
254
"--enable-mm-embeds",
255
]
256
)
257
+ elif enable_prompt_embeds:
258
+ common_args.append("--enable-prompt-embeds")
259
260
compilation_config = {
261
"mode": CompilationMode.VLLM_COMPILE,
@@ -262,7 +265,9 @@ def _compare_sp(
265
"fuse_gemm_comms": fuse_gemm_comms,
266
"fuse_norm_quant": fuse_norm_quant,
267
"fuse_act_quant": fuse_act_quant,
268
+ "fuse_allreduce_rms": False,
269
"eliminate_noops": True,
270
+ "sp_min_token_num": 0,
271
},
272
"use_inductor_graph_partition": use_inductor_graph_partition,
273
}
@@ -354,6 +359,49 @@ def test_tp_sp_generation(
359
num_gpus_available,
360
use_inductor_graph_partition,
361
fuse_gemm_comms=fuse_gemm_comms,
362
+ enable_prompt_embeds=False,
363
+ method="generate",
364
+ is_multimodal=False,
365
+ )
366
+
367
+
368
+# Focused regression test for the SP + prompt_embeds graph-rewrite path.
369
+# Covers pp_size=1 (SP only) and pp_size=2 (SP + PP); kept small on purpose so
370
+# we don't double the matrix of `test_tp_sp_generation` above.
371
+SP_PROMPT_EMBEDS_PARALLEL_SETUPS = [
372
+ ParallelSetup(
373
+ tp_size=2,
374
+ pp_size=pp_size,
375
+ fuse_norm_quant=False,
376
+ fuse_act_quant=False,
377
+ eager_mode=False,
378
+ chunked_prefill=False,
379
+ )
380
+ for pp_size in [1, 2]
381
+]
382
+
383
+
384
+@pytest.mark.parametrize("parallel_setup", SP_PROMPT_EMBEDS_PARALLEL_SETUPS)
385
+@pytest.mark.parametrize("use_inductor_graph_partition", [True, False])
386
+@create_new_process_for_each_test()
387
+def test_tp_sp_generation_prompt_embeds(
388
+ parallel_setup: ParallelSetup,
389
+ num_gpus_available,
390
+ use_inductor_graph_partition: bool,
391
+):
392
+ if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
393
+ pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
394
+
395
+ _compare_sp(
396
+ "hmellor/tiny-random-LlamaForCausalLM",
397
+ parallel_setup,
398
+ distributed_backend="mp",
399
+ runner="auto",
400
+ test_options=SPTestOptions(multi_node_only=False, load_format=None),
401
+ num_gpus_available=num_gpus_available,
402
+ use_inductor_graph_partition=use_inductor_graph_partition,
403
+ fuse_gemm_comms=False,
404
+ enable_prompt_embeds=True,
405
method="generate",
406
is_multimodal=False,
407
)
vllm/compilation/passes/fusion/sequence_parallelism.py
+55
-20
@@ -15,7 +15,10 @@ import vllm.ir.ops
15
from vllm.config import VllmConfig
16
from vllm.config.utils import Range
17
from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
18
-from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
18
+from vllm.distributed.parallel_state import (
19
+ get_tensor_model_parallel_rank,
20
+ get_tensor_model_parallel_world_size,
21
+)
22
from vllm.logger import init_logger
23
from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
kFp8StaticTensorSym,
@@ -122,6 +125,7 @@ class _SequenceParallelPatternHelper:
125
self.device = device
126
self.tp_group = get_tp_group()
127
self.tp_size = get_tensor_model_parallel_world_size()
128
+ self.tp_rank = get_tensor_model_parallel_rank()
129
130
def _all_reduce(self, x: torch.Tensor) -> torch.Tensor:
131
return tensor_model_parallel_all_reduce(x)
@@ -209,17 +213,35 @@ class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
213
mm_1: torch.Tensor,
214
rms_norm_weights: torch.Tensor,
215
) -> tuple[torch.Tensor, torch.Tensor]:
212
- # pattern matcher replaces from top-to-bottom,
213
- # so residual is still the full size here.
214
- # once the seqpar pattern with the previous rmsnorm is replaced
216
+ # The pattern matcher replaces from the end of the graph
217
+ # (last layer first). At the time each match is replaced,
218
+ # the preceding layer has NOT been replaced yet, so
219
+ # `residual` is still full-size and the slice below is
220
+ # correct. Once the preceding layer IS replaced, its
221
+ # residual output shrinks to [local_len, H], and this
222
+ # slice becomes semantically incorrect (e.g. for rank > 0,
223
+ # the indices would be out of bounds). However, since the
224
+ # symbolic output shape equals the input shape,
225
+ # NoOpEliminationPass (called at the end of
226
+ # SequenceParallelismPass.__call__) removes these slices
227
+ # before the graph is ever executed or compiled.
228
reduce_scatter = self._reduce_scatter(mm_1)
216
- residual = residual[0 : reduce_scatter.size(0), ...]
229
+ local_len = reduce_scatter.size(0)
230
+ # when the preceding VocabParallelEmbedding is excluded
231
+ # from the FX graph (e.g., passing `inputs_embeds` directly in VLMs),
232
+ # the FirstAllReduceRMSNorm pattern is never matched. we must
233
+ # perform a proper TP-aware slice here. simply using `[0:local_len]`
234
+ # would incorrectly cause all ranks to process rank 0's chunk.
235
+ residual = residual[
236
+ self.tp_rank * local_len : self.tp_rank * local_len + local_len, ...
237
+ ]
238
rmsnorm = vllm.ir.ops.fused_add_rms_norm(
239
reduce_scatter, residual, rms_norm_weights, self.epsilon
240
)
241
all_gather = self._all_gather(rmsnorm[0])
221
- # shape of residual changes but that's fine,
222
- # next node is already slicing it, now becomes a noop
242
+ # residual output is now [local_len, H]; the next layer's
243
+ # slice on it is semantically incorrect until
244
+ # NoOpEliminationPass removes it.
245
return all_gather, rmsnorm[1]
246
247
pm.register_replacement(
@@ -309,19 +331,29 @@ class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
331
rms_norm_weights: torch.Tensor,
332
scale: torch.Tensor,
333
) -> tuple[torch.Tensor, torch.Tensor]:
312
- # pattern matcher replaces from top-to-bottom,
313
- # so residual is still the full size here.
314
- # add a temporary slice which will become a noop
315
- # once the seqpar pattern with the previous rmsnorm is replaced
334
+ # See MiddleAllReduceRMSNormPattern.replacement for a
335
+ # detailed explanation of the temporary slice below:
336
+ # it is correct when first inserted, becomes semantically
337
+ # incorrect after the preceding layer is replaced, and is
338
+ # removed by NoOpEliminationPass before the graph is compiled.
339
reduce_scatter = self._reduce_scatter(mm_1)
317
- residual = residual[0 : reduce_scatter.size(0), ...]
340
+ local_len = reduce_scatter.size(0)
341
+ # when the preceding VocabParallelEmbedding is excluded
342
+ # from the FX graph (e.g., passing `inputs_embeds` directly in VLMs),
343
+ # the FirstAllReduceRMSNorm pattern is never matched. we must
344
+ # perform a proper TP-aware slice here. simply using `[0:local_len]`
345
+ # would incorrectly cause all ranks to process rank 0's chunk.
346
+ residual = residual[
347
+ self.tp_rank * local_len : self.tp_rank * local_len + local_len, ...
348
+ ]
349
rms, residual_out = vllm.ir.ops.fused_add_rms_norm(
350
reduce_scatter, residual, rms_norm_weights, self.epsilon
351
)
352
quant, _ = self.quant_matcher(rms, scale)
353
all_gather = self._all_gather(quant)
323
- # shape of residual changes but that's fine,
324
- # next node is already slicing it, now becomes a noop
354
+ # residual output is now [local_len, H]; the next layer's
355
+ # slice on it is semantically incorrect until
356
+ # NoOpEliminationPass removes it.
357
return all_gather, residual_out
358
359
pm.register_replacement(
@@ -485,12 +517,15 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
517
gets split across TP ranks, causing size mismatches at subgraph
518
boundaries.
519
488
- This pass splits up the residual tensor across TP ranks and hence
489
- divides its size. Because the pattern matcher starts at the end of
490
- the graph, the replacement contains a slice that temporarily conforms
491
- the input residual to the correct size. After all patterns have been
492
- matched, we use a NoOpEliminationPass to clean up what have now
493
- become no-op slices.
520
+ This pass splits up the residual tensor across TP ranks and hence divides
521
+ its size. The pattern matcher starts at the end of the graph (last layer
522
+ first), so when each replacement inserts a residual slice, the preceding
523
+ layer has not been replaced yet and the slice is correct. Once the
524
+ preceding layer IS replaced, its residual output shrinks and the slice
525
+ becomes semantically incorrect (out-of-bounds indices for rank > 0).
526
+ The graph is never executed in this intermediate state —
527
+ NoOpEliminationPass removes these slices based on symbolic shape equality
528
+ (input shape == output shape) before the graph is compiled.
529
"""
530
531
@enable_fake_mode
vllm/v1/worker/gpu_model_runner.py
+13
-14
@@ -3129,7 +3129,7 @@ class GPUModelRunner(
3129
3130
return tuple(tasks)
3131
3132
- def sync_and_slice_intermediate_tensors(
3132
+ def sync_and_gather_intermediate_tensors(
3133
self,
3134
num_tokens: int,
3135
intermediate_tensors: IntermediateTensors | None,
@@ -3140,24 +3140,23 @@ class GPUModelRunner(
3140
tp = self.vllm_config.parallel_config.tensor_parallel_size
3141
is_rs = is_residual_scattered_for_sp(self.vllm_config, num_tokens)
3142
3143
- # When sequence parallelism is enabled, the "residual" tensor is sharded
3144
- # across tensor parallel ranks, so each rank only needs its own slice.
3143
+ # When sequence parallelism is enabled, the "residual" tensor is
3144
+ # sharded across TP ranks. All-gather it here because downstream
3145
+ # QKV + Attention needs the full residual before the SP split point.
3146
if sync_self:
3147
assert intermediate_tensors is not None
3148
for k, v in intermediate_tensors.items():
3149
is_scattered = k == "residual" and is_rs
3149
- copy_len = num_tokens // tp if is_scattered else num_tokens
3150
- self.intermediate_tensors[k][:copy_len].copy_(
3151
- v[:copy_len], non_blocking=True
3150
+ if is_scattered:
3151
+ local_len = num_tokens // tp
3152
+ v = get_tp_group().all_gather(v[:local_len], dim=0)
3153
+
3154
+ self.intermediate_tensors[k][:num_tokens].copy_(
3155
+ v[:num_tokens], non_blocking=True
3156
)
3157
3158
return IntermediateTensors(
3155
- {
3156
- k: v[: num_tokens // tp]
3157
- if k == "residual" and is_rs
3158
- else v[:num_tokens]
3159
- for k, v in self.intermediate_tensors.items()
3160
- }
3159
+ {k: v[:num_tokens] for k, v in self.intermediate_tensors.items()}
3160
)
3161
3162
def eplb_step(self, is_dummy: bool = False, is_profile: bool = False) -> None:
@@ -3373,7 +3372,7 @@ class GPUModelRunner(
3372
intermediate_tensors = None
3373
else:
3374
assert intermediate_tensors is not None
3376
- intermediate_tensors = self.sync_and_slice_intermediate_tensors(
3375
+ intermediate_tensors = self.sync_and_gather_intermediate_tensors(
3376
num_input_tokens, intermediate_tensors, True
3377
)
3378
@@ -5589,7 +5588,7 @@ class GPUModelRunner(
5588
)
5589
)
5590
5592
- intermediate_tensors = self.sync_and_slice_intermediate_tensors(
5591
+ intermediate_tensors = self.sync_and_gather_intermediate_tensors(
5592
num_tokens_padded, None, False
5593
)
5594