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 6, 18:15 UTC
[Misc] Clean up Kimi-audio whisper encoder loading (#36903)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Mar 14, 2026 at 23:37 UTC
a8e8d62dd80f53444ae62191fa0bd3901a02c7e7
3 files changed
+89
-116
vllm/model_executor/model_loader/default_loader.py
+15
-4
@@ -52,6 +52,9 @@ class DefaultModelLoader(BaseModelLoader):
52
revision: str | None
53
"""The optional model revision."""
54
55
+ subfolder: str | None = None
56
+ """The subfolder inside the model repo."""
57
+
58
prefix: str = ""
59
"""A prefix to prepend to all weights."""
60
@@ -81,6 +84,7 @@ class DefaultModelLoader(BaseModelLoader):
84
def _prepare_weights(
85
self,
86
model_name_or_path: str,
87
+ subfolder: str | None,
88
revision: str | None,
89
fall_back_to_pt: bool,
90
allow_patterns_overrides: list[str] | None,
@@ -143,11 +147,15 @@ class DefaultModelLoader(BaseModelLoader):
147
self.load_config.download_dir,
148
allow_patterns,
149
revision,
150
+ subfolder=subfolder,
151
ignore_patterns=self.load_config.ignore_patterns,
152
)
153
else:
154
hf_folder = model_name_or_path
155
156
+ if subfolder is not None:
157
+ hf_folder = os.path.join(hf_folder, subfolder)
158
+
159
hf_weights_files: list[str] = []
160
for pattern in allow_patterns:
161
hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
@@ -166,8 +174,9 @@ class DefaultModelLoader(BaseModelLoader):
174
download_safetensors_index_file_from_hf(
175
model_name_or_path,
176
index_file,
169
- self.load_config.download_dir,
170
- revision,
177
+ cache_dir=self.load_config.download_dir,
178
+ subfolder=subfolder,
179
+ revision=revision,
180
)
181
hf_weights_files = filter_duplicate_safetensors_files(
182
hf_weights_files, hf_folder, index_file
@@ -189,6 +198,7 @@ class DefaultModelLoader(BaseModelLoader):
198
extra_config = self.load_config.model_loader_extra_config
199
hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
200
source.model_or_path,
201
+ source.subfolder,
202
source.revision,
203
source.fall_back_to_pt,
204
source.allow_patterns_overrides,
@@ -269,8 +279,9 @@ class DefaultModelLoader(BaseModelLoader):
279
280
def download_model(self, model_config: ModelConfig) -> None:
281
self._prepare_weights(
272
- model_config.model,
273
- model_config.revision,
282
+ model_name_or_path=model_config.model,
283
+ subfolder=None,
284
+ revision=model_config.revision,
285
fall_back_to_pt=True,
286
allow_patterns_overrides=None,
287
)
vllm/model_executor/model_loader/weight_utils.py
+13
-1
@@ -472,6 +472,7 @@ def download_weights_from_hf(
472
cache_dir: str | None,
473
allow_patterns: list[str],
474
revision: str | None = None,
475
+ subfolder: str | None = None,
476
ignore_patterns: str | list[str] | None = None,
477
) -> str:
478
"""Download model weights from Hugging Face Hub.
@@ -484,6 +485,8 @@ def download_weights_from_hf(
485
weight files. Files matched by any of the patterns will be
486
downloaded.
487
revision (Optional[str]): The revision of the model.
488
+ subfolder (Optional[str]): The subfolder within the model repository
489
+ to download weights from.
490
ignore_patterns (Optional[Union[str, list[str]]]): The patterns to
491
filter out the weight files. Files matched by any of the patterns
492
will be ignored.
@@ -498,7 +501,11 @@ def download_weights_from_hf(
501
# so we only have to call snapshot_download once.
502
try:
503
fs = HfFileSystem()
501
- file_list = fs.ls(model_name_or_path, detail=False, revision=revision)
504
+ file_list = fs.ls(
505
+ os.path.join(model_name_or_path, subfolder or ""),
506
+ detail=False,
507
+ revision=revision,
508
+ )
509
510
# If downloading safetensors and an index file exists, use the
511
# specific file names from the index to avoid downloading
@@ -510,6 +517,7 @@ def download_weights_from_hf(
517
filename=SAFE_WEIGHTS_INDEX_NAME,
518
cache_dir=cache_dir,
519
revision=revision,
520
+ subfolder=subfolder,
521
)
522
with open(index_path) as f:
523
weight_map = json.load(f)["weight_map"]
@@ -570,6 +578,7 @@ def download_safetensors_index_file_from_hf(
578
model_name_or_path: str,
579
index_file: str,
580
cache_dir: str | None,
581
+ subfolder: str | None = None,
582
revision: str | None = None,
583
) -> None:
584
"""Download hf safetensors index file from Hugging Face Hub.
@@ -579,6 +588,8 @@ def download_safetensors_index_file_from_hf(
588
index_file (str): The safetensors index file name
589
cache_dir (Optional[str]): The cache directory to store the model
590
weights. If None, will use HF defaults.
591
+ subfolder (Optional[str]): The subfolder within the model repository
592
+ to download weights from.
593
revision (Optional[str]): The revision of the model.
594
"""
595
# Use file lock to prevent multiple processes from
@@ -591,6 +602,7 @@ def download_safetensors_index_file_from_hf(
602
filename=index_file,
603
cache_dir=cache_dir,
604
revision=revision,
605
+ subfolder=subfolder,
606
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
607
)
608
# If file not found on remote or locally, we should not fail since
vllm/model_executor/models/kimi_audio.py
+61
-111
@@ -3,15 +3,12 @@
3
4
"""Inference-only Kimi-Audio model compatible with HuggingFace weights."""
5
6
-import os
6
from collections.abc import Iterable, Mapping, Sequence
7
from typing import Any, ClassVar, Literal
8
9
import numpy as np
10
import torch
11
import torch.nn as nn
13
-from huggingface_hub import snapshot_download
14
-from safetensors import safe_open
12
from transformers import BatchFeature
13
from transformers import WhisperConfig as HFWhisperConfig
14
@@ -19,9 +16,8 @@ from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
16
from vllm.config.multimodal import BaseDummyOptions
17
from vllm.inputs.data import PromptType, TokensPrompt
18
from vllm.model_executor.layers.logits_processor import LogitsProcessor
22
-from vllm.model_executor.model_loader.weight_utils import (
23
- default_weight_loader,
24
-)
19
+from vllm.model_executor.model_loader import DefaultModelLoader
20
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
21
from vllm.model_executor.models.interfaces import (
22
SupportsMultiModal,
23
SupportsPP,
@@ -64,15 +60,6 @@ from vllm.v1.sample.metadata import SamplingMetadata
60
KIMIA_WHISPER_SUBFOLDER = "whisper-large-v3"
61
62
67
-def _get_whisper_local_path(repo_id: str):
68
- if os.path.exists(repo_id):
69
- repo_local_path = repo_id
70
- else:
71
- repo_local_path = snapshot_download(repo_id, local_files_only=True)
72
-
73
- return os.path.join(repo_local_path, KIMIA_WHISPER_SUBFOLDER)
74
-
75
-
63
def _get_feat_extract_output_lengths(input_lengths: torch.Tensor) -> torch.Tensor:
64
"""Compute output lengths after Whisper feature extraction.
65
@@ -93,7 +80,6 @@ class KimiAudioWhisperEncoder(WhisperEncoder):
80
# packed_modules_mapping for Q/K/V fusion during weight loading
81
packed_modules_mapping = {
82
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
96
- "kv_proj": ["k_proj", "v_proj"],
83
}
84
85
def __init__(
@@ -104,19 +90,49 @@ class KimiAudioWhisperEncoder(WhisperEncoder):
90
model_path = vllm_config.model_config.model
91
92
# Load WhisperConfig from the subfolder
107
- whisper_dir = _get_whisper_local_path(model_path)
108
- whisper_config = HFWhisperConfig.from_pretrained(whisper_dir)
109
-
110
- # Temporarily replace hf_config for WhisperEncoder.__init__()
111
- original_config = vllm_config.model_config.hf_config
112
- vllm_config.model_config.hf_config = whisper_config
93
+ whisper_config = HFWhisperConfig.from_pretrained(
94
+ model_path,
95
+ subfolder=KIMIA_WHISPER_SUBFOLDER,
96
+ )
97
98
super().__init__(
115
- vllm_config=vllm_config, prefix=prefix, init_in_fp32=init_in_fp32
99
+ vllm_config=vllm_config.with_hf_config(whisper_config),
100
+ prefix=prefix,
101
+ init_in_fp32=init_in_fp32,
102
)
103
118
- # Restore original config
119
- vllm_config.model_config.hf_config = original_config
104
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
105
+ stacked_params_mapping = [
106
+ # (param_name, shard_name, shard_id)
107
+ ("qkv_proj", "q_proj", "q"),
108
+ ("qkv_proj", "k_proj", "k"),
109
+ ("qkv_proj", "v_proj", "v"),
110
+ ]
111
+ params_dict = dict(self.named_parameters())
112
+ loaded_params: set[str] = set()
113
+ for name, loaded_weight in weights:
114
+ for param_name, weight_name, shard_id in stacked_params_mapping:
115
+ if weight_name not in name:
116
+ continue
117
+ name = name.replace(weight_name, param_name)
118
+ # Skip loading extra bias for GPTQ models.
119
+ if name.endswith(".bias") and name not in params_dict:
120
+ continue
121
+
122
+ param = params_dict[name]
123
+ weight_loader = param.weight_loader
124
+ weight_loader(param, loaded_weight, shard_id)
125
+ break
126
+ else:
127
+ # Skip loading extra bias for GPTQ models.
128
+ if name.endswith(".bias") and name not in params_dict:
129
+ continue
130
+
131
+ param = params_dict[name]
132
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
133
+ weight_loader(param, loaded_weight)
134
+ loaded_params.add(name)
135
+ return loaded_params
136
137
138
# -----------------------------------------------------------------------------
@@ -374,6 +390,8 @@ class KimiAudioForConditionalGeneration(
390
391
hf_to_vllm_mapper = WeightsMapper(
392
orig_to_new_prefix={
393
+ # audio tower
394
+ "model.encoder.": "audio_tower.",
395
# Audio projector (VQ-Adaptor)
396
"model.vq_adaptor.layers.0.": "multi_modal_projector.vq_adaptor_layers_0.",
397
"model.vq_adaptor.layers.3.": "multi_modal_projector.vq_adaptor_layers_3.",
@@ -384,7 +402,11 @@ class KimiAudioForConditionalGeneration(
402
"model.embed_tokens.": "language_model.model.embed_tokens.",
403
"model.norm.": "language_model.model.norm.",
404
"lm_head.": "language_model.lm_head.",
387
- }
405
+ },
406
+ orig_to_new_substr={
407
+ ".fc1.": ".mlp.fc1.",
408
+ ".fc2.": ".mlp.fc2.",
409
+ },
410
)
411
412
# Audio placeholder token sequence
@@ -401,6 +423,14 @@ class KimiAudioForConditionalGeneration(
423
self.multimodal_config = vllm_config.model_config.multimodal_config
424
self.model_path = vllm_config.model_config.model
425
426
+ self.secondary_weights = [
427
+ DefaultModelLoader.Source(
428
+ model_or_path=vllm_config.model_config.model,
429
+ subfolder="whisper-large-v3",
430
+ revision=None,
431
+ )
432
+ ]
433
+
434
self.audio_tower = KimiAudioWhisperEncoder(
435
vllm_config=vllm_config,
436
prefix=maybe_prefix(prefix, "audio_tower"),
@@ -577,99 +607,19 @@ class KimiAudioForConditionalGeneration(
607
"""Load weights, skipping MIMO layers (TTS-only) for ASR."""
608
# Filter out MIMO/TTS weights since we only do ASR (speech-to-text)
609
skipped_patterns = [
610
+ # Audio tower
611
+ "model.",
612
+ # MIMO/TTS
613
"mimo_layers.",
614
"mimo_output.",
615
"mimo_norm.",
583
- "audio_decoder.",
584
- ]
585
-
586
- # Filter weights
587
- filtered_weights = [
588
- (name, param)
589
- for name, param in weights
590
- if not any(pattern in name for pattern in skipped_patterns)
591
- ]
592
-
593
- # Separate main weights (non-Whisper) from Whisper weights
594
- main_weights = [
595
- (name, param)
596
- for name, param in filtered_weights
597
- if not name.startswith("audio_tower.")
616
]
617
618
# Load main model weights (LLM + projector) with mapper
601
- loader = AutoWeightsLoader(self)
602
- loaded = loader.load_weights(main_weights, mapper=self.hf_to_vllm_mapper)
603
-
604
- # Load Whisper encoder weights from subfolder
605
- whisper_dir = _get_whisper_local_path(self.model_path)
606
- whisper_path = os.path.join(whisper_dir, "model.safetensors")
607
- if os.path.exists(whisper_path):
608
- whisper_loaded = self._load_whisper_weights_from_file(whisper_path)
609
- loaded.update(whisper_loaded)
610
-
619
+ loader = AutoWeightsLoader(self, skip_prefixes=skipped_patterns)
620
+ loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
621
return loaded
622
613
- def _load_whisper_weights_from_file(self, whisper_path: str) -> set[str]:
614
- """Load Whisper encoder weights from safetensors file with transformations."""
615
- if not os.path.exists(whisper_path):
616
- return set()
617
-
618
- # Step 1: Load raw weights from safetensors file
619
- whisper_weights = []
620
- with safe_open(whisper_path, framework="pt") as f:
621
- for key in f.keys(): # noqa: SIM118
622
- if key.startswith("model.encoder.") and "embed_positions" not in key:
623
- new_key = key.replace("model.encoder.", "")
624
- whisper_weights.append((new_key, f.get_tensor(key)))
625
-
626
- # Step 2: Apply fc → mlp mapping using WeightsMapper
627
- fc_mapper = WeightsMapper(
628
- orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."}
629
- )
630
- whisper_mapped = list(fc_mapper.apply(whisper_weights))
631
-
632
- # Step 3: Apply Q/K/V fusion manually
633
- stacked_params_mapping = [
634
- (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
635
- (".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
636
- (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
637
- ]
638
-
639
- params_dict = dict(self.audio_tower.named_parameters())
640
- whisper_loaded: set[str] = set()
641
-
642
- for name, loaded_weight in whisper_mapped:
643
- fused = False
644
- for param_name, weight_name, shard_id in stacked_params_mapping:
645
- if weight_name not in name:
646
- continue
647
- fused_name = name.replace(weight_name, param_name)
648
- if fused_name not in params_dict:
649
- continue
650
-
651
- param = params_dict[fused_name]
652
- param.weight_loader(param, loaded_weight, shard_id)
653
- whisper_loaded.add(f"audio_tower.{fused_name}")
654
- fused = True
655
- break
656
-
657
- if not fused:
658
- if name.endswith(".bias") and name not in params_dict:
659
- continue
660
- if name not in params_dict:
661
- continue
662
-
663
- param = params_dict[name]
664
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
665
- weight_loader(param, loaded_weight)
666
- whisper_loaded.add(f"audio_tower.{name}")
667
-
668
- # Add embed_positions which is initialized randomly
669
- whisper_loaded.add("audio_tower.embed_positions.weight")
670
-
671
- return whisper_loaded
672
-
623
@classmethod
624
def get_speech_to_text_config(
625
cls, model_config: ModelConfig, task_type: str