forked from tinygrad/tinygrad
508 lines
29 KiB
Python
508 lines
29 KiB
Python
import atexit, functools, math, pathlib
|
|
from tinygrad import Tensor, Device, dtypes
|
|
from tinygrad.dtype import AddrSpace
|
|
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
|
|
from tinygrad.renderer import Estimates
|
|
from tinygrad.helpers import getenv, all_same, DEBUG, ceildiv
|
|
from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
|
|
from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
|
|
|
|
TILE_M, TILE_N, TILE_K = 256, 256, 64
|
|
|
|
# ** FP8 GEMM custom kernel
|
|
|
|
@functools.cache
|
|
def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str, scale_mode:int=3) -> UOp:
|
|
# scale_mode: 0=no scale, 1=x only, 2=w only, 3=both
|
|
n_scales = (1 if scale_mode & 1 else 0) + (1 if scale_mode & 2 else 0) + (1 if scale_mode & 4 else 0)
|
|
scales, extra = args[:n_scales], args[n_scales:]
|
|
M, K = A.shape[0]*A.shape[1], A.shape[2]
|
|
N, K2 = B.shape[(1 if B.ndim == 3 else 0):]
|
|
assert K == K2, f"{A.shape} {B.shape}"
|
|
block_size = 256
|
|
threads = UOp.special(64 * 8, "lidx0")
|
|
workgroups = UOp.special((M // block_size) * (N // block_size), "gidx0")
|
|
sink_inputs = (C.base, A.base, B.base) + tuple(s.base for s in scales) + (threads, workgroups)
|
|
sink = UOp.sink(*sink_inputs,
|
|
arg=KernelInfo(f"hk_fp8_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
|
|
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
|
|
src = (kittens_path/"gemm_fp8.cpp").read_text()
|
|
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
|
|
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
|
|
f"-DSCALE_MODE={scale_mode}"]).compile_cached(src)
|
|
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
|
|
UOp(Ops.BINARY, arg=lib)))
|
|
|
|
# ** FP8 AtB GEMM custom kernel
|
|
|
|
@functools.cache
|
|
def custom_hk_fp8_atb_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str, scale_mode:int=5) -> UOp:
|
|
# C = A.T @ B, A and B are physically [K, M] and [K, N].
|
|
n_scales = (1 if scale_mode & 1 else 0) + (1 if scale_mode & 2 else 0) + (1 if scale_mode & 4 else 0)
|
|
scales = args[:n_scales]
|
|
K, M = A.shape[0]*A.shape[1], A.shape[2]
|
|
K2, N = B.shape[0]*B.shape[1], B.shape[2]
|
|
assert K == K2, f"{A.shape} {B.shape}"
|
|
block_m, block_n, block_k, num_warps = 256, 256, 128, 8
|
|
assert M % block_m == 0 and N % block_n == 0 and K % block_k == 0, f"invalid fp8 atb tile {(block_m, block_n, block_k)} for {(M, N, K)}"
|
|
threads = UOp.special(64 * num_warps, "lidx0")
|
|
workgroups = UOp.special((M // block_m) * (N // block_n), "gidx0")
|
|
sink_inputs = (C.base, A.base, B.base) + tuple(s.base for s in scales) + (threads, workgroups)
|
|
sink = UOp.sink(*sink_inputs,
|
|
arg=KernelInfo(f"hk_fp8_atb_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
|
|
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
|
|
src = (kittens_path/"gemm_fp8_atb.cpp").read_text()
|
|
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
|
|
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
|
|
f"-DSCALE_MODE={scale_mode}"]).compile_cached(src)
|
|
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
|
|
UOp(Ops.BINARY, arg=lib)))
|
|
|
|
def hk_fp8_atb_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, g_amax:Tensor|None=None) -> Tensor:
|
|
assert a.dtype == b.dtype == FP8_DTYPE, f"expected fp8, got {a.dtype} {b.dtype}"
|
|
assert a.ndim == b.ndim == 3 and a.shape[:2] == b.shape[:2], f"{a.shape} {b.shape}"
|
|
batch, rows, M = a.shape
|
|
N = b.shape[2]
|
|
assert M % TILE_M == 0 and N % TILE_N == 0 and (batch * rows) % 128 == 0, \
|
|
f"fp8 atb shape {a.shape} {b.shape} must produce (M,N,K) multiples of ({TILE_M},{TILE_N},128)"
|
|
is_multi = isinstance(a.device, tuple)
|
|
reduce_out = False
|
|
if is_multi:
|
|
ndev = len(a.device)
|
|
if a.uop.axis in (0, 1) or b.uop.axis in (0, 1): inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device), 0, True
|
|
elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=dtypes.bfloat16, device=a.device), 2
|
|
elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=dtypes.bfloat16, device=a.device), 1
|
|
else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device), 0, True
|
|
out = Tensor(inv.uop.unshard(out_axis), device=a.device)
|
|
dname = a.device[0]
|
|
else:
|
|
out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device)
|
|
dname = a.device
|
|
dname = dname.split(":")[0]
|
|
scales = tuple(s for s in (x_scale, g_amax) if s is not None)
|
|
scale_mode = (1 if x_scale is not None else 0) | (4 if g_amax is not None else 0)
|
|
out = Tensor.custom_kernel(out, a, b, *scales, fxn=functools.partial(custom_hk_fp8_atb_gemm, dname=dname, scale_mode=scale_mode))[0]
|
|
if reduce_out: out = out.sum(0)
|
|
return out.squeeze(0) if out.ndim == 3 else out
|
|
|
|
# ** MXFP8 GEMM custom kernel
|
|
|
|
@functools.cache
|
|
def custom_hk_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:UOp, dname:str) -> UOp:
|
|
# mxfp8 block-scaled gemm: A(M,K) @ B(N,K).T, e8m0 1x32 microscales packed (k_iters,dim) uint32
|
|
M, K = A.shape[0]*A.shape[1], A.shape[2]
|
|
N, K2 = B.shape[(1 if B.ndim == 3 else 0):]
|
|
assert K == K2, f"{A.shape} {B.shape}"
|
|
block_size = 256
|
|
threads = UOp.special(64 * 8, "lidx0")
|
|
workgroups = UOp.special((M // block_size) * (N // block_size), "gidx0")
|
|
e_a = extra[0].base if len(extra) >= 1 else scale_A.base
|
|
e_b = extra[1].base if len(extra) >= 2 else scale_B.base
|
|
sink_inputs = (C.base, A.base, B.base, scale_A.base, scale_B.base, e_a, e_b, threads, workgroups)
|
|
sink = UOp.sink(*sink_inputs,
|
|
arg=KernelInfo(f"hk_mxfp8_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
|
|
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
|
|
src = (kittens_path/"gemm_mxfp8.cpp").read_text()
|
|
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
|
|
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}"]).compile_cached(src)
|
|
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
|
|
UOp(Ops.BINARY, arg=lib)))
|
|
|
|
# ** MXFP4 GEMM custom kernel
|
|
|
|
@functools.cache
|
|
def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp, tile_m:int, tile_n:int) -> UOp:
|
|
from extra.gemm.gemm_mxfp4 import build_kernel
|
|
M, half_k = math.prod(A.shape[:-1]), A.shape[-1]
|
|
N, half_k_b = math.prod(B.shape[:-1]), B.shape[-1]
|
|
K = half_k * 2
|
|
assert half_k == half_k_b and math.prod(C.shape[:-1]) == M and C.shape[-1] == N
|
|
threads = UOp.special(256, "lidx0")
|
|
groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
|
|
lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
|
|
sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
|
|
arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
|
|
insts = build_kernel(M, N, K, tile_m, tile_n)
|
|
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))
|
|
|
|
def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
|
|
# 1x32 block scaling along the last axis
|
|
*batch, K = x.shape
|
|
scale_K = K // 32
|
|
amax = x.detach().float().reshape(*batch, scale_K, 32).abs().max(axis=-1)
|
|
e8 = (amax.maximum(1e-38).log2().floor() + 127).clamp(0, 254).cast(dtypes.uint8)
|
|
qscale = (127.0 - e8.cast(dtypes.float32)).exp2().reshape(*batch, scale_K, 1).expand(*batch, scale_K, 32).reshape(*batch, K)
|
|
x_scaled = x.float() * qscale
|
|
x_clamped = x_scaled + (x_scaled.detach().clamp(-448.0, 448.0) - x_scaled.detach()) # STE
|
|
packed = mx_pack(e8) if len(batch) == 1 and scale_K % 4 == 0 else None
|
|
return x_clamped.cast(FP8_DTYPE), e8, packed
|
|
|
|
def _mxfp4_shuffle_weight(x:Tensor) -> Tensor:
|
|
# shuffle_weight(x, layout=(16, 16)) on the packed uint8 buffer.
|
|
if x.ndim == 3:
|
|
ndev, rows, half_k = x.shape
|
|
return x.reshape(ndev, rows//16, 16, half_k//32, 2, 16).permute(0, 1, 3, 4, 2, 5).reshape(ndev, rows, half_k).contiguous()
|
|
rows, half_k = x.shape
|
|
return x.reshape(rows//16, 16, half_k//32, 2, 16).permute(0, 2, 3, 1, 4).reshape(rows, half_k).contiguous()
|
|
|
|
def _mxfp4_shuffle_scales(x:Tensor) -> Tensor:
|
|
# e8m0_shuffle: each 256x8 scale tile is arranged for the raw MFMA scale loads.
|
|
if x.ndim == 3:
|
|
ndev, rows, scale_k = x.shape
|
|
return x.reshape(ndev, rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 1, 4, 6, 3, 5, 2).reshape(ndev, rows, scale_k).contiguous()
|
|
rows, scale_k = x.shape
|
|
return x.reshape(rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 3, 5, 2, 4, 1).reshape(rows, scale_k).contiguous()
|
|
|
|
def quantize_mxfp4(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
|
|
# OCP MXFP4: 1x32 blocks, E2M1 values packed low-nibble first, and E8M0 scales.
|
|
*batch, K = x.shape
|
|
rows = math.prod(batch)
|
|
assert x.ndim >= 2 and K % 256 == 0 and rows % 32 == 0, \
|
|
f"mxfp4 quantization needs rows%32 and K%256, got {x.shape}"
|
|
xb = x.float().reshape(*batch, K//32, 32)
|
|
amax = xb.abs().max(axis=-1)
|
|
|
|
# even scale rounding: round the fp32 significand before choosing 2^(floor(log2)-2).
|
|
amax_rounded = ((amax.bitcast(dtypes.uint32) + 0x200000) & 0xFF800000).bitcast(dtypes.float32)
|
|
scale_exp = (amax_rounded.maximum(2**-126).log2().floor() - 2).clamp(-127, 127)
|
|
e8 = (scale_exp + 127).cast(dtypes.uint8)
|
|
scaled = xb * (-scale_exp).exp2().reshape(*batch, K//32, 1)
|
|
|
|
mag = scaled.abs()
|
|
code = sum(x.cast(dtypes.uint8) for x in
|
|
(mag > .25, mag >= .75, mag > 1.25, mag >= 1.75, mag > 2.5, mag >= 3.5, mag > 5.0))
|
|
code = code | ((scaled < 0).cast(dtypes.uint8) << 3)
|
|
code = code.reshape(*batch, K)
|
|
packed = code[..., 0::2] | (code[..., 1::2] << 4)
|
|
if isinstance(x.device, tuple) and x.uop.axis == x.ndim-2 and x.shape[x.uop.axis] == len(x.device):
|
|
axis = x.uop.axis
|
|
order = (axis, *range(axis), *range(axis+1, e8.ndim))
|
|
e8_local = e8.permute(order)
|
|
return packed, e8, _mxfp4_shuffle_scales(e8_local.reshape(e8_local.shape[0], -1, K//32))
|
|
return packed, e8, _mxfp4_shuffle_scales(e8.reshape(rows, K//32))
|
|
|
|
def mx_pack(e8:Tensor) -> Tensor:
|
|
rows, scale_K = e8.shape
|
|
return e8.reshape(rows, scale_K // 4, 4).bitcast(dtypes.uint32).reshape(rows, scale_K // 4).permute(1, 0).contiguous()
|
|
|
|
def _mx_block_scale(e8:Tensor) -> Tensor:
|
|
# dequant scale 2^(e8-127) broadcast back to element shape
|
|
rows, scale_K = e8.shape
|
|
return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(rows, scale_K, 1).expand(rows, scale_K, 32).reshape(rows, scale_K*32)
|
|
|
|
def _mx_block_scale_3d(e8:Tensor) -> Tensor:
|
|
# batched (E, rows, scale_K) dequant scale 2^(e8-127) broadcast to (E, rows, scale_K*32)
|
|
E, rows, scale_K = e8.shape
|
|
return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(E, rows, scale_K, 1).expand(E, rows, scale_K, 32).reshape(E, rows, scale_K*32)
|
|
|
|
counters = {"used":0, "todos":[]}
|
|
def todo(msg:str) -> bool: counters["todos"].append(msg); return False
|
|
def _asm_gemm_report():
|
|
print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used')
|
|
if DEBUG >= 2 and counters["todos"]:
|
|
from collections import Counter
|
|
for msg, cnt in Counter(counters["todos"]).most_common(): print(f' {cnt:3d}x {msg}')
|
|
atexit.register(_asm_gemm_report)
|
|
|
|
def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
|
|
if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
|
|
if a.dtype not in {dtypes.bfloat16, dtypes.float16, FP8_DTYPE}: return todo(f"only bfloat16/float16/fp8, got {a.dtype}")
|
|
batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
|
|
N = b.shape[1]
|
|
if isinstance(a.device, tuple):
|
|
if a.ndim == 2 and a.uop.axis == 0 and b.uop.axis is None: M //= len(a.device)
|
|
elif a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
|
|
elif a.ndim == 2 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
|
|
elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
|
|
elif a.ndim == 3 and a.uop.axis == 1 and b.uop.axis is None: M //= len(a.device)
|
|
elif a.ndim == 3 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
|
|
elif a.ndim == 3 and a.uop.axis == 2 and b.uop.axis == 0: K //= len(a.device)
|
|
else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
|
|
dname = a.device[0]
|
|
else: dname = a.device
|
|
arch = Device[dname].renderer.target.arch
|
|
if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
|
|
if (M % TILE_M != 0 or N % TILE_N != 0 or K % TILE_K != 0) and arch == "gfx950":
|
|
return todo(f"GEMM shape ({M},{N},{K}) not a multiple of ({TILE_M},{TILE_N},{TILE_K})")
|
|
return True
|
|
|
|
# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
|
|
# note: this can be removed after we have GEMM on mixins
|
|
|
|
def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
|
|
M, K = A.shape[0]*A.shape[1], A.shape[2]
|
|
K2, N = B.shape[(1 if B.ndim == 3 else 0):]
|
|
assert K == K2
|
|
m = UOp.range(M, 1)
|
|
n = UOp.range(N, 2)
|
|
k = UOp.range(K, 0, AxisType.REDUCE)
|
|
mul = (A.flatten().index((m*UOp.const(K)+k))*
|
|
B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
|
|
red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
|
|
store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n)
|
|
return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
|
|
|
|
# ** bf16 A @ B.T kernel in C
|
|
|
|
@functools.cache
|
|
def custom_hk_bf16_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str) -> UOp:
|
|
M, K = A.shape[0]*A.shape[1], A.shape[2]
|
|
N, K2 = B.shape[(1 if B.ndim == 3 else 0):]
|
|
assert K == K2, f"{A.shape} {B.shape}"
|
|
block_m, block_n, block_k, num_warps = 256, 256, 64, 8
|
|
assert M % block_m == 0 and N % block_n == 0 and K % block_k == 0, f"invalid bf16 tile {(block_m, block_n, block_k)} for {(M, N, K)}"
|
|
threads = UOp.special(64 * num_warps, "lidx0")
|
|
workgroups = UOp.special((M // block_m) * (N // block_n), "gidx0")
|
|
b_extra = args[0].base if len(args) >= 1 else B.base
|
|
sink = UOp.sink(C.base, A.base, B.base, b_extra, threads, workgroups,
|
|
arg=KernelInfo(f"hk_bf16_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K+M*N)*A.dtype.itemsize)))
|
|
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
|
|
src = (kittens_path/"gemm_bf16.cpp").read_text()
|
|
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
|
|
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}"]).compile_cached(src)
|
|
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
|
|
UOp(Ops.BINARY, arg=lib)))
|
|
|
|
@functools.cache
|
|
def custom_hk_bf16_atb_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
|
|
K, M = A.shape[0]*A.shape[1], A.shape[2]
|
|
K2, N = B.shape[0]*B.shape[1], B.shape[2]
|
|
assert K == K2, f"{A.shape} {B.shape}"
|
|
block_m, block_n, block_k, num_warps = 256, 256, 64, 8
|
|
assert M % block_m == 0 and N % block_n == 0 and K % block_k == 0, f"invalid bf16 atb tile {(block_m, block_n, block_k)} for {(M, N, K)}"
|
|
threads = UOp.special(64 * num_warps, "lidx0")
|
|
workgroups = UOp.special((M // block_m) * (N // block_n), "gidx0")
|
|
sink = UOp.sink(C.base, A.base, B.base, threads, workgroups,
|
|
arg=KernelInfo(f"hk_bf16_atb_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K+M*N)*A.dtype.itemsize)))
|
|
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
|
|
src = (kittens_path/"gemm_bf16_atb.cpp").read_text()
|
|
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
|
|
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}"]).compile_cached(src)
|
|
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
|
|
UOp(Ops.BINARY, arg=lib)))
|
|
|
|
def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
|
|
assert a.dtype == b.dtype == dtypes.bfloat16, f"expected bf16, got {a.dtype} {b.dtype}"
|
|
assert a.ndim == b.ndim == 3 and a.shape[:2] == b.shape[:2], f"{a.shape} {b.shape}"
|
|
batch, rows, M = a.shape
|
|
N = b.shape[2]
|
|
assert M % TILE_M == 0 and N % TILE_N == 0 and (batch * rows) % TILE_K == 0, \
|
|
f"atb shape {a.shape} {b.shape} must produce (M,N,K) multiples of ({TILE_M},{TILE_N},{TILE_K})"
|
|
is_multi = isinstance(a.device, tuple)
|
|
reduce_out = False
|
|
if is_multi:
|
|
ndev = len(a.device)
|
|
if a.uop.axis in (0, 1) or b.uop.axis in (0, 1): inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device), 0, True
|
|
elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=a.dtype, device=a.device), 2
|
|
elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=a.dtype, device=a.device), 1
|
|
else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device), 0, True
|
|
out = Tensor(inv.uop.unshard(out_axis), device=a.device)
|
|
dname = a.device[0]
|
|
else:
|
|
out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device)
|
|
dname = a.device
|
|
dname = dname.split(":")[0]
|
|
out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_hk_bf16_atb_gemm, dname=dname))[0]
|
|
if reduce_out: out = out.sum(0)
|
|
return out.squeeze(0) if out.ndim == 3 else out
|
|
|
|
# ** backward gemm, might use the asm gemm
|
|
|
|
def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=False, has_w_post:bool=False):
|
|
inputs = kernel.src[1:]
|
|
if inputs[1].dtype == FP8_DTYPE:
|
|
out, a, b = inputs[:3]
|
|
i = 3
|
|
s_x = inputs[i]; i += 1
|
|
has_w = n_scales >= 2
|
|
s_w = inputs[i] if has_w else None; i += has_w
|
|
s_g_amax = inputs[i] if n_scales == 3 else None; i += (n_scales == 3)
|
|
grad_amax_state = inputs[i] if has_grad_amax else None; i += has_grad_amax
|
|
next_grad_amax_state = inputs[i] if has_grad_amax else None; i += has_grad_amax
|
|
w_post = inputs[i] if has_w_post else None
|
|
a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
|
|
s_x_t = Tensor(s_x, device=a.device)
|
|
s_w_t = Tensor(s_w, device=a.device) if has_w else None
|
|
s_g_amax_t = Tensor(s_g_amax, device=a.device) if s_g_amax is not None else None
|
|
w_post_t = Tensor(w_post, device=a.device) if has_w_post else None
|
|
g_t = g_t[:a.shape[0]]
|
|
from extra.llama_kernels.cast_amax import _grad_fp8_mailbox
|
|
from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
|
|
gbase = gradient.base if hasattr(gradient, "base") else gradient
|
|
mailbox_entry = _grad_fp8_mailbox.pop(gbase, None) or _grad_fp8_mailbox.pop(gradient, None)
|
|
if mailbox_entry is not None:
|
|
g_fp8_u, grad_amax_u = mailbox_entry
|
|
g_fp8 = Tensor(g_fp8_u, device=a.device)[:a.shape[0]]
|
|
g_amax = Tensor(grad_amax_u, device=a.device)
|
|
else:
|
|
assert grad_amax_state is not None, "fp8 matmul bwd needs either a mailbox entry or a grad_amax_state"
|
|
if getenv("CURRENT_GRAD_SCALE", 0):
|
|
g_fp8, _, g_amax = quantize_fp8(g_t, amax_state=None)
|
|
elif getenv("FUSED_GRAD_QUANTIZE", 0):
|
|
grad_amax_t = Tensor(grad_amax_state, device=a.device)
|
|
g_amax = grad_amax_t
|
|
g_fp8, _ = quantize_fp8_delayed(g_t, g_amax, Tensor(next_grad_amax_state, device=a.device))
|
|
else:
|
|
grad_amax_t = Tensor(grad_amax_state, device=a.device)
|
|
g_amax = grad_amax_t
|
|
g_fp8, _, new_grad_amax = quantize_fp8(g_t, amax_state=g_amax)
|
|
store_effect = next_grad_amax_state.store(new_grad_amax.uop)
|
|
g_fp8 = Tensor(g_fp8.contiguous().uop.after(store_effect), device=a.device)
|
|
# dgrad: applies grad/activation amax scales in the GEMM epilogue; w_scale is already inverse.
|
|
assert s_g_amax_t is None, "fp8 GEMM bwd through g_amax scaling is unsupported"
|
|
grad_a = asm_gemm(g_fp8, b_t, x_scale=s_x_t, w_scale=s_w_t, g_amax=g_amax) if has_w else asm_gemm(g_fp8, b_t, x_scale=s_x_t, g_amax=g_amax)
|
|
# wgrad: no w_scale
|
|
grad_b = hk_fp8_atb_gemm(g_fp8, a_t, x_scale=s_x_t, g_amax=g_amax)
|
|
# wgrad: rescale if not scalar
|
|
if w_post_t is not None:
|
|
grad_b = grad_b / w_post_t.reshape(*w_post_t.shape, *([1]*(grad_b.ndim - w_post_t.ndim)))
|
|
# one None per input: (out, a, b, x_scale[, w_scale][, grad_amax][, w_post_scale])
|
|
ret = (None, grad_a.uop, grad_b.uop) + tuple(None for _ in inputs[3:])
|
|
return ret
|
|
else:
|
|
hk_bf16 = len(inputs) == 4 and inputs[1].dtype == dtypes.bfloat16
|
|
if hk_bf16:
|
|
out, a, b_t, b = inputs
|
|
assert all_same([gradient.device, a.device, b_t.device, b.device, out.device])
|
|
else:
|
|
assert len(inputs) == 3, f"regular gemm must have exactly 3 sources, got: {len(inputs)}"
|
|
out, a, b = inputs
|
|
assert all_same([gradient.device, a.device, b.device, out.device])
|
|
a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
|
|
g_t = g_t[:a.shape[0]]
|
|
if hk_bf16 and g_t.dtype != b_t.dtype: g_t = g_t.cast(b_t.dtype)
|
|
if can_use_asm_gemm(g_t, b_t.T): grad_a = asm_gemm(g_t, b_t.T).uop
|
|
else: grad_a = (g_t @ b_t.T).uop
|
|
if hk_bf16:
|
|
grad_b = hk_bf16_atb_gemm(a_t, g_t).uop
|
|
else:
|
|
a_t_flat, g_t_flat = a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1), g_t.reshape(-1, g_t.shape[-1])
|
|
if can_use_asm_gemm(a_t_flat, g_t_flat): grad_b = asm_gemm(a_t_flat, g_t_flat).uop
|
|
else: grad_b = (a_t_flat @ g_t_flat).uop
|
|
# hk_bf16 uses b.T, writes gradients only for a and b
|
|
return (None, grad_a, None, grad_b) if hk_bf16 else (None, grad_a, grad_b)
|
|
|
|
# ** mxfp8 gemm backward
|
|
|
|
def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=False):
|
|
inputs = kernel.src[1:] # (out, a_q, b_q, a_si, b_si, a_e8, b_e8, [w_post])
|
|
aq, bq = Tensor(inputs[1], device=inputs[1].device), Tensor(inputs[2], device=inputs[2].device)
|
|
ae8, be8 = Tensor(inputs[5], device=inputs[5].device), Tensor(inputs[6], device=inputs[6].device)
|
|
wp = Tensor(inputs[7], device=inputs[7].device) if has_w_post else None
|
|
|
|
a_phys = (aq.reshape(-1, aq.shape[-1]).cast(dtypes.bfloat16) * _mx_block_scale(ae8)).cast(dtypes.bfloat16)
|
|
b_phys = (bq.cast(dtypes.bfloat16) * _mx_block_scale(be8)).cast(dtypes.bfloat16)
|
|
|
|
g = Tensor(gradient, device=aq.device)[:aq.shape[0]].reshape(aq.shape[0]*aq.shape[1], bq.shape[0]).cast(dtypes.bfloat16)
|
|
grad_a = asm_gemm(g, b_phys, mx=True)
|
|
grad_b = asm_gemm(g.T, a_phys, mx=True, a_pretranspose=g)
|
|
|
|
grad_a = (grad_a * _mx_block_scale(ae8)).reshape(aq.shape)
|
|
if not w_stored: grad_b = grad_b * _mx_block_scale(be8)
|
|
if wp is not None: grad_b = grad_b / wp.reshape(-1, 1)
|
|
return (None, grad_a.uop, grad_b.uop) + tuple(None for _ in inputs[3:])
|
|
|
|
# ** mxfp4 gemm backward
|
|
|
|
def custom_mxfp4_gemm_bw(gradient:UOp, kernel:UOp):
|
|
# The raw kernel consumes quantized buffers, while the final two inputs retain the BF16 operands for STE gradients.
|
|
inputs = kernel.src[1:] # (out, a_q, b_q, scale_a, scale_b, a, w)
|
|
assert len(inputs) == 7
|
|
a, w = Tensor(inputs[5], device=inputs[5].device), Tensor(inputs[6], device=inputs[6].device)
|
|
g = Tensor(gradient, device=a.device)[:a.shape[0]].cast(dtypes.bfloat16)
|
|
grad_a = asm_gemm(g, w, mxfp4=True)
|
|
a_flat, g_flat = a.reshape(-1, a.shape[-1]), g.reshape(-1, g.shape[-1])
|
|
grad_w = asm_gemm(g_flat.T, a_flat, mxfp4=True)
|
|
return (None, None, None, None, None, grad_a.uop, grad_w.uop)
|
|
|
|
# ** main gemm function
|
|
|
|
def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None, grad_amax_state:Tensor|None=None,
|
|
next_grad_amax_state:Tensor|None=None,
|
|
w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_amax:Tensor|None=None,
|
|
a_pretranspose:Tensor|None=None, mxfp4:bool=False) -> Tensor:
|
|
assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
|
|
if mxfp4:
|
|
assert not mx and mx_scales is None, "mxfp4 owns quantization; mx/mx_scales are for mxfp8"
|
|
assert a.dtype == dtypes.bfloat16, f"cannot quantize {a.dtype} to mxfp4"
|
|
counters["used"] += 1
|
|
unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
|
|
if unfold_batch:
|
|
orig_batch = a.shape[0]
|
|
a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
|
|
squeeze = a.ndim == 2
|
|
if squeeze: a = a.unsqueeze(0)
|
|
out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE or mxfp4 else a.dtype
|
|
|
|
batch, M, K = a.shape
|
|
N = b.shape[1]
|
|
is_multi = isinstance(a.device, tuple)
|
|
if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
|
|
if (m_sharded:=is_multi and a.uop.axis == 1): M //= len(a.device)
|
|
n_sharded = is_multi and b.uop.axis == 1
|
|
|
|
if is_multi:
|
|
if n_sharded:
|
|
out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.unshard(2), device=a.device)
|
|
elif m_sharded:
|
|
out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(1), device=a.device)
|
|
else:
|
|
out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(0),
|
|
device=a.device)
|
|
else:
|
|
out = Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device)
|
|
|
|
renderer = Device[dname:=(a.device[0] if is_multi else a.device)].renderer
|
|
dname, arch = dname.split(":")[0], renderer.target.arch
|
|
if arch.startswith("gfx950") and getenv("USE_ASM", 1):
|
|
if mxfp4:
|
|
tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if (batch*M) % tm == N % tn == 0)
|
|
fxn = functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n)
|
|
w = b.T
|
|
if k_sharded:
|
|
ndev = len(a.device)
|
|
a_q, _, scale_a = quantize_mxfp4(a.reshape(batch, M, ndev, K))
|
|
b_q, _, scale_b = quantize_mxfp4(w.reshape(w.shape[0], ndev, K))
|
|
b_q = _mxfp4_shuffle_weight(b_q.permute(1, 0, 2))
|
|
else:
|
|
a_q, _, scale_a = quantize_mxfp4(a.reshape(batch*M, K))
|
|
b_q, _, scale_b = quantize_mxfp4(w)
|
|
a_q, b_q = a_q.reshape(batch, M, K//2).contiguous(), _mxfp4_shuffle_weight(b_q)
|
|
out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, a, w, fxn=fxn, grad_fxn=custom_mxfp4_gemm_bw)[0]
|
|
elif mx:
|
|
# mxfp8 1x32 block scaling
|
|
if mx_scales is not None:
|
|
a_si, a_e8, b_si, b_e8 = mx_scales
|
|
a_q, b_q = a.reshape(-1, a.shape[-1]), b.T
|
|
elif (a_pretranspose is not None and getenv("FUSED_GRAD_QUANTIZE", 0) and a_pretranspose.dtype == dtypes.bfloat16
|
|
and a_pretranspose.shape[0] % 32 == 0 and a_pretranspose.shape[1] % 256 == 0):
|
|
from extra.llama_kernels.transpose_quantize_mxfp8 import transpose_quantize_mxfp8
|
|
a_q, a_e8, a_si = transpose_quantize_mxfp8(a_pretranspose)
|
|
b_q, b_e8, b_si = quantize_mxfp8(b.T)
|
|
else:
|
|
a_q, a_e8, a_si = quantize_mxfp8(a.reshape(-1, a.shape[-1]))
|
|
b_q, b_e8, b_si = quantize_mxfp8(b.T)
|
|
has_w_post = w_post_scale is not None
|
|
fxn = functools.partial(custom_hk_mxfp8_gemm, dname=dname)
|
|
grad_fxn = functools.partial(custom_mx_gemm_bw, has_w_post=has_w_post, w_stored=mx_w_stored)
|
|
extra = [w_post_scale] if w_post_scale is not None else []
|
|
out = Tensor.custom_kernel(out, a_q.reshape(a.shape), b_q, a_si, b_si, a_e8, b_e8, *extra, fxn=fxn, grad_fxn=grad_fxn)[0]
|
|
# fp8 gemm computes [email protected], kernel multiplies output by x_scale * w_scale before bf16 store
|
|
elif a.dtype == FP8_DTYPE:
|
|
scales = tuple(s for s in (x_scale, w_scale, g_amax) if s is not None)
|
|
scale_mode = (1 if x_scale is not None else 0) | (2 if w_scale is not None else 0) | (4 if g_amax is not None else 0)
|
|
assert (grad_amax_state is None) == (next_grad_amax_state is None)
|
|
extra = ([grad_amax_state, next_grad_amax_state] if grad_amax_state is not None else []) + ([w_post_scale] if w_post_scale is not None else [])
|
|
fxn = functools.partial(custom_hk_fp8_gemm, dname=dname, scale_mode=scale_mode)
|
|
bw = functools.partial(custom_gemm_bw, n_scales=len(scales), has_grad_amax=grad_amax_state is not None, has_w_post=w_post_scale is not None)
|
|
out = Tensor.custom_kernel(out, a, b.T, *scales, *extra, fxn=fxn, grad_fxn=bw)[0]
|
|
elif a.dtype == dtypes.bfloat16:
|
|
out = Tensor.custom_kernel(out, a, b.T, b, fxn=functools.partial(custom_hk_bf16_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
|
|
else:
|
|
out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
|
|
if k_sharded: out = out.sum(0)
|
|
out = out.squeeze(0) if squeeze else out
|
|
if unfold_batch: out = out.reshape(orig_batch, -1, out.shape[-1])
|
|
if w_post_scale is not None: out = (out * w_post_scale.reshape(*([1]*(out.ndim-1)), -1)).cast(out.dtype)
|
|
return out
|