@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3e3ceb1961

[Perf] Avoid more GPU<->CPU syncs in multimodal encoders (#52369)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Nick Hill committed Aug 14, 2026 at 10:40 UTC 3e3ceb1961d44252ff4a4f8e22364186d70c54f7

9 files changed +81 -72

vllm/model_executor/models/diffusion_gemma.py

+5 -2

@@ -867,7 +867,7 @@ class DiffusionGemmaModelState(ModelState):

867 self.diffusion_states.add_request(req_index)

868 if not new_req_data.req_id.startswith("_warmup_"):

869 prompt_len = len(new_req_data.prompt_token_ids)

870 - self.diffusion_states.prompt_len[req_index] = prompt_len

870 + self.diffusion_states.prompt_len[req_index].fill_(prompt_len)

871

872 def remove_request(self, req_id: str) -> None:

873 idx = self._req_id_to_index.pop(req_id, None)

@@ -1276,7 +1276,10 @@ class DiffusionSampler:

1276 # before canvas padding so phantom positions stay uniform.

1277 if num_decode > 0:

1278 top_k, top_p = self.sampling_states.get_top_k_top_p(

1279 - decode_slots.repeat_interleave(valid_canvas_len), decode_slots_np

1279 + decode_slots.repeat_interleave(

1280 + valid_canvas_len, output_size=int(valid_canvas_len_np.sum())

1281 + ),

1282 + decode_slots_np,

1283 )

1284 if top_k is not None or top_p is not None:

1285 logits = apply_top_k_top_p(logits.float(), top_k, top_p)

vllm/model_executor/models/ernie45_vl.py

+9 -5

@@ -72,6 +72,7 @@ from vllm.multimodal.processing import (

72 )

73 from vllm.sequence import IntermediateTensors

74 from vllm.utils.tensor_schema import TensorSchema, TensorShape

75 +from vllm.utils.torch_utils import async_tensor_h2d

76 from vllm.v1.attention.backends.registry import AttentionBackendEnum

78 from .ernie45_vl_moe import Ernie4_5_VLMoeForCausalLM

@@ -446,6 +447,9 @@ class Ernie4_5_VisionTransformer(nn.Module):

447 pos_ids = torch.cat(pos_ids, dim=0)

448 max_grid_size = grid_thw[:, 1:].max()

449 rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)

450 + # `pos_ids` is built on the host; stage it over non-blocking so the

451 + # gather below doesn't index a device tensor with a CPU one.

452 + pos_ids = pos_ids.to(rotary_pos_emb_full.device, non_blocking=True)

453 rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)

454 return rotary_pos_emb

455

@@ -530,7 +534,7 @@ class Ernie4_5_VisionTransformer(nn.Module):

534 else:

535 max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)

536

533 - cu_seqlens = cu_seqlens.to(device)

537 + cu_seqlens = cu_seqlens.to(device, non_blocking=True)

538

539 return {

540 "rotary_pos_emb": rotary_pos_emb,

@@ -806,8 +810,8 @@ class VariableResolutionResamplerModel(nn.Module):

810 b_offset + (temp_offset + 1) * spatial_size,

811 )

812 )

809 - slice_offsets = torch.tensor(np.concatenate(slice_offsets, axis=-1)).to(

810 - x.device

813 + slice_offsets = async_tensor_h2d(

814 + np.concatenate(slice_offsets, axis=-1), device=x.device

815 )

816

817 slice_offsets2 = []

@@ -823,8 +827,8 @@ class VariableResolutionResamplerModel(nn.Module):

827 b_offset + (temp_offset + 1) * spatial_size,

828 )

829 )

826 - slice_offsets2 = torch.tensor(np.concatenate(slice_offsets2, axis=-1)).to(

827 - x.device

830 + slice_offsets2 = async_tensor_h2d(

831 + np.concatenate(slice_offsets2, axis=-1), device=x.device

832 )

833

834 x_timestep_1 = torch.index_select(x, dim=0, index=slice_offsets)

vllm/model_executor/models/glm4_1v.py

+18 -37

@@ -546,14 +546,6 @@ class Glm4vVisionEmbeddings(nn.Module):

546 0, hidden_size, device=device, dtype=pos_embed_weight.dtype

547 )

548 else:

549 - # Convert inputs to tensors if needed

550 - if isinstance(lengths, list):

