@MULTIMODAL_REGISTRY.register_processor(
DeepseekVL2MultiModalProcessor,
info=DeepseekVL2ProcessingInfo,
dummy_inputs=DeepseekVL2DummyInputsBuilder,
)
class DeepseekVLV2ForCausalLM(
nn.Module, SupportsMultiModal, SupportsPP, SupportsEncoderCudaGraph
):
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_prefix={
"language.": "language_model.",
}
)
@classmethod
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
if modality.startswith("image"):
return "<image>"
raise ValueError("Only image modality is supported")
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
config: DeepseekVLV2Config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
multimodal_config = vllm_config.model_config.multimodal_config
self.config = config
self.model_config = vllm_config.model_config
self.multimodal_config = multimodal_config
self.vision_config = config.vision_config
self.projector_config = config.projector_config
self.text_config = config.text_config
model_config = vllm_config.model_config
tokenizer = cached_tokenizer_from_config(model_config)
self.image_token_id: int = tokenizer.vocab[_IMAGE_TOKEN]
with self._mark_tower_model(vllm_config, "image"):
self.vision = self._init_vision_module(
self.vision_config, quant_config, maybe_prefix(prefix, "vision")
)
self.projector = MlpProjector(self.projector_config)
self.tile_tag = config.tile_tag
self.global_view_pos = config.global_view_pos
# special token for image token sequence format
embed_std = 1 / torch.sqrt(
torch.tensor(self.projector_config.n_embed, dtype=torch.float32)
)
if self.tile_tag == "2D":
# <|view_seperator|>, <|\n|>
self.image_newline = nn.Parameter(
torch.randn(self.projector_config.n_embed) * embed_std
)
# This is a typo in original implementation
self.view_seperator = nn.Parameter(
torch.randn(self.projector_config.n_embed) * embed_std
)
else:
raise ValueError(
f"Only 2D tile_tag is supported currently, got: {self.tile_tag}"
)
with self._mark_language_model(vllm_config):
self.language_model = init_vllm_registered_model(
vllm_config=vllm_config,
hf_config=self.text_config,
prefix=maybe_prefix(prefix, "language"),
)
self.make_empty_intermediate_tensors = (
self.language_model.make_empty_intermediate_tensors
)
def _get_parent_and_attr(self, root: torch.nn.Module, dotted_name: str):
"""Return (parent_module, final_attr_name) for a dotted module path."""
names = dotted_name.split(".")
parent = root
for n in names[:-1]:
parent = getattr(parent, n)
return parent, names[-1]
# patch for timm ViT instance to support tensor parallel
def patch_vit_for_tp(
self, vit: torch.nn.Module, quant_config: QuantizationConfig | None
):
try:
import timm
except ImportError as e:
raise ImportError("Please install timm") from e
for name, module in vit.named_modules():
if isinstance(module, nn.Linear):
parent, attr_name = self._get_parent_and_attr(vit, name)
if isinstance(parent, timm.layers.Mlp) and attr_name == "fc1":
new_linear = replace_linear_class(
module, "colwise", quant_config, prefix=name
)
setattr(parent, attr_name, new_linear)
elif isinstance(parent, timm.layers.Mlp) and attr_name == "fc2":
new_linear = replace_linear_class(
module, "rowwise", quant_config, prefix=name
)
setattr(parent, attr_name, new_linear)
return vit
def _init_vision_module(
self,
vision_config: VisionEncoderConfig,
quant_config: QuantizationConfig | None,
prefix: str = "",
) -> nn.Module:
# TODO: refactor vision model through timm wrapper from transformers
try:
import timm
except ImportError as e:
raise ImportError("Please install timm") from e
with set_default_torch_dtype(torch.float16):
model = timm.create_model(
"vit_so400m_patch14_siglip_384.webli",
pretrained=False,
num_classes=0,
dynamic_img_size=True,
dynamic_img_pad=True,
)
if get_tensor_model_parallel_world_size() > 1:
model = self.patch_vit_for_tp(model, quant_config)
model = model.to(dtype=torch.get_default_dtype())
return model
def _parse_and_validate_image_input(
self, **kwargs: object
) -> DeepseekVL2ImageInputs | None:
pixel_values = kwargs.pop("pixel_values", None)
images_spatial_crop = kwargs.pop("images_spatial_crop", None)
image_embeds = kwargs.pop("image_embeds", None)
if pixel_values is None and image_embeds is None:
return None
if pixel_values is not None:
expected_h = expected_w = self.vision_config.image_size
return DeepseekVL2ImagePixelInputs(
type="pixel_values",
data=pixel_values,
images_spatial_crop=images_spatial_crop,
resolve_bindings={
"h": expected_h,
"w": expected_w,
},
)
if image_embeds is not None:
return DeepseekVL2VImageEmbeddingInputs(
type="image_embeds",
data=image_embeds,
)
raise AssertionError("This line should be unreachable.")
def _pixel_values_to_embedding(
self,
pixel_values: torch.Tensor,
images_spatial_crop: torch.Tensor,
) -> list[torch.Tensor]:
# [batch_all_tiles, vit_seq_len, c]
images_feature = self.vision.forward_features(pixel_values)
# [batch_all_tiles, hw, D]
images_embeds = self.projector(images_feature)
_, hw, n_dim = images_embeds.shape
h = w = int(hw**0.5)
# fill image token based on self.tile_tag & self.global_view_pos
tile_index = 0
vision_embeddings = []
for jdx in range(images_spatial_crop.size(0)):
# extra global & local features
num_width_tiles, num_height_tiles = images_spatial_crop[jdx]
if num_width_tiles == 0 or num_height_tiles == 0:
break
num_tiles_in_image = num_width_tiles * num_height_tiles
# [hw, D]
global_features = images_embeds[tile_index]
# [num_height_tiles * num_width_tiles, hw, D]
local_features = images_embeds[
tile_index + 1 : tile_index + 1 + num_tiles_in_image
]
tile_index += num_tiles_in_image + 1
# format global and local features
# ----------------- global view add newline -----------------
# [hw, D] -> [h, w, D]
global_features = global_features.view(h, w, n_dim)
# [D] -> [h, 1, D]
new_lines_in_global = repeat(self.image_newline, "d -> h 1 d", h=h)
# cat([h, w, D], [h, 1, D], dim=1) -> [h, w + 1, D]
global_features = torch.cat([global_features, new_lines_in_global], dim=1)
# [h, w + 1, D] -> [h * (w + 1), D]
global_features = global_features.view(-1, n_dim)
# ----------------- local view add newline -----------------
# [num_height_tiles * num_width_tiles, h * w, D] ->
# [num_height_tiles * h, num_width_tiles * w, D]
local_features = rearrange(
local_features,
"(th tw) (h w) d -> (th h) (tw w) d",
th=num_height_tiles,
tw=num_width_tiles,
h=h,
w=w,
)
# [D] -> [num_height_tiles * h, 1, D]
new_lines_in_local = repeat(
self.image_newline, "d -> (th h) 1 d", th=num_height_tiles, h=h
)
# [num_height_tiles * h, num_width_tiles * w + 1, D]
local_features = torch.cat([local_features, new_lines_in_local], dim=1)
# [num_height_tiles * h, num_width_tiles * w + 1, D]
# --> [(num_height_tiles * h) * (num_width_tiles * w + 1), D]
local_features = local_features.view(-1, n_dim)
# merge global and local tiles
if self.global_view_pos == "head":
global_local_features = torch.cat(
[
global_features,
self.view_seperator[None, :],
local_features,
]
)
else:
global_local_features = torch.cat(
[
local_features,
self.view_seperator[None, :],
global_features,
]
)
vision_embeddings.append(global_local_features)
return vision_embeddings
def _process_image_input(
self, image_input: DeepseekVL2ImageInputs
) -> torch.Tensor | list[torch.Tensor]:
if image_input["type"] == "image_embeds":
return image_input["data"]
pixel_values = image_input["data"]
images_spatial_crop = image_input["images_spatial_crop"]
return self._pixel_values_to_embedding(
pixel_values=pixel_values, images_spatial_crop=images_spatial_crop
)
# -- SupportsEncoderCudaGraph protocol methods --
@functools.cached_property
def _proj_hw(self) -> tuple[int, int]:
"""Return (h, w) of the projector output per tile."""
h = w = math.ceil(
(self.vision_config.image_size // self.vision_config.patch_size)
/ self.projector_config.downsample_ratio
)
return h, w
def get_encoder_cudagraph_config(self):
from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig
return EncoderCudaGraphConfig(
modalities=["image"],
buffer_keys=["pixel_values"],
out_hidden_size=self.projector_config.n_embed,
)
def get_encoder_cudagraph_budget_range(
self, vllm_config: "VllmConfig"
) -> tuple[int, int]:
h, w = self._proj_hw
# Smallest valid image: 1 global + 1×1 local tile → 2 tiles total.
min_budget = h * (w + 1) * 2 + 1
max_budget = min(
vllm_config.scheduler_config.max_num_batched_tokens,
self.model_config.max_model_len,
)
return min_budget, max_budget
def get_encoder_cudagraph_item_specs(self, mm_kwargs: dict[str, Any]):
from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
h, w = self._proj_hw
specs = []
for row in mm_kwargs["images_spatial_crop"].tolist():
tw, th = int(row[0]), int(row[1])
if tw == 0 or th == 0:
break
n_tiles = 1 + tw * th
output_tokens = h * (w + 1) + th * h * (tw * w + 1) + 1
specs.append(
EncoderItemSpec(input_size=n_tiles, output_tokens=output_tokens)
)
return specs
def select_encoder_cudagraph_items(
self, mm_kwargs: dict[str, Any], indices: list[int]
) -> dict[str, Any]:
pixel_values = mm_kwargs["pixel_values"]
images_spatial_crop = mm_kwargs["images_spatial_crop"]
if not indices:
return {
"pixel_values": pixel_values[:0],
"images_spatial_crop": images_spatial_crop[:0],
}
# Build per-image tile counts (1 global + tw*th local).
tiles_per_image = []
for row in images_spatial_crop.tolist():
tw, th = int(row[0]), int(row[1])
if tw == 0 or th == 0:
break
tiles_per_image.append(1 + tw * th)
cum_tiles = [0]
for t in tiles_per_image:
cum_tiles.append(cum_tiles[-1] + t)
selected_pv = torch.cat(
[pixel_values[cum_tiles[i] : cum_tiles[i + 1]] for i in indices]
)
return {
"pixel_values": selected_pv,
"images_spatial_crop": images_spatial_crop[list(indices)],
}
def prepare_encoder_cudagraph_capture_inputs(
self,
token_budget: int,
max_batch_size: int,
max_frames_per_batch: int,
device: torch.device,
dtype: torch.dtype,
path: str = "default",
axis_keys: tuple[Hashable, ...] | None = None,
):
from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphCaptureInputs
h, w = self._proj_hw
# Each tile contributes at least h*w output tokens (before newlines).
# Using h*w (not h*(w+1)) guarantees we never under-allocate.
max_tiles = (token_budget + h * w - 1) // (h * w)
dummy_pixel_values = torch.randn(
max_tiles,
3,
self.vision_config.image_size,
self.vision_config.image_size,
device=device,
dtype=dtype,
)
return EncoderCudaGraphCaptureInputs(
values={"pixel_values": dummy_pixel_values}
)
def prepare_encoder_cudagraph_replay_buffers(
self,
mm_kwargs: dict[str, Any],
max_batch_size: int,
max_frames_per_batch: int,
path: str = "default",
):
from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers
return EncoderCudaGraphReplayBuffers(
values={"pixel_values": mm_kwargs["pixel_values"]}
)
def encoder_cudagraph_forward(
self, values: dict[str, torch.Tensor], path: str = "default"
) -> torch.Tensor:
pixel_values = values["pixel_values"]
features = self.vision.forward_features(pixel_values)
return self.projector(features)
def postprocess_encoder_output(
self,
outputs: dict[str, torch.Tensor],
indices: list[int],
per_item_out_tokens: list[int],
dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
clone: bool = False,
batch_mm_kwargs: dict[str, Any] | None = None,
) -> None:
"""Assemble per-image embeddings from raw per-tile projector output.
``outputs["default"]`` has shape ``[N_tiles_captured, hw, n_embed]``.
This method re-applies newlines and the view-separator that the
eager ``_pixel_values_to_embedding`` path inserts in-line.
"""
assert batch_mm_kwargs is not None
output = outputs["default"]
images_spatial_crop = batch_mm_kwargs["images_spatial_crop"]
_, hw, n_dim = output.shape
h = w = int(hw**0.5)
tile_offset = 0
for rank, img_idx in enumerate(indices):
tw, th = (
int(images_spatial_crop[rank][0]),
int(images_spatial_crop[rank][1]),
)
n_tiles = 1 + tw * th
tiles = output[tile_offset : tile_offset + n_tiles]
tile_offset += n_tiles
# Global tile: add one newline per row.
global_f = tiles[0].view(h, w, n_dim)
newline_g = self.image_newline.view(1, 1, n_dim).expand(h, 1, n_dim)
global_f = torch.cat([global_f, newline_g], dim=1).view(-1, n_dim)
# Local tiles: rearrange grid, then add one newline per row.
local_f = tiles[1:].view(th, tw, h, w, n_dim)
local_f = local_f.permute(0, 2, 1, 3, 4).reshape(th * h, tw * w, n_dim)
newline_l = self.image_newline.view(1, 1, n_dim).expand(th * h, 1, n_dim)
local_f = torch.cat([local_f, newline_l], dim=1).view(-1, n_dim)
if self.global_view_pos == "head":
emb = torch.cat([global_f, self.view_seperator[None], local_f])
else:
emb = torch.cat([local_f, self.view_seperator[None], global_f])
if isinstance(dest, dict):
dest[img_idx] = emb.clone() if clone else emb
else:
dest[rank] = emb.clone() if clone else emb
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
image_input = self._parse_and_validate_image_input(**kwargs)
if image_input is None:
return []
vision_embeddings = self._process_image_input(image_input)
return vision_embeddings
def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
**kwargs: object,
):
if intermediate_tensors is not None:
inputs_embeds = None
hidden_states = self.language_model(
input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
)
return hidden_states
def compute_logits(
self,
hidden_states: torch.Tensor,
) -> torch.Tensor | None:
return self.language_model.compute_logits(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self)
autoloaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
return autoloaded_weights