| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Kimi-K3 multimodal model implementation for vLLM.""" |
| 4 | |
| 5 | from collections.abc import Iterable |
| 6 | from typing import cast |
| 7 | |
| 8 | import torch |
| 9 | from torch import nn |
| 10 | |
| 11 | from vllm.config import VllmConfig |
| 12 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 13 | from vllm.model_executor.layers.quantization.compressed_tensors import ( |
| 14 | compressed_tensors, |
| 15 | ) |
| 16 | from vllm.model_executor.models.interfaces import ( |
| 17 | HasInnerState, |
| 18 | IsHybrid, |
| 19 | SupportsEagle3, |
| 20 | SupportsMultiModal, |
| 21 | SupportsPP, |
| 22 | SupportsQuant, |
| 23 | ) |
| 24 | from vllm.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs |
| 25 | from vllm.model_executor.models.kimi_k25_vit import ( |
| 26 | KimiK25MultiModalProjector, |
| 27 | MoonViT3dPretrainedModel, |
| 28 | vision_tower_forward, |
| 29 | ) |
| 30 | from vllm.model_executor.models.utils import ( |
| 31 | AutoWeightsLoader, |
| 32 | WeightsMapper, |
| 33 | init_vllm_registered_model, |
| 34 | maybe_prefix, |
| 35 | ) |
| 36 | from vllm.model_executor.models.vision import is_vit_use_data_parallel |
| 37 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 38 | from vllm.multimodal.inputs import NestedTensors |
| 39 | from vllm.platforms import current_platform |
| 40 | from vllm.sequence import IntermediateTensors |
| 41 | from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config |
| 42 | |
| 43 | from ..common.mm_preprocess import ( |
| 44 | KimiK3DummyInputsBuilder, |
| 45 | KimiK3MultiModalProcessor, |
| 46 | KimiK3ProcessingInfo, |
| 47 | ) |
| 48 | from .linear import KimiLinearForCausalLM |
| 49 | |
| 50 | |
| 51 | @MULTIMODAL_REGISTRY.register_processor( |
| 52 | KimiK3MultiModalProcessor, |
| 53 | info=KimiK3ProcessingInfo, |
| 54 | dummy_inputs=KimiK3DummyInputsBuilder, |
| 55 | ) |
| 56 | class KimiK3ForConditionalGeneration( |
| 57 | nn.Module, |
| 58 | SupportsMultiModal, |
| 59 | SupportsPP, |
| 60 | SupportsQuant, |
| 61 | SupportsEagle3, |
| 62 | HasInnerState, |
| 63 | IsHybrid, |
| 64 | ): |
| 65 | """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text.""" |
| 66 | |
| 67 | supports_encoder_tp_data = True |
| 68 | |
| 69 | hf_to_vllm_mapper = WeightsMapper( |
| 70 | orig_to_new_prefix={ |
| 71 | "language_model.layers.": "language_model.model.layers.", |
| 72 | "mm_projector.proj.0": "mm_projector.linear_1", |
| 73 | "mm_projector.proj.2": "mm_projector.linear_2", |
| 74 | } |
| 75 | ) |
| 76 | |
| 77 | @classmethod |
| 78 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 79 | if modality == "image": |
| 80 | return "<|kimi_image_placeholder|>" |
| 81 | raise ValueError(f"Unsupported modality: {modality}") |
| 82 | |
| 83 | def __init__( |
| 84 | self, |
| 85 | vllm_config: VllmConfig, |
| 86 | prefix: str = "", |
| 87 | ) -> None: |
| 88 | super().__init__() |
| 89 | model_config = vllm_config.model_config |
| 90 | config: KimiK3Config = model_config.hf_config |
| 91 | self.config = config |
| 92 | quant_config = vllm_config.quant_config |
| 93 | |
| 94 | multimodal_config = model_config.multimodal_config |
| 95 | assert multimodal_config is not None |
| 96 | self.use_data_parallel = is_vit_use_data_parallel( |
| 97 | config.vision_config.num_attention_heads |
| 98 | ) |
| 99 | self.hidden_size = config.text_config.hidden_size |
| 100 | self.device = current_platform.current_device() |
| 101 | |
| 102 | with self._mark_tower_model(vllm_config, "image"): |
| 103 | self.vision_tower = MoonViT3dPretrainedModel( |
| 104 | config.vision_config, |
| 105 | quant_config=self._maybe_ignore_quant_config(quant_config), |
| 106 | prefix=maybe_prefix(prefix, "vision_tower"), |
| 107 | ) |
| 108 | if self._maybe_ignore_quant_config(quant_config) is not None: |
| 109 | self.vision_tower = self.vision_tower.to(device=self.device) |
| 110 | else: |
| 111 | self.vision_tower = self.vision_tower.to( |
| 112 | device=self.device, dtype=model_config.dtype |
| 113 | ) |
| 114 | |
| 115 | self.mm_projector = KimiK25MultiModalProjector( |
| 116 | config=config.vision_config, |
| 117 | use_data_parallel=self.use_data_parallel, |
| 118 | quant_config=self._maybe_ignore_quant_config(quant_config), |
| 119 | prefix=maybe_prefix(prefix, "mm_projector"), |
| 120 | ) |
| 121 | self.mm_projector = self.mm_projector.to( |
| 122 | device=self.device, dtype=model_config.dtype |
| 123 | ) |
| 124 | |
| 125 | self.quant_config = quant_config |
| 126 | with self._mark_language_model(vllm_config): |
| 127 | self.language_model = init_vllm_registered_model( |
| 128 | vllm_config=vllm_config, |
| 129 | hf_config=config.text_config, |
| 130 | prefix=maybe_prefix(prefix, "language_model"), |
| 131 | architectures=["KimiLinearForCausalLM"], |
| 132 | ) |
| 133 | self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] |
| 134 | self.language_model.make_empty_intermediate_tensors |
| 135 | ) |
| 136 | self.media_placeholder: int = self.config.media_placeholder_token_id |
| 137 | |
| 138 | def _maybe_ignore_quant_config( |
| 139 | self, quant_config: QuantizationConfig | None |
| 140 | ) -> QuantizationConfig | None: |
| 141 | if isinstance(quant_config, compressed_tensors.CompressedTensorsConfig): |
| 142 | return None |
| 143 | return quant_config |
| 144 | |
| 145 | def _parse_and_validate_media_input( |
| 146 | self, **kwargs: object |
| 147 | ) -> KimiK25MediaPixelInputs | None: |
| 148 | pixel_values = kwargs.pop("pixel_values", None) |
| 149 | grid_thws = kwargs.pop("grid_thws", None) |
| 150 | if pixel_values is None: |
| 151 | return None |
| 152 | |
| 153 | if isinstance(pixel_values, list): |
| 154 | pixel_values = torch.cat(cast(list[torch.Tensor], pixel_values), dim=0) |
| 155 | if not isinstance(pixel_values, torch.Tensor): |
| 156 | raise TypeError( |
| 157 | "pixel_values must be a tensor or a list of tensors, " |
| 158 | f"got {type(pixel_values)}" |
| 159 | ) |
| 160 | |
| 161 | if len(pixel_values.shape) == 5 or len(pixel_values.shape) == 3: |
| 162 | pixel_values = pixel_values.reshape( |
| 163 | pixel_values.shape[0] * pixel_values.shape[1], *pixel_values.shape[2:] |
| 164 | ) |
| 165 | |
| 166 | target_dtype = next(self.vision_tower.parameters()).dtype |
| 167 | pixel_values = pixel_values.to(target_dtype) |
| 168 | assert isinstance(grid_thws, torch.Tensor), ( |
| 169 | f"expect grid_thws to be a tensor, got {type(grid_thws)}" |
| 170 | ) |
| 171 | grid_thws = grid_thws.reshape(-1, grid_thws.shape[-1]) |
| 172 | assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, ( |
| 173 | f"unexpected shape for grid_thws: {grid_thws.shape}" |
| 174 | ) |
| 175 | |
| 176 | return KimiK25MediaPixelInputs( |
| 177 | type="pixel_values", |
| 178 | pixel_values=pixel_values, |
| 179 | grid_thws=grid_thws, |
| 180 | ) |
| 181 | |
| 182 | def _process_media_input( |
| 183 | self, media_input: KimiK25MediaPixelInputs |
| 184 | ) -> list[torch.Tensor]: |
| 185 | media_features = vision_tower_forward( |
| 186 | self.vision_tower, |
| 187 | media_input["pixel_values"], |
| 188 | media_input["grid_thws"], |
| 189 | mm_projector=self.mm_projector, |
| 190 | use_data_parallel=self.use_data_parallel, |
| 191 | ) |
| 192 | return media_features |
| 193 | |
| 194 | def embed_multimodal(self, **kwargs: object) -> NestedTensors | None: |
| 195 | media_input = self._parse_and_validate_media_input(**kwargs) |
| 196 | if media_input is None: |
| 197 | return None |
| 198 | return self._process_media_input(media_input) |
| 199 | |
| 200 | def forward( |
| 201 | self, |
| 202 | input_ids: torch.Tensor, |
| 203 | positions: torch.Tensor, |
| 204 | intermediate_tensors: IntermediateTensors | None = None, |
| 205 | inputs_embeds: torch.Tensor | None = None, |
| 206 | **kwargs: object, |
| 207 | ) -> IntermediateTensors: |
| 208 | if intermediate_tensors is not None: |
| 209 | inputs_embeds = None |
| 210 | hidden_states = self.language_model( |
| 211 | input_ids=input_ids, |
| 212 | positions=positions, |
| 213 | intermediate_tensors=intermediate_tensors, |
| 214 | inputs_embeds=inputs_embeds, |
| 215 | ) |
| 216 | return hidden_states |
| 217 | |
| 218 | def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: |
| 219 | return self.language_model.compute_logits(hidden_states) |
| 220 | |
| 221 | def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): |
| 222 | return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs( |
| 223 | input_buffers, **kwargs |
| 224 | ) |
| 225 | |
| 226 | def get_seqlen_agnostic_capture_inputs(self, batch_size: int): |
| 227 | return self.language_model.mamba_cache.get_seqlen_agnostic_capture_inputs( |
| 228 | batch_size |
| 229 | ) |
| 230 | |
| 231 | @classmethod |
| 232 | def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig): |
| 233 | text_config = vllm_config.model_config.hf_config.text_config |
| 234 | temp_vllm_config = vllm_config.with_hf_config(text_config) |
| 235 | return KimiLinearForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config) |
| 236 | |
| 237 | @classmethod |
| 238 | def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig): |
| 239 | text_config = vllm_config.model_config.hf_config.text_config |
| 240 | temp_vllm_config = vllm_config.with_hf_config(text_config) |
| 241 | return KimiLinearForCausalLM.get_mamba_state_shape_from_config(temp_vllm_config) |
| 242 | |
| 243 | @classmethod |
| 244 | def get_mamba_state_copy_func(cls): |
| 245 | return KimiLinearForCausalLM.get_mamba_state_copy_func() |
| 246 | |
| 247 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): |
| 248 | loader = AutoWeightsLoader(self) |
| 249 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |