| 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 | 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 |
| 12 | from transformers import BatchFeature |
| 13 | from transformers import WhisperConfig as HFWhisperConfig |
| 14 | |
| 15 | from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig |
| 16 | from vllm.config.multimodal import MultiModalDummyOptions |
| 17 | from vllm.config.speech_to_text import SpeechToTextParams |
| 18 | from vllm.inputs import PromptType, TokensPrompt |
| 19 | from vllm.model_executor.model_loader import DefaultModelLoader |
| 20 | from vllm.model_executor.models.interfaces import ( |
| 21 | MultiModalEmbeddings, |
| 22 | SupportsMultiModal, |
| 23 | SupportsPP, |
| 24 | SupportsTranscription, |
| 25 | ) |
| 26 | from vllm.model_executor.models.utils import ( |
| 27 | AutoWeightsLoader, |
| 28 | WeightsMapper, |
| 29 | init_vllm_registered_model, |
| 30 | maybe_prefix, |
| 31 | ) |
| 32 | from vllm.model_executor.models.whisper import WhisperEncoder |
| 33 | from vllm.model_executor.models.whisper_utils import ISO639_1_SUPPORTED_LANGS |
| 34 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 35 | from vllm.multimodal.inputs import MultiModalFieldConfig |
| 36 | from vllm.multimodal.parse import ( |
| 37 | AudioItem, |
| 38 | DictEmbeddingItems, |
| 39 | ModalityData, |
| 40 | ModalityDataItems, |
| 41 | MultiModalDataItems, |
| 42 | MultiModalDataParser, |
| 43 | ) |
| 44 | from vllm.multimodal.processing import ( |
| 45 | BaseDummyInputsBuilder, |
| 46 | BaseProcessingInfo, |
| 47 | PromptReplacement, |
| 48 | ) |
| 49 | from vllm.multimodal.processing.processor import ( |
| 50 | BaseMultiModalProcessor, |
| 51 | HFMultiModalInputs, |
| 52 | ProcessorInputs, |
| 53 | ) |
| 54 | from vllm.sequence import IntermediateTensors |
| 55 | from vllm.tokenizers import cached_get_tokenizer |
| 56 | from vllm.tokenizers.kimi_audio import KimiAudioTokenizer |
| 57 | from vllm.transformers_utils.processor import cached_feature_extractor_from_config |
| 58 | from vllm.transformers_utils.processors.kimi_audio import KimiAudioProcessor |
| 59 | |
| 60 | # Kimi-Audio constants |
| 61 | KIMIA_WHISPER_SUBFOLDER = "whisper-large-v3" |
| 62 | |
| 63 | |
| 64 | def _get_feat_extract_output_lengths(input_lengths: torch.Tensor) -> torch.Tensor: |
| 65 | """Compute output lengths after Whisper feature extraction. |
| 66 | |
| 67 | Whisper processes audio through multiple conv layers with stride=2, |
| 68 | producing 13 output features per 100 input samples. |
| 69 | """ |
| 70 | input_lengths_leave = input_lengths % 100 |
| 71 | feat_lengths = (input_lengths_leave - 1) // 2 + 1 |
| 72 | output_lengths = ( |
| 73 | ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 |
| 74 | ) |
| 75 | return output_lengths |
| 76 | |
| 77 | |
| 78 | class KimiAudioWhisperEncoder(WhisperEncoder): |
| 79 | """WhisperEncoder for Kimi-Audio with packed_modules_mapping.""" |
| 80 | |
| 81 | # packed_modules_mapping for Q/K/V fusion during weight loading |
| 82 | packed_modules_mapping = { |
| 83 | "qkv_proj": ["q_proj", "k_proj", "v_proj"], |
| 84 | } |
| 85 | |
| 86 | hf_to_vllm_mapper = WeightsMapper( |
| 87 | orig_to_new_stacked={ |
| 88 | ".self_attn.q_proj": (".self_attn.qkv_proj", "q"), |
| 89 | ".self_attn.k_proj": (".self_attn.qkv_proj", "k"), |
| 90 | ".self_attn.v_proj": (".self_attn.qkv_proj", "v"), |
| 91 | }, |
| 92 | ) |
| 93 | |
| 94 | def __init__( |
| 95 | self, *, vllm_config: VllmConfig, prefix: str = "", init_in_fp32: bool = False |
| 96 | ): |
| 97 | # Load Whisper config from subfolder (authoritative source) |
| 98 | # Kimi-Audio stores Whisper config in whisper-large-v3/config.json |
| 99 | model_path = vllm_config.model_config.model |
| 100 | |
| 101 | # Load WhisperConfig from the subfolder |
| 102 | whisper_config = HFWhisperConfig.from_pretrained( |
| 103 | model_path, |
| 104 | subfolder=KIMIA_WHISPER_SUBFOLDER, |
| 105 | revision=vllm_config.model_config.revision, |
| 106 | ) |
| 107 | |
| 108 | super().__init__( |
| 109 | vllm_config=vllm_config.with_hf_config(whisper_config), |
| 110 | prefix=prefix, |
| 111 | init_in_fp32=init_in_fp32, |
| 112 | ) |
| 113 | |
| 114 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 115 | loader = AutoWeightsLoader(self) |
| 116 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 117 | |
| 118 | |
| 119 | # ----------------------------------------------------------------------------- |
| 120 | # Processing Info, Dummy Inputs, and MultiModal Processor |
| 121 | # (Following Qwen3ASR pattern - same file as model) |
| 122 | # ----------------------------------------------------------------------------- |
| 123 | |
| 124 | |
| 125 | class KimiAudioProcessingInfo(BaseProcessingInfo): |
| 126 | """Processing info for vLLM registry.""" |
| 127 | |
| 128 | def get_hf_processor(self, **kwargs: object) -> KimiAudioProcessor: |
| 129 | feature_extractor = cached_feature_extractor_from_config( |
| 130 | self.ctx.model_config, |
| 131 | subfolder=KIMIA_WHISPER_SUBFOLDER, |
| 132 | ) |
| 133 | |
| 134 | return KimiAudioProcessor( |
| 135 | feature_extractor=feature_extractor, |
| 136 | tokenizer=self.get_tokenizer(), |
| 137 | ) |
| 138 | |
| 139 | def get_feature_extractor(self, **kwargs: object): |
| 140 | return cached_feature_extractor_from_config( |
| 141 | self.ctx.model_config, subfolder=KIMIA_WHISPER_SUBFOLDER |
| 142 | ) |
| 143 | |
| 144 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 145 | return {"audio": 1} |
| 146 | |
| 147 | def get_data_parser(self) -> "KimiAudioMultiModalDataParser": |
| 148 | feature_extractor = self.get_feature_extractor() |
| 149 | return KimiAudioMultiModalDataParser( |
| 150 | target_sr=feature_extractor.sampling_rate, |
| 151 | expected_hidden_size=self._get_expected_hidden_size(), |
| 152 | ) |
| 153 | |
| 154 | |
| 155 | class KimiAudioDummyInputsBuilder(BaseDummyInputsBuilder[KimiAudioProcessingInfo]): |
| 156 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 157 | return "" |
| 158 | |
| 159 | def get_dummy_mm_data( |
| 160 | self, |
| 161 | seq_len: int, |
| 162 | mm_counts: Mapping[str, int], |
| 163 | mm_options: Mapping[str, Any] | None = None, |
| 164 | ) -> dict[str, Any]: |
| 165 | num_audios = mm_counts.get("audio", 0) |
| 166 | if num_audios == 0: |
| 167 | return {} |
| 168 | |
| 169 | feature_extractor = self.info.get_feature_extractor() |
| 170 | target_audio_length = ( |
| 171 | min(feature_extractor.chunk_length, 30) * feature_extractor.sampling_rate |
| 172 | ) |
| 173 | |
| 174 | return { |
| 175 | "audio": self._get_dummy_audios( |
| 176 | length=target_audio_length, num_audios=num_audios |
| 177 | ), |
| 178 | } |
| 179 | |
| 180 | |
| 181 | # Field config for Kimi-Audio multimodal data |
| 182 | _KIMIAUDIO_FIELD_CONFIG = { |
| 183 | "whisper_input_features": MultiModalFieldConfig.batched("audio"), |
| 184 | "feature_attention_mask": MultiModalFieldConfig.batched("audio", keep_on_cpu=True), |
| 185 | } |
| 186 | |
| 187 | |
| 188 | class KimiAudioMultiModalDataParser(MultiModalDataParser): |
| 189 | """Custom data parser for Kimi-Audio multimodal data.""" |
| 190 | |
| 191 | def _parse_audio_data( |
| 192 | self, |
| 193 | data: dict[str, torch.Tensor] | ModalityData[AudioItem], |
| 194 | ) -> ModalityDataItems[Any, Any] | None: |
| 195 | if isinstance(data, dict): |
| 196 | return DictEmbeddingItems( |
| 197 | data, |
| 198 | modality="audio", |
| 199 | required_fields={"whisper_input_features", "feature_attention_mask"}, |
| 200 | fields_factory=lambda hf_inputs: _KIMIAUDIO_FIELD_CONFIG, |
| 201 | ) |
| 202 | |
| 203 | return super()._parse_audio_data(data) |
| 204 | |
| 205 | |
| 206 | class KimiAudioMultiModalProcessor(BaseMultiModalProcessor[KimiAudioProcessingInfo]): |
| 207 | """vLLM multi-modal processor wrapper for Kimi-Audio.""" |
| 208 | |
| 209 | def get_dummy_inputs( |
| 210 | self, |
| 211 | seq_len: int, |
| 212 | mm_counts: Mapping[str, int], |
| 213 | mm_options: MultiModalDummyOptions, |
| 214 | ) -> ProcessorInputs: |
| 215 | builder = self.dummy_inputs |
| 216 | dummy_mm_data = builder.get_dummy_mm_data(seq_len, mm_counts, mm_options) |
| 217 | dummy_mm_items = self.info.parse_mm_data(dummy_mm_data) |
| 218 | |
| 219 | num_audios = mm_counts.get("audio", 0) |
| 220 | dummy_tokens = ( |
| 221 | [198] |
| 222 | if num_audios == 0 |
| 223 | else [ |
| 224 | KimiAudioProcessor.KIMIA_MEDIA_BEGIN, |
| 225 | KimiAudioProcessor.KIMIA_TEXT_BLANK, |
| 226 | KimiAudioProcessor.KIMIA_MEDIA_END, |
| 227 | ] |
| 228 | * num_audios |
| 229 | ) |
| 230 | |
| 231 | return ProcessorInputs(prompt=dummy_tokens, mm_data_items=dummy_mm_items) |
| 232 | |
| 233 | def _get_hf_mm_inputs( |
| 234 | self, |
| 235 | mm_items: MultiModalDataItems, |
| 236 | hf_kwargs: Mapping[str, object], |
| 237 | ) -> HFMultiModalInputs: |
| 238 | hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) |
| 239 | |
| 240 | # Convert audio format: [(array, sr), ...] -> [array, ...] |
| 241 | # KimiAudioProcessor expects raw numpy arrays |
| 242 | mm_data = hf_inputs.hf_data |
| 243 | if audios := mm_data.pop("audio", []): |
| 244 | assert isinstance(audios, Sequence) |
| 245 | audio_arrays = [] |
| 246 | for aud in audios: |
| 247 | if isinstance(aud, (tuple, list)) and len(aud) == 2: |
| 248 | # Format: (audio_array, sampling_rate) |
| 249 | audio_arrays.append(aud[0]) |
| 250 | elif isinstance(aud, np.ndarray): |
| 251 | audio_arrays.append(aud) |
| 252 | else: |
| 253 | audio_arrays.append(aud) |
| 254 | mm_data["audio"] = audio_arrays |
| 255 | |
| 256 | return hf_inputs |
| 257 | |
| 258 | def _get_mm_fields_config( |
| 259 | self, |
| 260 | hf_inputs: BatchFeature, |
| 261 | hf_processor_mm_kwargs: Mapping[str, object], |
| 262 | ) -> Mapping[str, Any]: |
| 263 | """Get multi-modal field configuration.""" |
| 264 | return _KIMIAUDIO_FIELD_CONFIG |
| 265 | |
| 266 | def _get_prompt_updates( |
| 267 | self, |
| 268 | mm_items, |
| 269 | hf_processor_mm_kwargs, |
| 270 | out_mm_kwargs, |
| 271 | ) -> Sequence[PromptReplacement]: |
| 272 | """Get prompt updates for audio tokens.""" |
| 273 | # Get audio feature lengths from processed output |
| 274 | out_mm_data = out_mm_kwargs.get_data() |
| 275 | feature_attention_mask = out_mm_data.get("feature_attention_mask") |
| 276 | |
| 277 | if feature_attention_mask is not None: |
| 278 | audio_output_lens = _get_feat_extract_output_lengths( |
| 279 | feature_attention_mask.sum(-1) |
| 280 | ) |
| 281 | audio_output_lengths = audio_output_lens.tolist() |
| 282 | else: |
| 283 | audio_output_lengths = [] |
| 284 | |
| 285 | def get_replacement_kimiaudio(item_idx: int): |
| 286 | num_features = ( |
| 287 | audio_output_lengths[item_idx] |
| 288 | if item_idx < len(audio_output_lengths) |
| 289 | else 376 |
| 290 | ) |
| 291 | if num_features == 0: |
| 292 | num_features = 376 # Default Kimi-Audio sequence length |
| 293 | # Return the placeholder token ID repeated num_features times |
| 294 | return [KimiAudioProcessor.KIMIA_TEXT_BLANK] * num_features |
| 295 | |
| 296 | # Use the token ID as target (as a list) |
| 297 | return [ |
| 298 | PromptReplacement( |
| 299 | modality="audio", |
| 300 | target=[KimiAudioProcessor.KIMIA_TEXT_BLANK], |
| 301 | replacement=get_replacement_kimiaudio, |
| 302 | ), |
| 303 | ] |
| 304 | |
| 305 | |
| 306 | # ----------------------------------------------------------------------------- |
| 307 | # Model Definition |
| 308 | # ----------------------------------------------------------------------------- |
| 309 | |
| 310 | |
| 311 | class KimiAudioMultiModalProjector(nn.Module): |
| 312 | """Projects Whisper features to LLM embedding space. |
| 313 | |
| 314 | Kimi-Audio VQ-Adaptor architecture: |
| 315 | Custom Whisper (5120) → Linear[5120→3584] → Linear[3584→3584] → LayerNorm |
| 316 | """ |
| 317 | |
| 318 | def __init__( |
| 319 | self, |
| 320 | whisper_dim: int = 5120, # Kimi-Audio custom Whisper encoder dim |
| 321 | llm_dim: int = 3584, |
| 322 | prefix: str = "", |
| 323 | ): |
| 324 | super().__init__() |
| 325 | self.whisper_dim = whisper_dim |
| 326 | self.llm_dim = llm_dim |
| 327 | |
| 328 | # VQ-Adaptor layers (exact checkpoint structure) |
| 329 | # layers.0: Linear[5120 → 3584] |
| 330 | self.vq_adaptor_layers_0 = nn.Linear(whisper_dim, llm_dim) |
| 331 | # layers.3: Linear[3584 → 3584] |
| 332 | self.vq_adaptor_layers_3 = nn.Linear(llm_dim, llm_dim) |
| 333 | # layers.4: LayerNorm[3584] |
| 334 | self.vq_adaptor_layers_4 = nn.LayerNorm(llm_dim) |
| 335 | |
| 336 | def forward(self, audio_features: torch.Tensor) -> torch.Tensor: |
| 337 | # Project: [B, T, 5120] → [B, T, 3584] |
| 338 | hidden = self.vq_adaptor_layers_0(audio_features) |
| 339 | hidden = torch.nn.functional.gelu(hidden) |
| 340 | hidden = self.vq_adaptor_layers_3(hidden) |
| 341 | hidden = self.vq_adaptor_layers_4(hidden) |
| 342 | return hidden |
| 343 | |
| 344 | |
| 345 | @MULTIMODAL_REGISTRY.register_processor( |
| 346 | KimiAudioMultiModalProcessor, |
| 347 | info=KimiAudioProcessingInfo, |
| 348 | dummy_inputs=KimiAudioDummyInputsBuilder, |
| 349 | ) |
| 350 | class KimiAudioForConditionalGeneration( |
| 351 | nn.Module, |
| 352 | SupportsMultiModal, |
| 353 | SupportsPP, |
| 354 | SupportsTranscription, |
| 355 | ): |
| 356 | """Kimi-Audio model for ASR transcription.""" |
| 357 | |
| 358 | # Kimi-Audio supports a subset of Whisper's supported languages |
| 359 | supported_languages: ClassVar[Mapping[str, str]] = { |
| 360 | k: ISO639_1_SUPPORTED_LANGS[k] |
| 361 | for k in ["zh", "en", "ja", "ko", "de", "fr", "es", "it", "pt", "ru", "ar"] |
| 362 | } |
| 363 | supports_transcription: ClassVar[Literal[True]] = True |
| 364 | |
| 365 | hf_to_vllm_mapper = WeightsMapper( |
| 366 | orig_to_new_prefix={ |
| 367 | # audio tower |
| 368 | "model.encoder.": "audio_tower.", |
| 369 | # Audio projector (VQ-Adaptor) |
| 370 | "model.vq_adaptor.layers.0.": "multi_modal_projector.vq_adaptor_layers_0.", |
| 371 | "model.vq_adaptor.layers.3.": "multi_modal_projector.vq_adaptor_layers_3.", |
| 372 | "model.vq_adaptor.layers.4.": "multi_modal_projector.vq_adaptor_layers_4.", |
| 373 | # Language model |
| 374 | "model.layers.": "language_model.model.layers.", |
| 375 | # Embeddings and output |
| 376 | "model.embed_tokens.": "language_model.model.embed_tokens.", |
| 377 | "model.norm.": "language_model.model.norm.", |
| 378 | "lm_head.": "language_model.lm_head.", |
| 379 | # MIMO/TTS weights and any `model.` residue no rule above |
| 380 | # claimed: this model only does ASR (speech-to-text). |
| 381 | "model.": None, |
| 382 | "mimo_layers.": None, |
| 383 | "mimo_output.": None, |
| 384 | "mimo_norm.": None, |
| 385 | }, |
| 386 | orig_to_new_substr={ |
| 387 | ".fc1.": ".mlp.fc1.", |
| 388 | ".fc2.": ".mlp.fc2.", |
| 389 | }, |
| 390 | ) |
| 391 | |
| 392 | # Audio placeholder token sequence |
| 393 | AUDIO_PLACEHOLDER = "<|im_media_begin|><|im_kimia_text_blank|><|im_media_end|>" |
| 394 | |
| 395 | @classmethod |
| 396 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 397 | return cls.AUDIO_PLACEHOLDER if modality.startswith("audio") else None |
| 398 | |
| 399 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 400 | super().__init__() |
| 401 | self.config = vllm_config.model_config.hf_config |
| 402 | self.quant_config = vllm_config.quant_config |
| 403 | self.multimodal_config = vllm_config.model_config.multimodal_config |
| 404 | self.model_path = vllm_config.model_config.model |
| 405 | |
| 406 | self.secondary_weights = [ |
| 407 | DefaultModelLoader.Source( |
| 408 | model_or_path=vllm_config.model_config.model, |
| 409 | subfolder="whisper-large-v3", |
| 410 | revision=vllm_config.model_config.revision, |
| 411 | ) |
| 412 | ] |
| 413 | |
| 414 | with self._mark_tower_model(vllm_config, "audio"): |
| 415 | self.audio_tower = KimiAudioWhisperEncoder( |
| 416 | vllm_config=vllm_config, |
| 417 | prefix=maybe_prefix(prefix, "audio_tower"), |
| 418 | ) |
| 419 | self.multi_modal_projector = KimiAudioMultiModalProjector( |
| 420 | whisper_dim=getattr(self.config, "kimia_adaptor_input_dim", 5120), |
| 421 | llm_dim=self.config.hidden_size, |
| 422 | prefix=maybe_prefix(prefix, "multi_modal_projector"), |
| 423 | ) |
| 424 | |
| 425 | with self._mark_language_model(vllm_config): |
| 426 | self.language_model = init_vllm_registered_model( |
| 427 | vllm_config=vllm_config.with_hf_config( |
| 428 | self.config, architectures=["Qwen2ForCausalLM"] |
| 429 | ), |
| 430 | prefix=maybe_prefix(prefix, "language_model"), |
| 431 | ) |
| 432 | |
| 433 | self.make_empty_intermediate_tensors = ( |
| 434 | self.language_model.make_empty_intermediate_tensors |
| 435 | ) |
| 436 | |
| 437 | def _parse_and_validate_audio_input( |
| 438 | self, **kwargs: object |
| 439 | ) -> dict[str, torch.Tensor] | None: |
| 440 | whisper_input_features = kwargs.pop("whisper_input_features", None) |
| 441 | if whisper_input_features is None: |
| 442 | return None |
| 443 | |
| 444 | return {"whisper_input_features": whisper_input_features} |
| 445 | |
| 446 | def _process_audio_input( |
| 447 | self, audio_input: dict[str, torch.Tensor] |
| 448 | ) -> torch.Tensor: |
| 449 | input_features = audio_input["whisper_input_features"] |
| 450 | |
| 451 | # KimiAudioWhisperEncoder expects list of tensors |
| 452 | if input_features.dim() == 3: |
| 453 | input_features = input_features.unbind(dim=0) |
| 454 | |
| 455 | # Run through Whisper encoder |
| 456 | audio_features = self.audio_tower(input_features) |
| 457 | |
| 458 | # Reshape for 4x downsampling (Whisper outputs at 50Hz, need 12.5Hz) |
| 459 | B, T, D = audio_features.shape |
| 460 | if T % 4 != 0: |
| 461 | pad_len = 4 - (T % 4) |
| 462 | audio_features = torch.nn.functional.pad(audio_features, (0, 0, 0, pad_len)) |
| 463 | T = audio_features.shape[1] # Update T after padding |
| 464 | |
| 465 | audio_features = audio_features.reshape(B, T // 4, D * 4) |
| 466 | |
| 467 | # Project to LLM dimension |
| 468 | audio_embeds = self.multi_modal_projector(audio_features) |
| 469 | return audio_embeds |
| 470 | |
| 471 | def embed_multimodal(self, **kwargs: object) -> list[torch.Tensor] | None: |
| 472 | audio_input = self._parse_and_validate_audio_input(**kwargs) |
| 473 | if audio_input is None: |
| 474 | return [] |
| 475 | |
| 476 | audio_embeds = self._process_audio_input(audio_input) |
| 477 | |
| 478 | # audio_embeds shape: [batch_size, seq_len, hidden_dim] |
| 479 | # Return as list of 2D tensors, one per batch item |
| 480 | if audio_embeds.dim() == 3: |
| 481 | # Unbind batch dimension: [B, T, D] -> list of B tensors [T, D] |
| 482 | return list(audio_embeds.unbind(dim=0)) |
| 483 | else: |
| 484 | # Single sample: [T, D] -> wrap in list |
| 485 | return [audio_embeds] |
| 486 | |
| 487 | def embed_input_ids( |
| 488 | self, |
| 489 | input_ids: torch.Tensor, |
| 490 | multimodal_embeddings: MultiModalEmbeddings | None = None, |
| 491 | *, |
| 492 | is_multimodal: torch.Tensor | None = None, |
| 493 | ) -> torch.Tensor: |
| 494 | """Embed input IDs and fuse with audio embeddings. |
| 495 | |
| 496 | Kimi-Audio fusion: inputs_embeds = (text_emb + audio_emb) × √2 |
| 497 | |
| 498 | For PP compatibility, we use the is_multimodal mask from vLLM engine |
| 499 | which is correctly computed per pipeline stage. |
| 500 | """ |
| 501 | # Get text embeddings |
| 502 | inputs_embeds = self.language_model.model.embed_tokens(input_ids) |
| 503 | |
| 504 | if multimodal_embeddings is None or len(multimodal_embeddings) == 0: |
| 505 | return inputs_embeds |
| 506 | |
| 507 | # is_multimodal must be provided for PP to work correctly |
| 508 | if is_multimodal is None or not is_multimodal.any(): |
| 509 | return inputs_embeds |
| 510 | |
| 511 | # multimodal_embeddings[0] contains audio embeddings |
| 512 | audio_embeds = multimodal_embeddings[0] |
| 513 | |
| 514 | # Handle different tensor structures |
| 515 | if isinstance(audio_embeds, (list, tuple)): |
| 516 | audio_embeds = torch.cat(audio_embeds, dim=0) |
| 517 | elif audio_embeds.dim() == 3: |
| 518 | audio_embeds = audio_embeds.reshape(-1, audio_embeds.shape[-1]) |
| 519 | |
| 520 | # In PP, audio_embeds count should match is_multimodal.sum() |
| 521 | # For now, use embeddings sequentially |
| 522 | # (works for non-PP, PP needs vLLM infra fix) |
| 523 | num_mm_tokens = is_multimodal.sum().item() |
| 524 | num_audio_embeds = audio_embeds.shape[0] |
| 525 | |
| 526 | # Use the minimum of available embeddings and positions |
| 527 | # This ensures we don't access out-of-bounds |
| 528 | num_to_use = min(num_audio_embeds, num_mm_tokens) |
| 529 | |
| 530 | # Get positions for the tokens we'll actually process |
| 531 | mm_positions = is_multimodal.nonzero(as_tuple=True)[0] |
| 532 | actual_mm_mask = torch.zeros_like(is_multimodal) |
| 533 | actual_mm_mask[mm_positions[:num_to_use]] = True |
| 534 | |
| 535 | # Use corresponding embeddings |
| 536 | used_audio_embeds = audio_embeds[:num_to_use] |
| 537 | |
| 538 | # Save text embeddings at multimodal positions |
| 539 | text_at_mm_positions = inputs_embeds[actual_mm_mask].clone() |
| 540 | |
| 541 | # Replace text with audio at multimodal positions |
| 542 | inputs_embeds[actual_mm_mask] = used_audio_embeds.to(dtype=inputs_embeds.dtype) |
| 543 | |
| 544 | # Apply Kimi-Audio's unique fusion formula: (text + audio) × √2 |
| 545 | inputs_embeds[actual_mm_mask] = ( |
| 546 | inputs_embeds[actual_mm_mask] + text_at_mm_positions |
| 547 | ) * (2**0.5) |
| 548 | |
| 549 | return inputs_embeds |
| 550 | |
| 551 | def forward( |
| 552 | self, |
| 553 | input_ids: torch.Tensor | None, |
| 554 | positions: torch.Tensor, |
| 555 | intermediate_tensors: IntermediateTensors | None = None, |
| 556 | inputs_embeds: torch.Tensor | None = None, |
| 557 | **kwargs: object, |
| 558 | ) -> torch.Tensor | IntermediateTensors: |
| 559 | if intermediate_tensors is not None: |
| 560 | inputs_embeds = None |
| 561 | |
| 562 | hidden_states = self.language_model.model( |
| 563 | input_ids, |
| 564 | positions, |
| 565 | intermediate_tensors, |
| 566 | inputs_embeds=inputs_embeds, |
| 567 | ) |
| 568 | |
| 569 | return hidden_states |
| 570 | |
| 571 | def compute_logits( |
| 572 | self, |
| 573 | hidden_states: torch.Tensor, |
| 574 | ) -> torch.Tensor | None: |
| 575 | return self.language_model.compute_logits(hidden_states) |
| 576 | |
| 577 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 578 | loader = AutoWeightsLoader(self) |
| 579 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 580 | |
| 581 | @classmethod |
| 582 | def get_speech_to_text_config( |
| 583 | cls, model_config: ModelConfig, task_type: str |
| 584 | ) -> SpeechToTextConfig: |
| 585 | """Get speech-to-text config with custom processor.""" |
| 586 | # Load feature extractor for config values |
| 587 | feature_extractor = cached_feature_extractor_from_config( |
| 588 | model_config, |
| 589 | subfolder=KIMIA_WHISPER_SUBFOLDER, |
| 590 | ) |
| 591 | |
| 592 | return SpeechToTextConfig( |
| 593 | max_audio_clip_s=feature_extractor.chunk_length, |
| 594 | sample_rate=feature_extractor.sampling_rate, |
| 595 | ) |
| 596 | |
| 597 | @classmethod |
| 598 | def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType: |
| 599 | audio = stt_params.audio |
| 600 | model_config = stt_params.model_config |
| 601 | task_type = stt_params.task_type |
| 602 | request_prompt = stt_params.request_prompt |
| 603 | |
| 604 | tokenizer = cached_get_tokenizer( |
| 605 | model_config.tokenizer, |
| 606 | tokenizer_cls=KimiAudioTokenizer, |
| 607 | tokenizer_mode=model_config.tokenizer_mode, |
| 608 | revision=model_config.tokenizer_revision, |
| 609 | trust_remote_code=model_config.trust_remote_code, |
| 610 | ) |
| 611 | |
| 612 | if task_type not in ("transcribe", "translate"): |
| 613 | raise ValueError( |
| 614 | f"Unsupported task_type '{task_type}'. " |
| 615 | "Supported task types are 'transcribe' and 'translate'." |
| 616 | ) |
| 617 | |
| 618 | # Incorporate request_prompt as context/instruction if provided |
| 619 | user_content = ( |
| 620 | f"{request_prompt}\n{cls.AUDIO_PLACEHOLDER}" |
| 621 | if request_prompt |
| 622 | else cls.AUDIO_PLACEHOLDER |
| 623 | ) |
| 624 | |
| 625 | prompt = ( |
| 626 | f"<|im_kimia_user_msg_start|>{user_content}" |
| 627 | f"<|im_msg_end|><|im_kimia_assistant_msg_start|>" |
| 628 | ) |
| 629 | |
| 630 | prompt_token_ids = tokenizer.encode(prompt) |
| 631 | |
| 632 | return TokensPrompt( |
| 633 | prompt_token_ids=prompt_token_ids, |
| 634 | multi_modal_data={"audio": audio}, |
| 635 | ) |
| 636 | |
| 637 | @classmethod |
| 638 | def post_process_output(cls, text: str) -> str: |
| 639 | if not text: |
| 640 | return "" |
| 641 | return text.strip() |