@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 654a71fc3c

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 9, 19:30 UTC

[torch.compile] Improve Cold Start for MoEs (#32805)

Signed-off-by: Richard Zou <zou3519@gmail.com>

Richard Zou committed Jan 22, 2026 at 10:44 UTC 654a71fc3c6d0665b8b4805b0219e35dc485416e

3 files changed +71 -9

tests/kernels/moe/test_moe.py

+5 -1

@@ -23,7 +23,7 @@ from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe

23 from vllm._aiter_ops import rocm_aiter_ops

24 from vllm.config import VllmConfig, set_current_vllm_config

25 from vllm.distributed.parallel_state import init_distributed_environment

26 -from vllm.forward_context import set_forward_context

26 +from vllm.forward_context import get_forward_context, set_forward_context

27 from vllm.model_executor.layers.fused_moe import (

28 fused_topk,

29 )

@@ -713,6 +713,10 @@ def test_mixtral_moe(

713

714 vllm_moe.experts.quant_method.process_weights_after_loading(vllm_moe.experts)

715

716 + # need to override the forward context for unittests, otherwise it assumes

717 + # we're running the model forward pass (the model specified in vllm_config)

718 + get_forward_context().remaining_moe_layers = None

719 +

720 # Run forward passes for both MoE blocks

721 hf_states, _ = hf_moe.forward(hf_inputs)

722 vllm_states = vllm_moe.forward(vllm_inputs)

vllm/forward_context.py

+34 -1

@@ -210,6 +210,30 @@ class ForwardContext:

210 # If True, bypass the compiled model call, e.g. by using .forward() directly

211 skip_compiled: bool = False

212

213 + # For torch.compile cold start times, we need to avoid hard-coding

214 + # any strings into the graph. Right now, the vllm.moe_forward

215 + # and vllm.moe_forward_shared custom operators hard-code strings into

216 + # the graph.

217 + #

218 + # The workaround is to store a list of the strings that each of those

219 + # custom ops needs, in reverse order, in the ForwardContext.

220 + # The ForwardContext object is alive for the duration of the forward pass.

221 + # When the custom op needs the string, pop the string from this list.

222 + #

223 + # This assumes that the custom operators will always be executed in

224 + # order and that torch.compile will not try to reorder these

225 + # operations with respect to each other.

226 + #

227 + # TODO(https://github.com/vllm-project/vllm/issues/31985):

228 + # There are longer-term solutions, like unwrapping the moe custom operator,

229 + # that aren't ready yet.

230 + # We could also treat the string as a "symbolic input" to the graph but

231 + # the PyTorch-side bits for that aren't ready yet either.

232 + #

233 + # If this value is None (like in some tests), then we end up baking the string

234 + # into the graph. Otherwise, the moe custom ops will pop a string from this list.

235 + remaining_moe_layers: list[str] | None = None

236 +

237 additional_kwargs: dict[str, Any] = field(default_factory=dict)

238

239 def __post_init__(self):

@@ -245,8 +269,17 @@ def create_forward_context(

269 additional_kwargs: dict[str, Any] | None = None,

270 skip_compiled: bool = False,

271 ):

272 + no_compile_layers = vllm_config.compilation_config.static_forward_context

273 + from vllm.model_executor.layers.fused_moe.layer import FusedMoE

274 +

275 + remaining_moe_layers = [

276 + name for name, layer in no_compile_layers.items() if isinstance(layer, FusedMoE)

277 + ]

278 + remaining_moe_layers.reverse()

279 +

280 return ForwardContext(

249 - no_compile_layers=vllm_config.compilation_config.static_forward_context,

281 + no_compile_layers=no_compile_layers,

282 + remaining_moe_layers=remaining_moe_layers,

283 virtual_engine=virtual_engine,

284 attn_metadata=attn_metadata,

285 dp_metadata=dp_metadata,

vllm/model_executor/layers/fused_moe/layer.py

+32 -7

@@ -22,7 +22,11 @@ from vllm.distributed import (

22 tensor_model_parallel_all_reduce,

23 )

24 from vllm.distributed.eplb.eplb_state import EplbLayerState, EplbState

25 -from vllm.forward_context import ForwardContext, get_forward_context

25 +from vllm.forward_context import (

26 + ForwardContext,

27 + get_forward_context,

28 + is_forward_context_available,

29 +)

30 from vllm.logger import init_logger

31 from vllm.model_executor.custom_op import CustomOp

32 from vllm.model_executor.layers.fused_moe.config import (

@@ -1564,6 +1568,15 @@ class FusedMoE(CustomOp):

1568 states = self.maybe_all_reduce_tensor_model_parallel(states)

1569 return states

1570

1571 + def encode_layer_name() -> str:

1572 + # Can be unavailable or None in unittests

1573 + if (

1574 + is_forward_context_available()

1575 + and get_forward_context().remaining_moe_layers is not None

1576 + ):

1577 + return "from_forward_context"

1578 + return self.layer_name

1579 +

1580 if self.shared_experts is None:

1581 if current_platform.is_tpu() or current_platform.is_cpu():

1582 # TODO: Once the OOM issue for the TPU backend is resolved, we

@@ -1573,7 +1586,7 @@ class FusedMoE(CustomOp):

1586 assert not isinstance(fused_output, tuple)

1587 else:

1588 fused_output = torch.ops.vllm.moe_forward(

1576 - hidden_states, router_logits, self.layer_name

1589 + hidden_states, router_logits, encode_layer_name()

1590 )

1591 return reduce_output(fused_output)[..., :og_hidden_states]

1592 else:

@@ -1586,7 +1599,7 @@ class FusedMoE(CustomOp):

1599 )

1600 else:

1601 shared_output, fused_output = torch.ops.vllm.moe_forward_shared(

1589 - hidden_states, router_logits, self.layer_name

1602 + hidden_states, router_logits, encode_layer_name()

1603 )

1604 return (

1605 reduce_output(shared_output)[..., :og_hidden_states],

@@ -1936,13 +1949,26 @@ class FusedMoE(CustomOp):

1949 return s

1950

1951

1952 +def get_layer_from_name(layer_name: str) -> FusedMoE:

1953 + forward_context: ForwardContext = get_forward_context()

1954 + if layer_name == "from_forward_context":

1955 + if not forward_context.remaining_moe_layers:

1956 + raise AssertionError(

1957 + "We expected the number of MOE layers in `remaining_moe_layers` "

1958 + "to be equal to the number of "

1959 + "{vllm.moe_forward, vllm.moe_forward_shared} calls."

1960 + )

1961 + layer_name = forward_context.remaining_moe_layers.pop()

1962 + self = cast(FusedMoE, forward_context.no_compile_layers[layer_name])

1963 + return self

1964 +

1965 +

1966 def moe_forward(

1967 hidden_states: torch.Tensor,

1968 router_logits: torch.Tensor,

1969 layer_name: str,

1970 ) -> torch.Tensor:

1944 - forward_context: ForwardContext = get_forward_context()

1945 - self = forward_context.no_compile_layers[layer_name]

1971 + self = get_layer_from_name(layer_name)

1972 assert self.shared_experts is None

1973 return self.forward_impl(hidden_states, router_logits)

1974

@@ -1969,8 +1995,7 @@ def moe_forward_shared(

1995 router_logits: torch.Tensor,

1996 layer_name: str,

1997 ) -> tuple[torch.Tensor, torch.Tensor]:

1972 - forward_context: ForwardContext = get_forward_context()

1973 - self = forward_context.no_compile_layers[layer_name]

1998 + self = get_layer_from_name(layer_name)

1999 assert self.shared_experts is not None

2000 return self.forward_impl(hidden_states, router_logits)

2001