[compile] Turn on TP/SP when use_inductor_graph_partition=True

Signed-off-by: angelayi <[email protected]>

Author: angelayi

vllm/compilation/collective_fusion.py

@@ -435,7 +435,8 @@ def is_applicable(self, shape: int | None) -> bool:
         # This pass is applied on top of the sequence parallelism pass.
         # It inherits the same applicability condition as `SequenceParallelismPass`.
         # See `SequenceParallelismPass.is_applicable` for more details.
-        if self.splitting_ops is None or self.splitting_ops == []:
+        splitting_ops = self.compilation_config.splitting_ops
+        if not splitting_ops or self.compilation_config.use_inductor_graph_partition:
             return True
         tp_size = get_tensor_model_parallel_world_size()
         return shape is not None and shape % tp_size == 0

vllm/compilation/pass_manager.py

@@ -74,6 +74,8 @@ def __call__(self, graph: fx.Graph):
             if pass_.is_applicable(shape):
                 pass_(graph)
                 VllmInductorPass.dump_prefix += 1
+            else:
+                logger.debug(f"Skipping {pass_} with shape {shape}")
 
         # post-cleanup goes before fix_functionalization
         # because it requires a functional graph

vllm/compilation/sequence_parallelism.py

@@ -496,7 +496,8 @@ def is_applicable(self, shape: int | None) -> bool:
         # 2. For specific shape provided during compilation (e.g., from
         #    `compile_sizes`), which must be divisible by the tensor-parallel
         #    size.
-        if self.splitting_ops is None or self.splitting_ops == []:
+        splitting_ops = self.compilation_config.splitting_ops
+        if not splitting_ops or self.compilation_config.use_inductor_graph_partition:
             return True
         tp_size = get_tensor_model_parallel_world_size()
         return shape is not None and shape % tp_size == 0

vllm/compilation/vllm_inductor_pass.py

@@ -28,12 +28,10 @@ class VllmInductorPass(InductorPass):
     """Keep track of pass index for debug dump ordering."""
 
     def __init__(self, config: VllmConfig):
+        self.compilation_config = config.compilation_config
         self.pass_config = config.compilation_config.pass_config
-        self.splitting_ops = config.compilation_config.splitting_ops
-        self.model_dtype = config.model_config.dtype if config.model_config \
-            else None
-        self.device = config.device_config.device if config.device_config \
-            else None
+        self.model_dtype = config.model_config.dtype if config.model_config else None
+        self.device = config.device_config.device if config.device_config else None
         self.pass_name = self.__class__.__name__
 
     @staticmethod