@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a566ea7e8f

[Perf] Avoid more h2d copies from non-pinned tensors (#54660)

Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Nick Hill committed Sep 1, 2026 at 20:37 UTC a566ea7e8f5975b4d38c0092d7adb0acbda3a890

5 files changed +72 -64

vllm/model_executor/models/qwen2_5_vl.py

+7 -6

@@ -805,7 +805,7 @@ class Qwen2_5_VisionTransformer(nn.Module):

805 window_index_thw, cu_seqlens_window_thw = self.get_window_index_thw(t, h, w)

806 cos_thw, sin_thw = self.rotary_pos_emb_thw(t, h, w)

807

808 - window_index_thw_dev = window_index_thw.to(cos_thw.device, non_blocking=True)

808 + window_index_thw_dev = async_tensor_h2d(window_index_thw, cos_thw.device)

809 cos_thw = cos_thw[window_index_thw_dev, :, :]

810 cos_thw = cos_thw.flatten(start_dim=0, end_dim=1)

811 sin_thw = sin_thw[window_index_thw_dev, :, :]

@@ -1053,7 +1053,7 @@ class Qwen2_5_VisionTransformer(nn.Module):

1053 )

1054 rotary_pos_emb_cos = rotary_pos_emb_cos.to(device=device, non_blocking=True)

1055 rotary_pos_emb_sin = rotary_pos_emb_sin.to(device=device, non_blocking=True)

1056 - window_index = window_index.to(device=device, non_blocking=True)

1056 + window_index = async_tensor_h2d(window_index, device)

1057 reverse_indices = reverse_indices.to(device=device, non_blocking=True)

1058

1059 metadata["rotary_pos_emb_cos"] = rotary_pos_emb_cos

@@ -1475,8 +1475,8 @@ class Qwen2_5_VLForConditionalGeneration(

1475 grid_thw_list = grid_thw.tolist()

1476 image_embeds_out = []

1477 for emb, size in zip(image_embeds_split, grid_thw_list):

1478 - positions = compute_mrope_for_media(size, merge_size).to(

1479 - emb.device, non_blocking=True

1478 + positions = async_tensor_h2d(

1479 + compute_mrope_for_media(size, merge_size), emb.device

1480 )

1481 emb = torch.cat([emb, positions], dim=1)

1482 image_embeds_out.append(emb)

@@ -1558,12 +1558,13 @@ class Qwen2_5_VLForConditionalGeneration(

1558 spatial_merge_size=self.visual.spatial_merge_size,

1559 q=self.video_pruning_rate,

1560 )

1561 - positions = compute_mrope_for_media(

1561 + positions_cpu = compute_mrope_for_media(

1562 size,

1563 merge_size,

1564 tokens_per_second=tokens_per_second,

1565 video_second_per_grid=video_second_per_grid_t.item(),

1566 - ).to(emb.device, non_blocking=True)

1566 + )

1567 + positions = async_tensor_h2d(positions_cpu, emb.device)

1568

1569 with gpu_sync_allowed():

1570 emb = emb[retention_mask]

vllm/model_executor/models/qwen2_vl.py

+2 -1

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

85 )

86 from vllm.sequence import IntermediateTensors

87 from vllm.utils.tensor_schema import TensorSchema, TensorShape

88 +from vllm.utils.torch_utils import async_tensor_h2d

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

90 from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers

@@ -707,7 +708,7 @@ class Qwen2VisionTransformer(nn.Module):

708 return {

709 "rotary_pos_emb_cos": rotary_pos_emb_cos,

710 "rotary_pos_emb_sin": rotary_pos_emb_sin,

710 - "cu_seqlens": cu_seqlens.to(device=device, non_blocking=True),

711 + "cu_seqlens": async_tensor_h2d(cu_seqlens, device),

712 "max_seqlen": max_seqlen,

713 }

714

vllm/model_executor/models/transformers/multimodal.py

+3 -2

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

60 )

61 from vllm.sequence import IntermediateTensors

62 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

63 +from vllm.utils.torch_utils import async_tensor_h2d

65 if TYPE_CHECKING:

66 from transformers import BatchFeature, PreTrainedModel

@@ -1180,8 +1181,8 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):

1181 # grid_thw fields are registered keep_on_cpu; restore the on-device

1182 # placement that HF get_image_features implementations expect.

1183 for key, value in kwargs.items():

1183 - if isinstance(value, torch.Tensor):

1184 - kwargs[key] = value.to(pixel_values.device, non_blocking=True)

1184 + if isinstance(value, torch.Tensor) and value.is_cpu:

1185 + kwargs[key] = async_tensor_h2d(value, pixel_values.device)

1186

