@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6c18a54648

[Perf] Avoid h2d copies from non-pinned CPU tensors (#54299)

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

Nick Hill committed Aug 28, 2026 at 21:06 UTC 6c18a5464852d9ddf4c59df8ec278ce1683a917e

14 files changed +81 -84

vllm/model_executor/layers/pooler/tokwise/methods.py

+2 -1

@@ -11,6 +11,7 @@ from vllm.config import get_current_vllm_config

11 from vllm.config.pooler import TokenPoolingType

12 from vllm.model_executor.layers.pooler import PoolingParamsUpdate

13 from vllm.tasks import PoolingTask

14 +from vllm.utils.torch_utils import async_tensor_h2d

15 from vllm.v1.pool.metadata import PoolingMetadata

17 TokenPoolingMethodOutputItem: TypeAlias = torch.Tensor | None

@@ -131,7 +132,7 @@ class StepPool(AllPool):

132

133 if step_tag_id is not None:

134 idx_cpu = (token_id_cpu == step_tag_id).nonzero(as_tuple=True)[0]

134 - idx = idx_cpu.to(data.device, non_blocking=True)

135 + idx = async_tensor_h2d(idx_cpu, data.device)

136 data = data[idx]

137

138 pooled_data.append(data)

vllm/model_executor/model_loader/utils.py

+11 -10

@@ -158,13 +158,13 @@ def device_loading_context(module: torch.nn.Module, target_device: torch.device)

158 yield module

159 return

160

161 - original_device_states: dict[str, torch.device] = {}

161 + cpu_params: set[str] = set()

162 uva_offloaded_parameters: list[str] = []

163

164 - # Store original device states and move parameters to GPU if they're on CPU

164 + # Store which parameters are on CPU and move them to the GPU

165 for name, p in module.named_parameters():

166 if p.device.type == "cpu":

167 - original_device_states[name] = p.device

167 + cpu_params.add(name)

168 p.data = p.data.to(target_device)

169 if getattr(p, "_vllm_is_uva_offloaded", False):

170 uva_offloaded_parameters.append(name)

@@ -178,20 +178,21 @@ def device_loading_context(module: torch.nn.Module, target_device: torch.device)

178 is_pin_memory_available()

179 and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY

180 )

181 - # Restore parameters to their original devices, ignoring new parameters

181 + # Restore the CPU-resident parameters, ignoring new parameters.

182 for name, p in module.named_parameters():

183 - if name in original_device_states:

184 - original_device: torch.device = original_device_states[name]

185 - p.data = p.data.to(original_device)

183 + if name in cpu_params:

184 + p.data = torch.empty_like(

185 + p.data, device="cpu", pin_memory=use_pin_memory

186 + ).copy_(p.data)

187

188 # parameter is UVA offloaded, but was replaced with a new device tensor

189 # re-offload it to CPU using UVA

190 if name in uva_offloaded_parameters and not getattr(

191 p, "_vllm_is_uva_offloaded", False

192 ):

192 - cpu_data = p.data.to(device="cpu")

193 - if use_pin_memory:

194 - cpu_data = cpu_data.pin_memory()

193 + cpu_data = torch.empty_like(

194 + p.data, device="cpu", pin_memory=use_pin_memory

195 + ).copy_(p.data)

196 p.data = get_accelerator_view_from_cpu_tensor(cpu_data)

197 p._vllm_is_uva_offloaded = True

198

vllm/model_executor/models/ernie45_vl.py

+3 -5

@@ -425,10 +425,11 @@ class Ernie4_5_VisionTransformer(nn.Module):

425 return self.patch_embed.proj.weight.device

426

427 def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:

428 + device = self.device

429 pos_ids = []

430 for t, h, w in grid_thw:

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

431 - wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)

431 + hpos_ids = torch.arange(h, device=device).unsqueeze(1).expand(-1, w)

432 + wpos_ids = torch.arange(w, device=device).unsqueeze(0).expand(h, -1)

