| """Preallocated native-expanded components; setup is separate from graph-safe execution.""" |
| from __future__ import annotations |
|
|
| from types import ModuleType |
| from typing import Literal |
|
|
| import torch |
|
|
| Mode = Literal["grouped", "small", "auto"] |
| DenseBackend = Literal["cutlass", "cublas", "auto"] |
|
|
|
|
| def dense_backend(groups: int, rows: int, n: int, k: int, dtype: torch.dtype) -> Literal["cutlass", "cublas"]: |
| """Frozen measured dispatch, no startup tuning; unmeasured shapes use CUTLASS.""" |
| if dtype == torch.bfloat16: |
| if groups == 1 and (n, k) == (5120, 15360) and rows in (8, 32, 128): |
| return "cublas" |
| if rows in (5, 10, 15, 20): |
| if groups == 1 and (n, k) in ((1280, 5120), (512, 5120), (5120, 4096)): |
| return "cublas" |
| if groups == 4 and (n, k) == (1024, 4096): |
| return "cublas" |
| return "cutlass" |
|
|
|
|
| def _weight(w: torch.Tensor, shape: tuple[int, ...], dtype: torch.dtype, device: torch.device) -> None: |
| if tuple(w.shape) != shape or w.dtype != dtype or w.device != device or not w.is_contiguous(): |
| raise ValueError(f"expected contiguous {shape} {dtype} on {device}; got {w.shape} {w.dtype} {w.device}") |
|
|
|
|
| def _input(x: torch.Tensor, shape: tuple[int, ...], dtype: torch.dtype, device: torch.device) -> None: |
| _weight(x, shape, dtype, device) |
|
|
|
|
| class ExpandedExperts: |
| """A single TP channel-shard partial. Output/status/scratch are overwritten on next call. |
| |
| Inputs/weights must remain alive until all captured graphs using this object are destroyed. |
| One instance per in-flight stream. Router weights are never narrowed to FP16. |
| """ |
|
|
| def __init__(self, ext: ModuleType, gate: torch.Tensor, up: torch.Tensor, down: torch.Tensor, |
| rows: int, top_k: int = 3, limit: float = 10.0, mode: Mode = "auto", trained_qdq: bool = True, input_is_qdq: bool = False) -> None: |
| if gate.ndim != 3 or gate.dtype not in (torch.float16, torch.bfloat16) or not gate.is_cuda: |
| raise ValueError("gate must be CUDA BF16/FP16 [E,J,H]") |
| e, j, h = gate.shape |
| if not 1 <= e <= 128 or min(h, j) <= 0 or h % 8 or j % 8 or rows <= 0 or top_k <= 0: |
| raise ValueError("E must be 1..128; positive H/J must be multiples of 8; rows/top_k positive") |
| if mode == "auto": |
| mode = "small" if (h, j, top_k) == (5120, 1152, 3) and rows <= 3 else "grouped" |
| if mode not in ("grouped", "small") or not 0 < limit < float("inf"): |
| raise ValueError("unknown mode or invalid SwiGLU limit") |
| self.ext, self.dtype, self.device = ext, gate.dtype, gate.device |
| _weight(gate, (e, j, h), self.dtype, self.device) |
| _weight(up, (e, j, h), self.dtype, self.device) |
| _weight(down, (e, h, j), self.dtype, self.device) |
| self.gate, self.up, self.down = gate, up, down |
| self.rows, self.top_k, self.limit, self.mode = rows, top_k, limit, mode |
| if trained_qdq and (top_k != 3 or h % 32 or j % 32): |
| raise ValueError("native QDQ requires top-3 and 32-aligned H/J") |
| self.trained_qdq, self.input_is_qdq = trained_qdq, input_is_qdq |
| self.h, self.j = h, j |
| r, g = rows * top_k, min(e, rows * top_k) |
| def new(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor: |
| return torch.empty(shape, dtype=dtype, device=self.device) |
| self.state = [new((r, h), self.dtype), new((r,), torch.int32), new((r,), torch.int32), |
| new((e,), torch.int32), new((e,), torch.int32), new((1,), torch.int32), |
| new((2*g, 3), torch.int32), new((g, 3), torch.int32), |
| new((4, 2*g), torch.int64), new((4, g), torch.int64), |
| new((4, 2*g), torch.int64), new((4, g), torch.int64), |
| new((r, j), torch.float32), new((r, j), torch.float32), |
| new((r, j), self.dtype), new((r, h), torch.float32)] |
| self.qx = new((rows, h), self.dtype) |
| self.qa = new((r, j), self.dtype) |
| self.output = new((rows, h), torch.float32) |
| self.slot_output = new((rows, top_k, h), torch.float32) |
| self.status = self.state[5] |
| self.nblocks = ext.blocks(gate) |
|
|
| def project(self, x: torch.Tensor, ids: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: |
| _input(x, (self.rows, self.h), self.dtype, self.device) |
| _input(ids, (self.rows, self.top_k), torch.int32, self.device) |
| _input(weights, (self.rows, self.top_k), torch.float32, self.device) |
| s = self.state |
| if self.trained_qdq and not self.input_is_qdq: |
| self.ext.qdq(x, self.qx) |
| x = self.qx |
| if self.mode == "grouped": |
| self.ext.prepare(x, ids, self.gate, self.up, self.down, s) |
| self.ext.gemm(x, s[6], s[8], s[10], self.nblocks) |
| self._activate(weights) |
| self.ext.gemm(x, s[7], s[9], s[11], self.nblocks) |
| else: |
| self.ext.small_prepare(ids, s[2], s[1], s[5], self.gate.shape[0]) |
| self.ext.small(x, self.gate, self.up, ids, s[12], s[13], True) |
| self._activate(weights) |
| self.ext.small(s[14], self.down, self.down, ids, s[15], s[15], False) |
| return s[15] |
|
|
| def run(self, x: torch.Tensor, ids: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: |
| self.project(x, ids, weights) |
| s = self.state |
| if self.trained_qdq: |
| self.ext.native_combine(s[15], ids, s[2], s[5], self.output, False) |
| else: |
| self.ext.combine(s[15], s[2], weights, s[5], self.output) |
| return self.output |
|
|
| def run_slots(self, x: torch.Tensor, ids: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: |
| """FP32 [T,3,H] in original route-slot order, ready for N1's exact TP2 epilogue.""" |
| self.project(x, ids, weights) |
| self.ext.slot_scatter(self.state[15], self.state[2], self.status, self.slot_output) |
| return self.slot_output |
|
|
| def _activate(self, weights: torch.Tensor) -> None: |
| s = self.state |
| if self.trained_qdq: |
| self.ext.native_activate(s[12], s[13], s[1], weights, self.qa, self.limit) |
| self.ext.qdq(self.qa, s[14]) |
| else: |
| self.ext.activate(s[12], s[13], s[14], self.limit) |
|
|
| def combine_after_tp(self, ids: torch.Tensor) -> torch.Tensor: |
| """Call after the integrator sums BOTH ranks' per-route state[15] FP32 partials. |
| |
| This is the native BF16 down-output boundary, before ascending-expert FP32 sum. |
| run() alone returns an unrounded shard partial for diagnostics. |
| """ |
| if not self.trained_qdq: |
| raise RuntimeError("combine_after_tp requires the native arithmetic profile") |
| _input(ids, (self.rows, self.top_k), torch.int32, self.device) |
| self.ext.native_combine(self.state[15], ids, self.state[2], self.status, self.output, True) |
| return self.output |
|
|
| @property |
| def workspace_bytes(self) -> int: |
| return sum(t.numel() * t.element_size() for t in (*self.state, self.qx, self.qa, self.output, self.slot_output)) |
|
|
|
|
| class Dense: |
| """Uniform grouped GEMM, [G,M,K] @ [G,N,K]^T -> FP32 [G,M,N].""" |
|
|
| def __init__(self, ext: ModuleType, weight: torch.Tensor, rows: int, |
| backend: DenseBackend = "auto") -> None: |
| if weight.ndim != 3 or not weight.is_cuda or weight.dtype not in (torch.float16, torch.bfloat16): |
| raise ValueError("weight must be CUDA BF16/FP16 [groups,N,K]") |
| g, n, k = weight.shape |
| if min(g, n, k, rows) <= 0 or n % 8 or k % 8: |
| raise ValueError("positive dimensions and 8-aligned K/N required") |
| _weight(weight, (g, n, k), weight.dtype, weight.device) |
| if backend == "auto": |
| backend = dense_backend(g, rows, n, k, weight.dtype) |
| if backend not in ("cutlass", "cublas"): |
| raise ValueError("unknown dense backend") |
| if backend == "cublas": |
| reduced = (torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction |
| if weight.dtype == torch.bfloat16 else |
| torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction) |
| if reduced: |
| raise ValueError("cuBLAS FP32 profile requires reduced_precision_reduction=False at setup") |
| self.backend = backend |
| self.ext, self.weight, self.rows = ext, weight, rows |
| self.input = torch.empty((g, rows, k), dtype=weight.dtype, device=weight.device) |
| self.output = torch.empty((g, rows, n), dtype=torch.float32, device=weight.device) |
| self.sizes = torch.empty((g, 3), dtype=torch.int32, device=weight.device) |
| self.ptr = torch.empty((4, g), dtype=torch.int64, device=weight.device) |
| self.ld = torch.empty_like(self.ptr) |
| self.nblocks = ext.blocks(weight) |
| ext.dense_prepare(self.input, weight, self.output, self.sizes, self.ptr, self.ld, False) |
| self.wt = weight.transpose(1, 2) |
| self.a0, self.b0, self.c0 = self.input[0], self.wt[0], self.output[0] |
|
|
| def run(self, x: torch.Tensor) -> torch.Tensor: |
| _input(x, tuple(self.input.shape), self.weight.dtype, self.weight.device) |
| self.input.copy_(x) |
| if self.backend == "cublas": |
| if self.input.shape[0] == 1: |
| torch.mm(self.a0, self.b0, out_dtype=torch.float32, out=self.c0) |
| else: |
| torch.bmm(self.input, self.wt, out_dtype=torch.float32, out=self.output) |
| else: |
| self.ext.gemm(self.input, self.sizes, self.ptr, self.ld, self.nblocks) |
| return self.output |
|
|
|
|
| class GroupedWoA: |
| """Grouped attention output projection in the model's flat [T,G*C] layout. |
| |
| Avoids the transpose/contiguous passes needed by grouped bmm. Weights [G,N,K] |
| are output-major, output [T,G*N] is FP32, with native BF16 rounding caller-owned. |
| """ |
|
|
| def __init__(self, ext: ModuleType, weight: torch.Tensor, rows: int, |
| backend: DenseBackend = "auto") -> None: |
| if weight.ndim != 3 or not weight.is_cuda or weight.dtype not in (torch.float16, torch.bfloat16): |
| raise ValueError("weight must be CUDA BF16/FP16 [groups,N,K]") |
| g, n, k = weight.shape |
| if min(g, n, k, rows)<=0 or n%8 or k%8: |
| raise ValueError("positive dimensions and 8-aligned K/N required") |
| _weight(weight, (g, n, k), weight.dtype, weight.device) |
| if backend == "auto": |
| backend = dense_backend(g, rows, n, k, weight.dtype) |
| if backend not in ("cutlass", "cublas"): |
| raise ValueError("unknown dense backend") |
| if backend == "cublas": |
| reduced = (torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction |
| if weight.dtype == torch.bfloat16 else |
| torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction) |
| if reduced: |
| raise ValueError("cuBLAS FP32 profile requires reduced_precision_reduction=False at setup") |
| self.backend = backend |
| self.ext, self.weight, self.rows = ext, weight, rows |
| self.input = torch.empty((rows, g*k), dtype=weight.dtype, device=weight.device) |
| self.output = torch.empty((rows, g*n), dtype=torch.float32, device=weight.device) |
| self.sizes = torch.empty((g, 3), dtype=torch.int32, device=weight.device) |
| self.ptr = torch.empty((4, g), dtype=torch.int64, device=weight.device) |
| self.ld = torch.empty_like(self.ptr) |
| self.nblocks = ext.blocks(weight) |
| ext.dense_prepare(self.input, weight, self.output, self.sizes, self.ptr, self.ld, True) |
| self.av = self.input.view(rows, g, k).transpose(0, 1) |
| self.cv = self.output.view(rows, g, n).transpose(0, 1) |
| self.wt = weight.transpose(1, 2) |
|
|
| def run(self, x: torch.Tensor) -> torch.Tensor: |
| _input(x, tuple(self.input.shape), self.weight.dtype, self.weight.device) |
| self.input.copy_(x) |
| if self.backend == "cublas": |
| torch.bmm(self.av, self.wt, out_dtype=torch.float32, out=self.cv) |
| else: |
| self.ext.gemm(self.input, self.sizes, self.ptr, self.ld, self.nblocks) |
| return self.output |
|
|
|
|
| class HeadEpilogue: |
| """Local vocab winner for an existing TP2 winner collective. Does not implement confidence. |
| |
| Native run() adds FP32 base and exact FP32 Markov bias without BF16 rounding. |
| Equal logits select the lowest global ID. BF16 engine arithmetic is an explicit |
| run_engine_bf16_control() operation only. No bias is dropped. |
| Nonfinite candidates propagate a NaN winner value, requiring caller failure. |
| """ |
|
|
| def __init__(self, ext: ModuleType, rows: int, vocab: int, device: torch.device, offset: int = 0) -> None: |
| if rows <= 0 or vocab <= 0 or offset < 0: |
| raise ValueError("positive rows/vocab and nonnegative offset required") |
| self.ext, self.rows, self.vocab, self.device, self.offset = ext, rows, vocab, device, offset |
| chunks = (vocab+1023)//1024 |
| self.values = torch.empty((rows, chunks), dtype=torch.float32, device=device) |
| self.ids = torch.empty((rows, chunks), dtype=torch.int64, device=device) |
| self.winner = torch.empty((rows,), dtype=torch.int64, device=device) |
| self.value = torch.empty((rows,), dtype=torch.float32, device=device) |
|
|
| def run(self, base: torch.Tensor, bias: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| return self._run(base, bias, engine_bf16_control=False) |
|
|
| def run_engine_bf16_control(self, base: torch.Tensor, bias: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| """Engine comparison only: BF16(base), BF16(bias), BF16(sum), then argmax.""" |
| return self._run(base, bias, engine_bf16_control=True) |
|
|
| def _run(self, base: torch.Tensor, bias: torch.Tensor, *, engine_bf16_control: bool) -> tuple[torch.Tensor, torch.Tensor]: |
| for t in (base, bias): |
| _input(t, (self.rows, self.vocab), torch.float32, self.device) |
| head = self.ext.head_engine_bf16_control if engine_bf16_control else self.ext.head |
| head(base, bias, self.values, self.ids, self.offset) |
| self.ext.head_reduce(self.values, self.ids, self.value, self.winner) |
| return self.value, self.winner |
|
|
