[Bugfix][DSV4.1] Avoid host sync in ViT CUDA graph replay metadata (#58499)
Kevin H. Luu committed
Sep 26, 2026 at 08:37 UTC
7d8c5fe9a95dd7a9734f05fdbad57c1ad61f77fe
2 files changed
+24
-15
vllm/models/deepseek_v4/common/vision.py
+10
-7
@@ -38,6 +38,7 @@ from vllm.model_executor.models.vision import (
38
get_load_balance_assignment,
39
is_vit_use_data_parallel,
40
)
41
+from vllm.utils.torch_utils import async_tensor_h2d
42
43
44
def _compute_vision_cos_sin(
@@ -219,8 +220,8 @@ class DeepseekV4ViT(nn.Module):
220
) -> torch.Tensor:
221
x = self.patch_embed(patches)
222
cos, sin = get_vision_cos_sin(n_vit_h, n_vit_w, self.rope_dim, self.rope_theta)
222
- cos = cos.to(device=x.device)
223
- sin = sin.to(device=x.device)
223
+ cos = async_tensor_h2d(cos, x.device)
224
+ sin = async_tensor_h2d(sin, x.device)
225
for block in self.blocks:
226
x = block(x, cos, sin)
227
return self.norm(x)
@@ -279,8 +280,8 @@ def build_packed_vit_metadata(
280
cu_seqlens.append(cu_seqlens[-1] + n)
281
max_seqlen = max(max_seqlen, n)
282
if cos_list:
282
- cos = torch.cat(cos_list).to(device)
283
- sin = torch.cat(sin_list).to(device)
283
+ cos = async_tensor_h2d(torch.cat(cos_list), device)
284
+ sin = async_tensor_h2d(torch.cat(sin_list), device)
285
else:
286
cos = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device)
287
sin = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device)
@@ -289,7 +290,7 @@ def build_packed_vit_metadata(
290
return {
291
"vit_cos": cos,
292
"vit_sin": sin,
292
- "cu_seqlens": torch.tensor(cu_seqlens, dtype=torch.int32, device=device),
293
+ "cu_seqlens": async_tensor_h2d(cu_seqlens, device, dtype=torch.int32),
294
# Read on the host by the attention wrapper; keep on CPU.
295
"max_seqlen": torch.tensor(max_seqlen, dtype=torch.int32),
296
}
@@ -327,8 +328,10 @@ def build_packed_merge_metadata(
328
mask_list.append(valid)
329
offset += n_h * n_w
330
if idx_list:
330
- merge_idx = torch.cat(idx_list).to(device)
331
- merge_mask = torch.cat(mask_list).unsqueeze(-1).to(device=device, dtype=dtype)
331
+ merge_idx = async_tensor_h2d(torch.cat(idx_list), device)
332
+ merge_mask = async_tensor_h2d(
333
+ torch.cat(mask_list).unsqueeze(-1), device, dtype=dtype
334
+ )
335
else:
336
merge_idx = torch.zeros((0, r * r), dtype=torch.int64, device=device)
337
merge_mask = torch.zeros((0, r * r, 1), dtype=dtype, device=device)
vllm/models/deepseek_v41/common/vl_cudagraph.py
+14
-8
@@ -25,6 +25,7 @@ from vllm.models.deepseek_v4.common.vision import (
25
build_packed_merge_metadata,
26
build_packed_vit_metadata,
27
)
28
+from vllm.utils.torch_utils import async_tensor_h2d
29
from vllm.v1.worker.encoder_cudagraph_defs import (
30
EncoderCudaGraphCaptureInputs,
31
EncoderCudaGraphConfig,
@@ -370,18 +371,23 @@ class DeepseekV4VLEncoderCudaGraphMixin:
371
r = self.config.vision_downsample_ratio
372
vit_grid = self._get_grid_list(batch_mm_kwargs, "vit_grid")
373
llm_grid = self._get_grid_list(batch_mm_kwargs, "llm_grid")
373
- types = batch_mm_kwargs["types"].to(aligner_out.device)
374
+ types = batch_mm_kwargs["types"]
375
376
n_rows = sum(-(-h // r) * (-(-w // r)) for h, w in vit_grid)
377
span_lens = [lh * (lw + 1) + 2 for lh, lw in llm_grid]
378
378
- # Batched span assembly: one masked fill per role across all items.
379
- dtype = aligner_out.dtype
380
- span = aligner_out.new_empty(sum(span_lens), self.config.hidden_size)
381
- span[types == IMAGE] = aligner_out[:n_rows]
382
- span[types == IMAGE_START] = self.image_start.to(dtype)
383
- span[types == IMAGE_END] = self.image_end.to(dtype)
384
- span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)
379
+ # Batched span assembly: one gather over [aligner rows, START,
380
+ # NEW_LINE, END] for all items. The index is built from the host-side
381
+ # ``types`` so that neither the copy nor the gather syncs.
382
+ is_image = types == IMAGE
383
+ src_idx = torch.where(
384
+ is_image,
385
+ is_image.cumsum(0) - 1,
386
+ n_rows + (types == IMAGE_NEW_LINE).long() + 2 * (types == IMAGE_END).long(),
387
+ )
388
+ delims = torch.stack([self.image_start, self.image_newline, self.image_end])
389
+ src = torch.cat([aligner_out[:n_rows], delims.to(aligner_out.dtype)])
390
+ span = src[async_tensor_h2d(src_idx, aligner_out.device)]
391
392
# Freshly allocated, so later replays cannot clobber the results.
393
offset = 0