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 5, 09:45 UTC
[torch.compile] Allow usage of Opaque Objects in PyTorch 2.11 (#39286)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Apr 9, 2026 at 01:21 UTC
ba4a78eb5d2ea30477b58a0bb8109b129f35c8b1
5 files changed
+69
-45
vllm/compilation/compiler_interface.py
+4
-44
@@ -16,6 +16,7 @@ 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.env_override import _apply_constrain_to_fx_strides_patch
20
from vllm.logger import init_logger
21
from vllm.utils.hashing import safe_hash
22
from vllm.utils.torch_utils import is_torch_equal_or_newer
@@ -225,48 +226,6 @@ def _patch_standalone_compile_atomic_save() -> None:
226
logger.debug("Patched %s.save for atomic writes (torch < 2.10)", cls.__name__)
227
228
228
-def _patch_constrain_to_fx_strides() -> contextlib.AbstractContextManager:
229
- """Context manager that patches inductor's ``constrain_to_fx_strides``
230
- to handle opaque (non-tensor) arguments.
231
-
232
- The original calls ``.stride()`` on every FX arg's meta value, which
233
- crashes on ``FakeScriptObject`` (the compile-time proxy for hoisted
234
- opaque types). The patched version skips args whose meta value is
235
- not a ``torch.Tensor``.
236
-
237
- Returns ``nullcontext`` on torch < 2.11.
238
- Upstream issue: https://github.com/pytorch/pytorch/issues/175973
239
- """
240
- if not is_torch_equal_or_newer("2.11.0.dev"):
241
- return contextlib.nullcontext()
242
-
243
- import torch._inductor.ir as _ir
244
- import torch._inductor.lowering as _lowering
245
- from torch._inductor.virtualized import V as _V
246
-
247
- def _patched(fx_node, *args, **kwargs):
248
- def apply_constraint(arg, fx_arg):
249
- if isinstance(arg, _ir.IRNode):
250
- meta_val = fx_arg.meta.get("val")
251
- if isinstance(meta_val, torch.Tensor):
252
- stride_order = _ir.get_stride_order(
253
- meta_val.stride(), _V.graph.sizevars.shape_env
254
- )
255
- return _ir.ExternKernel.require_stride_order(arg, stride_order)
256
- return arg
257
- if isinstance(arg, dict):
258
- return {key: apply_constraint(arg[key], fx_arg[key]) for key in arg}
259
- return arg
260
-
261
- args = tuple(
262
- apply_constraint(arg, fx_arg) for arg, fx_arg in zip(args, fx_node.args)
263
- )
264
- kwargs = {k: apply_constraint(v, fx_node.kwargs[k]) for k, v in kwargs.items()}
265
- return args, kwargs
266
-
267
- return patch.object(_lowering, "constrain_to_fx_strides", _patched)
268
-
269
-
229
class InductorStandaloneAdaptor(CompilerInterface):
230
"""
231
The adaptor for the Inductor compiler.
@@ -304,6 +263,7 @@ class InductorStandaloneAdaptor(CompilerInterface):
263
compile_range: Range,
264
key: str | None = None,
265
) -> tuple[Callable[..., Any] | None, Any | None]:
266
+ _apply_constrain_to_fx_strides_patch()
267
compilation_counter.num_inductor_compiles += 1
268
current_config = {}
269
if compiler_config is not None:
@@ -387,7 +347,7 @@ class InductorStandaloneAdaptor(CompilerInterface):
347
else:
348
fake_mode_ctx = contextlib.nullcontext()
349
390
- with pregrad_ctx, fake_mode_ctx, _patch_constrain_to_fx_strides():
350
+ with pregrad_ctx, fake_mode_ctx:
351
compiled_graph = standalone_compile(graph, example_inputs, **compile_kwargs)
352
353
if use_aot:
@@ -502,6 +462,7 @@ class InductorAdaptor(CompilerInterface):
462
compile_range: Range,
463
key: str | None = None,
464
) -> tuple[Callable[..., Any] | None, Any | None]:
465
+ _apply_constrain_to_fx_strides_patch()
466
compilation_counter.num_inductor_compiles += 1
467
from torch._inductor.compile_fx import compile_fx
468
@@ -630,7 +591,6 @@ class InductorAdaptor(CompilerInterface):
591
stack.enter_context(
592
torch._functorch.config.patch(enable_remote_autograd_cache=False)
593
)
633
- stack.enter_context(_patch_constrain_to_fx_strides())
594
595
# Clear the tracing context before calling compile_fx.
596
# vLLM calls compile_fx from within a PiecewiseCompileInterpreter
vllm/compilation/wrapper.py
+7
@@ -143,6 +143,13 @@ class TorchCompileWithNoGuardsWrapper:
143
144
compiled_ptr = self.check_invariants_and_forward
145
146
+ # Apply the constrain_to_fx_strides patch before first compilation.
147
+ # This covers STOCK_TORCH_COMPILE and DYNAMO_ONCE paths. The VLLM
148
+ # compile paths call this from their own compile() methods too.
149
+ from vllm.env_override import _apply_constrain_to_fx_strides_patch
150
+
151
+ _apply_constrain_to_fx_strides_patch()
152
+
153
aot_context = nullcontext()
154
if envs.VLLM_USE_AOT_COMPILE:
155
if hasattr(torch._dynamo.config, "enable_aot_compile"):
vllm/env_override.py
+54
@@ -500,6 +500,60 @@ if is_torch_equal("2.9.0"):
500
# This mirrors the fix in https://github.com/pytorch/pytorch/pull/177558
501
# and can be removed once torch >=2.12 is the minimum supported version.
502
503
+# ===================================================
504
+# torch >= 2.11 Inductor constrain_to_fx_strides monkeypatch
505
+# ===================================================
506
+# Inductor's constrain_to_fx_strides calls .stride() on every FX arg's meta
507
+# value, which crashes on FakeScriptObject (the compile-time proxy for hoisted
508
+# opaque types). The patched version skips args whose meta value is not a
509
+# torch.Tensor.
510
+# Upstream issue: https://github.com/pytorch/pytorch/issues/175973
511
+
512
+
513
+_constrain_to_fx_strides_patched = False
514
+
515
+
516
+def _apply_constrain_to_fx_strides_patch():
517
+ """Patch lowering.constrain_to_fx_strides globally. Safe to call
518
+ multiple times; only the first call does anything.
519
+ Only applies for torch >= 2.11 and < 2.12."""
520
+ global _constrain_to_fx_strides_patched
521
+ if _constrain_to_fx_strides_patched:
522
+ return
523
+ _constrain_to_fx_strides_patched = True
524
+
525
+ if not is_torch_equal_or_newer("2.11.0.dev") or is_torch_equal_or_newer(
526
+ "2.12.0.dev"
527
+ ):
528
+ return
529
+
530
+ import torch._inductor.ir as _ir
531
+ import torch._inductor.lowering as _lowering
532
+ from torch._inductor.virtualized import V as _V
533
+
534
+ def _patched(fx_node, *args, **kwargs):
535
+ def apply_constraint(arg, fx_arg):
536
+ if isinstance(arg, _ir.IRNode):
537
+ meta_val = fx_arg.meta.get("val")
538
+ if isinstance(meta_val, torch.Tensor):
539
+ stride_order = _ir.get_stride_order(
540
+ meta_val.stride(), _V.graph.sizevars.shape_env
541
+ )
542
+ return _ir.ExternKernel.require_stride_order(arg, stride_order)
543
+ return arg
544
+ if isinstance(arg, dict):
545
+ return {key: apply_constraint(arg[key], fx_arg[key]) for key in arg}
546
+ return arg
547
+
548
+ args = tuple(
549
+ apply_constraint(arg, fx_arg) for arg, fx_arg in zip(args, fx_node.args)
550
+ )
551
+ kwargs = {k: apply_constraint(v, fx_node.kwargs[k]) for k, v in kwargs.items()}
552
+ return args, kwargs
553
+
554
+ _lowering.constrain_to_fx_strides = _patched
555
+
556
+
557
if is_torch_equal_or_newer("2.10.0") and not is_torch_equal_or_newer("2.12.0"):
558
import builtins as _builtins
559
import pickle
vllm/utils/torch_utils.py
+1
-1
@@ -706,7 +706,7 @@ def is_torch_equal(target: str) -> bool:
706
return Version(importlib.metadata.version("torch")) == Version(target)
707
708
709
-HAS_OPAQUE_TYPE = is_torch_equal_or_newer("2.12.0.dev")
709
+HAS_OPAQUE_TYPE = is_torch_equal_or_newer("2.11.0.dev")
710
711
if HAS_OPAQUE_TYPE:
712
from torch._opaque_base import OpaqueBase
vllm/v1/worker/gpu_model_runner.py
+3
@@ -4857,6 +4857,9 @@ class GPUModelRunner(
4857
self.vllm_config.compilation_config.mode
4858
== CompilationMode.STOCK_TORCH_COMPILE
4859
):
4860
+ from vllm.env_override import _apply_constrain_to_fx_strides_patch
4861
+
4862
+ _apply_constrain_to_fx_strides_patch()
4863
backend = self.vllm_config.compilation_config.init_backend(self.vllm_config)
4864
compilation_counter.stock_torch_compile_count += 1
4865
self.model.compile(fullgraph=True, backend=backend)