[Perf][KV Connector] Pin token_indices before non_blocking H2D in hf3fs helper (#54293)
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 22:20 UTC
a758a9f671b6098c0b927272f2cc90b2c1dc48fa
1 file changed
+7
-6
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py
+7
-6
@@ -4,6 +4,7 @@ import torch
4
5
from vllm.logger import init_logger
6
from vllm.triton_utils import tl, triton
7
+from vllm.utils.torch_utils import async_tensor_h2d
8
9
10
@triton.jit
@@ -140,9 +141,9 @@ def scatter_kv_caches(
141
assert len(kv_cache_strides) == 4
142
143
device = src_tensor.device
143
- token_indices_tensor = torch.tensor(
144
- token_indices, dtype=torch.int32, device="cpu"
145
- ).to(device, non_blocking=True)
144
+ token_indices_tensor = async_tensor_h2d(
145
+ token_indices, device=device, dtype=torch.int32
146
+ )
147
148
grid = (num_layers, num_tokens_in_block)
149
BLOCK_SIZE = 128
@@ -196,9 +197,9 @@ def gather_kv_caches(
197
assert len(kv_cache_strides) == 4
198
199
device = dst_tensor.device
199
- token_indices_tensor = torch.tensor(
200
- token_indices, dtype=torch.int32, device="cpu"
201
- ).to(device, non_blocking=True)
200
+ token_indices_tensor = async_tensor_h2d(
201
+ token_indices, device=device, dtype=torch.int32
202
+ )
203
204
grid = (num_layers, num_tokens_in_block)
205
BLOCK_SIZE = 128