File size: 1,681 Bytes
d656a30 | 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 | """The swappable attention contract and the registry (docs/ARCHITECTURE.md).
Every variant (mha, gqa, mla, gdn) subclasses `Attention`, registers itself under
a name, and is picked per layer through `ModelConfig.layer_types`. The rest of
the model never needs to know which one it got.
"""
from torch import Tensor, nn
from tinyllm.config import ModelConfig
ATTENTION: dict[str, type["Attention"]] = {}
def register(name: str):
"""Class decorator: @register("mha") adds the class to ATTENTION."""
def deco(cls: type["Attention"]) -> type["Attention"]:
if name in ATTENTION:
raise ValueError(f"attention variant {name!r} already registered")
ATTENTION[name] = cls
return cls
return deco
class Attention(nn.Module):
def __init__(self, cfg: ModelConfig, layer_idx: int):
super().__init__()
self.cfg = cfg
self.layer_idx = layer_idx
def forward(self, x: Tensor, rope: tuple[Tensor, Tensor], cache=None) -> Tensor:
"""x: (B, T, d_model) -> (B, T, d_model).
rope: (cos, sin) for the T positions in x.
cache: per-layer state for incremental decoding. None during training.
"""
raise NotImplementedError
def cache_bytes_per_token(self, dtype_bytes: int = 2) -> int:
"""Decode-time state this layer adds per generated token (bf16 = 2 bytes)."""
raise NotImplementedError
def build_attention(cfg: ModelConfig, layer_idx: int) -> Attention:
name = cfg.layer_types[layer_idx]
if name not in ATTENTION:
raise KeyError(f"unknown attention {name!r}; registered: {sorted(ATTENTION)}")
return ATTENTION[name](cfg, layer_idx)
|