| """Training-only tensor-core products with FP32 outputs and frozen-weight dX. |
| |
| Packed routed weights are decoded per tile in registers. No decoded expert |
| bank, CPU routing list, or host activation copy is used. This backend has its |
| own agreement receipt; it never replaces the qualified reference evaluator. |
| """ |
| from __future__ import annotations |
|
|
| from typing import Any, Callable, cast |
|
|
| import torch |
| from torch import Tensor, nn |
| import triton |
| import triton.language as tl |
| from triton.language.extra import libdevice |
|
|
| from .ops import PackedLinear, check_cancelled |
|
|
|
|
| @triton.jit |
| def _qdq(X, Y, N: tl.constexpr, GROUPS: tl.constexpr): |
| g = tl.program_id(0) * GROUPS + tl.arange(0, GROUPS) |
| k = tl.arange(0, 32) |
| v = tl.load(X + g[:, None] * 32 + k[None, :], g[:, None] * 32 + k[None, :] < N, 0).to(tl.float32) |
| a = tl.maximum(tl.max(tl.abs(v), 1), 1.e-4) * (1. / 448.) |
| bits = a.to(tl.int32, bitcast=True) |
| e = ((bits >> 23) & 255) - 127 + ((bits & 0x7fffff) != 0).to(tl.int32) |
| scale = ((e + 127) << 23).to(tl.float32, bitcast=True) |
| q = tl.minimum(tl.maximum(v / scale[:, None], -448.), 448.) |
| qb = q.to(tl.int32, bitcast=True) |
| qe = tl.maximum(((qb >> 23) & 255) - 127, -6) |
| quantum = ((qe - 3 + 127) << 23).to(tl.float32, bitcast=True) |
| q = libdevice.rint(q / quantum) * quantum |
| tl.store(Y + g[:, None] * 32 + k[None, :], q * scale[:, None], g[:, None] * 32 + k[None, :] < N) |
|
|
|
|
| class _QDQ(torch.autograd.Function): |
| @staticmethod |
| def forward(ctx: Any, x: Tensor) -> Tensor: |
| x = x.contiguous() |
| if x.shape[-1] % 32: |
| raise ValueError("fast QDQ requires a multiple of 32 channels") |
| y = torch.empty_like(x) |
| _qdq[(triton.cdiv(x.numel(), 512),)](x, y, x.numel(), GROUPS=16, num_warps=4, |
| enable_fp_fusion=False) |
| return y |
|
|
| @staticmethod |
| def backward(ctx: Any, grad: Tensor) -> Tensor: |
| return grad |
|
|
|
|
| def fast_qdq(x: Tensor) -> Tensor: |
| return cast(Callable[[Tensor], Tensor], _QDQ.apply)(x) |
|
|
|
|
| @triton.jit |
| def _decode(W, S, e, n, k, N: tl.constexpr, K: tl.constexpr): |
| packed = tl.load(W + e.to(tl.int64) * N * (K // 2) + n * (K // 2) + k // 2, |
| (n < N) & (k < K), 0).to(tl.int32) |
| code = (packed >> (k % 2 * 4)) & 15 |
| mag = code & 7 |
| val = tl.where(mag < 2, mag.to(tl.float32) * .5, |
| (2 + (mag & 1)).to(tl.float32) * tl.exp2(((mag >> 1) - 2).to(tl.float32))) |
| val = tl.where((code & 8) != 0, -val, val) |
| scale = tl.load(S + e.to(tl.int64) * N * (K // 32) + n * (K // 32) + k // 32, |
| (n < N) & (k < K), 127).to(tl.int32) |
| return val * tl.exp2((scale - 127).to(tl.float32)) |
|
|
|
|
| @triton.jit |
| def _grouped(X, W, S, OFF, TILES, Y, E: tl.constexpr, N: tl.constexpr, K: tl.constexpr, |
| TRANSPOSE: tl.constexpr, BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, |
| EP: tl.constexpr): |
| task = tl.program_id(0) |
| es = tl.arange(0, EP) |
| ends = tl.load(TILES + es, es < E, 2147483647) |
| e = tl.sum((ends <= task).to(tl.int32), 0) |
| if e >= E: |
| return |
| first_tile = tl.load(TILES + e - 1, e > 0, 0) |
| lo = tl.load(OFF + e) |
| hi = tl.load(OFF + e + 1) |
| m = lo + (task - first_tile) * BM + tl.arange(0, BM) |
| out = tl.program_id(1) * BN + tl.arange(0, BN) |
| r = tl.arange(0, BK) |
| acc = tl.zeros((BM, BN), tl.float32) |
| if TRANSPOSE: |
| for start in range(tl.cdiv(N, BK)): |
| n = start * BK + r |
| a = tl.load(X + m[:, None] * N + n[None, :], (m[:, None] < hi) & (n[None, :] < N), 0) |
| b = _decode(W, S, e, n[:, None], out[None, :], N, K) |
| acc += tl.dot(a.to(tl.float32), b, input_precision="tf32x3") |
| tl.store(Y + m[:, None] * K + out[None, :], acc, (m[:, None] < hi) & (out[None, :] < K)) |
| else: |
| for start in range(tl.cdiv(K, BK)): |
| k = start * BK + r |
| a = tl.load(X + m[:, None] * K + k[None, :], (m[:, None] < hi) & (k[None, :] < K), 0) |
| b = _decode(W, S, e, out[None, :], k[:, None], N, K) |
| acc += tl.dot(a.to(tl.bfloat16), b.to(tl.bfloat16)) |
| tl.store(Y + m[:, None] * N + out[None, :], acc, (m[:, None] < hi) & (out[None, :] < N)) |
|
|
|
|
| def grouped_product(x: Tensor, weight: Tensor, scale: Tensor, offsets: Tensor, tiles: Tensor, |
| transpose: bool = False) -> Tensor: |
| check_cancelled() |
| e, n, k2 = weight.shape |
| k = k2 * 2 |
| x = x.contiguous() |
| width = k if transpose else n |
| out = torch.empty(x.shape[0], width, device=x.device, dtype=torch.float32) |
| bm = 16 |
| |
| grid = (triton.cdiv(x.shape[0], bm) + e, triton.cdiv(width, 64)) |
| _grouped[grid](x, weight, scale, offsets, tiles, out, e, n, k, transpose, |
| BM=bm, BN=64, BK=32, EP=triton.next_power_of_2(e), num_warps=4) |
| return out |
|
|
|
|
| class _Grouped(torch.autograd.Function): |
| @staticmethod |
| def forward(ctx: Any, x: Tensor, weight: Tensor, scale: Tensor, offsets: Tensor, tiles: Tensor) -> Tensor: |
| ctx.save_for_backward(weight, scale, offsets, tiles) |
| return grouped_product(x, weight, scale, offsets, tiles) |
|
|
| @staticmethod |
| def backward(ctx: Any, grad: Tensor) -> tuple[Tensor, None, None, None, None]: |
| w, s, offsets, tiles = ctx.saved_tensors |
| return grouped_product(grad, w, s, offsets, tiles, transpose=True), None, None, None, None |
|
|
|
|
| class ExpertBank(nn.Module): |
| """Packed banks share storage with the original per-expert buffer views.""" |
|
|
| def __init__(self, experts: nn.ModuleList): |
| super().__init__() |
| for name in ("w1", "w2", "w3"): |
| layers = [cast(PackedLinear, cast(nn.ModuleDict, ex)[name]) for ex in experts] |
| ws = torch.stack([ex.weight for ex in layers]) |
| ss = torch.stack([ex.scale for ex in layers]) |
| self.register_buffer(name + "_weight", ws) |
| self.register_buffer(name + "_scale", ss) |
| for i, ex in enumerate(layers): |
| ex.weight, ex.scale = ws[i], ss[i] |
|
|
| def product(self, name: str, x: Tensor, offsets: Tensor, tiles: Tensor) -> Tensor: |
| return cast(Callable[..., Tensor], _Grouped.apply)( |
| x, getattr(self, name + "_weight"), getattr(self, name + "_scale"), offsets, tiles) |
|
|
|
|
| def linear_product(x: Tensor, w: Tensor, quantized: bool) -> Tensor: |
| check_cancelled() |
| flat = x.reshape(-1, x.shape[-1]) |
| if quantized: |
| |
| |
| out = torch.mm(flat.bfloat16(), w.bfloat16().T, out_dtype=torch.float32) |
| else: |
| out = torch.mm(flat.float(), w.float().T) |
| return out.reshape(*x.shape[:-1], w.shape[0]) |
|
|
|
|
| class _Linear(torch.autograd.Function): |
| @staticmethod |
| def forward(ctx: Any, x: Tensor, w: Tensor, quantized: bool) -> Tensor: |
| ctx.save_for_backward(w) |
| return linear_product(x, w, quantized) |
|
|
| @staticmethod |
| def backward(ctx: Any, grad: Tensor) -> tuple[Tensor, None, None]: |
| (w,) = ctx.saved_tensors |
| check_cancelled() |
| dx = torch.mm(grad.reshape(-1, grad.shape[-1]).float(), w.float()) |
| return dx.reshape(*grad.shape[:-1], w.shape[1]), None, None |
|
|
|
|
| def fast_linear(x: Tensor, w: Tensor, quantized: bool = False) -> Tensor: |
| return cast(Callable[..., Tensor], _Linear.apply)(x, w, quantized) |
|
|
