107 lines
3.2 KiB
Python
107 lines
3.2 KiB
Python
"""Physical KV cache buffers.
|
|
|
|
Layer 1 — ``KVStorage``: flat token-level K/V GPU buffers [n_layers, size, n_kv_heads, head_dim]
|
|
Layer 2 — ``ReqToTokenPool``: index table [req_idx, pos] → physical token slot
|
|
Layer 3 — ``KVCache``: pure dataclass passed to the model for direct buffer access
|
|
|
|
These classes have no knowledge of tasks, allocation policies, or scheduling.
|
|
They are the "dumb" physical storage layer.
|
|
"""
|
|
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from typing import List, Optional
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
|
|
class ReqToTokenPool:
|
|
"""Maps [req_idx, pos] → physical token slot in KV storage.
|
|
|
|
Each row is one request; each column is a sequence position. The value
|
|
at [req_idx, pos] is the flat index into the KV storage buffers.
|
|
"""
|
|
|
|
def __init__(self, size: int, max_context_len: int, device: torch.device):
|
|
self.size = size
|
|
self.max_context_len = max_context_len
|
|
self.req_to_token = torch.zeros(
|
|
(size, max_context_len), dtype=torch.int32, device=device
|
|
)
|
|
self.free_slots = list(range(size))
|
|
self._lock = threading.Lock()
|
|
|
|
def alloc(self, num_reqs: int) -> Optional[List[int]]:
|
|
with self._lock:
|
|
if num_reqs > len(self.free_slots):
|
|
return None
|
|
slots = self.free_slots[:num_reqs]
|
|
self.free_slots = self.free_slots[num_reqs:]
|
|
return slots
|
|
|
|
def free(self, req_indices: List[int]):
|
|
with self._lock:
|
|
self.free_slots.extend(req_indices)
|
|
|
|
def write(self, indices, values):
|
|
self.req_to_token[indices] = values
|
|
|
|
|
|
class KVStorage:
|
|
"""Token-level KV cache storage.
|
|
|
|
Buffers: ``[n_layers, size, n_kv_heads, head_dim]``. Each token occupies
|
|
one slot indexed by ``ReqToTokenPool``.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
size: int,
|
|
n_layers: int,
|
|
n_kv_heads: int,
|
|
head_dim: int,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
):
|
|
self.size = size
|
|
self.k_buffer = torch.empty(
|
|
(n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
|
|
)
|
|
self.v_buffer = torch.empty(
|
|
(n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
|
|
)
|
|
|
|
def get_key_buffer(self, layer_id: int) -> Tensor:
|
|
return self.k_buffer[layer_id]
|
|
|
|
def get_value_buffer(self, layer_id: int) -> Tensor:
|
|
return self.v_buffer[layer_id]
|
|
|
|
def set_kv_buffer(self, layer_id: int, loc: Tensor, k: Tensor, v: Tensor) -> None:
|
|
self.k_buffer[layer_id, loc] = k
|
|
self.v_buffer[layer_id, loc] = v
|
|
|
|
|
|
@dataclass
|
|
class KVCache:
|
|
"""Pure data struct passed to model for KV cache I/O.
|
|
|
|
The attention layer does raw buffer indexing — no methods, no abstraction.
|
|
"""
|
|
|
|
k_buffer: Tensor
|
|
v_buffer: Tensor
|
|
req_to_token: Tensor
|
|
req_pool_indices: Tensor
|
|
seq_lens: Tensor
|
|
out_cache_loc: Tensor
|
|
max_len: int = 0
|
|
kv_indptr: Optional[Tensor] = None
|
|
qo_indptr: Optional[Tensor] = None
|
|
q_tile_to_batch: Optional[Tensor] = None
|
|
q_tile_to_index: Optional[Tensor] = None
|
|
decode_o_part: Optional[Tensor] = None
|
|
decode_ml_part: Optional[Tensor] = None
|
|
decode_out: Optional[Tensor] = None
|