433 hpos_ids = (

434 hpos_ids.reshape(

435 h // self.spatial_merge_size,

@@ -453,9 +454,6 @@ class Ernie4_5_VisionTransformer(nn.Module):

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

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

456 rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)

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

457 - # gather below doesn't index a device tensor with a CPU one.

458 - pos_ids = pos_ids.to(rotary_pos_emb_full.device, non_blocking=True)

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

458 return rotary_pos_emb

459

vllm/model_executor/models/glm4_1v.py

+3 -3

@@ -709,10 +709,11 @@ class Glm4vVisionTransformer(nn.Module):

709 def rot_pos_emb(

710 self, grid_thw: list[list[int]]

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

712 + device = self.device

713 pos_ids = []

714 for t, h, w in grid_thw:

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

715 - wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)

715 + hpos_ids = torch.arange(h, device=device).unsqueeze(1).expand(-1, w)

716 + wpos_ids = torch.arange(w, device=device).unsqueeze(0).expand(h, -1)

717 hpos_ids = (

718 hpos_ids.reshape(

719 h // self.spatial_merge_size,

@@ -740,7 +741,6 @@ class Glm4vVisionTransformer(nn.Module):

741 # Use pre-computed cos_sin_cache from RotaryEmbedding

742 cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)

743

743 - pos_ids = pos_ids.to(cos.device, non_blocking=True)

744 cos_combined = cos[pos_ids].flatten(1)

745 sin_combined = sin[pos_ids].flatten(1)

746 return cos_combined, sin_combined, pos_ids

vllm/model_executor/models/idefics2_vision_model.py

+4 -1

@@ -39,6 +39,7 @@ from vllm.model_executor.layers.linear import (

39 )

40 from vllm.model_executor.layers.quantization import QuantizationConfig

41 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

42 +from vllm.utils.torch_utils import PIN_MEMORY

44 from .utils import AutoWeightsLoader, WeightsMapper

45 from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model

@@ -85,7 +86,9 @@ class Idefics2VisionEmbeddings(nn.Module):

86 1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side

87 )

88 position_ids = torch.full(

88 - size=(batch_size, max_nb_patches_h * max_nb_patches_w), fill_value=0

89 + size=(batch_size, max_nb_patches_h * max_nb_patches_w),

90 + fill_value=0,

91 + pin_memory=PIN_MEMORY,

92 )

94 with gpu_sync_allowed():

vllm/model_executor/models/phi4mm_audio.py

+3 -6

@@ -32,6 +32,7 @@ from vllm.model_executor.models.phi4mm_utils import (

32 get_offset,

33 unfold_tensor,

34 )

35 +from vllm.utils.torch_utils import async_tensor_h2d

38 class ConformerEncoderLayer(nn.Module):

@@ -586,9 +587,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

587 seq_len, batch_size, self.chunk_size, self.left_chunk

588 )

589 device = xs_pad.device

589 - enc_streaming_mask = enc_streaming_mask.contiguous().to(

590 - device, non_blocking=True

591 - )

590 + enc_streaming_mask = async_tensor_h2d(enc_streaming_mask.contiguous(), device)

591 xs_pad = xs_pad.to(device)

592

593 input_tensor = xs_pad

@@ -921,9 +920,7 @@ class ConformerEncoder(TransformerEncoderBase):

920 enc_streaming_mask = self._streaming_mask(

921 max_audio_length, batch_size, self.chunk_size, self.left_chunk

922 )

924 - enc_streaming_mask = enc_streaming_mask.contiguous().to(

925 - device, non_blocking=True

926 - )

923 + enc_streaming_mask = async_tensor_h2d(enc_streaming_mask.contiguous(), device)

924 if mask is None:

925 return enc_streaming_mask

926

vllm/model_executor/models/qwen2_5_vl.py

+5 -7

@@ -715,8 +715,11 @@ class Qwen2_5_VisionTransformer(nn.Module):

715 return self.patch_embed.proj.weight.device

716

717 def rotary_pos_emb_thw(self, t, h, w):

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

719 - wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)

718 + # Use pre-computed cos_sin_cache from RotaryEmbedding

719 + cos, sin = self.rotary_pos_emb.get_cos_sin(max(h, w))

720 +

