@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4e26d3b09e
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 11, 23:45 UTC

[Compile] Conditional compilation. Introduce compile_ranges (#24252)

Signed-off-by: Luka Govedič <lgovedic@redhat.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Signed-off-by: ilmarkov <markovilya197@gmail.com> Signed-off-by: Luka Govedič <luka.govedic@gmail.com> Signed-off-by: ProExpertProg <lgovedic@redhat.com> Co-authored-by: Luka Govedič <lgovedic@redhat.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Luka Govedič <luka.govedic@gmail.com>

Ilya Markov committed Dec 5, 2025 at 19:17 UTC 4e26d3b09e9895279bf127f225c5f45922612261
15 files changed +583 -269
tests/compile/distributed/test_fusions_e2e.py
+19 -4
@@ -298,10 +298,14 @@ def test_tp2_attn_quant_allreduce_rmsnorm(
298 r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes",
299 log_holder.text,
300 )
301 - assert len(log_matches) == 2, log_holder.text
301 + # 2 for each compile range
302 + # (global compile range can be split due to fuse_allreduce_rmsnorm)
303 + num_compile_ranges = len(compilation_config.get_compile_ranges())
304 + assert num_compile_ranges in [1, 2]
305
303 - assert int(log_matches[0]) == matches.attention_fusion
304 - assert int(log_matches[1]) == matches.attention_fusion
306 + assert len(log_matches) == 2 * num_compile_ranges, log_holder.text
307 +
308 + assert all(int(log_match) == matches.attention_fusion for log_match in log_matches)
309
310 log_matches = re.findall(
311 r"collective_fusion.py:\d+] Replaced (\d+) patterns",
@@ -312,6 +316,12 @@ def test_tp2_attn_quant_allreduce_rmsnorm(
316 assert int(log_matches[0]) == matches.allreduce_fusion
317 assert int(log_matches[1]) == matches.allreduce_fusion
318
319 + log_matches = re.findall(
320 + r"pass_manager.py:\d+] Skipping .*AllReduceFusionPass.* with compile range",
321 + log_holder.text,
322 + )
323 + assert len(log_matches) == 2 * (num_compile_ranges - 1), log_holder.text
324 +
325
326 @multi_gpu_test(num_gpus=2)
327 @pytest.mark.parametrize(
@@ -446,7 +456,6 @@ def run_model(compile_config: int | CompilationConfig, model: str, **model_kwarg
456 # No cudagraphs by default
457 if compilation_config.cudagraph_mode is None:
458 compilation_config.cudagraph_mode = CUDAGraphMode.NONE
449 -
459 llm = LLM(
460 model=model,
461 compilation_config=compilation_config,
@@ -459,3 +468,9 @@ def run_model(compile_config: int | CompilationConfig, model: str, **model_kwarg
468 prompt = output.prompt
469 generated_text = output.outputs[0].text
470 print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
471 +
472 + # Get the compile ranges split points after vllm config post init
473 + # in order to compute compile ranges correctly
474 + compilation_config.compile_ranges_split_points = (
475 + llm.llm_engine.vllm_config.compilation_config.compile_ranges_split_points
476 + )
tests/compile/test_compile_ranges.py new
+168
@@ -0,0 +1,168 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any
4 +
5 +import torch
6 +from torch import fx as fx
7 +from torch import nn
8 +
9 +# This import automatically registers `torch.ops.silly.attention`
10 +import tests.compile.silly_attention # noqa
11 +from vllm.compilation.counter import compilation_counter
12 +from vllm.compilation.decorators import support_torch_compile
13 +from vllm.compilation.inductor_pass import (
14 + InductorPass,
15 + get_pass_context,
16 +)
17 +from vllm.config import (
18 + VllmConfig,
19 + set_current_vllm_config,
20 +)
21 +from vllm.config.compilation import CompilationConfig, CompilationMode
22 +from vllm.config.scheduler import SchedulerConfig
23 +from vllm.config.utils import Range
24 +from vllm.forward_context import set_forward_context
25 +
26 +BATCH_SIZE = 64
27 +MLP_SIZE = 128
28 +
29 +
30 +@support_torch_compile
31 +class TestModel(nn.Module):
32 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
33 + super().__init__()
34 +
35 + def forward(self, x: torch.Tensor) -> torch.Tensor:
36 + x = x + x
37 + attn_output = torch.empty_like(x)
38 + torch.ops.silly.attention(x, x, x, attn_output)
39 + x = attn_output
40 + x = x * 3
41 + return x
42 +
43 +
44 +@torch.inference_mode
45 +def run_model(vllm_config: VllmConfig, model: nn.Module, batch_sizes: list[int]):
46 + with set_forward_context({}, vllm_config=vllm_config):
47 + model(torch.randn(BATCH_SIZE, MLP_SIZE))
48 + for batch_size in batch_sizes:
49 + model(torch.randn(batch_size, MLP_SIZE))
50 +
51 +
52 +class PostGradRangeChecker(InductorPass):
53 + def __init__(self, ranges: list[Range]):
54 + self.ranges = ranges
55 + self.num_calls = 0
56 +
57 + def __call__(self, graph: fx.Graph):
58 + compile_range = get_pass_context().compile_range
59 + assert compile_range in self.ranges, (
60 + f"Compile range {compile_range} not in {self.ranges}"
61 + )
62 + self.num_calls += 1
63 +
64 + def uuid(self) -> str:
65 + state: dict[str, Any] = {}
66 + return InductorPass.hash_dict(state)
67 +
68 +
69 +def test_compile_ranges(use_fresh_inductor_cache):
70 + post_grad_range_checker = PostGradRangeChecker(
71 + [
72 + Range(start=1, end=8),
73 + Range(start=16, end=16),
74 + Range(start=9, end=32),
75 + Range(start=64, end=64),
76 + Range(start=33, end=8192),
77 + ]
78 + )
79 + torch.set_default_device("cuda")
80 + vllm_config = VllmConfig(
81 + scheduler_config=SchedulerConfig(
82 + max_num_batched_tokens=8192,
83 + ),
84 + compilation_config=CompilationConfig(
85 + mode=CompilationMode.VLLM_COMPILE,
86 + compile_ranges_split_points=[8, 32],
87 + compile_sizes=[16, 64, 128],
88 + inductor_compile_config={
89 + "post_grad_custom_post_pass": post_grad_range_checker,
90 + },
91 + ),
92 + )
93 +
94 + with set_current_vllm_config(vllm_config):
95 + model = TestModel(vllm_config=vllm_config, prefix="").eval()
96 + # Number of compilations: 3 for each compile range + 2 compile sizes
97 + batch_sizes = [1, 4, 16, 24, 48, 64, 8192]
98 +
99 + with compilation_counter.expect(
100 + num_graphs_seen=1,
101 + num_piecewise_graphs_seen=1,
102 + num_backend_compilations=5,
103 + ):
104 + run_model(vllm_config, model, batch_sizes)
105 + assert post_grad_range_checker.num_calls == 5
106 +
107 +
108 +def test_compile_config_get_compile_ranges():
109 + compilation_config = CompilationConfig(
110 + compile_ranges_split_points=[8, 32],
111 + )
112 + VllmConfig(
113 + scheduler_config=SchedulerConfig(
114 + max_num_batched_tokens=8192,
115 + ),
116 + compilation_config=compilation_config,
117 + )
118 + assert compilation_config.get_compile_ranges() == [
119 + Range(start=1, end=8),
120 + Range(start=9, end=32),
121 + Range(start=33, end=8192),
122 + ]
123 +
124 +
125 +def test_inductor_cache_compile_ranges(monkeypatch, use_fresh_inductor_cache):
126 + # To force multiple compilations, we disable the compile cache
127 + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
128 +
129 + post_grad_range_checker = PostGradRangeChecker(
130 + ranges=[
131 + Range(start=1, end=8),
132 + Range(start=9, end=8192),
133 + ]
134 + )
135 + scheduler_config = SchedulerConfig(
136 + max_num_batched_tokens=8192,
137 + )
138 + torch.set_default_device("cuda")
139 +
140 + def create_vllm_config():
141 + return VllmConfig(
142 + scheduler_config=scheduler_config,
143 + compilation_config=CompilationConfig(
144 + mode=CompilationMode.VLLM_COMPILE,
145 + compile_ranges_split_points=[8],
146 + inductor_compile_config={
147 + "post_grad_custom_post_pass": post_grad_range_checker,
148 + },
149 + ),
150 + )
151 +
152 + vllm_config_1 = create_vllm_config()
153 + with set_current_vllm_config(vllm_config_1):
154 + model1 = TestModel(vllm_config=vllm_config_1, prefix="").eval()
155 + batch_sizes = [1, 16]
156 + run_model(vllm_config_1, model1, batch_sizes)
157 + assert post_grad_range_checker.num_calls == 2
158 +
159 + post_grad_range_checker.num_calls = 0
160 + # Create a new vllm config with the new pass context
161 + vllm_config_2 = create_vllm_config()
162 + with set_current_vllm_config(vllm_config_2):
163 + model2 = TestModel(vllm_config=vllm_config_2, prefix="").eval()
164 + batch_sizes = [4, 32]
165 + run_model(vllm_config_2, model2, batch_sizes)
166 + # Check that cache is used, so the number of calls
167 + # should be 0
168 + assert post_grad_range_checker.num_calls == 0
tests/conftest.py
+14
@@ -67,6 +67,9 @@ from vllm.transformers_utils.utils import maybe_model_redirect
67 from vllm.utils.collection_utils import is_list_of
68 from vllm.utils.torch_utils import set_default_torch_num_threads
69
70 +from torch._inductor.utils import fresh_cache
71 +
72 +
73 if TYPE_CHECKING:
74 from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
75 from transformers.generation.utils import GenerateOutput
@@ -1465,3 +1468,14 @@ def clean_gpu_memory_between_tests():
1468 if torch.cuda.is_available():
1469 torch.cuda.empty_cache()
1470 gc.collect()
1471 +
1472 +
1473 +@pytest.fixture
1474 +def use_fresh_inductor_cache():
1475 + """
1476 + Use a fresh inductor cache for the test.
1477 + This is useful to ensure that the test is not affected by the
1478 + previous test calls.
1479 + """
1480 + with fresh_cache():
1481 + yield
vllm/compilation/backends.py
+52 -95
@@ -26,7 +26,7 @@ from vllm.compilation.partition_rules import (
26 should_split,
27 )
28 from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
29 -from vllm.config.utils import hash_factors
29 +from vllm.config.utils import Range, hash_factors
30 from vllm.logger import init_logger
31 from vllm.logging_utils import lazy
32 from vllm.platforms import current_platform
@@ -90,7 +90,7 @@ class CompilerManager:
90 """
91
92 def __init__(self, compilation_config: CompilationConfig):
93 - self.cache: dict[tuple[int | None, int, str], Any] = dict()
93 + self.cache: dict[tuple[Range, int, str], Any] = dict()
94 self.is_cache_updated = False
95 self.compilation_config = compilation_config
96 self.compiler = make_compiler(compilation_config)
@@ -99,11 +99,11 @@ class CompilerManager:
99 return self.compiler.compute_hash(vllm_config)
100
101 @contextmanager
102 - def compile_context(self, runtime_shape: int | None = None):
102 + def compile_context(self, compile_range: Range):
103 """Provide compilation context for the duration of compilation to set
104 any torch global properties we want to scope to a single Inductor
105 compilation (e.g. partition rules, pass context)."""
106 - with pass_context(runtime_shape):
106 + with pass_context(compile_range):
107 if self.compilation_config.use_inductor_graph_partition:
108 with inductor_partition_rule_context(
109 self.compilation_config.splitting_ops
@@ -159,29 +159,21 @@ class CompilerManager:
159 graph: fx.GraphModule,
160 example_inputs: list[Any],
161 graph_index: int,
162 - runtime_shape: int | None = None,
162 + compile_range: Range,
163 ) -> Callable | None:
164 - if (runtime_shape, graph_index, self.compiler.name) not in self.cache:
164 + if (compile_range, graph_index, self.compiler.name) not in self.cache:
165 return None
166 - handle = self.cache[(runtime_shape, graph_index, self.compiler.name)]
166 + handle = self.cache[(compile_range, graph_index, self.compiler.name)]
167 compiled_graph = self.compiler.load(
168 - handle, graph, example_inputs, graph_index, runtime_shape
168 + handle, graph, example_inputs, graph_index, compile_range
169 + )
170 + logger.debug(
171 + "Directly load the %s-th graph for compile range %sfrom %s via handle %s",
172 + graph_index,
173 + str(compile_range),
174 + self.compiler.name,
175 + handle,
176 )
170 - if runtime_shape is None:
171 - logger.debug(
172 - "Directly load the %s-th graph for dynamic shape from %s via handle %s",
173 - graph_index,
174 - self.compiler.name,
175 - handle,
176 - )
177 - else:
178 - logger.debug(
179 - "Directly load the %s-th graph for shape %s from %s via handle %s",
180 - graph_index,
181 - str(runtime_shape),
182 - self.compiler.name,
183 - handle,
184 - )
177 return compiled_graph
178
179 def compile(
@@ -190,9 +182,9 @@ class CompilerManager:
182 example_inputs,
183 additional_inductor_config,
184 compilation_config: CompilationConfig,
185 + compile_range: Range,
186 graph_index: int = 0,
187 num_graphs: int = 1,
195 - runtime_shape: int | None = None,
188 ) -> Any:
189 if graph_index == 0:
190 # before compiling the first graph, record the start time
@@ -204,7 +196,7 @@ class CompilerManager:
196 compiled_graph = None
197
198 # try to load from the cache
207 - compiled_graph = self.load(graph, example_inputs, graph_index, runtime_shape)
199 + compiled_graph = self.load(graph, example_inputs, graph_index, compile_range)
200 if compiled_graph is not None:
201 if graph_index == num_graphs - 1:
202 # after loading the last graph for this shape, record the time.
@@ -212,19 +204,12 @@ class CompilerManager:
204 now = time.time()
205 elapsed = now - compilation_start_time
206 compilation_config.compilation_time += elapsed
215 - if runtime_shape is None:
216 - logger.info(
217 - "Directly load the compiled graph(s) for dynamic shape "
218 - "from the cache, took %.3f s",
219 - elapsed,
220 - )
221 - else:
222 - logger.info(
223 - "Directly load the compiled graph(s) for shape %s "
224 - "from the cache, took %.3f s",
225 - str(runtime_shape),
226 - elapsed,
227 - )
207 + logger.info(
208 + "Directly load the compiled graph(s) for compile range %s "
209 + "from the cache, took %.3f s",
210 + str(compile_range),
211 + elapsed,
212 + )
213 return compiled_graph
214
215 # no compiler cached the graph, or the cache is disabled,
@@ -233,14 +218,15 @@ class CompilerManager:
218 # Let compile_fx generate a key for us
219 maybe_key = None
220 else:
236 - maybe_key = f"artifact_shape_{runtime_shape}_subgraph_{graph_index}"
237 -
238 - with self.compile_context(runtime_shape):
221 + maybe_key = "artifact_compile_range_"
222 + maybe_key += f"{compile_range.start}_{compile_range.end}"
223 + maybe_key += f"_subgraph_{graph_index}"
224 + with self.compile_context(compile_range):
225 compiled_graph, handle = self.compiler.compile(
226 graph,
227 example_inputs,
228 additional_inductor_config,
243 - runtime_shape,
229 + compile_range,
230 maybe_key,
231 )
232
@@ -248,55 +234,34 @@ class CompilerManager:
234
235 # store the artifact in the cache
236 if is_compile_cache_enabled(additional_inductor_config) and handle is not None:
251 - self.cache[(runtime_shape, graph_index, self.compiler.name)] = handle
237 + self.cache[(compile_range, graph_index, self.compiler.name)] = handle
238 compilation_counter.num_cache_entries_updated += 1
239 self.is_cache_updated = True
240 if graph_index == 0:
241 # adds some info logging for the first graph
256 - if runtime_shape is None:
257 - logger.info_once(
258 - "Cache the graph for dynamic shape for later use", scope="local"
259 - )
260 - else:
261 - logger.info_once(
262 - "Cache the graph of shape %s for later use",
263 - str(runtime_shape),
264 - scope="local",
265 - )
266 - if runtime_shape is None:
267 - logger.debug(
268 - "Store the %s-th graph for dynamic shape from %s via handle %s",
269 - graph_index,
270 - self.compiler.name,
271 - handle,
272 - )
273 - else:
274 - logger.debug(
275 - "Store the %s-th graph for shape %s from %s via handle %s",
276 - graph_index,
277 - str(runtime_shape),
278 - self.compiler.name,
279 - handle,
242 + logger.info_once(
243 + "Cache the graph of compile range %s for later use",
244 + str(compile_range),
245 )
246 + logger.debug(
247 + "Store the %s-th graph for compile range%s from %s via handle %s",
248 + graph_index,
249 + str(compile_range),
250 + self.compiler.name,
251 + handle,
252 + )
253
254 # after compiling the last graph, record the end time
255 if graph_index == num_graphs - 1:
256 now = time.time()
257 elapsed = now - compilation_start_time
258 compilation_config.compilation_time += elapsed
287 - if runtime_shape is None:
288 - logger.info_once(
289 - "Compiling a graph for dynamic shape takes %.2f s",
290 - elapsed,
291 - scope="local",
292 - )
293 - else:
294 - logger.info_once(
295 - "Compiling a graph for shape %s takes %.2f s",
296 - runtime_shape,
297 - elapsed,
298 - scope="local",
299 - )
259 + logger.info_once(
260 + "Compiling a graph for compile range %s takes %.2f s",
261 + str(compile_range),
262 + elapsed,
263 + scope="local",
264 + )
265
266 return compiled_graph
267
@@ -427,19 +392,7 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):
392 sym_shape_indices = [
393 i for i, x in enumerate(args) if isinstance(x, torch.SymInt)
394 ]
430 - global compilation_start_time
395
432 - compiled_graph_for_dynamic_shape = (
433 - self.vllm_backend.compiler_manager.compile(
434 - submod,
435 - args,
436 - self.vllm_backend.inductor_config,
437 - self.compilation_config,
438 - graph_index=index,
439 - num_graphs=len(self.compile_submod_names),
440 - runtime_shape=None,
441 - )
442 - )
396 # Lazy import here to avoid circular import
397 from .piecewise_backend import PiecewiseBackend
398
@@ -449,7 +402,6 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):
402 index,
403 len(self.compile_submod_names),
404 sym_shape_indices,
452 - compiled_graph_for_dynamic_shape,
405 self.vllm_backend,
406 )
407
@@ -589,8 +541,13 @@ class VllmBackend:
541 )
542 else:
543 # Config should automatically wrap all inductor passes
592 - assert isinstance(self.inductor_config[self.pass_key], InductorPass)
593 - self.pass_manager.add(self.inductor_config[self.pass_key])
544 + assert isinstance(
545 + self.compilation_config.inductor_compile_config[self.pass_key],
546 + InductorPass,
547 + )
548 + self.pass_manager.add(
549 + self.compilation_config.inductor_compile_config[self.pass_key]
550 + )
551 self.inductor_config[self.pass_key] = self.pass_manager
552
553 def __call__(
vllm/compilation/collective_fusion.py
+60 -86
@@ -10,6 +10,7 @@ from torch._inductor.pattern_matcher import PatternMatcherPass
10 from torch.distributed._symmetric_memory import enable_symm_mem_for_group
11
12 from vllm.config import VllmConfig
13 +from vllm.config.utils import Range
14 from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
15 from vllm.distributed.parallel_state import (
16 get_tensor_model_parallel_rank,
@@ -431,7 +432,7 @@ class AsyncTPPass(VllmPatternMatcherPass):
432
433 self.dump_patterns(config, self.patterns)
434
434 - def is_applicable(self, shape: int | None) -> bool:
435 + def is_applicable_for_range(self, compile_range: Range) -> bool:
436 # This pass is applied on top of the sequence parallelism pass.
437 # It inherits the same applicability condition as `SequenceParallelismPass`.
438 # See `SequenceParallelismPass.is_applicable` for more details.
@@ -441,7 +442,7 @@ class AsyncTPPass(VllmPatternMatcherPass):
442 ):
443 return True
444 tp_size = get_tensor_model_parallel_world_size()
444 - return shape is not None and shape % tp_size == 0
445 + return compile_range.is_single_size() and compile_range.end % tp_size == 0
446
447 @VllmInductorPass.time_and_log
448 def __call__(self, graph: fx.Graph):
@@ -505,91 +506,60 @@ if flashinfer_comm is not None:
506 num_tokens, hidden_size = allreduce_in.shape
507 element_size = allreduce_in.element_size()
508 current_tensor_size = num_tokens * hidden_size * element_size
509 + max_tensor_size = max_token_num * hidden_size * element_size
510 + assert current_tensor_size <= max_tensor_size, (
511 + f"Current tensor size {current_tensor_size} is larger than "
512 + f"max token num {max_token_num} * hidden size {hidden_size} * "
513 + f"element size {element_size}"
514 + )
515 + device_capability = current_platform.get_device_capability().to_int()
516 + # Get one shot input size limit for the current world size
517 + # for the current device capability
518 + max_one_shot_size = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get(
519 + device_capability, {}
520 + ).get(world_size, None)
521 + # Use one shot if no max size is specified
522 + use_oneshot = (
523 + max_one_shot_size is None or current_tensor_size <= max_one_shot_size * MiB
524 + )
525
509 - if num_tokens <= max_token_num:
510 - device_capability = current_platform.get_device_capability().to_int()
511 - # Get one shot input size limit for the current world size
512 - # for the current device capability
513 - max_one_shot_size_mb = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get(
514 - device_capability, {}
515 - ).get(world_size, None)
516 - # Use one shot if no max size for one shot is specified
517 - use_oneshot = (
518 - max_one_shot_size_mb is None
519 - or current_tensor_size <= max_one_shot_size_mb * MiB
520 - )
521 -
522 - assert _FI_WORKSPACE_TENSOR is not None, (
523 - "Flashinfer must be enabled when using flashinfer"
524 - )
525 - if norm_out is None:
526 - norm_out = allreduce_in
527 - residual_out = residual
528 - else:
529 - # return residual_out as allreduce_out with zeroed residual_in
530 - # as flashinfer does not support rms_norm
531 - # and allreduce_out together
532 - residual_out = allreduce_in
533 - # For the sizes that are smaller than the max size,
534 - # we only use flashinfer one shot allreduce
535 - flashinfer_comm.trtllm_allreduce_fusion(
536 - allreduce_in=allreduce_in,
537 - token_num=allreduce_in.shape[0],
538 - residual_in=residual,
539 - residual_out=residual_out,
540 - norm_out=norm_out,
541 - rms_gamma=rms_gamma,
542 - rms_eps=rms_eps,
543 - world_rank=world_rank,
544 - world_size=world_size,
545 - hidden_dim=allreduce_in.shape[-1],
546 - workspace_ptrs=_FI_WORKSPACE_TENSOR,
547 - launch_with_pdl=launch_with_pdl,
548 - use_oneshot=use_oneshot,
549 - trigger_completion_at_end=trigger_completion_at_end,
550 - fp32_acc=fp32_acc,
551 - pattern_code=pattern_code,
552 - allreduce_out=None,
553 - quant_out=quant_out,
554 - scale_out=scale_out,
555 - # in vllm we only support swizzled layout
556 - layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
557 - scale_factor=scale_factor,
558 - )
526 + assert _FI_WORKSPACE_TENSOR is not None, (
527 + "Flashinfer must be enabled when using flashinfer"
528 + )
529 + if norm_out is None:
530 + norm_out = allreduce_in
531 + residual_out = residual
532 else:
560 - allreduce_out = tensor_model_parallel_all_reduce(allreduce_in)
561 - if scale_factor is not None and scale_out is None:
562 - # Do fused rms norm static fp8 quant fused op
563 - if norm_out is None:
564 - torch.ops._C.fused_add_rms_norm_static_fp8_quant(
565 - quant_out,
566 - allreduce_out,
567 - residual,
568 - rms_gamma,
569 - scale_factor,
570 - rms_eps,
571 - )
572 - else:
573 - torch.ops._C.rms_norm_static_fp8_quant(
574 - quant_out, allreduce_out, rms_gamma, scale_factor, rms_eps
575 - )
576 - else:
577 - if norm_out is None:
578 - torch.ops._C.fused_add_rms_norm(
579 - allreduce_out, residual, rms_gamma, rms_eps
580 - )
581 - norm_out = allreduce_out
582 - else:
583 - torch.ops._C.rms_norm(norm_out, allreduce_out, rms_gamma, rms_eps)
584 - if scale_factor is not None and scale_out is not None:
585 - torch.ops._C.scaled_fp4_quant(
586 - quant_out, norm_out, scale_out, scale_factor
587 - )
588 - if scale_factor is None or norm_out is not None:
589 - # we need to return allreduce output
590 - # in cases of non quant fused AR + RMS norm
591 - # and fused AR + RMS norm + quant without fused add
592 - allreduce_in.copy_(allreduce_out)
533 + # return residual_out as allreduce_out with zeroed residual_in
534 + # as flashinfer does not support rms_norm
535 + # and allreduce_out together
536 + residual_out = allreduce_in
537 + # For the sizes that are smaller than the max size,
538 + # we only use flashinfer one shot allreduce
539 + flashinfer_comm.trtllm_allreduce_fusion(
540 + allreduce_in=allreduce_in,
541 + token_num=allreduce_in.shape[0],
542 + residual_in=residual,
543 + residual_out=residual_out,
544 + norm_out=norm_out,
545 + rms_gamma=rms_gamma,
546 + rms_eps=rms_eps,
547 + world_rank=world_rank,
548 + world_size=world_size,
549 + hidden_dim=allreduce_in.shape[-1],
550 + workspace_ptrs=_FI_WORKSPACE_TENSOR,
551 + launch_with_pdl=launch_with_pdl,
552 + use_oneshot=use_oneshot,
553 + trigger_completion_at_end=trigger_completion_at_end,
554 + fp32_acc=fp32_acc,
555 + pattern_code=pattern_code,
556 + allreduce_out=None,
557 + quant_out=quant_out,
558 + scale_out=scale_out,
559 + # in vllm we only support swizzled layout
560 + layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
561 + scale_factor=scale_factor,
562 + )
563
564 def call_trtllm_fused_allreduce_norm_fake(
565 allreduce_in: torch.Tensor,
@@ -1128,7 +1098,8 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
1098 if max_size is None:
1099 # Flashinfer doesn't support current world size
1100 logger.warning(
1131 - "Flashinfer allreduce fusion is not supported for world size %s",
1101 + "Flashinfer allreduce fusion is not supported for world size %s"
1102 + " or max size is not provided",
1103 self.tp_size,
1104 )
1105 return
@@ -1216,6 +1187,9 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
1187
1188 self.disabled = False
1189
1190 + def is_applicable_for_range(self, compile_range: Range) -> bool:
1191 + return compile_range.end <= self.max_token_num
1192 +
1193 @VllmInductorPass.time_and_log
1194 def __call__(self, graph: fx.Graph):
1195 if self.disabled:
vllm/compilation/compiler_interface.py
+20 -20
@@ -15,6 +15,7 @@ import torch.fx as fx
15 import vllm.envs as envs
16 from vllm.compilation.counter import compilation_counter
17 from vllm.config import VllmConfig
18 +from vllm.config.utils import Range
19 from vllm.utils.hashing import safe_hash
20 from vllm.utils.torch_utils import is_torch_equal_or_newer
21
@@ -63,16 +64,16 @@ class CompilerInterface:
64 graph: fx.GraphModule,
65 example_inputs: list[Any],
66 compiler_config: dict[str, Any],
66 - runtime_shape: int | None = None,
67 + compile_range: Range,
68 key: str | None = None,
69 ) -> tuple[Callable | None, Any | None]:
70 """
71 Compile the graph with the given example inputs and compiler config,
71 - with a runtime shape. If the `runtime_shape` is None, it means
72 - the `example_inputs` have a dynamic shape. Otherwise, the
73 - `runtime_shape` specifies the shape of the inputs. Right now we only
74 - support one variable shape for all inputs, which is the batchsize
75 - (number of tokens) during inference.
72 + with a range. The `compile_range` specifies the range of the inputs,
73 + it could be concrete size (if compile_sizes is provided), e.g. [4, 4]
74 + or a range [5, 8].
75 + Right now we only support one variable in ranges for all inputs,
76 + which is the batchsize (number of tokens) during inference.
77
78 Dynamo will make sure `graph(*example_inputs)` is valid.
79
@@ -98,7 +99,7 @@ class CompilerInterface:
99 graph: fx.GraphModule,
100 example_inputs: list[Any],
101 graph_index: int,
101 - runtime_shape: int | None = None,
102 + compile_range: Range,
103 ) -> Callable:
104 """
105 Load the compiled function from the handle.
@@ -212,20 +213,20 @@ class InductorStandaloneAdaptor(CompilerInterface):
213 graph: fx.GraphModule,
214 example_inputs: list[Any],
215 compiler_config: dict[str, Any],
215 - runtime_shape: int | None = None,
216 + compile_range: Range,
217 key: str | None = None,
218 ) -> tuple[Callable | None, Any | None]:
219 compilation_counter.num_inductor_compiles += 1
220 current_config = {}
221 if compiler_config is not None:
222 current_config.update(compiler_config)
222 - set_inductor_config(current_config, runtime_shape)
223 + set_inductor_config(current_config, compile_range)
224 set_functorch_config()
225
225 - if isinstance(runtime_shape, int):
226 + if compile_range.is_single_size():
227 dynamic_shapes = "from_example_inputs"
228 else:
228 - dynamic_shapes = "from_tracing_context"
229 + dynamic_shapes = "from_graph"
230
231 from torch._inductor import standalone_compile
232
@@ -235,7 +236,6 @@ class InductorStandaloneAdaptor(CompilerInterface):
236 dynamic_shapes=dynamic_shapes,
237 options={"config_patches": current_config},
238 )
238 -
239 # Save the compiled artifact to disk in the specified path
240 assert key is not None
241 path = os.path.join(self.cache_dir, key)
@@ -251,7 +251,7 @@ class InductorStandaloneAdaptor(CompilerInterface):
251 graph: fx.GraphModule,
252 example_inputs: list[Any],
253 graph_index: int,
254 - runtime_shape: int | None = None,
254 + compile_range: Range,
255 ) -> Callable:
256 assert isinstance(handle, tuple)
257 assert isinstance(handle[0], str)
@@ -315,7 +315,7 @@ class InductorAdaptor(CompilerInterface):
315 graph: fx.GraphModule,
316 example_inputs: list[Any],
317 compiler_config: dict[str, Any],
318 - runtime_shape: int | None = None,
318 + compile_range: Range,
319 key: str | None = None,
320 ) -> tuple[Callable | None, Any | None]:
321 compilation_counter.num_inductor_compiles += 1
@@ -329,7 +329,7 @@ class InductorAdaptor(CompilerInterface):
329 current_config["fx_graph_cache"] = True
330 current_config["fx_graph_remote_cache"] = False
331
332 - set_inductor_config(current_config, runtime_shape)
332 + set_inductor_config(current_config, compile_range)
333 set_functorch_config()
334
335 # inductor can inplace modify the graph, so we need to copy it
@@ -512,7 +512,7 @@ class InductorAdaptor(CompilerInterface):
512 graph: fx.GraphModule,
513 example_inputs: list[Any],
514 graph_index: int,
515 - runtime_shape: int | None = None,
515 + compile_range: Range,
516 ) -> Callable:
517 assert isinstance(handle, tuple)
518 assert isinstance(handle[0], str)
@@ -608,9 +608,9 @@ class InductorAdaptor(CompilerInterface):
608 return contextlib.nullcontext()
609
610
611 -def set_inductor_config(config, runtime_shape):
612 - if isinstance(runtime_shape, int):
613 - # for a specific batchsize, tuning triton kernel parameters
611 +def set_inductor_config(config, compile_range: Range):
612 + if compile_range.is_single_size():
613 + # for a specific batch size, tuning triton kernel parameters
614 # can be beneficial
615 config["max_autotune"] = envs.VLLM_ENABLE_INDUCTOR_MAX_AUTOTUNE
616 config["coordinate_descent_tuning"] = (
@@ -630,7 +630,7 @@ class EagerAdaptor(CompilerInterface):
630 graph: fx.GraphModule,
631 example_inputs: list[Any],
632 compiler_config: dict[str, Any],
633 - runtime_shape: int | None = None,
633 + compile_range: Range,
634 key: str | None = None,
635 ) -> tuple[Callable | None, Any | None]:
636 compilation_counter.num_eager_compiles += 1
vllm/compilation/inductor_pass.py
+6 -5
@@ -14,6 +14,7 @@ import torch
14 from torch import fx
15 from torch._subclasses.fake_tensor import FakeTensorMode, unset_fake_temporarily
16
17 +from vllm.config.utils import Range
18 from vllm.utils.torch_utils import is_torch_equal_or_newer
19
20 if is_torch_equal_or_newer("2.6"):
@@ -28,8 +29,8 @@ _pass_context = None
29
30
31 class PassContext:
31 - def __init__(self, runtime_shape: int | None):
32 - self.runtime_shape = runtime_shape
32 + def __init__(self, compile_range: Range):
33 + self.compile_range: Range = compile_range
34
35
36 def get_pass_context() -> PassContext:
@@ -39,13 +40,13 @@ def get_pass_context() -> PassContext:
40
41
42 @contextmanager
42 -def pass_context(runtime_shape: int | None):
43 +def pass_context(compile_range: Range):
44 """A context manager that stores the current pass context,
45 usually it is a list of sizes to specialize.
46 """
47 global _pass_context
48 prev_context = _pass_context
48 - _pass_context = PassContext(runtime_shape)
49 + _pass_context = PassContext(compile_range)
50 try:
51 yield
52 finally:
@@ -96,7 +97,7 @@ class InductorPass(CustomGraphPass):
97 encoded = json.dumps(dict_, sort_keys=True).encode("utf-8")
98 return hashlib.sha256(encoded).hexdigest()
99
99 - def is_applicable(self, shape: int | None):
100 + def is_applicable_for_range(self, compile_range: Range):
101 return True
102
103
vllm/compilation/pass_manager.py
+12 -4
@@ -24,7 +24,11 @@ if current_platform.is_cuda():
24 from .collective_fusion import AllReduceFusionPass, AsyncTPPass
25
26 from .fix_functionalization import FixFunctionalizationPass
27 -from .inductor_pass import CustomGraphPass, InductorPass, get_pass_context
27 +from .inductor_pass import (
28 + CustomGraphPass,
29 + InductorPass,
30 + get_pass_context,
31 +)
32 from .noop_elimination import NoOpEliminationPass
33
34 logger = init_logger(__name__)
@@ -70,13 +74,13 @@ class PostGradPassManager(CustomGraphPass):
74 def __call__(self, graph: fx.Graph):
75 VllmInductorPass.dump_prefix = 0 # reset dump index
76
73 - shape = get_pass_context().runtime_shape
77 + compile_range = get_pass_context().compile_range
78 for pass_ in self.passes:
75 - if pass_.is_applicable(shape):
79 + if pass_.is_applicable_for_range(compile_range):
80 pass_(graph)
81 VllmInductorPass.dump_prefix += 1
82 else:
79 - logger.debug("Skipping %s with shape %s", pass_, shape)
83 + logger.debug("Skipping %s with compile range %s", pass_, compile_range)
84
85 # post-cleanup goes before fix_functionalization
86 # because it requires a functional graph
@@ -133,4 +137,8 @@ class PostGradPassManager(CustomGraphPass):
137 state["passes"].append(pass_.uuid())
138 state["passes"].append(self.fix_functionalization.uuid())
139
140 + # Include the compile range in the uuid to ensure that inductor
141 + # recompiles the graph for the new dynamic compile range.
142 + state["compile_range"] = str(get_pass_context().compile_range)
143 +
144 return InductorPass.hash_dict(state)
vllm/compilation/piecewise_backend.py
+84 -41
@@ -7,18 +7,18 @@ from typing import Any
7
8 import torch.fx as fx
9
10 -import vllm.envs as envs
10 from vllm.compilation.backends import VllmBackend
11 from vllm.compilation.monitor import end_monitoring_torch_compile
12 from vllm.config import VllmConfig
13 +from vllm.config.compilation import Range
14 from vllm.logger import init_logger
15
16 logger = init_logger(__name__)
17
18
19 @dataclasses.dataclass
20 -class ConcreteSizeEntry:
21 - runtime_shape: int
20 +class RangeEntry:
21 + compile_range: Range
22 compiled: bool = False
23 runnable: Callable = None # type: ignore
24
@@ -31,7 +31,6 @@ class PiecewiseBackend:
31 piecewise_compile_index: int,
32 total_piecewise_compiles: int,
33 sym_shape_indices: list[int],
34 - compiled_graph_for_general_shape: Callable,
34 vllm_backend: VllmBackend,
35 ):
36 """
@@ -55,67 +54,111 @@ class PiecewiseBackend:
54
55 self.is_full_graph = total_piecewise_compiles == 1
56
58 - self.compile_sizes: set[int] = set(self.compilation_config.compile_sizes)
57 + self.compile_ranges = self.compilation_config.get_compile_ranges()
58 + log_string = f"PiecewiseBackend: compile_ranges: {self.compile_ranges}"
59 + logger.debug_once(log_string)
60
60 - self.first_run_finished = False
61 -
62 - self.compiled_graph_for_general_shape = compiled_graph_for_general_shape # noqa
61 + self.compile_sizes = self.compilation_config.compile_sizes
62 + log_string = f"PiecewiseBackend: compile_sizes: {self.compile_sizes}"
63 + logger.debug_once(log_string)
64
65 self.sym_shape_indices = sym_shape_indices
66
66 - self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG"
67 -
68 - # the entries for different shapes that we need to compile
69 - self.concrete_size_entries: dict[int, ConcreteSizeEntry] = {}
67 + # the entries for ranges that we need to either
68 + self.range_entries: dict[Range, RangeEntry] = {}
69
71 - # to_be_compiled_sizes tracks the remaining sizes to compile,
70 + # to_be_compiled_ranges tracks the remaining ranges to compile,
71 # and updates during the compilation process, so we need to copy it
73 - self.to_be_compiled_sizes: set[int] = self.compile_sizes.copy()
72 + self.to_be_compiled_ranges: set[Range] = set(self.compile_ranges)
73
74 # We only keep compilation management inside this class directly.
76 - for shape in self.compile_sizes:
77 - self.concrete_size_entries[shape] = ConcreteSizeEntry(
78 - runtime_shape=shape,
79 - runnable=self.compiled_graph_for_general_shape,
75 + for size in self.compile_sizes:
76 + range = Range(start=size, end=size)
77 + if range not in self.compile_ranges:
78 + self.range_entries[range] = RangeEntry(
79 + compile_range=range,
80 + )
81 + self.to_be_compiled_ranges.add(range)
82 +
83 + for range in self.compile_ranges:
84 + self.range_entries[range] = RangeEntry(
85 + compile_range=range,
86 )
87
88 def check_for_ending_compilation(self):
83 - if self.is_last_graph and not self.to_be_compiled_sizes:
89 + if self.is_last_graph and not self.to_be_compiled_ranges:
90 # no specific sizes to compile
91 # save the hash of the inductor graph for the next run
92 self.vllm_backend.compiler_manager.save_to_file()
93 end_monitoring_torch_compile(self.vllm_config)
94
89 - def __call__(self, *args) -> Any:
90 - if not self.first_run_finished:
91 - self.first_run_finished = True
92 - self.check_for_ending_compilation()
93 - return self.compiled_graph_for_general_shape(*args)
94 -
95 - runtime_shape = args[self.sym_shape_indices[0]]
96 -
97 - if runtime_shape not in self.concrete_size_entries:
98 - # we don't need to do anything for this shape
99 - return self.compiled_graph_for_general_shape(*args)
100 -
101 - entry = self.concrete_size_entries[runtime_shape]
95 + def _fakify_args(self, args: list[Any]) -> list[Any]:
96 + # We need to pass fake example_inputs, otherwise torch.compile
97 + # will fakify the example_inputs potentially causing some non dynamic
98 + # dimension to be be duck shaped to other existing shapes that have hints
99 + # matching their values.
100 + # This is problem because it can lead to unintended specializations!
101 + # if the new wrongly dynamic dim is specialized
102 + # it will force specializing the whole shape
103 + # torch.compile probably should not accept
104 + # non fake tensors as example inputs!
105 + # See issue https://github.com/vllm-project/vllm/issues/27899
106 + fake_example_inputs = []
107 + for node in self.graph.graph.nodes:
108 + # All place holders come first
109 + if node.op == "placeholder":
110 + fake_example_inputs.append(node.meta["example_value"])
111 + else:
112 + break
113 + assert len(fake_example_inputs) == len(args)
114 + return fake_example_inputs
115 +
116 + def _maybe_compile_for_range_entry(self, range_entry: RangeEntry, args) -> Any:
117 + if not range_entry.compiled:
118 + range_entry.compiled = True
119 + self.to_be_compiled_ranges.remove(range_entry.compile_range)
120
103 - if not entry.compiled:
104 - entry.compiled = True
105 - self.to_be_compiled_sizes.remove(runtime_shape)
121 # args are real arguments
107 - entry.runnable = self.vllm_backend.compiler_manager.compile(
122 + # fakify for range, real args for concrete size.
123 + # For concrete size, we clear the shape env in
124 + # compiler_manager.compile() so no need to fakify.
125 + args = (
126 + self._fakify_args(args)
127 + if not range_entry.compile_range.is_single_size()
128 + else args
129 + )
130 + range_entry.runnable = self.vllm_backend.compiler_manager.compile(
131 self.graph,
132 args,
133 self.vllm_backend.inductor_config,
134 self.compilation_config,
135 + compile_range=range_entry.compile_range,
136 graph_index=self.piecewise_compile_index,
137 num_graphs=self.total_piecewise_compiles,
114 - runtime_shape=runtime_shape,
138 )
139
117 - # finished compilations for all required shapes
118 - if self.is_last_graph and not self.to_be_compiled_sizes:
119 - self.check_for_ending_compilation()
140 + self.check_for_ending_compilation()
141 +
142 + def _find_range_for_shape(self, runtime_shape: int) -> Range | None:
143 + # First we try to find the range entry for the concrete compile size
144 + # If not found, we search for the range entry
145 + # that contains the runtime shape.
146 + if runtime_shape in self.compile_sizes:
147 + return self.range_entries[Range(start=runtime_shape, end=runtime_shape)]
148 + else:
149 + for range in self.compile_ranges:
150 + if runtime_shape in range:
151 + return self.range_entries[range]
152 + return None
153 +
154 + def __call__(self, *args) -> Any:
155 + runtime_shape = args[self.sym_shape_indices[0]]
156 + range_entry = self._find_range_for_shape(runtime_shape)
157 +
158 + assert range_entry is not None, (
159 + f"Shape out of considered range: {runtime_shape} "
160 + "[1, max_num_batched_tokens]"
161 + )
162
121 - return entry.runnable(*args)
163 + self._maybe_compile_for_range_entry(range_entry, args)
164 + return range_entry.runnable(*args)
vllm/compilation/sequence_parallelism.py
+3 -2
@@ -9,6 +9,7 @@ import torch.fx as fx
9 from torch._inductor.pattern_matcher import PatternMatcherPass
10
11 from vllm.config import VllmConfig
12 +from vllm.config.compilation import Range
13 from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
14 from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
15 from vllm.logger import init_logger
@@ -333,7 +334,7 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
334
335 self.dump_patterns(config, self.patterns)
336
336 - def is_applicable(self, shape: int | None) -> bool:
337 + def is_applicable_for_range(self, compile_range: Range) -> bool:
338 # When sequence parallelism is enabled, the residual tensor from RMSNorm
339 # needs to be split along the sequence dimension. However, this dimension
340 # is symbolic during piecewise compilation, and splitting symbolic shapes
@@ -353,7 +354,7 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
354 ):
355 return True
356 tp_size = get_tensor_model_parallel_world_size()
356 - return shape is not None and shape % tp_size == 0
357 + return (compile_range.is_single_size()) and (compile_range.end % tp_size == 0)
358
359 @VllmInductorPass.time_and_log
360 def __call__(self, graph: fx.Graph):
vllm/config/compilation.py
+37 -1
@@ -13,7 +13,13 @@ from pydantic.dataclasses import dataclass
13
14 import vllm.envs as envs
15 from vllm.compilation.inductor_pass import CallableInductorPass, InductorPass
16 -from vllm.config.utils import config, get_hash_factors, handle_deprecated, hash_factors
16 +from vllm.config.utils import (
17 + Range,
18 + config,
19 + get_hash_factors,
20 + handle_deprecated,
21 + hash_factors,
22 +)
23 from vllm.logger import init_logger
24 from vllm.platforms import current_platform
25 from vllm.utils.import_utils import resolve_obj_by_qualname
@@ -173,6 +179,9 @@ class PassConfig:
179 """
180
181 MiB = 1024 * 1024
182 + FI_SUPPORTED_WORLD_SIZES = [2, 4, 8]
183 + if world_size not in FI_SUPPORTED_WORLD_SIZES:
184 + return None
185 max_size_mb = self.fi_allreduce_fusion_max_size_mb
186 if max_size_mb is None:
187 max_size_mb = self.default_fi_allreduce_fusion_max_size_mb().get(world_size)
@@ -379,6 +388,8 @@ class CompilationConfig:
388 [vllm.config.CompilationConfig.cudagraph_copy_inputs]
389 - Inductor compilation:
390 - [`compile_sizes`][vllm.config.CompilationConfig.compile_sizes]
391 + - [`compile_ranges_split_points`]
392 + [vllm.config.CompilationConfig.compile_ranges_split_points]
393 - [`inductor_compile_config`]
394 [vllm.config.CompilationConfig.inductor_compile_config]
395 - [`inductor_passes`][vllm.config.CompilationConfig.inductor_passes]
@@ -492,6 +503,21 @@ class CompilationConfig:
503 to integers, it also supports "cudagraph_capture_sizes" to
504 specify the sizes for cudagraph capture."""
505
506 + compile_ranges_split_points: list[int] | None = None
507 + """Split points that represent compile ranges for inductor.
508 + The compile ranges are
509 + [1, split_points[0]],
510 + [split_points[0] + 1, split_points[1]], ...,
511 + [split_points[-1] + 1, max_num_batched_tokens].
512 + Compile sizes are also used single element ranges,
513 + the range is represented as [compile_sizes[i], compile_sizes[i]].
514 +
515 + If a range overlaps with the compile size, graph for compile size
516 + will be prioritized, i.e. if we have a range [1, 8] and a compile size 4,
517 + graph for compile size 4 will be compiled and used instead of the graph
518 + for range [1, 8].
519 + """
520 +
521 inductor_compile_config: dict = field(default_factory=dict)
522 """Additional configurations for inductor.
523 - None: use default configurations."""
@@ -1153,3 +1179,13 @@ class CompilationConfig:
1179 self.bs_to_padded_graph_size[bs] = start
1180 else:
1181 self.bs_to_padded_graph_size[bs] = end
1182 +
1183 + def get_compile_ranges(self) -> list[Range]:
1184 + """Get the compile ranges for the compilation config."""
1185 + if self.compile_ranges_split_points is None:
1186 + return []
1187 + split_points = sorted(set(self.compile_ranges_split_points))
1188 + return [
1189 + Range(start=s + 1, end=e)
1190 + for s, e in zip([0] + split_points[:-1], split_points)
1191 + ]
vllm/config/utils.py
+33 -1
@@ -10,7 +10,7 @@ import json
10 import pathlib
11 import textwrap
12 from collections.abc import Iterable, Mapping, Sequence, Set
13 -from dataclasses import MISSING, Field, field, fields, is_dataclass, replace
13 +from dataclasses import MISSING, Field, dataclass, field, fields, is_dataclass, replace
14 from itertools import pairwise
15 from typing import TYPE_CHECKING, Any, Protocol, TypeVar
16
@@ -322,3 +322,35 @@ def handle_deprecated(
322
323 for new_name in new_names:
324 setattr(config, new_name, old_val)
325 +
326 +
327 +@dataclass
328 +class Range:
329 + """
330 + A range of numbers.
331 + Inclusive of start, inclusive of end.
332 + """
333 +
334 + start: int
335 + end: int
336 +
337 + def is_single_size(self) -> bool:
338 + return self.start == self.end
339 +
340 + def __contains__(self, size: int) -> bool:
341 + # Inclusive of start, inclusive of end
342 + return self.start <= size <= self.end
343 +
344 + def __eq__(self, other: object) -> bool:
345 + if not isinstance(other, Range):
346 + return False
347 + return self.start == other.start and self.end == other.end
348 +
349 + def __hash__(self) -> int:
350 + return hash((self.start, self.end))
351 +
352 + def __str__(self) -> str:
353 + return f"({self.start}, {self.end})"
354 +
355 + def __repr__(self) -> str:
356 + return self.__str__()
vllm/config/vllm.py
+48
@@ -725,6 +725,8 @@ class VllmConfig:
725 "--kv-sharing-fast-prefill requires changes on model side for "
726 "correctness and to realize prefill savings. "
727 )
728 + # TODO: Move after https://github.com/vllm-project/vllm/pull/26847 lands
729 + self._set_compile_ranges()
730
731 if self.model_config and self.model_config.is_encoder_decoder:
732 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -1126,6 +1128,52 @@ class VllmConfig:
1128 # complete the remaining process.
1129 self.compilation_config.post_init_cudagraph_sizes()
1130
1131 + def _set_compile_ranges(self):
1132 + """
1133 + Set the compile ranges for the compilation config.
1134 + """
1135 + compilation_config = self.compilation_config
1136 + computed_compile_ranges_split_points = []
1137 +
1138 + # The upper bound of the compile ranges is the max_num_batched_tokens
1139 + max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
1140 + if max_num_batched_tokens is not None:
1141 + computed_compile_ranges_split_points.append(max_num_batched_tokens)
1142 +
1143 + # Add the compile ranges for flashinfer
1144 + if compilation_config.pass_config.fuse_allreduce_rms:
1145 + tp_size = self.parallel_config.tensor_parallel_size
1146 + max_size = compilation_config.pass_config.flashinfer_max_size(tp_size)
1147 + if max_size is not None:
1148 + max_token_num = max_size // (
1149 + self.model_config.get_hidden_size()
1150 + * self.model_config.dtype.itemsize
1151 + )
1152 + if (
1153 + max_num_batched_tokens is not None
1154 + and max_token_num < max_num_batched_tokens
1155 + ):
1156 + computed_compile_ranges_split_points.append(max_token_num)
1157 + else:
1158 + logger.debug(
1159 + "Max num batched tokens below allreduce-rms fusion threshold, "
1160 + "allreduce-rms fusion will be enabled for all num_tokens."
1161 + )
1162 +
1163 + if compilation_config.compile_ranges_split_points is not None:
1164 + for x in compilation_config.compile_ranges_split_points:
1165 + assert isinstance(x, int)
1166 + assert x > 0, f"Invalid compile range split point: {x}"
1167 + if (
1168 + max_num_batched_tokens is not None
1169 + and x < max_num_batched_tokens
1170 + and x > 1
1171 + ):
1172 + computed_compile_ranges_split_points.append(x)
1173 + compilation_config.compile_ranges_split_points = sorted(
1174 + computed_compile_ranges_split_points
1175 + )
1176 +
1177 def recalculate_max_model_len(self, max_model_len: int):
1178 # Can only be called in try_verify_and_update_config
1179 model_config = self.model_config
vllm/v1/worker/gpu_worker.py
+26 -9
@@ -15,6 +15,7 @@ import torch.nn as nn
15
16 import vllm.envs as envs
17 from vllm.config import CUDAGraphMode, VllmConfig
18 +from vllm.config.compilation import CompilationMode
19 from vllm.distributed import (
20 ensure_model_parallel_initialized,
21 init_distributed_environment,
@@ -407,15 +408,31 @@ class Worker(WorkerBase):
408 self.model_runner.initialize_kv_cache(kv_cache_config)
409
410 def compile_or_warm_up_model(self) -> None:
410 - # warm up sizes that are not in cudagraph capture sizes,
411 - # but users still want to compile for better performance,
412 - # e.g. for the max-num-batched token size in chunked prefill.
413 - compile_sizes = self.vllm_config.compilation_config.compile_sizes
414 - warmup_sizes = compile_sizes.copy() if compile_sizes is not None else []
415 - if not self.model_config.enforce_eager:
416 - capture_sizes = self.vllm_config.compilation_config.cudagraph_capture_sizes
417 - if capture_sizes is not None:
418 - warmup_sizes = [x for x in warmup_sizes if x not in capture_sizes]
411 + warmup_sizes = []
412 +
413 + if self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE:
414 + # warm up sizes that are not in cudagraph capture sizes,
415 + # but users still want to compile for better performance,
416 + # e.g. for the max-num-batched token size in chunked prefill.
417 + compile_sizes = self.vllm_config.compilation_config.compile_sizes
418 + warmup_sizes = compile_sizes.copy() if compile_sizes is not None else []
419 + cg_capture_sizes: list[int] = []
420 +
421 + if self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:
422 + cg_sizes = self.vllm_config.compilation_config.cudagraph_capture_sizes
423 + cg_capture_sizes = [] if cg_sizes is None else cg_sizes
424 + warmup_sizes = [x for x in warmup_sizes if x not in cg_capture_sizes]
425 +
426 + compile_ranges = self.vllm_config.compilation_config.get_compile_ranges()
427 + # For each compile_range, if none of the batch sizes
428 + # in warmup_sizes or cudagraph_capture_sizes are in the range,
429 + # add the end of the range to ensure compilation/warmup.
430 + all_sizes = set(cg_capture_sizes)
431 + all_sizes.update([x for x in warmup_sizes if isinstance(x, int)])
432 + for compile_range in compile_ranges:
433 + if not any(x in compile_range for x in all_sizes):
434 + warmup_sizes.append(compile_range.end)
435 +
436 # We skip EPLB here since we don't want to record dummy metrics
437 for size in sorted(warmup_sizes, reverse=True):
438 logger.info("Compile and warming up model for size %d", size)
vllm/v1/worker/utils.py
+1 -1
@@ -337,7 +337,7 @@ def is_residual_scattered_for_sp(
337 The residual tensor is scattered across tensor parallel ranks when sequence
338 parallelism and tensor parallelism is enabled.
339
340 - This follows the same logic as SequenceParallelismPass.is_applicable():
340 + This follows the same logic as SequenceParallelismPass.is_applicable_for_range():
341 - In full-graph compilation mode (no splitting ops or using inductor graph
342 partition), SP is always applied
343 - Otherwise, SP is only applied for specific shapes in compile_sizes