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 2, 02:15 UTC
[Attention][CPU] Standardize kv layout to blocks first (#44393)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Jun 3, 2026 at 19:03 UTC
823d271c0dc7dce43f6a9716be3c99aaf82ced37
2 files changed
+14
-7
tests/kernels/attention/test_cpu_attn.py
+6
-3
@@ -258,10 +258,13 @@ def varlen_with_paged_kv(
258
259
# KV cache for CPU attention
260
cache_dtype = torch.uint8 if is_fp8 else dtype
261
- packed_key_cache = torch.empty(
262
- num_blocks, num_kv_heads, block_size, head_size, dtype=cache_dtype
261
+ packed_key_value_cache = torch.empty(
262
+ num_blocks, num_kv_heads, block_size, head_size * 2, dtype=cache_dtype
263
)
264
- packed_value_cache = torch.empty_like(packed_key_cache)
264
+ packed_key_value_cache = packed_key_value_cache.view(
265
+ (num_blocks, num_kv_heads, block_size * 2, -1)
266
+ )
267
+ packed_key_cache, packed_value_cache = packed_key_value_cache.chunk(2, dim=2)
268
269
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
270
dim=0, dtype=torch.int32
vllm/v1/attention/backends/cpu_attn.py
+8
-4
@@ -93,7 +93,7 @@ class CPUAttentionBackend(AttentionBackend):
93
head_size: int,
94
cache_dtype_str: str = "auto",
95
) -> tuple[int, ...]:
96
- return 2, num_blocks, num_kv_heads, block_size, head_size
96
+ return num_blocks, num_kv_heads, block_size, 2 * head_size
97
98
@classmethod
99
def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
@@ -308,7 +308,7 @@ class CPUAttentionBackendImpl(AttentionImpl):
308
key: shape = [num_tokens, num_kv_heads, head_size]
309
value: shape = [num_tokens, num_kv_heads, head_size]
310
kv_cache: shape =
311
- [2, num_blocks, num_kv_heads, block_size, head_size]
311
+ [num_blocks, num_kv_heads, block_size, 2 * head_size]
312
attn_metadata: Metadata for attention.
313
Returns:
314
shape = [num_tokens, num_heads * head_size]
@@ -338,8 +338,12 @@ class CPUAttentionBackendImpl(AttentionImpl):
338
)
339
340
# For decoder and cross-attention, use KV cache, size are
341
- # [num_blocks, num_kv_heads, block_size, head_size]
342
- key_cache, value_cache = kv_cache.unbind(0)
341
+ # [num_blocks, num_kv_heads, block_size, 2 * head_size]
342
+ # Make a view [num_blocks, num_kv_heads, block_size * 2, head_size]
343
+ # Then slice KV at dim 2
344
+ num_blocks, num_kv_heads, block_size, _ = kv_cache.size()
345
+ kv_cache = kv_cache.view((num_blocks, num_kv_heads, block_size * 2, -1))
346
+ key_cache, value_cache = kv_cache.chunk(2, dim=2)
347
348
# key and value may be None in the case of cross attention. They are
349
# calculated once based on the output from the encoder and then cached