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 14, 05:30 UTC
[NIXL] use Host buffer to support TP_ratio > 1 for XPU (#27140)
Signed-off-by: Chendi Xue <chendi.xue@intel.com> Signed-off-by: Chendi.Xue <chendi.xue@intel.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@gmail.com>
Chendi.Xue committed
Oct 22, 2025 at 10:28 UTC
7c4767f1eb29f331c817de7e41eb39a06f83420c
2 files changed
+34
-5
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+24
-5
@@ -755,6 +755,7 @@ class NixlConnectorWorker:
755
self._use_flashinfer = attn_backend == _Backend.FLASHINFER
756
self._use_pallas = attn_backend == _Backend.PALLAS
757
self.kv_cache_layout = get_kv_cache_layout()
758
+ self.host_buffer_kv_cache_layout = self.kv_cache_layout
759
logger.debug("Detected attention backend %s", self.backend_name)
760
logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
761
@@ -875,6 +876,20 @@ class NixlConnectorWorker:
876
for layer_name, kv_cache in kv_caches.items():
877
kv_shape = kv_cache.shape
878
kv_dtype = kv_cache.dtype
879
+ if (
880
+ self.kv_cache_layout == "NHD"
881
+ and self.vllm_config.kv_transfer_config is not None
882
+ and self.vllm_config.kv_transfer_config.enable_permute_local_kv
883
+ ):
884
+ logger.info_once(
885
+ "'enable_permute_local_kv' flag is enabled while "
886
+ "device KV Layout is NHD. Init host buffer with"
887
+ " HND to better support Decode/Prefill TP_ratio > 1."
888
+ )
889
+ # Since NHD will not support Decode/Prefill TP_ratio > 1,
890
+ # we can leverage host_buffer for permute
891
+ self.host_buffer_kv_cache_layout = "HND"
892
+ kv_shape = tuple(kv_shape[i] for i in [0, 1, 3, 2, 4])
893
xfer_buffers[layer_name] = torch.empty(
894
kv_shape, dtype=kv_dtype, device="cpu"
895
)
@@ -1110,7 +1125,9 @@ class NixlConnectorWorker:
1125
num_blocks=self.num_blocks,
1126
block_lens=self.block_len_per_layer,
1127
attn_backend_name=self.backend_name,
1113
- kv_cache_layout=self.kv_cache_layout,
1128
+ kv_cache_layout=self.kv_cache_layout
1129
+ if not self.use_host_buffer
1130
+ else self.host_buffer_kv_cache_layout,
1131
)
1132
ready_event, stop_event = threading.Event(), threading.Event()
1133
self._nixl_handshake_listener_t = threading.Thread(
@@ -1273,7 +1290,12 @@ class NixlConnectorWorker:
1290
assert not self._use_pallas or tp_ratio == 1, (
1291
"TPU (pallas_v1) DOES NOT support heterogeneous TP yet."
1292
)
1276
- if not self.use_mla and nixl_agent_meta.kv_cache_layout != self.kv_cache_layout:
1293
+ kv_cache_layout = (
1294
+ self.kv_cache_layout
1295
+ if not self.use_host_buffer
1296
+ else self.host_buffer_kv_cache_layout
1297
+ )
1298
+ if not self.use_mla and nixl_agent_meta.kv_cache_layout != kv_cache_layout:
1299
if (
1300
self.kv_transfer_config.enable_permute_local_kv
1301
and nixl_agent_meta.kv_cache_layout == "HND"
@@ -1299,9 +1321,6 @@ class NixlConnectorWorker:
1321
)
1322
remote_block_size = remote_block_len // (self.slot_size_per_layer[0])
1323
else:
1302
- if tp_ratio > 1 and self.device_type == "xpu":
1303
- # XPU uses NHD, hence it does not support splitting on H
1304
- raise ValueError("Heterogeneous TP is not supported on XPU")
1324
# When MLA is not used, this is a list of the same block length
1325
for block_len in nixl_agent_meta.block_lens:
1326
assert block_len == remote_block_len, (
vllm/platforms/xpu.py
+10
@@ -144,6 +144,8 @@ class XPUPlatform(Platform):
144
# check and update parallel config
145
parallel_config = vllm_config.parallel_config
146
parallel_config.worker_cls = "vllm.v1.worker.xpu_worker.XPUWorker"
147
+ if vllm_config.kv_transfer_config is not None:
148
+ vllm_config.kv_transfer_config.enable_permute_local_kv = True
149
150
if parallel_config.distributed_executor_backend is None:
151
if parallel_config.world_size > 1:
@@ -245,6 +247,10 @@ class XPUPlatform(Platform):
247
) -> None:
248
"""Copy blocks from src_cache to dst_cache on XPU."""
249
_src_cache = src_cache[:, src_block_indices]
250
+ if _src_cache.shape[2:] != dst_cache.shape[2:]:
251
+ # To support TP_ratio, HOST KV might be initiated with HND
252
+ # while XPU device KV is with NHD
253
+ _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
@@ -257,4 +263,8 @@ class XPUPlatform(Platform):
263
) -> None:
264
"""Copy blocks from XPU to host (CPU)."""
265
_src_cache = src_cache[:, src_block_indices]
266
+ if _src_cache.shape[2:] != dst_cache.shape[2:]:
267
+ # XPU device KV is with NHD while HOST KV
268
+ # might be initiated with HND for TP_ratio support
269
+ _src_cache = _src_cache.permute(0, 1, 3, 2, 4)
270
dst_cache[:, dst_block_indices] = _src_cache.cpu()