| """KX lanes' kernels behind N1's ExpertBank interface (lane branches offload/kx-mxfp4, offload/kx-grouped; imported |
| from offload/dspark/kernels/** when that tree is present, refused loudly when not). |
| |
| KXRegisterBank: KX-MXFP4's register-fed SM86 W4A16 kernel (regfed.RegisterExperts, fused gate/up, device expert |
| schedule) for the projections; N1's swiglu_qdq for the official epilogue between them (full-precision expf, bf16 |
| rounding points, per-32 QDQ); per-slot fp32 down outputs (SLOT mode), so the official combine stays N1's. |
| The kernel's sticky status word (bit 0: an invalid expert id, NaN outputs) is zeroed at build and checked by check(). |
| """ |
| from __future__ import annotations |
|
|
| from typing import TYPE_CHECKING |
|
|
| import torch |
|
|
| from . import ops as NO |
|
|
| if TYPE_CHECKING: |
| from ..kernels.fp8dense import DenseFP8 |
|
|
|
|
| class KXRegisterBank: |
| def __init__(self, bank: NO.ExpertBank, dev: torch.device, max_rows: int = 21): |
| """``bank``: the rank's packed shards on the HOST (KX admission validates on the CPU before any GPU |
| allocation); RegisterExperts streams them one expert at a time into its fragment layout on ``dev``.""" |
| try: |
| from ..kernels.mxfp4.regfed import RegisterExperts |
| except ImportError as e: |
| raise RuntimeError("KX-MXFP4 kernels (offload/dspark/kernels/mxfp4, branch offload/kx-mxfp4) not " |
| "present in this tree") from e |
| if bank.kind != "packed" or bank.dev.type != "cpu": |
| raise ValueError("KXRegisterBank takes a packed bank on the host") |
| self.name, self.kind, self.E, self.H, self.I, self.dev = bank.name + "/kx", "packed", bank.E, bank.H, bank.I, dev |
| with torch.cuda.device(dev): |
| self.gate = RegisterExperts.from_checkpoint(bank.w13[:, 0], bank.s13[:, 0], dev) |
| self.up = RegisterExperts.from_checkpoint(bank.w13[:, 1], bank.s13[:, 1], dev) |
| bank.w13 = bank.s13 = None |
| self.down = RegisterExperts.from_checkpoint(bank.w2, bank.s2, dev) |
| bank.w2 = bank.s2 = None |
| self.status = torch.zeros(1, dtype=torch.int32, device=dev) |
| self.schedule = torch.empty(max_rows * 3, 8, dtype=torch.int32, device=dev) |
| self.max_rows = max_rows |
|
|
| def tensors(self) -> list[torch.Tensor]: |
| return [t for b in (self.gate, self.up, self.down) for t in (b.codes, b.scales)] |
|
|
| def nbytes(self) -> int: |
| return sum(t.numel() * t.element_size() for t in self.tensors()) |
|
|
| def forward_slots(self, xq: torch.Tensor, ids: torch.Tensor, rw: torch.Tensor, h: torch.Tensor | None = None, |
| out: torch.Tensor | None = None) -> torch.Tensor: |
| T = xq.shape[0] |
| if T > self.max_rows: |
| raise ValueError(f"{self.name}: {T} rows > {self.max_rows} (the kernel admits 64 routed slots)") |
| ids32 = ids.to(torch.int32).contiguous() |
| sched = self.schedule[:T * 3] |
| g = torch.empty(T, 3, self.I, dtype=torch.float32, device=xq.device) |
| u = torch.empty_like(g) |
| self.gate.linear_into(xq.contiguous(), ids32, g, self.status, sched, up=self.up, up_out=u) |
| hq = NO.swiglu_qdq(g, u, rw) |
| y = torch.empty(T, 3, self.H, dtype=torch.float32, device=xq.device) |
| self.down.linear_into(hq, ids32, y, self.status, sched, per_slot=True) |
| if out is not None: |
| out.copy_(y) |
| return out |
| return y |
|
|
| def check(self) -> None: |
| s = int(self.status.item()) |
| if s: |
| raise RuntimeError(f"{self.name}: KX kernel status {s:#x} (invalid expert ids: NaN outputs)") |
|
|
|
|
| class KXFusedBank(KXRegisterBank): |
| """KX delivery selector with caller-owned, pre-capture SLOT workspaces. |
| |
| A bank belongs to one rank/capture stream, like the parent drafter. Different |
| simultaneous streams must use separate banks and workspaces. |
| ``variant_rows``: select every row count's kernel variant as for that many rows |
| (batch-invariant mode: the B1 block); None keeps the measured per-row profile. |
| """ |
|
|
| def __init__(self, bank: NO.ExpertBank, dev: torch.device, stage: int, variant_rows: int | None = None): |
| from ..kernels.fusedmoe import FusedMoE, Workspace, select_variant |
|
|
| super().__init__(bank, dev, max_rows=20) |
| self.pipelines = { |
| rows: FusedMoE.from_register_banks(self.gate, self.up, self.down, |
| variant=select_variant(stage, rows if variant_rows is None |
| else variant_rows)) |
| for rows in range(1, 21) |
| } |
| self.workspaces = {rows: Workspace.allocate(rows, self.I, dev) for rows in self.pipelines} |
| self.slot_buffers = {rows: torch.empty(rows, 4, self.H, dtype=torch.float32, device=dev) |
| for rows in self.pipelines} |
| self.outputs = {rows: buf[:, :3] for rows, buf in self.slot_buffers.items()} |
| self.combined_outputs = {rows: torch.empty(rows, self.H, dtype=torch.bfloat16, device=dev) |
| for rows in self.pipelines} |
|
|
| def forward_slots(self, xq: torch.Tensor, ids: torch.Tensor, rw: torch.Tensor, |
| h: torch.Tensor | None = None, out: torch.Tensor | None = None) -> torch.Tensor: |
| rows = xq.shape[0] |
| if rows not in self.pipelines: |
| raise ValueError(f"{self.name}: fused pipeline admits 1..20 rows, got {rows}") |
| result = self.outputs[rows] if out is None else out |
| self.pipelines[rows].forward_into(xq, ids, rw, result, self.status, self.workspaces[rows]) |
| return result |
|
|
|
|
| class KXSelectiveDense: |
| """Keep BF16 as baseline; use measured FP8 only at explicitly selected M. |
| |
| Selected FP8 calls allocate no tensor storage and make no GPU-dependent choice. |
| The legacy BF16 fallback retains its original graph-captured allocation policy. |
| """ |
|
|
| def __init__(self, baseline: NO.NativeDense, candidate: DenseFP8, rows: tuple[int, ...], |
| splits: int | None = None): |
| """``splits``: one split-K count for every selected M (batch-invariant mode: B1's); None keeps the |
| measured per-M count. Row tiles stay per M: they partition rows, not the K reduction.""" |
| from ..kernels.fp8dense import DenseFP8 |
|
|
| if not isinstance(candidate, DenseFP8): |
| raise TypeError("candidate must be DenseFP8") |
| self.baseline, self.candidate = baseline, candidate |
| self.name, self.kind = baseline.name, baseline.kind |
| self.w, self.sc = baseline.w, baseline.sc |
| self.in_features = candidate.groups * candidate.inputs |
| self.out_features = candidate.groups * candidate.outputs |
| self.configs = {m: candidate.configuration(m) for m in rows} |
| if splits is not None: |
| self.configs = {m: (splits, tiles) for m, (_, tiles) in self.configs.items()} |
| self.scratch = { |
| m: (torch.empty(splits, m, candidate.groups, candidate.outputs, |
| dtype=torch.float32, device=self.w.device) if splits > 1 else None) |
| for m, (splits, _) in self.configs.items() |
| } |
| self.outputs = {(m, f32): torch.empty(m, self.out_features, |
| dtype=torch.float32 if f32 else torch.bfloat16, device=self.w.device) |
| for m in rows for f32 in (False, True)} |
|
|
| def nbytes(self) -> int: |
| return self.baseline.nbytes() + sum(t.numel() * t.element_size() |
| for t in (self.candidate.codes, self.candidate.scales)) |
|
|
| def __call__(self, xq: torch.Tensor, out_f32: bool = False, |
| out: torch.Tensor | None = None) -> torch.Tensor: |
| rows = xq.shape[0] |
| if out is not None and out.dtype != (torch.float32 if out_f32 else torch.bfloat16): |
| raise ValueError(f"{self.name}: caller output dtype disagrees with out_f32") |
| if rows not in self.configs: |
| if self.candidate.groups != 1: |
| if out_f32 or out is not None: |
| raise ValueError("grouped wo_a baseline returns BF16 without caller output") |
| groups, inputs, outputs = self.candidate.groups, self.candidate.inputs, self.candidate.outputs |
| z = torch.bmm(xq.view(rows, groups, inputs).transpose(0, 1), |
| self.w.view(groups, outputs, inputs).transpose(1, 2)) |
| return z.transpose(0, 1).reshape(rows, -1) |
| return self.baseline(xq, out_f32=out_f32, out=out) |
| result = self.outputs[rows, out_f32] if out is None else out |
| splits, tiles = self.configs[rows] |
| self.candidate.linear_into(xq, result, scratch=self.scratch[rows], splits=splits, row_tiles=tiles) |
| return result |
|
|
|
|
| class KXSharedPair: |
| """Measured T20 BF16 shared gate/up pair; smaller shapes retain FP8 calls.""" |
|
|
| def __init__(self, gate: KXSelectiveDense, up: KXSelectiveDense): |
| if gate.kind != "bf16" or up.kind != "bf16" or gate.w.shape != up.w.shape or gate.w.device != up.w.device: |
| raise ValueError("shared pair requires matching BF16 baseline matrices") |
| self.gate, self.up = gate, up |
| self.weight = torch.cat((gate.w, up.w), dim=0) |
| self.output = torch.empty(20, self.weight.shape[0], dtype=torch.bfloat16, device=self.weight.device) |
| self.projections = (self.output[:, :gate.out_features], self.output[:, gate.out_features:]) |
|
|
| def nbytes(self) -> int: |
| return self.weight.numel() * self.weight.element_size() |
|
|
| def __call__(self, xq: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| if xq.shape[0] != 20: |
| return self.gate(xq), self.up(xq) |
| if torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction: |
| raise RuntimeError("paired shared BF16 reduced-precision reduction is enabled") |
| torch.mm(xq, self.weight.t(), out=self.output) |
| return self.projections |
|
|
