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, 11:30 UTC
[torch.compile] Undo the fast_moe_cold_start hack in torch>=2.11 (#35475)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Mar 1, 2026 at 16:44 UTC
e82fbeec7b360af4fb908bf67a659b22f93266d3
4 files changed
+109
-10
vllm/config/vllm.py
+7
-1
@@ -883,7 +883,13 @@ class VllmConfig:
883
self.compilation_config.pass_config.enable_sp = False
884
self.compilation_config.pass_config.fuse_gemm_comms = False
885
886
- if self.compilation_config.fast_moe_cold_start is None:
886
+ from vllm.utils.torch_utils import HAS_OPAQUE_TYPE
887
+
888
+ if HAS_OPAQUE_TYPE:
889
+ # On torch >= 2.11 the hoisted OpaqueObject approach supersedes
890
+ # fast_moe_cold_start, so force it off.
891
+ self.compilation_config.fast_moe_cold_start = False
892
+ elif self.compilation_config.fast_moe_cold_start is None:
893
# resolve default behavior: try to be as safe as possible
894
# this config is unsafe if any spec decoding draft model has a MOE.
895
# We'll conservatively turn it off if we see spec decoding.
vllm/env_override.py
+41
@@ -482,3 +482,44 @@ 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
vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py
+26
-9
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from contextlib import nullcontext
4
+from typing import TYPE_CHECKING
5
6
import torch
7
import torch.nn.functional as F
@@ -30,6 +31,8 @@ from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
31
from vllm.platforms import current_platform
32
from vllm.utils.math_utils import cdiv
33
from vllm.utils.torch_utils import (
34
+ HAS_OPAQUE_TYPE,
35
+ ModuleName,
36
aux_stream,
37
current_stream,
38
direct_register_custom_op,
@@ -56,13 +59,27 @@ def get_layer_from_name(layer_name: str) -> torch.nn.Module:
59
return forward_context.no_compile_layers[layer_name]
60
61
62
+# On torch >= 2.11, layer_name is a hoisted ModuleName opaque object;
63
+# on older versions it remains a plain str.
64
+if TYPE_CHECKING:
65
+ from typing import TypeAlias
66
+
67
+ _layer_name_type: TypeAlias = str | ModuleName
68
+else:
69
+ _layer_name_type = ModuleName if HAS_OPAQUE_TYPE else str
70
+
71
+
72
+def _resolve_layer_name(layer_name: str | ModuleName) -> str:
73
+ return layer_name.value if isinstance(layer_name, ModuleName) else layer_name
74
+
75
+
76
def _moe_forward(
77
hidden_states: torch.Tensor,
78
router_logits: torch.Tensor,
79
shared_experts_input: torch.Tensor | None,
63
- layer_name: str,
80
+ layer_name: _layer_name_type,
81
) -> torch.Tensor:
65
- layer = get_layer_from_name(layer_name)
82
+ layer = get_layer_from_name(_resolve_layer_name(layer_name))
83
# TODO(bnell): this can be removed after MK migration is complete.
84
layer.ensure_moe_quant_config_init()
85
return layer.runner.forward_impl(
@@ -74,7 +91,7 @@ def _moe_forward_fake(
91
hidden_states: torch.Tensor,
92
router_logits: torch.Tensor,
93
shared_experts_input: torch.Tensor | None,
77
- layer_name: str,
94
+ layer_name: _layer_name_type,
95
) -> torch.Tensor:
96
return torch.empty_like(hidden_states)
97
@@ -83,9 +100,9 @@ def _moe_forward_shared(
100
hidden_states: torch.Tensor,
101
router_logits: torch.Tensor,
102
shared_experts_input: torch.Tensor | None,
86
- layer_name: str,
103
+ layer_name: _layer_name_type,
104
) -> tuple[torch.Tensor, torch.Tensor]:
88
- layer = get_layer_from_name(layer_name)
105
+ layer = get_layer_from_name(_resolve_layer_name(layer_name))
106
# TODO(bnell): this can be removed after MK migration is complete.
107
layer.ensure_moe_quant_config_init()
108
return layer.runner.forward_impl(
@@ -97,7 +114,7 @@ def _moe_forward_shared_fake(
114
hidden_states: torch.Tensor,
115
router_logits: torch.Tensor,
116
shared_experts_input: torch.Tensor | None,
100
- layer_name: str,
117
+ layer_name: _layer_name_type,
118
) -> tuple[torch.Tensor, torch.Tensor]:
119
# Output shapes:
120
# - fused_out: same as hidden_states (routed experts use transformed size)
@@ -105,12 +122,10 @@ def _moe_forward_shared_fake(
122
# hidden_states
123
# (For latent MoE: shared experts use original hidden_size, not latent size)
124
fused_out = torch.empty_like(hidden_states)
108
-
125
if shared_experts_input is not None:
126
shared_out = torch.empty_like(shared_experts_input)
127
else:
128
shared_out = torch.empty_like(hidden_states)
113
-
129
return shared_out, fused_out
130
131
@@ -367,7 +382,9 @@ class DefaultMoERunner(MoERunner):
382
assert len(trunc_sizes) == 1
383
return func(states, trunc_sizes[0])
384
370
- def _encode_layer_name(self) -> str:
385
+ def _encode_layer_name(self) -> str | ModuleName:
386
+ if HAS_OPAQUE_TYPE:
387
+ return ModuleName(self.layer_name)
388
# Can be unavailable or None in unittests
389
if (
390
is_forward_context_available()
vllm/utils/torch_utils.py
+35
@@ -740,6 +740,41 @@ def is_torch_equal(target: str) -> bool:
740
return Version(importlib.metadata.version("torch")) == Version(target)
741
742
743
+HAS_OPAQUE_TYPE = is_torch_equal_or_newer("2.11.0.dev")
744
+
745
+if HAS_OPAQUE_TYPE:
746
+ from torch._opaque_base import OpaqueBase
747
+else:
748
+ OpaqueBase = object # type: ignore[misc, assignment]
749
+
750
+
751
+class ModuleName(OpaqueBase): # type: ignore[misc]
752
+ """Wraps a module name string for use as a torch opaque type.
753
+
754
+ When torch >= 2.11, this is registered as a hoisted value-type opaque
755
+ object so that torch.compile lifts it as a graph input instead of baking
756
+ it as a constant. This avoids per-layer recompilation for MOE ops.
757
+ """
758
+
759
+ def __init__(self, value: str):
760
+ self.value = value
761
+
762
+ def __eq__(self, other):
763
+ return isinstance(other, ModuleName) and self.value == other.value
764
+
765
+ def __hash__(self):
766
+ return hash(self.value)
767
+
768
+ def __fx_repr__(self):
769
+ return (f"ModuleName({self.value!r})", {ModuleName})
770
+
771
+
772
+if HAS_OPAQUE_TYPE:
773
+ from torch._library.opaque_object import register_opaque_type
774
+
775
+ register_opaque_type(ModuleName, typ="value", hoist=True)
776
+
777
+
778
# Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform
779
def supports_xccl() -> bool:
780
return torch.distributed.is_xccl_available()