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 3, 09:15 UTC
[BugFix][CPU][Spec Decode] Fix Eagle implementation on CPU backend (#42468)
Signed-off-by: Ofir Zafrir <ofir.zafrir@intel.com>
Ofir Zafrir committed
May 19, 2026 at 06:16 UTC
8f16c4a5c0feb01f106e5981f22ae8808a94a28b
3 files changed
+28
-21
vllm/v1/sample/ops/topk_topp_sampler.py
+4
@@ -313,6 +313,10 @@ def apply_top_k_top_p(
313
if p is None and k is None:
314
return logits
315
316
+ # Keep CPU logits on the PyTorch path to avoid invoking Triton kernels.
317
+ if current_platform.is_cpu():
318
+ return apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=True)
319
+
320
if HAS_TRITON and logits.shape[0] >= 8:
321
return apply_top_k_top_p_triton(logits, k, p)
322
vllm/v1/worker/cpu_model_runner.py
+2
@@ -111,6 +111,8 @@ class CPUModelRunner(GPUModelRunner):
111
logger.info_once("Loading drafter model...")
112
self.drafter.load_model(self.model)
113
114
+ self._setup_eagle3_aux_hidden_state_outputs()
115
+
116
def get_model(self) -> nn.Module:
117
return self.model
118
vllm/v1/worker/gpu_model_runner.py
+22
-21
@@ -5055,27 +5055,7 @@ class GPUModelRunner(
5055
)
5056
eplb_models += 1
5057
5058
- if self.use_aux_hidden_state_outputs:
5059
- if not supports_eagle3(self.get_model()):
5060
- raise RuntimeError(
5061
- "Model does not support EAGLE3 interface but "
5062
- "aux_hidden_state_outputs was requested"
5063
- )
5064
-
5065
- # Try to get auxiliary layers from speculative config,
5066
- # otherwise use model's default layers
5067
- aux_layers = self._get_eagle3_aux_layers_from_config()
5068
- if aux_layers:
5069
- logger.info(
5070
- "Using auxiliary layers from speculative config: %s",
5071
- aux_layers,
5072
- )
5073
- else:
5074
- aux_layers = (
5075
- self.model.get_eagle3_default_aux_hidden_state_layers()
5076
- )
5077
-
5078
- self.model.set_aux_hidden_state_layers(aux_layers)
5058
+ self._setup_eagle3_aux_hidden_state_outputs()
5059
5060
# Resolve the MoE model, unwrapping VLM wrappers if needed.
5061
# VLM models (e.g. KimiK25ForConditionalGeneration) wrap the
@@ -5190,6 +5170,27 @@ class GPUModelRunner(
5170
5171
get_offloader().post_init()
5172
5173
+ def _setup_eagle3_aux_hidden_state_outputs(self) -> None:
5174
+ if not self.use_aux_hidden_state_outputs:
5175
+ return
5176
+
5177
+ if not supports_eagle3(self.get_model()):
5178
+ raise RuntimeError(
5179
+ "Model does not support EAGLE3 interface but "
5180
+ "aux_hidden_state_outputs was requested"
5181
+ )
5182
+ # Try to get auxiliary layers from speculative config,
5183
+ # otherwise use model's default layers
5184
+ aux_layers = self._get_eagle3_aux_layers_from_config()
5185
+ if aux_layers:
5186
+ logger.info(
5187
+ "Using auxiliary layers from speculative config: %s", aux_layers
5188
+ )
5189
+ else:
5190
+ aux_layers = self.model.get_eagle3_default_aux_hidden_state_layers()
5191
+
5192
+ self.model.set_aux_hidden_state_layers(aux_layers)
5193
+
5194
def _get_eagle3_aux_layers_from_config(self) -> tuple[int, ...] | None:
5195
"""Extract Eagle3 auxiliary layer indices from speculative config.
5196