@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b539f988e1

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 9, 13:33 UTC

[Models] Kimi-K2.5 (#33131)

Signed-off-by: wanglinian <wanglinian@stu.pku.edu.cn> Signed-off-by: wangln19 <96399074+wangln19@users.noreply.github.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: youkaichao <youkaichao@gmail.com> Signed-off-by: Roger Wang <hey@rogerw.io> Co-authored-by: wanglinian <wanglinian@stu.pku.edu.cn> Co-authored-by: wangln19 <96399074+wangln19@users.noreply.github.com> Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: youkaichao <youkaichao@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Roger Wang committed Jan 26, 2026 at 22:50 UTC b539f988e1eeffe1c39bebbeaba892dc529eefaf

16 files changed +1799 -8

docs/models/supported_models.md

@@ -687,6 +687,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen

687 | `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + IE+ + VE+ | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ |

688 | `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + IE+ + VE+ | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ |

689 | `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I+ | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |

690 +| `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I+ | `moonshotai/Kimi-K2.5` | | ✅︎ |

691 | `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I+ | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ |

692 | `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I+ | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |

693 | `Llama4ForConditionalGeneration` | Llama 4 | T + I+ | `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. | ✅︎ | ✅︎ |

tests/models/registry.py

@@ -775,6 +775,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {

775 )

776 },

777 ),

778 + "KimiK25ForConditionalGeneration": _HfExamplesInfo(

779 + "moonshotai/Kimi-K2.5",

780 + trust_remote_code=True,

781 + is_available_online=False,

782 + ),

783 "LightOnOCRForConditionalGeneration": _HfExamplesInfo(

784 "lightonai/LightOnOCR-1B-1025"

785 ),

vllm/entrypoints/chat_utils.py

+182 -5

@@ -46,6 +46,9 @@ from vllm.multimodal.inputs import (

46 MultiModalBatchedField,

47 MultiModalFlatField,

48 MultiModalSharedField,

49 + VisionChunk,

50 + VisionChunkImage,

51 + VisionChunkVideo,

52 )

53 from vllm.multimodal.processing import BaseMultiModalProcessor

54 from vllm.multimodal.utils import MEDIA_CONNECTOR_REGISTRY, MediaConnector

@@ -336,7 +339,9 @@ ChatTemplateContentFormatOption = Literal["auto", "string", "openai"]

339 ChatTemplateContentFormat = Literal["string", "openai"]

340

341

339 -ModalityStr = Literal["image", "audio", "video", "image_embeds", "audio_embeds"]

342 +ModalityStr = Literal[

343 + "image", "audio", "video", "image_embeds", "audio_embeds", "vision_chunk"

344 +]

345 _T = TypeVar("_T")

346

347

@@ -449,6 +454,78 @@ def _get_embeds_data(

454 raise NotImplementedError(type(data_items))

455

456

457 +def rebuild_mm_uuids_from_mm_data(

458 + mm_uuids: MultiModalUUIDDict,

459 + mm_data: MultiModalDataDict,

460 +) -> MultiModalUUIDDict:

461 + """Rebuild mm_uuids after vision_chunk processing.

462 +

463 + When videos are split into chunks, the original UUIDs need to be updated

464 + to reflect the new UUIDs generated for each chunk.

465 +

466 + Args:

467 + mm_uuids: Original UUIDs dictionary

468 + mm_data: Processed multimodal data with vision_chunk items

469 +

470 + Returns:

471 + Updated UUIDs dictionary with chunk UUIDs

472 + """

473 + vision_chunks = mm_data.get("vision_chunk")

474 + if vision_chunks is None:

475 + return mm_uuids

476 +

477 + new_uuids = dict(mm_uuids)

478 + vision_chunk_uuids = []

479 +

480 + for item in vision_chunks:

481 + # vision_chunk items are always dicts (VisionChunkImage/VisionChunkVideo)

482 + assert isinstance(item, dict)

483 + uuid_val = item.get("uuid")

484 + if uuid_val is not None:

485 + vision_chunk_uuids.append(uuid_val)

486 +

487 + if vision_chunk_uuids:

488 + new_uuids["vision_chunk"] = vision_chunk_uuids

489 +

490 + return new_uuids

491 +

492 +

493 +def build_video_prompts_from_mm_data(

494 + mm_data: MultiModalDataDict,

495 +) -> list[str]:

496 + """Build video prompts from vision_chunk data.

497 +

498 + Collects prompts from video chunks and groups them by video_idx.

499 +

500 + Args:

501 + mm_data: Processed multimodal data with vision_chunk items

502 +

503 + Returns:

504 + List of video prompts, one per video.

505 + """

506 + vision_chunks = mm_data.get("vision_chunk")

507 + if vision_chunks is None:

508 + return []

509 +

510 + # Group chunks by video_idx

511 + video_prompts_dict: dict[int, list[str]] = defaultdict(list)

512 +

513 + for item in vision_chunks:

514 + # vision_chunk items are always dicts (VisionChunkImage/VisionChunkVideo)

515 + assert isinstance(item, dict)

516 + if item.get("type") == "video_chunk":

517 + video_idx = item.get("video_idx", 0)

518 + prompt = item.get("prompt", "")

519 + video_prompts_dict[video_idx].append(prompt)

520 +

521 + # Build prompts in video order

522 + video_prompts = []

523 + for video_idx in sorted(video_prompts_dict.keys()):

524 + video_prompts.append("".join(video_prompts_dict[video_idx]))

525 +

526 + return video_prompts

527 +

528 +

529 class BaseMultiModalItemTracker(ABC, Generic[_T]):

530 """

531 Tracks multi-modal items in a given request and ensures that the number

@@ -462,6 +539,13 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):

539 self._model_config = model_config

540

541 self._items_by_modality = defaultdict[str, list[_T]](list)

542 + # Track original modality for each vision_chunk item (image or video)

543 + self._modality_order = defaultdict[str, list[str]](list)

544 +

545 + @cached_property

546 + def use_unified_vision_chunk_modality(self) -> bool:

547 + """Check if model uses unified vision_chunk modality for images/videos."""

548 + return getattr(self._model_config.hf_config, "use_unified_vision_chunk", False)

549

550 @property

551 def model_config(self) -> ModelConfig:

@@ -499,11 +583,31 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):

583 media.

584 """

585 input_modality = modality.replace("_embeds", "")

502 - num_items = len(self._items_by_modality[modality]) + 1

586 + original_modality = modality

587 + use_vision_chunk = (

588 + self.use_unified_vision_chunk_modality

589 + and original_modality in ["video", "image"]

590 + )

591 +

592 + # If use_unified_vision_chunk_modality is enabled,

593 + # map image/video to vision_chunk

594 + if use_vision_chunk:

595 + # To avoid validation fail

596 + # because models with use_unified_vision_chunk_modality=True

597 + # will only accept vision_chunk modality.

598 + input_modality = "vision_chunk"

599 + num_items = len(self._items_by_modality[input_modality]) + 1

600 + else:

601 + num_items = len(self._items_by_modality[original_modality]) + 1

602

603 self.mm_processor.validate_num_items(input_modality, num_items)

604

506 - self._items_by_modality[modality].append(item)

605 + # Track original modality for vision_chunk items

606 + if use_vision_chunk:

607 + self._items_by_modality[input_modality].append(item) # type: ignore

608 + self._modality_order["vision_chunk"].append(original_modality)

609 + else:

610 + self._items_by_modality[original_modality].append(item)

611

612 return self.model_cls.get_placeholder_str(modality, num_items)

613

@@ -515,6 +619,7 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):

619 def _resolve_items(

620 items_by_modality: dict[str, list[tuple[object, str | None]]],

621 mm_processor: BaseMultiModalProcessor,

622 + vision_chunk_modality_order: dict[str, list[str]],

623 ) -> tuple[MultiModalDataDict, MultiModalUUIDDict]:

624 if "image" in items_by_modality and "image_embeds" in items_by_modality:

625 raise ValueError("Mixing raw image and embedding inputs is not allowed")

@@ -546,6 +651,74 @@ def _resolve_items(

651 if "video" in items_by_modality:

652 mm_data["video"] = [data for data, uuid in items_by_modality["video"]]

653 mm_uuids["video"] = [uuid for data, uuid in items_by_modality["video"]]

654 + if "vision_chunk" in items_by_modality:

655 + # Process vision_chunk items - extract from (data, modality) tuples

656 + # and convert to VisionChunk types with proper UUID handling

657 + vision_chunk_items = items_by_modality["vision_chunk"]

658 + modality_order = vision_chunk_modality_order.get("vision_chunk", [])

659 + mm_uuids["vision_chunk"] = [

660 + uuid for data, uuid in items_by_modality["vision_chunk"]

661 + ]

662 +

663 + # Filter out None items (from asyncio.sleep(0) placeholders)

664 + filtered_items = [

665 + (idx, item)

666 + for idx, item in enumerate(vision_chunk_items)

667 + if item is not None

668 + ]

669 +

670 + assert len(filtered_items) == len(modality_order), (

671 + f"vision_chunk items ({len(filtered_items)}) and "

672 + f"modality_order ({len(modality_order)}) must have same length"

673 + )

674 +

675 + processed_chunks: list[VisionChunk] = []

676 + video_idx = 0

677 + for i, (idx, item) in enumerate(filtered_items):

678 + inner_modality = modality_order[i]

679 + data, uuid = item

680 + uuid_val = uuid if idx < len(mm_uuids["vision_chunk"]) else None

681 + if inner_modality == "image":

682 + # Cast data to proper type for image

683 + # Use .media (PIL.Image) directly to avoid redundant

684 + # bytes→PIL conversion in media_processor

685 + if hasattr(data, "media"):

686 + image_data = data.media # type: ignore[union-attr]

687 + processed_chunks.append(

688 + VisionChunkImage(type="image", image=image_data, uuid=uuid_val)

689 + )

690 + else:

691 + processed_chunks.append(data) # type: ignore[arg-type]

692 + elif inner_modality == "video":

693 + # For video, we may need to split into chunks

694 + # if processor supports it

695 + # For now, just wrap as a video chunk placeholder

696 + if hasattr(mm_processor, "split_video_chunks") and data is not None:

697 + try:

698 + video_uuid = uuid_val or random_uuid()

699 + # video await result is (video_data, video_meta) tuple

700 + if isinstance(data, tuple) and len(data) >= 1:

701 + video_data = data[0]

702 + else:

703 + video_data = data

704 + video_chunks = mm_processor.split_video_chunks(video_data)

705 + for i, vc in enumerate(video_chunks):

706 + processed_chunks.append(

707 + VisionChunkVideo(

708 + type="video_chunk",

709 + video_chunk=vc["video_chunk"],

710 + uuid=f"{video_uuid}-{i}",

711 + video_idx=video_idx,

712 + prompt=vc["prompt"],

713 + )

714 + )

715 + video_idx += 1

716 + except Exception as e:

717 + logger.warning("Failed to split video chunks: %s", e)

718 + processed_chunks.append(data) # type: ignore[arg-type]

719 + else:

720 + processed_chunks.append(data) # type: ignore[arg-type]

721 + mm_data["vision_chunk"] = processed_chunks

722

723 return mm_data, mm_uuids

724

@@ -557,7 +730,9 @@ class MultiModalItemTracker(BaseMultiModalItemTracker[tuple[object, str | None]]

730 if not self._items_by_modality:

731 return None, None

732

560 - return _resolve_items(dict(self._items_by_modality), self.mm_processor)

733 + return _resolve_items(

734 + dict(self._items_by_modality), self.mm_processor, self._modality_order

735 + )

736

737 def create_parser(self) -> "BaseMultiModalContentParser":

738 return MultiModalContentParser(self)

@@ -577,7 +752,9 @@ class AsyncMultiModalItemTracker(

752 for modality, coros in self._items_by_modality.items()

753 }

754

580 - return _resolve_items(resolved_items_by_modality, self.mm_processor)

755 + return _resolve_items(

756 + resolved_items_by_modality, self.mm_processor, self._modality_order

757 + )

758

759 def create_parser(self) -> "BaseMultiModalContentParser":

760 return AsyncMultiModalContentParser(self)

vllm/envs.py

@@ -778,6 +778,7 @@ environment_variables: dict[str, Callable[[], Any]] = {

778 ),

779 # Backend for Video IO

780 # - "opencv": Default backend that uses OpenCV stream buffered backend.

781 + # - "identity": Returns raw video bytes for model processor to handle.

782 #

783 # Custom backend implementations can be registered

784 # via `@VIDEO_LOADER_REGISTRY.register("my_custom_video_loader")` and

vllm/model_executor/models/kimi_k25.py new

+581

@@ -0,0 +1,581 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +# ruff: noqa: E501

4 +"""

5 +Kimi-K2.5 Model Implementation for vLLM.

6 +

7 +Kimi-K2.5 extends Kimi-K2 with vision support

8 +

9 +This module defines:

10 +- KimiK25ProcessingInfo/KimiK25MultiModalProcessor: Processing logic

11 +- KimiK25ForConditionalGeneration: Main model class

12 +"""

13 +

14 +import copy

15 +from collections.abc import Iterable, Mapping, Sequence

16 +from dataclasses import dataclass

17 +from typing import Annotated, Any, Literal

18 +

19 +import torch

20 +from torch import nn

21 +from transformers import BatchFeature

22 +from transformers.processing_utils import ProcessorMixin

23 +

24 +from vllm.config import VllmConfig

25 +from vllm.config.multimodal import BaseDummyOptions

26 +from vllm.distributed import get_pp_group

27 +from vllm.logger import init_logger

28 +from vllm.model_executor.layers.fused_moe import SharedFusedMoE

29 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

30 +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

31 +from vllm.model_executor.model_loader.weight_utils import (

32 + default_weight_loader,

33 + maybe_remap_kv_scale_name,

34 +)

35 +from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model

36 +from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP

37 +from vllm.model_executor.models.kimi_k25_vit import (

38 + KimiK25MultiModalProjector,

39 + MoonViT3dPretrainedModel,

40 + vision_tower_forward,

41 +)

42 +from vllm.multimodal import MULTIMODAL_REGISTRY

43 +from vllm.multimodal.inputs import (

44 + MultiModalDataDict,

45 + MultiModalFieldConfig,

46 + MultiModalKwargsItems,

47 + NestedTensors,

48 + VisionChunk,

49 + VisionChunkImage,

50 + VisionChunkVideo,

51 +)

52 +from vllm.multimodal.parse import MultiModalDataItems, VisionChunkProcessorItems

53 +from vllm.multimodal.processing import (

54 + BaseDummyInputsBuilder,

55 + BaseMultiModalProcessor,

56 + BaseProcessingInfo,

57 + InputProcessingContext,

58 + PromptReplacement,

59 + PromptUpdate,

60 +)

61 +from vllm.sequence import IntermediateTensors

62 +from vllm.transformers_utils.configs import KimiK25Config

63 +from vllm.transformers_utils.processor import cached_get_image_processor

64 +from vllm.utils.tensor_schema import TensorSchema, TensorShape

65 +

66 +from .utils import PPMissingLayer, is_pp_missing_parameter, maybe_prefix

67 +

68 +logger = init_logger(__name__)

69 +

70 +

71 +# Dummy input dimensions for profiling.

72 +@dataclass

73 +class MaxImageTokenMeta:

74 + width: int = 3000

75 + height: int = 3000

76 +

77 +

78 +class KimiK25MediaPixelInputs(TensorSchema):

79 + """

80 + Media input schema for K2-VL model.

81 +

82 + Dimensions:

83 + - np: Number of patches (flattened from all media items)

84 + - ps: Patch size

85 + - nm: Number of media items

86 + """

87 +

88 + type: Literal["pixel_values"] = "pixel_values"

89 +

90 + pixel_values: Annotated[

91 + torch.Tensor | list[torch.Tensor],

92 + TensorShape("np", 3, "ps", "ps"),

93 + ]

94 +

95 + grid_thws: Annotated[torch.Tensor, TensorShape("nm", 3)]

96 +

97 +

98 +class MoonshotKimiVAutoProcessor(ProcessorMixin):

99 + attributes = ["tokenizer"]

100 + tokenizer_class = "AutoTokenizer"

101 +

102 + def __init__(self, media_processor=None, tokenizer=None):

103 + super().__init__(tokenizer)

104 + self.media_processor = media_processor

105 +

106 + # We do not support str input for text here

107 + def __call__(

108 + self,

109 + vision_chunks: list[VisionChunk] | None = None,

110 + *,

111 + text: list[int],

112 + **kwargs,

113 + ) -> BatchFeature:

114 + """

115 + Args:

116 + vision_chunks: List of VisionChunk items to be processed.

117 + For image: VisionChunkImage with type='image', image=PIL.Image

118 + For video_chunk: VisionChunkVideo with type='video_chunk', video_chunk=list[PIL.Image]

119 + text: The token ids to be fed to a model (required).

120 + Returns:

121 + [`BatchFeature`]: A [`BatchFeature`] with the following fields:

122 +

123 + - **input_ids** -- list of token ids to be fed to a model.

124 + - **pixel_values** -- Pixel values to be fed to a model. Returned when `vision_chunks` is not `None`.

125 + - **grid_thws** -- list of image 3D grid in LLM. Returned when `vision_chunks` is not `None`.

126 + """

127 + mm_inputs = {}

128 + if vision_chunks is not None:

129 + assert isinstance(vision_chunks, list)

130 + mm_inputs = self.media_processor.preprocess(vision_chunks)

131 + # XXX: _apply_hf_processor_text_mm will call tolist() on input_ids

132 + return BatchFeature(

133 + data={

134 + "input_ids": torch.tensor([text]),

135 + **mm_inputs,

136 + }

137 + )

138 +

139 +

140 +class KimiK25ProcessingInfo(BaseProcessingInfo):

141 + """Processing information for Kimi-K2.5 model.

142 +

143 + Provides configuration and utilities for processing both

144 + images and video-chunks.

145 + """

146 +

147 + def __init__(self, ctx: InputProcessingContext) -> None:

148 + super().__init__(ctx)

149 + self.hf_config = self.get_hf_config()

150 + self.media_token_id = self.hf_config.media_placeholder_token_id

151 + media_processor = cached_get_image_processor(

152 + self.ctx.model_config.model, trust_remote_code=True

153 + )

154 + self.media_processor = media_processor

155 + self.hf_processor = MoonshotKimiVAutoProcessor(

156 + media_processor=self.media_processor,

157 + tokenizer=self.get_tokenizer(),

158 + )

159 + self.media_tokens_calculator = self.media_processor.media_tokens_calculator

160 +

161 + def get_hf_processor(self):

162 + return self.hf_processor

163 +

164 + def get_hf_config(self):

165 + return self.ctx.get_hf_config(KimiK25Config)

166 +

167 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

168 + # None means unlimited

169 + return {"vision_chunk": None}

170 +

171 +

172 +class KimiK25DummyInputsBuilder(BaseDummyInputsBuilder[KimiK25ProcessingInfo]):

173 + """Builds dummy inputs for Kimi-K2.5 model profiling."""

174 +

175 + def __init__(self, info: KimiK25ProcessingInfo) -> None:

176 + super().__init__(info)

177 + self.media_token_id = self.info.media_token_id

178 + self.frame_per_chunk = self.info.media_processor.num_frames_per_chunk

179 +

180 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> list[int]:

181 + num_media = mm_counts.get("vision_chunk", 0)

182 + return [self.media_token_id] * num_media

183 +

184 + def get_dummy_mm_items(self):

185 + dummy_videos = self._get_dummy_images(

186 + height=MaxImageTokenMeta.height,

187 + width=MaxImageTokenMeta.width,

188 + num_images=self.frame_per_chunk,

189 + )

190 +

191 + video_chunk_dummy_item = VisionChunkVideo(

192 + type="video_chunk", video_chunk=dummy_videos

193 + )

194 + video_chunk_num_tokens = self.info.media_tokens_calculator(

195 + video_chunk_dummy_item

196 + )

197 +

198 + image_dummy_item = VisionChunkImage(

199 + type="image",

200 + image=self._get_dummy_images(

201 + height=MaxImageTokenMeta.height,

202 + width=MaxImageTokenMeta.width,

203 + num_images=1,

204 + )[0],

205 + )

206 + image_num_tokens = self.info.media_tokens_calculator(image_dummy_item)

207 + # return the larger one

208 + if video_chunk_num_tokens >= image_num_tokens:

209 + return [video_chunk_dummy_item]

210 + else:

211 + return [image_dummy_item]

212 +

213 + def get_dummy_mm_data(

214 + self,

215 + seq_len: int,

216 + mm_counts: Mapping[str, int],

217 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

218 + ) -> MultiModalDataDict:

219 + # TODO: Support mm_options for vision_chunk to allow user configuration

220 + dummy_items = self.get_dummy_mm_items()

221 + return {"vision_chunk": dummy_items}

222 +

223 +

224 +class KimiK25MultiModalProcessor(BaseMultiModalProcessor[KimiK25ProcessingInfo]):

225 + """Multi-modal processor for Kimi-K2.5.

226 +

227 + Handles both image and video-chunk modalities.

228 + """

229 +

230 + def _get_mm_fields_config(

231 + self,

232 + hf_inputs: BatchFeature,

233 + hf_processor_mm_kwargs: Mapping[str, object],

234 + ) -> Mapping[str, MultiModalFieldConfig]:

235 + """Indicates how to slice media input into multiple items.

236 +

237 + pixel_values: [N, 3, patch_size, patch_size], all patches collected from B medias

238 + grid_thws: [B,3], each item: [N_t, N_h ,N_w], indicates the grid size in time/height/width direction

239 + for current item.

240 +

241 + by multiplying [N_t, N_h ,N_w], we get the number of patches for each media item, thus we can slice

242 + pixel_values by pixel_values[start:start + N_t*N_h*N_w] to get patches of one item.

243 +

244 + """

245 + grid_thws = hf_inputs.get("grid_thws", torch.empty((0, 3)))

246 + grid_sizes = grid_thws.prod(-1)

247 +

248 + return dict(

249 + pixel_values=MultiModalFieldConfig.flat_from_sizes(

250 + "vision_chunk", grid_sizes

251 + ),

252 + grid_thws=MultiModalFieldConfig.batched("vision_chunk"),

253 + )

254 +

255 + def _get_prompt_updates(

256 + self,

257 + mm_items: MultiModalDataItems,

258 + hf_processor_mm_kwargs: Mapping[str, Any],

259 + out_mm_kwargs: MultiModalKwargsItems,

260 + ) -> Sequence[PromptUpdate]:

261 + hf_config = self.info.get_hf_config()

262 + media_token_id = hf_config.media_placeholder_token_id

263 +

264 + def get_replacement(item_idx: int):

265 + media = mm_items.get_items("vision_chunk", (VisionChunkProcessorItems,))

266 + num_media_token = self.info.media_tokens_calculator(media[item_idx])

267 + return [media_token_id] * num_media_token

268 +

269 + return [

270 + PromptReplacement(

271 + modality="vision_chunk",

272 + target=[media_token_id],

273 + replacement=get_replacement,

274 + ),

275 + ]

276 +

277 + def split_video_chunks(self, video):

278 + return self.info.media_processor.split_video_chunks(video)

279 +

280 +

281 +@MULTIMODAL_REGISTRY.register_processor(

282 + KimiK25MultiModalProcessor,

283 + info=KimiK25ProcessingInfo,

284 + dummy_inputs=KimiK25DummyInputsBuilder,

285 +)

286 +class KimiK25ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

287 + """Kimi-K2.5 model for conditional generation.

288 +

289 + Supports both image and video-chunk modalities.

290 + Video-chunks are temporal segments (typically 4 frames) that are

291 + processed with temporal pooling.

292 + """

293 +

294 + supports_encoder_tp_data = True

295 +

296 + @classmethod

297 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

298 + # Kimi-K2.5 uses video_chunk for all media types

299 + if modality == "image":

301 + elif modality == "video":

302 + # return a placeholder, to be replaced in the future.

303 + return "<|kimi_k25_video_placeholder|>"

304 +

305 + raise ValueError(f"Unsupported modality: {modality}")

306 +

307 + def __init__(

308 + self,

309 + vllm_config: VllmConfig,

310 + prefix: str = "",

311 + ) -> None:

312 + super().__init__()

313 + model_config = vllm_config.model_config

314 + config: KimiK25Config = model_config.hf_config

315 + self.config = config

316 + quant_config = vllm_config.quant_config

317 +

318 + # Check for MoonViT config compatibility

319 + self.use_data_parallel = (

320 + model_config.multimodal_config.mm_encoder_tp_mode == "data"

321 + )

322 + self.hidden_size = config.text_config.hidden_size

323 + self.device = torch.cuda.current_device()

324 + # Build vision tower directly with KimiK25VisionConfig

325 + self.vision_tower = MoonViT3dPretrainedModel(

326 + config.vision_config,

327 + prefix=maybe_prefix(prefix, "vision_tower"),

328 + )

329 + self.vision_tower = self.vision_tower.to(

330 + device=self.device, dtype=model_config.dtype

331 + )

332 +

333 + self.mm_projector = KimiK25MultiModalProjector(

334 + config=config.vision_config,

335 + use_data_parallel=self.use_data_parallel,

336 + prefix=maybe_prefix(prefix, "mm_projector"),

337 + )

338 + self.mm_projector = self.mm_projector.to(

339 + device=self.device, dtype=model_config.dtype

340 + )

341 +

342 + self.quant_config = quant_config

343 + sub_vllm_config = copy.deepcopy(vllm_config)

344 + sub_vllm_config.model_config.hf_config = (

345 + sub_vllm_config.model_config.hf_config.text_config

346 + )

347 + self.language_model = DeepseekV2Model(

348 + vllm_config=sub_vllm_config,

349 + prefix=maybe_prefix(prefix, "language_model"),

350 + )

351 + if get_pp_group().is_last_rank:

352 + self.lm_head = ParallelLMHead(

353 + config.vocab_size,

354 + config.text_config.hidden_size,

355 + prefix=maybe_prefix(prefix, "lm_head"),

356 + )

357 + else:

358 + self.lm_head = PPMissingLayer()

359 + self.make_empty_intermediate_tensors = (

360 + self.language_model.make_empty_intermediate_tensors

361 + )

362 + logit_scale = getattr(config, "logit_scale", 1.0)

363 + self.logits_processor = LogitsProcessor(config.vocab_size, scale=logit_scale)

364 + self.media_placeholder: int = self.config.media_placeholder_token_id

365 +

366 + def _parse_and_validate_media_input(

367 + self, **kwargs: object

368 + ) -> KimiK25MediaPixelInputs | None:

369 + pixel_values = kwargs.pop("pixel_values", None)

370 + grid_thws = kwargs.pop("grid_thws", None)

371 + if pixel_values is None:

372 + return None

373 +

374 + if isinstance(pixel_values, list):

375 + pixel_values = torch.cat(pixel_values, dim=0)

376 +

377 + if len(pixel_values.shape) == 5 or len(pixel_values.shape) == 3:

378 + pixel_values = pixel_values.reshape(

379 + pixel_values.shape[0] * pixel_values.shape[1], *pixel_values.shape[2:]

380 + )

381 +

382 + # The batch dimension of pixel_values has been flattened into shape[0]

383 + target_dtype = next(self.vision_tower.parameters()).dtype

384 + pixel_values = pixel_values.to(target_dtype)

385 + assert isinstance(grid_thws, torch.Tensor), (

386 + f"expect grid_thws to be a tensor, get {type(grid_thws)}"

387 + )

388 + # In some cases (e.g. with merger), grid_thws has an extra middle dimension

389 + grid_thws = grid_thws.reshape(-1, grid_thws.shape[-1])

390 + assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, (

391 + f"unexpected shape for grid_thws: {grid_thws.shape}"

392 + )

393 +

394 + return KimiK25MediaPixelInputs(

395 + type="pixel_values",

396 + pixel_values=pixel_values,

397 + grid_thws=grid_thws,

398 + )

399 +

400 + def _process_media_input(

401 + self, media_input: KimiK25MediaPixelInputs

402 + ) -> list[torch.Tensor]:

403 + # NOTE(moyan): This forward will automatically batch the forward pass internally

404 + media_features = vision_tower_forward(

405 + self.vision_tower,

406 + media_input["pixel_values"],

407 + media_input["grid_thws"],

408 + mm_projector=self.mm_projector,

409 + use_data_parallel=self.use_data_parallel,

410 + )

411 + return media_features

412 +

413 + def embed_multimodal(self, **kwargs: object) -> NestedTensors | None:

414 + # Validate the multimodal input keyword arguments

415 + media_input = self._parse_and_validate_media_input(**kwargs)

416 + if media_input is None:

417 + return None

418 +

419 + # Run multimodal inputs through encoder and projector

420 + vision_embeddings = self._process_media_input(media_input)

421 + return vision_embeddings

422 +

423 + def get_language_model(self) -> torch.nn.Module:

424 + return self.language_model

425 +

426 + def forward(

427 + self,

428 + input_ids: torch.Tensor,

429 + positions: torch.Tensor,

430 + intermediate_tensors: IntermediateTensors | None = None,

431 + inputs_embeds: torch.Tensor | None = None,

432 + **kwargs: object,

433 + ) -> IntermediateTensors:

434 + if intermediate_tensors is not None:

435 + inputs_embeds = None

436 + hidden_states = self.language_model(

437 + input_ids=input_ids,

438 + positions=positions,

439 + intermediate_tensors=intermediate_tensors,

440 + inputs_embeds=inputs_embeds,

441 + )

442 +

443 + return hidden_states

444 +

445 + def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:

446 + logits = self.logits_processor(self.lm_head, hidden_states, **kwargs)

447 + return logits

448 +

449 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

450 + # Params for weights, fp8 weight scales, fp8 activation scales

451 + # (param_name, weight_name, expert_id, shard_id)

452 + config = self.config.text_config

453 + if not getattr(config, "n_routed_experts", None):

454 + return []

455 + return SharedFusedMoE.make_expert_params_mapping(

456 + self,

457 + ckpt_gate_proj_name="gate_proj",

458 + ckpt_down_proj_name="down_proj",

459 + ckpt_up_proj_name="up_proj",

460 + num_experts=config.n_routed_experts,

461 + num_redundant_experts=0,

462 + )

463 +

464 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

465 + config = self.config.text_config

466 + _KEYS_TO_MODIFY_MAPPING = {

467 + "language_model.lm_head": "lm_head",

468 + "language_model.model": "language_model",

469 + # mm_projector -> mm_projector mapping

470 + # "mm_projector": "mm_projector",

471 + "mm_projector.proj.0": "mm_projector.linear_1",

472 + "mm_projector.proj.2": "mm_projector.linear_2",

473 + }

474 + stacked_params_mapping = [

475 + (".gate_up_proj", ".gate_proj", 0),

476 + (".gate_up_proj", ".up_proj", 1),

477 + ]

478 + if getattr(config, "kv_lora_rank", None) and getattr(

479 + config, "q_lora_rank", None

480 + ):

481 + stacked_params_mapping += [

482 + (".fused_qkv_a_proj", ".q_a_proj", 0),

483 + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),

484 + ]

485 + expert_params_mapping = self.get_expert_mapping()

486 +

487 + params_dict = dict(self.named_parameters())

488 +

489 + for args in weights:

490 + name, loaded_weight = args[:2]

491 + kwargs = args[2] if len(args) > 2 else {}

492 + if "rotary_emb.inv_freq" in name:

493 + continue

494 +

495 + spec_layer = get_spec_layer_idx_from_weight_name(config, name)

496 + if spec_layer is not None:

497 + continue # skip spec decode layers for main model

498 +

499 + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:

500 + continue

501 +

502 + for key_to_modify, new_key in _KEYS_TO_MODIFY_MAPPING.items():

503 + if key_to_modify in name:

504 + name = name.replace(key_to_modify, new_key)

505 +

506 + use_default_weight_loading = False

507 + if "vision" in name:

508 + if self.vision_tower is not None:

509 + use_default_weight_loading = True

510 + else:

511 + for param_name, weight_name, shard_id in stacked_params_mapping:

512 + if weight_name not in name:

513 + continue

514 + if ("mlp.experts." in name) and name not in params_dict:

515 + continue

516 + name = name.replace(weight_name, param_name)

517 + if name.endswith(".bias") and name not in params_dict:

518 + continue

519 +

520 + if is_pp_missing_parameter(name, self):

521 + continue

522 +

523 + param = params_dict[name]

524 + weight_loader = param.weight_loader

525 + weight_loader(param, loaded_weight, shard_id, **kwargs)

526 + break

527 + else:

528 + for _, (

529 + param_name,

530 + weight_name,

531 + expert_id,

532 + shard_id,

533 + ) in enumerate(expert_params_mapping):

534 + if weight_name not in name:

535 + continue

536 + name = name.replace(weight_name, param_name)

537 +

538 + if is_pp_missing_parameter(name, self):

539 + continue

540 +

541 + param = params_dict[name]

542 + weight_loader = param.weight_loader

543 + weight_loader(

544 + param,

545 + loaded_weight,

546 + name,

547 + expert_id=expert_id,

548 + shard_id=shard_id,

549 + **kwargs,

550 + )

551 + break

552 + else:

553 + use_default_weight_loading = True

554 +

555 + if use_default_weight_loading:

556 + if name.endswith(".bias") and name not in params_dict:

557 + continue

558 + name = maybe_remap_kv_scale_name(name, params_dict)

559 + if name is None:

560 + continue

561 +

562 + if is_pp_missing_parameter(name, self):

563 + continue

564 +

565 + param = params_dict[name]

566 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

567 + weight_loader(param, loaded_weight, **kwargs)

568 +

569 +

570 +def get_spec_layer_idx_from_weight_name(

571 + config: KimiK25Config, weight_name: str

572 +) -> int | None:

573 + if hasattr(config, "num_nextn_predict_layers") and (

574 + config.num_nextn_predict_layers > 0

575 + ):

576 + layer_idx = config.num_hidden_layers

577 + for i in range(config.num_nextn_predict_layers):

578 + # might start with language_model.model.layers

579 + if f"model.layers.{layer_idx + i}." in weight_name:

580 + return layer_idx + i

581 + return None

vllm/model_executor/models/kimi_k25_vit.py new

+678

@@ -0,0 +1,678 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""

4 +Vision tower implementation for Kimi-K2.5 model.

5 +

6 +This module provides the vision encoder components for Kimi-K2.5,

7 +including 3D patch embedding, RoPE position embedding, and

8 +temporal pooling for video chunks.

9 +"""

10 +

11 +from collections.abc import Sequence

12 +from copy import deepcopy

13 +from typing import Any

14 +

15 +import numpy as np

16 +import torch

17 +import torch.nn as nn

18 +import torch.nn.functional as F

19 +from transformers.activations import GELUActivation

20 +

21 +from vllm.distributed import divide, get_tensor_model_parallel_world_size

22 +from vllm.logger import init_logger

23 +from vllm.model_executor.layers.activation import get_act_fn

24 +from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention

25 +from vllm.model_executor.layers.linear import (

26 + ColumnParallelLinear,

27 + QKVParallelLinear,

28 + ReplicatedLinear,

29 + RowParallelLinear,

30 +)

31 +from vllm.model_executor.models.utils import maybe_prefix

32 +from vllm.model_executor.models.vision import (

33 + is_vit_use_data_parallel,

34 + run_dp_sharded_mrope_vision_model,

35 +)

36 +from vllm.transformers_utils.configs.kimi_k25 import KimiK25VisionConfig

37 +

38 +logger = init_logger(__name__)

39 +

40 +

41 +def _apply_rope_input_validation(x, freqs_cis):

42 + assert x.ndim == freqs_cis.ndim + 1, (x.shape, freqs_cis.shape)

43 + assert x.shape[:-2] == freqs_cis.shape[:-1], (x.shape, freqs_cis.shape)

44 + assert x.shape[-1] == 2 * freqs_cis.shape[-1], (x.shape, freqs_cis.shape)

45 + assert freqs_cis.dtype == torch.complex64, freqs_cis.dtype

46 +

47 +

48 +def get_rope_shape_decorate(func):

49 + _get_rope_shape_first_call_flag = set()

50 +

51 + def wrapper(org, interpolation_mode, shape):

52 + key = (org.requires_grad, torch.is_grad_enabled(), interpolation_mode)

53 + if key not in _get_rope_shape_first_call_flag:

54 + _get_rope_shape_first_call_flag.add(key)

55 + _ = func(org, interpolation_mode, shape=(64, 64))

56 + return func(org, interpolation_mode, shape)

57 +

58 + return wrapper

59 +

60 +

61 +@get_rope_shape_decorate

62 +@torch.compile(dynamic=True)

63 +def get_rope_shape(org, interpolation_mode, shape):

64 + return (

65 + F.interpolate(

66 + org.permute((2, 0, 1)).unsqueeze(0),

67 + size=shape,

68 + mode=interpolation_mode,

69 + )

70 + .squeeze(0)

71 + .permute((1, 2, 0))

72 + .flatten(end_dim=1)

73 + )

74 +

75 +

76 +def apply_rope(

77 + xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor

78 +) -> tuple[torch.Tensor, torch.Tensor]:

79 + """

80 + Args: (The leading dimensions of all inputs should be the same)

81 + xq: query, tensor of shape (..., num_heads, head_dim)

82 + xk: key, tensor of shape (..., num_heads, head_dim)

83 + freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64.

84 + Returns:

85 + xq_out, xk_out: tensors of shape (..., num_heads, head_dim)

86 + """

87 + _apply_rope_input_validation(xq, freqs_cis)

88 + _apply_rope_input_validation(xk, freqs_cis)

89 +

90 + freqs_cis = freqs_cis.unsqueeze(-2) # ..., 1, head_dim/2

91 + # ..., num_heads, head_dim/2

92 + xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2))

93 + xk_ = torch.view_as_complex(xk.float().view(*xq.shape[:-1], -1, 2))

94 + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2) # ..., num_heads, head_dim

95 + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2) # ..., num_heads, head_dim

96 + return xq_out.type_as(xq), xk_out.type_as(xk)

97 +

98 +

99 +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):

100 + """Generate 1D sincos positional embedding from grid positions."""

101 + assert embed_dim % 2 == 0

102 + omega = np.arange(embed_dim // 2, dtype=np.float32)

103 + omega /= embed_dim / 2.0

104 + omega = 1.0 / 10000**omega # (D/2,)

105 +

106 + pos = pos.reshape(-1) # (M,)

107 + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product

108 +

109 + emb_sin = np.sin(out) # (M, D/2)

110 + emb_cos = np.cos(out) # (M, D/2)

111 +

112 + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)

113 + return emb

114 +

115 +

116 +def get_1d_sincos_pos_embed(embed_dim, t_size, cls_token=False):

117 + """Generate 1D sincos positional embedding."""

118 + grid_t = np.arange(t_size, dtype=np.float32)

119 + pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid_t)

120 + if cls_token:

121 + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)

122 + return pos_embed

123 +

124 +

125 +class Learnable2DInterpPosEmbDivided_fixed(nn.Module):

126 + """2D learnable position embedding with temporal extension."""

127 +

128 + def __init__(

129 + self,

130 + height: int,

131 + width: int,

132 + num_frames: int,

133 + dim: int,

134 + interpolation_mode: str = "bicubic",

135 + ) -> None:

136 + super().__init__()

137 + self.height = height

138 + self.width = width

139 + self.num_frames = num_frames

140 + self.dim = dim

141 + self.interpolation_mode = interpolation_mode

142 + self.weight = nn.Parameter(torch.empty(height, width, dim))

143 + self.register_buffer(

144 + "time_weight",

145 + torch.from_numpy(get_1d_sincos_pos_embed(self.dim, self.num_frames))

146 + .float()

147 + .unsqueeze(1),

148 + persistent=False,

149 + )

150 +

151 + self.reset_parameters()

152 +

153 + def reset_parameters(self):

154 + nn.init.normal_(self.weight)

155 +

156 + def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor:

157 + pos_embs = []

158 + for t, h, w in grid_thws.tolist():

159 + assert t <= self.num_frames, f"t:{t} > self.num_frames:{self.num_frames}"

160 + if (h, w) == self.weight.shape[:-1]:

161 + pos_emb_2d = self.weight.flatten(end_dim=1)

162 + else:

163 + pos_emb_2d = get_rope_shape(

164 + self.weight,

165 + interpolation_mode=self.interpolation_mode,

166 + shape=(h, w),

167 + )

168 +

169 + if t == 1:

170 + pos_emb_3d = pos_emb_2d

171 + else:

172 + pos_emb_3d = (

173 + pos_emb_2d.unsqueeze(0).repeat(t, 1, 1) + self.time_weight[0:t]

174 + )

175 +

176 + pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1]))

177 +

178 + out = x + torch.cat(pos_embs)

179 + return out

180 +

181 +

182 +class MoonVision3dPatchEmbed(nn.Module):

183 + """3D patch embedding for vision tower."""

184 +

185 + def __init__(

186 + self,

187 + out_dim: int,

188 + in_dim: int = 3,

189 + patch_size: int | tuple[int, int] = (14, 14),

190 + pos_emb_height: int = 14,

191 + pos_emb_width: int = 14,

192 + pos_emb_time: int = 4,

193 + pos_emb_type: str = "divided_fixed",

194 + ):

195 + super().__init__()

196 + assert isinstance(patch_size, int | Sequence), (

197 + f"Invalid patch_size type: {type(patch_size)}"

198 + )

199 + if isinstance(patch_size, int):

200 + patch_size = (patch_size, patch_size)

201 + assert len(patch_size) == 2, (

202 + f"Expected patch_size to be a tuple of 2, got {patch_size}"

203 + )

204 + self.patch_size = patch_size

205 +

206 + self.proj = nn.Conv2d(

207 + in_dim, out_dim, kernel_size=patch_size, stride=patch_size

208 + )

209 +

210 + if pos_emb_type == "divided_fixed":

211 + self.pos_emb = Learnable2DInterpPosEmbDivided_fixed(

212 + height=pos_emb_height,

213 + width=pos_emb_width,

214 + num_frames=pos_emb_time,

215 + dim=out_dim,

216 + )

217 + else:

218 + raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}")

219 +

220 + def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor:

221 + x = self.proj(x).view(x.size(0), -1)

222 + # apply positional embedding

223 + x = self.pos_emb(x, grid_thws)

224 + return x

225 +

226 +

227 +class Rope2DPosEmbRepeated(nn.Module):

228 + """2D rotary position embedding with multi-resolution support."""

229 +

230 + def __init__(self, dim: int, max_height: int, max_width: int, theta_base=10000):

231 + super().__init__()

232 + self.dim = dim

233 + assert self.dim % 4 == 0, "dim must be divisible by 4"

234 + self.max_height = max_height

235 + self.max_width = max_width

236 + self.theta_base = theta_base

237 +

238 + def extra_repr(self):

239 + return (

240 + f"dim={self.dim}, max_height={self.max_height}, "

241 + f"max_width={self.max_width}, theta_base={self.theta_base}"

242 + )

243 +

244 + def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor:

245 + """Calculate the cis(freqs) for each position in the 2D grid."""

246 + N = self.max_height * self.max_width

247 + flat_pos = torch.arange(0, N).float().to(device)

248 + x_pos = flat_pos % self.max_width

249 + y_pos = flat_pos // self.max_width

250 + dim_range = (

251 + torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(device)

252 + ) # C/4

253 + freqs = 1.0 / (self.theta_base ** (dim_range / self.dim))

254 + x_freqs = torch.outer(x_pos, freqs).float() # N, C/4

255 + y_freqs = torch.outer(y_pos, freqs).float() # N, C/4

256 + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4

257 + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4

258 + # N, C/4, 2

259 + freqs_cis = torch.cat(

260 + [x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1

261 + )

262 + # max_height, max_width, C/2

263 + freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1)

264 + return freqs_cis

265 +

266 + def get_freqs_cis(

267 + self, grid_thws: torch.Tensor, device: torch.device

268 + ) -> torch.Tensor:

269 + """

270 + Args:

271 + grid_thws (torch.Tensor): grid time, height and width

272 +

273 + Returns:

274 + freqs_cis: tensor of shape (sum(t * height * width), dim//2)

275 + """

276 + if not hasattr(self, "freqs_cis"):

277 + self.register_buffer(

278 + "freqs_cis", self._precompute_freqs_cis(device), persistent=False

279 + )

280 +

281 + shapes = grid_thws.tolist()

282 + assert all(

283 + 1 <= h <= self.max_height and 1 <= w <= self.max_width for t, h, w in shapes

284 + ), (

285 + shapes,

286 + self.max_height,

287 + self.max_width,

288 + )

289 + freqs_cis = torch.cat(

290 + [

291 + self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1)

292 + for t, h, w in shapes

293 + ],

294 + dim=0,

295 + )

296 + return freqs_cis

297 +

298 +

299 +class MLP2(nn.Module):

300 + """Two-layer MLP with tensor parallel support."""

301 +

302 + def __init__(

303 + self,

304 + dims: list[int],

305 + activation,

306 + bias: bool = True,

307 + prefix: str = "",

308 + use_data_parallel: bool = False,

309 + ):

310 + super().__init__()

311 + assert len(dims) == 3

312 + self.use_data_parallel = use_data_parallel

313 + self.fc0 = ColumnParallelLinear(

314 + dims[0],

315 + dims[1],

316 + bias=bias,

317 + prefix=maybe_prefix(prefix, "fc0"),

318 + disable_tp=self.use_data_parallel,

319 + )

320 + self.fc1 = RowParallelLinear(

321 + dims[1],

322 + dims[2],

323 + bias=bias,

324 + prefix=maybe_prefix(prefix, "fc1"),

325 + disable_tp=self.use_data_parallel,

326 + )

327 + self.activation = activation

328 +

329 + def forward(self, x: torch.Tensor) -> torch.Tensor:

330 + x, _ = self.fc0(x)

331 + x = self.activation(x)

332 + x, _ = self.fc1(x)

333 + return x

334 +

335 +

336 +class MoonViTEncoderLayer(nn.Module):

337 + """Single encoder layer for MoonViT with TP/DP support."""

338 +

339 + def __init__(

340 + self,

341 + num_heads: int,

342 + hidden_dim: int,

343 + mlp_dim: int,

344 + prefix: str = "",

345 + *,

346 + activation=F.gelu,

347 + attn_bias: bool = False,

348 + ):

349 + super().__init__()

350 + self.use_data_parallel = is_vit_use_data_parallel()

351 +

352 + self.num_heads = num_heads

353 + self.hidden_dim = hidden_dim

354 + self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads

355 + self.tp_size = (

356 + 1 if self.use_data_parallel else get_tensor_model_parallel_world_size()

357 + )

358 + self.num_attention_heads_per_partition = divide(num_heads, self.tp_size)

359 +

360 + self.norm0 = nn.LayerNorm(hidden_dim)

361 + self.norm1 = nn.LayerNorm(hidden_dim)

362 + self.mlp = MLP2(

363 + [hidden_dim, mlp_dim, hidden_dim],

364 + activation,

365 + prefix=f"{prefix}.mlp",

366 + use_data_parallel=self.use_data_parallel,

367 + )

368 + self.wqkv = QKVParallelLinear(

369 + hidden_size=hidden_dim,

370 + head_size=self.hidden_size_per_attention_head,

371 + total_num_heads=num_heads,

372 + total_num_kv_heads=num_heads,

373 + bias=attn_bias,

374 + prefix=f"{prefix}.wqkv",

375 + disable_tp=self.use_data_parallel,

376 + )

377 + self.wo = RowParallelLinear(

378 + hidden_dim,

379 + hidden_dim,

380 + bias=attn_bias,

381 + prefix=f"{prefix}.wo",

382 + disable_tp=self.use_data_parallel,

383 + )

384 + self.attn = MMEncoderAttention(

385 + num_heads=self.num_attention_heads_per_partition,

386 + head_size=self.hidden_size_per_attention_head,

387 + scale=self.hidden_size_per_attention_head**-0.5,

388 + prefix=f"{prefix}.attn",

389 + )

390 +

391 + def attention_qkvpacked(

392 + self,

393 + x: torch.Tensor,

394 + cu_seqlens: torch.Tensor,

395 + rope_freqs_cis: torch.Tensor | None = None,

396 + ):

397 + """Compute self-attention with packed QKV.

398 +

399 + Args:

400 + x (torch.Tensor): (seqlen, hidden_dim)

401 + cu_seqlens (torch.Tensor): cumulative sequence lengths

402 + """

403 + seq_length = x.size(0)

404 + xqkv, _ = self.wqkv(x)

405 +

406 + qkv_shape = xqkv.size()[:-1] + (

407 + 3,

408 + self.num_attention_heads_per_partition,

409 + self.hidden_size_per_attention_head,

410 + )

411 + # xqkv: (seqlen, 3, nheads, headdim)

412 + xqkv = xqkv.view(*qkv_shape)

413 + xq, xk, xv = torch.unbind(xqkv, dim=-3)

414 +

415 + xq, xk = apply_rope(xq, xk, rope_freqs_cis)

416 +

417 + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()

418 + attn_out = self.attn(

419 + xq.unsqueeze(0),

420 + xk.unsqueeze(0),

421 + xv.unsqueeze(0),

422 + cu_seqlens=cu_seqlens,

423 + max_seqlen=max_seqlen,

424 + )

425 + attn_out = attn_out.reshape(

426 + seq_length,

427 + self.num_attention_heads_per_partition

428 + * self.hidden_size_per_attention_head,

429 + )

430 + attn_out, _ = self.wo(attn_out)

431 + return attn_out

432 +

433 + def forward(

434 + self,

435 + hidden_states: torch.Tensor,

436 + cu_seqlens: torch.Tensor,

437 + rope_freqs_cis: torch.Tensor | None = None,

438 + ):

439 + residual = hidden_states

440 + hidden_states = self.norm0(hidden_states)

441 +

442 + hidden_states = self.attention_qkvpacked(

443 + hidden_states, cu_seqlens, rope_freqs_cis

444 + )

445 + hidden_states = residual + hidden_states

446 +

447 + residual = hidden_states

448 + hidden_states = self.norm1(hidden_states)

449 + hidden_states = self.mlp(hidden_states)

450 + hidden_states = residual + hidden_states

451 +

452 + return hidden_states

453 +

454 +

455 +class MoonViT3dEncoder(nn.Module):

456 + """Full encoder stack for MoonViT 3D."""

457 +

458 + def __init__(

459 + self,

460 + hidden_dim: int,

461 + num_layers: int,

462 + block_cfg: dict,

463 + video_attn_type: str = "spatial_temporal",

464 + prefix: str = "",

465 + ) -> None:

466 + super().__init__()

467 +

468 + assert video_attn_type == "spatial_temporal", (

469 + f'video_attn_type must be "spatial_temporal", got {video_attn_type}'

470 + )

471 + self.video_attn_type = video_attn_type

472 + self.rope_2d = Rope2DPosEmbRepeated(

473 + block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512

474 + )

475 + self.blocks = nn.ModuleList(

476 + [

477 + MoonViTEncoderLayer(

478 + **block_cfg,

479 + prefix=f"{prefix}.blocks.{layer_idx}",

480 + )

481 + for layer_idx in range(num_layers)

482 + ]

483 + )

484 + self.final_layernorm = nn.LayerNorm(hidden_dim)

485 +

486 + def forward(

487 + self,

488 + hidden_states: torch.Tensor,

489 + grid_thws: torch.Tensor,

490 + ) -> torch.Tensor:

491 + rope_freqs_cis = self.rope_2d.get_freqs_cis(

492 + grid_thws=grid_thws, device=hidden_states.device

493 + )

494 +

495 + lengths = torch.cat(

496 + (

497 + torch.zeros(1, dtype=grid_thws.dtype, device=grid_thws.device),

498 + grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2],

499 + )

500 + )

501 +

502 + cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0, dtype=torch.int32)

503 +

504 + for block in self.blocks:

505 + hidden_states = block(

506 + hidden_states, cu_seqlens, rope_freqs_cis=rope_freqs_cis

507 + )

508 +

509 + hidden_states = self.final_layernorm(hidden_states)

510 +

511 + return hidden_states

512 +

513 +

514 +def tpool_patch_merger(

515 + x: torch.Tensor,

516 + grid_thws: torch.Tensor,

517 + merge_kernel_size: tuple[int, int] = (2, 2),

518 +) -> list[torch.Tensor]:

519 + """Temporal pooling patch merger."""

520 + kh, kw = merge_kernel_size

521 + lengths = (grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2]).tolist()

522 + seqs = x.split(lengths, dim=0)

523 +

524 + outputs = []

525 + for seq, (t, h, w) in zip(seqs, grid_thws.tolist()):

526 + nh, nw = h // kh, w // kw

527 + # Reshape: (t*h*w, d) -> (t, nh, kh, nw, kw, d)

528 + v = seq.view(t, nh, kh, nw, kw, -1)

529 + # Temporal pooling first (reduces tensor size before permute)

530 + v = v.mean(dim=0) # (nh, kh, nw, kw, d)

531 + # Spatial rearrangement: (nh, kh, nw, kw, d) -> (nh, nw, kh, kw, d)

532 + out = v.permute(0, 2, 1, 3, 4).reshape(nh * nw, kh * kw, -1)

533 + outputs.append(out)

534 +

535 + return outputs

536 +

537 +

538 +class MoonViT3dPretrainedModel(nn.Module):

539 + """Main vision tower model.

540 +

541 + Uses KimiK25VisionConfig directly from transformers_utils/configs/kimi_k25.py.

542 + """

543 +

544 + def __init__(

545 + self,

546 + config: KimiK25VisionConfig,

547 + prefix: str = "",

548 + ):

549 + super().__init__()

550 + config = deepcopy(config)

551 + self.config = config # Required for run_dp_sharded_mrope_vision_model

552 + self.merge_kernel_size = config.merge_kernel_size

553 + self.patch_size = config.patch_size

554 + self.merge_type = config.merge_type

555 +

556 + self.patch_embed = MoonVision3dPatchEmbed(

557 + out_dim=config.hidden_size,

558 + patch_size=config.patch_size,

559 + pos_emb_height=config.init_pos_emb_height,

560 + pos_emb_width=config.init_pos_emb_width,

561 + pos_emb_time=config.init_pos_emb_time,

562 + pos_emb_type=config.pos_emb_type,

563 + )

564 +

565 + self.encoder = MoonViT3dEncoder(

566 + hidden_dim=config.hidden_size,

567 + num_layers=config.num_hidden_layers,

568 + block_cfg={

569 + "num_heads": config.num_attention_heads,

570 + "hidden_dim": config.hidden_size,

571 + "mlp_dim": config.intermediate_size,

572 + "activation": get_act_fn("gelu_pytorch_tanh"),

573 + "attn_bias": True,

574 + },

575 + video_attn_type=config.video_attn_type,

576 + prefix=maybe_prefix(prefix, "encoder"),

577 + )

578 +

579 + def forward(

580 + self, pixel_values: torch.Tensor, grid_thws: torch.Tensor

581 + ) -> torch.Tensor:

582 + """

583 + Args:

584 + pixel_values (torch.Tensor): The input pixel values.

585 + grid_thws (torch.Tensor): Temporal, height and width.

586 +

587 + Returns:

588 + torch.Tensor: The output tokens.

589 + """

590 + hidden_states = self.patch_embed(pixel_values, grid_thws)

591 + hidden_states = self.encoder(hidden_states, grid_thws)

592 + if (

593 + self.merge_type == "sd2_tpool"

594 + ): # spatial downsampling 2x with temporal pooling all

595 + hidden_states = tpool_patch_merger(

596 + hidden_states, grid_thws, merge_kernel_size=self.merge_kernel_size

597 + )

598 + else:

599 + raise NotImplementedError(f"Not support {self.merge_type}")

600 +

601 + return hidden_states

602 +

603 +

604 +@torch.inference_mode()

605 +def mm_projector_forward(mm_projector: torch.nn.Module, vt_output: list[torch.Tensor]):

606 + """Apply MM projector to vision tower outputs."""

607 + num_embedding_list = [x.shape[0] for x in vt_output]

608 + batched = torch.cat(vt_output, dim=0)

609 + proj_out = mm_projector(batched)

610 + proj_out = proj_out.reshape(-1, proj_out.shape[-1])

611 + proj_out = torch.split(proj_out, num_embedding_list)

612 + return proj_out

613 +

614 +

615 +@torch.inference_mode()

616 +def vision_tower_forward(

617 + vision_tower: Any,

618 + pixel_values: torch.Tensor,

619 + grid_thw: torch.Tensor,

620 + mm_projector: Any,

621 + use_data_parallel: bool,

622 +) -> list[torch.Tensor]:

623 + """DP-sharded vision tower forward with mrope.

624 +

625 + Uses vLLM's standard data parallelism utility to shard the batch

626 + across available GPUs, enabling parallel processing of vision features.

627 + """

628 + if use_data_parallel:

629 + grid_thw_list = grid_thw.tolist()

630 + vt_outputs = run_dp_sharded_mrope_vision_model(

631 + vision_model=vision_tower,

632 + pixel_values=pixel_values,

633 + grid_thw_list=grid_thw_list,

634 + rope_type="rope_2d",

635 + )

636 + else:

637 + vt_outputs = vision_tower(pixel_values, grid_thw)

638 + tensors = mm_projector_forward(mm_projector, list(vt_outputs))

639 + return list(tensors)

640 +

641 +

642 +class KimiK25MultiModalProjector(nn.Module):

643 + """Multi-modal projector with patch merging for Kimi-K2.5."""

644 +

645 + def __init__(

646 + self,

647 + config: KimiK25VisionConfig,

648 + use_data_parallel: bool = False,

649 + prefix: str = "",

650 + ):

651 + super().__init__()

652 + self.use_data_parallel = use_data_parallel

653 +

654 + # Hidden size after patch merging

655 + merge_h, merge_w = config.merge_kernel_size

656 + self.hidden_size = config.hidden_size * merge_h * merge_w

657 +

658 + self.pre_norm = torch.nn.LayerNorm(config.hidden_size, eps=1e-5)

659 + self.linear_1 = ReplicatedLinear(

660 + self.hidden_size,

661 + self.hidden_size,

662 + bias=True,

663 + prefix=maybe_prefix(prefix, "linear_1"),

664 + )

665 + self.linear_2 = ReplicatedLinear(

666 + self.hidden_size,

667 + config.mm_hidden_size,

668 + bias=True,

669 + prefix=maybe_prefix(prefix, "linear_2"),

670 + )

671 + self.act = GELUActivation()

672 +

673 + def forward(self, image_features: torch.Tensor) -> torch.Tensor:

674 + hidden_states = self.pre_norm(image_features).view(-1, self.hidden_size)

675 + hidden_states, _ = self.linear_1(hidden_states)

676 + hidden_states = self.act(hidden_states)

677 + hidden_states, _ = self.linear_2(hidden_states)

678 + return hidden_states

vllm/model_executor/models/registry.py

@@ -360,6 +360,7 @@ _MULTIMODAL_MODELS = {

360 ),

361 "RForConditionalGeneration": ("rvl", "RForConditionalGeneration"),

362 "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"), # noqa: E501

363 + "KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"), # noqa: E501

364 "LightOnOCRForConditionalGeneration": (

365 "lightonocr",

366 "LightOnOCRForConditionalGeneration",

vllm/multimodal/inputs.py

+26 -1

@@ -20,6 +20,7 @@ from typing import (

20 )

22 import numpy as np

23 +from PIL.Image import Image

24 from typing_extensions import NotRequired, TypeVar

26 from vllm.utils.collection_utils import full_groupby, is_list_of

@@ -29,7 +30,6 @@ from vllm.utils.jsontree import json_map_leaves

30 if TYPE_CHECKING:

31 import torch

32 import torch.types

32 - from PIL.Image import Image

33 from transformers.feature_extraction_utils import BatchFeature

35 from .media import MediaWithBytes

@@ -105,6 +105,28 @@ The number of data items allowed per modality is restricted by

105 """

106

107

108 +class VisionChunkImage(TypedDict):

109 + """Represents an image wrapped as a vision chunk."""

110 +

111 + type: Literal["image"]

112 + image: Image

113 + uuid: str | None

114 +

115 +

116 +class VisionChunkVideo(TypedDict):

117 + """Represents a video chunk with metadata."""

118 +

119 + type: Literal["video_chunk"]

120 + video_chunk: list[Image]

121 + uuid: str | None

122 + prompt: str

123 + video_idx: int

124 +

125 +

126 +VisionChunk = VisionChunkImage | VisionChunkVideo

127 +"""A vision chunk is either an image or a video chunk."""

128 +

129 +

130 @final

131 class MultiModalDataBuiltins(TypedDict, total=False):

132 """Type annotations for modality types predefined by vLLM."""

@@ -118,6 +140,9 @@ class MultiModalDataBuiltins(TypedDict, total=False):

140 audio: ModalityData[AudioItem]

141 """The input audio(s)."""

142

143 + vision_chunk: ModalityData[VisionChunk]

144 + """The input visual atom(s) - unified modality for images and video chunks."""

145 +

146

147 MultiModalDataDict: TypeAlias = Mapping[str, ModalityData[Any]]

148 """

vllm/multimodal/parse.py

+19

@@ -384,6 +384,13 @@ class VideoEmbeddingItems(EmbeddingItems):

384 super().__init__(data, "video", expected_hidden_size)

385

386

387 +class VisionChunkProcessorItems(ProcessorBatchItems[Any]):

388 + """Processor items for vision chunks (unified image and video chunks)."""

389 +

390 + def __init__(self, data: Sequence[Any]) -> None:

391 + super().__init__(data, "vision_chunk")

392 +

393 +

394 _D = TypeVar("_D", bound=ModalityDataItems[Any, Any])

395

396

@@ -652,11 +659,23 @@ class MultiModalDataParser:

659

660 return VideoProcessorItems(new_videos, metadata=metadata_lst)

661

662 + def _parse_vision_chunk_data(

663 + self,

664 + data: ModalityData[Any],

665 + ) -> ModalityDataItems[Any, Any] | None:

666 + """Parse vision chunk data (unified image and video chunks)."""

667 + if data is None or self._is_empty(data):

668 + return None

669 + if self.is_embeddings(data):

670 + raise ValueError("Do not support embedding data for vision_chunk right now")

671 + return VisionChunkProcessorItems(data)

672 +

673 def _get_subparsers(self) -> Mapping[str, ModalityDataParser]:

674 return {

675 "audio": self._parse_audio_data,

676 "image": self._parse_image_data,

677 "video": self._parse_video_data,

678 + "vision_chunk": self._parse_vision_chunk_data,

679 }

680

681 def parse_mm_data(self, mm_data: MultiModalDataDict) -> MultiModalDataItems:

vllm/multimodal/video.py

+21

@@ -235,6 +235,27 @@ class VideoLoader:

235 VIDEO_LOADER_REGISTRY = ExtensionManager()

236

237

238 +@VIDEO_LOADER_REGISTRY.register("identity")

239 +class IdentityVideoLoader(VideoLoader):

240 + """IdentityVideoLoader returns raw video bytes without decoding.

241 +

242 + This allows the model processor to handle video decoding and

243 + is required for models like Kimi-K2.5 that need custom video chunk splitting.

244 +

245 + NOTE: This is temporary for Kimi-K2.5 testing. Remember to change back

246 + to opencv before release if needed.

247 + """

248 +

249 + @classmethod

250 + def load_bytes(

251 + cls,

252 + data: bytes,

253 + num_frames: int = -1,

254 + **kwargs: Any,

255 + ) -> tuple[Any, Any]:

256 + return data, None

257 +

258 +

259 @VIDEO_LOADER_REGISTRY.register("opencv")

260 class OpenCVVideoBackend(VideoLoader):

261 def get_cv2_video_api(self):

vllm/reasoning/__init__.py

+2 -2

@@ -53,8 +53,8 @@ _REASONING_PARSERS_TO_REGISTER = {

53 "HunyuanA13BReasoningParser",

54 ),

55 "kimi_k2": (

56 - "deepseek_r1_reasoning_parser",

57 - "DeepSeekR1ReasoningParser",

56 + "kimi_k2_reasoning_parser",

57 + "KimiK2ReasoningParser",

58 ),

59 "minimax_m2": (

60 "minimax_m2_reasoning_parser",

vllm/reasoning/kimi_k2_reasoning_parser.py new

+80

@@ -0,0 +1,80 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from collections.abc import Sequence

5 +from typing import TYPE_CHECKING, Any

6 +

7 +from transformers import PreTrainedTokenizerBase

8 +

9 +from vllm.entrypoints.openai.engine.protocol import DeltaMessage

10 +from vllm.logger import init_logger

11 +from vllm.reasoning import ReasoningParser

12 +from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser

13 +

14 +from .identity_reasoning_parser import IdentityReasoningParser

15 +

16 +if TYPE_CHECKING:

17 + from vllm.entrypoints.openai.chat_completion.protocol import (

18 + ChatCompletionRequest,

19 + )

20 +else:

21 + ChatCompletionRequest = Any

22 +

23 +

24 +logger = init_logger(__name__)

25 +

26 +

27 +class KimiK2ReasoningParser(ReasoningParser):

28 + """

29 + Kimi K2 parser that delegates to either DeepSeekR1ReasoningParser or

30 + IdentityReasoningParser based on `thinking` and `separate_reasoning`.

31 +

32 + Unlike DeepSeekV3ReasoningParser which defaults to NOT thinking,

33 + KimiK2ReasoningParser defaults to thinking mode (uses DeepSeekR1ReasoningParser).

34 + """

35 +

36 + def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs):

37 + super().__init__(tokenizer, *args, **kwargs)

38 +

39 + chat_kwargs = kwargs.pop("chat_template_kwargs", {}) or {}

40 + # Key difference: default to True instead of False

41 + thinking = bool(chat_kwargs.pop("thinking", True))

42 +

43 + if thinking:

44 + self._parser = DeepSeekR1ReasoningParser(tokenizer, *args, **kwargs)

45 + else:

46 + self._parser = IdentityReasoningParser(tokenizer, *args, **kwargs)

47 +

48 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:

49 + return self._parser.is_reasoning_end(input_ids)

50 +

51 + def is_reasoning_end_streaming(

52 + self, input_ids: list[int], delta_ids: list[int]

53 + ) -> bool:

54 + return self._parser.is_reasoning_end_streaming(input_ids, delta_ids)

55 +

56 + def extract_content_ids(self, input_ids: list[int]) -> list[int]:

57 + return self._parser.extract_content_ids(input_ids)

58 +

59 + def extract_reasoning(

60 + self, model_output: str, request: "ChatCompletionRequest"

61 + ) -> tuple[str | None, str | None]:

62 + return self._parser.extract_reasoning(model_output, request)

63 +

64 + def extract_reasoning_streaming(

65 + self,

66 + previous_text: str,

67 + current_text: str,

68 + delta_text: str,

69 + previous_token_ids: Sequence[int],

70 + current_token_ids: Sequence[int],

71 + delta_token_ids: Sequence[int],

72 + ) -> DeltaMessage | None:

73 + return self._parser.extract_reasoning_streaming(

74 + previous_text,

75 + current_text,

76 + delta_text,

77 + previous_token_ids,

78 + current_token_ids,

79 + delta_token_ids,

80 + )

vllm/renderers/hf.py

+70

@@ -20,9 +20,11 @@ from vllm.entrypoints.chat_utils import (

20 ChatTemplateContentFormatOption,

21 ChatTemplateResolutionError,

22 ConversationMessage,

23 + build_video_prompts_from_mm_data,

24 load_chat_template,

25 parse_chat_messages,

26 parse_chat_messages_async,

27 + rebuild_mm_uuids_from_mm_data,

28 )

29 from vllm.inputs import TextPrompt, TokensPrompt

30 from vllm.logger import init_logger

@@ -547,6 +549,40 @@ class HfRenderer(RendererLike):

549 **kwargs,

550 )

551

552 + # NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5

553 + # model which uses unified vision chunks for both images and videos.

554 + if (

555 + getattr(model_config.hf_config, "use_unified_vision_chunk", False)

556 + and mm_uuids is not None

557 + and mm_data is not None

558 + ):

559 + mm_uuids = rebuild_mm_uuids_from_mm_data(mm_uuids, mm_data)

560 +

561 + # get video placehoder, replace it with runtime video-chunk prompts

562 + video_placeholder = getattr(

563 + model_config.hf_config, "video_placeholder", None

564 + )

565 + if video_placeholder and isinstance(prompt_raw, str):

566 + video_prompts = build_video_prompts_from_mm_data(mm_data)

567 +

568 + # replace in order

569 + prompt_raw_parts = prompt_raw.split(video_placeholder)

570 + if len(prompt_raw_parts) == len(video_prompts) + 1:

571 + prompt_raw = "".join(

572 + [

573 + prompt_raw_parts[i] + video_prompts[i]

574 + for i in range(len(video_prompts))

575 + ]

576 + )

577 + prompt_raw += prompt_raw_parts[-1]

578 + else:

579 + logger.warning(

580 + "Number of video placeholders (%d) does not match "

581 + "number of videos (%d) in the request.",

582 + len(prompt_raw_parts) - 1,

583 + len(video_prompts),

584 + )

585 +

586 prompt = (

587 TextPrompt(prompt=prompt_raw)

588 if isinstance(prompt_raw, str)

@@ -587,6 +623,40 @@ class HfRenderer(RendererLike):

623 **kwargs,

624 )

625

626 + # NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5

627 + # model which uses unified vision chunks for both images and videos.

628 + if (

629 + getattr(model_config.hf_config, "use_unified_vision_chunk", False)

630 + and mm_uuids is not None

631 + and mm_data is not None

632 + ):

633 + mm_uuids = rebuild_mm_uuids_from_mm_data(mm_uuids, mm_data)

634 +

635 + # get video placehoder, replace it with runtime video-chunk prompts

636 + video_placeholder = getattr(

637 + model_config.hf_config, "video_placeholder", None

638 + )

639 + if video_placeholder and isinstance(prompt_raw, str):

640 + video_prompts = build_video_prompts_from_mm_data(mm_data)

641 +

642 + # replace in order

643 + prompt_raw_parts = prompt_raw.split(video_placeholder)

644 + if len(prompt_raw_parts) == len(video_prompts) + 1:

645 + prompt_raw = "".join(

646 + [

647 + prompt_raw_parts[i] + video_prompts[i]

648 + for i in range(len(video_prompts))

649 + ]

650 + )

651 + prompt_raw += prompt_raw_parts[-1]

652 + else:

653 + logger.warning(

654 + "Number of video placeholders (%d) does not match "

655 + "number of videos (%d) in the request.",

656 + len(prompt_raw_parts) - 1,

657 + len(video_prompts),

658 + )

659 +

660 prompt = (

661 TextPrompt(prompt=prompt_raw)

662 if isinstance(prompt_raw, str)

vllm/transformers_utils/config.py

@@ -81,6 +81,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(

81 isaac="IsaacConfig",

82 kimi_linear="KimiLinearConfig",

83 kimi_vl="KimiVLConfig",

84 + kimi_k25="KimiK25Config",

85 RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct)

86 RefinedWebModel="RWConfig", # For tiiuae/falcon-7b(-instruct)

87 jais="JAISConfig",

vllm/transformers_utils/configs/__init__.py

@@ -38,6 +38,7 @@ _CLASS_TO_MODULE: dict[str, str] = {

38 "MoonViTConfig": "vllm.transformers_utils.configs.moonvit",

39 "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear",

40 "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl",

41 + "KimiK25Config": "vllm.transformers_utils.configs.kimi_k25",

42 "NemotronConfig": "vllm.transformers_utils.configs.nemotron",

43 "NemotronHConfig": "vllm.transformers_utils.configs.nemotron_h",

44 "Olmo3Config": "vllm.transformers_utils.configs.olmo3",

@@ -77,6 +78,7 @@ __all__ = [

78 "MoonViTConfig",

79 "KimiLinearConfig",

80 "KimiVLConfig",

81 + "KimiK25Config",

82 "NemotronConfig",

83 "NemotronHConfig",

84 "Olmo3Config",

vllm/transformers_utils/configs/kimi_k25.py new

+129

@@ -0,0 +1,129 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""

4 +Kimi-K2.5 Model Configuration.

5 +

6 +This configuration supports video-chunk as an internal modality type.

7 +A video-chunk is the smallest independently processable unit of video.

8 +"""

9 +

10 +from transformers import DeepseekV3Config

11 +from transformers.configuration_utils import PretrainedConfig

12 +

13 +

14 +class KimiK25VisionConfig(PretrainedConfig):

15 + model_type = "kimi_k25_vision"

16 +

17 + def __init__(

18 + self,

19 + # Vision Tower

20 + patch_size: int = 14,

21 + init_pos_emb_height: int = 64,

22 + init_pos_emb_width: int = 64,

23 + init_pos_emb_time: int = 4,

24 + pos_emb_type: str = "divided_fixed",

25 + num_attention_heads: int = 16,

26 + num_hidden_layers: int = 27,

27 + hidden_size: int = 1152,

28 + intermediate_size: int = 4304,

29 + merge_kernel_size: tuple[int, int] = (2, 2),

30 + video_attn_type: str = "spatial_temporal",

31 + merge_type: str = "sd2_tpool",

32 + # MM Projector

33 + mm_projector_type: str = "patchmerger",

34 + mm_hidden_size: int | None = None,

35 + projector_hidden_act: str = "gelu",

36 + projector_ln_eps: float = 1e-5,

37 + **kwargs,

38 + ):

39 + super().__init__(**kwargs)

40 + # Vision Tower

41 + self.patch_size = patch_size

42 + self.init_pos_emb_height = init_pos_emb_height

43 + self.init_pos_emb_width = init_pos_emb_width

44 + self.init_pos_emb_time = init_pos_emb_time

45 + self.pos_emb_type = pos_emb_type

46 + self.num_attention_heads = num_attention_heads

47 + self.num_hidden_layers = num_hidden_layers

48 + self.hidden_size = hidden_size

49 + self.intermediate_size = intermediate_size

50 + self.merge_kernel_size = merge_kernel_size

51 + self.video_attn_type = video_attn_type

52 + self.merge_type = merge_type

53 + # MM Projector

54 + self.mm_projector_type = mm_projector_type

55 + if mm_hidden_size is not None:

56 + self.mm_hidden_size = mm_hidden_size

57 + else:

58 + self.mm_hidden_size = hidden_size

59 + self.projector_hidden_act = projector_hidden_act

60 + self.projector_ln_eps = projector_ln_eps

61 +

62 +

63 +class KimiK25Config(PretrainedConfig):

64 + """Kimi-K2.5 model configuration.

65 +

66 + Kimi-K2.5 extends Kimi-K2 with vision support using video-chunks.

67 + A video-chunk consists of multiple consecutive frames

68 + that are processed together with temporal pooling.

69 +

70 + Args:

71 + vision_config: Configuration for the vision tower and projector.

72 + text_config: Configuration for the text model (DeepseekV3).

73 + ignore_index: The ignore index for the loss function.

74 + media_placeholder_token_id: The token ID for media placeholders.

75 + pad_token_id: The token ID for padding.

76 + """

77 +

78 + model_type = "kimi_k25"

79 +

80 + def __init__(

81 + self,

82 + vision_config: dict | KimiK25VisionConfig | None = None,

83 + text_config: dict | DeepseekV3Config | None = None,

84 + ignore_index: int = -100,

85 + media_placeholder_token_id: int = 163605,

86 + pad_token_id: int = 0,

87 + use_unified_vision_chunk: bool = False,

88 + video_placeholder: str = "<|kimi_k25_video_placeholder|>",

89 + **kwargs,

90 + ):

91 + # Vision config

92 + if vision_config is None:

93 + vision_config = KimiK25VisionConfig()

94 + elif isinstance(vision_config, dict):

95 + vision_config = KimiK25VisionConfig(**vision_config)

96 + self.vision_config: KimiK25VisionConfig = vision_config

97 +

98 + # Text config

99 + if text_config is None:

100 + text_config = DeepseekV3Config()

101 + elif isinstance(text_config, dict):

102 + text_config = DeepseekV3Config(**text_config)

103 + self.text_config: DeepseekV3Config = text_config

104 +

105 + # Set mm_hidden_size to text hidden size if not explicitly set

106 + if self.vision_config.mm_hidden_size == self.vision_config.hidden_size:

107 + self.vision_config.mm_hidden_size = self.text_config.hidden_size

108 +

109 + # Other config

110 + self.ignore_index = ignore_index

111 + self.media_placeholder_token_id = media_placeholder_token_id

112 + self.use_unified_vision_chunk = use_unified_vision_chunk

113 + self.video_placeholder = video_placeholder

114 +

115 + # Propagate quantization config from text model

116 + if getattr(self.text_config, "quantization_config", None) is not None:

117 + self.quantization_config = self.text_config.quantization_config

118 +

119 + super().__init__(pad_token_id=pad_token_id, **kwargs)

120 +

121 + @property

122 + def hidden_size(self) -> int:

123 + """Get hidden size from text config for compatibility."""

124 + return self.text_config.hidden_size

125 +

126 + @property

127 + def vocab_size(self) -> int:

128 + """Get vocab size from text config for compatibility."""

129 + return self.text_config.vocab_size