File size: 5,454 Bytes
d011ff0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | """Request-local inference cache; overflow preserves crop-and-recompute semantics."""
import torch
from torch.nn import functional as F
from .rotary import RotaryAttention, rotate_pairs
class CacheSession:
def __init__(self, model):
if model.training:
raise ValueError("KV caching requires model.eval(); training stays uncached")
self.rotary = model.config.position_encoding == "rope"
if not self.rotary and not hasattr(model, "position_embedding"):
raise ValueError("A learned-position model requires its position table")
for block in model.blocks:
attention = block.attention
if self.rotary != isinstance(attention, RotaryAttention):
raise ValueError("Cache position encoding does not match attention modules")
kv_heads = getattr(attention, "n_kv_head", attention.n_head)
width = model.config.n_embd + 2 * kv_heads * (model.config.n_embd // attention.n_head)
if kv_heads < 1 or attention.n_head % kv_heads or attention.qkv.out_features != width:
raise ValueError("Cache projection/head shapes do not match the model")
self.model = model
self.layers = None
self.length = 0
self.batch_size = None
@torch.no_grad()
def _forward(self, token_ids, past, offset):
model = self.model
if model.training:
raise ValueError("The cached model must remain in evaluation mode")
batch, time = token_ids.shape
positions = torch.arange(offset, offset + time, device=token_ids.device)
x = model.token_embedding(token_ids)
if not self.rotary:
x = x + model.position_embedding(positions)[None]
x = model.embedding_dropout(x)
layers = []
for index, block in enumerate(model.blocks):
attention = block.attention
dim = model.config.n_embd // attention.n_head
kv_heads = getattr(attention, "n_kv_head", attention.n_head)
q, k, v = attention.qkv(block.ln_attention(x)).split(
(model.config.n_embd, kv_heads * dim, kv_heads * dim), dim=-1)
q = q.view(batch, time, attention.n_head, dim).transpose(1, 2)
k, v = (value.view(batch, time, kv_heads, dim).transpose(1, 2) for value in (k, v))
if self.rotary:
cos, sin = attention.cos[offset:offset + time], attention.sin[offset:offset + time]
q, k = rotate_pairs(q, cos, sin), rotate_pairs(k, cos, sin)
if past is not None:
k = torch.cat((past[index][0], k), dim=-2)
v = torch.cat((past[index][1], v), dim=-2)
else:
# Do not retain Q's storage through views of the combined projection.
k = k.clone(memory_format=torch.contiguous_format)
v = v.clone(memory_format=torch.contiguous_format)
layers.append((k, v))
attention_k, attention_v = k, v
if kv_heads != attention.n_head:
repeats = attention.n_head // kv_heads
attention_k = k.repeat_interleave(repeats, dim=1)
attention_v = v.repeat_interleave(repeats, dim=1)
# A single decode query has no future keys. A rectangular causal
# mask aligned at the upper left would wrongly hide its cached past.
attended = F.scaled_dot_product_attention(
q, attention_k, attention_v, dropout_p=0.0, is_causal=past is None)
attended = attended.transpose(1, 2).contiguous().view(batch, time, model.config.n_embd)
x = x + attention.output_dropout(attention.output(attended))
x = x + block.ffn(block.ln_ffn(x))
return model.lm_head(model.final_norm(x)[:, -1:, :]), layers
def prefill(self, token_ids):
if token_ids.ndim != 2 or not 0 < token_ids.shape[1] <= self.model.config.block_size:
raise ValueError("Prefill requires a nonempty (B,T) prompt within the context window")
self.layers, self.length, self.batch_size = None, 0, None
logits, layers = self._forward(token_ids, None, 0)
self.layers, self.length, self.batch_size = layers, token_ids.shape[1], token_ids.shape[0]
return logits
def decode(self, token_ids):
if self.layers is None:
raise ValueError("Call prefill before decode")
if token_ids.shape != (self.batch_size, 1):
raise ValueError("Decode accepts exactly one new token per existing batch row")
if self.length >= self.model.config.block_size:
raise ValueError("Full cache: explicitly prefill the cropped window instead")
logits, layers = self._forward(token_ids, self.layers, self.length)
self.layers = layers
self.length += 1
return logits
def storage(self):
tensors = [] if self.layers is None else [value for pair in self.layers for value in pair]
unique = {value.untyped_storage().data_ptr(): value.untyped_storage().nbytes() for value in tensors}
return {"positions": self.length,
"logical_bytes": sum(value.numel() * value.element_size() for value in tensors),
"storage_bytes": sum(unique.values()),
"dtype": str(tensors[0].dtype) if tensors else None,
"allocation": "dynamic concatenation; no preallocated unused capacity"}
|