iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Perf][KV Connector] Pin token_indices before non_blocking H2D in hf3fs helper (#54293)

Browse Source 
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>
This commit is contained in:
Keyi Li
2026-08-28 22:20:17 -07:00
committed by GitHub
co-authored by Keyi Li Claude Opus 4.7
parent f4f3bcd4b4
commit a758a9f671
1 changed files with 7 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py
+7 -6
View File
@@ -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