From a758a9f671b6098c0b927272f2cc90b2c1dc48fa Mon Sep 17 00:00:00 2001 From: Keyi Li <94494390+JasonKeyiL@users.noreply.github.com> Date: Fri, 28 Aug 2026 22:20:17 -0700 Subject: [PATCH] [Perf][KV Connector] Pin token_indices before non_blocking H2D in hf3fs helper (#54293) Signed-off-by: Keyi Li Co-authored-by: Keyi Li Co-authored-by: Claude Opus 4.7 --- .../v1/hf3fs/utils/gather_scatter_helper.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py index aebaecd40b..4ba3ecee90 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py @@ -4,6 +4,7 @@ import torch from vllm.logger import init_logger from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import async_tensor_h2d @triton.jit @@ -140,9 +141,9 @@ def scatter_kv_caches( assert len(kv_cache_strides) == 4 device = src_tensor.device - token_indices_tensor = torch.tensor( - token_indices, dtype=torch.int32, device="cpu" - ).to(device, non_blocking=True) + token_indices_tensor = async_tensor_h2d( + token_indices, device=device, dtype=torch.int32 + ) grid = (num_layers, num_tokens_in_block) BLOCK_SIZE = 128 @@ -196,9 +197,9 @@ def gather_kv_caches( assert len(kv_cache_strides) == 4 device = dst_tensor.device - token_indices_tensor = torch.tensor( - token_indices, dtype=torch.int32, device="cpu" - ).to(device, non_blocking=True) + token_indices_tensor = async_tensor_h2d( + token_indices, device=device, dtype=torch.int32 + ) grid = (num_layers, num_tokens_in_block) BLOCK_SIZE = 128