721 + hpos_ids = torch.arange(h, device=cos.device).unsqueeze(1).expand(-1, w)

722 + wpos_ids = torch.arange(w, device=cos.device).unsqueeze(0).expand(h, -1)

723 hpos_ids = (

724 hpos_ids.reshape(

725 h // self.spatial_merge_size,

@@ -738,12 +741,7 @@ class Qwen2_5_VisionTransformer(nn.Module):

741 .flatten()

742 )

743 pos_ids = torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)

741 - max_size = max(h, w)

742 -

743 - # Use pre-computed cos_sin_cache from RotaryEmbedding

744 - cos, sin = self.rotary_pos_emb.get_cos_sin(max_size)

744

746 - pos_ids = pos_ids.to(cos.device, non_blocking=True)

745 cos_combined = cos[pos_ids].flatten(1)

746 sin_combined = sin[pos_ids].flatten(1)

747

vllm/model_executor/models/qwen2_vl.py

+3 -3

@@ -609,11 +609,12 @@ class Qwen2VisionTransformer(nn.Module):

609 def rot_pos_emb(

610 self, grid_thw: list[list[int]]

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

612 + device = self.device

613 pos_ids = []

614 max_grid_size = 0

615 for t, h, w in grid_thw:

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

616 - wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)

616 + hpos_ids = torch.arange(h, device=device).unsqueeze(1).expand(-1, w)

617 + wpos_ids = torch.arange(w, device=device).unsqueeze(0).expand(h, -1)

618 hpos_ids = (

619 hpos_ids.reshape(

620 h // self.spatial_merge_size,

@@ -641,7 +642,6 @@ class Qwen2VisionTransformer(nn.Module):

642 # Use pre-computed cos_sin_cache from RotaryEmbedding

643 cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)

644

644 - pos_ids = pos_ids.to(cos.device, non_blocking=True)

645 cos_combined = cos[pos_ids].flatten(1)

646 sin_combined = sin[pos_ids].flatten(1)

647 return cos_combined, sin_combined

vllm/model_executor/models/qwen3_asr.py

+3 -2

@@ -94,6 +94,7 @@ from vllm.transformers_utils.processors.qwen3_asr import (

94 Qwen3ASRProcessor,

95 )

96 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

97 +from vllm.utils.torch_utils import async_tensor_h2d

99 logger = init_logger(__name__)

100 _ASR_TEXT_TAG = "<asr_text>"

@@ -465,8 +466,8 @@ class Qwen3ASRForConditionalGeneration(

466 input_features = audio_input["input_features"]

467 # audio_feature_lengths is keep_on_cpu; the audio tower derives

468 # device placement from feature_lens, so move it explicitly.

468 - audio_feature_lengths = audio_input["audio_feature_lengths"].to(

469 - input_features.device, non_blocking=True

469 + audio_feature_lengths = async_tensor_h2d(

470 + audio_input["audio_feature_lengths"], input_features.device

471 )

472

473 audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths)

vllm/model_executor/models/qwen3_omni_moe_thinker.py

+7 -7

@@ -842,9 +842,14 @@ class Qwen3Omni_VisionTransformer(nn.Module):

842 return self.patch_embed.proj.weight.device

843

844 def rot_pos_emb(self, grid_thw):

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

846 +

847 + # Use pre-computed cos_sin_cache from RotaryEmbedding

848 + cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)

849 +

850 pos_ids = []

851 for t, h, w in grid_thw:

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

852 + hpos_ids = torch.arange(h, device=cos.device).unsqueeze(1).expand(-1, w)

853 hpos_ids = hpos_ids.reshape(

854 h // self.spatial_merge_size,

855 self.spatial_merge_size,

@@ -854,7 +859,7 @@ class Qwen3Omni_VisionTransformer(nn.Module):

859 hpos_ids = hpos_ids.permute(0, 2, 1, 3)

860 hpos_ids = hpos_ids.flatten()

861

857 - wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)

862 + wpos_ids = torch.arange(w, device=cos.device).unsqueeze(0).expand(h, -1)

863 wpos_ids = wpos_ids.reshape(

864 h // self.spatial_merge_size,

865 self.spatial_merge_size,

@@ -865,12 +870,7 @@ class Qwen3Omni_VisionTransformer(nn.Module):

870 wpos_ids = wpos_ids.flatten()

871 pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))

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

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