1187 # The underlying HuggingFace `get_image_features` implementations

1188 # contain model-internal syncs (e.g. Idefics3 filters all-zero

vllm/multimodal/inputs.py

+37 -51

@@ -293,21 +293,39 @@ def _nested_tensors_are_cpu(tensors: NestedTensors) -> bool:

293 )

294

295

296 +def _is_dense(t: torch.Tensor) -> bool:

297 + """Whether `t`'s elements fill a contiguous storage range, possibly

298 + permuted (e.g. a transposed matrix). Such layouts can be copied with

299 + pitched cudaMemcpy2D/3DAsync directly from pinned memory, whereas

300 + gapped layouts (e.g. strided slices) stage through a pageable temp."""

301 + if not t.is_contiguous():

302 + expected_stride = 1

303 + for stride, size in sorted(zip(t.stride(), t.shape)):

304 + if size > 1:

305 + if stride != expected_stride:

306 + return False

307 + expected_stride *= size

308 + return True

309 +

310 +

311 def _nested_tensors_h2d(

312 tensors: NestedTensors,

313 device: torch.types.Device,

314 + *,

315 + pin_memory: bool = False,

316 ) -> NestedTensors:

317 if device is None:

318 return tensors

319

303 - return json_map_leaves(

304 - (

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

306 - if isinstance(x, torch.Tensor)

307 - else x

308 - ),

309 - tensors,

310 - )

320 + def _h2d(x: torch.Tensor) -> torch.Tensor:

321 + if not isinstance(x, torch.Tensor):

322 + return x

323 + if pin_memory and x.is_cpu and not (x.is_pinned() and _is_dense(x)):

324 + # Ensure tensor is pinned and dense for non_blocking H2D copy.

325 + x = x.new_empty(x.shape, pin_memory=True).copy_(x)

326 + return x.to(device=device, non_blocking=True)

327 +

328 + return json_map_leaves(_h2d, tensors)

329

330

331 BatchedTensorInputs: TypeAlias = dict[str, NestedTensors]

@@ -452,10 +470,7 @@ class BaseMultiModalField(ABC):

470

471 @abstractmethod

472 def _reduce_data(

455 - self,

456 - batch: list[NestedTensors],

457 - *,

458 - pin_memory: bool,

473 + self, batch: list[NestedTensors], *, pin_memory: bool

474 ) -> NestedTensors:

475 raise NotImplementedError

476

@@ -491,7 +506,7 @@ class BaseMultiModalField(ABC):

506 # concept, and `pin_memory()` rejects device tensors outright.

507 pin_memory = False

508 out = self._reduce_data(batch, pin_memory=pin_memory)

494 - return _nested_tensors_h2d(out, device=device)

509 + return _nested_tensors_h2d(out, device=device, pin_memory=pin_memory)

510

511

512 @dataclass(frozen=True, kw_only=True)

@@ -511,10 +526,7 @@ class MultiModalBatchedField(BaseMultiModalField):

526 return [field_factory(item) for item in data]

527

528 def _reduce_data(

514 - self,

515 - batch: list[NestedTensors],

516 - *,

517 - pin_memory: bool,

529 + self, batch: list[NestedTensors], *, pin_memory: bool

530 ) -> NestedTensors:

531 if len(batch) > 0 and is_list_of(batch, torch.Tensor, check="all"):

532 batch = cast(list[torch.Tensor], batch)

@@ -522,20 +534,11 @@ class MultiModalBatchedField(BaseMultiModalField):

534 # An optimization when `batch` contains only one tensor:

535 # - produce exactly same result as `torch.stack(batch)`

536 # - will achieve zero-copy if the tensor is contiguous

525 - out = batch[0].unsqueeze(0)

526 - if not pin_memory:

527 - return out.contiguous()

528 - # Avoid extra copy - pinning unpinned memory will make it contiguous

529 - if not out.is_contiguous() and out.is_pinned():

530 - out = out.contiguous()

531 - return out.pin_memory()

537 + return batch[0].unsqueeze(0)

538 first_shape = batch[0].shape

539 if all(elem.shape == first_shape for elem in batch):

