| """Standalone PSSA model for the Hub: numpy-only inference, bf16 safetensors.""" |
| import math |
| from dataclasses import dataclass |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
|
|
| try: |
| from pssa_kernels import ScanFn |
| except Exception: |
| ScanFn = None |
|
|
|
|
| def _scan(x, delta, bp, cp, A, h0): |
| """h_t = exp(delta_t路A) h_{t-1} + delta_t路b_t路x_t, y_t = <h_t, c_t>.""" |
| if ScanFn is not None and x.is_cuda: |
| return ScanFn.apply(x, delta, bp, cp, A, h0) |
| B, T, D = x.shape |
| h = h0 |
| ys = [] |
| for t in range(T): |
| h = torch.exp(delta[:, t, :, None] * A) * h + delta[:, t, :, None] * bp[:, t, None, :] * x[:, t, :, None] |
| ys.append((h * cp[:, t, None, :]).sum(-1)) |
| return torch.stack(ys, 1), h |
|
|
|
|
| @dataclass |
| class PSSAConfig: |
| vocab_size: int = 16384 |
| d_model: int = 768 |
| depth: int = 18 |
| d_state: int = 16 |
| d_key: int = 32 |
| mem_slots: int = 512 |
| adapter_rank: int = 16 |
| mlp_hidden: int = 1536 |
| chunk: int = 256 |
| seq_len: int = 512 |
| tau_mem: float = 0.1 |
| value_cap: float = 8.0 |
|
|
|
|
| def _xavier(m): |
| if isinstance(m, nn.Linear): |
| bound = math.sqrt(6 / (m.in_features + m.out_features)) |
| nn.init.uniform_(m.weight, -bound, bound) |
| if m.bias is not None: |
| nn.init.zeros_(m.bias) |
|
|
|
|
| def _poincare(x): |
| """q -> q / (1 + |q|): maps R^n into the open unit ball.""" |
| return x / (1.0 + x.norm(dim=-1, keepdim=True)) |
|
|
|
|
| class PSSABlock(nn.Module): |
| """SSM scan -> hyperbolic memory read -> gate -> plastic adapter -> MLP.""" |
|
|
| def __init__(self, cfg): |
| super().__init__() |
| d, ds, dk, r = cfg.d_model, cfg.d_state, cfg.d_key, cfg.adapter_rank |
| self.cfg = cfg |
| self.norm = nn.Parameter(torch.ones(d)) |
| self.norm_b = nn.Parameter(torch.zeros(d)) |
| self.w_delta = nn.Linear(d, d, bias=False) |
| self.w_b = nn.Linear(d, ds, bias=False) |
| self.w_c = nn.Linear(d, ds, bias=False) |
| ratio = torch.arange(ds).float() / (ds - 1) |
| tau = torch.exp(math.log(1.5) + ratio * (math.log(200.0) - math.log(1.5))) |
| self.a_raw = nn.Parameter(torch.log(torch.expm1(1.0 / tau)).unsqueeze(0).repeat(d, 1)) |
| self.w_qx = nn.Linear(d, dk, bias=False) |
| self.w_qh = nn.Linear(d, dk, bias=False) |
| self.w_gate = nn.Linear(d, d, bias=False) |
| self.w_proj = nn.Linear(d, d, bias=False) |
| self.ad_down = nn.Linear(d, r, bias=False) |
| self.ad_up = nn.Linear(r, d, bias=False) |
| self.mlp1 = nn.Linear(d, cfg.mlp_hidden, bias=False) |
| self.mlp2 = nn.Linear(cfg.mlp_hidden, d, bias=False) |
| self.mk_raw = nn.Parameter(torch.randn(cfg.mem_slots, dk) * 0.05) |
| self.mv = nn.Parameter(torch.randn(cfg.mem_slots, d) * 0.02) |
| for m in (self.w_delta, self.w_b, self.w_c, self.w_qx, self.w_qh, |
| self.w_gate, self.ad_down, self.mlp1): |
| _xavier(m) |
| nn.init.zeros_(self.w_proj.weight) |
| nn.init.zeros_(self.ad_up.weight) |
| nn.init.zeros_(self.mlp2.weight) |
|
|
| def forward(self, x, h0): |
| cfg = self.cfg |
| xn = self.norm * (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-5)) + self.norm_b |
| delta = F.softplus(self.w_delta(xn).float()) |
| bp = self.w_b(xn).float() |
| cp = self.w_c(xn).float() |
| A = -F.softplus(self.a_raw) |
| y, h = _scan(xn, delta, bp, cp, A, h0) |
|
|
| q = _poincare((self.w_qx(xn) + self.w_qh(y)).float()) |
| mk = _poincare(self.mk_raw) |
| qsq = (q * q).sum(-1, keepdim=True) |
| mksq = (mk * mk).sum(-1) |
| sq = (qsq + mksq.view(1, 1, -1) - 2 * (q @ mk.T)).clamp_min(0) |
| den = ((1 - qsq) * (1 - mksq).view(1, 1, -1)).clamp_min(1e-12) |
| dist = 2 * torch.asinh(torch.sqrt(sq / den)) |
| w = torch.softmax(-dist / cfg.tau_mem, -1) |
| mv = cfg.value_cap * self.mv / (cfg.value_cap + self.mv.norm(dim=-1, keepdim=True)) |
| m = w @ mv |
|
|
| g = torch.sigmoid(self.w_gate(xn)) |
| z = y / math.sqrt(cfg.d_state) + g * self.w_proj(m.to(x.dtype)) |
| z = z + self.ad_up(F.silu(self.ad_down(xn))) |
| z = z + self.mlp2(F.silu(self.mlp1(z))) |
| return z, h |
|
|
|
|
| class PSSAModel(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| self.cfg = cfg |
| self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model) |
| self.blocks = nn.ModuleList([PSSABlock(cfg) for _ in range(cfg.depth)]) |
| self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) |
| self.scale = 1.0 / math.sqrt(cfg.depth) |
|
|
| def init_state(self, batch_size, device=None): |
| dev = device or self.embed.weight.device |
| return [torch.zeros(batch_size, self.cfg.d_model, self.cfg.d_state, device=dev) |
| for _ in range(self.cfg.depth)] |
|
|
| def forward_step(self, tokens, states): |
| x = self.embed(tokens) |
| for i, blk in enumerate(self.blocks): |
| z, h = blk(x, states[i]) |
| x = z if i == 0 else x + z * self.scale |
| states[i] = h |
| return self.head(x) / math.sqrt(self.cfg.d_model), states |
|
|
| def forward(self, tokens): |
| states = self.init_state(tokens.shape[0], tokens.device) |
| outs = [] |
| for s in range(0, tokens.shape[1], self.cfg.chunk): |
| logits, states = self.forward_step(tokens[:, s:s + self.cfg.chunk], states) |
| outs.append(logits) |
| return torch.cat(outs, 1) |
|
|
|
|
| def from_pretrained(model_dir: str, device: str = "cpu"): |
| """Load the Hub export: (model, tokenizer). Works on CPU with no Triton install.""" |
| from safetensors.torch import load_file |
| from tokenizers import Tokenizer |
|
|
| cfg = PSSAConfig() |
| sd = load_file(f"{model_dir}/model.safetensors") |
| sd = {k: v.float() for k, v in sd.items()} |
| m = PSSAModel(cfg) |
| m.load_state_dict(sd) |
| m.eval().to(device) |
| return m, Tokenizer.from_file(f"{model_dir}/tokenizer.json") |
|
|
|
|
| @torch.no_grad() |
| def generate(model, tokenizer, prompt, max_new=100, temp=0.7, top_k=40, eot=0, device="cpu"): |
| ids = tokenizer.encode(prompt).ids |
| states = model.init_state(1, device) |
| logits, states = model.forward_step(torch.tensor([ids], device=device), states) |
| out = [] |
| for _ in range(max_new): |
| lg = logits[:, -1, :].float() / temp |
| if top_k: |
| kth = torch.topk(lg, top_k).values[:, [-1]] |
| lg = lg.masked_fill(lg < kth, -float("inf")) |
| nxt = torch.multinomial(torch.softmax(lg, -1), 1) |
| if int(nxt) == eot: |
| break |
| out.append(int(nxt)) |
| logits, states = model.forward_step(nxt, states) |
| text = tokenizer.decode([i for i in out if i != eot]) |
| return prompt + text |
|
|