869 -

870 - # Use pre-computed cos_sin_cache from RotaryEmbedding

871 - cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)

873

873 - pos_ids = pos_ids.to(cos.device, non_blocking=True)

874 cos_combined = cos[pos_ids].flatten(1)

875 sin_combined = sin[pos_ids].flatten(1)

876

vllm/model_executor/models/siglip.py

+2 -1

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

56 from vllm.sequence import IntermediateTensors

57 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

58 from vllm.utils.tensor_schema import TensorSchema, TensorShape

59 +from vllm.utils.torch_utils import async_tensor_h2d

61 from .clip import dual_encoder_has_text_tokens, merge_dual_encoder_text_and_vision

62 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant

@@ -1010,7 +1011,7 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):

1011 flip_indices_cpu = (

1012 starts_cpu[sequence_ids_cpu] + ends_cpu[sequence_ids_cpu]

1013 ) - (1 + current_positions_cpu)

1013 - flip_indices = flip_indices_cpu.to(features.device, non_blocking=True)

1014 + flip_indices = async_tensor_h2d(flip_indices_cpu, features.device)

1015

1016 return features[flip_indices]

1017

vllm/v1/attention/backends/flashinfer.py

+19 -33

@@ -917,16 +917,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

917 arch,

918 )

919 # Preparing persistent buffers

920 - # Since we do not have explicit synchronization in ModelRunnerV2, we do not pin

921 - # reused CPU buffers to avoid a race condition between step N async copies to

922 - # GPU and step N+1 buffer updates.

923 - self.pin_memory = not vllm_config.use_v2_model_runner and PIN_MEMORY

924 - self.paged_kv_indptr = self._make_buffer(max_num_reqs + 1)

925 - self.paged_kv_indptr_cpu_buffer = torch.zeros_like(

926 - self.paged_kv_indptr.cpu, pin_memory=self.pin_memory

927 - ) # Extra buffer for mutable paged_kv_indptr.cpu in cuda graph mode

928 - self.paged_kv_indices = self._make_buffer(max_num_pages)

929 - self.paged_kv_last_page_len = self._make_buffer(max_num_reqs)

920 + self.paged_kv_indptr = CpuGpuBuffer(

921 + max_num_reqs + 1, dtype=torch.int32, device=self.device, pin_memory=False

922 + )

923 + self.paged_kv_indices = torch.zeros(

924 + max_num_pages, dtype=torch.int32, device=self.device

925 + )

926 + self.paged_kv_last_page_len = CpuGpuBuffer(

927 + max_num_reqs, dtype=torch.int32, device=self.device, pin_memory=False

928 + )

929

930 @property

931 def kv_cache_layout(self) -> KVCacheLayout:

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

969 return FlashInferBackend.get_dtype_for_flashinfer("fp8_e4m3")

970 return self.kv_cache_spec.dtype

971

973 - def _make_buffer(

974 - self, *size: int | torch.SymInt, dtype: torch.dtype = torch.int32

975 - ) -> CpuGpuBuffer:

976 - return CpuGpuBuffer(

977 - *size,

978 - dtype=dtype,

979 - device=self.device,

980 - pin_memory=self.pin_memory,

981 - with_numpy=True,

982 - )

983 -

972 @override # type: ignore[misc]

973 @classmethod

974 def get_cudagraph_support(

@@ -1197,7 +1185,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1185 if decode_wrapper is None:

1186 if use_cudagraph:

1187 paged_kv_indptr = self.paged_kv_indptr.gpu[: batch_size + 1]

1200 - paged_kv_indices = self.paged_kv_indices.gpu

1188 + paged_kv_indices = self.paged_kv_indices

1189 paged_kv_last_page_len = self.paged_kv_last_page_len.gpu[:batch_size]

1190 else:

1191 paged_kv_indptr = None

@@ -1260,20 +1248,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1248 dtype=np.int32,

1249 out=self.paged_kv_indptr.np[1 : num_reqs + 1],

1250 )

