From 9ea7de020bf233702e335a1732001a5fc4389125 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Sun, 12 Jul 2026 17:15:01 +0800 Subject: [PATCH] llama: fuse grad scale in gemm epilogue (#16992) --- extra/gemm/cdna_asm_gemm.py | 40 +++++++++++-------- extra/llama_kernels/cast_amax/__init__.py | 16 ++++---- .../cast_amax/cast_amax_bwd_w13.cpp | 9 ++++- extra/thunder/amd/gemm_fp8.cpp | 5 ++- extra/thunder/amd/gemm_fp8_atb.cpp | 7 ++-- 5 files changed, 45 insertions(+), 32 deletions(-) diff --git a/extra/gemm/cdna_asm_gemm.py b/extra/gemm/cdna_asm_gemm.py index 0214fbe30e..67a4debbf9 100644 --- a/extra/gemm/cdna_asm_gemm.py +++ b/extra/gemm/cdna_asm_gemm.py @@ -57,7 +57,7 @@ def custom_hk_fp8_atb_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str, scale_mode 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_scale:Tensor|None=None) -> Tensor: +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 @@ -78,8 +78,8 @@ def hk_fp8_atb_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, g_scale:Tensor 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_scale) if s is not None) - scale_mode = (1 if x_scale is not None else 0) | (4 if g_scale is not None else 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 @@ -250,13 +250,13 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool= 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 = inputs[i] if n_scales == 3 else None; i += (n_scales == 3) + 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 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_t = Tensor(s_g, device=a.device) if s_g is not None 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 @@ -264,27 +264,33 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool= 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, inv_scale_u = mailbox_entry + g_fp8_u, grad_amax_u = mailbox_entry g_fp8 = Tensor(g_fp8_u, device=a.device)[:a.shape[0]] - g_scale = Tensor(inv_scale_u, device=a.device) + 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_scale, _ = quantize_fp8(g_t, amax_state=None) + g_fp8, _, g_amax = quantize_fp8(g_t, amax_state=None) elif getenv("FUSED_GRAD_QUANTIZE", 0): - g_fp8, g_scale, _, store_effect = quantize_fp8_delayed(g_t, Tensor(grad_amax_state, device=a.device)) + grad_amax_t = Tensor(grad_amax_state, device=a.device) + # Snapshot delayed amax before updating state; GEMM must use the same scale as quantize. + g_amax = grad_amax_t.empty_like().assign(grad_amax_t) + g_fp8, _, new_grad_amax, _ = quantize_fp8_delayed(g_t, g_amax) + store_effect = grad_amax_state.store(new_grad_amax.uop) assert g_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {g_fp8.uop.op}" g_fp8 = Tensor(g_fp8.uop.replace(src=g_fp8.uop.src + (store_effect,)), device=a.device) else: grad_amax_t = Tensor(grad_amax_state, device=a.device) - g_fp8, g_scale, new_grad_amax = quantize_fp8(g_t, amax_state=grad_amax_t) + # Snapshot delayed amax before updating state; GEMM must use the same scale as quantize. + g_amax = grad_amax_t.empty_like().assign(grad_amax_t) + g_fp8, _, new_grad_amax = quantize_fp8(g_t, amax_state=g_amax) store_effect = grad_amax_state.store(new_grad_amax.uop) g_fp8 = Tensor(g_fp8.contiguous().uop.after(store_effect), device=a.device) - # dgrad: uses g_scale * x_scale * w_scale (only when scalar) - if s_g_t is not None: g_scale = g_scale * s_g_t - grad_a = asm_gemm(g_fp8, b_t, x_scale=s_x_t, w_scale=s_w_t, g_scale=g_scale) if has_w else asm_gemm(g_fp8, b_t, x_scale=s_x_t, w_scale=g_scale) + # 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_scale=g_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))) @@ -337,7 +343,7 @@ 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, + 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) -> Tensor: assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}" counters["used"] += 1 @@ -390,8 +396,8 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N 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 a@b.T, 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_scale) 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_scale is not None else 0) + 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) extra = ([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) diff --git a/extra/llama_kernels/cast_amax/__init__.py b/extra/llama_kernels/cast_amax/__init__.py index 3531340f35..0f3362d6e9 100644 --- a/extra/llama_kernels/cast_amax/__init__.py +++ b/extra/llama_kernels/cast_amax/__init__.py @@ -3,21 +3,21 @@ 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 FP8_MAX, NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, alloc_local, scalar_amax, dname_of +from extra.llama_kernels import NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, alloc_local, scalar_amax, dname_of -# module-level mailbox: grad_xw13 UOp -> (grad_xw13_fp8 UOp, inv_scale UOp) +# module-level mailbox: grad_xw13 UOp -> (grad_xw13_fp8 UOp, delayed amax UOp) # lets cdna_asm_gemm's bwd reuse the fp8 companion produced by the fused silu_mul bwd kernel # instead of doing a redundant bf16 -> fp8 quantize. _grad_fp8_mailbox:dict[UOp, tuple[UOp, UOp]] = {} @functools.cache -def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp, +def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp, grad_amax:UOp, xw13:UOp, grad_x2:UOp, amax_state:UOp, grad_amax_state:UOp, dname:str) -> UOp: hidden = xw13.shape[2] // 2 n_elems = xw13.shape[0] * xw13.shape[1] * hidden threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0") mem = n_elems * 2 * 3 + n_elems * 2 + NUM_WG * 4 + 4 - sink = UOp.sink(grad_xw13_fp8.base, grad_amax_buf.base, + sink = UOp.sink(grad_xw13_fp8.base, grad_amax_buf.base, grad_amax.base, xw13.base, grad_x2.base, amax_state.base, grad_amax_state.base, threads, workgroups, arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=mem))) src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden) @@ -44,19 +44,19 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp): grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3, device, axis) grad_amax_buf = alloc_local((NUM_WG,), dtypes.float32, device, axis) grad_amax_state_t = Tensor(grad_amax_state, device=device) + grad_amax = grad_amax_state_t.empty_like() fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device)) - grad_xw13_fp8, grad_amax_buf, *_ = Tensor.custom_kernel( - grad_xw13_fp8, grad_amax_buf, + grad_xw13_fp8, grad_amax_buf, grad_amax, *_ = Tensor.custom_kernel( + grad_xw13_fp8, grad_amax_buf, grad_amax, Tensor(xw13, device=device), Tensor(gradient, device=device).cast(dtypes.bfloat16), Tensor(amax_state, device=device), grad_amax_state_t, fxn=fxn) grad_xw13_uop = grad_xw13_fp8.uop.cast(dtypes.bfloat16) - inv_scale = (grad_amax_state_t.float() + 1e-8) / FP8_MAX new_grad_amax = scalar_amax(grad_amax_buf) store_effect = grad_amax_state_t.uop.store(new_grad_amax.uop) assert grad_xw13_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {grad_xw13_fp8.uop.op}" grad_xw13_fp8_uop = grad_xw13_fp8.uop.replace(src=grad_xw13_fp8.uop.src + (store_effect,)) # Stash fp8 companion for cdna_asm_gemm's bwd to attach to grad_a. - _grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8_uop, inv_scale.uop) + _grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8_uop, grad_amax.uop) return (None, None, grad_xw13_uop, None, None) def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor) -> tuple[Tensor, Tensor]: diff --git a/extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp b/extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp index c787fea4b0..459c1137a3 100644 --- a/extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp +++ b/extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp @@ -21,15 +21,17 @@ constexpr float FP8_MAX = 448.0f; static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC"); static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC"); -// fused silu*mul backward, two outputs in a single HBM pass: +// fused silu*mul backward, three outputs in a single HBM pass: // 1) fp8 grad_xw13_fp8 — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd) // 2) fp32 grad_amax_buf — per-WG partial |grad_xw13|, reduced into next step's grad_amax_state +// 3) fp32 grad_amax_out — delayed grad amax used for quantize/GEMM epilogue scale // grad_amax_state is read for the fp8 scale. The store of new_grad_amax into grad_amax_state's // buffer is built in Python as a separate effect and threaded into grad_a via .after(store). extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void fused_silu_mul_bwd_w13( __hip_fp8_storage_t* __restrict__ grad_xw13_fp8_out, // fp8, 2*N_ELEMS float* __restrict__ grad_amax_buf, // fp32, NUM_WG per-WG partials + float* __restrict__ grad_amax_out, // fp32 scalar delayed grad amax const __hip_bfloat16* __restrict__ xw13, // bf16, 2*N_ELEMS const __hip_bfloat16* __restrict__ grad_x2, // bf16, N_ELEMS const float* __restrict__ amax_state, // fp32 scalar (fwd x2 amax) @@ -43,9 +45,12 @@ fused_silu_mul_bwd_w13( const int stride_elems = NUM_WG * THREADS_PER_WG * VEC; const float scale = FP8_MAX / (static_cast(*amax_state) + 1e-8f); - const float g_scale = FP8_MAX / (static_cast(*grad_amax_state) + 1e-8f); + const float grad_amax = static_cast(*grad_amax_state); + const float g_scale = FP8_MAX / (grad_amax + 1e-8f); float local_max = 0.0f; + if (wg == 0 && tid == 0) *grad_amax_out = grad_amax; + for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) { const int outer = base / HIDDEN; const int inner = base % HIDDEN; diff --git a/extra/thunder/amd/gemm_fp8.cpp b/extra/thunder/amd/gemm_fp8.cpp index acf65aa92c..cfca410514 100644 --- a/extra/thunder/amd/gemm_fp8.cpp +++ b/extra/thunder/amd/gemm_fp8.cpp @@ -106,7 +106,7 @@ __global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_pt , float *w_scale_ptr #endif #if SCALE_MODE & 4 - , float *g_scale_ptr + , float *g_amax_ptr #endif ) { constexpr int M = GEMM_M, N = GEMM_N, K = GEMM_K; @@ -358,7 +358,8 @@ __global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_pt scale *= *w_scale_ptr; #endif #if SCALE_MODE & 4 - scale *= *g_scale_ptr; + float g_scale = (*g_amax_ptr + 1e-08f) * (1.0f / 448.0f); + scale *= g_scale; #endif mul(cA, cA, scale); diff --git a/extra/thunder/amd/gemm_fp8_atb.cpp b/extra/thunder/amd/gemm_fp8_atb.cpp index 7eca22c3c7..117a43cb7f 100644 --- a/extra/thunder/amd/gemm_fp8_atb.cpp +++ b/extra/thunder/amd/gemm_fp8_atb.cpp @@ -78,7 +78,7 @@ __device__ inline static void load_st_to_rt(RT &dst, const ST &src) { constexpr int NUM_WARPS = 8; using G = kittens::group; -// SCALE_MODE bits: 1=x_scale, 2=w_scale, 4=g_scale +// SCALE_MODE bits: 1=x_scale, 2=w_scale, 4=g_amax #ifndef SCALE_MODE #define SCALE_MODE 5 #endif @@ -91,7 +91,7 @@ __global__ __launch_bounds__(512, 2) void hk_fp8_atb_gemm(bf16 *C_ptr, fp8e4m3 * , float *w_scale_ptr #endif #if SCALE_MODE & 4 - , float *g_scale_ptr + , float *g_amax_ptr #endif ) { constexpr int M = GEMM_M, N = GEMM_N, K = GEMM_K; @@ -335,7 +335,8 @@ __global__ __launch_bounds__(512, 2) void hk_fp8_atb_gemm(bf16 *C_ptr, fp8e4m3 * scale *= *w_scale_ptr; #endif #if SCALE_MODE & 4 - scale *= *g_scale_ptr; + float g_scale = (*g_amax_ptr + 1e-08f) * (1.0f / 448.0f); + scale *= g_scale; #endif mul(cA, cA, scale);