[Perf] Pin CPU tensors before non_blocking H2D in three MM paths (#54292)
Signed-off-by: Keyi Li <likey6688@gmail.com> Co-authored-by: Keyi Li <likey6688@gmail.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Keyi Li committed
Aug 28, 2026 at 20:13 UTC
46a83642f686c4ea7efa39d5c1e1b50f373d901a
3 files changed
+5
-11
vllm/model_executor/layers/attention/mm_encoder_attention.py
+1
-4
@@ -259,10 +259,7 @@ class MMEncoderAttention(CustomOp):
259
sequence_lengths = add_padding_to_seqlens(
260
sequence_lengths, len(sequence_lengths), 0
261
)
262
- sequence_lengths = torch.from_numpy(sequence_lengths).to(
263
- device, non_blocking=True
264
- )
265
- return sequence_lengths
262
+ return async_tensor_h2d(sequence_lengths, device=device)
263
264
@classmethod
265
def maybe_recompute_cu_seqlens(
vllm/model_executor/models/kanana_v.py
+2
-4
@@ -38,6 +38,7 @@ from vllm.multimodal.processing import (
38
from vllm.sequence import IntermediateTensors
39
from vllm.utils.import_utils import resolve_obj_by_qualname
40
from vllm.utils.tensor_schema import TensorSchema, TensorShape
41
+from vllm.utils.torch_utils import async_tensor_h2d
42
43
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
44
from .qwen2_vl import Qwen2VisionTransformer
@@ -323,10 +324,7 @@ class CustomQwen2VLVE(Qwen2VisionTransformer):
324
grid_thw_np[:, 0],
325
).cumsum(axis=0, dtype=np.int32)
326
cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
326
- cu_seqlens = torch.from_numpy(cu_seqlens).to(
327
- self.device,
328
- non_blocking=True,
329
- )
327
+ cu_seqlens = async_tensor_h2d(cu_seqlens, device=self.device)
328
329
# Shape to (S, B, D) with batch dimension 1 as expected by the blocks.
330
x = x.unsqueeze(1)
vllm/model_executor/models/kimi_k25_vit.py
+2
-3
@@ -37,6 +37,7 @@ from vllm.model_executor.models.vision import (
37
)
38
from vllm.platforms import current_platform
39
from vllm.transformers_utils.configs.kimi_k25 import KimiK25VisionConfig
40
+from vllm.utils.torch_utils import async_tensor_h2d
41
42
logger = init_logger(__name__)
43
@@ -821,9 +822,7 @@ class MoonViT3dPretrainedModel(nn.Module):
822
merge_gather_idx = build_image_merge_gather_idx(
823
grid_thw_list, self.merge_kernel_size
824
)
824
- metadata["merge_gather_idx"] = torch.from_numpy(merge_gather_idx).to(
825
- device=device, non_blocking=True
826
- )
825
+ metadata["merge_gather_idx"] = async_tensor_h2d(merge_gather_idx, device=device)
826
return metadata
827
828