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
[bugfix] Fix SP + PP without specifying compile size (#26955)
Signed-off-by: angelayi <yiangela7@gmail.com>
Angela Yi committed
Oct 15, 2025 at 20:05 UTC
e19b16dde62ecf2dcbbaac7d18a276bdda6935e6
2 files changed
+21
-3
tests/distributed/test_sequence_parallel.py
+9
@@ -18,6 +18,7 @@ import pytest
18
from vllm.config.compilation import CompilationMode
19
from vllm.config.model import RunnerOption
20
from vllm.logger import init_logger
21
+from vllm.utils import is_torch_equal_or_newer
22
23
from ..models.registry import HF_EXAMPLE_MODELS
24
from ..utils import compare_two_settings, create_new_process_for_each_test
@@ -159,6 +160,7 @@ def _compare_sp(
160
runner: RunnerOption,
161
test_options: SPTestOptions,
162
num_gpus_available: int,
163
+ use_inductor_graph_partition: bool,
164
*,
165
method: Literal["generate", "encode"],
166
is_multimodal: bool,
@@ -243,6 +245,7 @@ def _compare_sp(
245
"enable_fusion": enable_fusion,
246
"enable_noop": True,
247
},
248
+ "use_inductor_graph_partition": use_inductor_graph_partition,
249
}
250
251
tp_sp_args = [
@@ -297,6 +300,7 @@ SP_TEST_MODELS = [
300
if model_id in SP_TEST_MODELS
301
],
302
)
303
+@pytest.mark.parametrize("use_inductor_graph_partition", [True, False])
304
@create_new_process_for_each_test()
305
def test_tp_sp_generation(
306
model_id: str,
@@ -305,7 +309,11 @@ def test_tp_sp_generation(
309
runner: RunnerOption,
310
test_options: SPTestOptions,
311
num_gpus_available,
312
+ use_inductor_graph_partition: bool,
313
):
314
+ if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
315
+ pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
316
+
317
_compare_sp(
318
model_id,
319
parallel_setup,
@@ -313,6 +321,7 @@ def test_tp_sp_generation(
321
runner,
322
test_options,
323
num_gpus_available,
324
+ use_inductor_graph_partition,
325
method="generate",
326
is_multimodal=False,
327
)
vllm/v1/worker/utils.py
+12
-3
@@ -328,8 +328,12 @@ def is_residual_scattered_for_sp(
328
"""Check if the residual tensor is scattered for sequence parallelism.
329
330
The residual tensor is scattered across tensor parallel ranks when sequence
331
- parallelism and tensor parallelism is enabled, and the number of
332
- input tokens is one of the compilation sizes.
331
+ parallelism and tensor parallelism is enabled.
332
+
333
+ This follows the same logic as SequenceParallelismPass.is_applicable():
334
+ - In full-graph compilation mode (no splitting ops or using inductor graph
335
+ partition), SP is always applied
336
+ - Otherwise, SP is only applied for specific shapes in compile_sizes
337
"""
338
if not vllm_config.compilation_config.pass_config.enable_sequence_parallelism:
339
return False
@@ -343,5 +347,10 @@ def is_residual_scattered_for_sp(
347
# to be a multiple of tensor_parallel_size (tp) earlier.
348
assert num_input_tokens % tp == 0
349
346
- # Currently, SP is only enabled for static size fx graphs.
350
+ if (
351
+ not vllm_config.compilation_config.splitting_ops
352
+ or vllm_config.compilation_config.use_inductor_graph_partition
353
+ ):
354
+ return True
355
+
356
return num_input_tokens in vllm_config.compilation_config.compile_sizes