1263 - # NOTE(woosuk): Because self.paged_kv_indptr_cpu can be modified

1264 - # after this line (e.g., for cuda graphs), we need to copy the data to

1265 - # self.paged_kv_indptr_buffer to avoid race condition.

1266 - self.paged_kv_indptr_cpu_buffer[: num_reqs + 1] = self.paged_kv_indptr.cpu[

1267 - : num_reqs + 1

1268 - ]

1251 paged_kv_indptr = self.paged_kv_indptr.gpu[: num_reqs + 1]

1270 - paged_kv_indptr.copy_(

1271 - self.paged_kv_indptr_cpu_buffer[: num_reqs + 1], non_blocking=True

1272 - )

1252 + paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: num_reqs + 1]

1253 + if PIN_MEMORY:

1254 + paged_kv_indptr_cpu = paged_kv_indptr_cpu.pin_memory()

1255 + paged_kv_indptr.copy_(paged_kv_indptr_cpu, non_blocking=True)

1256

1257 # write self.paged_kv_indices inplace

1258 num_actual_pages = self.paged_kv_indptr.np[num_reqs]

1276 - paged_kv_indices = self.paged_kv_indices.gpu[:num_actual_pages]

1259 + paged_kv_indices = self.paged_kv_indices[:num_actual_pages]

1260 _copy_page_indices_kernel[(num_reqs,)](

1261 paged_kv_indices,

1262 block_table_tensor,

@@ -1289,8 +1272,11 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1272 page_size,

1273 paged_kv_last_page_len_np,

1274 )

1275 + paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[:num_reqs]

1276 + if PIN_MEMORY:

1277 + paged_kv_last_page_len_cpu = paged_kv_last_page_len_cpu.pin_memory()

1278 self.paged_kv_last_page_len.gpu[:num_reqs].copy_(

1293 - self.paged_kv_last_page_len.cpu[:num_reqs], non_blocking=True

1279 + paged_kv_last_page_len_cpu, non_blocking=True

1280 )

1281 return paged_kv_indices

1282

vllm/v1/outputs.py

+12 -4

@@ -11,6 +11,7 @@ import numpy as np

11 import torch

13 from vllm.compilation.cuda_graph import CUDAGraphStat

14 +from vllm.utils.torch_utils import PIN_MEMORY

15 from vllm.v1.core.sched.output import SchedulerOutput

17 if TYPE_CHECKING:

@@ -182,11 +183,18 @@ class LogprobsTensors(NamedTuple):

183 """Create empty LogprobsTensors on CPU."""

184

185 logprob_token_ids = torch.empty(

185 - (num_positions, num_tokens_per_position), dtype=torch.int32, device="cpu"

186 + (num_positions, num_tokens_per_position),

187 + dtype=torch.int32,

188 + device="cpu",

189 + pin_memory=PIN_MEMORY,

190 )

187 - logprobs = torch.empty_like(logprob_token_ids, dtype=torch.float32)

188 - selected_token_ranks = torch.empty(

189 - num_positions, dtype=torch.int32, device="cpu"

191 + logprobs = logprob_token_ids.new_empty(

192 + (num_positions, num_tokens_per_position),

193 + dtype=torch.float32,

194 + pin_memory=PIN_MEMORY,

195 + )

196 + selected_token_ranks = logprob_token_ids.new_empty(

197 + num_positions, pin_memory=PIN_MEMORY

198 )

199 return LogprobsTensors(

200 logprob_token_ids=logprob_token_ids,

vllm/v1/worker/gpu/mm/encoder_runner.py

+4 -1

@@ -215,7 +215,10 @@ class EncoderRunner:

215 num_computed_tokens = num_computed_tokens + draft_lookahead

216

217 is_mm_embed = torch.zeros(

218 - total_num_scheduled_tokens, dtype=torch.bool, device="cpu"

218 + total_num_scheduled_tokens,

219 + dtype=torch.bool,

220 + device="cpu",

221 + pin_memory=PIN_MEMORY,

222 )

223

224 # Whether to gather media embeddings this step.