diff --git a/extra/llama_kernels/__init__.py b/extra/llama_kernels/__init__.py index c033bd9b02..7d50601a0c 100644 --- a/extra/llama_kernels/__init__.py +++ b/extra/llama_kernels/__init__.py @@ -20,11 +20,6 @@ def local_abs_max(x:Tensor) -> Tensor: fxn = _local_abs_max_fxn(param.uop, x.device) return Tensor(fxn[0].uop.call(x.uop).gettuple(0)) -def scalar_amax(amax_buf:Tensor) -> Tensor: - if isinstance(amax_buf.device, tuple): - return local_abs_max(amax_buf).detach() - return amax_buf.max().detach() - def shard_shape(shape:tuple, axis:int, ndev:int) -> list: s = list(shape) s[axis] //= ndev diff --git a/extra/llama_kernels/cast_amax/__init__.py b/extra/llama_kernels/cast_amax/__init__.py index 6767df3172..7eed8e628a 100644 --- a/extra/llama_kernels/cast_amax/__init__.py +++ b/extra/llama_kernels/cast_amax/__init__.py @@ -3,7 +3,7 @@ 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 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, dname_of # 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 @@ -11,13 +11,13 @@ from extra.llama_kernels import NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, _grad_fp8_mailbox:dict[UOp, tuple[UOp, UOp]] = {} @functools.cache -def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp, grad_amax:UOp, +def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_next: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, grad_amax.base, + mem = n_elems * 2 * 3 + n_elems * 2 + 4 + 4 + sink = UOp.sink(grad_xw13_fp8.base, grad_amax_next.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) @@ -25,14 +25,14 @@ def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp, grad_amax:UOp, UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib))) @functools.cache -def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, grad_amax_state:UOp, +def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_out:UOp, xw13:UOp, amax_state:UOp, grad_amax_state:UOp, next_grad_amax_state:UOp, dname:str) -> UOp: # NOTE: grad_amax_state is plumbed through as an unused fwd input so the bwd kernel can read it via kernel.src 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 * 2 + n_elems + NUM_WG * 4 - sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups, + mem = n_elems * 2 * 2 + n_elems + 4 + sink = UOp.sink(fp8_out.base, amax_out.base, xw13.base, amax_state.base, threads, workgroups, arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem))) src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_fwd_w13.cpp", n_elems, hidden) return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), @@ -43,17 +43,16 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp): device = xw13.device axis = xw13.axis if isinstance(device, tuple) else None grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3, device, axis) - grad_amax_buf = alloc_local((NUM_WG,), dtypes.float32, device, axis) + grad_amax_next = Tensor.zeros((), dtype=dtypes.float32, device=device).contiguous() grad_amax_state_t = Tensor(grad_amax_state, device=device) fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device)) grad_amax = grad_amax_state_t.empty_like() - grad_xw13_fp8, grad_amax_buf, grad_amax, *_ = Tensor.custom_kernel( - grad_xw13_fp8, grad_amax_buf, grad_amax, + grad_xw13_fp8, grad_amax_next, grad_amax, *_ = Tensor.custom_kernel( + grad_xw13_fp8, grad_amax_next, 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) - new_grad_amax = scalar_amax(grad_amax_buf) - store_effect = next_grad_amax_state.store(new_grad_amax.uop) + store_effect = next_grad_amax_state.store(grad_amax_next.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. @@ -70,8 +69,8 @@ def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_ HIDDEN = H2 // 2 axis = xw13.uop.axis if isinstance(xw13.device, tuple) else None fp8_out = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype, xw13.device, axis) - amax_buf = alloc_local((NUM_WG,), dtypes.float32, xw13.device, axis) + amax_out = Tensor.zeros((), dtype=dtypes.float32, device=xw13.device).contiguous() fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname_of(xw13.device)) - fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, grad_amax_state, next_grad_amax_state, + fp8_out, amax_out, *_ = Tensor.custom_kernel(fp8_out, amax_out, xw13, amax_state, grad_amax_state, next_grad_amax_state, fxn=fxn, grad_fxn=_fused_quantize_bwd_w13) - return fp8_out, scalar_amax(amax_buf) + return fp8_out, amax_out 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 459c1137a3..9620f6117e 100644 --- a/extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp +++ b/extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp @@ -23,14 +23,14 @@ static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC"); // 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 +// 2) fp32 grad_amax_next — scalar |grad_xw13| via global atomic max // 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_next, // fp32 scalar, initialized to 0 before launch 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 @@ -92,5 +92,6 @@ fused_silu_mul_bwd_w13( if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]); __syncthreads(); } - if (tid == 0) grad_amax_buf[wg] = sdata[0]; + if (tid == 0 && sdata[0] > *grad_amax_next) + atomicMax(reinterpret_cast(grad_amax_next), __float_as_int(sdata[0])); } diff --git a/extra/llama_kernels/cast_amax/cast_amax_fwd_w13.cpp b/extra/llama_kernels/cast_amax/cast_amax_fwd_w13.cpp index 58c13f4310..445d6d4b21 100644 --- a/extra/llama_kernels/cast_amax/cast_amax_fwd_w13.cpp +++ b/extra/llama_kernels/cast_amax/cast_amax_fwd_w13.cpp @@ -24,7 +24,7 @@ static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC (so VEC loads extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void fused_silu_mul_cast_amax_w13( __hip_fp8_storage_t* __restrict__ fp8_out, // fp8, N_ELEMS - float* __restrict__ amax_buf, // fp32, NUM_WG (per-WG amaxes) + float* __restrict__ amax_out, // fp32 scalar, initialized to 0 before launch const __hip_bfloat16* __restrict__ xw13, // bf16, 2*N_ELEMS const float* __restrict__ amax_state) // fp32 scalar { @@ -67,7 +67,7 @@ fused_silu_mul_cast_amax_w13( *reinterpret_cast(&fp8_out[base]) = *reinterpret_cast(out); } - // LDS tree reduction: per-workgroup amax + // LDS tree reduction: per-workgroup amax, then global atomic into the scalar. sdata[tid] = local_max; __syncthreads(); for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) { @@ -75,5 +75,5 @@ fused_silu_mul_cast_amax_w13( __syncthreads(); } - if (tid == 0) amax_buf[wg] = sdata[0]; + if (tid == 0 && sdata[0] > *amax_out) atomicMax(reinterpret_cast(amax_out), __float_as_int(sdata[0])); } diff --git a/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py b/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py index 55f928ca08..8a4a4e84f6 100644 --- a/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py +++ b/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py @@ -3,19 +3,19 @@ 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, alloc_like, alloc_local, scalar_amax, dname_of, compile_hip +from extra.llama_kernels import NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, dname_of, compile_hip def _src() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8.cpp").read_text() def _src_bwd() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8_bwd.cpp").read_text() @functools.cache -def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp, +def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_out:UOp, x:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp: MBS, SEQ, HIDDEN = x.shape n_elems = MBS * SEQ * HIDDEN threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0") - mem = n_elems * 2 + n_elems + MBS * SEQ * 4 + n_elems + HIDDEN * 2 + NUM_WG * 4 + 4 - sink = UOp.sink(fp8_out.base, x_normed_out.base, rrms_out.base, amax_buf.base, + mem = n_elems * 2 + n_elems + MBS * SEQ * 4 + n_elems + HIDDEN * 2 + 4 + 4 + sink = UOp.sink(fp8_out.base, x_normed_out.base, rrms_out.base, amax_out.base, x.base, weight.base, amax_state.base, threads, workgroups, arg=KernelInfo(f"fused_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}", estimates=Estimates(ops=6*n_elems, mem=mem))) @@ -26,13 +26,13 @@ def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp, UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines)))) @functools.cache -def _custom_fwd_add(fp8_out:UOp, h_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp, +def _custom_fwd_add(fp8_out:UOp, h_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_out:UOp, x:UOp, residual:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp: MBS, SEQ, HIDDEN = x.shape n_elems = MBS * SEQ * HIDDEN threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0") - mem = n_elems * 2 * 4 + MBS * SEQ * 4 + HIDDEN * 2 + NUM_WG * 4 + 4 - sink = UOp.sink(fp8_out.base, h_out.base, x_normed_out.base, rrms_out.base, amax_buf.base, + mem = n_elems * 2 * 4 + MBS * SEQ * 4 + HIDDEN * 2 + 4 + 4 + sink = UOp.sink(fp8_out.base, h_out.base, x_normed_out.base, rrms_out.base, amax_out.base, x.base, residual.base, weight.base, amax_state.base, threads, workgroups, arg=KernelInfo(f"fused_add_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}", estimates=Estimates(ops=7*n_elems, mem=mem))) @@ -85,7 +85,7 @@ def _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_st return grad_total.uop, grad_weight_uop def _fused_bwd(gradient:UOp, kernel:UOp): - # NOTE: fwd inputs (fp8_out, x_normed_out, rrms_out, amax_buf, x, weight, amax_state) + # NOTE: fwd inputs (fp8_out, x_normed_out, rrms_out, amax_out, x, weight, amax_state) _, x_normed_u, rrms_u, _, x_u, weight_u, amax_state_u = kernel.src[1:] grad_x, grad_w = _bwd_common(gradient, None, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel) return (None, None, None, None, grad_x, grad_w, None) @@ -123,11 +123,11 @@ def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, e fp8_out = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype, x.device, axis) x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis) rrms_out = alloc_like((MBS, SEQ), dtypes.float32, x.device, axis) - amax_buf = alloc_local((NUM_WG,), dtypes.float32, x.device, axis) + amax_out = Tensor.zeros((), dtype=dtypes.float32, device=x.device).contiguous() fxn = functools.partial(_custom_fwd, dname=dname_of(x.device), eps_val=eps) - fp8_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel( - fp8_out, x_normed_out, rrms_out, amax_buf, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_bwd) - return fp8_out, scalar_amax(amax_buf), x_normed_out, rrms_out + fp8_out, x_normed_out, rrms_out, amax_out, *_ = Tensor.custom_kernel( + fp8_out, x_normed_out, rrms_out, amax_out, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_bwd) + return fp8_out, amax_out, x_normed_out, rrms_out def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor, amax_state:Tensor, eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: @@ -143,9 +143,9 @@ def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor, h_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis) x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis) rrms_out = alloc_like((MBS, SEQ), dtypes.float32, x.device, axis) - amax_buf = alloc_local((NUM_WG,), dtypes.float32, x.device, axis) + amax_out = Tensor.zeros((), dtype=dtypes.float32, device=x.device).contiguous() fxn = functools.partial(_custom_fwd_add, dname=dname_of(x.device), eps_val=eps) - fp8_out, h_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel( - fp8_out, h_out, x_normed_out, rrms_out, amax_buf, x, residual, weight, amax_state, + fp8_out, h_out, x_normed_out, rrms_out, amax_out, *_ = Tensor.custom_kernel( + fp8_out, h_out, x_normed_out, rrms_out, amax_out, x, residual, weight, amax_state, fxn=fxn, grad_fxn=_fused_add_bwd) - return fp8_out, scalar_amax(amax_buf), h_out, x_normed_out, rrms_out + return fp8_out, amax_out, h_out, x_normed_out, rrms_out diff --git a/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/fused_rmsnorm_mul_quantize_fp8.cpp b/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/fused_rmsnorm_mul_quantize_fp8.cpp index 95750f32cf..1ccebf20cb 100644 --- a/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/fused_rmsnorm_mul_quantize_fp8.cpp +++ b/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/fused_rmsnorm_mul_quantize_fp8.cpp @@ -7,7 +7,7 @@ // fp8 = fp8_sat(y * (FP8_MAX / amax_state)) // Also writes: // rrms[row] — saved for the rmsnorm backward -// amax_buf[wg] — per-WG |y| partials, reduced later to update amax_state +// amax_out — scalar |y| via global atomic max // // Layout: one WG per row, ROWS_PER_WG rows per WG via grid-stride (ROWS = N_ELEMS / HIDDEN). // Each thread handles HIDDEN / THREADS_PER_WG elements per row. @@ -48,7 +48,7 @@ fused_add_rmsnorm_mul_quantize_fp8( __hip_bfloat16* __restrict__ h_out, // bf16, ROWS*HIDDEN — x + residual (saved for downstream) __hip_bfloat16* __restrict__ x_normed_out, // bf16, ROWS*HIDDEN float* __restrict__ rrms_out, // fp32, ROWS - float* __restrict__ amax_buf, // fp32, NUM_WG + float* __restrict__ amax_out, // fp32 scalar, initialized to 0 before launch const __hip_bfloat16* __restrict__ x, // bf16, ROWS*HIDDEN const __hip_bfloat16* __restrict__ residual, // bf16, ROWS*HIDDEN — added into x before rmsnorm const __hip_bfloat16* __restrict__ weight, // bf16, HIDDEN @@ -60,7 +60,7 @@ fused_rmsnorm_mul_quantize_fp8( __hip_fp8_storage_t* __restrict__ fp8_out, // fp8, ROWS*HIDDEN __hip_bfloat16* __restrict__ x_normed_out, // bf16, ROWS*HIDDEN (saved for rmsnorm bwd) float* __restrict__ rrms_out, // fp32, ROWS (fp32 to match rmsnorm_bwd.cpp expectation) - float* __restrict__ amax_buf, // fp32, NUM_WG per-WG partials + float* __restrict__ amax_out, // fp32 scalar, initialized to 0 before launch const __hip_bfloat16* __restrict__ x, // bf16, ROWS*HIDDEN const __hip_bfloat16* __restrict__ weight, // bf16, HIDDEN (per-hidden scale) const float* __restrict__ amax_state) // fp32 scalar @@ -144,12 +144,12 @@ fused_rmsnorm_mul_quantize_fp8( __syncthreads(); // before next row's sum_sq reduce reuses sdata } - // Final per-WG amax reduce. + // Final per-WG amax reduce, then global atomic into the scalar. sdata[tid] = local_max; __syncthreads(); for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) { if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]); __syncthreads(); } - if (tid == 0) amax_buf[wg] = sdata[0]; + if (tid == 0 && sdata[0] > *amax_out) atomicMax(reinterpret_cast(amax_out), __float_as_int(sdata[0])); } diff --git a/extra/llama_kernels/quantize_fp8_delayed/__init__.py b/extra/llama_kernels/quantize_fp8_delayed/__init__.py index f6eab12b7f..015b8e1ec2 100644 --- a/extra/llama_kernels/quantize_fp8_delayed/__init__.py +++ b/extra/llama_kernels/quantize_fp8_delayed/__init__.py @@ -3,14 +3,13 @@ from tinygrad import Tensor, dtypes from tinygrad.dtype import AddrSpace from tinygrad.helpers import prod from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType -from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax +from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like @functools.cache -def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_state:UOp) -> UOp: +def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:UOp, device=None) -> UOp: VEC = 8 n_elems = prod(x.shape) assert n_elems % (NUM_WG * THREADS_PER_WG * VEC) == 0 - assert amax_partial.shape[0] == NUM_WG x = x.reshape(n_elems) fp8_out = fp8_out.reshape(n_elems) @@ -46,8 +45,13 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_st lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)).barrier()) step //= 2 - amax_store = amax_partial[tid.eq(0).where(wg, UOp.invalid())].store(lds[0]) - return amax_store.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=())) + device = device[0].split(":")[0] if isinstance(device, tuple) else device.split(":")[0] + if device in {"AMD", "NULL"}: atomic_arg = "if ({2} > {3}) __hip_atomic_fetch_max((int*){0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);" + else: raise NotImplementedError(f"no atomic max for device {device}") + amax_idx = amax_out.reshape((1,)).index(UOp.const(dtypes.index, 0)) + max_val = lds[0].load() + atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg) + return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=())) @functools.cache def _custom_quantize_fp8_scalar(fp8_out:UOp, x:UOp, amax_state:UOp) -> UOp: @@ -71,7 +75,7 @@ def _quantize_fp8_delayed_bwd(gradient:UOp, kernel:UOp): def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> tuple[Tensor, Tensor, Tensor, UOp]: # NOTE: one-pass bf16 -> fp8 quantize with delayed scaling. Returns (fp8, inv_scale, new_amax, store_effect). - # Fused kernel reads x once and writes fp8 + per-WG |x| partials (then a small reduce produces scalar new_amax). + # Fused kernel reads x once and writes fp8 + scalar amax via global atomic max. # store_effect writes new_amax into amax_state's buffer — the caller must thread it into a realized # output via `.after(store_effect)`. Calling `amax_state.assign(new_amax)` inside a grad_fxn does # NOT work because .assign mutates only the temp Tensor's .uop, not the original layer-owned buffer. @@ -80,14 +84,13 @@ def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) fp8_out = alloc_like(x.shape, fp8_dtype, x.device, axis) n_elems = prod(x.uop.shard_shape) assert n_elems % NUM_WG == 0, f"{n_elems=} must divide over {NUM_WG=}" - amax_partial = alloc_local((NUM_WG,), dtypes.float32, x.device, axis) - fxn = _custom_quantize_fp8_with_amax - fp8_out, amax_partial, *_ = Tensor.custom_kernel(fp8_out, amax_partial, x, amax_state, - fxn=fxn, grad_fxn=_quantize_fp8_delayed_bwd) - new_amax = scalar_amax(amax_partial) + amax_out = Tensor.zeros((), dtype=dtypes.float32, device=x.device).contiguous() + fxn = functools.partial(_custom_quantize_fp8_with_amax, device=x.device) + fp8_out, amax_out, *_ = Tensor.custom_kernel(fp8_out, amax_out, x, amax_state, + fxn=fxn, grad_fxn=_quantize_fp8_delayed_bwd) inv_scale = (amax_state.float() + 1e-8) / FP8_MAX - store_effect = amax_state.uop.store(new_amax.uop) - return fp8_out, inv_scale, new_amax, store_effect + store_effect = amax_state.uop.store(amax_out.uop) + return fp8_out, inv_scale, amax_out, store_effect def quantize_fp8_scalar(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> Tensor: # NOTE: pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation. diff --git a/test/backend/test_llama_kernels.py b/test/backend/test_llama_kernels.py index 569ed1c18d..e0cd4c5559 100644 --- a/test/backend/test_llama_kernels.py +++ b/test/backend/test_llama_kernels.py @@ -66,6 +66,7 @@ def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None: assert new_amax.allclose(ref_new_amax, atol=0, rtol=0).item(), \ f"amax mismatch: got={new_amax.item()} ref={ref_new_amax.item()} diff={abs(new_amax.item()-ref_new_amax.item())}" +@unittest.skipUnless(Device.DEFAULT == "AMD", "requires atomic max") class TestQuantizeFP8(unittest.TestCase): def setUp(self): ren = Device[Device.DEFAULT].renderer