[MM][CG] Support ViT full CUDA graph for DeepSeek-VL2 (#46005)
Signed-off-by: littlecircle0730 <littlecircle0730@gmail.com> Signed-off-by: littlecircle0730 <43994952+littlecircle0730@users.noreply.github.com> Co-authored-by: Shanshan Shen <467638484@qq.com>
This commit is contained in:
littlecircle0730
co-authored by
Shanshan Shen
Shanshan Shen
parent
0951fc183c
commit
71536bd5f8
@@ -12,6 +12,7 @@ For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tili
|
||||
| Architecture | Models | CG for Image | CG for Video | Multi-Path Graph |
|
||||
| ------------ | ------ | ------------ | ------------ | --------------- |
|
||||
| `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ |
|
||||
| `DeepseekVLV2ForCausalLM` | `DeepSeek-VL2` | ✅︎ | ❌︎ | ❌︎ |
|
||||
| `DeepseekV41ForCausalLM` | `DeepSeek-V4.1-Flash` | ✅︎ | ❌︎ | ❌︎ |
|
||||
| `Ernie4_5_VLMoeForConditionalGeneration` | `ERNIE-4.5-VL` | ✅︎ | ❌︎ | ❌︎ |
|
||||
| `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ |
|
||||
|
||||
@@ -2274,6 +2274,7 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [
|
||||
"stepvl",
|
||||
"glm4_1v",
|
||||
"deepseek_ocr",
|
||||
"deepseek_vl_v2",
|
||||
"ernie45_vl",
|
||||
"minicpmv2_5_vl",
|
||||
"minicpmv2_6_vl",
|
||||
|
||||
@@ -70,6 +70,10 @@ def gemma3_chat_template(content: str) -> str:
|
||||
return f"<bos><start_of_turn>user\n{content}<end_of_turn>\n<start_of_turn>model\n"
|
||||
|
||||
|
||||
def deepseek_vl2_chat_template(content: str) -> str:
|
||||
return f"<|User|>: {content}\n\n<|Assistant|>:"
|
||||
|
||||
|
||||
def ernie45_vl_chat_template(content: str) -> str:
|
||||
return (
|
||||
f"<|begin_of_sentence|>User: {content}"
|
||||
@@ -356,6 +360,15 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
|
||||
},
|
||||
skip=True, # TODO: Re-enable this once OOM issues are resolved on CI.
|
||||
),
|
||||
"deepseek_vl2": VitCudagraphTestConfig(
|
||||
model="deepseek-ai/deepseek-vl2-tiny",
|
||||
modalities=["image"],
|
||||
image_prompt=deepseek_vl2_chat_template("<image>\nWhat is in this image?"),
|
||||
vllm_runner_kwargs={
|
||||
"hf_overrides": {"architectures": ["DeepseekVLV2ForCausalLM"]},
|
||||
},
|
||||
marks=[pytest.mark.core_model],
|
||||
),
|
||||
"gemma4": VitCudagraphTestConfig(
|
||||
model="google/gemma-4-E2B-it",
|
||||
image_prompt=(
|
||||
|
||||
@@ -160,6 +160,8 @@ def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager:
|
||||
Skips the parts of __init__ that require a real VllmConfig / model
|
||||
by patching the attributes directly after construction.
|
||||
"""
|
||||
from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig
|
||||
|
||||
mgr = object.__new__(EncoderCudaGraphManager)
|
||||
mgr.token_budgets = sorted(budgets)
|
||||
mgr.path_token_budgets = {"default": mgr.token_budgets}
|
||||
@@ -1044,6 +1046,278 @@ class TestGetInputModality:
|
||||
assert model.get_input_modality(mm_kwargs) == "video"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GPU tests — DeepSeek-VL2 batched-tile encoder CUDA graph
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# DeepSeek-VL2 mock constants (kept small for fast capture)
|
||||
_DSV2_H = 2 # projector output side per tile (h = w)
|
||||
_DSV2_DIM = 16 # hidden size
|
||||
_DSV2_IMG_SIZE = 4 # each tile is 4×4 pixels
|
||||
_DSV2_BUDGETS = [16, 64]
|
||||
_DSV2_MAX_BATCH = 4
|
||||
|
||||
|
||||
def _dsv2_output_tokens(tw: int, th: int) -> int:
|
||||
"""Expected output tokens for an image with (tw, th) local tiles."""
|
||||
h = w = _DSV2_H
|
||||
return h * (w + 1) + th * h * (tw * w + 1) + 1
|
||||
|
||||
|
||||
def _make_deepseek_vl2_mm_kwargs(
|
||||
spatial_crops: list[list[int]],
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> dict[str, Any]:
|
||||
"""spatial_crops: [[tw, th], ...] per image."""
|
||||
total_tiles = sum(1 + int(tw) * int(th) for tw, th in spatial_crops)
|
||||
pixel_values = torch.randn(
|
||||
total_tiles, 3, _DSV2_IMG_SIZE, _DSV2_IMG_SIZE, device=device, dtype=dtype
|
||||
)
|
||||
return {
|
||||
"pixel_values": pixel_values,
|
||||
"images_spatial_crop": torch.tensor(
|
||||
spatial_crops, dtype=torch.long, device=device
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
class MockDeepseekVL2Model(torch.nn.Module, SupportsEncoderCudaGraph):
|
||||
"""Mock of DeepseekVLV2ForCausalLM for CUDA graph tests.
|
||||
|
||||
Simulates the batched-tile ViT pattern: pixel_values is
|
||||
[N_tiles, 3, H, W], output is assembled with newlines and
|
||||
view-separator in postprocess_encoder_output.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
hw = _DSV2_H * _DSV2_H
|
||||
in_features = 3 * _DSV2_IMG_SIZE * _DSV2_IMG_SIZE
|
||||
self.proj = torch.nn.Linear(in_features, hw * _DSV2_DIM)
|
||||
self.image_newline = torch.nn.Parameter(torch.randn(_DSV2_DIM))
|
||||
self.view_seperator = torch.nn.Parameter(torch.randn(_DSV2_DIM))
|
||||
self.global_view_pos = "head"
|
||||
|
||||
def get_max_frames_per_video(self) -> int:
|
||||
return 1
|
||||
|
||||
def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
|
||||
return EncoderCudaGraphConfig(
|
||||
modalities=["image"],
|
||||
buffer_keys=["pixel_values"],
|
||||
out_hidden_size=_DSV2_DIM,
|
||||
)
|
||||
|
||||
def get_encoder_cudagraph_budget_range(self, vllm_config) -> tuple[int, int]:
|
||||
return (_dsv2_output_tokens(1, 1), 128)
|
||||
|
||||
def get_encoder_cudagraph_item_specs(
|
||||
self, mm_kwargs: dict[str, Any]
|
||||
) -> list[EncoderItemSpec]:
|
||||
h = w = _DSV2_H
|
||||
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
|
||||
specs.append(
|
||||
EncoderItemSpec(
|
||||
input_size=1 + tw * th,
|
||||
output_tokens=h * (w + 1) + th * h * (tw * w + 1) + 1,
|
||||
)
|
||||
)
|
||||
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],
|
||||
}
|
||||
|
||||
tiles_per_image = [
|
||||
1 + int(row[0]) * int(row[1]) for row in images_spatial_crop.tolist()
|
||||
]
|
||||
cum = [0]
|
||||
for t in tiles_per_image:
|
||||
cum.append(cum[-1] + t)
|
||||
|
||||
selected_pv = torch.cat([pixel_values[cum[i] : cum[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,
|
||||
) -> EncoderCudaGraphCaptureInputs:
|
||||
hw = _DSV2_H * _DSV2_H
|
||||
max_tiles = (token_budget + hw - 1) // hw
|
||||
dummy = torch.randn(
|
||||
max_tiles, 3, _DSV2_IMG_SIZE, _DSV2_IMG_SIZE, device=device, dtype=dtype
|
||||
)
|
||||
return EncoderCudaGraphCaptureInputs(values={"pixel_values": dummy})
|
||||
|
||||
def prepare_encoder_cudagraph_replay_buffers(
|
||||
self,
|
||||
mm_kwargs: dict[str, Any],
|
||||
max_batch_size: int,
|
||||
max_frames_per_batch: int,
|
||||
path: str = "default",
|
||||
) -> EncoderCudaGraphReplayBuffers:
|
||||
return EncoderCudaGraphReplayBuffers(
|
||||
values={"pixel_values": mm_kwargs["pixel_values"]}
|
||||
)
|
||||
|
||||
def _tile_forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
||||
"""Project tiles: [N, 3, H, W] → [N, hw, DIM]."""
|
||||
n = pixel_values.shape[0]
|
||||
flat = pixel_values.view(n, -1)
|
||||
hw = _DSV2_H * _DSV2_H
|
||||
return self.proj(flat).view(n, hw, _DSV2_DIM)
|
||||
|
||||
def encoder_cudagraph_forward(
|
||||
self, values: dict[str, torch.Tensor], path: str = "default"
|
||||
) -> torch.Tensor:
|
||||
return self._tile_forward(values["pixel_values"])
|
||||
|
||||
def encoder_eager_forward(
|
||||
self, mm_kwargs: dict[str, Any], path: str = "default"
|
||||
) -> torch.Tensor:
|
||||
# Manager always runs `postprocess_encoder_output` after this, whether
|
||||
# the batch went through eager or CUDA graph replay, so this must
|
||||
# return the same raw per-tile shape as `encoder_cudagraph_forward`.
|
||||
return self._tile_forward(mm_kwargs["pixel_values"])
|
||||
|
||||
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:
|
||||
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 = int(images_spatial_crop[rank][0])
|
||||
th = int(images_spatial_crop[rank][1])
|
||||
n_tiles = 1 + tw * th
|
||||
tiles = output[tile_offset : tile_offset + n_tiles]
|
||||
tile_offset += n_tiles
|
||||
|
||||
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_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
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
|
||||
class TestDeepseekVL2CudaGraph:
|
||||
"""CUDA graph tests for the batched-tile encoder pattern (DeepSeek-VL2)."""
|
||||
|
||||
def setup_method(self):
|
||||
self.device = torch.device("cuda:0")
|
||||
self.dtype = torch.float16
|
||||
self.model = MockDeepseekVL2Model().to(self.device).half()
|
||||
self.mgr = _make_manager_for_gpu(
|
||||
self.model, _DSV2_BUDGETS, _DSV2_MAX_BATCH, self.device, self.dtype
|
||||
)
|
||||
self.graph_pool = current_platform.graph_pool_handle()
|
||||
self.mgr.capture(graph_pool=self.graph_pool)
|
||||
|
||||
def test_capture_creates_one_graph_per_budget(self):
|
||||
assert set(self.mgr.budget_graphs["default"].keys()) == set(_DSV2_BUDGETS)
|
||||
|
||||
def test_execute_returns_one_tensor_per_image(self):
|
||||
mm_kwargs = _make_deepseek_vl2_mm_kwargs(
|
||||
[[1, 1], [1, 1]], self.device, self.dtype
|
||||
)
|
||||
result = self.mgr.execute(mm_kwargs)
|
||||
assert result is not None
|
||||
assert len(result) == 2
|
||||
|
||||
def test_output_tokens_1x1_tile(self):
|
||||
# global=h*(w+1)=2*3=6, local=1*2*(1*2+1)=6, sep=1 → 13
|
||||
expected = _dsv2_output_tokens(1, 1)
|
||||
mm_kwargs = _make_deepseek_vl2_mm_kwargs([[1, 1]], self.device, self.dtype)
|
||||
result = self.mgr.execute(mm_kwargs)
|
||||
assert result is not None
|
||||
assert result[0].shape == (expected, _DSV2_DIM)
|
||||
|
||||
def test_output_tokens_2x1_tile(self):
|
||||
# global=6, local=1*2*(2*2+1)=10, sep=1 → 17
|
||||
expected = _dsv2_output_tokens(2, 1)
|
||||
mm_kwargs = _make_deepseek_vl2_mm_kwargs([[2, 1]], self.device, self.dtype)
|
||||
result = self.mgr.execute(mm_kwargs)
|
||||
assert result is not None
|
||||
assert result[0].shape == (expected, _DSV2_DIM)
|
||||
|
||||
def test_output_tokens_multi_image(self):
|
||||
# Image 0: tw=1,th=1 → 13; Image 1: tw=2,th=1 → 17
|
||||
e0 = _dsv2_output_tokens(1, 1)
|
||||
e1 = _dsv2_output_tokens(2, 1)
|
||||
mm_kwargs = _make_deepseek_vl2_mm_kwargs(
|
||||
[[1, 1], [2, 1]], self.device, self.dtype
|
||||
)
|
||||
result = self.mgr.execute(mm_kwargs)
|
||||
assert result is not None
|
||||
assert result[0].shape == (e0, _DSV2_DIM)
|
||||
assert result[1].shape == (e1, _DSV2_DIM)
|
||||
|
||||
def test_eager_fallback_when_tokens_exceed_all_budgets(self):
|
||||
# tw=4,th=4: global=6, local=4*2*(4*2+1)=72, sep=1 → 79 > max budget 64
|
||||
tw, th = 4, 4
|
||||
expected = _dsv2_output_tokens(tw, th)
|
||||
assert expected > _DSV2_BUDGETS[-1], "test precondition: must exceed max budget"
|
||||
mm_kwargs = _make_deepseek_vl2_mm_kwargs([[tw, th]], self.device, self.dtype)
|
||||
result = self.mgr.execute(mm_kwargs)
|
||||
assert result is not None
|
||||
assert len(result) == 1
|
||||
assert result[0].shape == (expected, _DSV2_DIM)
|
||||
assert self.mgr.graph_misses >= 1
|
||||
|
||||
def test_graph_hit_counter(self):
|
||||
mm_kwargs = _make_deepseek_vl2_mm_kwargs(
|
||||
[[1, 1], [1, 1]], self.device, self.dtype
|
||||
)
|
||||
self.mgr.execute(mm_kwargs)
|
||||
assert self.mgr.graph_hits == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GPU tests — video capture, replay, fallback, and mixed image+video
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -4,9 +4,10 @@
|
||||
# adapted from https://github.com/deepseek-ai/DeepSeek-VL2/blob/faf18023f24b962b32d9f0a2d89e402a8d383a78/deepseek_vl2/models/modeling_deepseek_vl_v2.py
|
||||
"""Inference-only Deepseek-VL2 model compatible with HuggingFace weights."""
|
||||
|
||||
import functools
|
||||
import math
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from typing import Annotated, Literal, TypeAlias
|
||||
from collections.abc import Hashable, Iterable, Mapping, Sequence
|
||||
from typing import Annotated, Any, Literal, TypeAlias
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
@@ -51,7 +52,12 @@ from vllm.transformers_utils.configs.deepseek_vl2 import (
|
||||
from vllm.utils.tensor_schema import TensorSchema, TensorShape
|
||||
from vllm.utils.torch_utils import set_default_torch_dtype
|
||||
|
||||
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
|
||||
from .interfaces import (
|
||||
MultiModalEmbeddings,
|
||||
SupportsEncoderCudaGraph,
|
||||
SupportsMultiModal,
|
||||
SupportsPP,
|
||||
)
|
||||
from .utils import (
|
||||
AutoWeightsLoader,
|
||||
WeightsMapper,
|
||||
@@ -319,7 +325,9 @@ class DeepseekVL2MultiModalProcessor(
|
||||
info=DeepseekVL2ProcessingInfo,
|
||||
dummy_inputs=DeepseekVL2DummyInputsBuilder,
|
||||
)
|
||||
class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
|
||||
class DeepseekVLV2ForCausalLM(
|
||||
nn.Module, SupportsMultiModal, SupportsPP, SupportsEncoderCudaGraph
|
||||
):
|
||||
hf_to_vllm_mapper = WeightsMapper(
|
||||
orig_to_new_prefix={
|
||||
"language.": "language_model.",
|
||||
@@ -340,6 +348,7 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
|
||||
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
|
||||
@@ -583,6 +592,186 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
|
||||
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:
|
||||
|
||||
Reference in New Issue
Block a user