A high-throughput and memory-efficient inference and serving engine for LLMs
[Attention] Extract KV-cache update from CPU attention backend (#40470)
Signed-off-by: Diego Maniloff <diego.maniloff@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Diego Maniloff committed
Jun 8, 2026 at 11:43 UTC
ba94a3b9989666f950e1f784d18f2033c63c6cad
1 file changed
+31
-22
vllm/v1/attention/backends/cpu_attn.py
+31
-22
@@ -42,6 +42,8 @@ _CPU_ARCH_PREFER_MIXED_BATCH = (
42
43
44
class CPUAttentionBackend(AttentionBackend):
45
+ forward_includes_kv_cache_update: bool = False
46
+
47
supported_dtypes: ClassVar[list[torch.dtype]] = [
48
torch.float16,
49
torch.bfloat16,
@@ -106,7 +108,6 @@ class CPUAttentionBackend(AttentionBackend):
108
109
@dataclass
110
class CPUAttentionMetadata:
109
- isa: str
111
num_actual_tokens: int # Number of tokens excluding padding.
112
max_query_len: int
113
query_start_loc: torch.Tensor
@@ -217,7 +218,6 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
218
)
219
220
attn_metadata = CPUAttentionMetadata(
220
- isa=self.isa,
221
num_actual_tokens=num_actual_tokens,
222
max_query_len=max_query_len,
223
query_start_loc=query_start_loc,
@@ -345,26 +345,6 @@ class CPUAttentionBackendImpl(AttentionImpl):
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
350
- # in KV cache.
351
- if (
352
- self.kv_sharing_target_layer_name is None
353
- and key is not None
354
- and value is not None
355
- ):
356
- ops.cpu_attn_reshape_and_cache(
357
- key,
358
- value,
359
- key_cache,
360
- value_cache,
361
- attn_metadata.slot_mapping,
362
- attn_metadata.isa,
363
- k_scale=layer._k_scale_float,
364
- v_scale=layer._v_scale_float,
365
- kv_cache_dtype=self.kv_cache_dtype,
366
- )
367
-
348
if attn_metadata.use_sdpa_prefill:
349
assert self.sinks is None, "Attention sink is unsupported in SDPA prefill"
350
num_decode_tokens = attn_metadata.num_decode_tokens
@@ -401,6 +381,35 @@ class CPUAttentionBackendImpl(AttentionImpl):
381
382
return output
383
384
+ def do_kv_cache_update(
385
+ self,
386
+ layer: torch.nn.Module,
387
+ key: torch.Tensor,
388
+ value: torch.Tensor,
389
+ kv_cache: torch.Tensor,
390
+ slot_mapping: torch.Tensor,
391
+ ) -> None:
392
+ if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
393
+ return
394
+
395
+ num_blocks, num_kv_heads, block_size, _ = kv_cache.size()
396
+ kv_cache = kv_cache.view((num_blocks, num_kv_heads, block_size * 2, -1))
397
+ key_cache, value_cache = kv_cache.chunk(2, dim=2)
398
+ isa = _get_attn_isa(
399
+ key.dtype, key_cache.shape[2], self.head_size, self.kv_cache_dtype
400
+ )
401
+ ops.cpu_attn_reshape_and_cache(
402
+ key,
403
+ value,
404
+ key_cache,
405
+ value_cache,
406
+ slot_mapping,
407
+ isa,
408
+ k_scale=layer._k_scale_float,
409
+ v_scale=layer._v_scale_float,
410
+ kv_cache_dtype=self.kv_cache_dtype,
411
+ )
412
+
413
def _run_sdpa_forward(
414
self,
415
query: torch.Tensor,