551 - lengths = async_tensor_h2d(lengths, device=device, dtype=torch.long)

552 - if not isinstance(image_shapes, torch.Tensor):

553 - image_shapes = async_tensor_h2d(

554 - image_shapes, device=device, dtype=torch.long

555 - )

556 -

549 # Prepare 2D position embedding

550 orig_size_sq = pos_embed_weight.shape[0]

551 orig_size = int(orig_size_sq**0.5)

@@ -564,32 +556,23 @@ class Glm4vVisionEmbeddings(nn.Module):

556 .to(device=device, dtype=torch.float32)

557 )

558

567 - # Calculate target dimensions for each patch

568 - # Add bounds checking for data parallel mode

569 - if len(lengths) > image_shapes.shape[0]:

570 - # In data parallel mode, some GPUs might not have all

571 - # image shapes

572 - # Use available image shapes, cycling if necessary

573 - target_h_list = []

574 - target_w_list = []

575 - for i in range(len(lengths)):

576 - # Cycle through available shapes

577 - shape_idx = i % image_shapes.shape[0]

578 - target_h_list.append(image_shapes[shape_idx, 1].repeat(lengths[i]))

579 - target_w_list.append(image_shapes[shape_idx, 2].repeat(lengths[i]))

580 - target_h = torch.cat(target_h_list).to(

581 - device=device, dtype=torch.float32

582 - )

583 - target_w = torch.cat(target_w_list).to(

584 - device=device, dtype=torch.float32

585 - )

586 - else:

587 - target_h = torch.cat(

588 - [image_shapes[i, 1].repeat(lengths[i]) for i in range(len(lengths))]

589 - ).to(device=device, dtype=torch.float32)

590 - target_w = torch.cat(

591 - [image_shapes[i, 2].repeat(lengths[i]) for i in range(len(lengths))]

592 - ).to(device=device, dtype=torch.float32)

559 + # Calculate target dimensions for each patch. `lengths` and

560 + # `image_shapes` are host data, so expand them with numpy and move

561 + # the result across once rather than per-image.

562 + # Shapes are cycled: in data parallel mode some GPUs might not

563 + # have all image shapes.

564 + shapes_np = np.asarray(image_shapes)

565 + shape_idx = np.arange(len(lengths)) % shapes_np.shape[0]

566 + target_h = async_tensor_h2d(

567 + np.repeat(shapes_np[shape_idx, 1], lengths),

568 + device=device,

569 + dtype=torch.float32,

570 + )

571 + target_w = async_tensor_h2d(

572 + np.repeat(shapes_np[shape_idx, 2], lengths),

573 + device=device,

574 + dtype=torch.float32,

575 + )

576

577 # Normalize coordinates to [-1, 1] range for grid_sample

578 h_coords = h_coords.to(device=device, dtype=torch.float32)

@@ -818,9 +801,7 @@ class Glm4vVisionTransformer(nn.Module):

801 )

802

803 lengths = [h * w] * t

821 - image_shapes = async_tensor_h2d(

822 - [[t, h, w]], dtype=torch.long, device=device

823 - )

804 + image_shapes = [[t, h, w]]

805

806 # Build the coordinates on the host (cheap integer math) but move

807 # them across pinned + non-blocking, so the consumer's

vllm/model_executor/models/interns1.py

+2 -2

