"""Stateless Phase 2a native drafter. All windows use absolute positions. Forward semantics follow the sha-pinned official model.py read by V1. Inputs are committed target taps or explicitly marked engine main_x fallback rows. The stages never detach/reuse a conditioning cache across optimizer steps. """ from __future__ import annotations from dataclasses import dataclass import hashlib import json from typing import Any, Mapping, Protocol, cast import torch from torch import Tensor, nn from torch.nn import functional as F from torch.utils.checkpoint import checkpoint from ...native.formats import decode_fp8_block from ...native.official import OfficialShards from .ops import Dense, PackedLinear, collapse, expand, frozen_linear, grouped_linear, mixes, qdq, rms, rope class Source(Protocol): def get(self, name: str) -> Tensor: ... def info(self, name: str) -> dict[str, object]: ... def verify(self, name: str) -> None: ... @dataclass(frozen=True) class Config: dim: int = 5120 inter: int = 2304 stages: int = 3 vocab: int = 129280 block: int = 5 noise: int = 128799 hc: int = 4 heads: int = 64 hd: int = 512 rd: int = 64 qr: int = 1280 groups: int = 8 outr: int = 1024 experts: int = 128 topk: int = 3 markov: int = 256 window: int = 128 eps: float = 1e-20 hc_eps: float = 1e-6 sinkhorn: int = 20 theta: float = 10000.0 route_scale: float = 1.5 gate_temp: float = 1.0 limit: float = 10.0 def key(name: str) -> str: return name.replace(".", "__") class Tensors(nn.Module): def __init__(self, source: Source, device: torch.device, reference_math: bool = False): super().__init__() self.source, self.device = source, device self.reference_math = reference_math def read(self, name: str) -> Tensor: self.source.verify(name) return self.source.get(name).to(self.device) def add(self, name: str, train: bool = False, bf16: bool = False) -> Tensor: t = self.read(name).to(torch.bfloat16 if bf16 else torch.float32) if self.reference_math and not bf16: t = t.cpu() if not bool(torch.isfinite(t).all()): raise ValueError(f"nonfinite {name}") if train: self.register_parameter(key(name), nn.Parameter(t)) else: self.register_buffer(key(name), t) return t def __getitem__(self, name: str) -> Tensor: return cast(Tensor, getattr(self, key(name))) def dense(self, name: str, rank: int, acc: torch.dtype, quant_input: bool = True) -> Dense: w = self.read(name + ".weight") if self.source.info(name + ".weight")["dtype"] == "F8_E4M3": w = decode_fp8_block(w, self.read(name + ".scale")) return Dense(w, rank, quant_input, acc) class Stage(nn.Module): wo_a_base: Tensor def __init__(self, c: Config, index: int, source: Source, device: torch.device, rank: int, acc: torch.dtype, reference_math: bool = False): super().__init__() self.c, self.index, self.acc = c, index, acc self.prefix = p = f"mtp.{index}." self.fast_training = False self.t = t = Tensors(source, device, reference_math) self.dense = nn.ModuleDict() for name in ("attn.wq_a", "attn.wq_b", "attn.wkv", "attn.wo_b", "ffn.shared_experts.w1", "ffn.shared_experts.w2", "ffn.shared_experts.w3"): self.dense[key(name)] = t.dense(p + name, rank, acc) # wo_a is grouped, so each output group has its own low-rank update. full = t.dense(p + "attn.wo_a", 0, acc, False).weight self.register_buffer("wo_a_base", (full.cpu() if reference_math else full).reshape(c.groups, c.outr, -1)) self.wo_a = nn.ModuleList(Dense(w, rank, False, acc) for w in full.chunk(c.groups, 0)) for name in ("attn_norm", "ffn_norm", "attn.q_norm", "attn.kv_norm"): t.add(p + name + ".weight", train=True) for name in ("hc_attn_fn", "hc_attn_scale", "hc_attn_base", "hc_ffn_fn", "hc_ffn_scale", "hc_ffn_base", "attn.attn_sink", "ffn.gate.weight", "ffn.gate.bias"): t.add(p + name) self.experts = nn.ModuleList() for e in range(c.experts): ex = nn.ModuleDict() for name in ("w1", "w2", "w3"): ep = f"{p}ffn.experts.{e}.{name}" ex[name] = PackedLinear(t.read(ep + ".weight"), t.read(ep + ".scale"), acc) self.experts.append(ex) def weight(self, name: str) -> Tensor: return self.t[self.prefix + name] def norm(self, x: Tensor, name: str) -> Tensor: return rms(x, self.weight(name + ".weight"), self.c.eps) def linear(self, name: str, x: Tensor) -> Tensor: return cast(Tensor, self.dense[key(name)](x)) def kv(self, x: Tensor, positions: Tensor) -> Tensor: z = self.norm(self.linear("attn.wkv", x), "attn.kv_norm") return qdq(rope(z, positions, self.c.rd, self.c.theta), self.fast_training) def attention(self, x: Tensor, mx: Tensor, positions: Tensor, mask: Tensor, P: Tensor ) -> tuple[Tensor, Tensor]: c = self.c kv_ctx = self.kv(mx, positions) # Match the reference's ring-slot order (important for summation bands). slots = positions.remainder(c.window).argsort(-1) ctx = kv_ctx.gather(1, slots.unsqueeze(-1).expand_as(kv_ctx)) valid_ctx = mask.gather(1, slots) query_pos = P[:, None] + 1 + torch.arange(c.block, device=x.device) q = self.linear("attn.wq_b", self.norm(self.linear("attn.wq_a", x), "attn.q_norm")) q = rope(q.reshape(*q.shape[:-1], c.heads, c.hd), query_pos, c.rd, c.theta) kv = torch.cat((ctx, self.kv(x, query_pos)), 1) valid = torch.cat((valid_ctx, torch.ones(x.shape[:2], device=x.device, dtype=torch.bool)), 1) scores = torch.einsum("bshd,btd->bsht", q.to(self.acc), kv.to(self.acc)) * c.hd ** -.5 scores = scores.masked_fill(~valid[:, None, None, :], -torch.inf) sink = self.weight("attn.attn_sink").to(self.acc)[None, None, :, None].expand(*scores.shape[:-1], 1) prob = torch.cat((scores, sink), -1).softmax(-1)[..., :-1] o = torch.einsum("bsht,btd->bshd", prob, kv.to(self.acc)).to(x.dtype) o = rope(o, query_pos, c.rd, c.theta, inverse=True).reshape(x.shape[0], c.block, c.groups, -1) o = qdq(o, self.fast_training) projected = grouped_linear(o, self.wo_a_base) updates = [F.linear(F.linear(o[:, :, g].to(cast(Dense, layer).a.device), cast(Dense, layer).a), cast(Dense, layer).b).to(o.device) for g, layer in enumerate(self.wo_a) if cast(Dense, layer).rank] if updates: projected = projected + torch.stack(updates, -2) return self.linear("attn.wo_b", projected.flatten(-2)), kv_ctx def expert(self, x: Tensor, e: nn.ModuleDict, w: Tensor | None = None) -> Tensor: gate, up = e["w1"](x).float(), e["w3"](x).float() h = F.silu(gate.clamp(max=self.c.limit)) * up.clamp(-self.c.limit, self.c.limit) if w is not None: h = h * w return cast(Tensor, e["w2"](h.to(x.dtype))) def moe(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: c = self.c shape = x.shape x = x.reshape(-1, c.dim) # The official router uses F.linear FP32 even in the FP64 oracle mode. if self.t.reference_math: raw = torch.cat([F.linear(x[b:b+c.block].float(), self.weight("ffn.gate.weight")) for b in range(0, x.shape[0], c.block)], 0) else: raw = F.linear(x.float(), self.weight("ffn.gate.weight")) score = F.softplus(raw / c.gate_temp).sqrt() biased = score + self.weight("ffn.gate.bias") idx = biased.topk(c.topk, -1).indices w = score.gather(1, idx) w = w / (w.sum(-1, keepdim=True) + 1e-20) * c.route_scale y = torch.zeros_like(x, dtype=torch.float32) if self.fast_training: y = self.fast_moe(x, idx, w) else: for eid in sorted(set(idx.detach().flatten().tolist())): rows, col = torch.where(idx == eid) # Per selected expert only; no full-bank dequantization in either pass. out = self.expert(x[rows], self.experts[eid], w[rows, col, None]) y = y.index_add(0, rows, out.float()) shared = nn.ModuleDict({n: self.dense[key("ffn.shared_experts." + n)] for n in ("w1", "w2", "w3")}) y = y + self.expert(x, shared) return y.reshape(shape).to(x.dtype), idx.reshape(*shape[:-1], c.topk), biased def fast_moe(self, x: Tensor, idx: Tensor, w: Tensor) -> Tensor: from .fast_gpu import ExpertBank bank = cast(ExpertBank, self.expert_bank) order = idx.flatten().argsort(stable=True) counts = torch.bincount(idx.flatten(), minlength=self.c.experts) offsets = F.pad(counts.cumsum(0), (1, 0)) tiles = ((counts + 15) // 16).cumsum(0) rows = order // self.c.topk xq = qdq(x, True)[rows] gate = bank.product("w1", xq, offsets, tiles) up = bank.product("w3", xq, offsets, tiles) h = F.silu(gate.clamp(max=self.c.limit)) * up.clamp(-self.c.limit, self.c.limit) h = qdq(h * w.flatten()[order, None], True) out = bank.product("w2", h, offsets, tiles) # Ascending expert order, as in the reference. CUDA index_add accumulates # only top-k terms; agreement is measured, never presumed bit-exact. return torch.zeros_like(x).index_add(0, rows, out) def hc_mix(self, x: Tensor, kind: str) -> tuple[Tensor, Tensor, Tensor]: c = self.c if self.t.reference_math and x.shape[0] > 1: parts = [self.hc_mix(x[b:b+1], kind) for b in range(x.shape[0])] return (torch.cat([p[0] for p in parts], 0), torch.cat([p[1] for p in parts], 0), torch.cat([p[2] for p in parts], 0)) return mixes(x, self.weight(f"hc_{kind}_fn"), self.weight(f"hc_{kind}_scale"), self.weight(f"hc_{kind}_base"), c.hc, c.sinkhorn, c.eps, c.hc_eps, self.acc) def forward(self, h: Tensor, pre: Tensor, mx: Tensor, positions: Tensor, mask: Tensor, P: Tensor, trace: dict[str, object] | None = None) -> tuple[Tensor, Tensor]: original_h, original_pre = h, pre if trace is None and bool((P == 0).all()): # Official DSparkBlock(start_pos=0) seeds context only; no block # transformation occurs. Stateless callers need no mutable seed. return h, pre apre, apost, acomb = self.hc_mix(h, "attn") aout, kv = self.attention(self.norm(collapse(h, pre), "attn_norm"), mx, positions, mask, P) h = expand(aout, h, apost, acomb) fpre, fpost, fcomb = self.hc_mix(h, "ffn") fin = self.norm(collapse(h, apre), "ffn_norm") fout, idx, biased = self.moe(fin) h = expand(fout, h, fpost, fcomb) h = torch.where((P == 0)[:, None, None, None], original_h, h) fpre = torch.where((P == 0)[:, None, None], original_pre, fpre) if trace is not None: trace.update(attn_out=aout.detach(), ffn_in=fin.detach(), ffn_out=fout.detach(), h_out=h.detach(), kv=kv.detach(), route={"idx": idx.detach(), "biased": biased.detach()}) return h, fpre class Drafter(nn.Module): def __init__(self, source: Source, config: Config = Config(), device: str = "cpu", rank: int = 32, acc: torch.dtype = torch.float32, checkpoint_stages: bool = True, reference_math: bool = False): super().__init__() torch.backends.cuda.matmul.allow_tf32 = False self.c, self.acc, self.checkpoint_stages = config, acc, checkpoint_stages self.reference_math = reference_math self.fast_training = False self.activation_device = "cpu" if reference_math else device identity = source.manifest if isinstance(source, OfficialShards) else getattr(source, "fingerprint", None) if identity is None: raise ValueError("tensor source must provide a stable frozen-weight identity") self.frozen_identity = hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest() self.t = t = Tensors(source, torch.device(device), reference_math) self.main_proj = t.dense("mtp.0.main_proj", 0, acc) t.add("mtp.0.main_norm.weight", train=True) t.add("embed.weight", bf16=True) t.add("head.weight", bf16=True) last = f"mtp.{config.stages-1}." for name in ("norm.weight", "markov_head.embed.weight", "markov_head.head.weight", "confidence_head.proj.weight"): t.add(last + name, train=True) self.stages = nn.ModuleList(Stage(config, i, source, t.device, rank, acc, reference_math) for i in range(config.stages)) def enable_fast_training(self) -> None: if self.reference_math or self.t.device.type != "cuda": raise ValueError("fast training requires a separate all-GPU model") if self.fast_training: return from .fast_gpu import ExpertBank torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False for module_stage in self.stages: stage = cast(Stage, module_stage) stage.expert_bank = ExpertBank(stage.experts) stage.fast_training = True for module in self.modules(): if isinstance(module, Dense): module.fast_training = True self.fast_training = True def last(self, name: str) -> Tensor: return self.t[f"mtp.{self.c.stages-1}." + name] def conditioning(self, taps: Tensor, engine_mx: Tensor, has_taps: Tensor, main_projection: Tensor | None = None) -> Tensor: if main_projection is not None: raise ValueError("projection cache refused: detached output does not prove frozen weights AND frozen inputs; use raw taps") # FP64 quantized GEMM sums retain the same per-row QDQ inputs when batched. flat = taps.reshape(-1, taps.shape[-1]) selected = has_taps.flatten().nonzero().flatten() mx = engine_mx.float().reshape(-1, self.c.dim) if selected.numel(): projected = self.main_proj(flat[selected].float()) native = rms(projected, self.t["mtp.0.main_norm.weight"], self.c.eps) mx = mx.index_copy(0, selected, native) return mx.reshape(*taps.shape[:-1], self.c.dim) def hidden(self, taps: Tensor, engine_mx: Tensor, has_taps: Tensor, positions: Tensor, mask: Tensor, anchor: Tensor, P: Tensor, trace: dict[str, Any] | None = None, main_projection: Tensor | None = None) -> tuple[Tensor, Tensor]: c = self.c mx = self.conditioning(taps, engine_mx, has_taps, main_projection) ids = torch.full((anchor.numel(), c.block), c.noise, dtype=torch.long, device=anchor.device) ids[:, 0] = anchor h = F.embedding(ids.to(self.t["embed.weight"].device), self.t["embed.weight"]).to(self.activation_device).float().unsqueeze(-2).repeat(1, 1, c.hc, 1) pre = torch.zeros(*h.shape[:-1], device=h.device, dtype=torch.float32) pre[..., 0] = 1 for i, stage in enumerate(self.stages): if self.checkpoint_stages and self.training and torch.is_grad_enabled() and trace is None: h, pre = checkpoint(stage, h, pre, mx, positions, mask, P, use_reentrant=False) else: st: dict[str, object] = {} h, pre = stage(h, pre, mx, positions, mask, P, st if trace is not None else None) if trace is not None: trace[f"s{i}"] = st head_x = collapse(h, pre) hidden = rms(head_x, self.last("norm.weight"), c.eps) if trace is not None: trace.update(head_x=head_x.detach(), head_hidden=hidden.detach(), main_x=mx.detach()) return hidden, head_x def hidden_from(self, inp: Mapping[str, Tensor], trace: dict[str, Any] | None = None) -> tuple[Tensor, Tensor]: return self.hidden(inp["taps"], inp["engine_mx"], inp["has_taps"], inp["positions"], inp["mask"], inp["anchor"], inp["P"], trace, inp.get("main_projection")) def logit_chunk(self, hidden: Tensor, previous: Tensor, start: int, end: int) -> Tensor: if self.fast_training: from .ops import fast_head base = fast_head(hidden.float(), self.t["head.weight"][start:end]) else: base = frozen_linear(hidden.float(), self.t["head.weight"][start:end], acc=torch.float32) emb = F.embedding(previous, self.last("markov_head.embed.weight")) return base + F.linear(emb.float(), self.last("markov_head.head.weight")[start:end]) def confidence(self, head_x: Tensor, previous: Tensor) -> Tensor: emb = F.embedding(previous, self.last("markov_head.embed.weight")) return F.linear(torch.cat((head_x, emb), -1).float(), self.last("confidence_head.proj.weight")).squeeze(-1) @torch.no_grad() def greedy(self, hidden: Tensor, head_x: Tensor, anchor: Tensor, chunk: int = 4096) -> tuple[Tensor, Tensor]: prev = anchor drafts, conf = [], [] for pos in range(self.c.block): best = torch.full_like(prev, -torch.inf, dtype=torch.float32) pick = torch.zeros_like(prev) for lo in range(0, self.c.vocab, chunk): logits = self.logit_chunk(hidden[:, pos], prev, lo, min(lo + chunk, self.c.vocab)) val, idx = logits.max(-1) better = val > best # equal ties retain the lower vocabulary id pick = torch.where(better, idx + lo, pick) best = torch.maximum(best, val) conf.append(self.confidence(head_x[:, pos], prev).sigmoid()) drafts.append(pick) prev = pick return torch.stack(drafts, -1), torch.stack(conf, -1) def load(root: str, device: str = "cpu", rank: int = 32, oracle: bool = False, reference_math: bool = False, native_manifest: str | None = None) -> Drafter: source = OfficialShards(root) if native_manifest is None else OfficialShards(root, native_manifest) return Drafter(source, device=device, rank=rank, acc=torch.float64 if oracle or reference_math else torch.float32, reference_math=reference_math)