File size: 10,020 Bytes
e1acc15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | """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
|