534 - out = torch.empty(

535 - (len(batch), *batch[0].shape),

536 - dtype=batch[0].dtype,

537 - device=batch[0].device,

538 - pin_memory=pin_memory,

540 + out = batch[0].new_empty(

541 + (len(batch), *first_shape), pin_memory=pin_memory

542 )

543 return torch.stack(batch, out=out)

544

@@ -567,10 +570,7 @@ class MultiModalFlatField(BaseMultiModalField):

570 return [field_factory(data[cast(slice, s)]) for s in self.slices]

571

572 def _reduce_data(

570 - self,

571 - batch: list[NestedTensors],

572 - *,

573 - pin_memory: bool,

573 + self, batch: list[NestedTensors], *, pin_memory: bool

574 ) -> NestedTensors:

575 if len(batch) > 0 and is_list_of(batch, torch.Tensor, check="all"):

576 batch = cast(list[torch.Tensor], batch)

@@ -578,13 +578,7 @@ class MultiModalFlatField(BaseMultiModalField):

578 # An optimization when `batch` contains only one tensor:

579 # - produce exactly same result as `torch.concat(batch)`

580 # - will achieve zero-copy if the tensor is contiguous

581 - out = batch[0]

582 - if not pin_memory:

583 - return out.contiguous()

584 - # Avoid extra copy - pinning unpinned memory will make it contiguous

585 - if not out.is_contiguous() and out.is_pinned():

586 - out = out.contiguous()

587 - return out.pin_memory()

581 + return batch[0]

582

583 dim = self.dim + (self.dim < 0) * len(batch[0].shape)

584

@@ -596,11 +590,8 @@ class MultiModalFlatField(BaseMultiModalField):

590 if all(_shape_before_after(elem) == first_shape for elem in batch):

591 shape_before, shape_after = first_shape

592 shape_concat = sum(item.shape[dim] for item in batch)

599 - out = torch.empty(

600 - (*shape_before, shape_concat, *shape_after),

601 - dtype=batch[0].dtype,

602 - device=batch[0].device,

603 - pin_memory=pin_memory,

593 + out = batch[0].new_empty(

594 + (*shape_before, shape_concat, *shape_after), pin_memory=pin_memory

595 )

596 return torch.concat(batch, dim=self.dim, out=out)

597

@@ -622,12 +613,7 @@ class MultiModalFlatField(BaseMultiModalField):

613 max_sizes.append(max(t.shape[d] for t in batch))

614

615 # Step 2: Create zero-initialized output tensor

625 - out = torch.zeros(

626 - max_sizes,

627 - dtype=batch[0].dtype,

628 - device=batch[0].device,

629 - pin_memory=pin_memory,

630 - )

616 + out = batch[0].new_zeros(max_sizes, pin_memory=pin_memory)

617

618 # Step 3: Slice-assign each tensor to its proper position

619 concat_offset = 0

vllm/v1/attention/backends/flashinfer.py

+23 -4

@@ -1569,6 +1569,17 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1569 prefill_start : num_reqs + 1

1570 ]

1571 assert paged_kv_indptr_prefill_cpu.shape[0] == num_prefills + 1

1572 + # plan() copies these to the GPU with non_blocking=True;

1573 + # stage them in pinned memory so the copies stay async

1574 + # (the reused buffers themselves are intentionally not pinned).

1575 + if PIN_MEMORY:

1576 + qo_indptr_prefill_cpu = qo_indptr_prefill_cpu.pin_memory()

1577 + paged_kv_indptr_prefill_cpu = (

1578 + paged_kv_indptr_prefill_cpu.pin_memory()

1579 + )

1580 + paged_kv_last_page_len_prefill_cpu = (

1581 + paged_kv_last_page_len_prefill_cpu.pin_memory()

1582 + )

1583 if self.use_dcp:

1584 assert isinstance(prefill_wrapper, BatchDCPPrefillWrapper)

1585 prefill_wrapper.plan(

@@ -1695,13 +1706,21 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1706 o_dtype = (

1707 FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype

1708 )

1709 + paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: num_input_tokens + 1]

1710 + paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[

1711 + :num_input_tokens

1712 + ]

1713 + # plan() copies these to the GPU with non_blocking=True;

1714 + # stage them in pinned memory so the copies stay async

1715 + # (the reused buffers themselves are intentionally not pinned).

1716 + if PIN_MEMORY:

1717 + paged_kv_indptr_cpu = paged_kv_indptr_cpu.pin_memory()

1718 + paged_kv_last_page_len_cpu = paged_kv_last_page_len_cpu.pin_memory()

1719 fast_plan_decode(

1720 decode_wrapper,

1700 - indptr_cpu=self.paged_kv_indptr.cpu[: num_input_tokens + 1],

1721 + indptr_cpu=paged_kv_indptr_cpu,

1722 indices=paged_kv_indices,

1702 - last_page_len_cpu=self.paged_kv_last_page_len.cpu[

1703 - :num_input_tokens

1704 - ],

1723 + last_page_len_cpu=paged_kv_last_page_len_cpu,

1724 num_qo_heads=self.num_qo_heads * self.dcp_world_size,

1725 num_kv_heads=self.num_kv_heads,

1726 head_dim=self.head_dim,