from dataclasses import dataclass, asdict import hashlib import math import numpy as np import torch from torch import nn from torch.nn import functional as F @dataclass(frozen=True) class Config: arm: str = 'mha_gated_ffn_router' vocab: int = 16384 width: int = 256 layers: int = 7 hidden: int = 896 q_heads: int = 8 kv_heads: int = 8 head_dim: int = 32 eps: float = 1e-6 seed: int = 20260907 backend: str = 'sdpa' dropout: float = 0. value_variance: bool = True variance_delivery: str = 'full_width_direct' def validate(self): assert self.arm == 'mha_gated_ffn_router' assert self.backend in ('sdpa','flash') assert self.width==self.q_heads*self.head_dim and self.q_heads==self.kv_heads assert self.head_dim%2==0 and self.layers>0 assert 0<=self.dropout<1 assert self.variance_delivery=='full_width_direct' def sphere(x,eps=1e-6): dtype=x.dtype y=x.double() if dtype==torch.float64 else x.float() return (y*torch.rsqrt(y.square().mean(-1,keepdim=True)+eps)).to(dtype) @dataclass class Layout: mask: torch.Tensor | None positions: torch.Tensor indices: torch.Tensor cu: torch.Tensor max_length: int def prepare_layout(segments,device,backend='sdpa'): """CPU metadata boundary: positive contiguous document IDs, padding0.""" if segments.device.type!='cpu':raise ValueError('prepare metadata on CPU before transfer') if segments.ndim!=2 or bool((segments<0).any()):raise ValueError('bad segment IDs') b,t=segments.shape # Vectorized host metadata; no per-token Python loop or device synchronization. raw=segments.numpy() valid=raw>0 starts=valid.copy();starts[:,1:] &= raw[:,1:]!=raw[:,:-1] cols=np.arange(t)[None,:] origins=np.maximum.accumulate(np.where(starts,cols,0),axis=1) positions=torch.from_numpy(np.where(valid,cols-origins,0).astype(np.int64)) indices=np.flatnonzero(valid.reshape(-1)) start_offsets=np.flatnonzero(starts.reshape(-1)[indices]) lengths=np.diff(np.append(start_offsets,len(indices))) mask=None if backend=='sdpa': causal=torch.arange(t)[None,:]<=torch.arange(t)[:,None] mask=((segments[:,:,None]==segments[:,None,:]) & (segments[:,:,None]>0) & causal).unsqueeze(1).to(device) cu=torch.from_numpy(np.concatenate(([0],np.cumsum(lengths))).astype(np.int32)) return Layout(mask,positions.to(device),torch.tensor(indices,dtype=torch.long,device=device),cu.to(device),int(max(lengths,default=0))) class Norm(nn.Module): def __init__(self,d,eps): super().__init__();self.weight=nn.Parameter(torch.ones(d));self.eps=eps def forward(self,x):return sphere(x,self.eps)*self.weight.to(x.dtype) class Router(nn.Module): def __init__(self,c): super().__init__();self.eps=c.eps self.query=nn.Linear(c.width,16,bias=False) self.key=nn.Linear(c.width,16,bias=False) def forward(self,current,s1,s2): q=self.query(sphere(current,self.eps)).float() k1,k2=self.key(s1).float(),self.key(s2).float() scores=torch.stack(((q*k1).sum(-1),(q*k2).sum(-1)),dim=-1)/4. return scores.softmax(-1) class DeltaReader(nn.Module): def __init__(self,c,out_width): super().__init__() self.aux1=nn.Linear(c.width,out_width,bias=False) self.aux2=nn.Linear(c.width,out_width,bias=False) self.router=Router(c) def forward(self,current,s1,s2): v1,v2=self.aux1(s1),self.aux2(s2) p=self.router(current,s1,s2).to(v1.dtype) return p[...,0:1]*v1+p[...,1:2]*v2 def spread_features(variance,eps=1e-6): """Separate direction and smooth magnitude; zero spread maps exactly to zero.""" work=variance.double() if variance.dtype==torch.float64 else variance.float() r=torch.log1p(work/0.1) mean_square=r.square().mean(-1,keepdim=True) energy=mean_square+eps inverse=torch.rsqrt(energy) direction=r*inverse # Rationalized difference avoids cancellation and is exactly zero at r=0. magnitude=torch.log1p(mean_square/((energy.sqrt()+math.sqrt(eps))*(1+math.sqrt(eps)))) return torch.cat((direction,magnitude),-1) class FFN(nn.Module): def __init__(self,c,enabled=True): super().__init__();self.c=c self.base=nn.Linear(c.width,2*c.hidden,bias=False) self.down=nn.Linear(c.hidden,c.width,bias=False) self.reader=DeltaReader(c,2*c.hidden) if enabled else None self.spread_input=nn.Linear(c.width+1,2*c.hidden,bias=False) if c.value_variance else None def forward(self,x,current,s1,s2,variance=None): pre=self.base(x) if self.reader is not None:pre=pre+self.reader(current,s1,s2) if self.spread_input is not None: pre=pre+self.spread_input(spread_features(variance,self.c.eps).to(x.dtype)) g,u=pre.chunk(2,-1) return self.down(F.dropout(F.silu(g)*u,self.c.dropout,self.training)) class AttentionGate(nn.Module): """Head-local elementwise write gate, kept as a compilable fusion boundary.""" def __init__(self,c): super().__init__();self.c=c self.proj=nn.Linear(c.width,c.q_heads*c.head_dim,bias=False) def forward(self,out,xq): b,t,h,d=out.shape gate=torch.sigmoid(self.proj(xq).view(b,t,h,d)) return out*gate class Block(nn.Module): def __init__(self,c,index): super().__init__();self.c=c self.an=Norm(c.width,c.eps);self.fn=Norm(c.width,c.eps) self.qk=nn.Linear(c.width,(c.q_heads+c.kv_heads)*c.head_dim,bias=False) self.v=nn.Linear(c.width,c.kv_heads*c.head_dim,bias=False) # Head-specific elementwise gate from the same pre-normalized query state. # Applied after SDPA and before heads are concatenated/output-projected. self.attn_gate=AttentionGate(c) self.out=nn.Linear(c.width,c.width,bias=False) self.ffn=FFN(c,index>0) def project_attention(self,h,cos,sin): c=self.c;b,t,d=h.shape xq=xv=self.an(h) q,k=self.qk(xq).split((c.q_heads*c.head_dim,c.kv_heads*c.head_dim),-1) q=q.view(b,t,c.q_heads,c.head_dim);k=k.view(b,t,c.kv_heads,c.head_dim) v=self.v(xv).view(b,t,c.kv_heads,c.head_dim) def rope(x): a,z=x.chunk(2,-1) co,si=cos.to(x.dtype),sin.to(x.dtype) return torch.cat((a*co-z*si,z*co+a*si),-1) q,k=rope(q),rope(k) if c.value_variance: q,k=F.pad(q,(0,c.head_dim)),F.pad(k,(0,c.head_dim)) # Keep double for the high precision CPU reference. vv=v.double() if v.dtype==torch.float64 else v.float() v=torch.cat((v,vv.square().to(v.dtype)),-1) return q,k,v,xq def finish_attention(self,h,paired,xq): c=self.c;b,t,d=h.shape if c.value_variance: out,second=paired.split(c.head_dim,-1) mean=out.double() if out.dtype==torch.float64 else out.float() moment=second.double() if second.dtype==torch.float64 else second.float() variance=(moment-mean.square()).clamp_min(0).reshape(b,t,c.width) else:out=paired;variance=None out=self.attn_gate(out,xq).reshape(b,t,d) a=sphere(h+self.out(out),c.eps) return a,self.fn(a),variance def forward(self,h,history,layout,cos,sin): c=self.c;b,t,d=h.shape q,k,v,xq=self.project_attention(h,cos,sin) attention_width=c.head_dim*(2 if c.value_variance else 1) if c.backend=='flash': from flash_attn import flash_attn_varlen_func idx=layout.indices qp=q.reshape(-1,c.q_heads,attention_width)[idx];kp=k.reshape(-1,c.kv_heads,attention_width)[idx];vp=v.reshape(-1,c.kv_heads,attention_width)[idx] y=flash_attn_varlen_func(qp,kp,vp,layout.cu,layout.cu,layout.max_length,layout.max_length,softmax_scale=c.head_dim**-.5,dropout_p=0.,causal=True,deterministic=True) # Scatter the paired moments once; calculate variance after unpacking. paired=torch.zeros_like(q.reshape(-1,c.q_heads,attention_width)).index_copy(0,idx,y).view(b,t,c.q_heads,attention_width) else: paired=F.scaled_dot_product_attention(q.transpose(1,2),k.transpose(1,2),v.transpose(1,2),attn_mask=layout.mask,dropout_p=0.,enable_gqa=False,scale=c.head_dim**-.5).transpose(1,2) a,normalized,variance=self.finish_attention(h,paired,xq) if variance is not None: # A causal document's first token has exactly one Value: variance is # identically zero. Do not normalize BF16 moment cancellation noise. variance=torch.where((layout.positions>0).unsqueeze(-1),variance,0.) if history is None:s1=s2=torch.zeros_like(a) else:s1,s2=history y=sphere(a+self.ffn(normalized,a,s1,s2,variance),c.eps) return y,a class DenseLM(nn.Module): def __init__(self,c): super().__init__();c.validate();self.cfg=c self.embed=nn.Embedding(c.vocab,c.width) self.blocks=nn.ModuleList([Block(c,i) for i in range(c.layers)]) self.norm=Norm(c.width,c.eps);self.lm_head=nn.Linear(c.width,c.vocab,bias=False) self.register_buffer('inv_freq',10000.**(-torch.arange(0,c.head_dim,2).float()/c.head_dim),persistent=True) for name,p in self.named_parameters(): if p.is_meta:continue if name.endswith('spread_input.weight'):nn.init.zeros_(p);continue if p.ndim==1:nn.init.ones_(p);continue seed=(int.from_bytes(hashlib.sha256(name.encode()).digest()[:8],'little')+c.seed)%(2**63-1) gen=torch.Generator(device=p.device).manual_seed(seed) scale=.02/math.sqrt(2*c.layers) if name.endswith(('down.weight','out.weight')) else .02 nn.init.normal_(p,std=scale,generator=gen) def hidden(self,ids,layout): h=sphere(self.embed(ids),self.cfg.eps) # Keep residual states and their differences in FP32 to avoid cancellation # before RMS normalization; autocast still handles matrix multiplications. history=None angles=layout.positions[...,None].float()*self.inv_freq cos,sin=angles.cos().to(h.dtype).unsqueeze(2),angles.sin().to(h.dtype).unsqueeze(2) for i,block in enumerate(self.blocks): before=h;h,attention_after=block(h,history,layout,cos,sin) if i+1=2],'weight_decay':.1},{'params':[p for p in model.parameters() if p.ndim<2],'weight_decay':0.}],lr=8e-4,betas=(.9,.95),eps=1e-8)