nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

llama: fuse grad scale in gemm epilogue (#16992)

Browse Source
This commit is contained in:
qazal
2026-07-12 18:15:01 +09:00
committed by GitHub
parent 73b38854b4
commit 9ea7de020b
5 changed files with 45 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/cdna_asm_gemm.py
+23 -17
View File
@@ -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 [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_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)
extra/llama_kernels/cast_amax/__init__.py
+8 -8
View File
@@ -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]:
extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp
+7 -2
View File
@@ -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<float>(*amax_state) + 1e-8f);
const float g_scale = FP8_MAX / (static_cast<float>(*grad_amax_state) + 1e-8f);
const float grad_amax = static_cast<float>(*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;
extra/thunder/amd/gemm_fp8.cpp
+3 -2
View File
@@ -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);
extra/thunder/amd/gemm_fp8_atb.cpp
+4 -3
View File
@@ -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<NUM_WARPS>;
// 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);