From e0c69d7a12c0e024fbb9e440478afdddc1b91a36 Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Wed, 24 Jun 2026 23:25:39 -0400 Subject: [PATCH] llama: fused grad quantize (#16731) --- extra/gemm/cdna_asm_gemm.py | 10 ++- .../transpose_quantize_mxfp8/__init__.py | 37 +++++++++++ .../transpose_quantize_mxfp8.cpp | 62 +++++++++++++++++++ 3 files changed, 107 insertions(+), 2 deletions(-) create mode 100644 extra/llama_kernels/transpose_quantize_mxfp8/__init__.py create mode 100644 extra/llama_kernels/transpose_quantize_mxfp8/transpose_quantize_mxfp8.cpp diff --git a/extra/gemm/cdna_asm_gemm.py b/extra/gemm/cdna_asm_gemm.py index cb47ce8769..2fa695b005 100644 --- a/extra/gemm/cdna_asm_gemm.py +++ b/extra/gemm/cdna_asm_gemm.py @@ -2899,7 +2899,7 @@ def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=F 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) + 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) @@ -2909,7 +2909,8 @@ def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=F # ** 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, - w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_scale:Tensor|None=None) -> Tensor: + w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_scale:Tensor|None=None, + a_pretranspose:Tensor|None=None) -> Tensor: assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}" counters["used"] += 1 unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0 @@ -2946,6 +2947,11 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N 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) diff --git a/extra/llama_kernels/transpose_quantize_mxfp8/__init__.py b/extra/llama_kernels/transpose_quantize_mxfp8/__init__.py new file mode 100644 index 0000000000..e5b43bebc8 --- /dev/null +++ b/extra/llama_kernels/transpose_quantize_mxfp8/__init__.py @@ -0,0 +1,37 @@ +from __future__ import annotations +import functools, pathlib +from tinygrad import Tensor, dtypes +from tinygrad.uop.ops import UOp, Ops, KernelInfo +from tinygrad.renderer import Estimates +from extra.llama_kernels import THREADS_PER_WG, alloc_like, dname_of, compile_hip + +TILE_N = THREADS_PER_WG # 256 +BLK = 32 + +@functools.cache +def _custom_transpose_quantize_mxfp8(q:UOp, e8:UOp, g:UOp, dname:str) -> UOp: + M, N = g.shape + num_wg = (M // BLK) * (N // TILE_N) + threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(num_wg, "gidx0") + mem = M * N * 2 + M * N + (M // BLK) * N # read bf16, write fp8 + e8 + sink = UOp.sink(q.base, e8.base, g.base, threads, workgroups, + arg=KernelInfo(f"transpose_quantize_mxfp8_{M}_{N}", estimates=Estimates(ops=M*N, mem=mem))) + src = (pathlib.Path(__file__).parent/"transpose_quantize_mxfp8.cpp").read_text() + defines = [f"-DM_DIM={M}", f"-DN_DIM={N}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"] + return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), + UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines)))) + +def transpose_quantize_mxfp8(g:Tensor) -> tuple[Tensor, Tensor, Tensor]: + # fused g.T quantize: returns (q, e8, si) == quantize_mxfp8(g.T) — q (N,M) fp8, e8 (N, M/32), si packed (M/128, N) + assert g.ndim == 2 and g.dtype == dtypes.bfloat16, f"{g.shape} {g.dtype}" + from extra.gemm.cdna_asm_gemm import FP8_DTYPE, mx_pack + M, N = g.shape + assert M % BLK == 0 and N % TILE_N == 0, f"M={M} must%{BLK}, N={N} must%{TILE_N}" + device = g.device + axis = g.uop.axis if isinstance(device, tuple) else None + out_axis = None if axis is None else (1 if axis == 0 else 0) + q = alloc_like((N, M), FP8_DTYPE, device, out_axis) + e8 = alloc_like((N, M // BLK), dtypes.uint8, device, out_axis) + fxn = functools.partial(_custom_transpose_quantize_mxfp8, dname=dname_of(device)) + q, e8, *_ = Tensor.custom_kernel(q, e8, g, fxn=fxn) + return q, e8, mx_pack(e8) diff --git a/extra/llama_kernels/transpose_quantize_mxfp8/transpose_quantize_mxfp8.cpp b/extra/llama_kernels/transpose_quantize_mxfp8/transpose_quantize_mxfp8.cpp new file mode 100644 index 0000000000..157e57e0ba --- /dev/null +++ b/extra/llama_kernels/transpose_quantize_mxfp8/transpose_quantize_mxfp8.cpp @@ -0,0 +1,62 @@ +#include +#include +#include + +#ifndef M_DIM +#define M_DIM 8192 +#endif +#ifndef N_DIM +#define N_DIM 14336 +#endif +#ifndef THREADS_PER_WG +#define THREADS_PER_WG 256 +#endif + +constexpr int BLK = 32; +constexpr int TILE_M = BLK; // one mxfp8 block along M per tile +constexpr int TILE_N = THREADS_PER_WG; // 256, one output column per thread +constexpr int LDS_STRIDE = TILE_N + 1; // +1 pad: stride 257 ≡ 1 (mod 32) -> conflict-free column reads +constexpr int N_TILES_N = N_DIM / TILE_N; +constexpr float FP8_MAX = 448.0f; + +static_assert(M_DIM % TILE_M == 0, "M_DIM must be a multiple of 32"); +static_assert(N_DIM % TILE_N == 0, "N_DIM must be a multiple of 256"); + +extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void +transpose_quantize_mxfp8(__hip_fp8_storage_t* __restrict__ q, // (N_DIM, M_DIM) + uint8_t* __restrict__ e8_out, // (N_DIM, M_DIM/32) + const __hip_bfloat16* __restrict__ g) // (M_DIM, N_DIM) +{ + __shared__ __hip_bfloat16 lds[TILE_M * LDS_STRIDE]; + const int tid = threadIdx.x; + const int tile_m = blockIdx.x / N_TILES_N; // which 32-block along M + const int tile_n = blockIdx.x % N_TILES_N; + + #pragma unroll + for (int mm = 0; mm < TILE_M; mm++) + lds[mm * LDS_STRIDE + tid] = g[(long long)(tile_m * TILE_M + mm) * N_DIM + (tile_n * TILE_N + tid)]; + __syncthreads(); + + float vals[TILE_M]; + float amax = 0.0f; + #pragma unroll + for (int mm = 0; mm < TILE_M; mm++) { + float v = (float)lds[mm * LDS_STRIDE + tid]; + vals[mm] = v; + amax = fmaxf(amax, fabsf(v)); + } + int e8 = (int)floorf(log2f(fmaxf(amax, 1e-38f))) + 127; + e8 = max(0, min(254, e8)); + float qscale = exp2f((float)(127 - e8)); + + const long long n = tile_n * TILE_N + tid; + __hip_fp8_storage_t out[TILE_M]; + #pragma unroll + for (int mm = 0; mm < TILE_M; mm++) + out[mm] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, vals[mm] * qscale)), __HIP_SATFINITE, __HIP_E4M3); + // 32 contiguous fp8 along M -> two 16-byte vector stores + long long obase = n * M_DIM + (long long)(tile_m * TILE_M); + *reinterpret_cast(&q[obase]) = *reinterpret_cast(&out[0]); + *reinterpret_cast(&q[obase + 16]) = *reinterpret_cast(&out[16]); + e8_out[n * (M_DIM / BLK) + tile_m] = (uint8_t)e8; +}