| """Native 1017 weights in the TP2 drafter: NativePacked / NativeExpanded / NativeHybrid arms over lane E's |
| ``TPDrafter`` (same collectives, graphs, contract and server), with the deployed K2 TP2 as the control. |
| |
| Arms (``ArmConfig``; presets in ARMS): |
| experts "k2" (EXL3 K2 coop kernel, the control) | "packed" (official MXFP4 bytes) | "expanded" (exact BF16) |
| ``expanded_stages``: stages whose experts are expanded when experts == "packed" (NativeHybrid) |
| dense "k5" (EXL3 K5, the control: E's attention / shared-expert path, engine semantics) | "fp8" (official |
| bytes + scales) | "bf16" (exact decode); with fp8/bf16 the attention and shared expert follow the |
| OFFICIAL semantics (act QDQ before every linear, bf16 linear outputs, the official KV QDQ) |
| kv_qdq "official" (act_quant per 32 over all 512 KV dims after RoPE, model.py) | "engine" (the V4 fp8_ds_mla |
| record the Spark's engine stores: 448 NoPE per-64 + 64 bf16 RoPE dims) [native dense only] |
| wo_a_qdq True: act_quant(o) before the wo_a einsum (model.py) | False: bf16 o (the engine's W8A16 branch) |
| combine "official": per-slot fp32 partials all-reduced, each expert rounded to bf16, ascending-id fp32 sum, + |
| bf16 shared, bf16 (MoE.forward) | "fused": one fp32 sum per rank, one all-reduce, one bf16 rounding |
| head "k6" only so far (the target-shared head stays as deployed: NDM's first condition); the BF16 head arm |
| is a separate step (NATIVE.md §5) |
| |
| Everything not listed (embedding, mHC, norms, routing, attention kernel, Markov, confidence) is E's code unchanged. |
| """ |
| from __future__ import annotations |
|
|
| import os |
| from collections.abc import Callable |
| from dataclasses import dataclass, field, replace |
|
|
| import torch |
|
|
| from ..worker import fused as FU |
| from ..worker import fused_e as FE |
| from ..worker.checkpoint import Checkpoint |
| from ..worker.exl3 import kernels as K |
| from ..worker.exl3.bank import Exl3Bank |
| from ..worker.exl3.plan import PROJ_DOWN, PROJ_GATE, PROJ_UP, Exl3MoEPlan |
| from ..worker.model import facc_ext |
| from ..worker.ops import rope_cos_sin |
| from ..worker.tp2 import (KERNEL_PLAN, Exl3Mat, MarkovShard, Rank, TPDrafter, coop_with_shared, col_shard, row_shard, |
| vocab_split) |
| from . import ops as NO |
| from .loader import LoadStats, load_bf16, load_dense, load_expert_banks |
| from .official import OfficialShards |
| from .kx import KXSharedPair |
| from .attention import attention_into |
|
|
| NATIVE_ROOT = os.environ.get("NATIVE_1017_ROOT", "/opt/dsv41/x86host-data/native-1017") |
| INSERT_SUBBATCH = 8 |
|
|
|
|
| from .arms import ARMS, ArmConfig, batch_invariant_from_env |
|
|
|
|
| @dataclass |
| class NativeStage: |
| """StageShard's fields (so E's attention / FFN code can run on a stage whose dense part is K5) plus the native |
| parts. Dense fields hold Exl3Mat (dense k5) or NativeDense (fp8 / bf16).""" |
| wq_a: object |
| wkv: object |
| wq_b: object |
| wo_b: object |
| wo_a8: torch.Tensor | None |
| wo_a_sc: torch.Tensor | None |
| wo_a_n: object |
| attn_sink: torch.Tensor |
| q_norm: torch.Tensor |
| kv_norm: torch.Tensor |
| attn_norm: torch.Tensor |
| ffn_norm: torch.Tensor |
| hc_attn: tuple |
| hc_ffn: tuple |
| gate_w: torch.Tensor |
| gate_bias: torch.Tensor |
| sh_w1: object |
| sh_w3: object |
| sh_w2: object |
| bank: Exl3Bank | None |
| groups: list |
| perm: torch.Tensor |
| nbank: NO.ExpertBank | None = None |
| qkv: object = None |
| gu: object = None |
| shared_ones: torch.Tensor | None = None |
| shared_hq: torch.Tensor | None = None |
| shared_pair: KXSharedPair | None = None |
| adapter_wo_a: Callable[[torch.Tensor], torch.Tensor] | None = None |
|
|
|
|
| class NativeRank: |
| """Rank r's tensors for an arm: the replicated / shared parts exactly as lane E's Rank builds them, the expert |
| banks and dense matrices in the arm's representation.""" |
|
|
| def __init__(self, r: int, dev: torch.device, ck: Checkpoint, sh: OfficialShards, cfg: ArmConfig, plan, |
| slots: int, max_model_len: int, max_rows: int, markov_cache_rows: int, tokens, tp: int, |
| stats: LoadStats, batch_invariant: bool = False): |
| s = ck.shape |
| |
| |
| invariant_rows = tuple(s.block * b for b in range(1, 20 // s.block + 1)) if batch_invariant else () |
| self.r, self.dev, self.s, self.tp, self.cfg = r, dev, s, tp, cfg |
| k_dense = ck.drafter_k |
| xplan = Exl3MoEPlan.from_dict(ck.plan_json) |
| heads_r = s.n_heads // tp |
| self.h0, self.h1 = r * heads_r, (r + 1) * heads_r |
| g_r = s.o_groups // tp |
| self.g0, self.g1 = r * g_r, (r + 1) * g_r |
| self.vrange = vocab_split(s.vocab, tp)[r] |
| self.v0, self.v1 = self.vrange |
| hd = s.head_dim |
| rows_per_group = s.o_lora |
| i_loc = s.inter // tp |
| self.i0, self.i1 = r * i_loc, (r + 1) * i_loc |
| nat = cfg.dense != "k5" |
|
|
| def exl3(name, want_k): |
| 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 tr, suh, svh, k |
|
|
| def mat(parts, k, name, ext=None, force_shape=0, fixed_bucket=None): |
| tr, suh, svh = parts |
| return Exl3Mat(tr, suh, svh, k, dev, plan, ext, force_shape, name, fixed_bucket) |
|
|
| def g(name, dtype=None): |
| t = ck.get(name) |
| return t.to(dev, dtype=dtype or t.dtype).contiguous() |
|
|
| def dense(name, rows=None, cols=None): |
| baseline = load_dense(sh, name, dev, cfg.dense, stats, rows, cols) |
| selected = (8,) if name == "mtp.0.main_proj" else (5,) if name.endswith("attn.wq_b") else () |
| if cfg.dense_profile == "measured-small" and name != "mtp.0.main_proj": |
| selected = (5, 10, 15) |
| if invariant_rows and name != "mtp.0.main_proj": |
| selected = invariant_rows |
| if cfg.dense_kernel != "kx-selective" or not selected: |
| return baseline |
| from ..kernels.fp8dense import DenseFP8 |
| from .kx import KXSelectiveDense |
| |
| candidate = DenseFP8.from_checkpoint( |
| sh.get(name + ".weight")[(slice(*rows) if rows else slice(None)), |
| (slice(*cols) if cols else slice(None))], |
| sh.get(name + ".scale"), dev, |
| groups=4 if name.endswith("attn.wo_a") else 1, |
| row_offset=rows[0] if rows else 0, col_offset=cols[0] if cols else 0) |
| if invariant_rows and name != "mtp.0.main_proj": |
| return KXSelectiveDense(baseline, candidate, selected, |
| splits=candidate.configuration(s.block)[0]) |
| return KXSelectiveDense(baseline, candidate, selected) |
|
|
| self.embed = g("embed.weight") |
| self.head, self.head_bf16 = None, None |
| self.mrows = None |
| hot_tokens = tokens |
| if cfg.native: |
| tokens = None |
| if cfg.head == "bf16": |
| hw = load_bf16(sh, "head.weight", dev, stats, rows=(self.v0, self.v1)) |
| self.head_bf16 = NO.NativeDense(f"head[{self.v0}:{self.v1}]", "bf16", hw, None) |
| tokens = None |
| else: |
| tr, suh, svh, hk = exl3("head", ck.head_k) |
| self.head = mat(col_shard(tr, suh, svh, self.v0, self.v1), hk, f"head[{self.v0}:{self.v1}]", |
| force_shape=3, fixed_bucket="head") |
| if nat: |
| self.main_proj = dense("mtp.0.main_proj") |
| else: |
| tr, suh, svh, mk = exl3("mtp.0.main_proj", k_dense["main_proj"]) |
| self.main_proj = mat((tr, suh, svh), mk, "main_proj", ext=facc_ext()) |
| 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) |
| self.markov = MarkovShard(g(last + "markov_head.embed.weight"), g(last + "markov_head.head.weight"), self.v0, |
| self.v1, tokens, markov_cache_rows) |
| if cfg.head == "bf16" and cfg.markov == "hot": |
| with torch.cuda.device(dev): |
| self.mrows = NO.MarkovRowsFP32(self.markov.w1, self.markov.w2, hot_tokens, cfg.markov_rows) |
| self.stages: list[NativeStage] = [] |
| any_k2 = False |
| for i in range(s.n_stages): |
| p = f"mtp.{i}." |
| L = s.n_target_layers + i |
| lp = xplan.drafter_layers[L] |
| kind = cfg.stage_expert_kind(i) |
| bank, groups, nbank = None, [], None |
| if kind == "k2": |
| any_k2 = True |
| bank = Exl3Bank(lp, s.hidden, s.inter, tp, r) |
| 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() |
| elif kind == "packed" and cfg.expert_kernel in ("kx-regfed", "kx-fused"): |
| from .kx import KXFusedBank, KXRegisterBank |
| host = load_expert_banks(sh, i, [(r, torch.device("cpu"))], tp, kind, stats)[r] |
| if cfg.expert_kernel != "kx-fused": |
| nbank = KXRegisterBank(host, dev) |
| elif invariant_rows: |
| nbank = KXFusedBank(host, dev, i, variant_rows=s.block) |
| else: |
| nbank = KXFusedBank(host, dev, i) |
| del host |
| else: |
| nbank = load_expert_banks(sh, i, [(r, dev)], tp, kind, stats)[r] |
| r0, r1 = self.g0 * rows_per_group, self.g1 * rows_per_group |
| if nat: |
| woa = dense(p + "attn.wo_a", rows=(r0, r1)) |
| st_dense = dict( |
| wq_a=dense(p + "attn.wq_a"), wkv=dense(p + "attn.wkv"), |
| wq_b=dense(p + "attn.wq_b", rows=(self.h0 * hd, self.h1 * hd)), |
| wo_b=dense(p + "attn.wo_b", cols=(r0, r1)), |
| wo_a8=woa.w if woa.kind == "fp8" else None, wo_a_sc=woa.sc if woa.kind == "fp8" else None, |
| wo_a_n=woa, |
| sh_w1=dense(p + "ffn.shared_experts.w1", rows=(self.i0, self.i1)), |
| sh_w3=dense(p + "ffn.shared_experts.w3", rows=(self.i0, self.i1)), |
| sh_w2=dense(p + "ffn.shared_experts.w2", cols=(self.i0, self.i1))) |
| else: |
| wq_a = exl3(p + "attn.wq_a", k_dense["attn.wq_a"]) |
| wkv = exl3(p + "attn.wkv", k_dense["attn.wq_a"]) |
| wq_b = exl3(p + "attn.wq_b", k_dense["attn.wq_b"]) |
| wo_b = exl3(p + "attn.wo_b", k_dense["attn.wo_b"]) |
| w1 = exl3(p + "ffn.shared_experts.w1", k_dense["ffn.shared_experts.w1"]) |
| w3 = exl3(p + "ffn.shared_experts.w3", k_dense["ffn.shared_experts.w1"]) |
| w2 = exl3(p + "ffn.shared_experts.w2", k_dense["ffn.shared_experts.w2"]) |
| woa = ck.get(p + "attn.wo_a.weight") |
| woa_sc = ck.get(p + "attn.wo_a.scale") |
| st_dense = dict( |
| wq_a=mat(wq_a[:3], wq_a[3], p + "wq_a"), wkv=mat(wkv[:3], wkv[3], p + "wkv"), |
| wq_b=mat(col_shard(*wq_b[:3], self.h0 * hd, self.h1 * hd), wq_b[3], p + "wq_b"), |
| wo_b=mat(row_shard(*wo_b[:3], r0, r1), wo_b[3], p + "wo_b"), |
| wo_a8=woa[r0:r1].contiguous().to(dev).view(torch.uint8), |
| wo_a_sc=woa_sc[r0 // 32:r1 // 32].to(dev, dtype=torch.float32).contiguous(), wo_a_n=None, |
| sh_w1=mat(col_shard(*w1[:3], self.i0, self.i1), w1[3], p + "sh_w1"), |
| sh_w3=mat(col_shard(*w3[:3], self.i0, self.i1), w3[3], p + "sh_w3"), |
| sh_w2=mat(row_shard(*w2[:3], self.i0, self.i1), w2[3], p + "sh_w2")) |
| st = NativeStage( |
| **st_dense, |
| attn_sink=g(p + "attn.attn_sink", torch.float32)[self.h0:self.h1].contiguous(), |
| 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), |
| bank=bank, groups=groups, perm=torch.tensor(lp.perm, dtype=torch.int64, device=dev), nbank=nbank) |
| self.stages.append(st) |
| if cfg.dense_kernel == "kx-selective" and cfg.dense_profile == "measured-small": |
| st.shared_ones = torch.ones(20, 1, dtype=torch.float32, device=dev) |
| st.shared_hq = torch.empty(20, i_loc, dtype=torch.bfloat16, device=dev) |
| if not invariant_rows: |
| st.shared_pair = KXSharedPair(st.sh_w1, st.sh_w3) |
| self.coop = K.get_coop_scratch(dev, max_rows, s.top_k, s.hidden, i_loc) if any_k2 else None |
| if any_k2: |
| k2 = next(st for st in self.stages if st.bank is not None) |
| with torch.cuda.device(dev): |
| K.prime(dev, k2.groups, self.coop, None, s.hidden, s.top_k, s.swiglu_limit, False, True) |
| self.cos_sin = rope_cos_sin(max_model_len + 2048, s.rope_dim, s.rope_theta).to(dev) |
| self.q_clamp = torch.tensor(max_model_len - 1, dtype=torch.int64, device=dev) |
| self.ring_kv = torch.zeros(slots + 1, s.n_stages, s.window, hd, dtype=torch.bfloat16, device=dev) |
| self.ring_pos = torch.full((slots + 1, s.window), -1, dtype=torch.int64, device=dev) |
| self.attention_out = torch.empty(max_rows, heads_r, hd, dtype=torch.bfloat16, device=dev) |
| self.arange_block = torch.arange(s.block + 1, dtype=torch.int64, device=dev) |
| self.pool = None |
| self.cap_stream = torch.cuda.Stream(dev) |
|
|
| def nbytes(self) -> dict: |
| def nb(t): |
| return 0 if t is None else t.numel() * t.element_size() |
|
|
| def mb(m): |
| return m.nbytes() if m is not None else 0 |
| d = {"embed_bf16": nb(self.embed), "head_shard": mb(self.head) + mb(self.head_bf16), |
| "main_proj": mb(self.main_proj)} |
| d["stages_dense_replicated"] = sum(mb(st.wq_a) + mb(st.wkv) for st in self.stages) |
| d["stages_dense_shards"] = sum(mb(st.wq_b) + mb(st.wo_b) + mb(st.sh_w1) + mb(st.sh_w3) + mb(st.sh_w2) |
| for st in self.stages) |
| d["stages_wo_a"] = sum((mb(st.wo_a_n) if st.wo_a_n is not None else nb(st.wo_a8) + nb(st.wo_a_sc)) |
| for st in self.stages) |
| d["stages_shared_pair_bf16"] = sum(mb(st.shared_pair) for st in self.stages) |
| d["stages_experts_k2"] = sum(sum(nb(t) for t in st.bank.t.values()) for st in self.stages |
| if st.bank is not None) |
| d["stages_experts_native"] = sum(st.nbank.nbytes() for st in self.stages if st.nbank is not None) |
| d["stages_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)) |
| d["markov_embed_bf16"] = nb(self.markov.w1) |
| d["markov_head_shard_bf16"] = nb(self.markov.w2) |
| d["markov_cache_shard_bf16"] = nb(self.markov.cache) + nb(self.markov.slot_of_token) |
| d["markov_rows_fp32"] = self.mrows.nbytes() if self.mrows is not None else 0 |
| d["rope_cos_sin_fp32"] = nb(self.cos_sin) |
| d["windows_total"] = nb(self.ring_kv) + nb(self.ring_pos) |
| d["attention_output"] = nb(self.attention_out) |
| return d |
|
|
|
|
| class NativeTPDrafter(TPDrafter): |
| """TPDrafter with an arm's weight representation. ``arm``: an ArmConfig or a key of ARMS.""" |
|
|
| def __init__(self, ck: Checkpoint, devs: list[torch.device], arm: ArmConfig | str = "native-hybrid", |
| native_root: str = NATIVE_ROOT, batch_invariant: int | None = None, **kw): |
| self.cfg = ARMS[arm] if isinstance(arm, str) else arm |
| self.cfg.check() |
| |
| |
| self.batch_invariant = batch_invariant_from_env() if batch_invariant is None else int(batch_invariant) |
| if self.batch_invariant not in (0, 1, 2): |
| raise ValueError(f"batch_invariant level {self.batch_invariant}: expected 0, 1 or 2") |
| if self.batch_invariant and not self.cfg.kx2_dispatch: |
| raise ValueError(f"{self.cfg.name}: batch-invariant dispatch is implemented for the kx2 dispatch only") |
| self.insert_rows = INSERT_SUBBATCH if self.batch_invariant == 2 else 0 |
| self.sh = OfficialShards(native_root) |
| self.load_stats = LoadStats() |
| super().__init__(ck, devs, **kw) |
|
|
| |
| def _make_rank(self, r, dev, ck, slots, max_model_len, max_rows, markov_cache_rows, tokens, tp): |
| if not self.cfg.native: |
| return super()._make_rank(r, dev, ck, slots, max_model_len, max_rows, markov_cache_rows, tokens, tp) |
| if self.batch_invariant: |
| return NativeRank(r, dev, ck, self.sh, self.cfg, self.plan, slots, max_model_len, max_rows, |
| markov_cache_rows, tokens, tp, self.load_stats, batch_invariant=True) |
| return NativeRank(r, dev, ck, self.sh, self.cfg, self.plan, slots, max_model_len, max_rows, |
| markov_cache_rows, tokens, tp, self.load_stats) |
|
|
| def comm_words(self) -> int: |
| |
| return 32 * 4 * self.s.hidden if self.cfg.combine == "official" else 32 * self.s.hidden |
|
|
| def _native_dense(self, rk) -> bool: |
| return isinstance(rk, NativeRank) and self.cfg.dense != "k5" |
|
|
| |
| def _main_x(self, rk, taps: torch.Tensor) -> torch.Tensor: |
| if not self._native_dense(rk): |
| return super()._main_x(rk, taps) |
| if self.insert_rows: |
| y = self._fixed_rows(rk.main_proj, NO.act_qdq(taps.to(torch.bfloat16))) |
| else: |
| y = rk.main_proj(NO.act_qdq(taps.to(torch.bfloat16))) |
| return FU.rms_norm(y, rk.main_norm, self.eps) |
|
|
| def _fixed_rows(self, linear, x: torch.Tensor) -> torch.Tensor: |
| """NB1 level 2: ``linear`` over fixed ``insert_rows``-row sub-batches (eager calls zero-pad the last one), so |
| every inserted row gets the bucket-8 arithmetic whatever the flush size or companions. Each sub-batch's |
| output is copied out before the next call (KX dense returns persistent per-M buffers).""" |
| n, step = x.shape[0], self.insert_rows |
| if n % step: |
| x = torch.cat((x, x.new_zeros(step - n % step, x.shape[1]))) |
| out = None |
| for a in range(0, x.shape[0], step): |
| y = linear(x[a:a + step]) |
| if out is None: |
| out = torch.empty(x.shape[0], y.shape[1], dtype=y.dtype, device=y.device) |
| out[a:a + step] = y |
| return out[:n] |
|
|
| def _kv_official(self, rk, kv16: torch.Tensor, w: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: |
| """model.py: kv_norm -> RoPE (bf16) -> act_quant(kv, 32) in place, over all 512 dims.""" |
| return NO.act_qdq(FU.rope(FU.rms_norm(kv16, w, self.eps), pos, rk.cos_sin)) |
|
|
| def _insert_rank(self, rk, slot_rows: torch.Tensor, pos_rows: torch.Tensor, taps: torch.Tensor) -> None: |
| if not self._native_dense(rk): |
| return super()._insert_rank(rk, slot_rows, pos_rows, taps) |
| self._insert_mx(rk, slot_rows, pos_rows, self._main_x(rk, taps)) |
|
|
| def insert_mx_eager(self, slot_rows: torch.Tensor, pos_rows: torch.Tensor, mx: torch.Tensor) -> None: |
| if not self.cfg.native or self.cfg.dense == "k5": |
| return super().insert_mx_eager(slot_rows, pos_rows, mx) |
| R = self.ALL |
| sr, pr, mxr = self.to_ranks(R, slot_rows), self.to_ranks(R, pos_rows), self.to_ranks(R, mx) |
| self._each(R, lambda rk: self._insert_mx(rk, sr[rk.r], pr[rk.r], mxr[rk.r].to(torch.bfloat16))) |
|
|
| def _insert_mx(self, rk, slot_rows: torch.Tensor, pos_rows: torch.Tensor, mx: torch.Tensor) -> None: |
| """Window KV of main_x rows (bf16): wkv (official linear) -> kv_norm -> RoPE -> KV QDQ -> ring.""" |
| mxq = NO.act_qdq(mx) |
| W = self.s.window |
| for i, st in enumerate(rk.stages): |
| kv = self._fixed_rows(st.wkv, mxq) if self.insert_rows else st.wkv(mxq) |
| if self.cfg.kv_qdq == "engine": |
| FU.kv_store(kv.float(), st.kv_norm, self.eps, pos_rows, rk.cos_sin, rk.ring_kv[:, i], slot_rows, |
| rk.ring_pos if i == 0 else None) |
| else: |
| rk.ring_kv[slot_rows, i, pos_rows % W] = self._kv_official(rk, kv, st.kv_norm, pos_rows) |
| if i == 0: |
| rk.ring_pos[slot_rows, pos_rows % W] = pos_rows |
|
|
| |
| def _attn_part(self, rk, i, st, stream, pre_in, pos, slots, lo, last, B, k): |
| if not self._native_dense(rk): |
| return super()._attn_part(rk, i, st, stream, pre_in, pos, slots, lo, last, B, k) |
| s = self.s |
| T = B * k |
| pre, post, comb = FU.hc_mixes(stream, *st.hc_attn, self.eps, s.hc_eps, s.hc_iters) |
| x = FU.collapse_norm(stream, pre_in, st.attn_norm, self.eps) |
| xq = NO.act_qdq(x) |
| qr = FU.rms_norm(st.wq_a(xq), st.q_norm, self.eps) |
| kv = st.wkv(xq) |
| if self.cfg.kv_qdq == "engine": |
| kvb = FU.kv_rope_qdq(FU.rms_norm(kv, st.kv_norm, self.eps), pos, rk.cos_sin) |
| else: |
| kvb = self._kv_official(rk, kv, st.kv_norm, pos) |
| hr = rk.h1 - rk.h0 |
| q = FU.rope(st.wq_b(NO.act_qdq(qr)).view(T, hr, s.head_dim), pos, rk.cos_sin) |
| o = attention_into(q.view(B, k, hr, s.head_dim), rk.ring_kv[:, i], rk.ring_pos, slots, lo, last, |
| kvb.view(B, k, s.head_dim), st.attn_sink, self.scale, |
| rk.attention_out[:T].view(B, k, hr, s.head_dim)) |
| o = FU.rope(o.view(T, hr, s.head_dim), pos, rk.cos_sin, inverse=True).view(T, -1) |
| if self.cfg.wo_a_qdq: |
| o = NO.act_qdq(o) |
| wa = st.wo_a_n |
| if st.adapter_wo_a is not None: |
| z = st.adapter_wo_a(o) |
| elif self.cfg.dense_kernel == "kx-selective" and self.cfg.dense_profile == "measured-small": |
| z = wa(o) |
| elif wa.kind == "fp8": |
| z = FU.wo_a_fp8(o, wa.w, wa.sc, rk.g1 - rk.g0) |
| else: |
| G = rk.g1 - rk.g0 |
| z = torch.bmm(o.view(T, G, -1).transpose(0, 1), wa.w.view(G, s.o_lora, -1).transpose(1, 2)) |
| z = z.transpose(0, 1).reshape(T, -1) |
| return st.wo_b(NO.act_qdq(z), out_f32=True), (pre, post, comb) |
|
|
| def _ffn_part(self, rk, st, stream, pre, trace=None, tag=""): |
| if not isinstance(rk, NativeRank) or (st.nbank is None and self.cfg.dense == "k5"): |
| return super()._ffn_part(rk, st, stream, pre, trace, tag) |
| s = self.s |
| fpre, fpost, fcomb = FU.hc_mixes(stream, *st.hc_ffn, self.eps, s.hc_eps, s.hc_iters) |
| xin, x16 = FE.collapse_norm2(stream, pre, st.ffn_norm, self.eps) |
| ids, sel, w, w16 = FE.route_fused(xin, st.gate_w, st.gate_bias, st.perm, s.top_k, s.route_scale) |
| if trace is not None and rk.r == 0: |
| trace[tag + ".experts"] = ids |
| trace[tag + ".route_w"] = w |
| trace[tag + ".ffn_in"] = xin |
| lim = s.swiglu_limit |
| xq = NO.act_qdq(xin) if (st.nbank is not None or self.cfg.dense != "k5") else None |
| |
| if self.cfg.dense == "k5": |
| g, u = st.sh_w1.gemm(x16), st.sh_w3.gemm(x16) |
| sh = st.sh_w2.gemm(FE.swiglu_fp16(g, u, lim)) |
| else: |
| g, u = st.shared_pair(xq) if st.shared_pair is not None else (st.sh_w1(xq), st.sh_w3(xq)) |
| if st.shared_hq is not None: |
| from ..kernels.fp8dense.epilogue import shared_activation_into |
| if st.shared_ones is None: |
| raise RuntimeError("shared activation workspace lacks initialized route weights") |
| hq = st.shared_hq[:xin.shape[0]] |
| shared_activation_into(g, u, st.shared_ones[:xin.shape[0]], hq) |
| else: |
| hq = NO.act_qdq(NO.swiglu_bf16(g, u, lim)) |
| sh = st.sh_w2(hq, out_f32=True) |
| if st.nbank is None: |
| return coop_with_shared(x16, sel, w16, st.groups, rk.coop, lim, sh, self.shared_scratch), \ |
| (fpre, fpost, fcomb) |
| T = xin.shape[0] |
| if self.cfg.combine == "official": |
| if self.cfg.expert_kernel == "kx-fused": |
| buf = st.nbank.slot_buffers[T] |
| else: |
| buf = torch.empty(T, 4, s.hidden, dtype=torch.float32, device=xin.device) |
| st.nbank.forward_slots(xq, ids, w, out=buf[:, :3]) |
| buf[:, 3] = sh |
| return (buf, ids), (fpre, fpost, fcomb) |
| slots = st.nbank.forward_slots(xq, ids, w) |
| return NO.ordered_slot_sum(slots, ids, False) + sh, (fpre, fpost, fcomb) |
|
|
| def forward_block(self, R, inp: dict, k: int, trace: dict | None = None): |
| if not self.cfg.native: |
| return super().forward_block(R, inp, k, trace) |
| s = self.s |
| bi = self._each(R, lambda rk: self._block_inputs(rk, inp[rk.r], k)) |
| B = next(iter(bi.values()))[0].shape[0] |
| T = B * k |
| stream = self._each(R, lambda rk: rk.embed[bi[rk.r][0].reshape(T)].unsqueeze(1) |
| .expand(T, s.hc_mult, s.hidden).contiguous()) |
| pre_in = {r: None for r in R} |
| for i in range(s.n_stages): |
| res = self._each(R, lambda rk: self._attn_part(rk, i, rk.stages[i], stream[rk.r], pre_in[rk.r], |
| bi[rk.r][1].reshape(T), bi[rk.r][2], bi[rk.r][5], |
| bi[rk.r][4], B, k)) |
| a = self.comm.all_reduce(R, {r: res[r][0] for r in R}) |
| mix = {r: res[r][1] for r in R} |
| stream = self._each(R, lambda rk: FU.hc_post(a[rk.r], stream[rk.r], mix[rk.r][1], mix[rk.r][2])) |
| res = self._each(R, lambda rk: self._ffn_part(rk, rk.stages[i], stream[rk.r], mix[rk.r][0], trace, |
| f"stage{i}")) |
| first = res[next(iter(R))][0] |
| if isinstance(first, tuple): |
| tot = self.comm.all_reduce(R, {r: res[r][0][0] for r in R}, out_dtype=torch.float32) |
| y = self._each(R, lambda rk: self._combine_official(rk, i, tot[rk.r], res[rk.r][0][1])) |
| else: |
| y = self.comm.all_reduce(R, {r: res[r][0] for r in R}) |
| fmix = {r: res[r][1] for r in R} |
| stream = self._each(R, lambda rk: FU.hc_post(y[rk.r], stream[rk.r], fmix[rk.r][1], fmix[rk.r][2])) |
| pre_in = {r: fmix[r][0] for r in R} |
| if trace is not None and 0 in R: |
| trace[f"stage{i}.attn_out"] = a[0] |
| trace[f"stage{i}.ffn_out"] = y[0] |
| trace[f"stage{i}.stream"] = stream[0] |
| hh = self._each(R, lambda rk: FU.collapse(stream[rk.r], pre_in[rk.r])) |
| return hh, bi |
|
|
| def _combine_official(self, rk, stage: int, parts: torch.Tensor, ids: torch.Tensor) -> torch.Tensor: |
| if self.cfg.expert_kernel == "kx-fused" and self.cfg.dense_profile == "measured-small": |
| from ..kernels.fusedmoe.epilogue import combine_official_into |
| out = rk.stages[stage].nbank.combined_outputs[parts.shape[0]] |
| combine_official_into(parts, ids, out) |
| return out |
| return NO.moe_combine_official(parts[:, :3], parts[:, 3], ids) |
|
|
|
|
| def _heads_bf16(self, R, hh: dict, bi: dict, B: int, k: int, n_out: int | None = None, margin: bool = False, |
| trace: dict | None = None): |
| """TPDrafter.heads with the official BF16 head (fp32 logits over the rank's vocabulary shard) and the exact fp32 |
| Markov bias (no cache): per position the shard's (max, lowest index) or top-2 meet in E's exchanges. Positions |
| >= n_out: draft 0, confidence 0 (CONTRACT 1.3).""" |
| 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}") |
|
|
| def logits_of(rk): |
| rows = FU.rms_norm(hh[rk.r], rk.norm, self.eps).view(B, k, s.hidden)[:, :n].reshape(B * n, s.hidden) |
| if rk.head_bf16 is not None: |
| return rk.head_bf16(rows.contiguous(), out_f32=True).view(B, n, rk.v1 - rk.v0) |
| |
| rows = rows.to(torch.float16) |
| logits = torch.empty(B * n, rk.v1 - rk.v0, dtype=torch.float32, device=rk.dev) |
| for c in range(0, B * n, self.HEAD_CHUNK): |
| rk.head.gemm(rows[c:c + self.HEAD_CHUNK], out=logits[c:c + self.HEAD_CHUNK]) |
| return logits.view(B, n, rk.v1 - rk.v0) |
| logits = self._each(R, logits_of) |
| out = self._each(R, lambda rk: torch.zeros(B, k, dtype=torch.int64, device=rk.dev)) |
| mg = self._each(R, lambda rk: torch.zeros(B, k, dtype=torch.float32, device=rk.dev)) if margin else None |
| anchors = {r: bi[r][3] for r in R} |
| if trace is not None: |
| trace["base_logits"] = torch.cat([logits[r].to(self.dev) for r in R], -1) if tuple(R) == self.ALL else None |
| for i in range(n): |
| prev = {r: anchors[r] if i == 0 else out[r][:, i - 1] for r in R} |
| full = self._each(R, lambda rk: NO.markov_logits_fp32(logits[rk.r][:, i].contiguous(), |
| prev[rk.r].contiguous(), rk.markov.w1, rk.markov.w2, |
| rk.mrows)) |
| if not margin: |
| loc = self._each(R, lambda rk: full[rk.r].max(dim=1)) |
| self.comm.argmax(R, self.ranks, {r: loc[r][0].contiguous() for r in R}, |
| {r: loc[r][1].contiguous() for r in R}, out, i) |
| else: |
| loc = self._each(R, lambda rk: (full[rk.r].topk(2, dim=-1).values.contiguous(), |
| full[rk.r].argmax(dim=-1))) |
| self.comm.margin(R, self.ranks, {r: loc[r][0] for r in R}, {r: loc[r][1] for r in R}, out, mg, i) |
| conf = None |
| if 0 in R: |
| rk = self.ranks[0] |
| with torch.cuda.device(rk.dev): |
| o = out[0] |
| emb = rk.markov.embed(torch.cat((anchors[0].unsqueeze(1), o[:, :n - 1]), 1)) |
| if n < k: |
| emb = torch.cat((emb, emb.new_zeros(B, k - n, emb.shape[-1])), 1) |
| feat = torch.cat((hh[0].view(B, k, s.hidden), emb), dim=-1).float() |
| conf = torch.sigmoid(feat @ rk.conf_w) |
| if n < k: |
| conf[:, n:] = 0.0 |
| return out, conf, mg |
|
|
|
|
| def _heads(self, R, hh, bi, B, k, n_out=None, margin=False, trace=None): |
| if self.cfg.native: |
| return _heads_bf16(self, R, hh, bi, B, k, n_out, margin, trace) |
| return TPDrafter.heads(self, R, hh, bi, B, k, n_out, margin, trace) |
|
|
|
|
| NativeTPDrafter.heads = _heads |
|
|
|
|
| def build(arm: str | ArmConfig, devs: list[torch.device], ck: Checkpoint, **kw) -> TPDrafter: |
| """The drafter for an arm: the plain TPDrafter for the control, NativeTPDrafter otherwise.""" |
| cfg = ARMS[arm] if isinstance(arm, str) else arm |
| if not cfg.native: |
| return TPDrafter(ck, devs, **kw) |
| return NativeTPDrafter(ck, devs, cfg, **kw) |
|
|
