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 8, 16:30 UTC
[torch.compile] Add an option to force-enable the MOE cold start optimization (#33735)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Feb 8, 2026 at 13:42 UTC
4df841fe7538cb8de281b9d78e37ba51ac35b5da
3 files changed
+18
-12
vllm/config/compilation.py
+9
-3
@@ -593,7 +593,7 @@ class CompilationConfig:
593
local_cache_dir: str = field(default=None, init=False) # type: ignore
594
"""local cache dir for each rank"""
595
596
- fast_moe_cold_start = True
596
+ fast_moe_cold_start: bool | None = None
597
"""Optimization for fast MOE cold start.
598
599
This is a bit of a hack that assumes that:
@@ -604,8 +604,14 @@ class CompilationConfig:
604
When the above two conditions hold, this option greatly decreases cold start
605
time for MOE models.
606
607
- If the above two conditions don't hold, then this option will lead to silent
608
- incorrectness. The only condition in which this doesn't hold is speculative
607
+ The options are:
608
+ - True: optimization is always on
609
+ - False: optimization is always off
610
+ - None: optimization is on usually but off for speculative decoding
611
+
612
+ If conditions 1&2 don't hold then this option will lead to silent
613
+ incorrectness.
614
+ The only condition in which this doesn't hold is speculative
615
decoding, where there is a draft model that may have MOEs in them.
616
617
NB: We're working on a longer-term solution that doesn't need these assumptions.
vllm/config/vllm.py
+8
@@ -806,6 +806,14 @@ class VllmConfig:
806
else:
807
self.compilation_config.custom_ops.append("+rms_norm")
808
809
+ if self.compilation_config.fast_moe_cold_start is None:
810
+ # resolve default behavior: try to be as safe as possible
811
+ # this config is unsafe if any spec decoding draft model has a MOE.
812
+ # We'll conservatively turn it off if we see spec decoding.
813
+ self.compilation_config.fast_moe_cold_start = (
814
+ self.speculative_config is None
815
+ )
816
+
817
if current_platform.support_static_graph_mode():
818
# if cudagraph_mode has full cudagraphs, we need to check support
819
if model_config := self.model_config:
vllm/forward_context.py
+1
-9
@@ -287,15 +287,7 @@ def create_forward_context(
287
skip_compiled: bool = False,
288
):
289
if vllm_config.compilation_config.fast_moe_cold_start:
290
- if vllm_config.speculative_config is None:
291
- all_moe_layers = vllm_config.compilation_config.static_all_moe_layers
292
- else:
293
- logger.warning_once(
294
- "vllm_config.compilation_config.fast_moe_cold_start is not "
295
- "compatible with speculative decoding so we are ignoring "
296
- "fast_moe_cold_start."
297
- )
298
- all_moe_layers = None
290
+ all_moe_layers = vllm_config.compilation_config.static_all_moe_layers
291
else:
292
all_moe_layers = None
293