File size: 2,069 Bytes
ecbc3d2 | 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 | """Incremental caches for Inkling generation.
Two kinds of per-layer state must persist across decode steps:
* ``KVCache`` — the appended key/value tensors for attention.
* ``ConvCache`` — the last ``kernel-1`` inputs of each depthwise short-convolution
(there are 4 per layer: k, v, post-attn, post-mlp).
A ``LayerCache`` bundles one KVCache + the 4 ConvCaches; ``make_cache`` builds one
per decoder layer. Absolute key positions are always ``arange(kv_len)`` because the
cache holds every key from position 0 (KVCache keeps the full history — correct for
both global and sliding layers, since the sliding-window constraint is enforced by
the attention mask).
"""
from __future__ import annotations
import mlx.core as mx
class ConvCache:
"""Holds the last ``kernel-1`` inputs of a short convolution."""
__slots__ = ("state",)
def __init__(self):
self.state = None # [B, kernel-1, C] or None
class KVCache:
"""Appends keys/values along the sequence axis (full history)."""
__slots__ = ("keys", "values")
def __init__(self):
self.keys = None # [B, heads, T, d]
self.values = None
@property
def offset(self) -> int:
return 0 if self.keys is None else self.keys.shape[2]
def update(self, k: mx.array, v: mx.array):
if self.keys is None:
self.keys, self.values = k, v
else:
self.keys = mx.concatenate([self.keys, k], axis=2)
self.values = mx.concatenate([self.values, v], axis=2)
return self.keys, self.values
class LayerCache:
__slots__ = ("kv", "k_conv", "v_conv", "attn_conv", "mlp_conv")
def __init__(self):
self.kv = KVCache()
self.k_conv = ConvCache()
self.v_conv = ConvCache()
self.attn_conv = ConvCache()
self.mlp_conv = ConvCache()
def make_cache(model) -> list[LayerCache]:
"""One LayerCache per text decoder layer."""
n = len(model.model.llm.layers) if hasattr(model, "model") else len(model.layers)
return [LayerCache() for _ in range(n)]
|