A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 18:30 UTC
[NIXL] Fix after virtual block_size for host_buffer with heter kv_layout (#29122)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
Chendi.Xue committed
Nov 21, 2025 at 10:55 UTC
460d02a417b440ce8b3b8d09c6f5214a2a346426
2 files changed
+13
-9
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+13
-1
@@ -1042,10 +1042,12 @@ class NixlConnectorWorker:
1042
NOT directly supported by NIXL (e.g., tpu)
1043
"""
1044
xfer_buffers: dict[str, torch.Tensor] = {}
1045
+ inv_order = [0, 1, 3, 2, 4]
1046
try:
1047
for layer_name, kv_cache in kv_caches.items():
1048
kv_shape = kv_cache.shape
1049
kv_dtype = kv_cache.dtype
1050
+ permute_shape = False
1051
if (
1052
self.kv_cache_layout == "NHD"
1053
and self.vllm_config.kv_transfer_config is not None
@@ -1059,10 +1061,20 @@ class NixlConnectorWorker:
1061
# Since NHD will not support Decode/Prefill TP_ratio > 1,
1062
# we can leverage host_buffer for permute
1063
self.host_buffer_kv_cache_layout = "HND"
1062
- kv_shape = tuple(kv_shape[i] for i in [0, 1, 3, 2, 4])
1064
+ kv_shape = (
1065
+ tuple(kv_shape[i] for i in inv_order)
1066
+ if not self.use_mla
1067
+ else kv_shape
1068
+ )
1069
+ permute_shape = not self.use_mla
1070
+
1071
xfer_buffers[layer_name] = torch.empty(
1072
kv_shape, dtype=kv_dtype, device="cpu"
1073
)
1074
+ if permute_shape:
1075
+ xfer_buffers[layer_name] = xfer_buffers[layer_name].permute(
1076
+ inv_order
1077
+ )
1078
except MemoryError as e:
1079
logger.error("NIXLConnectorWorker gets %s.", e)
1080
raise
vllm/platforms/xpu.py
-8
@@ -251,10 +251,6 @@ class XPUPlatform(Platform):
251
) -> None:
252
"""Copy blocks from src_cache to dst_cache on XPU."""
253
_src_cache = src_cache[:, src_block_indices]
254
- if _src_cache.shape[2:] != dst_cache.shape[2:]:
255
- # To support TP_ratio, HOST KV might be initiated with HND
256
- # while XPU device KV is with NHD
257
- _src_cache = _src_cache.permute(0, 1, 3, 2, 4)
254
dst_cache[:, dst_block_indices] = _src_cache.to(dst_cache.device)
255
256
@classmethod
@@ -267,8 +263,4 @@ class XPUPlatform(Platform):
263
) -> None:
264
"""Copy blocks from XPU to host (CPU)."""
265
_src_cache = src_cache[:, src_block_indices]
270
- if _src_cache.shape[2:] != dst_cache.shape[2:]:
271
- # XPU device KV is with NHD while HOST KV
272
- # might be initiated with HND for TP_ratio support
273
- _src_cache = _src_cache.permute(0, 1, 3, 2, 4)
266
dst_cache[:, dst_block_indices] = _src_cache.cpu()