| """The DSpark drafter of the served checkpoint on one RTX 3090 (SM86): weights, per-request window state, draft step. |
| |
| Semantics: NOTES.md (the engine's DSpark round, traced from the running image). Tensors: the exact checkpoint bytes |
| (checkpoint.Checkpoint refuses anything else). Representation changes are draft-only and listed in NOTES.md. |
| |
| Public surface (all device tensors, no host syncs inside, so the draft step and the insert capture into CUDA graphs): |
| * ``main_x(taps)`` [n, 15360] bf16 target taps -> [n, 5120] bf16 conditioning features |
| * ``insert(slot_rows, pos_rows, mx)`` write the context KV of accepted positions into request windows |
| * ``draft(slots, anchors, last_pos, k)`` -> (draft ids [B, k] int64, confidences [B, k] fp32) |
| * ``reset_slots(slots)`` empty request windows |
| Graphs: ``capture(batch_sizes, ks, n_outs)`` records one graph per (B, k, n_out) (and a margin twin); |
| ``draft_graphed`` replays it. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| import time |
| from dataclasses import dataclass, field |
|
|
| import torch |
|
|
| from . import env |
| from . import fused as FU |
| from .checkpoint import Checkpoint, DrafterShape |
| from .exl3 import kernels as K |
| from .exl3.bank import Exl3Bank |
| from .exl3.plan import PROJ_DOWN, PROJ_GATE, PROJ_UP, Exl3MoEPlan |
| from .ops import (MarkovBias, apply_rope, block_attention, hc_collapse, hc_mixes, hc_post, kv_record_qdq, rms_norm, |
| route, rope_cos_sin) |
|
|
| MUL1 = 0x83DCD12D - (1 << 32) |
|
|
|
|
| @dataclass |
| class Window: |
| """What a draft block attends to: the ring rows of ``slots`` whose position lies in [lo, last] (int64 [B] each). |
| ``kv`` [B, n_stages, W, D] / ``valid`` [B, W] are the gathered rows, filled only for the unfused (ops) path; the |
| fused path reads the ring in place (fused.block_attention_ring).""" |
| slots: torch.Tensor |
| lo: torch.Tensor |
| last: torch.Tensor |
| kv: torch.Tensor | None = None |
| valid: torch.Tensor | None = None |
|
|
|
|
| def facc_ext(): |
| """exllamav3_ext_facc: the same exllamav3 1.6.0 source built with fp32 MMA accumulation on sm_86 (stock 1.6.0 |
| accumulates exl3_gemm in fp16 on sm_86 only, EXL3_GEMM_H_ACC; ENV.md). Raises if it is not installed.""" |
| import exllamav3_ext_facc |
| return exllamav3_ext_facc |
|
|
|
|
| class Exl3Linear: |
| """One EXL3 matrix on exllamav3 1.6.0: fp16 input, fp32 output (the plugin's prompt-size path; NOTES.md). |
| |
| ``prescale``: divide the input by the power of two that brings max |x| under 2^14 and multiply the output back |
| (exact; the plugin does this for main_proj, whose input is the target's raw residual stream). |
| ``ext`` / ``force_shape``: call ``ext.exl3_gemm`` directly (``ext`` defaults to the stock exllamav3_ext) with a |
| fixed kernel shape instead of LinearEXL3's timing-based autotune: the kernel, and so each row's result, is then |
| the same for every row count the shape covers (1..16 for the 16-row shapes) and in every process.""" |
|
|
| def __init__(self, trellis: torch.Tensor, suh: torch.Tensor, svh: torch.Tensor, k: int, dev: torch.device, |
| prescale: bool = False, ext=None, force_shape: int = 0): |
| from exllamav3.modules.quant.exl3 import LinearEXL3 |
|
|
| self.in_features, self.out_features, self.k = trellis.shape[0] * 16, trellis.shape[1] * 16, k |
| self.trellis = trellis.to(dev).contiguous() |
| self.suh = suh.to(dev).contiguous() |
| self.svh = svh.to(dev).contiguous() |
| self.mul1 = torch.tensor(MUL1, dtype=torch.int32, device=dev) |
| self.lin = LinearEXL3(None, self.in_features, self.out_features, suh=self.suh, svh=self.svh, |
| trellis=self.trellis, mul1=self.mul1) |
| self.prescale = prescale |
| self.direct = ext is not None or force_shape > 0 |
| if self.direct: |
| if ext is None: |
| import exllamav3_ext as ext |
| self.ext = ext |
| self.force_shape = force_shape |
|
|
| def _gemm(self, x16: torch.Tensor, out: torch.Tensor | None) -> torch.Tensor: |
| if not self.direct: |
| if out is not None: |
| raise ValueError("out= needs the direct exl3_gemm path") |
| return self.lin.forward(x16, {}, out_dtype=torch.float32) |
| if out is None: |
| out = torch.empty(x16.shape[0], self.out_features, dtype=torch.float32, device=x16.device) |
| self.ext.exl3_gemm(x16, self.trellis, out, self.suh, torch.empty_like(x16), self.svh, self.force_shape, False, |
| True, 0) |
| return out |
|
|
| def __call__(self, x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor: |
| """x [m, in] (bf16 / fp32 are converted to fp16; fp16 is used as is) -> fp32 [m, out].""" |
| if self.prescale: |
| sc = torch.exp2(torch.ceil(torch.log2(x.abs().amax().float().clamp_min(1.0) / 16384.0)).clamp_min(0.0)) |
| y = self._gemm((x.float() / sc).to(torch.float16).contiguous(), out) |
| return y.mul_(sc) |
| if x.dtype != torch.float16: |
| x = x.to(torch.float16) |
| return self._gemm(x.contiguous(), out) |
|
|
| def nbytes(self) -> int: |
| return sum(t.numel() * t.element_size() for t in (self.trellis, self.suh, self.svh)) |
|
|
|
|
| @dataclass |
| class Stage: |
| wq_a: Exl3Linear |
| wkv: Exl3Linear |
| wq_b: Exl3Linear |
| wo_b: Exl3Linear |
| wo_a8: torch.Tensor |
| wo_a_sc: torch.Tensor |
| wo_a_t: torch.Tensor | None |
| |
| attn_sink: torch.Tensor |
| q_norm: torch.Tensor |
| kv_norm: torch.Tensor |
| attn_norm: torch.Tensor |
| ffn_norm: torch.Tensor |
| hc_attn: tuple[torch.Tensor, torch.Tensor, torch.Tensor] |
| hc_ffn: tuple[torch.Tensor, torch.Tensor, torch.Tensor] |
| gate_w: torch.Tensor |
| gate_bias: torch.Tensor |
| sh_w1: Exl3Linear |
| sh_w3: Exl3Linear |
| sh_w2: Exl3Linear |
| bank: Exl3Bank |
| groups: list |
| perm: torch.Tensor |
|
|
|
|
| def dequant_wo_a(weight: torch.Tensor, scale: torch.Tensor, groups: int) -> torch.Tensor: |
| """FP8 e4m3 [G*R, C] x UE8M0 block scales [G*R/bs, C/bs] -> bf16 [G, C, R] (exact: e4m3 x 2^k is a bf16 value).""" |
| rows, cols = weight.shape |
| br, bc = rows // scale.shape[0], cols // scale.shape[1] |
| s = scale.float().repeat_interleave(br, 0).repeat_interleave(bc, 1) |
| w = (weight.float() * s).to(torch.bfloat16) |
| if not torch.equal(w.float(), weight.float() * s): |
| raise RuntimeError("wo_a dequantization is not exact in bf16") |
| return w.view(groups, rows // groups, cols).transpose(1, 2).contiguous() |
|
|
|
|
| class DSparkDrafter: |
| """All weights resident on one GPU; per-request windows in ``slots`` fixed slots (slot ``slots`` = scratch).""" |
|
|
| HEAD_SHAPE = 3 |
| HEAD_CHUNK = 16 |
|
|
| def __init__(self, ck: Checkpoint, dev: torch.device, slots: int = 16, max_model_len: int = 327680, |
| max_rows: int = 32, kv_format: str = "v4", fused: bool = True, main_proj_fp32acc: bool = True, |
| markov_cache_rows: int = 8192): |
| if not ck.verified: |
| raise RuntimeError("checkpoint not identity-verified") |
| t0 = time.time() |
| self.dev = dev |
| self.s: DrafterShape = ck.shape |
| s = self.s |
| self.kv_format = kv_format |
| self.fused = fused |
| torch.backends.cuda.matmul.allow_tf32 = False |
| torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False |
| self.max_model_len = max_model_len |
| self.max_pos = max_model_len + 2048 |
| self.max_rows = max_rows |
| self.eps = s.rms_eps |
| self.scale = s.head_dim ** -0.5 |
| k_dense = ck.drafter_k |
| plan = Exl3MoEPlan.from_dict(ck.plan_json) |
|
|
| def lin(name: str, want_k: int | None, prescale: bool = False, ext=None, force_shape: int = 0) -> Exl3Linear: |
| tr, suh, svh, k = ck.exl3(name) |
| if want_k is not None and k != want_k: |
| raise RuntimeError(f"{name}: trellis K{k}, config says K{want_k}") |
| return Exl3Linear(tr, suh, svh, k, dev, prescale, ext, force_shape) |
|
|
| def g(name: str, dtype: torch.dtype | None = None) -> torch.Tensor: |
| t = ck.get(name) |
| return t.to(dev, dtype=dtype or t.dtype).contiguous() |
|
|
| self.embed = g("embed.weight") |
| self.head = lin("head", ck.head_k, force_shape=self.HEAD_SHAPE) |
| |
| |
| self.main_proj_fp32acc = main_proj_fp32acc |
| self.main_proj = lin("mtp.0.main_proj", k_dense["main_proj"], prescale=True, |
| ext=facc_ext() if main_proj_fp32acc else None) |
| self.main_norm = g("mtp.0.main_norm.weight") |
| last = f"mtp.{s.n_stages - 1}." |
| self.norm = g(last + "norm.weight") |
| self.conf_w = g(last + "confidence_head.proj.weight", torch.float32).view(-1) |
| tokens = None |
| tok_path = os.path.join(ck.root, "markov_top_tokens.json") |
| if os.path.exists(tok_path): |
| with open(tok_path) as f: |
| tokens = json.load(f) |
| |
| |
| |
| self.markov = MarkovBias(g(last + "markov_head.embed.weight"), g(last + "markov_head.head.weight"), tokens, |
| rows=markov_cache_rows) |
|
|
| self.stages: list[Stage] = [] |
| for i in range(s.n_stages): |
| p = f"mtp.{i}." |
| L = s.n_target_layers + i |
| lp = plan.drafter_layers[L] |
| if lp.num_experts != s.n_experts: |
| raise RuntimeError(f"plan layer {L}: {lp.num_experts} experts, config {s.n_experts}") |
| bank = Exl3Bank(lp, s.hidden, s.inter, 1, 0) |
| bank.allocate(dev) |
| for e in range(s.n_experts): |
| for proj, wn in ((PROJ_GATE, "w1"), (PROJ_UP, "w3"), (PROJ_DOWN, "w2")): |
| tr, suh, svh, _ = ck.exl3(f"{p}ffn.experts.{e}.{wn}") |
| bank.load_piece("trellis", proj, e, tr) |
| bank.load_piece("suh", proj, e, suh) |
| bank.load_piece("svh", proj, e, svh) |
| groups, _ = bank.tables() |
| st = Stage( |
| wq_a=lin(p + "attn.wq_a", k_dense["attn.wq_a"]), wkv=lin(p + "attn.wkv", k_dense["attn.wq_a"]), |
| wq_b=lin(p + "attn.wq_b", k_dense["attn.wq_b"]), wo_b=lin(p + "attn.wo_b", k_dense["attn.wo_b"]), |
| wo_a8=g(p + "attn.wo_a.weight").view(torch.uint8), wo_a_sc=g(p + "attn.wo_a.scale", torch.float32), |
| wo_a_t=None if fused else dequant_wo_a(ck.get(p + "attn.wo_a.weight").to(dev), |
| ck.get(p + "attn.wo_a.scale").to(dev), s.o_groups), |
| attn_sink=g(p + "attn.attn_sink", torch.float32), 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=(g(p + "hc_attn_fn"), g(p + "hc_attn_scale"), g(p + "hc_attn_base")), |
| hc_ffn=(g(p + "hc_ffn_fn"), g(p + "hc_ffn_scale"), g(p + "hc_ffn_base")), |
| gate_w=g(p + "ffn.gate.weight", torch.float32), gate_bias=g(p + "ffn.gate.bias", torch.float32), |
| sh_w1=lin(p + "ffn.shared_experts.w1", k_dense["ffn.shared_experts.w1"]), |
| sh_w3=lin(p + "ffn.shared_experts.w3", k_dense["ffn.shared_experts.w1"]), |
| sh_w2=lin(p + "ffn.shared_experts.w2", k_dense["ffn.shared_experts.w2"]), |
| bank=bank, groups=groups, perm=torch.tensor(lp.perm, dtype=torch.int64, device=dev), |
| ) |
| self.stages.append(st) |
| from .exl3.hostcopy import release_stage |
| release_stage() |
|
|
| self.coop = K.get_coop_scratch(dev, max_rows, s.top_k, s.hidden, s.inter) |
| K.prime(dev, self.stages[0].groups, self.coop, None, s.hidden, s.top_k, s.swiglu_limit, False, True) |
| self.cos_sin = rope_cos_sin(self.max_pos, s.rope_dim, s.rope_theta).to(dev) |
| |
| |
| self.q_clamp = torch.tensor(max_model_len - 1, dtype=torch.int64, device=dev) |
|
|
| |
| self.n_slots = slots |
| self.ring_kv = torch.zeros(slots + 1, s.n_stages, s.window, s.head_dim, dtype=torch.bfloat16, device=dev) |
| self.ring_pos = torch.full((slots + 1, s.window), -1, dtype=torch.int64, device=dev) |
| self.arange_block = torch.arange(s.block + 1, dtype=torch.int64, device=dev) |
| self.graphs: dict[tuple[int, int, int], dict] = {} |
| self.margin_graphs: dict[tuple[int, int, int], dict] = {} |
| self.reply_graphs: dict[tuple[int, int, int, bool], dict] = {} |
| self.insert_graphs: dict[int, dict] = {} |
| self.pool = None |
| torch.cuda.synchronize(dev) |
| self.load_seconds = time.time() - t0 |
|
|
| def wo_a_bf16(self, i: int) -> torch.Tensor: |
| """Stage i's wo_a as bf16 [groups, 4096, 1024] (exact dequantization; for tools on the fused model).""" |
| st = self.stages[i] |
| if st.wo_a_t is not None: |
| return st.wo_a_t |
| return dequant_wo_a(st.wo_a8.view(torch.float8_e4m3fn), st.wo_a_sc, self.s.o_groups) |
|
|
| |
| def _norm(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: |
| return FU.rms_norm(x, w, self.eps) if self.fused else rms_norm(x.to(torch.bfloat16), w, self.eps) |
|
|
| def _rope(self, x: torch.Tensor, pos: torch.Tensor, inverse: bool = False) -> torch.Tensor: |
| return FU.rope(x, pos, self.cos_sin, inverse) if self.fused else apply_rope(x, pos, self.cos_sin, inverse) |
|
|
| def _kv(self, kv: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: |
| """rope + (V4 record) of normed KV rows.""" |
| if self.kv_format != "v4": |
| return self._rope(kv, pos) |
| if self.fused: |
| return FU.kv_rope_qdq(kv, pos, self.cos_sin) |
| return kv_record_qdq(apply_rope(kv, pos, self.cos_sin)) |
|
|
| def _mixes(self, stream, fn, scale, base): |
| f = FU.hc_mixes if self.fused else hc_mixes |
| return f(stream, fn, scale, base, self.eps, self.s.hc_eps, self.s.hc_iters) |
|
|
| def _collapse_norm(self, stream, pre, w): |
| if self.fused: |
| return FU.collapse_norm(stream, pre, w, self.eps) |
| return rms_norm(hc_collapse(stream, pre), w, self.eps) |
|
|
| def _post(self, y, stream, post, comb): |
| return (FU.hc_post if self.fused else hc_post)(y, stream, post, comb) |
|
|
| def _collapse(self, stream, pre): |
| return FU.collapse(stream, pre) if self.fused else hc_collapse(stream, pre) |
|
|
| |
| def main_x(self, taps: torch.Tensor) -> torch.Tensor: |
| """[n, 3H] bf16 taps of target layers 37/38/39 -> main_x [n, H] bf16.""" |
| if self.fused and isinstance(self.main_proj, Exl3Linear): |
| x16, sc = FU.prescale_fp16(taps) |
| return FU.scaled_rms_norm(self.main_proj._gemm(x16, None), sc, self.main_norm, self.eps) |
| return self._norm(self.main_proj(taps), self.main_norm) |
|
|
| def context_kv(self, mx: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: |
| """Stage context KV [n_stages, n, D] (rope at pos, record-quantized) of main_x rows.""" |
| out = [] |
| for st in self.stages: |
| out.append(self._kv(self._norm(st.wkv(mx), st.kv_norm), pos)) |
| return torch.stack(out) |
|
|
| def insert(self, slot_rows: torch.Tensor, pos_rows: torch.Tensor, mx: torch.Tensor) -> None: |
| """Write the context KV of rows (slot, position, main_x) into the windows. Callers pass at most the last |
| `window` positions of a request per call (older ones would be overwritten anyway; duplicates are refused |
| upstream). Padded rows use slot == n_slots (scratch).""" |
| if self.fused and self.kv_format == "v4": |
| mx16 = mx.to(torch.float16) |
| for i, st in enumerate(self.stages): |
| FU.kv_store(st.wkv(mx16), st.kv_norm, self.eps, pos_rows, self.cos_sin, self.ring_kv[:, i], slot_rows, |
| self.ring_pos if i == 0 else None) |
| return |
| kv = self.context_kv(mx, pos_rows) |
| idx = pos_rows % self.s.window |
| self.ring_kv[slot_rows, :, idx] = kv.transpose(0, 1) |
| self.ring_pos[slot_rows, idx] = pos_rows |
|
|
| def reset_slots(self, slots: torch.Tensor) -> None: |
| self.ring_pos[slots] = -1 |
| self.ring_kv[slots] = 0 |
|
|
| |
| def _attention(self, st: Stage, i: int, x: torch.Tensor, pos: torch.Tensor, win: Window, B: int, k: int |
| ) -> torch.Tensor: |
| s = self.s |
| T = B * k |
| x16 = x.to(torch.float16) |
| qr = self._norm(st.wq_a(x16), st.q_norm) |
| kvb = self._kv(self._norm(st.wkv(x16), st.kv_norm), pos) |
| q = self._rope(st.wq_b(qr).to(torch.bfloat16).view(T, s.n_heads, s.head_dim), pos) |
| qb, kb = q.view(B, k, s.n_heads, s.head_dim), kvb.view(B, k, s.head_dim) |
| if win.kv is None: |
| o = FU.block_attention_ring(qb, self.ring_kv[:, i], self.ring_pos, win.slots, win.lo, win.last, kb, |
| st.attn_sink, self.scale) |
| else: |
| o = block_attention(qb, win.kv[:, i], win.valid, kb, st.attn_sink, self.scale) |
| o = self._rope(o.view(T, s.n_heads, s.head_dim), pos, inverse=True) |
| if st.wo_a_t is None: |
| z = FU.wo_a_fp8(o.view(T, -1), st.wo_a8, st.wo_a_sc, s.o_groups) |
| else: |
| og = o.view(T, s.o_groups, -1).transpose(0, 1) |
| z = torch.bmm(og, st.wo_a_t).transpose(0, 1).reshape(T, -1) |
| return st.wo_b(z).to(torch.bfloat16) |
|
|
| def _moe(self, st: Stage, x: torch.Tensor, trace: dict | None = None, tag: str = "") -> torch.Tensor: |
| s = self.s |
| ids, w = route(x, st.gate_w, st.gate_bias, s.top_k, s.route_scale) |
| if trace is not None: |
| trace[tag + ".experts"] = ids |
| trace[tag + ".route_w"] = w |
| x16 = x.to(torch.float16).contiguous() |
| sel = st.perm[ids].contiguous() |
| routed = K.run_decode_coop(x16, sel, w.to(torch.float16).contiguous(), st.groups, self.coop, |
| s.swiglu_limit, False, True) |
| g = st.sh_w1(x16) |
| u = st.sh_w3(x16) |
| lim = s.swiglu_limit |
| a = torch.nn.functional.silu(g.clamp(max=lim)) * u.clamp(min=-lim, max=lim) |
| sh = st.sh_w2(a.to(torch.float16)) |
| return (routed + sh).to(torch.bfloat16) |
|
|
| def forward_block(self, ids: torch.Tensor, pos: torch.Tensor, win: Window, trace: dict | None = None |
| ) -> torch.Tensor: |
| """The three stages over a query block. ids/pos [B, k]; win: the context window (block_inputs). |
| Returns head_hidden [B*k, H] bf16 (pre-norm). ``trace`` (eager only) collects per-stage tensors.""" |
| s = self.s |
| B, k = ids.shape |
| T = B * k |
| pos_f = pos.reshape(T) |
| x = self.embed[ids.reshape(T)] |
| stream = x.unsqueeze(1).expand(T, s.hc_mult, s.hidden).contiguous() |
| pre_in = None |
| for i, st in enumerate(self.stages): |
| pre, post, comb = self._mixes(stream, *st.hc_attn) |
| xin = self._collapse_norm(stream, pre_in, st.attn_norm) |
| a = self._attention(st, i, xin, pos_f, win, B, k) |
| stream = self._post(a, stream, post, comb) |
| fpre, fpost, fcomb = self._mixes(stream, *st.hc_ffn) |
| xin = self._collapse_norm(stream, pre, st.ffn_norm) |
| y = self._moe(st, xin, trace, f"stage{i}") |
| stream = self._post(y, stream, fpost, fcomb) |
| pre_in = fpre |
| if trace is not None: |
| trace[f"stage{i}.attn_out"] = a |
| trace[f"stage{i}.ffn_in"] = xin |
| trace[f"stage{i}.ffn_out"] = y |
| trace[f"stage{i}.stream"] = stream |
| return self._collapse(stream, pre_in) |
|
|
| def heads(self, hh: torch.Tensor, anchors: torch.Tensor, B: int, k: int, trace: dict | None = None, |
| n_out: int | None = None) -> tuple[torch.Tensor, torch.Tensor]: |
| """Greedy sequential Markov drafting + confidence over head_hidden [B*k, H], for the first ``n_out`` (default |
| k) positions of every request; later positions return draft 0 and confidence 0. Positions < n_out are |
| bit-identical for every n_out (CONTRACT 1.3): head rows go through the GEMM in chunks of <= HEAD_CHUNK rows on |
| one fixed kernel shape, Markov and argmax are per row, and the confidence product keeps its [B, k] shape |
| (zero Markov embeddings past n_out).""" |
| s = self.s |
| n = k if n_out is None else n_out |
| if not 1 <= n <= k: |
| raise ValueError(f"n_out {n} outside 1..{k}") |
| rows = self._norm(hh, self.norm).view(B, k, s.hidden)[:, :n].reshape(B * n, s.hidden).to(torch.float16) |
| l32 = torch.empty(B * n, s.vocab, dtype=torch.float32, device=self.dev) |
| for c in range(0, B * n, self.HEAD_CHUNK): |
| self.head(rows[c:c + self.HEAD_CHUNK], out=l32[c:c + self.HEAD_CHUNK]) |
| logits = l32.to(torch.bfloat16).view(B, n, s.vocab) |
| out = torch.zeros(B, k, dtype=torch.int64, device=self.dev) |
| if self.fused and trace is None: |
| for i in range(n): |
| FU.markov_argmax_into(logits[:, i], anchors if i == 0 else out[:, i - 1], self.markov, out, i) |
| emb = self.markov.embed(torch.cat((anchors.unsqueeze(1), out[:, :n - 1]), 1)) |
| else: |
| prev = anchors |
| embs = [] |
| for i in range(n): |
| e = self.markov.embed(prev) |
| li = logits[:, i] + self.markov(prev, e) |
| d = li.argmax(dim=-1) |
| if trace is not None: |
| trace[f"markov{i}.logits"] = li |
| out[:, i] = d |
| embs.append(e) |
| prev = d |
| emb = torch.stack(embs, 1) |
| if n < k: |
| emb = torch.cat((emb, emb.new_zeros(B, k - n, emb.shape[-1])), 1) |
| feat = torch.cat((hh.view(B, k, s.hidden), emb), dim=-1).float() |
| conf = torch.sigmoid(feat @ self.conf_w) |
| if n < k: |
| conf[:, n:] = 0.0 |
| if trace is not None: |
| trace["base_logits"] = logits |
| return out, conf |
|
|
| def block_inputs(self, slots: torch.Tensor, anchors: torch.Tensor, last_pos: torch.Tensor, k: int, |
| ctx_lo: torch.Tensor | None = None, gather: bool | None = None |
| ) -> tuple[torch.Tensor, torch.Tensor, Window]: |
| """Query ids/positions and the window: ring rows whose position is in [ctx_lo, last_pos]. ctx_lo defaults to |
| max(last_pos - W + 1, 0); the server passes max(lo, anchor_pos - W) with the slot's contiguous lower bound. |
| ``gather`` (default: not fused) also materializes the window rows and their mask.""" |
| s = self.s |
| ar = self.arange_block[:k] |
| ids = torch.where(ar.unsqueeze(0) == 0, anchors.unsqueeze(1), torch.full_like(ar, s.noise_token).unsqueeze(0)) |
| pos = torch.minimum(last_pos.unsqueeze(1) + 1 + ar.unsqueeze(0), self.q_clamp) |
| if ctx_lo is None: |
| ctx_lo = (last_pos - (s.window - 1)).clamp_min(0) |
| win = Window(slots, ctx_lo, last_pos) |
| if gather is None: |
| gather = not self.fused |
| if gather: |
| cpos = self.ring_pos[slots] |
| win.valid = (cpos >= ctx_lo.unsqueeze(1)) & (cpos <= last_pos.unsqueeze(1)) |
| win.kv = self.ring_kv[slots] |
| return ids, pos, win |
|
|
| def draft(self, slots: torch.Tensor, anchors: torch.Tensor, last_pos: torch.Tensor, k: int, |
| trace: dict | None = None, ctx_lo: torch.Tensor | None = None, n_out: int | None = None |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Eager draft step. slots/anchors/last_pos (/ctx_lo) int64 [B] on the device. The body always runs the |
| whole k-token block; ``n_out`` limits the head / Markov / confidence to the first n_out positions.""" |
| B = slots.shape[0] |
| ids, pos, win = self.block_inputs(slots, anchors, last_pos, k, ctx_lo) |
| hh = self.forward_block(ids, pos, win, trace) |
| if trace is not None: |
| trace["head_hidden"] = hh |
| return self.heads(hh, anchors, B, k, trace, n_out) |
|
|
| def set_query_clamp(self, max_model_len: int) -> None: |
| if not 0 < max_model_len <= self.max_model_len: |
| raise ValueError(f"max_model_len {max_model_len} outside 1..{self.max_model_len}") |
| self.q_clamp.fill_(max_model_len - 1) |
|
|
| def draft_with_margin(self, slots: torch.Tensor, anchors: torch.Tensor, last_pos: torch.Tensor, k: int, |
| ctx_lo: torch.Tensor | None = None, n_out: int | None = None |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Eager draft that also returns top1 - top2 of (base + Markov) logits per position (CONTRACT WANT_MARGIN; |
| 0 past n_out). Same drafts as draft(): the argmax runs on the same bf16 logits.""" |
| tr: dict = {} |
| ids, conf = self.draft(slots, anchors, last_pos, k, trace=tr, ctx_lo=ctx_lo, n_out=n_out) |
| n = k if n_out is None else n_out |
| mg = torch.zeros(slots.shape[0], k, dtype=torch.float32, device=self.dev) |
| for i in range(n): |
| v = tr[f"markov{i}.logits"].float().topk(2, dim=-1).values |
| mg[:, i] = v[:, 0] - v[:, 1] |
| return ids, conf, mg |
|
|
| |
| def capture(self, batch_sizes=(1, 2, 3, 4), ks=(4, 5), insert_rows=(8, 16, 32, 64, 128), n_outs=(4, 5), |
| margins: bool = True, pack_reply: bool = False) -> float: |
| """One graph per (B, k, n_out) (n_out in n_outs and <= k; n_out = k always) for the draft step, the same with |
| top1 - top2 margins (CONTRACT WANT_MARGIN; shadow mode asks for them) when ``margins``, and one per row |
| bucket for the insert. ``pack_reply`` adds optional float64 reply twins; |
| default False, so normal capture count and outputs remain unchanged. Returns seconds.""" |
| t0 = time.time() |
| dev = self.dev |
| if self.pool is None: |
| self.pool = torch.cuda.graph_pool_handle() |
| bmax = max(batch_sizes) |
| self.g_slots = torch.full((bmax,), self.n_slots, dtype=torch.int64, device=dev) |
| self.g_anchor = torch.zeros(bmax, dtype=torch.int64, device=dev) |
| self.g_last = torch.zeros(bmax, dtype=torch.int64, device=dev) |
| self.g_lo = torch.zeros(bmax, dtype=torch.int64, device=dev) |
| rmax = max(insert_rows) |
| self.gi_slot = torch.full((rmax,), self.n_slots, dtype=torch.int64, device=dev) |
| self.gi_pos = torch.zeros(rmax, dtype=torch.int64, device=dev) |
| self.gi_taps = torch.zeros(rmax, self.s.hidden * len(self.s.taps), dtype=torch.bfloat16, device=dev) |
| s1 = torch.cuda.Stream(dev) |
| for B in batch_sizes: |
| for k in ks: |
| for n in sorted({x for x in n_outs if x <= k} | {k}): |
| args = (self.g_slots[:B], self.g_anchor[:B], self.g_last[:B], k) |
| lo = self.g_lo[:B] |
| s1.wait_stream(torch.cuda.current_stream(dev)) |
| with torch.cuda.stream(s1): |
| for _ in range(2): |
| self.draft(*args, ctx_lo=lo, n_out=n) |
| torch.cuda.current_stream(dev).wait_stream(s1) |
| torch.cuda.synchronize(dev) |
| gr = torch.cuda.CUDAGraph() |
| |
| |
| with torch.cuda.graph(gr, pool=self.pool, stream=s1): |
| ids, conf = self.draft(*args, ctx_lo=lo, n_out=n) |
| self.graphs[(B, k, n)] = {"graph": gr, "ids": ids, "conf": conf} |
| if margins: |
| s1.wait_stream(torch.cuda.current_stream(dev)) |
| with torch.cuda.stream(s1): |
| for _ in range(2): |
| self.draft_with_margin(*args, ctx_lo=lo, n_out=n) |
| torch.cuda.current_stream(dev).wait_stream(s1) |
| torch.cuda.synchronize(dev) |
| gr = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(gr, pool=self.pool, stream=s1): |
| ids, conf, mg = self.draft_with_margin(*args, ctx_lo=lo, n_out=n) |
| self.margin_graphs[(B, k, n)] = {"graph": gr, "ids": ids, "conf": conf, "mg": mg} |
| if pack_reply: |
| for margin in ((False, True) if margins else (False,)): |
| gr = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(gr, pool=self.pool, stream=s1): |
| values = (self.draft_with_margin(*args, ctx_lo=lo, n_out=n) if margin |
| else self.draft(*args, ctx_lo=lo, n_out=n)) |
| packed = torch.cat([v.to(torch.float64) for v in values], 1) |
| self.reply_graphs[(B, k, n, margin)] = {"graph": gr, "packed": packed} |
| for n in insert_rows: |
| args = (self.gi_slot[:n], self.gi_pos[:n]) |
| s1.wait_stream(torch.cuda.current_stream(dev)) |
| with torch.cuda.stream(s1): |
| for _ in range(2): |
| self.insert(*args, self.main_x(self.gi_taps[:n])) |
| torch.cuda.current_stream(dev).wait_stream(s1) |
| torch.cuda.synchronize(dev) |
| gr = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(gr, pool=self.pool, stream=s1): |
| self.insert(*args, self.main_x(self.gi_taps[:n])) |
| self.insert_graphs[n] = {"graph": gr} |
| |
| self.reset_slots(torch.tensor([self.n_slots], device=dev)) |
| torch.cuda.synchronize(dev) |
| return time.time() - t0 |
|
|
| def has_graph(self, B: int, k: int, n_out: int | None = None, margin: bool = False) -> bool: |
| return (B, k, k if n_out is None else n_out) in (self.margin_graphs if margin else self.graphs) |
|
|
| def draft_graphed(self, slots: torch.Tensor, anchors: torch.Tensor, last_pos: torch.Tensor, k: int, |
| ctx_lo: torch.Tensor | None = None, n_out: int | None = None, margin: bool = False, |
| packed_reply: bool = False |
| ) -> tuple[torch.Tensor, ...]: |
| """Replay the (B, k, n_out) graph; returns views of the graph's output buffers (ids, conf[, margin]; copy |
| before the next replay). ``margin``: the draft_with_margin graph (same ids and confidences). |
| ``packed_reply`` returns a one-tuple holding the captured float64 wire staging tensor; |
| requires capture(pack_reply=True), including a margin twin when margin=True.""" |
| B = slots.shape[0] |
| n = k if n_out is None else n_out |
| g = (self.reply_graphs.get((B, k, n, margin)) if packed_reply |
| else (self.margin_graphs if margin else self.graphs).get((B, k, n))) |
| if g is None: |
| raise KeyError(f"no CUDA graph for batch {B}, depth {k}, n_out {n}, margin {margin}, packed {packed_reply}") |
| self.g_slots[:B].copy_(slots, non_blocking=True) |
| self.g_anchor[:B].copy_(anchors, non_blocking=True) |
| self.g_last[:B].copy_(last_pos, non_blocking=True) |
| if ctx_lo is None: |
| ctx_lo = (last_pos - (self.s.window - 1)).clamp_min(0) |
| self.g_lo[:B].copy_(ctx_lo, non_blocking=True) |
| g["graph"].replay() |
| if packed_reply: |
| return (g["packed"],) |
| return (g["ids"], g["conf"], g["mg"]) if margin else (g["ids"], g["conf"]) |
|
|
| def insert_taps_graphed(self, slot_rows: torch.Tensor, pos_rows: torch.Tensor, taps: torch.Tensor) -> None: |
| """main_x + insert for n rows through the smallest captured bucket >= n (padding rows hit the scratch slot).""" |
| n = slot_rows.shape[0] |
| sizes = sorted(self.insert_graphs) |
| b = next((x for x in sizes if x >= n), None) |
| if b is None: |
| raise ValueError(f"{n} insert rows exceed the largest graph bucket {sizes[-1]}") |
| self.gi_slot[:n].copy_(slot_rows, non_blocking=True) |
| self.gi_slot[n:b].fill_(self.n_slots) |
| self.gi_pos[:n].copy_(pos_rows, non_blocking=True) |
| self.gi_pos[n:b].zero_() |
| self.gi_taps[:n].copy_(taps, non_blocking=True) |
| self.gi_taps[n:b].zero_() |
| self.insert_graphs[b]["graph"].replay() |
|
|
| |
| def memory_table(self) -> dict: |
| def nb(t: torch.Tensor) -> int: |
| return t.numel() * t.element_size() |
|
|
| st_dense = sum(sum(l.nbytes() for l in (st.wq_a, st.wkv, st.wq_b, st.wo_b, st.sh_w1, st.sh_w3, st.sh_w2)) |
| for st in self.stages) |
| st_woa = sum(nb(st.wo_a8) + nb(st.wo_a_sc) + (nb(st.wo_a_t) if st.wo_a_t is not None else 0) |
| for st in self.stages) |
| st_exp = sum(sum(nb(t) for t in st.bank.t.values()) for st in self.stages) |
| st_small = sum(nb(t) for st in self.stages for t in (st.attn_sink, st.q_norm, st.kv_norm, st.attn_norm, |
| st.ffn_norm, st.gate_w, st.gate_bias, *st.hc_attn, |
| *st.hc_ffn)) |
| return { |
| "embed_bf16": nb(self.embed), |
| "head_exl3_k6": self.head.nbytes(), |
| "main_proj_exl3_k5": self.main_proj.nbytes(), |
| "stages_dense_exl3_k5": st_dense, |
| "stages_wo_a_fp8": st_woa, |
| "stages_experts_exl3_k2": st_exp, |
| "stages_norms_hc_gate_fp32": st_small, |
| "markov_w1_w2_bf16": nb(self.markov.w1) + nb(self.markov.w2), |
| "markov_cache_rows_bf16": nb(self.markov.cache) + nb(self.markov.slot_of_token), |
| "confidence_and_norms": nb(self.conf_w) + nb(self.norm) + nb(self.main_norm), |
| "rope_cos_sin_fp32": nb(self.cos_sin), |
| "windows_per_slot": nb(self.ring_kv[0]) + nb(self.ring_pos[0]), |
| "windows_total": nb(self.ring_kv) + nb(self.ring_pos), |
| "coop_scratch": sum(nb(t) for t in (self.coop.had_g, self.coop.had_u, self.coop.gu_g, self.coop.gu_u, |
| self.coop.act_out, self.coop.d_out, self.coop.ctr, *self.coop.out)), |
| "torch_allocated": torch.cuda.memory_allocated(self.dev), |
| "torch_reserved": torch.cuda.memory_reserved(self.dev), |
| "torch_max_allocated": torch.cuda.max_memory_allocated(self.dev), |
| } |
|
|
