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 13, 02:00 UTC
[Quantization] [Eagle] Add complete quantization support to the draft model in Eagle (#28435)
Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>
Shreyas Kulkarni committed
Nov 17, 2025 at 18:01 UTC
95ae50b7d1bf3b5b66ac39b19d3169bad5443f2e
4 files changed
+282
-29
tests/model_executor/test_eagle_quantization.py
new
+169
@@ -0,0 +1,169 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from unittest.mock import Mock, patch
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.config import LoadConfig, ModelConfig, SpeculativeConfig, VllmConfig
10
+from vllm.model_executor.models.utils import get_draft_quant_config
11
+from vllm.platforms import current_platform
12
+
13
+DEVICES = (
14
+ [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
15
+ if current_platform.is_cuda_alike()
16
+ else ["cpu"]
17
+)
18
+
19
+
20
+def test_get_draft_quant_config_with_draft_model():
21
+ mock_draft_model_config = Mock(spec=ModelConfig)
22
+ mock_load_config = Mock(spec=LoadConfig)
23
+ mock_speculative_config = Mock(spec=SpeculativeConfig)
24
+ mock_speculative_config.draft_model_config = mock_draft_model_config
25
+
26
+ mock_vllm_config = Mock(spec=VllmConfig)
27
+ mock_vllm_config.speculative_config = mock_speculative_config
28
+ mock_vllm_config.load_config = mock_load_config
29
+
30
+ mock_quant_config = Mock()
31
+ with patch.object(
32
+ VllmConfig, "get_quantization_config", return_value=mock_quant_config
33
+ ):
34
+ result = get_draft_quant_config(mock_vllm_config)
35
+
36
+ # Verify the function calls get_quantization_config with draft model config
37
+ VllmConfig.get_quantization_config.assert_called_once_with(
38
+ mock_draft_model_config, mock_load_config
39
+ )
40
+ assert result == mock_quant_config
41
+
42
+
43
+def test_get_draft_quant_config_without_draft_model():
44
+ mock_speculative_config = Mock(spec=SpeculativeConfig)
45
+ mock_speculative_config.draft_model_config = None
46
+
47
+ mock_vllm_config = Mock(spec=VllmConfig)
48
+ mock_vllm_config.speculative_config = mock_speculative_config
49
+ mock_vllm_config.load_config = Mock(spec=LoadConfig)
50
+
51
+ result = get_draft_quant_config(mock_vllm_config)
52
+
53
+ assert result is None
54
+
55
+
56
+@torch.inference_mode()
57
+@pytest.mark.parametrize("device", DEVICES)
58
+def test_fc_layer_quant_config_usage(dist_init, device) -> None:
59
+ import torch
60
+
61
+ from vllm.model_executor.layers.linear import ReplicatedLinear
62
+
63
+ if current_platform.is_cuda_alike():
64
+ torch.cuda.set_device(device)
65
+
66
+ torch.set_default_device(device)
67
+
68
+ input_size = 256
69
+ output_size = 128
70
+
71
+ fc_no_quant = ReplicatedLinear(
72
+ input_size=input_size,
73
+ output_size=output_size,
74
+ bias=False,
75
+ params_dtype=torch.float16,
76
+ quant_config=None,
77
+ prefix="fc",
78
+ )
79
+
80
+ assert fc_no_quant.quant_config is None
81
+ assert fc_no_quant.input_size == input_size
82
+ assert fc_no_quant.output_size == output_size
83
+
84
+ mock_quant_config = Mock()
85
+ fc_with_quant = ReplicatedLinear(
86
+ input_size=input_size,
87
+ output_size=output_size,
88
+ bias=False,
89
+ params_dtype=torch.float16,
90
+ quant_config=mock_quant_config,
91
+ prefix="fc",
92
+ )
93
+
94
+ assert fc_with_quant.quant_config == mock_quant_config
95
+
96
+ # Check forward pass
97
+ x = torch.randn(2, input_size, dtype=torch.float16)
98
+ output, _ = fc_no_quant(x)
99
+ assert output.shape == (2, output_size)
100
+
101
+
102
+def test_kv_cache_scale_name_handling():
103
+ # Mock a quant config that supports cache scales
104
+ mock_quant_config = Mock()
105
+ mock_quant_config.get_cache_scale = Mock(return_value="layers.0.self_attn.kv_scale")
106
+
107
+ # Condition check in load_weights
108
+ name = "layers.0.self_attn.k_proj.weight"
109
+ scale_name = mock_quant_config.get_cache_scale(name)
110
+
111
+ # Check if get_cache_scale is called and returns expected value
112
+ mock_quant_config.get_cache_scale.assert_called_once_with(name)
113
+ assert scale_name == "layers.0.self_attn.kv_scale"
114
+
115
+
116
+def test_kv_cache_scale_name_no_scale():
117
+ # Mock a quant config that returns None for get_cache_scale
118
+ mock_quant_config = Mock()
119
+ mock_quant_config.get_cache_scale = Mock(return_value=None)
120
+
121
+ name = "layers.0.mlp.gate_proj.weight"
122
+ scale_name = mock_quant_config.get_cache_scale(name)
123
+
124
+ # Should return None for weights that don't have cache scales
125
+ assert scale_name is None
126
+
127
+
128
+def test_maybe_remap_kv_scale_name():
129
+ from vllm.model_executor.model_loader.weight_utils import maybe_remap_kv_scale_name
130
+
131
+ params_dict = {
132
+ "layers.0.self_attn.kv_scale": Mock(),
133
+ "layers.1.self_attn.kv_scale": Mock(),
134
+ }
135
+
136
+ name = "layers.0.self_attn.some_scale"
137
+ remapped = maybe_remap_kv_scale_name(name, params_dict)
138
+
139
+ assert remapped in params_dict or remapped == name or remapped is None
140
+
141
+
142
+def test_load_weights_kv_scale_handling():
143
+ kv_scale_param = Mock()
144
+ kv_scale_param.weight_loader = Mock()
145
+
146
+ params_dict = {
147
+ "layers.0.self_attn.kv_scale": kv_scale_param,
148
+ }
149
+
150
+ mock_quant_config = Mock()
151
+ mock_quant_config.get_cache_scale = Mock(return_value="layers.0.self_attn.kv_scale")
152
+
153
+ # Load_weights logic for KV cache scales
154
+ name = "layers.0.self_attn.k_proj.weight"
155
+ loaded_weight_tensor = torch.tensor([1.0, 2.0])
156
+
157
+ if mock_quant_config is not None:
158
+ scale_name = mock_quant_config.get_cache_scale(name)
159
+ if scale_name:
160
+ param = params_dict[scale_name]
161
+ assert param is kv_scale_param
162
+ weight_to_load = (
163
+ loaded_weight_tensor
164
+ if loaded_weight_tensor.dim() == 0
165
+ else loaded_weight_tensor[0]
166
+ )
167
+
168
+ assert scale_name == "layers.0.self_attn.kv_scale"
169
+ assert weight_to_load == loaded_weight_tensor[0]
vllm/model_executor/models/llama_eagle.py
+41
-12
@@ -11,13 +11,22 @@ from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import VllmConfig
12
from vllm.distributed.parallel_state import get_pp_group
13
from vllm.logger import init_logger
14
+from vllm.model_executor.layers.linear import ReplicatedLinear
15
from vllm.model_executor.layers.logits_processor import LogitsProcessor
16
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
17
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
17
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
18
+from vllm.model_executor.model_loader.weight_utils import (
19
+ default_weight_loader,
20
+ maybe_remap_kv_scale_name,
21
+)
22
from vllm.model_executor.models.llama import LlamaDecoderLayer, LlamaForCausalLM
23
20
-from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight
24
+from .utils import (
25
+ AutoWeightsLoader,
26
+ get_draft_quant_config,
27
+ maybe_prefix,
28
+ process_eagle_weight,
29
+)
30
31
logger = init_logger(__name__)
32
@@ -40,14 +49,7 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
49
50
def get_quant_config(self, vllm_config: VllmConfig) -> QuantizationConfig | None:
51
"""Use drafter's quantization config instead of verifier's."""
43
- draft_model_config = vllm_config.speculative_config.draft_model_config
44
- draft_load_config = vllm_config.load_config
45
-
46
- return (
47
- VllmConfig.get_quantization_config(draft_model_config, draft_load_config)
48
- if draft_model_config
49
- else None
50
- )
52
+ return get_draft_quant_config(vllm_config)
53
54
55
@support_torch_compile
@@ -63,6 +65,9 @@ class LlamaModel(nn.Module):
65
self.config = vllm_config.speculative_config.draft_model_config.hf_config
66
self.vocab_size = self.config.vocab_size
67
68
+ # Get drafter's quantization config
69
+ self.quant_config = get_draft_quant_config(vllm_config)
70
+
71
self.embed_tokens = VocabParallelEmbedding(
72
self.config.vocab_size,
73
self.config.hidden_size,
@@ -80,8 +85,14 @@ class LlamaModel(nn.Module):
85
for i in range(self.config.num_hidden_layers)
86
]
87
)
83
- self.fc = torch.nn.Linear(
84
- self.config.hidden_size * 2, self.config.hidden_size, bias=False
88
+ self.fc = ReplicatedLinear(
89
+ input_size=self.config.hidden_size * 2,
90
+ output_size=self.config.hidden_size,
91
+ bias=False,
92
+ params_dtype=vllm_config.model_config.dtype,
93
+ quant_config=self.quant_config,
94
+ prefix=maybe_prefix(prefix, "fc"),
95
+ return_bias=False,
96
)
97
98
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
@@ -117,6 +128,24 @@ class LlamaModel(nn.Module):
128
params_dict = dict(self.named_parameters())
129
loaded_params: set[str] = set()
130
for name, loaded_weight in weights:
131
+ # Handle kv cache quantization scales
132
+ if self.quant_config is not None and (
133
+ scale_name := self.quant_config.get_cache_scale(name)
134
+ ):
135
+ # Loading kv cache quantization scales
136
+ param = params_dict[scale_name]
137
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
138
+ loaded_weight = (
139
+ loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
140
+ )
141
+ weight_loader(param, loaded_weight)
142
+ loaded_params.add(scale_name)
143
+ continue
144
+ # Remapping the name FP8 kv-scale
145
+ if "scale" in name:
146
+ name = maybe_remap_kv_scale_name(name, params_dict)
147
+ if name is None:
148
+ continue
149
for param_name, weight_name, shard_id in stacked_params_mapping:
150
if weight_name not in name:
151
continue
vllm/model_executor/models/llama_eagle3.py
+45
-17
@@ -11,19 +11,27 @@ from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import VllmConfig, get_current_vllm_config
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.layernorm import RMSNorm
14
-from vllm.model_executor.layers.linear import QKVParallelLinear
14
+from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
15
from vllm.model_executor.layers.logits_processor import LogitsProcessor
16
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
17
from vllm.model_executor.layers.vocab_parallel_embedding import (
18
ParallelLMHead,
19
VocabParallelEmbedding,
20
)
21
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
21
+from vllm.model_executor.model_loader.weight_utils import (
22
+ default_weight_loader,
23
+ maybe_remap_kv_scale_name,
24
+)
25
from vllm.model_executor.models.llama import LlamaDecoderLayer, LlamaForCausalLM
26
from vllm.multimodal import MULTIMODAL_REGISTRY
27
from vllm.multimodal.inputs import NestedTensors
28
26
-from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight
29
+from .utils import (
30
+ AutoWeightsLoader,
31
+ get_draft_quant_config,
32
+ maybe_prefix,
33
+ process_eagle_weight,
34
+)
35
36
logger = init_logger(__name__)
37
@@ -66,14 +74,7 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
74
75
def get_quant_config(self, vllm_config: VllmConfig) -> QuantizationConfig | None:
76
"""Use drafter's quantization config instead of verifier's."""
69
- draft_model_config = vllm_config.speculative_config.draft_model_config
70
- draft_load_config = vllm_config.load_config
71
-
72
- return (
73
- VllmConfig.get_quantization_config(draft_model_config, draft_load_config)
74
- if draft_model_config
75
- else None
76
- )
77
+ return get_draft_quant_config(vllm_config)
78
79
def _norm_before_residual(
80
self, hidden_states: torch.Tensor
@@ -140,6 +141,9 @@ class LlamaModel(nn.Module):
141
self.config = vllm_config.speculative_config.draft_model_config.hf_config
142
self.vocab_size = self.config.vocab_size
143
144
+ # Get drafter's quantization config
145
+ self.quant_config = get_draft_quant_config(vllm_config)
146
+
147
current_vllm_config = get_current_vllm_config()
148
149
self.embed_tokens = VocabParallelEmbedding(
@@ -160,13 +164,19 @@ class LlamaModel(nn.Module):
164
]
165
)
166
if hasattr(self.config, "target_hidden_size"):
163
- self.fc = torch.nn.Linear(
164
- self.config.target_hidden_size * 3, self.config.hidden_size, bias=False
165
- )
167
+ fc_input_size = self.config.target_hidden_size * 3
168
else:
167
- self.fc = torch.nn.Linear(
168
- self.config.hidden_size * 3, self.config.hidden_size, bias=False
169
- )
169
+ fc_input_size = self.config.hidden_size * 3
170
+ self.fc = ReplicatedLinear(
171
+ input_size=fc_input_size,
172
+ output_size=self.config.hidden_size,
173
+ bias=False,
174
+ params_dtype=vllm_config.model_config.dtype,
175
+ quant_config=self.quant_config,
176
+ prefix=maybe_prefix(prefix, "fc"),
177
+ return_bias=False,
178
+ )
179
+
180
self.norm = RMSNorm(
181
self.config.hidden_size,
182
eps=self.config.rms_norm_eps,
@@ -211,6 +221,24 @@ class LlamaModel(nn.Module):
221
for name, loaded_weight in weights:
222
if "midlayer." in name:
223
name = name.replace("midlayer.", "layers.0.")
224
+ # Handle kv cache quantization scales
225
+ if self.quant_config is not None and (
226
+ scale_name := self.quant_config.get_cache_scale(name)
227
+ ):
228
+ # Loading kv cache quantization scales
229
+ param = params_dict[scale_name]
230
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
231
+ loaded_weight = (
232
+ loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
233
+ )
234
+ weight_loader(param, loaded_weight)
235
+ loaded_params.add(scale_name)
236
+ continue
237
+ # Remapping the name FP8 kv-scale
238
+ if "scale" in name:
239
+ name = maybe_remap_kv_scale_name(name, params_dict)
240
+ if name is None:
241
+ continue
242
for param_name, weight_name, shard_id in stacked_params_mapping:
243
if weight_name not in name:
244
continue
vllm/model_executor/models/utils.py
+27
@@ -18,6 +18,9 @@ from vllm.distributed import (
18
get_tensor_model_parallel_world_size,
19
)
20
from vllm.logger import init_logger
21
+from vllm.model_executor.layers.quantization.base_config import (
22
+ QuantizationConfig,
23
+)
24
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
25
from vllm.model_executor.models.interfaces import supports_any_eagle
26
from vllm.multimodal import NestedTensors
@@ -715,6 +718,30 @@ def maybe_prefix(prefix: str, name: str) -> str:
718
return name if not prefix else f"{prefix}.{name}"
719
720
721
+def get_draft_quant_config(
722
+ vllm_config: VllmConfig,
723
+) -> QuantizationConfig | None:
724
+ """Get quantization config for Draft models.
725
+
726
+ Draft models should use their own quantization config instead of the verifier/target
727
+ model's config. This helper retrieves the draft model's quantization config.
728
+
729
+ Args:
730
+ vllm_config: The vLLM configuration object.
731
+
732
+ Returns:
733
+ The draft model's config if available, None otherwise.
734
+ """
735
+ draft_model_config = vllm_config.speculative_config.draft_model_config
736
+ draft_load_config = vllm_config.load_config
737
+
738
+ return (
739
+ VllmConfig.get_quantization_config(draft_model_config, draft_load_config)
740
+ if draft_model_config
741
+ else None
742
+ )
743
+
744
+
745
def extract_layer_index(layer_name: str, num_attn_module: int = 1) -> int:
746
"""
747
Extract the layer index from the module name.