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
[Optimization] Avoid repeated model architecture conversion for pooling models (#25261)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Sep 20, 2025 at 13:30 UTC
c60e6137f0bf2034853919b3a9d705d7e06b93cf
2 files changed
+37
-2
vllm/config/model.py
+21
-1
@@ -322,8 +322,28 @@ class ModelConfig:
322
factors.append(self.override_generation_config)
323
factors.append(self.rope_scaling)
324
factors.append(self.rope_theta)
325
+
326
# hf_config can control how the model looks!
326
- factors.append(self.hf_config.to_json_string())
327
+ try:
328
+ hf_config_json = self.hf_config.to_json_string(use_diff=False)
329
+ except TypeError:
330
+ from transformers import PretrainedConfig
331
+
332
+ from vllm.utils.jsontree import json_map_leaves
333
+
334
+ # Handle nested HF configs with unserializable values gracefully
335
+ hf_config_json = json.dumps(
336
+ json_map_leaves(
337
+ lambda v: v.to_dict()
338
+ if isinstance(v, PretrainedConfig) else str(v),
339
+ self.hf_config.to_dict(),
340
+ ),
341
+ indent=2,
342
+ sort_keys=True,
343
+ ) + "\n"
344
+
345
+ factors.append(hf_config_json)
346
+
347
str_factors = str(factors)
348
assert_hashable(str_factors)
349
return hashlib.sha256(str(factors).encode()).hexdigest()
vllm/model_executor/model_loader/utils.py
+16
-1
@@ -165,7 +165,11 @@ def device_loading_context(module: torch.nn.Module,
165
# New parameters or parameters already on target device are untouched
166
167
168
-def get_model_architecture(
168
+_MODEL_ARCH_BY_HASH = dict[str, tuple[type[nn.Module], str]]()
169
+"""Caches the outputs of `_get_model_architecture`."""
170
+
171
+
172
+def _get_model_architecture(
173
model_config: ModelConfig) -> tuple[type[nn.Module], str]:
174
architectures = getattr(model_config.hf_config, "architectures", [])
175
@@ -209,6 +213,17 @@ def get_model_architecture(
213
return model_cls, arch
214
215
216
+def get_model_architecture(
217
+ model_config: ModelConfig) -> tuple[type[nn.Module], str]:
218
+ key = model_config.compute_hash()
219
+ if key in _MODEL_ARCH_BY_HASH:
220
+ return _MODEL_ARCH_BY_HASH[key]
221
+
222
+ model_arch = _get_model_architecture(model_config)
223
+ _MODEL_ARCH_BY_HASH[key] = model_arch
224
+ return model_arch
225
+
226
+
227
def get_model_cls(model_config: ModelConfig) -> type[nn.Module]:
228
return get_model_architecture(model_config)[0]
229