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 14, 05:30 UTC
[Perf] Optimize memory peak during EAGLE model loading. (#24585)
Signed-off-by: Chen Ding <candy.dc@alibaba-inc.com>
Chen Ding committed
Sep 19, 2025 at 11:31 UTC
1a0a04dae94b7a768c0d59b4f687bcf5e12d3127
3 files changed
+26
-26
vllm/model_executor/models/deepseek_eagle.py
+8
-7
@@ -229,14 +229,15 @@ class EagleDeepseekV3ForCausalLM(DeepseekV3ForCausalLM):
229
return logits
230
231
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
232
+
233
+ def transform(inputs):
234
+ name, loaded_weight = inputs
235
+ if "lm_head" not in name:
236
+ name = "model." + name
237
+ return name, loaded_weight
238
+
239
loader = AutoWeightsLoader(
240
self,
241
skip_prefixes=None,
242
)
236
-
237
- model_weights = {}
238
- for name, loaded_weight in weights:
239
- if "lm_head" not in name:
240
- name = "model." + name
241
- model_weights[name] = loaded_weight
242
- loader.load_weights(model_weights.items())
243
+ loader.load_weights(map(transform, weights))
vllm/model_executor/models/llama4_eagle.py
+10
-12
@@ -205,23 +205,21 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM):
205
206
def load_weights(self, weights: Iterable[tuple[str,
207
torch.Tensor]]) -> None:
208
+
209
+ def transform(inputs):
210
+ name, loaded_weight = inputs
211
+ name, weight = self.permute_qk_weight_for_rotary(
212
+ name, loaded_weight)
213
+ if "lm_head" not in name:
214
+ name = "model." + name
215
+ return name, weight
216
+
217
loader = AutoWeightsLoader(
218
self,
219
# lm_head is tied with target model (Llama4ForCausalLM)
220
skip_prefixes=(["lm_head."]),
221
)
213
-
214
- model_weights = {}
215
- weights = [
216
- self.permute_qk_weight_for_rotary(name, loaded_weight)
217
- for name, loaded_weight in weights
218
- ]
219
- for name, loaded_weight in weights:
220
- if "lm_head" not in name:
221
- name = "model." + name
222
- model_weights[name] = loaded_weight
223
-
224
- loader.load_weights(model_weights.items())
222
+ loader.load_weights(map(transform, weights))
223
224
def get_input_embeddings(
225
self,
vllm/model_executor/models/llama_eagle.py
+8
-7
@@ -158,14 +158,15 @@ class EagleLlamaForCausalLM(LlamaForCausalLM):
158
return self.model(input_ids, positions, hidden_states)
159
160
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
161
+
162
+ def transform(inputs):
163
+ name, loaded_weight = inputs
164
+ if "lm_head" not in name:
165
+ name = "model." + name
166
+ return name, loaded_weight
167
+
168
loader = AutoWeightsLoader(
169
self,
170
skip_prefixes=None,
171
)
165
-
166
- model_weights = {}
167
- for name, loaded_weight in weights:
168
- if "lm_head" not in name:
169
- name = "model." + name
170
- model_weights[name] = loaded_weight
171
- loader.load_weights(model_weights.items())
172
+ loader.load_weights(map(transform, weights))