"""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 # NB1 level 2: the smallest insert graph bucket (one request's decode-time flush) 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 # NativeDense of the rank's wo_a groups (native dense) or None 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 # NB1 batch-invariant draft dispatch (default off): every draft-body shape B x block (B = 1..4, the fused # experts' 20-row limit) runs FP8 with B1's split-K, B1's expert variant, and no paired T20 shared GEMM. 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 # Full scale grid and explicit offsets preserve rank-1 scale blocks. 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 # Native Markov uses FP32 bias, never the deployed BF16 cache. 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 # no bf16 bias rows: the fp32 bias is exact (factorized or hot rows) 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() # NB1: None reads DSPARK_NATIVE_BATCH_INVARIANT (default 0 = off); 1 = draft dispatch, 2 = also insertion. # Only the kx2 dispatch implements it. 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) # ------------------------------------------------------------------------------------------ construction 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: # official combine all-reduces [rows, 3 slots + shared, hidden] fp32 per stage 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" # ------------------------------------------------------------------------------------------ context side 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))) # bf16, the official linear 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 # ------------------------------------------------------------------------------------------ draft step 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 # shared expert partial (fp32 [T, H] over the rank's 1152 channels) 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: # K2 routed experts + native shared: the coop kernel adds sh 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): # official combine: fp32 all-reduce of [T, 4, H], then epilogue 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) # Retain K6 weights but keep the head's FP32 accumulation/output. 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)