[CPU] fixes heterogeneous NIXL KV transfer into CPU_ATTN decode workers (#47871)
Signed-off-by: Spycsh <sihan.chen@intel.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Sihan Chen committed
Jul 20, 2026 at 15:33 UTC
4938d44a3b818fd443cf45aea92e670f16973ba0
2 files changed
+15
-13
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+1
-6
@@ -1935,13 +1935,8 @@ class NixlBaseConnectorWorker:
1935
indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
1936
1937
for _, cache_or_caches in self.device_kv_caches.items():
1938
- blocks_to_update = cache_or_caches.index_select(1, indices)
1938
current_platform.pack_kv_cache(
1940
- key=blocks_to_update[0],
1941
- value=blocks_to_update[1],
1942
- key_cache=cache_or_caches[0],
1943
- value_cache=cache_or_caches[1],
1944
- block_ids=block_ids,
1939
+ kv_cache=cache_or_caches,
1940
indices=indices,
1941
)
1942
vllm/platforms/cpu.py
+14
-7
@@ -461,11 +461,7 @@ class CpuPlatform(Platform):
461
@classmethod
462
def pack_kv_cache(
463
cls,
464
- key: torch.Tensor,
465
- value: torch.Tensor,
466
- key_cache: torch.Tensor,
467
- value_cache: torch.Tensor,
468
- block_ids: list[int],
464
+ kv_cache: torch.Tensor,
465
indices: torch.Tensor,
466
) -> None:
467
"""
@@ -476,15 +472,26 @@ class CpuPlatform(Platform):
472
from vllm._custom_ops import cpu_attn_reshape_and_cache
473
from vllm.v1.attention.backends.cpu_attn import _get_attn_isa
474
475
+ num_blocks, num_kv_heads, block_size, fused_head_size = kv_cache.shape
476
+ head_size = fused_head_size // 2
477
+
478
+ # Fused path used by heterogeneous NIXL CPU_ATTN post-processing.
479
+ blocks_to_update = kv_cache.index_select(0, indices)
480
+ key = blocks_to_update[..., :head_size]
481
+ value = blocks_to_update[..., head_size:]
482
+
483
+ key_cache, value_cache = kv_cache.view(
484
+ num_blocks, num_kv_heads, block_size * 2, head_size
485
+ ).chunk(2, dim=2)
486
+
487
dtype = key.dtype
488
# For CPU_ATTN, the shape is [N, num_kv_heads, block_size, head_size]
481
- _, _, block_size, head_size = key_cache.shape
489
key = key.permute(0, 2, 1, 3).flatten(0, 1)
490
value = value.permute(0, 2, 1, 3).flatten(0, 1)
491
492
isa = _get_attn_isa(dtype, block_size, head_size)
493
block_offsets = torch.arange(block_size, device="cpu", dtype=torch.long)
487
- num_blocks = len(block_ids)
494
+ num_blocks = indices.numel()
495
slot_mapping = (
496
block_offsets.reshape(1, block_size)
497
+ indices.reshape(num_blocks, 1) * block_size