1"""Per-request inference caches: (i) backbone per-layer KV, (ii) source-bank K/V projections. Isolated per2RequestBatch; workspace states are never stored across tokens (v1 contract)."""3from __future__ import annotations45from dataclasses import dataclass, field6from typing import Optional78import torch910from .attention import LayerKVCache111213@dataclass14class SourceBankCache:15 K: Optional[torch.Tensor] = None # [B, t, w]16 V: Optional[torch.Tensor] = None1718 @property19 def length(self) -> int:20 return 0 if self.K is None else self.K.shape[1]2122 def append(self, K: torch.Tensor, V: torch.Tensor):23 if self.K is None:24 self.K, self.V = K, V25 else:26 self.K = torch.cat([self.K, K], 1)27 self.V = torch.cat([self.V, V], 1)2829 def reset(self):30 self.K = self.V = None3132 def nbytes(self) -> int:33 return 0 if self.K is None else (self.K.numel() + self.V.numel()) * self.K.element_size()343536@dataclass37class RequestBatchCache:38 """One cache object per batch of concurrent requests (same length, no padding in v1)."""39 layers: list[LayerKVCache]40 bank: SourceBankCache = field(default_factory=SourceBankCache)4142 @property43 def length(self) -> int:44 return self.layers[0].length4546 def reset(self):47 for l in self.layers:48 l.reset()49 self.bank.reset()5051 def nbytes(self) -> int:52 n = self.bank.nbytes()53 for l in self.layers:54 if l.k is not None:55 n += (l.k.numel() + l.v.numel()) * l.k.element_size()56 return n57