Avoid long-blocking H2D copies in ViT (#51841)
Signed-off-by: Max Hu <hyoung2991@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Max Hu committed
Aug 12, 2026 at 10:54 UTC
b1b752042f622c692d5019c3ea122f2f7ee9d6ac
2 files changed
+23
-5
vllm/model_executor/models/qwen3_vl.py
+9
-1
@@ -705,7 +705,15 @@ class Qwen3_VisionTransformer(nn.Module):
705
else self.rot_pos_ids(h, w, self.spatial_merge_size).repeat(t, 1)
706
for t, h, w in grid_thw
707
]
708
- pos_ids = torch.cat(pos_ids, dim=0).to(self.device, non_blocking=True)
708
+ num_pos = sum(p.shape[0] for p in pos_ids)
709
+ pinned = torch.empty(
710
+ (num_pos, pos_ids[0].shape[1]),
711
+ dtype=pos_ids[0].dtype,
712
+ pin_memory=True,
713
+ )
714
+ pos_ids = torch.cat(pos_ids, dim=0, out=pinned).to(
715
+ self.device, non_blocking=True
716
+ )
717
718
# Use pre-computed cos_sin_cache from RotaryEmbedding
719
cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)
vllm/v1/worker/gpu_model_runner.py
+14
-4
@@ -2256,10 +2256,20 @@ class GPUModelRunner(
2256
2257
if self.uses_mrope:
2258
# Only relevant for models using M-RoPE (e.g, Qwen2-VL)
2259
- self.mrope_positions.gpu[:, :total_num_scheduled_tokens].copy_(
2260
- self.mrope_positions.cpu[:, :total_num_scheduled_tokens],
2261
- non_blocking=True,
2262
- )
2259
+ # Copy one row at a time. mrope_positions is allocated as
2260
+ # [3, max_num_tokens + 1] with a dummy trailing column to keep it
2261
+ # non-contiguous for torch.compile, so cpu[:, :N] is a strided view.
2262
+ # copy_() cannot express a strided source as a single
2263
+ # cudaMemcpyAsync, so it first gathers into a contiguous *pageable*
2264
+ # temporary, and a pageable H2D ignores non_blocking=True and
2265
+ # synchronizes the stream before the transfer starts. Each row is
2266
+ # contiguous within the pinned allocation, so per-row copies stay on
2267
+ # the pinned path and are genuinely asynchronous.
2268
+ for row in range(self.mrope_positions.gpu.shape[0]):
2269
+ self.mrope_positions.gpu[row, :total_num_scheduled_tokens].copy_(
2270
+ self.mrope_positions.cpu[row, :total_num_scheduled_tokens],
2271
+ non_blocking=True,
2272
+ )
2273
elif self.uses_xdrope_dim > 0:
2274
# Only relevant for models using XD-RoPE (e.g, HunYuan-VL)
2275
self.xdrope_positions.gpu[:, :total_num_scheduled_tokens].copy_(