@@ -424,7 +424,7 @@ class InternS1MultiModalProcessor(BaseMultiModalProcessor[InternS1ProcessingInfo

424 pixel_values=MultiModalFieldConfig.flat_from_sizes(

425 "image", image_num_patches

426 ),

427 - image_num_patches=MultiModalFieldConfig.batched("image"),

427 + image_num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),

428 image_embeds=MultiModalFieldConfig.batched("image"),

429 image_token_id=MultiModalFieldConfig.shared(

430 "image", num_images, keep_on_cpu=True

@@ -432,7 +432,7 @@ class InternS1MultiModalProcessor(BaseMultiModalProcessor[InternS1ProcessingInfo

432 pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(

433 "video", video_num_patches

434 ),

435 - video_num_patches=MultiModalFieldConfig.batched("video"),

435 + video_num_patches=MultiModalFieldConfig.batched("video", keep_on_cpu=True),

436 video_token_id=MultiModalFieldConfig.shared(

437 "video", num_videos, keep_on_cpu=True

438 ),

vllm/model_executor/models/keye.py

+25 -17

@@ -59,6 +59,7 @@ from vllm.multimodal.processing import (

59 )

60 from vllm.sequence import IntermediateTensors

61 from vllm.utils.tensor_schema import TensorSchema, TensorShape

62 +from vllm.utils.torch_utils import async_tensor_h2d

64 from .interfaces import (

65 MultiModalEmbeddings,

@@ -608,7 +609,9 @@ class KeyeSiglipEncoder(nn.Module):

609 [height_position_ids, width_position_ids],

610 dim=-1,

611 )

611 - max_grid_size = pids.max() + 1

612 + # The ids are built from the grids above, so `h`/`w` bound them

613 + # and the table size is known on the host.

614 + max_grid_size = max(max(h, w) for _, h, w in flatten_image_grid_thw)

615 rope_emb_max_grid = self.rotary_pos_emb(max_grid_size)

616 rope_emb = rope_emb_max_grid[pids].flatten(1)

617 rope_emb = rope_emb.repeat(1, 2)

@@ -694,19 +697,19 @@ class KeyeSiglipVisionTransformer(nn.Module):

697

698 last_hidden_state = self.post_layernorm(last_hidden_state)

699

697 - sample_hidden_state = list()

700 if cu_seqlens is None:

701 raise ValueError(

702 "cu_seqlens cannot be None for "

703 "SiglipVisionTransformer output processing."

704 )

703 - for i in range(cu_seqlens.shape[0] - 1):

704 - start = cu_seqlens[i]

705 - end = cu_seqlens[i + 1]

706 - tensor = last_hidden_state[:, start:end, :].squeeze(0)

707 - sample_hidden_state.append(tensor)

708 -

709 - return sample_hidden_state

705 + # `cu_seqlens` is the running sum of the per-image `t * h * w`, so the

706 + # split sizes are known on the host and slicing needs no device read.

707 + split_sizes = [

708 + int(np.prod(thw)) for thw in self.encoder.flatten_list(image_grid_thw)

709 + ]

710 + return [

711 + tensor.squeeze(0) for tensor in last_hidden_state.split(split_sizes, dim=1)

712 + ]

713

714

715 class KeyeSiglipVisionModel(nn.Module):

@@ -1307,13 +1310,17 @@ class BaseKeyeModule(nn.Module, SupportsMultiModal):

1310 )

1311 else:

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

1313 + # These are all built on the host, so stage them over

1314 + # non-blocking rather than forcing a sync per transfer.

1315 siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to(

1311 - pixel_values.device

1316 + pixel_values.device, non_blocking=True

1317 + )

1318 + cu_seqlens = async_tensor_h2d(

1319 + cu_seqlens, dtype=torch.int32, device=pixel_values.device

1320 )

1313 - cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32).to(

1314 - pixel_values.device

1321 + sample_indices = torch.concat(sample_indices, dim=0).to(

1322 + pixel_values.device, non_blocking=True

1323 )

1316 - sample_indices = torch.concat(sample_indices, dim=0).to(pixel_values.device)

1324

1325 image_embeds = self.visual(

1326 pixel_values=pixel_values,

@@ -1357,14 +1364,15 @@ class BaseKeyeModule(nn.Module, SupportsMultiModal):

1364 )

1365 else:

1366 pixel_values_videos = pixel_values_videos.type(self.visual.dtype)

1367 + # These are host-built; stage them over non-blocking.

1368 siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to(

1361 - pixel_values_videos.device

1369 + pixel_values_videos.device, non_blocking=True

1370 )

1363 - cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32).to(

1364 - pixel_values_videos.device

1371 + cu_seqlens = async_tensor_h2d(

1372 + cu_seqlens, dtype=torch.int32, device=pixel_values_videos.device

1373 )

1374 sample_indices = torch.concat(sample_indices, dim=0).to(

1367 - pixel_values_videos.device

1375 + pixel_values_videos.device, non_blocking=True

1376 )

1377

1378 video_embeds = self.visual(

vllm/model_executor/models/paddleocr_vl.py

+11 -5

@@ -68,6 +68,7 @@ from vllm.multimodal.processing import (

68 )

69 from vllm.sequence import IntermediateTensors

70 from vllm.utils.tensor_schema import TensorSchema, TensorShape

71 +from vllm.utils.torch_utils import async_tensor_h2d

72 from vllm.v1.attention.backends.registry import AttentionBackendEnum

74 from .ernie45 import Ernie4_5ForCausalLM

@@ -800,16 +801,18 @@ class SiglipEncoder(nn.Module):

801 [height_position_ids, width_position_ids],

802 dim=-1,

803 )

803 - max_grid_size = pids.max() + 1

804 + # The ids are built from the grids above, so `h`/`w` bound them

805 + # and the table size is known on the host.

806 + max_grid_size = max(max(h, w) for _, h, w in flatten_image_grid_thw)

807 rope_emb_max_grid = self.rotary_pos_emb(max_grid_size)

808 rotary_pos_emb = rope_emb_max_grid[pids].flatten(1)

809

810 if cu_seqlens is None:

811 raise ValueError("cu_seqlens cannot be None for SiglipEncoder.")

812 if not isinstance(cu_seqlens, torch.Tensor):

810 - cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)

813 + cu_seqlens = async_tensor_h2d(cu_seqlens, dtype=torch.int32, device=device)

814 else:

812 - cu_seqlens = cu_seqlens.to(device=device)

815 + cu_seqlens = cu_seqlens.to(device=device, non_blocking=True)

816

817 max_seqlen = None

818 if self.attn_backend in {

@@ -1131,10 +1134,13 @@ class PaddleOCRVLForConditionalGeneration(nn.Module, SupportsMultiModal, Support

1134 siglip_position_ids.append(image_position_ids)

1135 cu_seqlens.append(cu_seqlens[-1] + numel)

1136

1137 + # Both are built on the host; stage them over non-blocking.

1138 siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to(

1135 - pixel_values.device

1139 + pixel_values.device, non_blocking=True

1140 + )

1141 + cu_seqlens = async_tensor_h2d(

1142 + cu_seqlens, dtype=torch.int32, device=pixel_values.device

1143 )

1137 - cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32).to(pixel_values.device)

1144

1145 vision_outputs = self.visual(

1146 pixel_values=pixel_values,

vllm/model_executor/models/phi4mm.py

+1 -1

@@ -905,7 +905,7 @@ class Phi4MMMultiModalProcessor(BaseMultiModalProcessor[Phi4MMProcessingInfo]):

905 return dict(

906 input_image_embeds=MultiModalFieldConfig.batched("image"),

907 image_attention_mask=MultiModalFieldConfig.batched("image"),

908 - image_sizes=MultiModalFieldConfig.batched("image"),

908 + image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),

909 num_img_tokens=MultiModalFieldConfig.batched("image"),

910 input_audio_embeds=MultiModalFieldConfig.batched("audio"),

911 )

vllm/model_executor/models/ultravox.py

+5 -1

@@ -244,7 +244,11 @@ class UltravoxMultiModalProcessor(BaseMultiModalProcessor[UltravoxProcessingInfo

244 # higher than the number of audio samples

245 audio_features=MultiModalFieldConfig.flat_from_sizes("audio", num_chunks),

246 audio_token_len=MultiModalFieldConfig.flat_from_sizes("audio", num_chunks),

247 - audio_lens=MultiModalFieldConfig.flat_from_sizes("audio", num_chunks),

247 + # Only ever used to derive the encoder attention metadata on the

248 + # host, so keep it there.

249 + audio_lens=MultiModalFieldConfig.flat_from_sizes(

250 + "audio", num_chunks, keep_on_cpu=True

251 + ),

252 # num_chunks can convert audio_chunked to audio batch dimension

253 audio_num_chunks=MultiModalFieldConfig.batched("audio"),

254 audio_embeds=MultiModalFieldConfig.batched("audio"),

vllm/third_party/flash_linear_attention/ops/index.py

+5 -2

@@ -27,11 +27,14 @@ def prepare_chunk_indices(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Te

27 for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()

28 ]

29 )

30 - return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens)

30 + chunk_indices = torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1)

31 + return chunk_indices.to(

32 + device=cu_seqlens.device, dtype=cu_seqlens.dtype, non_blocking=True

33 + )

36 @tensor_cache

37 def prepare_chunk_offsets(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Tensor:

38 return torch.cat(

36 - [cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]

39 + [cu_seqlens.new_zeros(1), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]

40 ).cumsum(-1)