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

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

[torch.compile] Speed up MOE handling in forward_context (#33184)

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

Richard Zou committed Jan 27, 2026 at 18:17 UTC d9aa39a3bbc1cb48cff8bfe9d6156c11fc4c42d9

4 files changed +22 -18

tests/kernels/moe/test_moe.py

+1 -1

@@ -715,7 +715,7 @@ def test_mixtral_moe(

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

718 + get_forward_context().all_moe_layers = None

719

720 # Run forward passes for both MoE blocks

721 hf_states, _ = hf_moe.forward(hf_inputs)

vllm/config/compilation.py

@@ -597,6 +597,10 @@ class CompilationConfig:

597 Map from layer name to layer objects that need to be accessed outside

598 model code, e.g., Attention, FusedMOE when dp_size>1."""

599

600 + static_all_moe_layers: list[str] = field(default_factory=list, init=False)

601 + """The names of all the MOE layers in the model

602 + """

603 +

604 # Attention ops; used for piecewise cudagraphs

605 # Use PyTorch operator format: "namespace::name"

606 _attention_ops: ClassVar[list[str]] = [

vllm/forward_context.py

+8 -13

@@ -217,9 +217,11 @@ class ForwardContext:

217 # the graph.

218 #

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

220 - # custom ops needs, in reverse order, in the ForwardContext.

220 + # custom ops needs in the ForwardContext (all_moe_layers)

221 + # as well as a counter (moe_layer_index).

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

222 - # When the custom op needs the string, pop the string from this list.

223 + # When the custom op needs a layer string, get the next string

224 + # from all_moe_layers and increment the counter.

225 #

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

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

@@ -233,7 +235,8 @@ class ForwardContext:

235 #

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

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

236 - remaining_moe_layers: list[str] | None = None

238 + all_moe_layers: list[str] | None = None

239 + moe_layer_index: int = 0

240

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

242

@@ -271,17 +274,9 @@ def create_forward_context(

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

275 skip_compiled: bool = False,

276 ):

274 - no_compile_layers = vllm_config.compilation_config.static_forward_context

275 - from vllm.model_executor.layers.fused_moe.layer import FusedMoE

276 -

277 - remaining_moe_layers = [

278 - name for name, layer in no_compile_layers.items() if isinstance(layer, FusedMoE)

279 - ]

280 - remaining_moe_layers.reverse()

281 -

277 return ForwardContext(

283 - no_compile_layers=no_compile_layers,

284 - remaining_moe_layers=remaining_moe_layers,

278 + no_compile_layers=vllm_config.compilation_config.static_forward_context,

279 + all_moe_layers=vllm_config.compilation_config.static_all_moe_layers,

280 virtual_engine=virtual_engine,

281 attn_metadata=attn_metadata,

282 slot_mapping=slot_mapping or {},

vllm/model_executor/layers/fused_moe/layer.py

+9 -4

@@ -407,6 +407,7 @@ class FusedMoE(CustomOp):

407 if prefix in compilation_config.static_forward_context:

408 raise ValueError("Duplicate layer name: {}".format(prefix))

409 compilation_config.static_forward_context[prefix] = self

410 + compilation_config.static_all_moe_layers.append(prefix)

411 self.layer_name = prefix

412

413 self.enable_eplb = enable_eplb

@@ -1566,7 +1567,7 @@ class FusedMoE(CustomOp):

1567 # Can be unavailable or None in unittests

1568 if (

1569 is_forward_context_available()

1569 - and get_forward_context().remaining_moe_layers is not None

1570 + and get_forward_context().all_moe_layers is not None

1571 ):

1572 return "from_forward_context"

1573 return self.layer_name

@@ -1987,13 +1988,17 @@ class FusedMoE(CustomOp):

1988 def get_layer_from_name(layer_name: str) -> FusedMoE:

1989 forward_context: ForwardContext = get_forward_context()

1990 if layer_name == "from_forward_context":

1990 - if not forward_context.remaining_moe_layers:

1991 + all_moe_layers = forward_context.all_moe_layers

1992 + assert all_moe_layers is not None

1993 + moe_layer_index = forward_context.moe_layer_index

1994 + if moe_layer_index >= len(all_moe_layers):

1995 raise AssertionError(

1992 - "We expected the number of MOE layers in `remaining_moe_layers` "

1996 + "We expected the number of MOE layers in `all_moe_layers` "

1997 "to be equal to the number of "

1998 "{vllm.moe_forward, vllm.moe_forward_shared} calls."

1999 )

1996 - layer_name = forward_context.remaining_moe_layers.pop()

2000 + layer_name = all_moe_layers[moe_layer_index]

2001 + forward_context.moe_layer_index += 1

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

2003 return self

2004