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 14, 05:30 UTC
[compile] Enable sequence parallelism for full cuda graph without specifying compile sizes (#26681)
Signed-off-by: angelayi <yiangela7@gmail.com>
Angela Yi committed
Oct 13, 2025 at 18:15 UTC
b59dd19b55036050cf491501614fb46cedf5c5c1
5 files changed
+34
-5
vllm/compilation/collective_fusion.py
+9
-2
@@ -431,8 +431,15 @@ class AsyncTPPass(VllmPatternMatcherPass):
431
432
self.dump_patterns(config, self.patterns)
433
434
- def is_applicable_for_shape(self, shape: int | None) -> bool:
435
- # only do replace for specific shapes
434
+ def is_applicable(self, shape: int | None) -> bool:
435
+ # This pass is applied on top of the sequence parallelism pass.
436
+ # It inherits the same applicability condition as `SequenceParallelismPass`.
437
+ # See `SequenceParallelismPass.is_applicable` for more details.
438
+ if (
439
+ not self.compilation_config.splitting_ops
440
+ or self.compilation_config.use_inductor_graph_partition
441
+ ):
442
+ return True
443
tp_size = get_tensor_model_parallel_world_size()
444
return shape is not None and shape % tp_size == 0
445
vllm/compilation/inductor_pass.py
+1
-1
@@ -96,7 +96,7 @@ class InductorPass(CustomGraphPass):
96
encoded = json.dumps(dict_, sort_keys=True).encode("utf-8")
97
return hashlib.sha256(encoded).hexdigest()
98
99
- def is_applicable_for_shape(self, shape: int | None):
99
+ def is_applicable(self, shape: int | None):
100
return True
101
102
vllm/compilation/pass_manager.py
+3
-1
@@ -71,9 +71,11 @@ class PostGradPassManager(CustomGraphPass):
71
72
shape = get_pass_context().runtime_shape
73
for pass_ in self.passes:
74
- if pass_.is_applicable_for_shape(shape):
74
+ if pass_.is_applicable(shape):
75
pass_(graph)
76
VllmInductorPass.dump_prefix += 1
77
+ else:
78
+ logger.debug("Skipping %s with shape %s", pass_, shape)
79
80
# post-cleanup goes before fix_functionalization
81
# because it requires a functional graph
vllm/compilation/sequence_parallelism.py
+19
-1
@@ -482,7 +482,25 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
482
).register(self.patterns)
483
self.dump_patterns(config, self.patterns)
484
485
- def is_applicable_for_shape(self, shape: int | None) -> bool:
485
+ def is_applicable(self, shape: int | None) -> bool:
486
+ # When sequence parallelism is enabled, the residual tensor from RMSNorm
487
+ # needs to be split along the sequence dimension. However, this dimension
488
+ # is symbolic during piecewise compilation, and splitting symbolic shapes
489
+ # is not supported.
490
+ #
491
+ # This pass is therefore only applied when the sequence dimension is
492
+ # concrete:
493
+ # 1. In full-graph compilation mode (no Dynamo splitting ops are used).
494
+ # For this case we always pad num_tokens to be a multiple of
495
+ # tensor_parallel_size, so there's no need to check shape % tp_size == 0.
496
+ # 2. For specific shape provided during compilation (e.g., from
497
+ # `compile_sizes`), which must be divisible by the tensor-parallel
498
+ # size.
499
+ if (
500
+ not self.compilation_config.splitting_ops
501
+ or self.compilation_config.use_inductor_graph_partition
502
+ ):
503
+ return True
504
tp_size = get_tensor_model_parallel_world_size()
505
return shape is not None and shape % tp_size == 0
506
vllm/compilation/vllm_inductor_pass.py
+2
@@ -3,6 +3,7 @@
3
import functools
4
import operator
5
import time
6
+import weakref
7
from typing import ClassVar
8
9
import regex as re
@@ -28,6 +29,7 @@ class VllmInductorPass(InductorPass):
29
"""Keep track of pass index for debug dump ordering."""
30
31
def __init__(self, config: VllmConfig):
32
+ self.compilation_config = weakref.proxy(config.compilation_config)
33
self.pass_config = config.compilation_config.pass_config
34
self.model_dtype = config.model_config.dtype if config.model_config else None
35
self.device = config.device_config.device if config.device_config else None