xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-01-07 12:01:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

[NIXL] Fix after virtual block_size for host_buffer with heter kv_layout (#29122)

Browse Source 
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
This commit is contained in:
Chendi.Xue 2025-11-21 10:55:27 -06:00 committed by GitHub
parent b4c8fbaae2
commit 460d02a417
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 13 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
View File

@ -1042,10 +1042,12 @@ class NixlConnectorWorker:
NOT directly supported by NIXL (e.g., tpu)
"""
xfer_buffers: dict[str, torch.Tensor] = {}
inv_order = [0, 1, 3, 2, 4]
try:
for layer_name, kv_cache in kv_caches.items():
kv_shape = kv_cache.shape
kv_dtype = kv_cache.dtype
permute_shape = False
if (
self.kv_cache_layout == "NHD"
and self.vllm_config.kv_transfer_config is not None
@ -1059,10 +1061,20 @@ class NixlConnectorWorker:
# Since NHD will not support Decode/Prefill TP_ratio > 1,
# we can leverage host_buffer for permute
self.host_buffer_kv_cache_layout = "HND"
kv_shape = tuple(kv_shape[i] for i in [0, 1, 3, 2, 4])
kv_shape = (
tuple(kv_shape[i] for i in inv_order)
if not self.use_mla
else kv_shape
)
permute_shape = not self.use_mla
xfer_buffers[layer_name] = torch.empty(
kv_shape, dtype=kv_dtype, device="cpu"
)
if permute_shape:
xfer_buffers[layer_name] = xfer_buffers[layer_name].permute(
inv_order
)
except MemoryError as e:
logger.error("NIXLConnectorWorker gets %s.", e)
raise

8
vllm/platforms/xpu.py
View File

@ -251,10 +251,6 @@ class XPUPlatform(Platform):
) -> None:
"""Copy blocks from src_cache to dst_cache on XPU."""
_src_cache = src_cache[:, src_block_indices]
if _src_cache.shape[2:] != dst_cache.shape[2:]:
# To support TP_ratio, HOST KV might be initiated with HND
# while XPU device KV is with NHD
_src_cache = _src_cache.permute(0, 1, 3, 2, 4)
dst_cache[:, dst_block_indices] = _src_cache.to(dst_cache.device)
@classmethod
@ -267,8 +263,4 @@ class XPUPlatform(Platform):
) -> None:
"""Copy blocks from XPU to host (CPU)."""
_src_cache = src_cache[:, src_block_indices]
if _src_cache.shape[2:] != dst_cache.shape[2:]:
# XPU device KV is with NHD while HOST KV
# might be initiated with HND for TP_ratio support
_src_cache = _src_cache.permute(0, 1, 3, 2, 4)
dst_cache[:, dst_block_indices] = _src_cache.cpu()