@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7a228b5305
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 12, 16:00 UTC

Add option to use unbacked, and backed size obl dynamic shapes for more sounds compilation. (#26199)

Signed-off-by: Laith Sakka <lsakka@meta.com>

Laith Sakka committed Nov 24, 2025 at 07:12 UTC 7a228b5305f3257834c5079fd475f659bc4bf73d
8 files changed +442 -15
docs/design/debug_vllm_compile.md
+70
@@ -151,6 +151,76 @@ To avoid this, please either:
151 2. wrap the branching logic into a custom operator. TorchDynamo does not
152 trace into custom operators.
153
154 +## Debugging constraint violations and dynamic shapes guards issues
155 +
156 +Dynamic-shape guards are a specific category of Dynamo guards. They are constraints that `torch.compile`
157 +attaches to dynamic dimensions (e.g., `seq_len`) to ensure the compiled artifact remains valid.
158 +These guards typically appear when framework code, custom passes, or user code branches based on
159 +dynamic shape values.
160 +
161 +**Example:**
162 +
163 +```python
164 +if x > 10:
165 + # path A
166 +else:
167 + # path B
168 +```
169 +
170 +This creates a guard `x > 10` or `x <= 10` depending on which path was traced.
171 +
172 +**vLLM's Assumption:**
173 +vLLM assumes that all guards added by torch.compile are safe to drop and will not
174 +constrain the compiled graph to specific input shapes. When this assumption is violated,
175 +it can cause issues that users need to debug.
176 +Some side effects that indicates this assumption is violated are runtime errors
177 +or `ConstraintViolationErrors`.
178 +
179 +A `ConstraintViolationErrors` will be thrown if a dynamic shape gets constrained to
180 +a single value. If you encounter a constraint violation error or suspect that a dynamic
181 +shapes guard is being added incorrectly, you can use stricter dynamic shape modes to
182 +help debug the issue:
183 +
184 +```sh
185 +# Online - using unbacked mode
186 +vllm serve meta-llama/Llama-3.2-1B -O.dynamic_shapes_config.type=unbacked
187 +
188 +# Online - using backed_size_oblivious mode
189 +vllm serve meta-llama/Llama-3.2-1B -O.dynamic_shapes_config.type=backed_size_oblivious
190 +```
191 +
192 +```py
193 +# Offline - using unbacked mode
194 +from vllm.config.compilation import CompilationConfig, DynamicShapesConfig, DynamicShapesType
195 +LLM(model, compilation_config=CompilationConfig(
196 + dynamic_shapes_config=DynamicShapesConfig(type=DynamicShapesType.UNBACKED)
197 +))
198 +
199 +# Offline - using backed_size_oblivious mode
200 +from vllm.config.compilation import CompilationConfig, DynamicShapesConfig, DynamicShapesType
201 +LLM(model, compilation_config=CompilationConfig(
202 + dynamic_shapes_config=DynamicShapesConfig(type=DynamicShapesType.BACKED_SIZE_OBLIVIOUS)
203 +))
204 +```
205 +
206 +These modes are stricter and reduce or eliminate the need of dynamic shapes guarding, which can help isolate issues:
207 +
208 +- `unbacked`: Uses unbacked symints which don't allow guards, making it easier to identify where guards are being incorrectly added
209 +- `backed_size_oblivious`: Uses a mode that is more strict about guarding.
210 +
211 +For more details on dynamic shapes modes, see [Dynamic shapes and vLLM guard dropping](torch_compile.md#dynamic-shapes-and-vllm-guard-dropping).
212 +
213 +### Printing guards
214 +
215 +To see all guards that are being added during compilation, you can use `TORCH_LOGS=+dynamic`:
216 +
217 +```sh
218 +TORCH_LOGS=+dynamic vllm serve meta-llama/Llama-3.2-1B
219 +```
220 +
221 +Look for `[guard added]` in the logs to see where guards are being added. This can help you identify which operations are
222 +causing guards to be added incorrectly.
223 +
224 ## Debugging TorchInductor
225
226 TorchInductor takes a captured graph and then compiles it down to some Python code
docs/design/torch_compile.md
+104 -1
@@ -29,6 +29,109 @@ A unique aspect of vLLM's `torch.compile` integration, is that we guarantee all
29
30 By default, the cache saves compiled artifacts as binary files. If you would like to interact with the generated code for debugging purposes, set the field `compile_cache_save_format=unpacked` in the compilation config, or omit this and set the env variable `VLLM_COMPILE_CACHE_SAVE_FORMAT=unpacked`.
31
32 +## Dynamic shapes and vllm guard dropping
33 +
34 +`torch.compile` is designed to guard on dynamic shapes with no hesitation
35 +when needed. This contradicts with vLLM's `torch.compile` approach of
36 +dropping the guards since many of those guards could be material.
37 +
38 +`torch.compile` provides two kinds of dynamic shapes: `backed` and `unbacked`.
39 +`torch.compile` guards on `backed` dynamic shapes and does not provide a
40 +guarantee that no guards will be added to them. User code, dynamo,
41 +inductor, and autograd all can add guards. Moreover, for 0/1
42 +specializations, backed symbols are specialized unconditionally to 0, 1,
43 +or >=2 even without encountering a branching on those ranges.
44 +
45 +On the contrary, `unbacked` dynamic shapes are guaranteed not to be guarded
46 +on and are not 0/1 specialized. However, there is a possibility of
47 +throwing a data dependent error when a branch that requires their value is
48 +encountered and no explicit unbacked handling is defined. The framework is
49 +converging to a state where it won't throw DDE but rather pick general
50 +paths. One downside of using unbacked is missed optimization opportunities
51 +due to either perf bugs or picking general paths, also using a fixed
52 +non-example input-based hint (this will be fixed soon with override_hint
53 +API). An example of picking general paths is assuming input not contiguous
54 +in functions call contiguous() and reshape() when can't be symbolically proven
55 +with a change of introducing a clone.
56 +
57 +`backed_size_oblivious` is a flag that enables treating backed symbols as
58 +unbacked wherever explicit handling for unbacked is defined. With this
59 +mode, 0/1 specializations are mostly avoided in framework code and the
60 +default 0/1 specialization does not happen. However, there is still no
61 +guarantee that torch.compile won't guard, especially due to user code or
62 +custom passes. `backed_size_oblivious` is experimental in PyTorch compile
63 +and could be deprecated. That said, it's a safer option to use than
64 +`backed` and the probability of reducing performance is lower than
65 +`unbacked`.
66 +
67 +### Configuring Dynamic Shapes
68 +
69 +The `DynamicShapesConfig` allows you to control the dynamic shapes behavior by
70 +setting the `type` field. You can choose between three modes:
71 +`BACKED`(default), `UNBACKED` , and `BACKED_SIZE_OBLIVIOUS`.
72 +
73 +#### Offline Inference Example (Using LLM class)
74 +
75 +When using the `LLM` class for offline inference, you can configure dynamic
76 +shapes through the `compilation_config` parameter:
77 +
78 +```python
79 +from vllm import LLM, SamplingParams
80 +from vllm.config.compilation import CompilationConfig, DynamicShapesConfig, DynamicShapesType
81 +
82 +# Example: Using backed_size_oblivious (experimental, safer than backed)
83 +llm = LLM(
84 + model="meta-llama/Llama-3.2-1B",
85 + compilation_config=CompilationConfig(
86 + dynamic_shapes_config=DynamicShapesConfig(
87 + type=DynamicShapesType.BACKED_SIZE_OBLIVIOUS
88 + )
89 + )
90 +)
91 +
92 +# Example: Using unbacked (strongest guarantee against guards)
93 +llm = LLM(
94 + model="meta-llama/Llama-3.2-1B",
95 + compilation_config=CompilationConfig(
96 + dynamic_shapes_config=DynamicShapesConfig(
97 + type=DynamicShapesType.UNBACKED
98 + )
99 + )
100 +)
101 +
102 +# Generate outputs
103 +prompts = ["Hello, my name is", "The future of AI is"]
104 +sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
105 +outputs = llm.generate(prompts, sampling_params)
106 +```
107 +
108 +#### Online Serving Example (Using vllm serve)
109 +
110 +When using `vllm serve` for online serving, you can configure dynamic shapes
111 +through the `--compilation-config` flag:
112 +
113 +```bash
114 +# Example: Using unbacked
115 +vllm serve meta-llama/Llama-3.2-1B \
116 + --compilation-config '{"dynamic_shapes_config": {"type": "unbacked"}}'
117 +
118 +
119 +# Alternative: Using dot notation (simpler for single values)
120 +vllm serve meta-llama/Llama-3.2-1B -O.dynamic_shapes_config.type=unbacked
121 +```
122 +
123 +#### Choosing the Right Mode
124 +
125 +- **BACKED** (default): Use when you're willing to accept potential unsafe dropping of guards
126 +for maximal performance. Guard could be unsoundly added and then ignored.
127 +
128 +- **UNBACKED** Use when you need the strongest guarantee against guards.
129 + This is the most conservative option but may miss some optimization opportunities.
130 +
131 +- **BACKED_SIZE_OBLIVIOUS**: Use when you want a balance between avoiding guards
132 + and performance. This experimental mode is safer than BACKED but still not as
133 + conservative as UNBACKED.
134 +
135 ## Python Code Compilation
136
137 In the very verbose logs, we can see:
@@ -122,7 +225,7 @@ When all the shapes are known, `torch.compile` can compare different configs, an
225 triton_mm_4 0.0130 ms 100.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=16, BLOCK_N=32, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=5, num_warps=2
226 triton_mm_8 0.0134 ms 97.4% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=16, BLOCK_N=64, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=5, num_warps=4
227 triton_mm_12 0.0148 ms 87.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=16, BLOCK_N=128, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=4, num_warps=4
125 - mm 0.0160 ms 81.6%
228 + mm 0.0160 ms 81.6%
229 triton_mm_16 0.0165 ms 78.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=16, BLOCK_N=128, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=5, num_warps=8
230 triton_mm_3 0.0199 ms 65.4% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=16, BLOCK_N=32, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=5, num_warps=2
231 triton_mm_1 0.0203 ms 64.2% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=16, BLOCK_N=32, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=2, num_warps=2
tests/compile/test_dynamic_shapes_compilation.py new
+88
@@ -0,0 +1,88 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import gc
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm import LLM, SamplingParams
10 +from vllm.config.compilation import CompilationMode, DynamicShapesType
11 +from vllm.transformers_utils.tokenizer import get_tokenizer
12 +from vllm.utils.torch_utils import is_torch_equal_or_newer
13 +
14 +
15 +def get_test_models():
16 + """Get list of models to test based on PyTorch version"""
17 + # TODO "Qwen/Qwen3-4B-Instruct-2507" fails Fix issue and support it.
18 + return ["gpt2", "Qwen/Qwen2-7B-Instruct", "meta-llama/Llama-3.1-8B"]
19 +
20 +
21 +@pytest.mark.parametrize("model_name", get_test_models())
22 +@pytest.mark.parametrize(
23 + "shapes_type",
24 + [
25 + DynamicShapesType.BACKED,
26 + DynamicShapesType.UNBACKED,
27 + DynamicShapesType.BACKED_SIZE_OBLIVIOUS,
28 + ],
29 +)
30 +@pytest.mark.parametrize("use_aot_compile", ["0"])
31 +@pytest.mark.parametrize("use_bytecode_hook", [True, False])
32 +@pytest.mark.skipif(
33 + not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
34 +)
35 +def test_dynamic_shapes_compilation(
36 + monkeypatch, model_name, shapes_type, use_aot_compile, use_bytecode_hook
37 +):
38 + """Test that all dynamic shapes types compile successfully"""
39 + print(
40 + f"\nTesting model: {model_name} with {shapes_type.name}, "
41 + f"AOT compile: {use_aot_compile}, "
42 + f"Bytecode hook: {use_bytecode_hook}"
43 + )
44 + if use_bytecode_hook and shapes_type == DynamicShapesType.UNBACKED:
45 + pytest.skip("UNBACKED dynamic shapes require VLLM_USE_BYTECODE_HOOK=0")
46 +
47 + monkeypatch.setenv("VLLM_USE_AOT_COMPILE", use_aot_compile)
48 + monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0")
49 +
50 + prompt = "Hello, my name is"
51 +
52 + print(f"Testing {shapes_type.name} dynamic shapes...")
53 +
54 + # Initialize the model with specific dynamic shapes configuration
55 + model = LLM(
56 + model=model_name,
57 + compilation_config={
58 + "mode": CompilationMode.VLLM_COMPILE,
59 + "dynamic_shapes_config": {
60 + "type": shapes_type.value,
61 + },
62 + },
63 + )
64 +
65 + output = model.generate(prompt)
66 + result = output[0].outputs[0].text
67 + # Example of setting the sampling parameters
68 + tokenizer = get_tokenizer(model_name)
69 + yes_tokens = tokenizer.encode("yes", add_special_tokens=False)
70 + no_tokens = tokenizer.encode("no", add_special_tokens=False)
71 + allowed_ids = list(set(yes_tokens + no_tokens))
72 + sampling_params = SamplingParams(
73 + max_tokens=1, temperature=0, allowed_token_ids=allowed_ids
74 + )
75 +
76 + output = model.generate(
77 + "answer with yes or no is " + result + " rubbish for prompt " + prompt + "?",
78 + sampling_params=sampling_params,
79 + )
80 + result = output[0].outputs[0].text
81 + assert result == "yes"
82 +
83 + # Clean up GPU memory
84 + del model
85 + gc.collect()
86 + torch.cuda.empty_cache()
87 + torch.cuda.synchronize()
88 + print("GPU memory cleared")
vllm/compilation/decorators.py
+53 -10
@@ -24,6 +24,7 @@ from vllm.config import (
24 get_current_vllm_config,
25 set_current_vllm_config,
26 )
27 +from vllm.config.compilation import DynamicShapesType
28 from vllm.logger import init_logger
29 from vllm.sequence import IntermediateTensors
30 from vllm.utils.import_utils import resolve_obj_by_qualname
@@ -104,6 +105,7 @@ def support_torch_compile(
105 dynamic_arg_dims: dict[str, int | list[int]] | None = None,
106 mark_unbacked_dims: dict[str, int | list[int]] | None = None,
107 enable_if: Callable[[VllmConfig], bool] | None = None,
108 + shape_invariants: Callable[..., None] = lambda *args, **kwargs: None,
109 ) -> Callable[[_T], _T] | _T:
110 """
111 A decorator to add support for compiling the forward method of a class.
@@ -161,6 +163,14 @@ def support_torch_compile(
163 dim to be decorated with `mark_unbacked`. This is useful if we would like to
164 enforce that dynamo does not specialize on 0/1 values in the case of dummy input
165 such as for vision model compilation
166 +
167 + `shape_invariants` is a function that gets compiled right before forward.
168 + The function should have the torch._check calls that are needed to set
169 + the relationships between different input sizes. For example:
170 + torch._check(input_ids.size()[0] == inputs_embeds.size()[0])
171 + This enforces constraints on the symbolic shapes without hardcoding
172 + specific values. It is needed for some models to avoid data dependent
173 + errors.
174 """
175
176 def cls_decorator_helper(cls: _T) -> _T:
@@ -199,7 +209,11 @@ def support_torch_compile(
209 f"Argument {k} not found in the forward method of {cls}"
210 )
211 return _support_torch_compile(
202 - cls, inferred_dynamic_arg_dims, mark_unbacked_dims, enable_if
212 + cls,
213 + inferred_dynamic_arg_dims,
214 + mark_unbacked_dims,
215 + enable_if,
216 + shape_invariants,
217 )
218
219 if cls is not None:
@@ -242,6 +256,7 @@ def _support_torch_compile(
256 dynamic_arg_dims: dict[str, int | list[int]],
257 mark_unbacked_dims: dict[str, int | list[int]] | None = None,
258 enable_if: Callable[[VllmConfig], bool] | None = None,
259 + shape_invariants: Callable[..., None] = lambda *args, **kwargs: None,
260 ) -> _T:
261 """
262 A decorator to add support for compiling the forward method of a class.
@@ -276,11 +291,12 @@ def _support_torch_compile(
291 old_init(self, **kwargs)
292
293 self.vllm_config = vllm_config
294 + self.compilation_config = self.vllm_config.compilation_config
295 enable_compile = enable_if is None or enable_if(vllm_config)
296 # for CompilationMode.STOCK_TORCH_COMPILE , the upper level model runner
297 # will handle the compilation, so we don't need to do anything here.
298 self.do_not_compile = (
283 - vllm_config.compilation_config.mode
299 + self.compilation_config.mode
300 in [CompilationMode.NONE, CompilationMode.STOCK_TORCH_COMPILE]
301 or not supports_dynamo()
302 or _should_ignore_torch_compile(self.__class__)
@@ -289,29 +305,38 @@ def _support_torch_compile(
305 if self.do_not_compile:
306 return
307
308 + self._check_shape_invariants = shape_invariants
309 +
310 compilation_counter.num_models_seen += 1
311 self.compiled = False
312 TorchCompileWithNoGuardsWrapper.__init__(self)
313
314 cls.__init__ = __init__
315
298 - def _mark_dynamic_inputs(mod, *args, **kwargs):
316 + def _mark_dynamic_inputs(mod, type, *args, **kwargs):
317 + def mark_dynamic(arg, dims):
318 + if type == DynamicShapesType.UNBACKED:
319 + torch._dynamo.decorators.mark_unbacked(arg, dims)
320 + else:
321 + torch._dynamo.mark_dynamic(arg, dims)
322 +
323 sig = inspect.signature(mod.__class__.forward)
324 bound_args = sig.bind(mod, *args, **kwargs)
325 bound_args.apply_defaults()
326 for k, dims in dynamic_arg_dims.items():
327 arg = bound_args.arguments.get(k)
328 +
329 if arg is not None:
330 dims = [dims] if isinstance(dims, int) else dims
331 if isinstance(arg, torch.Tensor):
332 # In case dims is specified with negative indexing
333 dims = [arg.ndim + dim if dim < 0 else dim for dim in dims]
309 - torch._dynamo.mark_dynamic(arg, dims)
334 + mark_dynamic(arg, dims)
335 elif isinstance(arg, IntermediateTensors):
336 for tensor in arg.tensors.values():
337 # In case dims is specified with negative indexing
338 dims = [tensor.ndim + dim if dim < 0 else dim for dim in dims]
314 - torch._dynamo.mark_dynamic(tensor, dims)
339 + mark_dynamic(tensor, dims)
340 else:
341 raise ValueError(
342 "Unsupported dynamic dimensions"
@@ -338,6 +363,7 @@ def _support_torch_compile(
363 if getattr(self, "aot_compiled_fn", None) is not None:
364 return self.aot_compiled_fn(self, *args, **kwargs)
365
366 + ds_type = self.compilation_config.dynamic_shapes_config.type
367 cache_dir = None
368 aot_compilation_path = None
369 if envs.VLLM_USE_AOT_COMPILE:
@@ -352,6 +378,14 @@ def _support_torch_compile(
378 serialized backend artifacts), then we need to generate a new AOT
379 compile artifact from scratch.
380 """
381 + # Validate that AOT compile is not used with unbacked dynamic
382 + # shapes. aot_compile re-allocates backed symbols post dynamo!
383 + if ds_type == DynamicShapesType.UNBACKED:
384 + raise ValueError(
385 + "AOT compilation is not compatible with UNBACKED dynamic shapes. "
386 + "Please use BACKED or BACKED_SIZE_OBLIVIOUS dynamic shapes type "
387 + "when VLLM_USE_AOT_COMPILE is enabled."
388 + )
389 from .caching import compilation_config_hash_factors
390
391 factors: list[str] = compilation_config_hash_factors(self.vllm_config)
@@ -401,7 +435,12 @@ def _support_torch_compile(
435 # This is the path for the first compilation.
436
437 # the first compilation needs to have dynamic shapes marked
404 - _mark_dynamic_inputs(self, *args, **kwargs)
438 + _mark_dynamic_inputs(
439 + self,
440 + ds_type,
441 + *args,
442 + **kwargs,
443 + )
444
445 # here, it is the starting point of the `torch.compile` process
446 start_monitoring_torch_compile(self.vllm_config)
@@ -417,9 +456,7 @@ def _support_torch_compile(
456 # properly when any of these files change.
457
458 # 1. the file containing the top-level forward function
420 - self.vllm_config.compilation_config.traced_files.add(
421 - original_code_object.co_filename
422 - )
459 + self.compilation_config.traced_files.add(original_code_object.co_filename)
460
461 # 2. every time Dynamo sees a function call, it will inline
462 # the function by calling InliningInstructionTranslator.inline_call_
@@ -429,7 +466,7 @@ def _support_torch_compile(
466
467 def patched_inline_call(self_):
468 code = self_.f_code
432 - self.vllm_config.compilation_config.traced_files.add(code.co_filename)
469 + self.compilation_config.traced_files.add(code.co_filename)
470 return inline_call(self_)
471
472 # Disable the C++ compilation of symbolic shape guards. C++-fication
@@ -445,12 +482,18 @@ def _support_torch_compile(
482 # if the config doesn't exist
483 logger.debug("enable_cpp_symbolic_shape_guards config not available")
484
485 + # Prepare backed_size_oblivious config patch if needed
486 + fx_config_patches = {}
487 + if ds_type == DynamicShapesType.BACKED_SIZE_OBLIVIOUS:
488 + fx_config_patches["backed_size_oblivious"] = True
489 +
490 with (
491 patch.object(
492 InliningInstructionTranslator, "inline_call_", patched_inline_call
493 ),
494 torch._dynamo.config.patch(**dynamo_config_patches),
495 maybe_use_cudagraph_partition_wrapper(self.vllm_config),
496 + torch.fx.experimental._config.patch(**fx_config_patches),
497 _torch27_patch_tensor_subclasses(),
498 ):
499 if envs.VLLM_USE_AOT_COMPILE:
vllm/compilation/wrapper.py
+23 -1
@@ -6,6 +6,7 @@ import sys
6 from abc import abstractmethod
7 from contextlib import contextmanager
8 from types import CodeType
9 +from typing import Any
10
11 import torch
12 import torch._C._dynamo.guards
@@ -85,6 +86,12 @@ class TorchCompileWithNoGuardsWrapper:
86 since we drop all guards.
87 """
88
89 + def check_invariants_and_forward(self, *args, **kwargs):
90 + assert hasattr(self, "_check_shape_invariants")
91 + self._check_shape_invariants(*args, **kwargs)
92 +
93 + return self.forward(*args, **kwargs)
94 +
95 def __init__(self):
96 self.compiled = False
97
@@ -104,6 +111,21 @@ class TorchCompileWithNoGuardsWrapper:
111 # Drop all the guards.
112 options["guard_filter_fn"] = lambda x: [False for _ in x]
113
114 + # Validate that unbacked dynamic shapes require VLLM_USE_BYTECODE_HOOK=False
115 + from vllm.compilation.decorators import DynamicShapesType
116 +
117 + ds_type = vllm_config.compilation_config.dynamic_shapes_config.type
118 + compiled_ptr: Any = self.forward
119 + if ds_type == DynamicShapesType.UNBACKED:
120 + if envs.VLLM_USE_BYTECODE_HOOK:
121 + # reason is that bytecode does this hack torch._dynamo.eval_frame.
122 + # remove_from_cache(self.original_code_object()) to force a new
123 + # re-compilation.
124 + raise ValueError(
125 + "UNBACKED dynamic shapes require VLLM_USE_BYTECODE_HOOK=0. "
126 + )
127 + compiled_ptr = self.check_invariants_and_forward
128 +
129 if envs.VLLM_USE_AOT_COMPILE:
130 if hasattr(torch._dynamo.config, "enable_aot_compile"):
131 torch._dynamo.config.enable_aot_compile = True
@@ -114,7 +136,7 @@ class TorchCompileWithNoGuardsWrapper:
136 logger.warning(msg)
137
138 self._compiled_callable = torch.compile(
117 - self.forward,
139 + compiled_ptr,
140 fullgraph=True,
141 dynamic=False,
142 backend=backend,
vllm/config/compilation.py
+59 -1
@@ -192,6 +192,54 @@ class PassConfig:
192 self.enable_qk_norm_rope_fusion = False
193
194
195 +class DynamicShapesType(str, enum.Enum):
196 + """Types of dynamic shapes handling in torch.compile().
197 + see Dynamic shapes and vllm guard dropping in torch_compile.md
198 + for more details."""
199 +
200 + BACKED = "backed"
201 + """Use backed dynamic shapes. torch.compile() guards on backed dynamic
202 + shapes and may add guards. Symbols are specialized to 0, 1, or >=2 even
203 + without encountering branching on those ranges."""
204 +
205 + UNBACKED = "unbacked"
206 + """Use unbacked dynamic shapes. Guaranteed not to be guarded on and not
207 + 0/1 specialized, but may throw data dependent errors when branches require
208 + their value without explicit unbacked handling."""
209 +
210 + BACKED_SIZE_OBLIVIOUS = "backed_size_oblivious"
211 + """Experimental flag that treats backed symbols as unbacked when explicit
212 + unbacked handling is defined."""
213 +
214 +
215 +@config
216 +@dataclass
217 +class DynamicShapesConfig:
218 + """Configuration to control/debug torch compile dynamic shapes."""
219 +
220 + type: DynamicShapesType = DynamicShapesType.BACKED
221 + """Controls the type of dynamic shapes handling to use with torch.compile().
222 +
223 + - BACKED: Default PyTorch behavior with potential guards ignored.
224 + - UNBACKED: No guards guaranteed (most sound) but may throw
225 + data dependent errors.
226 + - BACKED_SIZE_OBLIVIOUS: Experimental safer alternative to
227 + backed/unbacked.
228 + """
229 +
230 + # TODO add a debug mode to fail
231 +
232 + def compute_hash(self) -> str:
233 + """
234 + Provide a hash for DynamicShapesConfig
235 + """
236 +
237 + from vllm.config.utils import get_hash_factors, hash_factors
238 +
239 + factors = get_hash_factors(self, {})
240 + return hash_factors(factors)
241 +
242 +
243 @config
244 @dataclass
245 class CompilationConfig:
@@ -322,7 +370,7 @@ class CompilationConfig:
370 If empty list [], no ops are excluded (suitable for full cudagraphs)."""
371 compile_mm_encoder: bool = False
372 """Whether or not to compile the multimodal encoder.
325 - Currently, this only works for `Qwen2_5_vl` on selected platforms.
373 + Currently, this only works for `Qwen2_5_vl` on selected platforms.
374 Disabled by default until more models are supported/tested to work."""
375
376 # Inductor capture
@@ -348,9 +396,11 @@ class CompilationConfig:
396 """Sizes to compile for inductor. In addition
397 to integers, it also supports "cudagraph_capture_sizes" to
398 specify the sizes for cudagraph capture."""
399 +
400 inductor_compile_config: dict = field(default_factory=dict)
401 """Additional configurations for inductor.
402 - None: use default configurations."""
403 +
404 inductor_passes: dict[str, str] = field(default_factory=dict)
405 """Additional passes for inductor. It is a dictionary
406 from pass name to pass function qualified name. We use function
@@ -460,8 +510,15 @@ class CompilationConfig:
510 max_num_seqs, and prevents capture of many large graphs (>512) that would
511 greatly increase startup time with limited performance benefit.
512 """
513 +
514 + dynamic_shapes_config: DynamicShapesConfig = field(
515 + default_factory=DynamicShapesConfig
516 + )
517 + """Configuration for dynamic shapes options"""
518 +
519 local_cache_dir: str = field(default=None, init=False) # type: ignore
520 """local cache dir for each rank"""
521 +
522 bs_to_padded_graph_size: list[int] = field(
523 default=None, # type: ignore
524 init=False,
@@ -530,6 +587,7 @@ class CompilationConfig:
587 from vllm.config.utils import get_hash_factors, hash_factors
588
589 factors = get_hash_factors(self, ignored_factors)
590 +
591 factors["pass_config"] = self.pass_config.compute_hash()
592 return hash_factors(factors)
593
vllm/model_executor/models/llama.py
+11 -1
@@ -354,7 +354,17 @@ class LlamaDecoderLayer(nn.Module):
354 return vllm_config.quant_config
355
356
357 -@support_torch_compile
357 +def llama_model_invariants(
358 + input_ids, positions, intermediate_tensors=None, inputs_embeds=None
359 +):
360 + """Shape invariants for Llama model compilation, those are translated to
361 + runtime assertions for unbacked dynamic shapes and are compiled away for
362 + backed"""
363 + if input_ids is not None:
364 + torch._check(positions.size()[0] == input_ids.size()[0])
365 +
366 +
367 +@support_torch_compile(shape_invariants=llama_model_invariants)
368 class LlamaModel(nn.Module):
369 def __init__(
370 self,
vllm/model_executor/models/qwen2.py
+34 -1
@@ -274,6 +274,38 @@ class Qwen2DecoderLayer(nn.Module):
274 return hidden_states, residual
275
276
277 +def qwen_2_model_invariants(
278 + input_ids: torch.Tensor,
279 + positions: torch.Tensor,
280 + intermediate_tensors: IntermediateTensors | None = None,
281 + inputs_embeds: torch.Tensor | None = None,
282 +):
283 + """Shape invariants for Qwen2Model Model, those are translated to
284 + runtime assertions for unbacked dynamic shapes and are compiled away for
285 + backed"""
286 + # All these should be equal.
287 + # input_ids.size()[0]
288 + # positions.size()[-1]
289 + # intermediate_tensors["hidden_states"].size()[0]
290 + # inputs_embeds.size()[0]
291 + torch._check(input_ids.size()[0] == positions.size()[-1])
292 + if intermediate_tensors is not None:
293 + torch._check(
294 + input_ids.size()[0] == intermediate_tensors["hidden_states"].size()[0]
295 + )
296 +
297 + if inputs_embeds is not None:
298 + torch._check(input_ids.size()[0] == inputs_embeds.size()[0])
299 +
300 + # Hidden dimensions should match (hidden_size)
301 + # intermediate_tensors["hidden_states"].size()[1]
302 + # inputs_embeds.size()[1]
303 + if inputs_embeds is not None and intermediate_tensors is not None:
304 + torch._check(
305 + inputs_embeds.size()[1] == intermediate_tensors["hidden_states"].size()[1]
306 + )
307 +
308 +
309 @support_torch_compile(
310 dynamic_arg_dims={
311 "input_ids": 0,
@@ -282,7 +314,8 @@ class Qwen2DecoderLayer(nn.Module):
314 "positions": -1,
315 "intermediate_tensors": 0,
316 "inputs_embeds": 0,
285 - }
317 + },
318 + shape_invariants=qwen_2_model_invariants,
319 )
320 class Qwen2Model(nn.Module):
321 def __init__(