llama: fused grad quantize (#16731)

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)

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)

extra/llama_kernels/transpose_quantize_mxfp8/transpose_quantize_mxfp8.cpp

@@ -0,0 +1,62 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_fp8.h>
+#include <hip/hip_bf16.h>
+
+#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<uint4*>(&q[obase])      = *reinterpret_cast<uint4*>(&out[0]);
+  *reinterpret_cast<uint4*>(&q[obase + 16]) = *reinterpret_cast<uint4*>(&out[16]);
+  e8_out[n * (M_DIM / BLK) + tile_m] = (uint8_t)e8;
+}