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 14, 05:30 UTC
Migrate KimiVLImagePixelInputs to TensorSchema (#21769)
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: Isotr0py <2037008807@qq.com>
Benji Beck committed
Aug 5, 2025 at 02:36 UTC
05fae021750be1049927299ea2317d742c03718a
1 file changed
+15
-9
vllm/model_executor/models/kimi_vl.py
+15
-9
@@ -46,7 +46,7 @@ import copy
46
import math
47
from collections.abc import Iterable, Mapping, Sequence
48
from dataclasses import dataclass
49
-from typing import Any, Literal, Optional, TypedDict, Union
49
+from typing import Annotated, Any, Literal, Optional, Union
50
51
import torch
52
from torch import nn
@@ -79,6 +79,7 @@ from vllm.multimodal.profiling import BaseDummyInputsBuilder
79
from vllm.sequence import IntermediateTensors
80
from vllm.transformers_utils.configs import KimiVLConfig, MoonViTConfig
81
from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekV2Config
82
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
83
84
from .utils import is_pp_missing_parameter, maybe_prefix
85
@@ -118,15 +119,22 @@ class KimiVLMultiModalProjector(nn.Module):
119
return hidden_states
120
121
121
-class KimiVLImagePixelInputs(TypedDict):
122
- type: Literal["pixel_values"]
123
- pixel_values: Union[torch.Tensor, list[torch.Tensor]]
122
+class KimiVLImagePixelInputs(TensorSchema):
123
"""
125
- Shape:`(num_patches, num_channels, patch_size, patch_size)`
124
+ Dimensions:
125
+ - nc: Number of channels
126
+ - np: Number of patches
127
+ - ps: Patch size
128
+ - ni: Number of images
129
"""
130
+ type: Literal["pixel_values"] = "pixel_values"
131
128
- image_grid_hws: torch.Tensor
129
- """Shape:`(num_images, 2)`"""
132
+ pixel_values: Annotated[
133
+ Union[torch.Tensor, list[torch.Tensor]],
134
+ TensorShape("np", 3, "ps", "ps"),
135
+ ]
136
+
137
+ image_grid_hws: Annotated[torch.Tensor, TensorShape("ni", 2)]
138
139
140
# TODO: support embeds too
@@ -348,8 +356,6 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal):
356
pixel_values = pixel_values.reshape(-1, num_channels, patch_size,
357
patch_size)
358
pixel_values = pixel_values.to(self.vision_tower.dtype)
351
- # image_grid_hws.shape = (N, 2)
352
- assert image_grid_hws.ndim == 2, f"unexpected shape for image_grid_hws: {image_grid_hws.shape}"
359
360
return KimiVLImagePixelInputs(
361
type="pixel_values",