1"""Causal decoder backbone: pre-norm blocks (GQA attention + SwiGLU), tied embeddings, single final norm."""2from __future__ import annotations34from typing import Optional56import torch7import torch.nn as nn89from .attention import CausalSelfAttention, LayerKVCache, RMSNorm, SwiGLU10from .config import PrismeConfig111213class DecoderBlock(nn.Module):14 def __init__(self, cfg: PrismeConfig):15 super().__init__()16 self.attn_norm = RMSNorm(cfg.d_model, cfg.norm_eps)17 self.attn = CausalSelfAttention(cfg.d_model, cfg.n_heads, cfg.n_kv_heads, cfg.rope_theta)18 self.ffn_norm = RMSNorm(cfg.d_model, cfg.norm_eps)19 self.ffn = SwiGLU(cfg.d_model, cfg.ffn_width)2021 def forward(self, x, positions, cache: Optional[LayerKVCache] = None, attn_mask=None):22 x = x + self.attn(self.attn_norm(x), positions, cache, attn_mask)23 x = x + self.ffn(self.ffn_norm(x))24 return x252627class Backbone(nn.Module):28 """Returns the final residual stream H (pre final-norm). The LM head is `head()`."""2930 def __init__(self, cfg: PrismeConfig):31 super().__init__()32 self.cfg = cfg33 self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)34 self.layers = nn.ModuleList([DecoderBlock(cfg) for _ in range(cfg.n_layers)])35 self.final_norm = RMSNorm(cfg.d_model, cfg.norm_eps)36 if not cfg.tie_embeddings:37 self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)38 self.apply(self._init)39 # scaled init for residual output projections (GPT-2 style)40 for n, p in self.named_parameters():41 if n.endswith("wo.weight") or n.endswith("w_down.weight"):42 nn.init.normal_(p, mean=0.0, std=cfg.init_std / (2 * cfg.n_layers) ** 0.5)4344 def _init(self, m):45 if isinstance(m, nn.Linear):46 nn.init.normal_(m.weight, mean=0.0, std=self.cfg.init_std)47 elif isinstance(m, nn.Embedding):48 nn.init.normal_(m.weight, mean=0.0, std=self.cfg.init_std)4950 def new_cache(self) -> list[LayerKVCache]:51 return [LayerKVCache() for _ in self.layers]5253 def forward(self, tokens: torch.Tensor, positions: Optional[torch.Tensor] = None,54 cache: Optional[list[LayerKVCache]] = None, attn_mask=None) -> torch.Tensor:55 B, T = tokens.shape56 if positions is None:57 start = 0 if cache is None else cache[0].length58 positions = torch.arange(start, start + T, device=tokens.device).unsqueeze(0).expand(B, T)59 x = self.embed(tokens)60 for i, layer in enumerate(self.layers):61 x = layer(x, positions, None if cache is None else cache[i], attn_mask)62 return x # [B, T, d], pre final-norm residual stream6364 def head(self, h: torch.Tensor) -> torch.Tensor:65 """logits = E @ RMSNorm(h) (tied) or lm_head(RMSNorm(h))."""66 hn = self.final_norm(h)67 if self.cfg.tie_embeddings:68 return hn @ self.embed.weight.t()69 return self.lm_head(hn)70