| """Straightforward reference of the same drafter for parity checks (D3a). Slow, fp32, one request at a time. |
| |
| Independent of ops.py / model.py on purpose: written from DeepSeek's reference inference/model.py (DSparkBlock, |
| Block.hc_mixes / hc_pre / hc_post, Gate, Expert, DSparkAttention, forward_head) with the running engine's choices |
| where they differ (NOTES.md): the SM12x V4 KV record, wo_a without activation quantization, EXL3 weights. Every EXL3 |
| matrix is decoded once to its exact rotated-basis fp16 weight (exllamav3 reconstruct, bit-identical to the CPU |
| decoder, D1) and applied in fp32 as y = H(H(x * suh) @ W) * svh; experts are decoded on demand. No exllamav3 GEMM, |
| coop MoE, Triton or batched kernel of the fast path is used here. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import math |
| import os |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| from .checkpoint import Checkpoint |
|
|
|
|
| def _hadamard128(dev: torch.device) -> torch.Tensor: |
| h = torch.ones((1, 1), dtype=torch.float64) |
| while h.shape[0] < 128: |
| h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)], 0) |
| return (h / math.sqrt(128)).to(dev, torch.float32) |
|
|
|
|
| class RefLinear: |
| def __init__(self, ck: Checkpoint, name: str, dev: torch.device, H: torch.Tensor, keep: bool = True): |
| import exllamav3_ext as ext |
|
|
| tr, suh, svh, k = ck.exl3(name) |
| self.name, self.k, self.dev, self.H = name, k, dev, H |
| self.in_f, self.out_f = tr.shape[0] * 16, tr.shape[1] * 16 |
| self.suh = suh.to(dev).float() |
| self.svh = svh.to(dev).float() |
| self._tr = tr |
| self.w = None |
| if keep: |
| self.w = self.decode() |
| self._ext = ext |
|
|
| def decode(self) -> torch.Tensor: |
| import exllamav3_ext as ext |
| w = torch.empty(self.in_f, self.out_f, dtype=torch.float16, device=self.dev) |
| ext.reconstruct(w, self._tr.to(self.dev), self.k, False, True) |
| return w |
|
|
| def _had(self, x: torch.Tensor) -> torch.Tensor: |
| m, n = x.shape |
| return (x.view(m, n // 128, 128) @ self.H).view(m, n) |
|
|
| def __call__(self, x: torch.Tensor) -> torch.Tensor: |
| """fp32 [m, in] -> fp32 [m, out] on the exact decoded weight.""" |
| w = self.w if self.w is not None else self.decode() |
| y = self._had(x.float() * self.suh) @ w.float() |
| return self._had(y) * self.svh |
|
|
|
|
| class RefDrafter: |
| def __init__(self, ck: Checkpoint, dev: torch.device, kv_format: str = "v4"): |
| s = ck.shape |
| self.s, self.dev, self.ck, self.kv_format = s, dev, ck, kv_format |
| H = _hadamard128(dev) |
| self.H = H |
| g = lambda n: ck.get(n).to(dev) |
| self.embed = g("embed.weight") |
| self.head = RefLinear(ck, "head", dev, H) |
| self.main_proj = RefLinear(ck, "mtp.0.main_proj", dev, H) |
| self.main_norm = g("mtp.0.main_norm.weight") |
| last = f"mtp.{s.n_stages - 1}." |
| self.norm = g(last + "norm.weight") |
| self.markov_embed = g(last + "markov_head.embed.weight") |
| self.markov_head = g(last + "markov_head.head.weight") |
| self.conf_w = g(last + "confidence_head.proj.weight").float() |
| self.stages = [] |
| for i in range(s.n_stages): |
| p = f"mtp.{i}." |
| wo_a = ck.get(p + "attn.wo_a.weight").to(dev).float() |
| sc = ck.get(p + "attn.wo_a.scale").to(dev).float() |
| wo_a = wo_a * sc.repeat_interleave(32, 0).repeat_interleave(32, 1) |
| self.stages.append({ |
| "wq_a": RefLinear(ck, p + "attn.wq_a", dev, H), "wkv": RefLinear(ck, p + "attn.wkv", dev, H), |
| "wq_b": RefLinear(ck, p + "attn.wq_b", dev, H), "wo_b": RefLinear(ck, p + "attn.wo_b", dev, H), |
| "wo_a": wo_a.view(s.o_groups, s.o_lora, -1), |
| "attn_sink": g(p + "attn.attn_sink").float(), "q_norm": g(p + "attn.q_norm.weight"), |
| "kv_norm": g(p + "attn.kv_norm.weight"), "attn_norm": g(p + "attn_norm.weight"), |
| "ffn_norm": g(p + "ffn_norm.weight"), |
| "hc_attn_fn": g(p + "hc_attn_fn"), "hc_attn_scale": g(p + "hc_attn_scale"), |
| "hc_attn_base": g(p + "hc_attn_base"), "hc_ffn_fn": g(p + "hc_ffn_fn"), |
| "hc_ffn_scale": g(p + "hc_ffn_scale"), "hc_ffn_base": g(p + "hc_ffn_base"), |
| "gate_w": g(p + "ffn.gate.weight").float(), "gate_bias": g(p + "ffn.gate.bias").float(), |
| "sh": [RefLinear(ck, p + f"ffn.shared_experts.{w}", dev, H) for w in ("w1", "w3", "w2")], |
| "prefix": p, |
| }) |
| |
| freqs = 1.0 / (s.rope_theta ** (torch.arange(0, s.rope_dim, 2, dtype=torch.float32) / s.rope_dim)) |
| self.freqs = freqs.to(dev) |
|
|
| |
| def norm_(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: |
| dtype = x.dtype |
| x = x.float() |
| var = x.square().mean(-1, keepdim=True) |
| x = x * torch.rsqrt(var + self.s.rms_eps) |
| return (w.float() * x).to(dtype) |
|
|
| def rope_(self, x: torch.Tensor, pos: torch.Tensor, inverse: bool = False) -> torch.Tensor: |
| """x [n, (h,) d]: rotate the last rope_dim channels (adjacent pairs as complex numbers).""" |
| rd = self.s.rope_dim |
| ang = pos.to(torch.float32).unsqueeze(-1) * self.freqs |
| fc = torch.polar(torch.ones_like(ang), ang) |
| if inverse: |
| fc = fc.conj() |
| y = x.clone() |
| z = torch.view_as_complex(x[..., -rd:].float().contiguous().unflatten(-1, (-1, 2))) |
| if z.dim() == 3: |
| fc = fc.unsqueeze(1) |
| y[..., -rd:] = torch.view_as_real(z * fc).flatten(-2).to(x.dtype) |
| return y |
|
|
| def kv_quant_(self, kv: torch.Tensor) -> torch.Tensor: |
| if self.kv_format != "v4": |
| return kv |
| out = kv.clone() |
| for gi in range(7): |
| blk = kv[..., gi * 64:(gi + 1) * 64].float() |
| amax = blk.abs().amax(-1, keepdim=True).clamp_min(448.0 * 2.0 ** -126) |
| sc = 2.0 ** torch.ceil(torch.log2(amax / 448.0)) |
| out[..., gi * 64:(gi + 1) * 64] = ((blk / sc).clamp(-448, 448).to(torch.float8_e4m3fn).float() * sc |
| ).to(kv.dtype) |
| return out |
|
|
| def hc_mixes_(self, x: torch.Tensor, fn: torch.Tensor, scale: torch.Tensor, base: torch.Tensor): |
| hc, eps = self.s.hc_mult, self.s.hc_eps |
| xf = x.flatten(1).float() |
| mixes = F.linear(xf, fn) * torch.rsqrt(xf.square().mean(-1, keepdim=True) + self.s.rms_eps) |
| pre = torch.sigmoid(mixes[:, :hc] * scale[0] + base[:hc]) + eps |
| post = 2 * torch.sigmoid(mixes[:, hc:2 * hc] * scale[1] + base[hc:2 * hc]) |
| comb = mixes[:, 2 * hc:].unflatten(-1, (hc, hc)) * scale[2] + base[2 * hc:].unflatten(-1, (hc, hc)) |
| comb = comb.softmax(-1) + eps |
| comb = comb / (comb.sum(-2, keepdim=True) + eps) |
| for _ in range(self.s.hc_iters - 1): |
| comb = comb / (comb.sum(-1, keepdim=True) + eps) |
| comb = comb / (comb.sum(-2, keepdim=True) + eps) |
| return pre, post, comb |
|
|
| @staticmethod |
| def hc_pre_(x: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor: |
| return torch.sum(pre_mix.unsqueeze(-1) * x.float(), dim=1).to(x.dtype) |
|
|
| @staticmethod |
| def hc_post_(y: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor) -> torch.Tensor: |
| out = post.unsqueeze(-1) * y.float().unsqueeze(-2) + torch.sum( |
| comb.unsqueeze(-1) * residual.float().unsqueeze(-2), dim=1) |
| return out.to(y.dtype) |
|
|
| |
| def main_x(self, taps: torch.Tensor) -> torch.Tensor: |
| return self.norm_(self.main_proj(taps).to(torch.bfloat16), self.main_norm) |
|
|
| def context_kv(self, mx: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: |
| out = [] |
| for st in self.stages: |
| kv = self.norm_(st["wkv"](mx).to(torch.bfloat16), st["kv_norm"]) |
| out.append(self.kv_quant_(self.rope_(kv, pos))) |
| return torch.stack(out) |
|
|
| |
| def attention_(self, st: dict, x: torch.Tensor, pos: torch.Tensor, ctx: torch.Tensor) -> torch.Tensor: |
| s = self.s |
| k = x.shape[0] |
| qr = self.norm_(st["wq_a"](x).to(torch.bfloat16), st["q_norm"]) |
| q = st["wq_b"](qr).to(torch.bfloat16).unflatten(-1, (s.n_heads, s.head_dim)) |
| q = self.rope_(q, pos) |
| kv = self.norm_(st["wkv"](x).to(torch.bfloat16), st["kv_norm"]) |
| kv = self.kv_quant_(self.rope_(kv, pos)) |
| keys = torch.cat([ctx, kv], 0).float() |
| o = torch.empty(k, s.n_heads, s.head_dim, dtype=torch.float32, device=x.device) |
| for t in range(k): |
| sc = (q[t].float() @ keys.t()) * (s.head_dim ** -0.5) |
| sc = torch.cat([sc, st["attn_sink"].unsqueeze(1)], 1) |
| p = sc.softmax(-1)[:, :-1] |
| o[t] = p @ keys |
| o = self.rope_(o.to(torch.bfloat16), pos, inverse=True) |
| o = o.view(k, s.o_groups, -1) |
| z = torch.einsum("sgd,grd->sgr", o.float(), st["wo_a"]).to(torch.bfloat16) |
| return st["wo_b"](z.flatten(1)).to(torch.bfloat16) |
|
|
| def moe_(self, st: dict, x: torch.Tensor) -> torch.Tensor: |
| s = self.s |
| scores = F.softplus(F.linear(x.float(), st["gate_w"])).sqrt() |
| idx = (scores + st["gate_bias"]).topk(s.top_k, dim=-1)[1] |
| w = scores.gather(1, idx) |
| w = w / (w.sum(-1, keepdim=True) + 1e-20) * s.route_scale |
| y = torch.zeros(x.shape, dtype=torch.float32, device=x.device) |
| lim = s.swiglu_limit |
| for t in range(x.shape[0]): |
| for j in range(s.top_k): |
| e = int(idx[t, j]) |
| p = st["prefix"] + f"ffn.experts.{e}." |
| w1, w3, w2 = (RefLinear(self.ck, p + n, self.dev, self.H) for n in ("w1", "w3", "w2")) |
| gt = w1(x[t:t + 1]).clamp(max=lim) |
| up = w3(x[t:t + 1]).clamp(-lim, lim) |
| a = F.silu(gt) * up * w[t, j] |
| y[t] += w2(a.to(x.dtype))[0] |
| w1, w3, w2 = st["sh"] |
| a = F.silu(w1(x).clamp(max=lim)) * w3(x).clamp(-lim, lim) |
| y += w2(a.to(x.dtype)) |
| return y.to(x.dtype), idx |
|
|
| def block(self, anchor: int, last_pos: int, ctx: list[torch.Tensor], k: int, trace: dict | None = None): |
| """One request's query block. ctx: per stage [n_ctx, 512] bf16 context KV (the valid window).""" |
| s = self.s |
| ids = torch.full((k,), s.noise_token, dtype=torch.int64, device=self.dev) |
| ids[0] = anchor |
| pos = torch.arange(last_pos + 1, last_pos + 1 + k, device=self.dev) |
| x = self.embed[ids].unsqueeze(1).repeat(1, s.hc_mult, 1) |
| pre_mix = torch.zeros(k, s.hc_mult, dtype=torch.float32, device=self.dev) |
| pre_mix[:, 0] = 1.0 |
| for i, st in enumerate(self.stages): |
| residual = x |
| a_pre, a_post, a_comb = self.hc_mixes_(x, st["hc_attn_fn"], st["hc_attn_scale"], st["hc_attn_base"]) |
| h = self.norm_(self.hc_pre_(x, pre_mix), st["attn_norm"]) |
| h = self.attention_(st, h, pos, ctx[i]) |
| x = self.hc_post_(h, residual, a_post, a_comb) |
| residual = x |
| f_pre, f_post, f_comb = self.hc_mixes_(x, st["hc_ffn_fn"], st["hc_ffn_scale"], st["hc_ffn_base"]) |
| h = self.norm_(self.hc_pre_(x, a_pre), st["ffn_norm"]) |
| y, idx = self.moe_(st, h) |
| x = self.hc_post_(y, residual, f_post, f_comb) |
| pre_mix = f_pre |
| if trace is not None: |
| trace[f"stage{i}.ffn_in"] = h |
| trace[f"stage{i}.ffn_out"] = y |
| trace[f"stage{i}.stream"] = x |
| trace[f"stage{i}.experts"] = idx |
| return self.hc_pre_(x, pre_mix) |
|
|
| def heads(self, hh: torch.Tensor, anchor: int, k: int, trace: dict | None = None): |
| """Greedy Markov drafting on bf16 logits + bias (the engine's dtypes); fp32 values kept for diagnostics.""" |
| base32 = self.head(self.norm_(hh, self.norm)) |
| prev = anchor |
| ids, embs, logits32 = [], [], [] |
| for i in range(k): |
| e = self.markov_embed[prev] |
| bias32 = self.markov_head.float() @ e.float() |
| l16 = base32[i].to(torch.bfloat16) + bias32.to(torch.bfloat16) |
| d = int(l16.argmax()) |
| ids.append(d) |
| embs.append(e) |
| logits32.append(base32[i] + bias32) |
| prev = d |
| feat = torch.cat([hh, torch.stack(embs)], -1).float() |
| conf = torch.sigmoid(feat @ self.conf_w.view(-1)) |
| if trace is not None: |
| trace["logits32"] = torch.stack(logits32) |
| trace["base32"] = base32 |
| return ids, conf |
|
|
