A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 02:45 UTC
[compile] Stop unconditionally patching constrain_to_fx_strides (#36152)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Mar 6, 2026 at 10:17 UTC
54756b61091e3c913436ddd00b9d99e11e7c9a8c
2 files changed
+44
-42
vllm/compilation/compiler_interface.py
+44
-1
@@ -225,6 +225,48 @@ def _patch_standalone_compile_atomic_save() -> None:
225
logger.debug("Patched %s.save for atomic writes (torch < 2.10)", cls.__name__)
226
227
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
+
270
class InductorStandaloneAdaptor(CompilerInterface):
271
"""
272
The adaptor for the Inductor compiler.
@@ -312,7 +354,7 @@ class InductorStandaloneAdaptor(CompilerInterface):
354
"torch._inductor.compile_fx._recursive_pre_grad_passes",
355
lambda gm, _: gm,
356
)
315
- with ctx:
357
+ with ctx, _patch_constrain_to_fx_strides():
358
compiled_graph = standalone_compile(graph, example_inputs, **compile_kwargs)
359
360
if use_aot:
@@ -555,6 +597,7 @@ class InductorAdaptor(CompilerInterface):
597
stack.enter_context(
598
torch._functorch.config.patch(enable_remote_autograd_cache=False)
599
)
600
+ stack.enter_context(_patch_constrain_to_fx_strides())
601
602
compiled_graph = compile_fx(
603
graph,
vllm/env_override.py
-41
@@ -482,44 +482,3 @@ if is_torch_equal("2.9.0"):
482
483
PythonWrapperCodegen.memory_plan_reuse = memory_plan_reuse_patched
484
GraphLowering._update_scheduler = _update_scheduler_patched
485
-
486
-# ===================================================
487
-# torch 2.11 Inductor constrain_to_fx_strides monkeypatch
488
-# ===================================================
489
-# Patch the inductor's `constrain_to_fx_strides` to handle opaque
490
-# (non-tensor) arguments. The original calls `.stride()` on every FX
491
-# arg's meta value, which crashes on FakeScriptObject (the compile-time
492
-# proxy for hoisted opaque types). The patched version skips args
493
-# whose meta value is not a torch.Tensor.
494
-# Upstream issue: https://github.com/pytorch/pytorch/issues/175973
495
-
496
-from vllm.utils.torch_utils import is_torch_equal_or_newer
497
-
498
-if is_torch_equal_or_newer("2.11.0.dev"):
499
- import torch._inductor.ir as _ir
500
- import torch._inductor.lowering as _lowering
501
- from torch._inductor.virtualized import V as _V
502
-
503
- _orig_constrain = _lowering.constrain_to_fx_strides
504
-
505
- def _patched_constrain_to_fx_strides(fx_node, *args, **kwargs):
506
- def apply_constraint(arg, fx_arg):
507
- if isinstance(arg, _ir.IRNode):
508
- meta_val = fx_arg.meta.get("val")
509
- if isinstance(meta_val, torch.Tensor):
510
- stride_order = _ir.get_stride_order(
511
- meta_val.stride(), _V.graph.sizevars.shape_env
512
- )
513
- return _ir.ExternKernel.require_stride_order(arg, stride_order)
514
- return arg
515
- if isinstance(arg, dict):
516
- return {key: apply_constraint(arg[key], fx_arg[key]) for key in arg}
517
- return arg
518
-
519
- args = tuple(
520
- apply_constraint(arg, fx_arg) for arg, fx_arg in zip(args, fx_node.args)
521
- )
522
- kwargs = {k: apply_constraint(v, fx_node.kwargs[k]) for k, v in kwargs.items()}
523
- return args, kwargs
524
-
525
- _lowering.constrain_to_fx_strides = _patched_constrain_to_fx_strides