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 4, 12:30 UTC
Replace shape_invariants with simpler apprach in dynamic_arg_dims utilizing shape_id property. (#36194)
Signed-off-by: Laith Sakka <lsakka@meta.com>
Laith Sakka committed
Apr 29, 2026 at 11:32 UTC
6f20f81cbf1d12dc9d499d25ea0a64ef4c816c00
5 files changed
+88
-104
tests/compile/test_dynamic_shapes_compilation.py
+1
-3
@@ -55,12 +55,10 @@ def test_dynamic_shapes_compilation(
55
evaluate_guards,
56
):
57
"""Test that all dynamic shapes types compile successfully"""
58
- if use_bytecode_hook and shapes_type == DynamicShapesType.UNBACKED:
59
- pytest.skip("UNBACKED dynamic shapes require VLLM_USE_BYTECODE_HOOK=0")
60
-
58
if evaluate_guards and shapes_type == DynamicShapesType.UNBACKED:
59
pytest.skip("unbacked dynamic shapes do not add guards")
60
61
+ # TODO is this still a requirement?
62
if evaluate_guards and use_aot_compile:
63
pytest.skip("evaluate_guards requires use_aot_compile=0")
64
vllm/compilation/decorators.py
+73
-33
@@ -32,6 +32,9 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer
32
33
from .monitor import monitor_profiling_run, monitor_torch_compile
34
35
+# shape_id parameter was added to mark_unbacked in PyTorch 2.11.0
36
+_SUPPORTS_SHAPE_ID = is_torch_equal_or_newer("2.11.0")
37
+
38
if TYPE_CHECKING:
39
# Only added on nightly/2.10 so wrap
40
try:
@@ -89,7 +92,7 @@ def support_torch_compile(
92
@overload
93
def support_torch_compile(
94
*,
92
- dynamic_arg_dims: dict[str, int | list[int]] | None,
95
+ dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None,
96
) -> Callable[[type[_T]], type[_T]]: ...
97
98
@@ -103,7 +106,7 @@ def support_torch_compile(
106
@overload
107
def support_torch_compile(
108
*,
106
- dynamic_arg_dims: dict[str, int | list[int]] | None,
109
+ dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None,
110
mark_unbacked_dims: dict[str, int | list[int]] | None,
111
) -> Callable[[type[_T]], type[_T]]: ...
112
@@ -115,11 +118,10 @@ def support_torch_compile(cls: type[_T]) -> type[_T]: ...
118
def support_torch_compile(
119
cls: type[_T] | None = None,
120
*,
118
- dynamic_arg_dims: dict[str, int | list[int]] | None = None,
121
+ dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None = None,
122
mark_unbacked_dims: dict[str, int | list[int]] | None = None,
123
enable_if: Callable[[VllmConfig], bool] | None = None,
124
is_encoder: bool = False,
122
- shape_invariants: Callable[..., None] = lambda *args, **kwargs: None,
125
) -> Callable[[type[_T]], type[_T]] | type[_T]:
126
"""
127
A decorator to add support for compiling the forward method of a class.
@@ -141,8 +143,12 @@ def support_torch_compile(
143
```
144
145
`dynamic_arg_dims` is a dictionary that maps argument names to the dynamic
144
- dimensions of the argument. The dynamic dimensions can be either a single
145
- integer or a list of integers.
146
+ dimensions of the argument. The value can be:
147
+ - int: a single dimension index (e.g., 0)
148
+ - list[int]: multiple dimension indices (e.g., [0, 1])
149
+ - dict[int, str]: dimension to shape_id mapping for shape relations
150
+ (e.g., {0: "b"}). Dimensions with the same shape_id share the same
151
+ unbacked symbol.
152
153
if `dynamic_arg_dims` is `None`, it is inferred from the type annotation
154
of the `forward` method, based on the following default rules:
@@ -189,7 +195,7 @@ def support_torch_compile(
195
torch._check(input_ids.size()[0] == inputs_embeds.size()[0])
196
This enforces constraints on the symbolic shapes without hardcoding
197
specific values. It is needed for some models to avoid data dependent
192
- errors.
198
+ errors and maximize perf when unbacked shapes are used.
199
"""
200
201
def cls_decorator_helper(cls: type[_T]) -> type[_T]:
@@ -229,13 +235,13 @@ def support_torch_compile(
235
raise ValueError(
236
f"Argument {k} not found in the forward method of {cls}"
237
)
238
+
239
return _support_torch_compile(
240
cls,
241
inferred_dynamic_arg_dims,
242
mark_unbacked_dims,
243
enable_if,
244
is_encoder,
238
- shape_invariants,
245
)
246
247
if cls is not None:
@@ -324,15 +330,13 @@ def _try_load_aot_compiled_fn(
330
331
def _support_torch_compile(
332
cls: type[_T],
327
- dynamic_arg_dims: dict[str, int | list[int]],
333
+ dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]],
334
mark_unbacked_dims: dict[str, int | list[int]] | None = None,
335
enable_if: Callable[[VllmConfig], bool] | None = None,
336
is_encoder: bool = False,
331
- shape_invariants: Callable[..., None] = lambda *args, **kwargs: None,
337
) -> type[_T]:
333
- """
334
- A decorator to add support for compiling the forward method of a class.
335
- """
338
+ """Internal implementation of support_torch_compile decorator."""
339
+
340
if TorchCompileWithNoGuardsWrapper in cls.__bases__:
341
# support decorating multiple times
342
return cls
@@ -392,7 +396,8 @@ def _support_torch_compile(
396
if self.do_not_compile:
397
return
398
395
- self._check_shape_invariants = shape_invariants
399
+ self._dynamic_arg_dims = dynamic_arg_dims
400
+
401
self.was_aot_compile_fn_loaded_from_disk = False
402
compilation_counter.num_models_seen += 1
403
self.compiled = False
@@ -409,48 +414,83 @@ def _support_torch_compile(
414
def _mark_dynamic_inputs(
415
mod: type[_T], ds_type: DynamicShapesType, *args: Any, **kwargs: Any
416
) -> None:
412
- def mark_dynamic(arg: torch.Tensor, dims: list[int]) -> None:
417
+ def mark_dynamic(
418
+ arg: torch.Tensor, dim_shape_pairs: list[tuple[int, str | None]]
419
+ ) -> None:
420
if ds_type == DynamicShapesType.UNBACKED:
421
if is_torch_equal_or_newer("2.10.0"):
415
- for dim in dims:
416
- torch._dynamo.decorators.mark_unbacked(
417
- arg, dim, hint_override=arg.size()[dim]
418
- )
422
+ for dim, shape_id in dim_shape_pairs:
423
+ if shape_id is not None:
424
+ if not _SUPPORTS_SHAPE_ID:
425
+ raise RuntimeError(
426
+ f"shape_id='{shape_id}' requires PyTorch >= 2.11.0"
427
+ )
428
+ torch._dynamo.decorators.mark_unbacked(
429
+ arg,
430
+ dim,
431
+ hint_override=arg.size()[dim],
432
+ shape_id=shape_id,
433
+ )
434
+ else:
435
+ torch._dynamo.decorators.mark_unbacked(
436
+ arg,
437
+ dim,
438
+ hint_override=arg.size()[dim],
439
+ )
440
else:
441
+ # For older versions, we can't use hint_override or shape_id
442
+ dims = [dim for dim, _ in dim_shape_pairs]
443
torch._dynamo.decorators.mark_unbacked(arg, dims)
444
else:
445
+ dims = [dim for dim, _ in dim_shape_pairs]
446
torch._dynamo.mark_dynamic(arg, dims)
447
448
sig = inspect.signature(mod.__class__.forward) # type: ignore[attr-defined]
449
bound_args = sig.bind(mod, *args, **kwargs)
450
bound_args.apply_defaults()
427
- for k, dims in dynamic_arg_dims.items():
451
+
452
+ # Normalize dynamic_arg_dims to dict[str, dict[int, str | None]]
453
+ normalized_dims: dict[str, dict[int, str | None]] = {}
454
+ for k, v in dynamic_arg_dims.items():
455
+ if isinstance(v, dict):
456
+ normalized_dims[k] = {dim: shape_id for dim, shape_id in v.items()}
457
+ elif isinstance(v, int):
458
+ normalized_dims[k] = {v: None}
459
+ else:
460
+ normalized_dims[k] = {d: None for d in v}
461
+
462
+ for k, dim_to_shape_id in normalized_dims.items():
463
arg = bound_args.arguments.get(k)
464
465
if arg is not None:
431
- dims = [dims] if isinstance(dims, int) else dims
466
+ dims = list(dim_to_shape_id.keys())
467
+
468
if isinstance(arg, torch.Tensor):
433
- # In case dims is specified with negative indexing
434
- dims = [arg.ndim + dim if dim < 0 else dim for dim in dims]
435
- mark_dynamic(arg, dims)
469
+ dim_shape_pairs = [
470
+ (arg.ndim + d if d < 0 else d, dim_to_shape_id.get(d))
471
+ for d in dims
472
+ ]
473
+ mark_dynamic(arg, dim_shape_pairs)
474
elif isinstance(arg, IntermediateTensors):
475
for tensor in arg.tensors.values():
438
- # In case dims is specified with negative indexing
439
- dims = [tensor.ndim + dim if dim < 0 else dim for dim in dims]
440
- mark_dynamic(tensor, dims)
476
+ dim_shape_pairs = [
477
+ (tensor.ndim + d if d < 0 else d, dim_to_shape_id.get(d))
478
+ for d in dims
479
+ ]
480
+ mark_dynamic(tensor, dim_shape_pairs)
481
else:
482
raise ValueError(
443
- "Unsupported dynamic dimensions"
444
- f" {dims} for argument {k} with type {type(arg)}."
483
+ f"Unsupported dynamic dimensions {dims} "
484
+ f"for argument {k} with type {type(arg)}."
485
)
486
+
487
if mark_unbacked_dims:
447
- for k, dims in mark_unbacked_dims.items():
488
+ for k, dims_val in mark_unbacked_dims.items():
489
arg = bound_args.arguments.get(k)
490
if arg is not None:
450
- dims = [dims] if isinstance(dims, int) else dims
491
+ dims = [dims_val] if isinstance(dims_val, int) else list(dims_val)
492
if isinstance(arg, torch.Tensor):
452
- # In case dims is specified with negative indexing
453
- dims = [arg.ndim + dim if dim < 0 else dim for dim in dims]
493
+ dims = [arg.ndim + d if d < 0 else d for d in dims]
494
if is_torch_equal_or_newer("2.10.0"):
495
for dim in dims:
496
torch._dynamo.decorators.mark_unbacked(
vllm/compilation/wrapper.py
+3
-19
@@ -53,12 +53,6 @@ class TorchCompileWithNoGuardsWrapper:
53
since we drop all guards.
54
"""
55
56
- def check_invariants_and_forward(self, *args: Any, **kwargs: Any) -> Any:
57
- assert hasattr(self, "_check_shape_invariants")
58
- self._check_shape_invariants(*args, **kwargs)
59
-
60
- return self.forward(*args, **kwargs)
61
-
56
def _call_with_optional_nvtx_range(
57
self, callable_fn: Callable[P, R], *args: P.args, **kwargs: P.kwargs
58
) -> Any:
@@ -115,6 +109,9 @@ class TorchCompileWithNoGuardsWrapper:
109
"compilation_config.dynamic_shapes_config.evaluate_guards "
110
"requires VLLM_USE_BYTECODE_HOOK=0. "
111
)
112
+ assert ds_type != DynamicShapesType.UNBACKED, (
113
+ "UNBACKED dynamic shapes do not add guards"
114
+ )
115
116
options["guard_filter_fn"] = lambda x: [
117
entry.guard_type == "SHAPE_ENV" for entry in x
@@ -130,19 +127,6 @@ class TorchCompileWithNoGuardsWrapper:
127
compiled_ptr: Any = self.forward
128
# Validate that unbacked dynamic shapes require VLLM_USE_BYTECODE_HOOK=False
129
133
- if ds_type == DynamicShapesType.UNBACKED:
134
- # reason is that bytecode does torch._dynamo.eval_frame.
135
- # remove_from_cache(self.original_code_object()) to force a new
136
- # re-compilation. And if we use
137
- # compiled_ptr = self.check_invariants_and_forward
138
- # it will reset all entries.
139
- assert not envs.VLLM_USE_BYTECODE_HOOK, (
140
- "UNBACKED dynamic shapes requires VLLM_USE_BYTECODE_HOOK=0. "
141
- )
142
- assert not self.evaluate_guards, "UNBACKED dynamic shapes do not add guards"
143
-
144
- compiled_ptr = self.check_invariants_and_forward
145
-
130
# Apply the constrain_to_fx_strides patch before first compilation.
131
# This covers STOCK_TORCH_COMPILE and DYNAMO_ONCE paths. The VLLM
132
# compile paths call this from their own compile() methods too.
vllm/model_executor/models/llama.py
+6
-11
@@ -337,20 +337,15 @@ class LlamaDecoderLayer(nn.Module):
337
return vllm_config.quant_config
338
339
340
-def llama_model_invariants(
341
- input_ids, positions, intermediate_tensors=None, inputs_embeds=None
342
-):
343
- """Shape invariants for Llama model compilation, those are translated to
344
- runtime assertions for unbacked dynamic shapes and are compiled away for
345
- backed"""
346
- if input_ids is not None:
347
- torch._check(positions.size()[0] == input_ids.size()[0])
348
-
349
-
340
@support_torch_compile(
341
# TODO[#32068]: Investigate recompilation
342
# mark_unbacked_dims={"input_ids": 0},
353
- shape_invariants=llama_model_invariants
343
+ dynamic_arg_dims={
344
+ "input_ids": {0: "b"},
345
+ "positions": {0: "b"},
346
+ "intermediate_tensors": {0: "b"},
347
+ "inputs_embeds": {0: "b"},
348
+ },
349
)
350
class LlamaModel(nn.Module, EagleModelMixin):
351
def __init__(
vllm/model_executor/models/qwen2.py
+5
-38
@@ -312,48 +312,15 @@ class Qwen2DecoderLayer(nn.Module):
312
return hidden_states, residual
313
314
315
-def qwen_2_model_invariants(
316
- input_ids: torch.Tensor,
317
- positions: torch.Tensor,
318
- intermediate_tensors: IntermediateTensors | None = None,
319
- inputs_embeds: torch.Tensor | None = None,
320
-):
321
- """Shape invariants for Qwen2Model Model, those are translated to
322
- runtime assertions for unbacked dynamic shapes and are compiled away for
323
- backed"""
324
- # All these should be equal.
325
- # input_ids.size()[0]
326
- # positions.size()[-1]
327
- # intermediate_tensors["hidden_states"].size()[0]
328
- # inputs_embeds.size()[0]
329
- torch._check(input_ids.size()[0] == positions.size()[-1])
330
- if intermediate_tensors is not None:
331
- torch._check(
332
- input_ids.size()[0] == intermediate_tensors["hidden_states"].size()[0]
333
- )
334
-
335
- if inputs_embeds is not None:
336
- torch._check(input_ids.size()[0] == inputs_embeds.size()[0])
337
-
338
- # Hidden dimensions should match (hidden_size)
339
- # intermediate_tensors["hidden_states"].size()[1]
340
- # inputs_embeds.size()[1]
341
- if inputs_embeds is not None and intermediate_tensors is not None:
342
- torch._check(
343
- inputs_embeds.size()[1] == intermediate_tensors["hidden_states"].size()[1]
344
- )
345
-
346
-
315
@support_torch_compile(
316
dynamic_arg_dims={
349
- "input_ids": 0,
317
+ "input_ids": {0: "b"},
318
# positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl,
319
# otherwise (seq_len, ).
352
- "positions": -1,
353
- "intermediate_tensors": 0,
354
- "inputs_embeds": 0,
355
- },
356
- shape_invariants=qwen_2_model_invariants,
320
+ "positions": {-1: "b"},
321
+ "intermediate_tensors": {0: "b"},
322
+ "inputs_embeds": {0: "b"},
323
+ }
324
)
325
class Qwen2Model(nn.Module, EagleModelMixin):
326
def __init__(