1"""Building blocks: RMSNorm, RoPE, grouped-query causal attention with optional KV cache, SwiGLU."""2from __future__ import annotations34import math5from typing import Optional67import torch8import torch.nn as nn9import torch.nn.functional as F101112class RMSNorm(nn.Module):13 def __init__(self, dim: int, eps: float = 1e-5):14 super().__init__()15 self.eps = eps16 self.weight = nn.Parameter(torch.ones(dim))1718 def forward(self, x: torch.Tensor) -> torch.Tensor:19 dtype = x.dtype20 xf = x.float()21 out = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)22 return (out * self.weight.float()).to(dtype)232425def rope_cos_sin(positions: torch.Tensor, head_dim: int, theta: float, device, dtype=torch.float32):26 """positions: [..., T] integer tensor -> cos, sin of shape [..., T, head_dim]."""27 inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))28 freqs = positions.to(torch.float32).unsqueeze(-1) * inv_freq # [..., T, hd/2]29 emb = torch.cat([freqs, freqs], dim=-1)30 return emb.cos().to(dtype), emb.sin().to(dtype)313233def rotate_half(x: torch.Tensor) -> torch.Tensor:34 x1, x2 = x.chunk(2, dim=-1)35 return torch.cat([-x2, x1], dim=-1)363738def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:39 """x: [B, H, T, hd]; cos/sin: [B or 1, T, hd] -> broadcast over heads."""40 cos = cos.unsqueeze(1)41 sin = sin.unsqueeze(1)42 return x * cos + rotate_half(x) * sin434445class SwiGLU(nn.Module):46 def __init__(self, dim: int, hidden: int):47 super().__init__()48 self.w_gate = nn.Linear(dim, hidden, bias=False)49 self.w_up = nn.Linear(dim, hidden, bias=False)50 self.w_down = nn.Linear(hidden, dim, bias=False)5152 def forward(self, x: torch.Tensor) -> torch.Tensor:53 return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))545556class LayerKVCache:57 """Per-layer KV cache. Tensors are [B, Hkv, T, hd]. One instance per request batch."""5859 def __init__(self):60 self.k: Optional[torch.Tensor] = None61 self.v: Optional[torch.Tensor] = None6263 @property64 def length(self) -> int:65 return 0 if self.k is None else self.k.shape[2]6667 def append(self, k: torch.Tensor, v: torch.Tensor):68 if self.k is None:69 self.k, self.v = k, v70 else:71 self.k = torch.cat([self.k, k], dim=2)72 self.v = torch.cat([self.v, v], dim=2)73 return self.k, self.v7475 def reset(self):76 self.k = None77 self.v = None787980class CausalSelfAttention(nn.Module):81 """GQA causal attention with RoPE. Supports full-sequence (training) and cached incremental decoding."""8283 def __init__(self, d_model: int, n_heads: int, n_kv_heads: int, rope_theta: float):84 super().__init__()85 self.n_heads = n_heads86 self.n_kv_heads = n_kv_heads87 self.head_dim = d_model // n_heads88 self.rope_theta = rope_theta89 self.wq = nn.Linear(d_model, n_heads * self.head_dim, bias=False)90 self.wk = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)91 self.wv = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)92 self.wo = nn.Linear(n_heads * self.head_dim, d_model, bias=False)9394 def forward(self, x: torch.Tensor, positions: torch.Tensor, cache: Optional[LayerKVCache] = None,95 attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor:96 """x: [B, T, d]; positions: [B, T] absolute positions (for RoPE and causal masking).9798 If cache is given, k/v are appended and attention runs over the full cached length.99 attn_mask: optional boolean [B, 1, T, T_total] (True = attend). If None and cache is None,100 a standard causal mask is used.101 """102 B, T, _ = x.shape103 q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)104 k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)105 v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)106 cos, sin = rope_cos_sin(positions, self.head_dim, self.rope_theta, x.device, x.dtype)107 q = apply_rope(q, cos, sin)108 k = apply_rope(k, cos, sin)109 if cache is not None:110 k, v = cache.append(k, v)111 rep = self.n_heads // self.n_kv_heads112 if rep > 1:113 k = k.repeat_interleave(rep, dim=1)114 v = v.repeat_interleave(rep, dim=1)115 if attn_mask is None and cache is None:116 out = F.scaled_dot_product_attention(q, k, v, is_causal=True)117 else:118 if attn_mask is None:119 # incremental decoding: query positions attend to all cached keys at or before them.120 T_total = k.shape[2]121 key_pos = torch.arange(T_total, device=x.device)122 # cached keys are stored in insertion order == absolute position order123 attn_mask = (key_pos[None, None, None, :] <= positions[:, None, :, None])124 out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)125 out = out.transpose(1, 2).reshape(B, T, self.n_heads * self.head_dim)126 return self.wo(out)127