@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9d37941017
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 7, 17:45 UTC

[torch.compile] Sequence Parallelism threshold compile ranges (#28672)

Signed-off-by: jasonlizhengjian <jasonlizhengjian@gmail.com> Signed-off-by: Jason Li <jasonlizhengjian@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

Jason Li committed Feb 25, 2026 at 21:00 UTC 9d379410179b649f4e7651940debc35c4ac7c0a5
8 files changed +524 -32
tests/compile/conftest.py new
+34
@@ -0,0 +1,34 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from contextlib import contextmanager
5 +from unittest.mock import MagicMock, patch
6 +
7 +import pytest
8 +
9 +from vllm.platforms.interface import DeviceCapability
10 +
11 +
12 +@pytest.fixture
13 +def mock_cuda_platform():
14 + """
15 + Fixture that returns a factory for creating mocked CUDA platforms.
16 +
17 + Usage:
18 + def test_something(mock_cuda_platform):
19 + with mock_cuda_platform(is_cuda=True, capability=(9, 0)):
20 + # test code
21 + """
22 +
23 + @contextmanager
24 + def _mock_platform(is_cuda: bool = True, capability: tuple[int, int] | None = None):
25 + mock_platform = MagicMock()
26 + mock_platform.is_cuda.return_value = is_cuda
27 + if capability is not None:
28 + mock_platform.get_device_capability.return_value = DeviceCapability(
29 + *capability
30 + )
31 + with patch("vllm.platforms.current_platform", mock_platform):
32 + yield mock_platform
33 +
34 + return _mock_platform
tests/compile/fusions_e2e/conftest.py
+80 -9
@@ -94,7 +94,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
94 run_model(full_compilation_config, model_name, **model_kwargs)
95
96 num_compile_ranges = len(full_compilation_config.get_compile_ranges())
97 - assert num_compile_ranges in [1, 2]
97 + assert num_compile_ranges in [1, 2, 3]
98
99 print(f"Compile ranges: {full_compilation_config.get_compile_ranges()}")
100 print("Fusion results:")
@@ -107,12 +107,33 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
107
108 # Now check the matches
109 for match_name in matches_check:
110 - num_ranges_activated = (
111 - 1 if match_name == "ar_rms_fusion" else num_compile_ranges
112 - )
113 - n_expected = tp_size * num_ranges_activated
114 -
110 log_matches = list(int(ms) for ms in log_matches_dict[match_name])
111 +
112 + # AR+RMS skips the largest range; SP skips the smallest.
113 + # When both are enabled, AR+RMS activation count is
114 + # model-dependent (hidden_size affects threshold), so derive
115 + # from log data.
116 + if (
117 + match_name == "ar_rms_fusion"
118 + and "sequence_parallel" in matches_check
119 + and num_compile_ranges >= 2
120 + ):
121 + assert (
122 + len(log_matches) >= tp_size and len(log_matches) % tp_size == 0
123 + ), (
124 + f"Expected multiple of {tp_size} ar_rms log entries, "
125 + f"found {len(log_matches)}"
126 + )
127 + num_ranges_activated = len(log_matches) // tp_size
128 + elif (
129 + match_name in ("ar_rms_fusion", "sequence_parallel")
130 + and num_compile_ranges >= 2
131 + ):
132 + num_ranges_activated = num_compile_ranges - 1
133 + else:
134 + num_ranges_activated = num_compile_ranges
135 +
136 + n_expected = tp_size * num_ranges_activated
137 assert len(log_matches) == n_expected, (
138 f"Could not find {n_expected} {match_name} "
139 f"(found {len(log_matches)}) in:\n {log_holder.text}"
@@ -122,8 +143,8 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
143
144 if match_name == "rms_quant_fusion" and "ar_rms_fusion" in matches_check:
145 # AR+rms+quant takes precedence over rms+quant if activated.
125 - # That means we get full matching where ar+rms+quant was not activated,
126 - # and less where it was
146 + # That means we get full matching where ar+rms+quant was not
147 + # activated, and less where it was (only the smallest range).
148 assert sum(m == expected_matches for m in log_matches) == tp_size * (
149 num_ranges_activated - 1
150 ), "Expecting full rms+quant fusion where ar+rms+quant not activated"
@@ -135,6 +156,43 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
156 f"Expecting at least {expected_matches - matches.ar_rms_fusion} "
157 f"where ar+rms+quant was activated"
158 )
159 + elif (
160 + match_name == "async_tp"
161 + and "sequence_parallel" in matches_check
162 + and num_compile_ranges >= 2
163 + ):
164 + # AsyncTP only finds patterns on ranges where SP ran.
165 + n_sp_ranges = num_compile_ranges - 1
166 + assert (
167 + sum(m == expected_matches for m in log_matches)
168 + == tp_size * n_sp_ranges
169 + ), (
170 + f"Expecting {expected_matches} async_tp on "
171 + f"{tp_size * n_sp_ranges} SP-range entries, "
172 + f"found: {log_matches}"
173 + )
174 + assert sum(m == 0 for m in log_matches) == tp_size, (
175 + f"Expecting 0 async_tp on {tp_size} small-range entries "
176 + f"(no SP), found: {log_matches}"
177 + )
178 + elif (
179 + match_name == "ar_rms_fusion"
180 + and "sequence_parallel" in matches_check
181 + and num_compile_ranges >= 2
182 + ):
183 + # SP consumes allreduce patterns first, so AR+RMS finds
184 + # full matches only on the smallest range (no SP).
185 + assert sum(m == expected_matches for m in log_matches) == tp_size, (
186 + f"Expecting {expected_matches} ar_rms on "
187 + f"{tp_size} small-range entries, found: {log_matches}"
188 + )
189 + assert sum(m == 0 for m in log_matches) == tp_size * (
190 + num_ranges_activated - 1
191 + ), (
192 + f"Expecting 0 ar_rms on "
193 + f"{tp_size * (num_ranges_activated - 1)} large-range "
194 + f"entries (SP took precedence), found: {log_matches}"
195 + )
196 else:
197 expected_matches_list = [expected_matches] * n_expected
198 assert sorted(log_matches) == expected_matches_list, (
@@ -142,7 +200,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
200 f"found: {sorted(log_matches)}"
201 )
202
145 - if match_name == "ar_rms_fusion":
203 + if match_name == "ar_rms_fusion" and num_compile_ranges >= 2:
204 log_matches = re.findall(
205 r"pass_manager.py:\d+] Skipping "
206 r".*AllReduceFusionPass.* with compile range",
@@ -155,4 +213,17 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
213 f"(found {len(log_matches)}) in:\n {log_holder.text}"
214 )
215
216 + if match_name == "sequence_parallel" and num_compile_ranges >= 2:
217 + log_matches = re.findall(
218 + r"pass_manager.py:\d+] Skipping "
219 + r".*SequenceParallelismPass.* with compile range",
220 + log_holder.text,
221 + )
222 +
223 + n_expected = tp_size * (num_compile_ranges - num_ranges_activated)
224 + assert len(log_matches) == n_expected, (
225 + f'Could not find {n_expected} "Skipping SequenceParallelismPass" '
226 + f"(found {len(log_matches)}) in:\n {log_holder.text}"
227 + )
228 +
229 return run
tests/compile/fusions_e2e/test_tp2_async_tp.py
+133
@@ -66,6 +66,9 @@ def test_tp2_async_tp_fp8_fusions(
66 enable_qk_norm_rope_fusion=True,
67 enable_sp=True,
68 fuse_gemm_comms=True,
69 + fuse_allreduce_rms=False,
70 + # Override threshold for testing (models have small hidden_size)
71 + sp_min_token_num=512,
72 ),
73 )
74
@@ -123,6 +126,9 @@ def test_tp2_async_tp_fusions(
126 enable_qk_norm_rope_fusion=True,
127 enable_sp=True,
128 fuse_gemm_comms=True,
129 + fuse_allreduce_rms=False,
130 + # Override threshold for testing (models have small hidden_size)
131 + sp_min_token_num=512,
132 ),
133 )
134
@@ -141,3 +147,130 @@ def test_tp2_async_tp_fusions(
147 matches_check,
148 tp_size=2,
149 )
150 +
151 +
152 +@multi_gpu_test(num_gpus=2)
153 +@pytest.mark.parametrize(
154 + "model_name, matches_fn, model_kwargs, hf_overrides",
155 + [llama3_8b_fp8, llama4_scout_fp8],
156 +)
157 +@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN])
158 +@pytest.mark.parametrize("n_layers", [4])
159 +@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
160 +@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
161 +def test_tp2_sp_ar_rms_fp8_fusions(
162 + model_name: str,
163 + matches_fn: Callable[[int], Matches],
164 + model_kwargs: dict,
165 + hf_overrides: Callable[[int], dict],
166 + attn_backend: AttentionBackendCase,
167 + n_layers: int,
168 + custom_ops: str,
169 + inductor_graph_partition: bool,
170 + run_e2e_fusion_test,
171 + monkeypatch,
172 +):
173 + matches = matches_fn(n_layers)
174 +
175 + if is_blackwell():
176 + # Disable FlashInfer scaled_mm FP8 as it's not supported in async tp patterns
177 + monkeypatch.setenv("VLLM_DISABLED_KERNELS", "FlashInferFP8ScaledMMLinearKernel")
178 +
179 + # Reduce size of model and skip weight loading time
180 + model_kwargs["hf_overrides"] = hf_overrides(n_layers)
181 + model_kwargs["load_format"] = "dummy"
182 + model_kwargs["max_model_len"] = 1024
183 +
184 + compilation_config = dict(
185 + use_inductor_graph_partition=inductor_graph_partition,
186 + custom_ops=custom_ops.split(","),
187 + pass_config=PassConfig(
188 + fuse_norm_quant=True,
189 + fuse_act_quant=True,
190 + fuse_attn_quant=True,
191 + enable_qk_norm_rope_fusion=True,
192 + enable_sp=True,
193 + fuse_gemm_comms=True,
194 + fuse_allreduce_rms=True,
195 + # Override threshold for testing (models have small hidden_size)
196 + sp_min_token_num=512,
197 + ),
198 + )
199 +
200 + matches_check = [
201 + "rms_quant_fusion",
202 + "act_quant_fusion",
203 + "norm_rope_fusion",
204 + "attn_quant_fusion",
205 + "ar_rms_fusion",
206 + "sequence_parallel",
207 + "async_tp",
208 + ]
209 +
210 + run_e2e_fusion_test(
211 + model_name,
212 + matches,
213 + model_kwargs,
214 + attn_backend,
215 + compilation_config,
216 + matches_check,
217 + tp_size=2,
218 + )
219 +
220 +
221 +@multi_gpu_test(num_gpus=2)
222 +@pytest.mark.parametrize(
223 + "model_name, matches_fn, model_kwargs, hf_overrides",
224 + [llama3_8b, qwen3_a3b],
225 +)
226 +@pytest.mark.parametrize("attn_backend", [TRITON_ATTN])
227 +@pytest.mark.parametrize("n_layers", [4])
228 +@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
229 +@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
230 +def test_tp2_sp_ar_rms_fusions(
231 + model_name: str,
232 + matches_fn: Callable[[int], Matches],
233 + model_kwargs: dict,
234 + hf_overrides: Callable[[int], dict],
235 + attn_backend: AttentionBackendCase,
236 + n_layers: int,
237 + custom_ops: str,
238 + inductor_graph_partition: bool,
239 + run_e2e_fusion_test,
240 +):
241 + matches = matches_fn(n_layers)
242 +
243 + # Reduce size of model and skip weight loading time
244 + model_kwargs["hf_overrides"] = hf_overrides(n_layers)
245 + model_kwargs["load_format"] = "dummy"
246 + model_kwargs["max_model_len"] = 1024
247 +
248 + compilation_config = dict(
249 + use_inductor_graph_partition=inductor_graph_partition,
250 + custom_ops=custom_ops.split(","),
251 + pass_config=PassConfig(
252 + enable_qk_norm_rope_fusion=True,
253 + enable_sp=True,
254 + fuse_gemm_comms=True,
255 + fuse_allreduce_rms=True,
256 + # Override threshold for testing (models have small hidden_size)
257 + sp_min_token_num=512,
258 + ),
259 + )
260 +
261 + matches_check = [
262 + "norm_rope_fusion",
263 + "ar_rms_fusion",
264 + "sequence_parallel",
265 + "async_tp",
266 + ]
267 +
268 + run_e2e_fusion_test(
269 + model_name,
270 + matches,
271 + model_kwargs,
272 + attn_backend,
273 + compilation_config,
274 + matches_check,
275 + tp_size=2,
276 + )
tests/compile/test_config.py
+1
@@ -421,6 +421,7 @@ def test_cudagraph_sizes_post_init(
421 fuse_norm_quant=True,
422 fuse_act_quant=True,
423 eliminate_noops=True,
424 + sp_min_token_num=512 if enable_sp else None,
425 ),
426 cudagraph_mode=cudagraph_mode,
427 )
tests/compile/test_sequence_parallelism_threshold.py new
+110
@@ -0,0 +1,110 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from vllm.compilation.passes.fusion.sequence_parallelism import (
7 + SP_MIN_HIDDEN_SIZE,
8 + SP_MIN_PER_GPU_SIZE_MB,
9 + get_sequence_parallelism_threshold,
10 +)
11 +
12 +
13 +class TestGetSequenceParallelismThreshold:
14 + """Tests for get_sequence_parallelism_threshold function."""
15 +
16 + def test_non_cuda_returns_none(self, mock_cuda_platform):
17 + """Non-CUDA platforms should return None."""
18 + with mock_cuda_platform(is_cuda=False):
19 + result = get_sequence_parallelism_threshold(
20 + hidden_size=8192, tp_size=2, element_size=2
21 + )
22 + assert result is None
23 +
24 + def test_unsupported_device_capability_returns_none(self, mock_cuda_platform):
25 + """Unsupported device capabilities (e.g., sm80) should return None."""
26 + with mock_cuda_platform(capability=(8, 0)):
27 + result = get_sequence_parallelism_threshold(
28 + hidden_size=8192, tp_size=2, element_size=2
29 + )
30 + assert result is None
31 +
32 + def test_small_hidden_size_returns_none(self, mock_cuda_platform):
33 + """H100 with hidden_size below threshold should return None."""
34 + with mock_cuda_platform(capability=(9, 0)):
35 + result = get_sequence_parallelism_threshold(
36 + hidden_size=4096,
37 + tp_size=2,
38 + element_size=2, # 4096 < 8192
39 + )
40 + assert result is None
41 +
42 + def test_h100_large_model_returns_threshold(self, mock_cuda_platform):
43 + """H100 with large enough hidden_size should return calculated threshold."""
44 + with mock_cuda_platform(capability=(9, 0)):
45 + hidden_size = 8192
46 + tp_size = 2
47 + element_size = 2 # float16/bfloat16
48 +
49 + result = get_sequence_parallelism_threshold(
50 + hidden_size=hidden_size,
51 + tp_size=tp_size,
52 + element_size=element_size,
53 + )
54 +
55 + # Verify calculation: (8 * 2 * 1024 * 1024) // (8192 * 2) = 1024
56 + MiB = 1024 * 1024
57 + expected = int(
58 + (SP_MIN_PER_GPU_SIZE_MB[90] * tp_size * MiB)
59 + // (hidden_size * element_size)
60 + )
61 + assert result == expected
62 + assert result == 1024
63 +
64 + @pytest.mark.parametrize(
65 + "hidden_size,tp_size,element_size,expected",
66 + [
67 + # Boundary: exactly at min hidden size threshold, tp_size=1
68 + # (8 * 1 * 1024 * 1024) // (8192 * 2) = 512
69 + (8192, 1, 2, 512),
70 + # Larger hidden size reduces token threshold
71 + # (8 * 1 * 1024 * 1024) // (16384 * 2) = 256
72 + (16384, 1, 2, 256),
73 + # Larger tp_size increases token threshold
74 + # (8 * 4 * 1024 * 1024) // (8192 * 2) = 2048
75 + (8192, 4, 2, 2048),
76 + # Larger element_size (fp32) reduces token threshold
77 + # (8 * 2 * 1024 * 1024) // (8192 * 4) = 512
78 + (8192, 2, 4, 512),
79 + ],
80 + )
81 + def test_threshold_calculation_variations(
82 + self, mock_cuda_platform, hidden_size, tp_size, element_size, expected
83 + ):
84 + """Test threshold calculation with various parameter combinations."""
85 + with mock_cuda_platform(capability=(9, 0)):
86 + result = get_sequence_parallelism_threshold(
87 + hidden_size=hidden_size,
88 + tp_size=tp_size,
89 + element_size=element_size,
90 + )
91 + assert result == expected
92 +
93 + def test_hidden_size_boundary(self, mock_cuda_platform):
94 + """Test behavior at the exact hidden_size boundary."""
95 + with mock_cuda_platform(capability=(9, 0)):
96 + # Just below threshold
97 + result = get_sequence_parallelism_threshold(
98 + hidden_size=SP_MIN_HIDDEN_SIZE[90] - 1,
99 + tp_size=2,
100 + element_size=2,
101 + )
102 + assert result is None
103 +
104 + # Exactly at threshold
105 + result = get_sequence_parallelism_threshold(
106 + hidden_size=SP_MIN_HIDDEN_SIZE[90],
107 + tp_size=2,
108 + element_size=2,
109 + )
110 + assert result is not None
vllm/compilation/passes/fusion/sequence_parallelism.py
+102 -21
@@ -27,6 +27,63 @@ from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherRMSNo
27
28 logger = init_logger(__name__)
29
30 +# Min hidden size per device capability for sequence parallelism
31 +# Only apply sequence parallelism for models with hidden_size >= threshold
32 +SP_MIN_HIDDEN_SIZE: dict[int, int] = {
33 + 90: 8192, # H100: only for models with hidden_size >= 8192
34 +}
35 +
36 +# Min size per GPU per device capability for sequence parallelism
37 +# Total min size = min_per_gpu_size * tp_size
38 +# This ensures the threshold scales appropriately with tensor parallelism
39 +SP_MIN_PER_GPU_SIZE_MB: dict[int, float] = {
40 + 90: 8, # 8MB per GPU for H100
41 +}
42 +
43 +
44 +def get_sequence_parallelism_threshold(
45 + hidden_size: int,
46 + tp_size: int,
47 + element_size: int,
48 +) -> int | None:
49 + """
50 + Calculate the minimum token threshold for applying sequence parallelism.
51 +
52 + Returns None if sequence parallelism should not be applied based on model size.
53 +
54 + Branching logic based on device capability:
55 + - Check if hidden_size >= SP_MIN_HIDDEN_SIZE[device_capability]
56 + - If not, returns None (SP disabled for small models on this device)
57 + - If yes, calculates threshold based on per-GPU size
58 +
59 + Formula: min_token_num = (min_per_gpu_size_mb * tp_size * MiB) //
60 + (hidden_size * element_size)
61 + """
62 + from vllm.platforms import current_platform
63 +
64 + if not current_platform.is_cuda():
65 + return None
66 +
67 + capability = current_platform.get_device_capability()
68 + if capability is None:
69 + return None
70 + device_capability = capability.to_int()
71 +
72 + # Check if device has configured thresholds
73 + min_hidden_size = SP_MIN_HIDDEN_SIZE.get(device_capability)
74 + min_per_gpu_size_mb = SP_MIN_PER_GPU_SIZE_MB.get(device_capability)
75 +
76 + if min_hidden_size is None or min_per_gpu_size_mb is None:
77 + return None
78 +
79 + # Only apply sequence parallelism for models meeting the size threshold
80 + if hidden_size < min_hidden_size:
81 + return None
82 +
83 + MiB = 1024 * 1024
84 + min_size = min_per_gpu_size_mb * MiB * tp_size
85 + return int(min_size // (hidden_size * element_size))
86 +
87
88 def get_first_out_wrapper(
89 fn: Callable[..., Sequence[torch.Tensor]],
@@ -309,6 +366,23 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
366 def __init__(self, config: VllmConfig) -> None:
367 super().__init__(config)
368
369 + # Get min_token_num threshold
370 + # Read min_token_num from config (calculated during config init)
371 + self.min_token_num = None
372 + if config.model_config is not None:
373 + pass_config = config.compilation_config.pass_config
374 + self.min_token_num = pass_config.sp_min_token_num
375 +
376 + if self.min_token_num is not None:
377 + # Take the min to avoid exceeding max_num_batched_tokens
378 + max_batched = config.scheduler_config.max_num_batched_tokens
379 + if max_batched is not None:
380 + self.min_token_num = min(self.min_token_num, max_batched)
381 + logger.debug_once(
382 + f"Sequence parallelism min token threshold: {self.min_token_num}",
383 + scope="global",
384 + )
385 +
386 # Used to clean up redundant views created temporarily
387 # to circumvent residual shape change issues
388 self.noop_cleanup = NoOpEliminationPass(config)
@@ -339,29 +413,36 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
413 self.dump_patterns(config, self.patterns)
414
415 def is_applicable_for_range(self, compile_range: Range) -> bool:
342 - # When sequence parallelism is enabled, the residual tensor from RMSNorm
343 - # needs to be split along the sequence dimension. However, this dimension
344 - # is symbolic during piecewise compilation, and splitting symbolic shapes
345 - # is not supported.
346 - #
347 - # This pass is therefore only applied when the sequence dimension is
348 - # concrete:
349 - # 1. In full-graph compilation mode (no Dynamo splitting ops are used).
350 - # For this case we always pad num_tokens to be a multiple of
351 - # tensor_parallel_size, so there's no need to check shape % tp_size == 0.
352 - # 2. For specific shape provided during compilation (e.g., from
353 - # `compile_sizes`), which must be divisible by the tensor-parallel
354 - # size.
416 + """
417 + Determines if sequence parallelism should be applied for the given
418 + compile range.
419 +
420 + SP is only beneficial for larger batch sizes where the communication
421 + overhead is amortized. For small batches, the overhead of splitting
422 + and gathering tensors across TP ranks outweighs the benefits.
423 +
424 + Returns False (SP disabled) when:
425 + - Using piecewise compilation with non-concrete or TP-indivisible sizes
426 + - min_token_num is None (SP disabled for this device/config)
427 + - The compile range starts below the minimum token threshold
428 + """
429 + # For piecewise compilation (not using inductor graph partition),
430 + # we need concrete sizes that are divisible by TP for correct splitting
431 if (
356 - not self.compilation_config.splitting_ops
357 - or self.compilation_config.use_inductor_graph_partition
432 + not self.compilation_config.use_inductor_graph_partition
433 + and self.compilation_config.splitting_ops
434 ):
359 - return True
360 - tp_size = get_tensor_model_parallel_world_size()
361 - result: bool = (compile_range.is_single_size()) and (
362 - compile_range.end % tp_size == 0
363 - )
364 - return result
435 + tp_size = get_tensor_model_parallel_world_size()
436 + if not compile_range.is_single_size() or compile_range.end % tp_size != 0:
437 + return False
438 +
439 + # min_token_num is None when SP is disabled for this device/config
440 + # (e.g., non-CUDA platform, unsupported GPU, or small hidden_size)
441 + if self.min_token_num is None:
442 + return False
443 +
444 + # Only apply SP when batch size meets the minimum threshold
445 + return compile_range.start >= self.min_token_num
446
447 @VllmInductorPass.time_and_log
448 def __call__(self, graph: fx.Graph) -> None:
vllm/config/compilation.py
+8 -1
@@ -118,7 +118,9 @@ class PassConfig:
118 eliminate_noops: bool = Field(default=True)
119 """Eliminate no-op ops."""
120 enable_sp: bool = Field(default=None)
121 - """Enable sequence parallelism."""
121 + """Enable sequence parallelism. Requires TP>1. Automatically disabled
122 + if the model's hidden_size is too small for SP to be beneficial
123 + (threshold is device-capability dependent)."""
124 fuse_gemm_comms: bool = Field(default=None)
125 """Enable async TP."""
126 fuse_allreduce_rms: bool = Field(default=None)
@@ -155,6 +157,11 @@ class PassConfig:
157 8: 1, # 1MB
158 },
159 }, where key is the device capability"""
160 + sp_min_token_num: int | None = None
161 + """The minimum number of tokens above which vllm should use
162 + sequence parallelism. Specified as an integer token count.
163 + Unspecified will fallback to default values which are compute
164 + capability and world size dependent."""
165
166 # TODO(luka) better pass enabling system.
167
vllm/config/vllm.py
+56 -1
@@ -853,8 +853,33 @@ class VllmConfig:
853 logger.warning("Sequence Parallelism requires TP>1, disabling")
854 self.compilation_config.pass_config.enable_sp = False
855 self.compilation_config.pass_config.fuse_gemm_comms = False
856 + else:
857 + # Compute SP threshold early; disable if None (model too
858 + # small) before +rms_norm gets forced into custom_ops.
859 + pass_config = self.compilation_config.pass_config
860 + if pass_config.sp_min_token_num is None:
861 + from vllm.compilation.passes.fusion.sequence_parallelism import (
862 + get_sequence_parallelism_threshold,
863 + )
864 +
865 + tp_size = self.parallel_config.tensor_parallel_size
866 + hidden_size = self.model_config.get_hidden_size()
867 + element_size = self.model_config.dtype.itemsize
868 + pass_config.sp_min_token_num = get_sequence_parallelism_threshold(
869 + hidden_size, tp_size, element_size
870 + )
871
857 - elif "-rms_norm" in self.compilation_config.custom_ops:
872 + if pass_config.sp_min_token_num is None:
873 + logger.warning(
874 + "Model hidden_size too small for the SP "
875 + "threshold heuristic, disabling. To force SP, "
876 + "set pass_config.sp_min_token_num manually."
877 + )
878 + self.compilation_config.pass_config.enable_sp = False
879 + self.compilation_config.pass_config.fuse_gemm_comms = False
880 +
881 + if self.compilation_config.pass_config.enable_sp:
882 + if "-rms_norm" in self.compilation_config.custom_ops:
883 logger.warning(
884 "RMS norm force disabled, sequence parallelism might break"
885 )
@@ -1456,6 +1481,36 @@ class VllmConfig:
1481 "allreduce-rms fusion will be enabled for all num_tokens."
1482 )
1483
1484 + # Add the compile ranges for sequence parallelism
1485 + if compilation_config.pass_config.enable_sp:
1486 + pass_config = compilation_config.pass_config
1487 +
1488 + # Calculate min_token_num if not explicitly provided
1489 + # User override works regardless of hidden_size
1490 + if pass_config.sp_min_token_num is None:
1491 + from vllm.compilation.passes.fusion.sequence_parallelism import (
1492 + get_sequence_parallelism_threshold,
1493 + )
1494 +
1495 + tp_size = self.parallel_config.tensor_parallel_size
1496 + hidden_size = self.model_config.get_hidden_size()
1497 + element_size = self.model_config.dtype.itemsize
1498 + pass_config.sp_min_token_num = get_sequence_parallelism_threshold(
1499 + hidden_size, tp_size, element_size
1500 + )
1501 +
1502 + min_token_num = pass_config.sp_min_token_num
1503 + max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
1504 + if min_token_num is not None and (
1505 + max_num_batched_tokens is not None
1506 + and min_token_num < max_num_batched_tokens
1507 + and min_token_num > 1
1508 + ):
1509 + # Add split point at min_token_num - 1 to ensure SP applies
1510 + # starting from min_token_num
1511 + # This creates ranges: [1, min-1] (no SP), [min, max] (SP applies)
1512 + computed_compile_ranges_split_points.append(min_token_num - 1)
1513 +
1514 if compilation_config.pass_config.fuse_rope_kvcache:
1515 max_token_num = (
1516 compilation_config.pass_config.rope_kvcache_fusion_max_token_num