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

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 13, 15:45 UTC

[Core][MM] Add mechanism to configure multimodal fields which should stay on CPU (#28168)

Signed-off-by: Lukas Geiger <lukas.geiger94@gmail.com>

Lukas Geiger committed Nov 7, 2025 at 12:14 UTC e0919f331d12dc5dbdefd0775bb6f94dd2fab4e2

7 files changed +68 -74

vllm/model_executor/models/interfaces.py

+6 -1

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

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

4 -from collections.abc import Callable, Iterable, Mapping, MutableSequence

4 +from collections.abc import Callable, Iterable, Mapping, MutableSequence, Set

5 from typing import (

6 TYPE_CHECKING,

7 ClassVar,

@@ -81,6 +81,11 @@ class SupportsMultiModal(Protocol):

81 `vllm.multimodal.utils.group_mm_kwargs_by_modality` to use.

82 """

84 + multimodal_cpu_fields: ClassVar[Set[str]] = frozenset()

85 + """

86 + A set indicating CPU-only multimodal fields.

87 + """

88 +

89 @classmethod

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

91 """

vllm/model_executor/models/qwen2_5_vl.py

+3 -12

@@ -1090,6 +1090,7 @@ class Qwen2_5_VLForConditionalGeneration(

1090 SupportsMRoPE,

1091 ):

1092 merge_by_field_config = True

1093 + multimodal_cpu_fields = {"image_grid_thw", "video_grid_thw"}

1094

1095 packed_modules_mapping = {

1096 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

@@ -1364,13 +1365,8 @@ class Qwen2_5_VLForConditionalGeneration(

1365 image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)

1366

1367 # Split concatenated embeddings for each image item.

1367 - # Using prod on grid_thw_list instead of grid_thw.prod avoids CUDA sync

1368 merge_size = self.visual.spatial_merge_size

1369 - sizes = (

1370 - torch.tensor(grid_thw_list, dtype=torch.long).prod(-1)

1371 - // (merge_size * merge_size)

1372 - ).tolist()

1373 -

1369 + sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()

1370 return image_embeds.split(sizes)

1371

1372 def _postprocess_image_embeds_evs(

@@ -1430,12 +1426,7 @@ class Qwen2_5_VLForConditionalGeneration(

1426

1427 # Split concatenated embeddings for each video item.

1428 merge_size = self.visual.spatial_merge_size

1433 - # Using prod on grid_thw_list instead of grid_thw.prod avoids CUDA sync

1434 - sizes = (

1435 - torch.tensor(grid_thw_list, dtype=torch.long).prod(-1)

1436 - // (merge_size * merge_size)

1437 - ).tolist()

1438 -

1429 + sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()

1430 return video_embeds.split(sizes)

1431

1432 def _postprocess_video_embeds_evs(

vllm/model_executor/models/qwen2_vl.py

+18 -20

@@ -798,21 +798,27 @@ class Qwen2VisionTransformer(nn.Module):

798 def forward(

799 self,

800 x: torch.Tensor,

801 - grid_thw: list[list[int]],

801 + grid_thw: torch.Tensor | list[list[int]],

802 ) -> torch.Tensor:

803 # patchify

804 x = x.to(device=self.device, dtype=self.dtype)

805 x = self.patch_embed(x)

806

807 + if isinstance(grid_thw, list):

808 + grid_thw_list = grid_thw

809 + grid_thw = torch.tensor(grid_thw, dtype=torch.int32)

810 + else:

811 + grid_thw_list = grid_thw.tolist()

812 +

813 # compute position embedding

808 - rotary_pos_emb = self.rot_pos_emb(grid_thw)

814 + rotary_pos_emb = self.rot_pos_emb(grid_thw_list)

815

816 # compute cu_seqlens

811 - grid_thw_ = torch.tensor(grid_thw, device=x.device, dtype=torch.long)

817 cu_seqlens = torch.repeat_interleave(

813 - grid_thw_[:, 1] * grid_thw_[:, 2], grid_thw_[:, 0]

818 + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]

819 ).cumsum(dim=0, dtype=torch.int32)

815 - cu_seqlens = F.pad(cu_seqlens, (1, 0), "constant", 0)

820 + cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])

821 + cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)

822

823 # transformers

824 x = x.unsqueeze(1)

@@ -1211,6 +1217,7 @@ class Qwen2VLForConditionalGeneration(

1217 nn.Module, SupportsMultiModal, SupportsLoRA, SupportsPP, SupportsMRoPE

1218 ):

1219 merge_by_field_config = True

1220 + multimodal_cpu_fields = {"image_grid_thw", "video_grid_thw"}

1221

1222 # To ensure correct weight loading and mapping.

1223 hf_to_vllm_mapper = WeightsMapper(

@@ -1458,7 +1465,6 @@ class Qwen2VLForConditionalGeneration(

1465 ) -> tuple[torch.Tensor, ...]:

1466 grid_thw = image_input["image_grid_thw"]

1467 assert grid_thw.ndim == 2

1461 - grid_thw_list = grid_thw.tolist()

1468

1469 if image_input["type"] == "image_embeds":

1470 image_embeds = image_input["image_embeds"]

@@ -1467,18 +1473,14 @@ class Qwen2VLForConditionalGeneration(

1473

1474 if self.use_data_parallel:

1475 return run_dp_sharded_mrope_vision_model(

1470 - self.visual, pixel_values, grid_thw_list, rope_type="rope_3d"

1476 + self.visual, pixel_values, grid_thw.tolist(), rope_type="rope_3d"

1477 )

1478 else:

1473 - image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)

1479 + image_embeds = self.visual(pixel_values, grid_thw=grid_thw)

1480

1481 # Split concatenated embeddings for each image item.

1482 merge_size = self.visual.spatial_merge_size

1477 - sizes = (

1478 - torch.tensor(grid_thw_list, dtype=torch.long).prod(-1)

1479 - // (merge_size * merge_size)

1480 - ).tolist()

1481 -

1483 + sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()

1484 return image_embeds.split(sizes)

1485

1486 def _process_video_input(

@@ -1486,26 +1488,22 @@ class Qwen2VLForConditionalGeneration(

1488 ) -> tuple[torch.Tensor, ...]:

1489 grid_thw = video_input["video_grid_thw"]

1490 assert grid_thw.ndim == 2

1489 - grid_thw_list = grid_thw.tolist()

1491

1492 if video_input["type"] == "video_embeds":

1493 video_embeds = video_input["video_embeds"]

1494 else:

1495 pixel_values_videos = video_input["pixel_values_videos"]

1496 if self.use_data_parallel:

1497 + grid_thw_list = grid_thw.tolist()

1498 return run_dp_sharded_mrope_vision_model(

1499 self.visual, pixel_values_videos, grid_thw_list, rope_type="rope_3d"

1500 )

1501 else:

1500 - video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw_list)

1502 + video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw)

1503

1504 # Split concatenated embeddings for each video item.

1505 merge_size = self.visual.spatial_merge_size

1504 - sizes = (

1505 - torch.tensor(grid_thw_list, dtype=torch.long).prod(-1)

1506 - // (merge_size * merge_size)

1507 - ).tolist()

1508 -

1506 + sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()

1507 return video_embeds.split(sizes)

1508

1509 def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

vllm/model_executor/models/qwen3_vl.py

+21 -34

@@ -414,16 +414,10 @@ class Qwen3_VisionTransformer(nn.Module):

414 def device(self) -> torch.device:

415 return self.patch_embed.proj.weight.device

416

417 - def rot_pos_emb(self, grid_thw):

417 + def rot_pos_emb(self, grid_thw: list[list[int]]):

418 pos_ids = []

419 - # Support both Tensor and list inputs for DP path

420 - if isinstance(grid_thw, list):

421 - grid_list = grid_thw

422 - max_grid_size = max(max(h, w) for _, h, w in grid_list)

423 - else:

424 - grid_list = grid_thw.tolist()

425 - max_grid_size = int(grid_thw[:, 1:].max().item())

426 - for t, h, w in grid_list:

419 + max_grid_size = max(max(h, w) for _, h, w in grid_thw)

420 + for t, h, w in grid_thw:

421 hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)

422 hpos_ids = hpos_ids.reshape(

423 h // self.spatial_merge_size,

@@ -527,24 +521,25 @@ class Qwen3_VisionTransformer(nn.Module):

521 def forward(

522 self,

523 x: torch.Tensor,

530 - grid_thw: list[list[int]],

524 + grid_thw: torch.Tensor | list[list[int]],

525 ) -> torch.Tensor:

526 hidden_states = x.to(device=self.device, dtype=self.dtype, non_blocking=True)

527 hidden_states = self.patch_embed(hidden_states)

528

535 - pos_embeds = self.fast_pos_embed_interpolate(grid_thw)

529 + if isinstance(grid_thw, list):

530 + grid_thw_list = grid_thw

531 + grid_thw = torch.tensor(grid_thw, dtype=torch.int32)

532 + else:

533 + grid_thw_list = grid_thw.tolist()

534 +

535 + pos_embeds = self.fast_pos_embed_interpolate(grid_thw_list)

536 hidden_states = hidden_states + pos_embeds

537 - rotary_pos_emb = self.rot_pos_emb(grid_thw)

537 + rotary_pos_emb = self.rot_pos_emb(grid_thw_list)

538 rotary_pos_emb = rotary_pos_emb.to(hidden_states.device, non_blocking=True)

539

540 - grid_thw_tensor = torch.tensor(grid_thw, dtype=torch.int32)

541 -

540 cu_seqlens = torch.repeat_interleave(

543 - grid_thw_tensor[:, 1] * grid_thw_tensor[:, 2], grid_thw_tensor[:, 0]

544 - ).cumsum(

545 - dim=0,

546 - dtype=grid_thw_tensor.dtype if torch.jit.is_tracing() else torch.int32,

547 - )

541 + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]

542 + ).cumsum(dim=0, dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32)

