File size: 2,187 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 | """Packed MXFP4 forward, decoded only in registers, FP64 reference sums.
Backward alone materializes one selected expert matrix in BF16. This kernel
does not expand a matrix/bank in device memory and has no expert gradients.
"""
from __future__ import annotations
import torch
from torch import Tensor
import triton # type: ignore[import-untyped]
import triton.language as tl # type: ignore[import-untyped]
@triton.jit # type: ignore[untyped-decorator]
def _packed(X, W, S, Y, K: tl.constexpr, N: tl.constexpr, # type: ignore[no-untyped-def]
BN: tl.constexpr, BK: tl.constexpr) -> None:
m = tl.program_id(0)
n = tl.program_id(1)*BN + tl.arange(0, BN)
kk = tl.arange(0, BK)
total = tl.full((BN,), 0., tl.float64)
for start in range(tl.cdiv(K, BK)):
k = start*BK+kk
a = tl.load(X+m*K+k, k < K, 0).to(tl.float32)
packed = tl.load(W+n[:, None]*(K//2)+k[None, :]//2,
(n[:, None] < N) & (k[None, :] < K), 0).to(tl.int32)
code = (packed >> ((k[None, :] % 2)*4)) & 15
mag = code & 7
value = tl.where(mag < 2, mag.to(tl.float32)*.5,
(2+(mag & 1)).to(tl.float32)*tl.exp2(((mag >> 1)-2).to(tl.float32)))
value = tl.where((code & 8) != 0, -value, value)
scale = tl.load(S+n[:, None]*(K//32)+k[None, :]//32,
(n[:, None] < N) & (k[None, :] < K), 127).to(tl.int32)
weight = value*tl.exp2((scale-127).to(tl.float32))
# QDQ activations and FP4 values have <=4 and <=3 significant bits;
# their product is exact in FP32. Only the sum needs FP64.
total += tl.sum((a[None, :]*weight).to(tl.float64), 1)
tl.store(Y+m*N+n, total.to(tl.float32), n < N)
def packed_matmul(x: Tensor, weight: Tensor, scale: Tensor) -> Tensor:
K = x.shape[-1]
n = weight.shape[0]
x2 = x.reshape(-1, K).contiguous()
out = torch.empty(x2.shape[0], n, device=x.device, dtype=torch.float32)
_packed[(x2.shape[0], triton.cdiv(n, 4))](x2, weight, scale, out, K, n, BN=4, BK=128,
num_warps=4, enable_fp_fusion=False)
return out.reshape(*x.shape[:-1], n)
|