汪志鹏
7271116973
Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
86 lines
3.3 KiB
Python
86 lines
3.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
"""Utility helpers for KV cache manipulation."""
|
||
|
||
import torch
|
||
from vllm.logger import init_logger
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
|
||
|
||
|
||
def normalize_layer_kv(
|
||
layer_kv: LayerKV,
|
||
*,
|
||
req_id: str = "",
|
||
layer_idx: int = -1,
|
||
) -> tuple[torch.Tensor, torch.Tensor] | None:
|
||
"""Normalize one layer KV cache to a ``(key_blocks, value_blocks)`` tuple.
|
||
|
||
In vLLM, different attention backends return paged-attention KV blocks
|
||
with different layouts. For example:
|
||
|
||
* **FlashAttention** (``vllm/v1/attention/backends/flash_attn.py``)
|
||
returns shape ``(2, num_blocks, block_size, num_kv_heads, head_size)``
|
||
– the key/value dimension is at **dim 0**.
|
||
* **FlashInfer** (``vllm/v1/attention/backends/flashinfer.py``)
|
||
returns shape ``(num_blocks, 2, block_size, num_kv_heads, head_size)``
|
||
– the key/value dimension is at **dim 1**.
|
||
|
||
This utility handles both layouts (and the tuple case) so that
|
||
downstream code can work with any backend.
|
||
|
||
Supported layouts:
|
||
|
||
* **Stacked tensor** ``[2, num_blocks, block_size, n_heads, head_dim]`` –
|
||
dim-0 selects key / value.
|
||
* **Stacked tensor** ``[num_blocks, 2, block_size, n_heads, head_dim]`` –
|
||
dim-1 selects key / value.
|
||
* **Tuple** ``(key_tensor, value_tensor)`` – returned as-is after
|
||
validation.
|
||
|
||
Args:
|
||
layer_kv: The raw KV cache (tensor or tuple) for the layer.
|
||
req_id: Request ID used only for diagnostic log messages.
|
||
layer_idx: Layer index used only for diagnostic log messages.
|
||
|
||
Returns:
|
||
``(key_blocks, value_blocks)`` if *layer_kv* is valid, ``None``
|
||
otherwise.
|
||
"""
|
||
if isinstance(layer_kv, torch.Tensor):
|
||
if layer_kv.ndim >= 3 and layer_kv.shape[0] == 2:
|
||
key_blocks = layer_kv[0]
|
||
value_blocks = layer_kv[1]
|
||
elif layer_kv.ndim >= 3 and layer_kv.shape[1] == 2:
|
||
key_blocks = layer_kv[:, 0]
|
||
value_blocks = layer_kv[:, 1]
|
||
else:
|
||
logger.warning(
|
||
f"Layer {layer_idx} for request {req_id} has invalid stacked KV shape: "
|
||
f"expected [2, ...] or [..., 2, ...] at dim 0/1, got {tuple(layer_kv.shape)}"
|
||
)
|
||
return None
|
||
elif isinstance(layer_kv, tuple):
|
||
if len(layer_kv) != 2:
|
||
logger.warning(f"Layer {layer_idx} for request {req_id} has KV pair length {len(layer_kv)} (expected 2)")
|
||
return None
|
||
key_blocks, value_blocks = layer_kv
|
||
if not isinstance(key_blocks, torch.Tensor) or not isinstance(value_blocks, torch.Tensor):
|
||
logger.warning(f"Layer {layer_idx} for request {req_id} has non-tensor KV pair entries")
|
||
return None
|
||
else:
|
||
logger.warning(f"Layer {layer_idx} for request {req_id} has unsupported KV type {type(layer_kv).__name__}")
|
||
return None
|
||
|
||
# ensure key/value blocks are at least 2D for block indexing
|
||
if key_blocks.ndim < 2 or value_blocks.ndim < 2:
|
||
logger.warning(
|
||
f"Layer {layer_idx} for request {req_id} has invalid KV block shape: "
|
||
f"got key={tuple(key_blocks.shape)} value={tuple(value_blocks.shape)}"
|
||
)
|
||
return None
|
||
|
||
return key_blocks, value_blocks
|