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, 01:15 UTC
[torch.compile] Don't do the fast moe cold start optimization if there is speculative decoding (#33624)
Signed-off-by: Richard Zou <zou3519@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Richard Zou committed
Feb 2, 2026 at 19:38 UTC
5eac9a1b341b93478d0d0d57239c92edd18ad19e
2 files changed
+32
-1
vllm/config/compilation.py
+18
@@ -595,6 +595,24 @@ class CompilationConfig:
595
local_cache_dir: str = field(default=None, init=False) # type: ignore
596
"""local cache dir for each rank"""
597
598
+ fast_moe_cold_start = True
599
+ """Optimization for fast MOE cold start.
600
+
601
+ This is a bit of a hack that assumes that:
602
+ 1. the only decoder forward pass being run is the current model
603
+ 2. the decoder forward pass runs all of the MOEs in the order in which they
604
+ are initialized
605
+
606
+ When the above two conditions hold, this option greatly decreases cold start
607
+ time for MOE models.
608
+
609
+ If the above two conditions don't hold, then this option will lead to silent
610
+ incorrectness. The only condition in which this doesn't hold is speculative
611
+ decoding, where there is a draft model that may have MOEs in them.
612
+
613
+ NB: We're working on a longer-term solution that doesn't need these assumptions.
614
+ """
615
+
616
# keep track of enabled and disabled custom ops
617
enabled_custom_ops: Counter[str] = field(default_factory=Counter, init=False)
618
"""custom ops that are enabled"""
vllm/forward_context.py
+14
-1
@@ -286,9 +286,22 @@ def create_forward_context(
286
additional_kwargs: dict[str, Any] | None = None,
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
299
+ else:
300
+ all_moe_layers = None
301
+
302
return ForwardContext(
303
no_compile_layers=vllm_config.compilation_config.static_forward_context,
291
- all_moe_layers=vllm_config.compilation_config.static_all_moe_layers,
304
+ all_moe_layers=all_moe_layers,
305
virtual_engine=virtual_engine,
306
attn_metadata=attn_metadata,
307
slot_mapping=slot_mapping or {},