@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 42fadebecb
Fork 0

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, 23:00 UTC

[Model] Add support for moonshotai/Kimi-Audio-7B-Instruct (#36127)

Signed-off-by: tunglinwood <tunglinwood@gmail.com> Signed-off-by: tunglinwood <tomwu.tunglin@gmail.com> Signed-off-by: tunglinwood <113751333+tunglinwood@users.noreply.github.com>

tunglinwood committed Mar 11, 2026 at 12:24 UTC 42fadebecb79290ad722f33f3094de23b121f33d
14 files changed +1446 -29
docs/models/supported_models.md
+2 -1
@@ -713,8 +713,9 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
713 | `KananaVForConditionalGeneration` | Kanana-V | T + I<sup>+</sup> | `kakaocorp/kanana-1.5-v-3b-instruct`, etc. | | ✅︎ |
714 | `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ |
715 | `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ |
716 -| `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
716 +| `KimiAudioForConditionalGeneration` | Kimi-Audio | T + A<sup>+</sup> | `moonshotai/Kimi-Audio-7B-Instruct` | | ✅︎ |
717 | `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I<sup>+</sup> | `moonshotai/Kimi-K2.5` | | ✅︎ |
718 +| `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
719 | `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I<sup>+</sup> | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ |
720 | `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I<sup>+</sup> | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |
721 | `Llama4ForConditionalGeneration` | Llama 4 | T + I<sup>+</sup> | `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc. | ✅︎ | ✅︎ |
examples/offline_inference/audio_language.py
+29
@@ -201,6 +201,34 @@ def run_granite_speech(question: str, audio_count: int) -> ModelRequestData:
201 )
202
203
204 +# Kimi-Audio-7B-Instruct
205 +def run_kimi_audio(question: str, audio_count: int) -> ModelRequestData:
206 + """Kimi-Audio-7B-Instruct for audio transcription and understanding."""
207 + model_name = "moonshotai/Kimi-Audio-7B-Instruct"
208 +
209 + engine_args = EngineArgs(
210 + model=model_name,
211 + trust_remote_code=True,
212 + max_model_len=4096,
213 + max_num_seqs=2,
214 + limit_mm_per_prompt={"audio": audio_count},
215 + )
216 +
217 + # Kimi-Audio uses <|im_kimia_text_blank|> as placeholder for audio features
218 + audio_placeholder = "<|im_kimia_text_blank|>" * audio_count
219 + # Default prompt for transcription
220 + if not question:
221 + question = "Please transcribe the audio"
222 + prompt = f"{audio_placeholder}{question}"
223 +
224 + # Stop at EOS token (151644) to prevent repetition
225 + return ModelRequestData(
226 + engine_args=engine_args,
227 + prompt=prompt,
228 + stop_token_ids=[151644],
229 + )
230 +
231 +
232 # MiDashengLM
233 def run_midashenglm(question: str, audio_count: int):
234 model_name = "mispeech/midashenglm-7b"
@@ -485,6 +513,7 @@ model_example_map = {
513 "glmasr": run_glmasr,
514 "funaudiochat": run_funaudiochat,
515 "granite_speech": run_granite_speech,
516 + "kimi_audio": run_kimi_audio,
517 "midashenglm": run_midashenglm,
518 "minicpmo": run_minicpmo,
519 "phi4_mm": run_phi4mm,
tests/models/multimodal/processing/test_common.py
+11 -7
@@ -198,13 +198,17 @@ def get_text_token_prompts(
198 mm_counts,
199 mm_options={},
200 )
201 - assert isinstance(inputs.prompt, str)
202 -
203 - text_prompt = inputs.prompt
204 - token_prompt = tokenizer.encode(
205 - text_prompt,
206 - add_special_tokens=_ADD_SPECIAL_TOKENS_OVERRIDES.get(model_type, True),
207 - )
201 + # Some models (e.g., Kimi-Audio) return token IDs directly instead of str
202 + if isinstance(inputs.prompt, list):
203 + text_prompt = None
204 + token_prompt = inputs.prompt
205 + else:
206 + assert isinstance(inputs.prompt, str)
207 + text_prompt = inputs.prompt
208 + token_prompt = tokenizer.encode(
209 + text_prompt,
210 + add_special_tokens=_ADD_SPECIAL_TOKENS_OVERRIDES.get(model_type, True),
211 + )
212
213 return text_prompt, token_prompt
214
tests/models/registry.py
+9 -4
@@ -857,6 +857,15 @@ _MULTIMODAL_EXAMPLE_MODELS = {
857 "Kwai-Keye/Keye-VL-1_5-8B",
858 trust_remote_code=True,
859 ),
860 + "MoonshotKimiaForCausalLM": _HfExamplesInfo(
861 + "moonshotai/Kimi-Audio-7B-Instruct",
862 + tokenizer_mode="kimi_audio",
863 + trust_remote_code=True,
864 + ),
865 + "KimiK25ForConditionalGeneration": _HfExamplesInfo(
866 + "moonshotai/Kimi-K2.5",
867 + trust_remote_code=True,
868 + ),
869 "KimiVLForConditionalGeneration": _HfExamplesInfo(
870 "moonshotai/Kimi-VL-A3B-Instruct",
871 extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"},
@@ -870,10 +879,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
879 )
880 },
881 ),
873 - "KimiK25ForConditionalGeneration": _HfExamplesInfo(
874 - "moonshotai/Kimi-K2.5",
875 - trust_remote_code=True,
876 - ),
882 "LightOnOCRForConditionalGeneration": _HfExamplesInfo(
883 "lightonai/LightOnOCR-1B-1025"
884 ),
tests/models/test_initialization.py
+6
@@ -103,6 +103,12 @@ def can_initialize(
103 "pickle error when loading `transformers.models.auto.CONFIG_MAPPING`"
104 )
105
106 + if model_arch == "MoonshotKimiaForCausalLM":
107 + pytest.skip(
108 + "Kimi-Audio requires SpeechToTextConfig "
109 + "which is not configured in test environment"
110 + )
111 +
112 if model_arch in ["DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM"]:
113 from vllm.platforms import current_platform
114
vllm/model_executor/models/kimi_audio.py new
+725
@@ -0,0 +1,725 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Inference-only Kimi-Audio model compatible with HuggingFace weights."""
5 +
6 +import os
7 +from collections.abc import Iterable, Mapping, Sequence
8 +from typing import Any, ClassVar, Literal
9 +
10 +import numpy as np
11 +import torch
12 +import torch.nn as nn
13 +from safetensors import safe_open
14 +from transformers import BatchFeature
15 +from transformers import WhisperConfig as HFWhisperConfig
16 +
17 +from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
18 +from vllm.inputs.data import PromptType, TokensPrompt
19 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
20 +from vllm.model_executor.model_loader.weight_utils import (
21 + default_weight_loader,
22 +)
23 +from vllm.model_executor.models.interfaces import (
24 + SupportsMultiModal,
25 + SupportsPP,
26 + SupportsTranscription,
27 +)
28 +from vllm.model_executor.models.utils import (
29 + AutoWeightsLoader,
30 + WeightsMapper,
31 + init_vllm_registered_model,
32 + maybe_prefix,
33 +)
34 +from vllm.model_executor.models.whisper import WhisperEncoder
35 +from vllm.model_executor.models.whisper_utils import ISO639_1_SUPPORTED_LANGS
36 +from vllm.multimodal import MULTIMODAL_REGISTRY
37 +from vllm.multimodal.inputs import MultiModalFieldConfig
38 +from vllm.multimodal.parse import (
39 + AudioItem,
40 + DictEmbeddingItems,
41 + ModalityData,
42 + ModalityDataItems,
43 + MultiModalDataParser,
44 +)
45 +from vllm.multimodal.processing import (
46 + BaseDummyInputsBuilder,
47 + BaseProcessingInfo,
48 + PromptReplacement,
49 +)
50 +from vllm.multimodal.processing.processor import BaseMultiModalProcessor
51 +from vllm.sequence import IntermediateTensors
52 +from vllm.tokenizers import cached_get_tokenizer
53 +from vllm.tokenizers.kimi_audio import KimiAudioTokenizer
54 +from vllm.transformers_utils.processor import cached_feature_extractor_from_config
55 +from vllm.transformers_utils.processors.kimi_audio import KimiAudioProcessor
56 +from vllm.v1.sample.metadata import SamplingMetadata
57 +
58 +# Kimi-Audio constants
59 +KIMIA_WHISPER_SUBFOLDER = "whisper-large-v3"
60 +
61 +
62 +def _get_feat_extract_output_lengths(input_lengths: torch.Tensor) -> torch.Tensor:
63 + """Compute output lengths after Whisper feature extraction.
64 +
65 + Whisper processes audio through multiple conv layers with stride=2,
66 + producing 13 output features per 100 input samples.
67 + """
68 + input_lengths_leave = input_lengths % 100
69 + feat_lengths = (input_lengths_leave - 1) // 2 + 1
70 + output_lengths = (
71 + ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
72 + )
73 + return output_lengths
74 +
75 +
76 +class KimiAudioWhisperEncoder(WhisperEncoder):
77 + """WhisperEncoder for Kimi-Audio with packed_modules_mapping."""
78 +
79 + # packed_modules_mapping for Q/K/V fusion during weight loading
80 + packed_modules_mapping = {
81 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
82 + "kv_proj": ["k_proj", "v_proj"],
83 + }
84 +
85 + def __init__(
86 + self, *, vllm_config: VllmConfig, prefix: str = "", init_in_fp32: bool = False
87 + ):
88 + # Load Whisper config from subfolder (authoritative source)
89 + # Kimi-Audio stores Whisper config in whisper-large-v3/config.json
90 + model_path = vllm_config.model_config.model
91 + whisper_config_path = os.path.join(model_path, KIMIA_WHISPER_SUBFOLDER)
92 +
93 + # Load WhisperConfig from the subfolder
94 + whisper_config = HFWhisperConfig.from_pretrained(whisper_config_path)
95 +
96 + # Temporarily replace hf_config for WhisperEncoder.__init__()
97 + original_config = vllm_config.model_config.hf_config
98 + vllm_config.model_config.hf_config = whisper_config
99 +
100 + super().__init__(
101 + vllm_config=vllm_config, prefix=prefix, init_in_fp32=init_in_fp32
102 + )
103 +
104 + # Restore original config
105 + vllm_config.model_config.hf_config = original_config
106 +
107 +
108 +# -----------------------------------------------------------------------------
109 +# Processing Info, Dummy Inputs, and MultiModal Processor
110 +# (Following Qwen3ASR pattern - same file as model)
111 +# -----------------------------------------------------------------------------
112 +
113 +
114 +class KimiAudioProcessingInfo(BaseProcessingInfo):
115 + """Processing info for vLLM registry."""
116 +
117 + def get_hf_config(self):
118 + return self.ctx.model_config.hf_config
119 +
120 + def get_hf_processor(self, **kwargs: object) -> KimiAudioProcessor:
121 + """Get KimiAudioProcessor with feature extractor and tokenizer."""
122 + # Use vLLM's cached loader for feature extractor
123 + feature_extractor = cached_feature_extractor_from_config(
124 + self.ctx.model_config,
125 + subfolder=KIMIA_WHISPER_SUBFOLDER,
126 + )
127 +
128 + # Use vLLM's standard tokenizer loading (respects tokenizer_mode)
129 + tokenizer = self.get_tokenizer()
130 +
131 + # Construct processor directly
132 + return KimiAudioProcessor(
133 + feature_extractor=feature_extractor,
134 + tokenizer=tokenizer,
135 + )
136 +
137 + def get_feature_extractor(self, **kwargs: object):
138 + """Get feature extractor using vLLM's cached loader."""
139 + return cached_feature_extractor_from_config(
140 + self.ctx.model_config, subfolder=KIMIA_WHISPER_SUBFOLDER
141 + )
142 +
143 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
144 + return {"audio": 1}
145 +
146 + def get_data_parser(self) -> "KimiAudioMultiModalDataParser":
147 + """Get data parser for audio inputs."""
148 + return KimiAudioMultiModalDataParser(
149 + expected_hidden_size=self._get_expected_hidden_size(),
150 + )
151 +
152 +
153 +class KimiAudioDummyInputsBuilder(BaseDummyInputsBuilder[KimiAudioProcessingInfo]):
154 + """Dummy inputs builder for vLLM registry."""
155 +
156 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> list[int]:
157 + """Return dummy text as token IDs directly."""
158 + num_audios = mm_counts.get("audio", 0)
159 + if num_audios == 0:
160 + return [198] # "Transcribe" tokenized
161 + # Return as token IDs directly to avoid tokenizer issues
162 + return [
163 + KimiAudioProcessor.KIMIA_MEDIA_BEGIN,
164 + KimiAudioProcessor.KIMIA_TEXT_BLANK,
165 + KimiAudioProcessor.KIMIA_MEDIA_END,
166 + ] * num_audios
167 +
168 + def get_dummy_mm_data(
169 + self,
170 + seq_len: int,
171 + mm_counts: Mapping[str, int],
172 + mm_options: Mapping[str, Any] | None = None,
173 + ) -> dict[str, Any]:
174 + num_audios = mm_counts.get("audio", 0)
175 + if num_audios == 0:
176 + return {}
177 +
178 + feature_extractor = self.info.get_feature_extractor()
179 + target_audio_length = (
180 + min(feature_extractor.chunk_length, 30) * feature_extractor.sampling_rate
181 + )
182 +
183 + return {
184 + "audio": self._get_dummy_audios(
185 + length=target_audio_length, num_audios=num_audios
186 + ),
187 + }
188 +
189 +
190 +# Field config for Kimi-Audio multimodal data
191 +_KIMIAUDIO_FIELD_CONFIG = {
192 + "whisper_input_features": MultiModalFieldConfig.batched("audio"),
193 + "feature_attention_mask": MultiModalFieldConfig.batched("audio"),
194 +}
195 +
196 +
197 +class KimiAudioMultiModalDataParser(MultiModalDataParser):
198 + """Custom data parser for Kimi-Audio multimodal data."""
199 +
200 + def __init__(self, **kwargs):
201 + # Whisper expects 16kHz audio
202 + super().__init__(target_sr=16000, **kwargs)
203 +
204 + def _parse_audio_data(
205 + self,
206 + data: dict[str, torch.Tensor] | ModalityData[AudioItem],
207 + ) -> ModalityDataItems[Any, Any] | None:
208 + if isinstance(data, dict):
209 + return DictEmbeddingItems(
210 + data,
211 + modality="audio",
212 + required_fields={"whisper_input_features", "feature_attention_mask"},
213 + fields_factory=lambda hf_inputs: _KIMIAUDIO_FIELD_CONFIG,
214 + )
215 +
216 + return super()._parse_audio_data(data)
217 +
218 +
219 +class KimiAudioMultiModalProcessor(BaseMultiModalProcessor[KimiAudioProcessingInfo]):
220 + """vLLM multi-modal processor wrapper for Kimi-Audio."""
221 +
222 + def _call_hf_processor(
223 + self,
224 + prompt: str,
225 + mm_data: Mapping[str, object],
226 + mm_kwargs: Mapping[str, object],
227 + tok_kwargs: Mapping[str, object],
228 + ) -> BatchFeature:
229 + """Call the HuggingFace processor."""
230 + # Convert mm_data format: {'audios': [...]} -> {'audio': ...}
231 + mm_data = dict(mm_data)
232 + audios = mm_data.pop("audios", [])
233 +
234 + # Convert audio format: [(array, sr), ...] -> [array, ...]
235 + # KimiAudioProcessor expects raw numpy arrays
236 + if audios:
237 + audio_arrays = []
238 + for aud in audios:
239 + if isinstance(aud, (tuple, list)) and len(aud) == 2:
240 + # Format: (audio_array, sampling_rate)
241 + audio_arrays.append(aud[0])
242 + elif isinstance(aud, np.ndarray):
243 + audio_arrays.append(aud)
244 + else:
245 + audio_arrays.append(aud)
246 + mm_data["audio"] = audio_arrays
247 +
248 + # Use the context's call_hf_processor for proper handling
249 + return self.info.ctx.call_hf_processor(
250 + self.info.get_hf_processor(**mm_kwargs),
251 + dict(text=prompt, **mm_data),
252 + dict(**mm_kwargs, **tok_kwargs),
253 + )
254 +
255 + def _get_mm_fields_config(
256 + self,
257 + hf_inputs: BatchFeature,
258 + hf_processor_mm_kwargs: Mapping[str, object],
259 + ) -> Mapping[str, Any]:
260 + """Get multi-modal field configuration."""
261 + return _KIMIAUDIO_FIELD_CONFIG
262 +
263 + def _get_prompt_updates(
264 + self,
265 + mm_items,
266 + hf_processor_mm_kwargs,
267 + out_mm_kwargs,
268 + ) -> Sequence[PromptReplacement]:
269 + """Get prompt updates for audio tokens."""
270 + # Get audio feature lengths from processed output
271 + out_mm_data = out_mm_kwargs.get_data()
272 + feature_attention_mask = out_mm_data.get("feature_attention_mask")
273 +
274 + if feature_attention_mask is not None:
275 + audio_output_lens = _get_feat_extract_output_lengths(
276 + feature_attention_mask.sum(-1)
277 + )
278 + audio_output_lengths = audio_output_lens.tolist()
279 + else:
280 + audio_output_lengths = []
281 +
282 + def get_replacement_kimiaudio(item_idx: int):
283 + num_features = (
284 + audio_output_lengths[item_idx]
285 + if item_idx < len(audio_output_lengths)
286 + else 376
287 + )
288 + if num_features == 0:
289 + num_features = 376 # Default Kimi-Audio sequence length
290 + # Return the placeholder token ID repeated num_features times
291 + return [KimiAudioProcessor.KIMIA_TEXT_BLANK] * num_features
292 +
293 + # Use the token ID as target (as a list)
294 + return [
295 + PromptReplacement(
296 + modality="audio",
297 + target=[KimiAudioProcessor.KIMIA_TEXT_BLANK],
298 + replacement=get_replacement_kimiaudio,
299 + ),
300 + ]
301 +
302 +
303 +# -----------------------------------------------------------------------------
304 +# Model Definition
305 +# -----------------------------------------------------------------------------
306 +
307 +
308 +class KimiAudioMultiModalProjector(nn.Module):
309 + """Projects Whisper features to LLM embedding space.
310 +
311 + Kimi-Audio VQ-Adaptor architecture:
312 + Custom Whisper (5120) → Linear[5120→3584] → Linear[3584→3584] → LayerNorm
313 + """
314 +
315 + def __init__(
316 + self,
317 + whisper_dim: int = 5120, # Kimi-Audio custom Whisper encoder dim
318 + llm_dim: int = 3584,
319 + prefix: str = "",
320 + ):
321 + super().__init__()
322 + self.whisper_dim = whisper_dim
323 + self.llm_dim = llm_dim
324 +
325 + # VQ-Adaptor layers (exact checkpoint structure)
326 + # layers.0: Linear[5120 → 3584]
327 + self.vq_adaptor_layers_0 = nn.Linear(whisper_dim, llm_dim)
328 + # layers.3: Linear[3584 → 3584]
329 + self.vq_adaptor_layers_3 = nn.Linear(llm_dim, llm_dim)
330 + # layers.4: LayerNorm[3584]
331 + self.vq_adaptor_layers_4 = nn.LayerNorm(llm_dim)
332 +
333 + def forward(self, audio_features: torch.Tensor) -> torch.Tensor:
334 + # Project: [B, T, 5120] → [B, T, 3584]
335 + hidden = self.vq_adaptor_layers_0(audio_features)
336 + hidden = torch.nn.functional.gelu(hidden)
337 + hidden = self.vq_adaptor_layers_3(hidden)
338 + hidden = self.vq_adaptor_layers_4(hidden)
339 + return hidden
340 +
341 +
342 +@MULTIMODAL_REGISTRY.register_processor(
343 + KimiAudioMultiModalProcessor,
344 + info=KimiAudioProcessingInfo,
345 + dummy_inputs=KimiAudioDummyInputsBuilder,
346 +)
347 +class KimiAudioForConditionalGeneration(
348 + nn.Module,
349 + SupportsMultiModal,
350 + SupportsPP,
351 + SupportsTranscription,
352 +):
353 + """Kimi-Audio model for ASR transcription."""
354 +
355 + # Kimi-Audio supports a subset of Whisper's supported languages
356 + supported_languages: ClassVar[Mapping[str, str]] = {
357 + k: ISO639_1_SUPPORTED_LANGS[k]
358 + for k in ["zh", "en", "ja", "ko", "de", "fr", "es", "it", "pt", "ru", "ar"]
359 + }
360 + supports_transcription: ClassVar[Literal[True]] = True
361 +
362 + hf_to_vllm_mapper = WeightsMapper(
363 + orig_to_new_prefix={
364 + # Audio projector (VQ-Adaptor)
365 + "model.vq_adaptor.layers.0.": "multi_modal_projector.vq_adaptor_layers_0.",
366 + "model.vq_adaptor.layers.3.": "multi_modal_projector.vq_adaptor_layers_3.",
367 + "model.vq_adaptor.layers.4.": "multi_modal_projector.vq_adaptor_layers_4.",
368 + # Language model
369 + "model.layers.": "language_model.model.layers.",
370 + # Embeddings and output
371 + "model.embed_tokens.": "language_model.model.embed_tokens.",
372 + "model.norm.": "language_model.model.norm.",
373 + "lm_head.": "language_model.lm_head.",
374 + }
375 + )
376 +
377 + # Audio placeholder token sequence
378 + AUDIO_PLACEHOLDER = "<|im_media_begin|><|im_kimia_text_blank|><|im_media_end|>"
379 +
380 + @classmethod
381 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
382 + return cls.AUDIO_PLACEHOLDER if modality.startswith("audio") else None
383 +
384 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
385 + super().__init__()
386 + self.config = vllm_config.model_config.hf_config
387 + self.quant_config = vllm_config.quant_config
388 + self.multimodal_config = vllm_config.model_config.multimodal_config
389 + self.model_path = vllm_config.model_config.model
390 +
391 + self.audio_tower = KimiAudioWhisperEncoder(
392 + vllm_config=vllm_config,
393 + prefix=maybe_prefix(prefix, "audio_tower"),
394 + )
395 +
396 + self.multi_modal_projector = KimiAudioMultiModalProjector(
397 + whisper_dim=getattr(self.config, "kimia_adaptor_input_dim", 5120),
398 + llm_dim=self.config.hidden_size,
399 + prefix=maybe_prefix(prefix, "multi_modal_projector"),
400 + )
401 +
402 + self.language_model = init_vllm_registered_model(
403 + vllm_config=vllm_config.with_hf_config(
404 + self.config, architectures=["Qwen2ForCausalLM"]
405 + ),
406 + prefix=maybe_prefix(prefix, "language_model"),
407 + )
408 +
409 + self.logits_processor = LogitsProcessor(
410 + self.config.vocab_size,
411 + self.config.vocab_size,
412 + )
413 +
414 + self.make_empty_intermediate_tensors = (
415 + self.language_model.make_empty_intermediate_tensors
416 + )
417 +
418 + def _parse_and_validate_audio_input(
419 + self, **kwargs: object
420 + ) -> dict[str, torch.Tensor] | None:
421 + whisper_input_features = kwargs.pop("whisper_input_features", None)
422 + if whisper_input_features is None:
423 + return None
424 +
425 + return {"whisper_input_features": whisper_input_features}
426 +
427 + def _process_audio_input(
428 + self, audio_input: dict[str, torch.Tensor]
429 + ) -> torch.Tensor:
430 + input_features = audio_input["whisper_input_features"]
431 +
432 + # KimiAudioWhisperEncoder expects list of tensors
433 + if input_features.dim() == 3:
434 + input_features = input_features.unbind(dim=0)
435 +
436 + # Run through Whisper encoder
437 + audio_features = self.audio_tower(input_features)
438 +
439 + # Reshape for 4x downsampling (Whisper outputs at 50Hz, need 12.5Hz)
440 + B, T, D = audio_features.shape
441 + if T % 4 != 0:
442 + pad_len = 4 - (T % 4)
443 + audio_features = torch.nn.functional.pad(audio_features, (0, 0, 0, pad_len))
444 + T = audio_features.shape[1] # Update T after padding
445 +
446 + audio_features = audio_features.reshape(B, T // 4, D * 4)
447 +
448 + # Project to LLM dimension
449 + audio_embeds = self.multi_modal_projector(audio_features)
450 + return audio_embeds
451 +
452 + def embed_multimodal(self, **kwargs: object) -> list[torch.Tensor] | None:
453 + audio_input = self._parse_and_validate_audio_input(**kwargs)
454 + if audio_input is None:
455 + return []
456 +
457 + audio_embeds = self._process_audio_input(audio_input)
458 +
459 + # audio_embeds shape: [batch_size, seq_len, hidden_dim]
460 + # Return as list of 2D tensors, one per batch item
461 + if audio_embeds.dim() == 3:
462 + # Unbind batch dimension: [B, T, D] -> list of B tensors [T, D]
463 + return list(audio_embeds.unbind(dim=0))
464 + else:
465 + # Single sample: [T, D] -> wrap in list
466 + return [audio_embeds]
467 +
468 + def embed_input_ids(
469 + self,
470 + input_ids: torch.Tensor,
471 + multimodal_embeddings: tuple[torch.Tensor, ...] | None = None,
472 + *,
473 + is_multimodal: torch.Tensor | None = None,
474 + handle_oov_mm_token: bool = False,
475 + ) -> torch.Tensor:
476 + """Embed input IDs and fuse with audio embeddings.
477 +
478 + Kimi-Audio fusion: inputs_embeds = (text_emb + audio_emb) × √2
479 +
480 + For PP compatibility, we use the is_multimodal mask from vLLM engine
481 + which is correctly computed per pipeline stage.
482 + """
483 + # Get text embeddings
484 + inputs_embeds = self.language_model.model.embed_tokens(input_ids)
485 +
486 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
487 + return inputs_embeds
488 +
489 + # is_multimodal must be provided for PP to work correctly
490 + if is_multimodal is None or not is_multimodal.any():
491 + return inputs_embeds
492 +
493 + # multimodal_embeddings[0] contains audio embeddings
494 + audio_embeds = multimodal_embeddings[0]
495 +
496 + # Handle different tensor structures
497 + if isinstance(audio_embeds, (list, tuple)):
498 + audio_embeds = torch.cat(audio_embeds, dim=0)
499 + elif audio_embeds.dim() == 3:
500 + audio_embeds = audio_embeds.reshape(-1, audio_embeds.shape[-1])
501 +
502 + # In PP, audio_embeds count should match is_multimodal.sum()
503 + # For now, use embeddings sequentially
504 + # (works for non-PP, PP needs vLLM infra fix)
505 + num_mm_tokens = is_multimodal.sum().item()
506 + num_audio_embeds = audio_embeds.shape[0]
507 +
508 + # Use the minimum of available embeddings and positions
509 + # This ensures we don't access out-of-bounds
510 + num_to_use = min(num_audio_embeds, num_mm_tokens)
511 +
512 + # Get positions for the tokens we'll actually process
513 + mm_positions = is_multimodal.nonzero(as_tuple=True)[0]
514 + actual_mm_mask = torch.zeros_like(is_multimodal)
515 + actual_mm_mask[mm_positions[:num_to_use]] = True
516 +
517 + # Use corresponding embeddings
518 + used_audio_embeds = audio_embeds[:num_to_use]
519 +
520 + # Save text embeddings at multimodal positions
521 + text_at_mm_positions = inputs_embeds[actual_mm_mask].clone()
522 +
523 + # Replace text with audio at multimodal positions
524 + inputs_embeds[actual_mm_mask] = used_audio_embeds.to(dtype=inputs_embeds.dtype)
525 +
526 + # Apply Kimi-Audio's unique fusion formula: (text + audio) × √2
527 + inputs_embeds[actual_mm_mask] = (
528 + inputs_embeds[actual_mm_mask] + text_at_mm_positions
529 + ) * (2**0.5)
530 +
531 + return inputs_embeds
532 +
533 + def forward(
534 + self,
535 + input_ids: torch.Tensor | None,
536 + positions: torch.Tensor,
537 + intermediate_tensors: IntermediateTensors | None = None,
538 + inputs_embeds: torch.Tensor | None = None,
539 + **kwargs: object,
540 + ) -> torch.Tensor | IntermediateTensors:
541 + if intermediate_tensors is not None:
542 + inputs_embeds = None
543 +
544 + hidden_states = self.language_model.model(
545 + input_ids,
546 + positions,
547 + intermediate_tensors,
548 + inputs_embeds=inputs_embeds,
549 + )
550 +
551 + return hidden_states
552 +
553 + def compute_logits(
554 + self,
555 + hidden_states: torch.Tensor,
556 + sampling_metadata: SamplingMetadata | None = None,
557 + ) -> torch.Tensor | None:
558 + logits = self.logits_processor(
559 + self.language_model.lm_head, hidden_states, sampling_metadata
560 + )
561 + return logits
562 +
563 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
564 + """Load weights, skipping MIMO layers (TTS-only) for ASR."""
565 + # Filter out MIMO/TTS weights since we only do ASR (speech-to-text)
566 + skipped_patterns = [
567 + "mimo_layers.",
568 + "mimo_output.",
569 + "mimo_norm.",
570 + "audio_decoder.",
571 + ]
572 +
573 + # Filter weights
574 + filtered_weights = [
575 + (name, param)
576 + for name, param in weights
577 + if not any(pattern in name for pattern in skipped_patterns)
578 + ]
579 +
580 + # Separate main weights (non-Whisper) from Whisper weights
581 + main_weights = [
582 + (name, param)
583 + for name, param in filtered_weights
584 + if not name.startswith("audio_tower.")
585 + ]
586 +
587 + # Load main model weights (LLM + projector) with mapper
588 + loader = AutoWeightsLoader(self)
589 + loaded = loader.load_weights(main_weights, mapper=self.hf_to_vllm_mapper)
590 +
591 + # Load Whisper encoder weights from subfolder
592 + whisper_path = os.path.join(
593 + self.model_path, f"{KIMIA_WHISPER_SUBFOLDER}/model.safetensors"
594 + )
595 + if os.path.exists(whisper_path):
596 + whisper_loaded = self._load_whisper_weights_from_file(whisper_path)
597 + loaded.update(whisper_loaded)
598 +
599 + return loaded
600 +
601 + def _load_whisper_weights_from_file(self, whisper_path: str) -> set[str]:
602 + """Load Whisper encoder weights from safetensors file with transformations."""
603 + if not os.path.exists(whisper_path):
604 + return set()
605 +
606 + # Step 1: Load raw weights from safetensors file
607 + whisper_weights = []
608 + with safe_open(whisper_path, framework="pt") as f:
609 + for key in f.keys(): # noqa: SIM118
610 + if key.startswith("model.encoder.") and "embed_positions" not in key:
611 + new_key = key.replace("model.encoder.", "")
612 + whisper_weights.append((new_key, f.get_tensor(key)))
613 +
614 + # Step 2: Apply fc → mlp mapping using WeightsMapper
615 + fc_mapper = WeightsMapper(
616 + orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."}
617 + )
618 + whisper_mapped = list(fc_mapper.apply(whisper_weights))
619 +
620 + # Step 3: Apply Q/K/V fusion manually
621 + stacked_params_mapping = [
622 + (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
623 + (".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
624 + (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
625 + ]
626 +
627 + params_dict = dict(self.audio_tower.named_parameters())
628 + whisper_loaded: set[str] = set()
629 +
630 + for name, loaded_weight in whisper_mapped:
631 + fused = False
632 + for param_name, weight_name, shard_id in stacked_params_mapping:
633 + if weight_name not in name:
634 + continue
635 + fused_name = name.replace(weight_name, param_name)
636 + if fused_name not in params_dict:
637 + continue
638 +
639 + param = params_dict[fused_name]
640 + param.weight_loader(param, loaded_weight, shard_id)
641 + whisper_loaded.add(f"audio_tower.{fused_name}")
642 + fused = True
643 + break
644 +
645 + if not fused:
646 + if name.endswith(".bias") and name not in params_dict:
647 + continue
648 + if name not in params_dict:
649 + continue
650 +
651 + param = params_dict[name]
652 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
653 + weight_loader(param, loaded_weight)
654 + whisper_loaded.add(f"audio_tower.{name}")
655 +
656 + # Add embed_positions which is initialized randomly
657 + whisper_loaded.add("audio_tower.embed_positions.weight")
658 +
659 + return whisper_loaded
660 +
661 + @classmethod
662 + def get_speech_to_text_config(
663 + cls, model_config: ModelConfig, task_type: str
664 + ) -> SpeechToTextConfig:
665 + """Get speech-to-text config with custom processor."""
666 + # Load feature extractor for config values
667 + feature_extractor = cached_feature_extractor_from_config(
668 + model_config,
669 + subfolder=KIMIA_WHISPER_SUBFOLDER,
670 + )
671 +
672 + return SpeechToTextConfig(
673 + max_audio_clip_s=feature_extractor.chunk_length,
674 + sample_rate=feature_extractor.sampling_rate,
675 + )
676 +
677 + @classmethod
678 + def get_generation_prompt(
679 + cls,
680 + audio: np.ndarray,
681 + model_config: ModelConfig,
682 + stt_config: SpeechToTextConfig,
683 + language: str | None,
684 + task_type: Literal["transcribe", "translate"],
685 + request_prompt: str,
686 + to_language: str | None,
687 + ) -> PromptType:
688 + tokenizer = cached_get_tokenizer(
689 + model_config.tokenizer,
690 + tokenizer_cls=KimiAudioTokenizer,
691 + tokenizer_mode=model_config.tokenizer_mode,
692 + revision=model_config.tokenizer_revision,
693 + trust_remote_code=model_config.trust_remote_code,
694 + )
695 +
696 + if task_type not in ("transcribe", "translate"):
697 + raise ValueError(
698 + f"Unsupported task_type '{task_type}'. "
699 + "Supported task types are 'transcribe' and 'translate'."
700 + )
701 +
702 + # Incorporate request_prompt as context/instruction if provided
703 + user_content = (
704 + f"{request_prompt}\n{cls.AUDIO_PLACEHOLDER}"
705 + if request_prompt
706 + else cls.AUDIO_PLACEHOLDER
707 + )
708 +
709 + prompt = (
710 + f"<|im_kimia_user_msg_start|>{user_content}"
711 + f"<|im_msg_end|><|im_kimia_assistant_msg_start|>"
712 + )
713 +
714 + prompt_token_ids = tokenizer.encode(prompt)
715 +
716 + return TokensPrompt(
717 + prompt_token_ids=prompt_token_ids,
718 + multi_modal_data={"audio": audio},
719 + )
720 +
721 + @classmethod
722 + def post_process_output(cls, text: str) -> str:
723 + if not text:
724 + return ""
725 + return text.strip()
vllm/model_executor/models/registry.py
+1
@@ -421,6 +421,7 @@ _MULTIMODAL_MODELS = {
421 "RForConditionalGeneration": ("rvl", "RForConditionalGeneration"),
422 "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"), # noqa: E501
423 "KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"), # noqa: E501
424 + "MoonshotKimiaForCausalLM": ("kimi_audio", "KimiAudioForConditionalGeneration"), # noqa: E501
425 "LightOnOCRForConditionalGeneration": (
426 "lightonocr",
427 "LightOnOCRForConditionalGeneration",
vllm/renderers/kimi_audio.py new
+49
@@ -0,0 +1,49 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any, cast
4 +
5 +from vllm.config import VllmConfig
6 +from vllm.tokenizers.kimi_audio import KimiAudioTokenizer
7 +from vllm.tokenizers.registry import get_tokenizer
8 +
9 +from .hf import HfRenderer, HfTokenizer
10 +
11 +
12 +class KimiAudioRenderer(HfRenderer):
13 + """Renderer for Kimi-Audio models.
14 +
15 + This renderer uses HfRenderer internally with a custom TikToken tokenizer.
16 + """
17 +
18 + @classmethod
19 + def from_config( # type: ignore[override]
20 + cls,
21 + config: VllmConfig,
22 + tokenizer_kwargs: dict[str, Any],
23 + ) -> "HfRenderer":
24 + """Create an HfRenderer instance for Kimi-Audio models."""
25 + model_config = config.model_config
26 + if model_config.skip_tokenizer_init:
27 + tokenizer = None
28 + else:
29 + # Extract tokenizer_name from kwargs (already processed by
30 + # tokenizer_args_from_config for ModelScope/GGUF/etc)
31 + tokenizer_name = tokenizer_kwargs.pop(
32 + "tokenizer_name", model_config.tokenizer
33 + )
34 + # Remove tokenizer_cls from kwargs to avoid duplicate argument
35 + tokenizer_kwargs = {
36 + k: v for k, v in tokenizer_kwargs.items() if k != "tokenizer_cls"
37 + }
38 + # Use get_tokenizer directly instead of cached_get_tokenizer
39 + # (KimiAudioTokenizer doesn't work with get_cached_tokenizer)
40 + tokenizer = cast(
41 + HfTokenizer,
42 + get_tokenizer(
43 + tokenizer_name,
44 + tokenizer_cls=KimiAudioTokenizer, # type: ignore[arg-type]
45 + **tokenizer_kwargs,
46 + ),
47 + )
48 +
49 + return HfRenderer(config, tokenizer)
vllm/renderers/registry.py
+9
@@ -19,6 +19,7 @@ _VLLM_RENDERERS = {
19 "deepseek_v32": ("deepseek_v32", "DeepseekV32Renderer"),
20 "hf": ("hf", "HfRenderer"),
21 "grok2": ("grok2", "Grok2Renderer"),
22 + "kimi_audio": ("kimi_audio", "KimiAudioRenderer"),
23 "mistral": ("mistral", "MistralRenderer"),
24 "qwen_vl": ("qwen_vl", "QwenVLRenderer"),
25 "terratorch": ("terratorch", "TerratorchRenderer"),
@@ -74,10 +75,18 @@ RENDERER_REGISTRY = RendererRegistry(
75
76 def renderer_from_config(config: "VllmConfig", **kwargs):
77 model_config = config.model_config
78 +
79 tokenizer_mode, tokenizer_name, args, kwargs = tokenizer_args_from_config(
80 model_config, **kwargs
81 )
82
83 + # Override tokenizer_mode for Kimi-Audio models
84 + if model_config.architecture == "MoonshotKimiaForCausalLM":
85 + tokenizer_mode = "kimi_audio"
86 + # Update model_config so other components (e.g., multimodal registry)
87 + # also use the correct tokenizer mode
88 + model_config.tokenizer_mode = "kimi_audio"
89 +
90 if (
91 model_config.tokenizer_mode == "auto"
92 and model_config.model_impl == "terratorch"
vllm/tokenizers/kimi_audio.py new
+410
@@ -0,0 +1,410 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tokenizer for Kimi-Audio using TikToken."""
4 +
5 +import contextlib
6 +import json
7 +from pathlib import Path
8 +from typing import Any, overload
9 +
10 +import pybase64
11 +import tiktoken
12 +from huggingface_hub import hf_hub_download
13 +from transformers import AddedToken, BatchEncoding
14 +from transformers.utils import chat_template_utils as hf_chat_utils
15 +
16 +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
17 +from vllm.logger import init_logger
18 +from vllm.tokenizers.protocol import TokenizerLike
19 +
20 +logger = init_logger(__name__)
21 +
22 +
23 +def _load_tiktoken_encoding(
24 + vocab_file: Path, special_tokens: dict[str, int]
25 +) -> tuple[Any, dict[str, int]]:
26 + """Load TikToken encoding from vocab file."""
27 + mergeable_ranks: dict[bytes, int] = {}
28 + with open(vocab_file, encoding="utf-8") as f:
29 + for line in f:
30 + line = line.strip()
31 + if not line:
32 + continue
33 + parts = line.split()
34 + if len(parts) == 2:
35 + token_b64 = parts[0]
36 + rank = int(parts[1])
37 + token_bytes = pybase64.b64decode(token_b64)
38 + mergeable_ranks[token_bytes] = rank
39 +
40 + tokenizer = tiktoken.Encoding(
41 + name=str(vocab_file),
42 + pat_str=r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}|"""
43 + r""" ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""",
44 + mergeable_ranks=mergeable_ranks,
45 + special_tokens=special_tokens,
46 + )
47 +
48 + return tokenizer, special_tokens
49 +
50 +
51 +class KimiAudioTokenizer(TokenizerLike):
52 + """TikToken tokenizer for Kimi-Audio."""
53 +
54 + @classmethod
55 + def from_pretrained(
56 + cls,
57 + path_or_repo_id: str | Path,
58 + *args,
59 + trust_remote_code: bool = False,
60 + revision: str | None = None,
61 + download_dir: str | None = None,
62 + **kwargs,
63 + ) -> "KimiAudioTokenizer":
64 + if args:
65 + logger.debug_once("Ignoring extra positional args for KimiAudioTokenizer.")
66 +
67 + path = Path(path_or_repo_id)
68 + if path.is_file():
69 + vocab_file = path
70 + elif path.is_dir():
71 + vocab_file = path / "tiktoken.model"
72 + if not vocab_file.is_file():
73 + vocab_file = path / "tokenizer.model"
74 + else:
75 + # Download from HuggingFace Hub
76 + repo_id = str(path_or_repo_id)
77 +
78 + # Try to download tiktoken.model or tokenizer.model
79 + try:
80 + vocab_path = hf_hub_download(
81 + repo_id=repo_id,
82 + filename="tiktoken.model",
83 + revision=revision,
84 + local_dir=download_dir,
85 + )
86 + vocab_file = Path(vocab_path)
87 + except Exception:
88 + try:
89 + vocab_path = hf_hub_download(
90 + repo_id=repo_id,
91 + filename="tokenizer.model",
92 + revision=revision,
93 + local_dir=download_dir,
94 + )
95 + vocab_file = Path(vocab_path)
96 + except Exception as exc:
97 + raise ValueError(
98 + f"Could not find tiktoken.model or tokenizer.model in {repo_id}"
99 + ) from exc
100 +
101 + # Also download tokenizer_config.json if available
102 + with contextlib.suppress(Exception):
103 + hf_hub_download(
104 + repo_id=repo_id,
105 + filename="tokenizer_config.json",
106 + revision=revision,
107 + local_dir=download_dir,
108 + )
109 +
110 + if not vocab_file.is_file():
111 + raise FileNotFoundError(f"tiktoken.model not found at {vocab_file}.")
112 +
113 + return cls(
114 + vocab_file=vocab_file,
115 + name_or_path=str(path_or_repo_id),
116 + truncation_side=kwargs.get("truncation_side", "left"),
117 + )
118 +
119 + def __init__(
120 + self,
121 + *,
122 + vocab_file: Path,
123 + name_or_path: str,
124 + truncation_side: str,
125 + ) -> None:
126 + super().__init__()
127 + self.name_or_path = name_or_path
128 + self._truncation_side = truncation_side
129 + self._vocab_file = vocab_file
130 +
131 + # Load special tokens from tokenizer_config.json
132 + special_tokens: dict[str, int] = {}
133 + tokenizer_config = vocab_file.parent / "tokenizer_config.json"
134 + if tokenizer_config.is_file():
135 + with open(tokenizer_config, encoding="utf-8") as f:
136 + config = json.load(f)
137 + # Extract special tokens from added_tokens_decoder
138 + added_tokens = config.get("added_tokens_decoder", {})
139 + for token_id_str, token_info in added_tokens.items():
140 + token_id = int(token_id_str)
141 + content = token_info.get("content", "")
142 + if content:
143 + special_tokens[content] = token_id
144 +
145 + self._tokenizer, self._special_tokens = _load_tiktoken_encoding(
146 + vocab_file, special_tokens
147 + )
148 +
149 + # Build token <-> ID mappings
150 + self._token_to_id: dict[str, int] = {}
151 + self._id_to_token: dict[int, str] = {}
152 + for token_bytes, token_id in self._tokenizer._mergeable_ranks.items():
153 + token_str = token_bytes.decode("utf-8", errors="replace")
154 + self._token_to_id[token_str] = token_id
155 + self._id_to_token[token_id] = token_str
156 +
157 + # Initialize added_tokens_decoder before adding special tokens
158 + self._added_tokens_decoder: dict[int, Any] = {}
159 +
160 + # Add Kimi-Audio special tokens
161 + self._add_kimiaudio_special_tokens()
162 +
163 + # Set default special token IDs (will be updated when special tokens are added)
164 + self._bos_token_id = 151643 # Kimi-Audio BOS
165 + self._eos_token_id = 151644 # Kimi-Audio EOS
166 + self._pad_token_id = self._eos_token_id
167 + self._unk_token_id = self._pad_token_id
168 +
169 + self._max_chars_per_token = max(
170 + (len(tok) for tok in self._token_to_id), default=10
171 + )
172 +
173 + def _add_kimiaudio_special_tokens(self) -> None:
174 + """Add Kimi-Audio special tokens to the tokenizer."""
175 + # Tokens should already be in self._special_tokens from tokenizer_config.json
176 + # Just add them to added_tokens_decoder for compatibility
177 + kimiaudio_special_tokens = {
178 + "<|im_media_begin|>": 151661,
179 + "<|im_media_end|>": 151663,
180 + "<|im_kimia_text_blank|>": 151666,
181 + "<|im_msg_end|>": 151645,
182 + "<|im_kimia_user_msg_start|>": 151670,
183 + "<|im_kimia_assistant_msg_start|>": 151671,
184 + }
185 +
186 + for token_str, token_id in kimiaudio_special_tokens.items():
187 + # Only add if not already present
188 + if token_id not in self._added_tokens_decoder:
189 + self._added_tokens_decoder[token_id] = AddedToken(
190 + token_str, single_word=True, normalized=False, special=True
191 + )
192 + # Also ensure it's in _token_to_id and _id_to_token
193 + if token_str not in self._token_to_id:
194 + self._token_to_id[token_str] = token_id
195 + if token_id not in self._id_to_token:
196 + self._id_to_token[token_id] = token_str
197 +
198 + def num_special_tokens_to_add(self) -> int:
199 + return 0
200 +
201 + @property
202 + def all_special_tokens(self) -> list[str]:
203 + return list(self._added_tokens_decoder.values())
204 +
205 + @property
206 + def all_special_ids(self) -> list[int]:
207 + return list(self._added_tokens_decoder.keys())
208 +
209 + @property
210 + def bos_token_id(self) -> int:
211 + return self._bos_token_id
212 +
213 + @property
214 + def eos_token_id(self) -> int:
215 + return self._eos_token_id
216 +
217 + @property
218 + def pad_token_id(self) -> int:
219 + return self._pad_token_id
220 +
221 + @property
222 + def is_fast(self) -> bool:
223 + return False
224 +
225 + @property
226 + def vocab_size(self) -> int:
227 + return self._tokenizer.n_vocab
228 +
229 + @property
230 + def max_token_id(self) -> int:
231 + return self._tokenizer.n_vocab - 1
232 +
233 + @property
234 + def max_chars_per_token(self) -> int:
235 + return self._max_chars_per_token
236 +
237 + @property
238 + def truncation_side(self) -> str:
239 + return self._truncation_side
240 +
241 + @property
242 + def added_tokens_decoder(self) -> dict[int, Any]:
243 + return self._added_tokens_decoder
244 +
245 + @added_tokens_decoder.setter
246 + def added_tokens_decoder(self, value: dict[int, Any]) -> None:
247 + """Set added tokens decoder and update special token IDs."""
248 + self._added_tokens_decoder = value
249 + # Update special token IDs if known tokens are added
250 + for token_id, token in value.items():
251 + token_str = str(token) if hasattr(token, "__str__") else token
252 + if "<|im_kimia_user_msg_start|>" in token_str:
253 + self._bos_token_id = token_id
254 + elif "<|im_msg_end|>" in token_str or "<|im_end|>" in token_str:
255 + self._eos_token_id = token_id
256 +
257 + def get_vocab(self) -> dict[str, int]:
258 + return dict(self._token_to_id)
259 +
260 + def __len__(self) -> int:
261 + """Return vocab size for compatibility with HF tokenizer interface."""
262 + return self._tokenizer.n_vocab
263 +
264 + def get_added_vocab(self) -> dict[str, int]:
265 + return {
266 + str(token): token_id
267 + for token_id, token in self._added_tokens_decoder.items()
268 + }
269 +
270 + def _maybe_truncate(self, tokens: list[int], max_length: int | None) -> list[int]:
271 + if max_length is None or len(tokens) <= max_length:
272 + return tokens
273 + if self.truncation_side == "left":
274 + return tokens[-max_length:]
275 + return tokens[:max_length]
276 +
277 + def encode(
278 + self,
279 + text: str,
280 + truncation: bool | None = None,
281 + max_length: int | None = None,
282 + add_special_tokens: bool = True,
283 + **kwargs,
284 + ) -> list[int]:
285 + del add_special_tokens
286 + # Allow Kimi-Audio special tokens to be encoded
287 + tokens = self._tokenizer.encode(
288 + text,
289 + allowed_special={
290 + "<|im_media_begin|>",
291 + "<|im_media_end|>",
292 + "<|im_kimia_text_blank|>",
293 + "<|im_msg_end|>",
294 + "<|im_kimia_user_msg_start|>",
295 + "<|im_kimia_assistant_msg_start|>",
296 + },
297 + )
298 + if truncation:
299 + tokens = self._maybe_truncate(tokens, max_length)
300 + return tokens
301 +
302 + def decode(self, ids: list[int] | int, skip_special_tokens: bool = False) -> str:
303 + """Decode token IDs to text, optionally skipping special tokens."""
304 + if isinstance(ids, int):
305 + ids = [ids]
306 + if skip_special_tokens:
307 + # Skip tokens that are in special_tokens (loaded from config)
308 + special_ids = set(self._special_tokens.values())
309 + ids = [token_id for token_id in ids if token_id not in special_ids]
310 + return self._tokenizer.decode(ids)
311 +
312 + @overload
313 + def convert_tokens_to_ids(self, tokens: str) -> int: ...
314 +
315 + @overload
316 + def convert_tokens_to_ids(self, tokens: list[str]) -> list[int]: ...
317 +
318 + def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]:
319 + if isinstance(tokens, str):
320 + return self._token_to_id.get(tokens, self._unk_token_id)
321 + return [self._token_to_id.get(token, self._unk_token_id) for token in tokens]
322 +
323 + def convert_ids_to_tokens(
324 + self, ids: list[int], skip_special_tokens: bool = False
325 + ) -> list[str]:
326 + tokens = []
327 + for token_id in ids:
328 + if skip_special_tokens and token_id in self._added_tokens_decoder:
329 + continue
330 + tokens.append(self._id_to_token.get(token_id, "<|unk|>"))
331 + return tokens
332 +
333 + def convert_tokens_to_string(self, tokens: list[str]) -> str:
334 + token_ids = self.convert_tokens_to_ids(tokens)
335 + return self.decode(token_ids, skip_special_tokens=False)
336 +
337 + def __call__(
338 + self,
339 + text: str | list[str],
340 + text_pair: str | None = None,
341 + add_special_tokens: bool = True,
342 + truncation: bool = False,
343 + max_length: int | None = None,
344 + **kwargs,
345 + ) -> BatchEncoding:
346 + if text_pair is not None:
347 + raise NotImplementedError(
348 + "text_pair is not supported for KimiAudioTokenizer."
349 + )
350 +
351 + if isinstance(text, list):
352 + input_ids_batch: list[list[int]] = [
353 + self.encode(
354 + item,
355 + truncation=truncation,
356 + max_length=max_length,
357 + add_special_tokens=add_special_tokens,
358 + )
359 + for item in text
360 + ]
361 + attention_mask_batch = [[1] * len(ids) for ids in input_ids_batch]
362 + return BatchEncoding(
363 + {"input_ids": input_ids_batch, "attention_mask": attention_mask_batch}
364 + )
365 +
366 + input_ids = self.encode(
367 + text,
368 + truncation=truncation,
369 + max_length=max_length,
370 + add_special_tokens=add_special_tokens,
371 + )
372 + attention_mask = [1] * len(input_ids)
373 + return BatchEncoding({"input_ids": input_ids, "attention_mask": attention_mask})
374 +
375 + def get_chat_template(
376 + self, chat_template: str | None, tools: list[dict[str, Any]] | None = None
377 + ) -> str | None:
378 + del tools
379 + return chat_template
380 +
381 + def apply_chat_template(
382 + self,
383 + messages: list[ChatCompletionMessageParam] | None = None,
384 + tools: list[dict[str, Any]] | None = None,
385 + chat_template: str | None = None,
386 + tokenize: bool = False,
387 + **kwargs,
388 + ) -> str | list[int]:
389 + # Handle both 'messages' (protocol) and 'conversation' (caller) parameter names
390 + conversation = messages if messages is not None else kwargs.get("conversation")
391 + if conversation is None:
392 + raise ValueError("Either 'messages' or 'conversation' must be provided.")
393 + template = self.get_chat_template(chat_template, tools=tools)
394 + if template is None:
395 + raise ValueError(
396 + "No chat template available. Provide `chat_template` explicitly."
397 + )
398 + # Use render_jinja_template instead of apply_chat_template
399 + # Note: render_jinja_template returns ([prompts], [generation_indices])
400 + rendered, _ = hf_chat_utils.render_jinja_template(
401 + conversation,
402 + chat_template=template,
403 + tools=tools,
404 + **kwargs,
405 + )
406 + # Extract the first (and usually only) prompt
407 + prompt = rendered[0] if rendered else ""
408 + if tokenize:
409 + return self.encode(prompt, add_special_tokens=False)
410 + return prompt
vllm/tokenizers/registry.py
+1
@@ -35,6 +35,7 @@ _VLLM_TOKENIZERS = {
35 "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"),
36 "grok2": ("grok2", "Grok2Tokenizer"),
37 "hf": ("hf", "CachedHfTokenizer"),
38 + "kimi_audio": ("kimi_audio", "KimiAudioTokenizer"),
39 "mistral": ("mistral", "MistralTokenizer"),
40 "qwen_vl": ("qwen_vl", "QwenVLTokenizer"),
41 }
vllm/transformers_utils/chat_templates/template_kimi_audio.jinja new
+13
@@ -0,0 +1,13 @@
1 +{% set messages = conversations[0] if conversations else [] -%}
2 +{% if messages and messages[0]['role'] == 'system' -%}
3 + {% set loop_messages = messages[1:] -%}
4 +{% else -%}
5 + {% set loop_messages = messages -%}
6 +{% endif -%}
7 +{% for message in loop_messages -%}
8 + {% if message['role'] == 'user' -%}
9 + <|im_kimia_user_msg_start|>{{ message['content'] }}<|im_msg_end|><|im_kimia_assistant_msg_start|>
10 + {%- elif message['role'] == 'assistant' -%}
11 + {{ message['content'] }}<|im_kimia_text_eos|>
12 + {%- endif -%}
13 +{% endfor -%}
vllm/transformers_utils/processors/__init__.py
+18 -17
@@ -10,23 +10,6 @@ reasons:
10
11 import importlib
12
13 -_CLASS_TO_MODULE: dict[str, str] = {
14 - "BagelProcessor": "vllm.transformers_utils.processors.bagel",
15 - "DeepseekVLV2Processor": "vllm.transformers_utils.processors.deepseek_vl2",
16 - "FireRedASR2Processor": "vllm.transformers_utils.processors.fireredasr2",
17 - "FunASRProcessor": "vllm.transformers_utils.processors.funasr",
18 - "GLM4VProcessor": "vllm.transformers_utils.processors.glm4v",
19 - "HunYuanVLProcessor": "vllm.transformers_utils.processors.hunyuan_vl",
20 - "HunYuanVLImageProcessor": "vllm.transformers_utils.processors.hunyuan_vl_image",
21 - "MistralCommonPixtralProcessor": "vllm.transformers_utils.processors.pixtral",
22 - "MistralCommonVoxtralProcessor": "vllm.transformers_utils.processors.voxtral",
23 - "OvisProcessor": "vllm.transformers_utils.processors.ovis",
24 - "Ovis2_5Processor": "vllm.transformers_utils.processors.ovis2_5",
25 - "QwenVLProcessor": "vllm.transformers_utils.processors.qwen_vl",
26 - "Qwen3ASRProcessor": "vllm.transformers_utils.processors.qwen3_asr",
27 -}
28 -
29 -
13 __all__ = [
14 "BagelProcessor",
15 "DeepseekVLV2Processor",
@@ -35,6 +18,7 @@ __all__ = [
18 "GLM4VProcessor",
19 "HunYuanVLProcessor",
20 "HunYuanVLImageProcessor",
21 + "KimiAudioProcessor",
22 "MistralCommonPixtralProcessor",
23 "MistralCommonVoxtralProcessor",
24 "OvisProcessor",
@@ -43,6 +27,23 @@ __all__ = [
27 "Qwen3ASRProcessor",
28 ]
29
30 +_CLASS_TO_MODULE: dict[str, str] = {
31 + "BagelProcessor": "vllm.transformers_utils.processors.bagel",
32 + "DeepseekVLV2Processor": "vllm.transformers_utils.processors.deepseek_vl2",
33 + "FireRedASR2Processor": "vllm.transformers_utils.processors.fireredasr2",
34 + "FunASRProcessor": "vllm.transformers_utils.processors.funasr",
35 + "GLM4VProcessor": "vllm.transformers_utils.processors.glm4v",
36 + "HunYuanVLProcessor": "vllm.transformers_utils.processors.hunyuan_vl",
37 + "HunYuanVLImageProcessor": "vllm.transformers_utils.processors.hunyuan_vl_image",
38 + "KimiAudioProcessor": "vllm.transformers_utils.processors.kimi_audio",
39 + "MistralCommonPixtralProcessor": "vllm.transformers_utils.processors.pixtral",
40 + "MistralCommonVoxtralProcessor": "vllm.transformers_utils.processors.voxtral",
41 + "OvisProcessor": "vllm.transformers_utils.processors.ovis",
42 + "Ovis2_5Processor": "vllm.transformers_utils.processors.ovis2_5",
43 + "QwenVLProcessor": "vllm.transformers_utils.processors.qwen_vl",
44 + "Qwen3ASRProcessor": "vllm.transformers_utils.processors.qwen3_asr",
45 +}
46 +
47
48 def __getattr__(name: str):
49 if name in _CLASS_TO_MODULE:
vllm/transformers_utils/processors/kimi_audio.py new
+163
@@ -0,0 +1,163 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# ruff: noqa
5 +# mypy: ignore-errors
6 +# coding=utf-8
7 +# Copyright 2026 The Moonshot AI team and the HuggingFace Inc. team. All rights reserved.
8 +#
9 +# Licensed under the Apache License, Version 2.0 (the "License");
10 +# you may not use this file except in compliance with the License.
11 +# You may obtain a copy of the License at
12 +#
13 +# http://www.apache.org/licenses/LICENSE-2.0
14 +#
15 +# Unless required by applicable law or agreed to in writing, software
16 +# distributed under the License is distributed on an "AS IS" BASIS,
17 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18 +# See the License for the specific language governing permissions and
19 +# limitations under the License.
20 +"""Processor for Kimi-Audio ASR model."""
21 +
22 +from collections.abc import Mapping
23 +from typing import Any
24 +
25 +import numpy as np
26 +import torch
27 +from transformers import AutoFeatureExtractor, BatchFeature, ProcessorMixin
28 +from transformers.audio_utils import AudioInput
29 +from transformers.tokenization_utils_base import TextInput
30 +
31 +from vllm.tokenizers.kimi_audio import KimiAudioTokenizer
32 +
33 +
34 +def _get_feat_extract_output_lengths(input_lengths: torch.Tensor) -> torch.Tensor:
35 + """Compute output lengths after Whisper feature extraction."""
36 + input_lengths_leave = input_lengths % 100
37 + feat_lengths = (input_lengths_leave - 1) // 2 + 1
38 + output_lengths = (
39 + ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
40 + )
41 + return output_lengths
42 +
43 +
44 +class KimiAudioProcessor(ProcessorMixin):
45 + r"""
46 + Constructs a Kimi-Audio processor.
47 +
48 + [`KimiAudioProcessor`] offers all the functionalities of [`WhisperFeatureExtractor`], and a tokenizer.
49 + See the [`~KimiAudioProcessor.__call__`] and [`~KimiAudioProcessor.decode`] for more information.
50 +
51 + Args:
52 + feature_extractor ([`WhisperFeatureExtractor`], *optional*):
53 + The audio feature extractor.
54 + tokenizer ([`PreTrainedTokenizer`], *optional*):
55 + The text tokenizer.
56 + """
57 +
58 + # Required for ProcessorMixin
59 + attributes = ["feature_extractor", "tokenizer"]
60 + feature_extractor_class = "AutoFeatureExtractor"
61 + tokenizer_class = "AutoTokenizer"
62 +
63 + # Special token IDs
64 + KIMIA_MEDIA_BEGIN: int = 151661
65 + KIMIA_MEDIA_END: int = 151663
66 + KIMIA_TEXT_BLANK: int = 151666
67 +
68 + # Audio processing constants
69 + AUDIO_SEQ_LEN: int = 376
70 +
71 + def __init__(self, feature_extractor=None, tokenizer=None, **kwargs):
72 + # Pass feature_extractor and tokenizer to parent ProcessorMixin
73 + super().__init__(
74 + feature_extractor=feature_extractor,
75 + tokenizer=tokenizer,
76 + **kwargs,
77 + )
78 +
79 + def check_argument_for_proper_class(self, attribute_name: str, argument: Any):
80 + """Override to skip class validation for custom tokenizer."""
81 + # Skip validation for tokenizer since KimiAudioTokenizer doesn't inherit
82 + # from PreTrainedTokenizerBase but is compatible
83 + if attribute_name == "tokenizer" and argument is not None:
84 + return
85 + # For other attributes, use default validation
86 + super().check_argument_for_proper_class(attribute_name, argument)
87 +
88 + def __call__(
89 + self,
90 + text: TextInput = None,
91 + audio: AudioInput = None,
92 + return_tensors: str = "pt",
93 + **kwargs,
94 + ) -> BatchFeature:
95 + """
96 + Main method to prepare for the model one or several sequences(s) and audio(s).
97 +
98 + Args:
99 + text (`str`, `List[str]`):
100 + The sequence or batch of sequences to be encoded.
101 + audio (`np.ndarray`, `List[np.ndarray]`):
102 + The audio or batch of audio to be prepared. Each audio can be a NumPy array.
103 + return_tensors (`str`):
104 + The type of tensors to return ("pt", "np", etc.)
105 + """
106 + if text is None:
107 + raise ValueError("You need to specify either a `text` input to process.")
108 +
109 + # Process audio if provided
110 + if audio is not None:
111 + # Ensure audio is a list
112 + if isinstance(audio, np.ndarray):
113 + audio = [audio]
114 +
115 + # Pad audio to hop length (required by WhisperFeatureExtractor)
116 + hop_length = self.feature_extractor.hop_length
117 + padded_audio = []
118 + for aud in audio:
119 + length = aud.shape[-1]
120 + if length % hop_length != 0:
121 + pad_length = hop_length - (length % hop_length)
122 + aud = np.pad(
123 + aud, (0, pad_length), mode="constant", constant_values=0
124 + )
125 + padded_audio.append(aud)
126 +
127 + # Use feature_extractor directly like Qwen3ASR does
128 + audio_inputs = self.feature_extractor(
129 + padded_audio,
130 + sampling_rate=16000,
131 + padding=True,
132 + return_attention_mask=True,
133 + return_tensors=return_tensors,
134 + )
135 + # Rename to match Kimi-Audio expectations
136 + if "input_features" in audio_inputs:
137 + audio_inputs["whisper_input_features"] = audio_inputs.pop(
138 + "input_features"
139 + )
140 + if "attention_mask" in audio_inputs:
141 + audio_inputs["feature_attention_mask"] = audio_inputs.pop(
142 + "attention_mask"
143 + )
144 + else:
145 + audio_inputs = {}
146 +
147 + # Handle text input - can be string or token IDs from vLLM processor
148 + if isinstance(text, list) and len(text) > 0 and isinstance(text[0], int):
149 + # Text is already token IDs (from vLLM processor) - just wrap
150 + text_inputs = {"input_ids": torch.tensor([text], dtype=torch.long)}
151 + else:
152 + # Text is string - tokenize
153 + if not isinstance(text, list):
154 + text = [text]
155 +
156 + text_inputs = self.tokenizer(
157 + text, return_tensors=return_tensors, padding=True
158 + )
159 +
160 + return BatchFeature(
161 + data={**text_inputs, **audio_inputs},
162 + tensor_type=return_tensors,
163 + )