543 cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])

544

545 hidden_states = hidden_states.unsqueeze(1)

@@ -1177,6 +1172,7 @@ class Qwen3VLForConditionalGeneration(

1172 nn.Module, SupportsMultiModal, SupportsLoRA, SupportsPP, SupportsMRoPE

1173 ):

1174 merge_by_field_config = True

1175 + multimodal_cpu_fields = {"image_grid_thw", "video_grid_thw"}

1176

1177 packed_modules_mapping = {

1178 "qkv_proj": [

@@ -1356,7 +1352,6 @@ class Qwen3VLForConditionalGeneration(

1352 ) -> tuple[torch.Tensor, ...]:

1353 grid_thw = image_input["image_grid_thw"]

1354 assert grid_thw.ndim == 2

1359 - grid_thw_list = grid_thw.tolist()

1355

1356 if image_input["type"] == "image_embeds":

1357 image_embeds = image_input["image_embeds"].type(self.visual.dtype)

@@ -1364,18 +1359,14 @@ class Qwen3VLForConditionalGeneration(

1359 pixel_values = image_input["pixel_values"].type(self.visual.dtype)

1360 if self.use_data_parallel:

1361 return run_dp_sharded_mrope_vision_model(

1367 - self.visual, pixel_values, grid_thw_list, rope_type="rope_3d"

1362 + self.visual, pixel_values, grid_thw.tolist(), rope_type="rope_3d"

1363 )

1364 else:

1370 - image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)

1365 + image_embeds = self.visual(pixel_values, grid_thw=grid_thw)

1366

1367 # Split concatenated embeddings for each image item.

1373 - # Using prod on grid_thw_list instead of grid_thw.prod avoids CUDA sync

1368 merge_size = self.visual.spatial_merge_size

1375 - sizes = (

1376 - torch.tensor(grid_thw_list, dtype=torch.long).prod(-1)

1377 - // (merge_size * merge_size)

1378 - ).tolist()

1369 + sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()

1370 return image_embeds.split(sizes)

1371

1372 def _process_video_input(

@@ -1383,7 +1374,6 @@ class Qwen3VLForConditionalGeneration(

1374 ) -> tuple[torch.Tensor, ...]:

1375 grid_thw = video_input["video_grid_thw"]

1376 assert grid_thw.ndim == 2

1386 - grid_thw_list = grid_thw.tolist()

1377

1378 if video_input["type"] == "video_embeds":

1379 video_embeds = video_input["video_embeds"].type(self.visual.dtype)

@@ -1392,19 +1382,16 @@ class Qwen3VLForConditionalGeneration(

1382 self.visual.dtype

1383 )

1384 if self.use_data_parallel:

1385 + grid_thw_list = grid_thw.tolist()

1386 return run_dp_sharded_mrope_vision_model(

1387 self.visual, pixel_values_videos, grid_thw_list, rope_type="rope_3d"

1388 )

1389 else:

1399 - video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw_list)

1390 + video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw)

1391

1392 # Split concatenated embeddings for each video item.

1402 - # Using prod on grid_thw_list instead of grid_thw.prod avoids CUDA sync

1393 merge_size = self.visual.spatial_merge_size

1404 - sizes = (

1405 - torch.tensor(grid_thw_list, dtype=torch.long).prod(-1)

1406 - // (merge_size * merge_size)

1407 - ).tolist()

1394 + sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()

1395 return video_embeds.split(sizes)

1396

1397 def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

vllm/multimodal/utils.py

+13 -7

@@ -3,7 +3,7 @@

4 import asyncio

5 import atexit

6 -from collections.abc import Iterable

6 +from collections.abc import Iterable, Set

7 from concurrent.futures import ThreadPoolExecutor

8 from itertools import groupby

9 from pathlib import Path

@@ -402,6 +402,7 @@ def group_mm_kwargs_by_modality(

402 device: torch.types.Device = None,

403 pin_memory: bool = False,

404 merge_by_field_config: bool | None = None,

405 + multimodal_cpu_fields: Set[str] = frozenset(),

406 ) -> Iterable[tuple[str, int, BatchedTensorInputs]]:

407 """Group consecutive `MultiModalKwargsItem`s from `mm_kwargs` with the same

408 modality together into the same `MultiModalKwargs` instance.

@@ -443,12 +444,17 @@ def group_mm_kwargs_by_modality(

444 )

445

446 if device is not None:

446 - mm_kwargs_group = json_map_leaves(

447 - lambda x: x.to(device=device, non_blocking=True)

448 - if isinstance(x, torch.Tensor)

449 - else x,

450 - mm_kwargs_group,

451 - )

447 + mm_kwargs_group = {

448 + k: json_map_leaves(

449 + lambda x: x.to(device=device, non_blocking=True)

450 + if isinstance(x, torch.Tensor)

451 + else x,

452 + v,

453 + )

454 + if k not in multimodal_cpu_fields

455 + else v

456 + for k, v in mm_kwargs_group.items()

457 + }

458 else:

459 mm_kwargs_group = MultiModalKwargs.as_kwargs(

460 MultiModalKwargs.batch(

vllm/v1/worker/gpu_model_runner.py

@@ -938,6 +938,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

938 device=self.device,

939 pin_memory=self.pin_memory,

940 merge_by_field_config=model.merge_by_field_config,

941 + multimodal_cpu_fields=model.multimodal_cpu_fields,

942 ):

943 mm_kwargs_combined.update(mm_kwargs_group)

944

@@ -1768,6 +1769,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1769 device=self.device,

1770 pin_memory=self.pin_memory,

1771 merge_by_field_config=model.merge_by_field_config,

1772 + multimodal_cpu_fields=model.multimodal_cpu_fields,

1773 ):

1774 curr_group_outputs = []

1775

@@ -1794,6 +1796,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1796 device=self.device,

1797 pin_memory=self.pin_memory,

1798 merge_by_field_config=model.merge_by_field_config,

1799 + multimodal_cpu_fields=model.multimodal_cpu_fields,

1800 )

1801 )

1802

@@ -1936,6 +1939,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1939 device=self.device,

1940 pin_memory=self.pin_memory,

1941 merge_by_field_config=model.merge_by_field_config,

1942 + multimodal_cpu_fields=model.multimodal_cpu_fields,

1943 ):

1944 # Add the grouped features to encoder_features dict

1945 # This allows the model to receive them as kwargs (e.g.,

@@ -3292,6 +3296,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3296 device=self.device,

3297 pin_memory=self.pin_memory,

3298 merge_by_field_config=model.merge_by_field_config,

3299 + multimodal_cpu_fields=model.multimodal_cpu_fields,

3300 )

3301 )

3302

vllm/v1/worker/tpu_model_runner.py

@@ -952,6 +952,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

952 device=self.device,

953 pin_memory=self.pin_memory,

954 merge_by_field_config=model.merge_by_field_config,

955 + multimodal_cpu_fields=model.multimodal_cpu_fields,

956 ):

957 # Run the encoder.

958 # `curr_group_outputs` is either of the following:

@@ -2037,6 +2038,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2038 device=self.device,

2039 pin_memory=self.pin_memory,

2040 merge_by_field_config=model.merge_by_field_config,

2041 + multimodal_cpu_fields=model.multimodal_cpu_fields,

2042 )

2043 )

2044