@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 32e45636e3
Fork 0

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, 15:00 UTC

[torch.compile]: Disable Sequence Parallelism (SP) for piecewise compilation (#38373)

Signed-off-by: SouthWest7 <am1ao@qq.com> Signed-off-by: Xinan Miao <1403572259@qq.com> Co-authored-by: SouthWest7 <am1ao@qq.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Wang Xingran <72983099+wangxingran222@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Xinan Miao committed Apr 27, 2026 at 01:44 UTC 32e45636e3d7e02615facc8c63645ce4ac1d7e11
9 files changed +223 -80
tests/compile/correctness_e2e/test_sequence_parallel.py
+2
@@ -261,6 +261,8 @@ def _compare_sp(
261 },
262 "use_inductor_graph_partition": use_inductor_graph_partition,
263 }
264 + if not use_inductor_graph_partition:
265 + compilation_config["splitting_ops"] = []
266
267 tp_sp_args = [
268 *common_args,
tests/compile/passes/distributed/test_async_tp.py
+17
@@ -19,6 +19,7 @@ from vllm.config import (
19 VllmConfig,
20 set_current_vllm_config,
21 )
22 +from vllm.config.utils import Range
23 from vllm.distributed import (
24 tensor_model_parallel_all_gather,
25 tensor_model_parallel_reduce_scatter,
@@ -288,6 +289,22 @@ def test_async_tp_pass_replace(
289 run_torch_spawn(async_tp_pass_on_test_model, num_processes)
290
291
292 +def test_async_tp_pass_requires_full_graph_compilation():
293 + vllm_config = VllmConfig()
294 + vllm_config.compilation_config.use_inductor_graph_partition = False
295 + vllm_config.compilation_config.splitting_ops = [
296 + "vllm::unified_attention_with_output"
297 + ]
298 +
299 + async_tp_pass = object.__new__(AsyncTPPass)
300 + async_tp_pass.compilation_config = vllm_config.compilation_config
301 +
302 + with pytest.raises(
303 + AssertionError, match="AsyncTPPass requires full-graph compilation"
304 + ):
305 + async_tp_pass.is_applicable_for_range(Range(start=8, end=8))
306 +
307 +
308 def async_tp_pass_on_test_model(
309 local_rank: int,
310 world_size: int,
tests/compile/passes/distributed/test_sequence_parallelism.py
+19
@@ -22,6 +22,7 @@ from vllm.config import (
22 get_current_vllm_config,
23 set_current_vllm_config,
24 )
25 +from vllm.config.utils import Range
26 from vllm.distributed import tensor_model_parallel_all_reduce
27 from vllm.distributed.parallel_state import (
28 init_distributed_environment,
@@ -216,6 +217,24 @@ def test_sequence_parallelism_pass(
217 run_torch_spawn(sequence_parallelism_pass_on_test_model, num_processes)
218
219
220 +def test_sequence_parallelism_pass_requires_full_graph_compilation():
221 + vllm_config = VllmConfig()
222 + vllm_config.compilation_config.use_inductor_graph_partition = False
223 + vllm_config.compilation_config.splitting_ops = [
224 + "vllm::unified_attention_with_output"
225 + ]
226 +
227 + sequence_parallelism_pass = object.__new__(SequenceParallelismPass)
228 + sequence_parallelism_pass.compilation_config = vllm_config.compilation_config
229 + sequence_parallelism_pass.min_token_num = 1
230 +
231 + with pytest.raises(
232 + AssertionError,
233 + match="SequenceParallelismPass requires full-graph compilation",
234 + ):
235 + sequence_parallelism_pass.is_applicable_for_range(Range(start=8, end=8))
236 +
237 +
238 def sequence_parallelism_pass_on_test_model(
239 local_rank: int,
240 world_size: int,
tests/compile/test_config.py
+118 -1
@@ -407,7 +407,7 @@ def test_should_split():
407 (None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256),
408 # max from list
409 ([1, 2, 4, 15], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15),
410 - # filtered out 15 due to SP
410 + # SP forces full-graph compilation, sizes are filtered by TP
411 ([1, 2, 4, 15], None, 2, True, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4),
412 # limited by the max_tokens
413 ([1, 2, 4, 15], None, 1, False, 8, CUDAGraphMode.FULL_AND_PIECEWISE, 4),
@@ -465,6 +465,123 @@ def test_cudagraph_sizes_post_init(
465 )
466
467
468 +@pytest.mark.skipif(
469 + not current_platform.support_static_graph_mode(),
470 + reason="Skip if not cudagraph mode supported",
471 +)
472 +@pytest.mark.parametrize(
473 + (
474 + "cudagraph_mode",
475 + "use_inductor_graph_partition",
476 + "expected_enable_sp",
477 + "expected_cudagraph_mode",
478 + "expected_piecewise_compile",
479 + "expected_capture_sizes",
480 + "expected_max_size",
481 + ),
482 + [
483 + (CUDAGraphMode.PIECEWISE, False, True, CUDAGraphMode.FULL, False, [2, 4], 4),
484 + (
485 + CUDAGraphMode.FULL_DECODE_ONLY,
486 + False,
487 + True,
488 + CUDAGraphMode.FULL_DECODE_ONLY,
489 + False,
490 + [2, 4],
491 + 4,
492 + ),
493 + (
494 + CUDAGraphMode.FULL_AND_PIECEWISE,
495 + False,
496 + True,
497 + CUDAGraphMode.FULL,
498 + False,
499 + [2, 4],
500 + 4,
501 + ),
502 + (
503 + CUDAGraphMode.FULL_AND_PIECEWISE,
504 + True,
505 + True,
506 + CUDAGraphMode.FULL_AND_PIECEWISE,
507 + True,
508 + [2, 4],
509 + 4,
510 + ),
511 + ],
512 +)
513 +def test_sequence_parallelism_requires_full_graph_compilation(
514 + cudagraph_mode: CUDAGraphMode,
515 + use_inductor_graph_partition: bool,
516 + expected_enable_sp: bool,
517 + expected_cudagraph_mode: CUDAGraphMode,
518 + expected_piecewise_compile: bool,
519 + expected_capture_sizes: list[int],
520 + expected_max_size: int,
521 +):
522 + with patch.object(current_platform, "device_count", return_value=2):
523 + vllm_config = VllmConfig(
524 + parallel_config=ParallelConfig(tensor_parallel_size=2),
525 + scheduler_config=SchedulerConfig(
526 + max_num_seqs=128,
527 + max_num_batched_tokens=2048,
528 + max_model_len=2048,
529 + is_encoder_decoder=False,
530 + ),
531 + )
532 + vllm_config.model_config = MagicMock(
533 + dtype=torch.float16,
534 + enforce_eager=False,
535 + is_moe=False,
536 + disable_cascade_attn=False,
537 + get_hidden_size=MagicMock(return_value=4096),
538 + )
539 + vllm_config.compilation_config = CompilationConfig(
540 + mode=CompilationMode.VLLM_COMPILE,
541 + cudagraph_capture_sizes=[1, 2, 4, 15],
542 + max_cudagraph_capture_size=None,
543 + compile_sizes=["cudagraph_capture_sizes"],
544 + use_inductor_graph_partition=use_inductor_graph_partition,
545 + pass_config=PassConfig(
546 + enable_sp=True,
547 + fuse_gemm_comms=True,
548 + fuse_norm_quant=True,
549 + fuse_act_quant=True,
550 + eliminate_noops=True,
551 + sp_min_token_num=512,
552 + ),
553 + cudagraph_mode=cudagraph_mode,
554 + )
555 + vllm_config.compilation_config.set_splitting_ops_for_v1(
556 + all2all_backend=vllm_config.parallel_config.all2all_backend,
557 + data_parallel_size=1,
558 + )
559 + vllm_config._set_compile_ranges()
560 + vllm_config._set_cudagraph_sizes()
561 +
562 + assert (
563 + vllm_config.compilation_config.use_inductor_graph_partition
564 + == use_inductor_graph_partition
565 + )
566 + assert (
567 + bool(vllm_config.compilation_config.splitting_ops) == expected_piecewise_compile
568 + )
569 + assert vllm_config.compilation_config.pass_config.enable_sp == expected_enable_sp
570 + assert (
571 + vllm_config.compilation_config.pass_config.fuse_gemm_comms == expected_enable_sp
572 + )
573 + assert vllm_config.compilation_config.cudagraph_mode == expected_cudagraph_mode
574 + assert (
575 + vllm_config.compilation_config.cudagraph_capture_sizes == expected_capture_sizes
576 + )
577 + assert (
578 + vllm_config.compilation_config.max_cudagraph_capture_size == expected_max_size
579 + )
580 + assert (
581 + 511 in vllm_config.compilation_config.compile_ranges_endpoints
582 + ) == expected_enable_sp
583 +
584 +
585 def test_cached_compilation_config(default_vllm_config):
586 import torch
587 from torch._inductor.utils import run_and_get_code
vllm/compilation/passes/fusion/collective_fusion.py
+7 -10
@@ -406,16 +406,13 @@ class AsyncTPPass(VllmPatternMatcherPass):
406 self.dump_patterns(config, self.patterns)
407
408 def is_applicable_for_range(self, compile_range: Range) -> bool:
409 - # This pass is applied on top of the sequence parallelism pass.
410 - # It inherits the same applicability condition as `SequenceParallelismPass`.
411 - # See `SequenceParallelismPass.is_applicable` for more details.
412 - if (
413 - not self.compilation_config.splitting_ops
414 - or self.compilation_config.use_inductor_graph_partition
415 - ):
416 - return True
417 - tp_size = get_tensor_model_parallel_world_size()
418 - return bool(compile_range.is_single_size() and compile_range.end % tp_size == 0)
409 + # This pass is applied on top of the sequence parallelism pass,
410 + # which is only supported in fullgraph compilation mode.
411 + assert (
412 + self.compilation_config.use_inductor_graph_partition
413 + or not self.compilation_config.splitting_ops
414 + ), "AsyncTPPass requires full-graph compilation"
415 + return True
416
417 @VllmInductorPass.time_and_log
418 def __call__(self, graph: fx.Graph) -> None:
vllm/compilation/passes/fusion/sequence_parallelism.py
+16 -26
@@ -341,22 +341,18 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
341 significantly reduce communication overhead and improve overall model
342 performance.
343
344 -
345 - This pass splits up the residual tensor across TP ranks and hence divides its size.
346 - Because the pattern matcher starts at the end of the graph, the replacement
347 - contains a slice that temporarily conforms the input residual to the correct size.
348 - After all patterns have been matched, we use a NoOpEliminationPass to clean up
349 - what have now become no-op slices.
350 -
351 - Note that an older version of the pass did not need this as it operated only on
352 - custom rms_norm and fused_rms_norm_add custom ops which did not complain about
353 - mismatched shapes during replacement. So this approach has the same assumption that
354 - correctness is only maintained if all rms_norm operations are split across ranks.
355 -
356 - Correctness-wise, this is approach strictly better than before - before,
357 - the graph was incorrect semantically and shape-wise during the pass.
358 - With this approach there's only semantic incorrectness during the pass.
359 - Both approaches restore a correct graph once all patterns are matched.
344 + This pass is only supported when compiling the whole graph (fullgraph
345 + mode, i.e. using Inductor graph partition or empty splitting_ops).
346 + Piecewise compilation is not supported because the residual tensor
347 + gets split across TP ranks, causing size mismatches at subgraph
348 + boundaries.
349 +
350 + This pass splits up the residual tensor across TP ranks and hence
351 + divides its size. Because the pattern matcher starts at the end of
352 + the graph, the replacement contains a slice that temporarily conforms
353 + the input residual to the correct size. After all patterns have been
354 + matched, we use a NoOpEliminationPass to clean up what have now
355 + become no-op slices.
356 """
357
358 @enable_fake_mode
@@ -419,19 +415,13 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
415 and gathering tensors across TP ranks outweighs the benefits.
416
417 Returns False (SP disabled) when:
422 - - Using piecewise compilation with non-concrete or TP-indivisible sizes
418 - min_token_num is None (SP disabled for this device/config)
419 - The compile range starts below the minimum token threshold
420 """
426 - # For piecewise compilation (not using inductor graph partition),
427 - # we need concrete sizes that are divisible by TP for correct splitting
428 - if (
429 - not self.compilation_config.use_inductor_graph_partition
430 - and self.compilation_config.splitting_ops
431 - ):
432 - tp_size = get_tensor_model_parallel_world_size()
433 - if not compile_range.is_single_size() or compile_range.end % tp_size != 0:
434 - return False
421 + assert (
422 + self.compilation_config.use_inductor_graph_partition
423 + or not self.compilation_config.splitting_ops
424 + ), "SequenceParallelismPass requires full-graph compilation"
425
426 # min_token_num is None when SP is disabled for this device/config
427 # (e.g., non-CUDA platform, unsupported GPU, or small hidden_size)
vllm/config/compilation.py
+19
@@ -1148,6 +1148,25 @@ class CompilationConfig:
1148 self.cudagraph_mode = CUDAGraphMode.FULL
1149 self.splitting_ops = []
1150
1151 + if (
1152 + not self.use_inductor_graph_partition
1153 + and (self.pass_config.enable_sp or self.pass_config.fuse_gemm_comms)
1154 + and self.splitting_ops
1155 + ):
1156 + logger.warning_once(
1157 + "Sequence parallelism requires full-graph compilation when "
1158 + "use_inductor_graph_partition is off. Setting splitting_ops "
1159 + "to an empty list to preserve SP and async TP."
1160 + )
1161 + self.splitting_ops = []
1162 + if self.cudagraph_mode.has_piecewise_cudagraphs():
1163 + logger.warning_once(
1164 + "Sequence parallelism is incompatible with piecewise "
1165 + "cudagraph when use_inductor_graph_partition is off. "
1166 + "Setting cudagraph_mode to FULL."
1167 + )
1168 + self.cudagraph_mode = CUDAGraphMode.FULL
1169 +
1170 # Disable CUDA graphs for DeepEP high-throughput since its not CG compatible
1171 if (
1172 all2all_backend == "deepep_high_throughput"
vllm/config/vllm.py
+17 -28
@@ -983,19 +983,16 @@ class VllmConfig:
983 )
984 self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
985
986 - # async tp is built on top of sequence parallelism
987 - # and requires it to be enabled.
988 - if self.compilation_config.pass_config.fuse_gemm_comms:
989 - self.compilation_config.pass_config.enable_sp = True
990 - if self.compilation_config.pass_config.enable_sp:
986 + # async tp is built on top of sequence parallelism and requires it.
987 + pass_config = self.compilation_config.pass_config
988 + if pass_config.fuse_gemm_comms:
989 + pass_config.enable_sp = True
990 + if pass_config.enable_sp:
991 if self.parallel_config.tensor_parallel_size == 1:
992 logger.warning("Sequence Parallelism requires TP>1, disabling")
993 - self.compilation_config.pass_config.enable_sp = False
994 - self.compilation_config.pass_config.fuse_gemm_comms = False
993 + pass_config.enable_sp = False
994 + pass_config.fuse_gemm_comms = False
995 else:
996 - # Compute SP threshold early; disable if None (model too
997 - # small for SP to be beneficial).
998 - pass_config = self.compilation_config.pass_config
996 if pass_config.sp_min_token_num is None:
997 from vllm.compilation.passes.fusion.sequence_parallelism import (
998 get_sequence_parallelism_threshold,
@@ -1015,8 +1012,8 @@ class VllmConfig:
1012 "threshold heuristic, disabling. To force SP, "
1013 "set pass_config.sp_min_token_num manually."
1014 )
1018 - self.compilation_config.pass_config.enable_sp = False
1019 - self.compilation_config.pass_config.fuse_gemm_comms = False
1015 + pass_config.enable_sp = False
1016 + pass_config.fuse_gemm_comms = False
1017
1018 from vllm.utils.torch_utils import HAS_OPAQUE_TYPE
1019
@@ -1098,6 +1095,7 @@ class VllmConfig:
1095 self.compilation_config.cudagraph_num_of_warmups = 1
1096
1097 self._set_cudagraph_sizes()
1098 +
1099 else:
1100 self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
1101
@@ -1171,8 +1169,8 @@ class VllmConfig:
1169 )
1170
1171 if self.compilation_config.pass_config.enable_sp:
1174 - # With pipeline parallelism or dynamo partitioning,
1175 - # native rms norm tracing errors due to incorrect residual shape.
1172 + # With pipeline parallelism, native rms norm tracing errors due to
1173 + # incorrect residual shape.
1174 # Use custom rms norm to unblock. In the future,
1175 # the pass will operate on higher-level IR to avoid the issue.
1176 # TODO: https://github.com/vllm-project/vllm/issues/27894
@@ -1183,24 +1181,15 @@ class VllmConfig:
1181 self.compilation_config.mode,
1182 )
1183
1186 - is_fullgraph = (
1187 - self.compilation_config.use_inductor_graph_partition
1188 - or len(self.compilation_config.splitting_ops or []) == 0
1189 - )
1190 - if self.parallel_config.pipeline_parallel_size > 1 or not is_fullgraph:
1184 + if self.parallel_config.pipeline_parallel_size > 1:
1185 if "-rms_norm" not in self.compilation_config.custom_ops:
1186 self.compilation_config.custom_ops.append("+rms_norm")
1187 else:
1194 - regime = (
1195 - "Dynamo partition"
1196 - if not is_fullgraph
1197 - else "pipeline parallelism"
1198 - )
1188 logger.warning_once(
1189 "Sequence parallelism not supported with "
1190 "native rms_norm when using %s, "
1191 "this will likely lead to an error.",
1203 - regime,
1192 + "pipeline parallelism",
1193 )
1194
1195 # final check of cudagraph mode after all possible updates
@@ -1212,9 +1201,9 @@ class VllmConfig:
1201 and not self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs() # noqa: E501
1202 ):
1203 logger.warning_once(
1215 - "No piecewise cudagraph for executing cascade attention."
1216 - " Will fall back to eager execution if a batch runs "
1217 - "into cascade attentions."
1204 + "No piecewise cudagraph for executing cascade attention. "
1205 + "Will fall back to eager execution if a batch runs into "
1206 + "cascade attentions."
1207 )
1208
1209 if self.compilation_config.cudagraph_mode.requires_piecewise_compilation():
vllm/v1/worker/utils.py
+8 -15
@@ -519,12 +519,8 @@ def is_residual_scattered_for_sp(
519 """Check if the residual tensor is scattered for sequence parallelism.
520
521 The residual tensor is scattered across tensor parallel ranks when sequence
522 - parallelism and tensor parallelism is enabled.
523 -
524 - This follows the same logic as SequenceParallelismPass.is_applicable_for_range():
525 - - In full-graph compilation mode (no splitting ops or using inductor graph
526 - partition), SP is always applied
527 - - Otherwise, SP is only applied for specific shapes in compile_sizes
522 + parallelism and tensor parallelism is enabled. SP is only supported in
523 + full-graph compilation mode.
524 """
525 if not vllm_config.compilation_config.pass_config.enable_sp:
526 return False
@@ -534,16 +530,13 @@ def is_residual_scattered_for_sp(
530 if tp == 1:
531 return False
532
533 + assert (
534 + vllm_config.compilation_config.use_inductor_graph_partition
535 + or not vllm_config.compilation_config.splitting_ops
536 + ), "Sequence parallelism requires full-graph compilation"
537 +
538 # When sequence parallelism is enabled, we always pad num_input_tokens
539 # to be a multiple of tensor_parallel_size (tp) earlier.
540 assert num_input_tokens % tp == 0
541
541 - if (
542 - not vllm_config.compilation_config.splitting_ops
543 - or vllm_config.compilation_config.use_inductor_graph_partition
544 - ):
545 - return True
546 - compile_sizes = vllm_config.compilation_config.compile_sizes
547 - if compile_sizes is None:
548 - return False
549 - return num_input_tokens in compile_sizes
542 + return True