From 0e55d08a68002339e1326b154bab45a0076bcc3d Mon Sep 17 00:00:00 2001 From: Rahul Tuli Date: Fri, 10 Oct 2025 10:19:28 -0400 Subject: [PATCH 1/2] Extend: fix from #25883 to llama_eagle.py Signed-off-by: Rahul Tuli --- vllm/model_executor/models/llama_eagle.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/vllm/model_executor/models/llama_eagle.py b/vllm/model_executor/models/llama_eagle.py index 8f4ba886777..fb52cf1b54c 100644 --- a/vllm/model_executor/models/llama_eagle.py +++ b/vllm/model_executor/models/llama_eagle.py @@ -12,6 +12,7 @@ from vllm.distributed.parallel_state import get_pp_group from vllm.logger import init_logger from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.llama import LlamaDecoderLayer, LlamaForCausalLM @@ -37,6 +38,17 @@ def __init__( del self.input_layernorm self.input_layernorm = nn.Identity() + def get_quant_config(self, vllm_config: VllmConfig) -> Optional[QuantizationConfig]: + """Use drafter's quantization config instead of verifier's.""" + draft_model_config = vllm_config.speculative_config.draft_model_config + draft_load_config = vllm_config.load_config + + return ( + VllmConfig.get_quantization_config(draft_model_config, draft_load_config) + if draft_model_config + else None + ) + @support_torch_compile class LlamaModel(nn.Module): From 1ba253e8d4ca7753427a63417d273e01e0e14e69 Mon Sep 17 00:00:00 2001 From: Rahul Tuli Date: Mon, 13 Oct 2025 09:24:40 -0400 Subject: [PATCH 2/2] fix precommit Signed-off-by: Rahul Tuli --- vllm/model_executor/models/llama_eagle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/models/llama_eagle.py b/vllm/model_executor/models/llama_eagle.py index fb52cf1b54c..3617294bd62 100644 --- a/vllm/model_executor/models/llama_eagle.py +++ b/vllm/model_executor/models/llama_eagle.py @@ -38,7 +38,7 @@ def __init__( del self.input_layernorm self.input_layernorm = nn.Identity() - def get_quant_config(self, vllm_config: VllmConfig) -> Optional[QuantizationConfig]: + def get_quant_config(self, vllm_config: VllmConfig) -> QuantizationConfig | None: """Use drafter's quantization config instead of verifier's.""" draft_model_config = vllm_config.speculative_config.draft_model_config draft_load_config = vllm_config.load_config