geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

llama: replace two stage amax with atomics (#17063)

Browse Source 
* atomic amax in c kernels

* quantize fp8 UOp kernel

* diff
This commit is contained in:
qazal
2026-07-17 19:27:10 +09:00
committed by GitHub
parent a836c3822a
commit 7dd3422c63
8 changed files with 59 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/llama_kernels/__init__.py
-5
View File
@@ -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
extra/llama_kernels/cast_amax/__init__.py
+14 -15
View File
@@ -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
extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp
+4 -3
View File
@@ -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<int32_t*>(grad_amax_next), __float_as_int(sdata[0]));
}
extra/llama_kernels/cast_amax/cast_amax_fwd_w13.cpp
+3 -3
View File
@@ -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<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(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<int32_t*>(amax_out), __float_as_int(sdata[0]));
}
extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py
+16 -16
View File
@@ -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
extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/fused_rmsnorm_mul_quantize_fp8.cpp
+5 -5
View File
@@ -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<int32_t*>(amax_out), __float_as_int(sdata[0]));
}
extra/llama_kernels/quantize_fp8_delayed/__init__.py
+16 -13
View File
@@ -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.
test/backend/test_llama_kernels.py
+1
View File
@@ -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