sunnypilot v2026.002.000

version: sunnypilot v2026.002.000 (staging) date: 2026-07-21T07:44:02 master commit: 2c334ede443d7391d27575af8c854a095ba702a8
2026-07-21 07:44:02 +00:00
commit de77c89cf6
@@ -0,0 +1,67 @@
+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 FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax, dname_of, compile_hip
+
+@functools.cache
+def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_state:UOp, dname:str) -> UOp:
+  n_elems = 1
+  for d in x.shape: n_elems *= d
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  mem = n_elems * 2 + n_elems + 4 + NUM_WG * 4
+  sink = UOp.sink(fp8_out.base, amax_partial.base, x.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", estimates=Estimates(ops=3*n_elems, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"quantize_fp8_with_amax.cpp").read_text()
+  defines = [f"-DN_ELEMS={n_elems}", f"-DNUM_WG={NUM_WG}", 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))))
+
+@functools.cache
+def _custom_quantize_fp8_scalar(fp8_out:UOp, x:UOp, amax_state:UOp, dname:str) -> UOp:
+  n_elems = 1
+  for d in x.shape: n_elems *= d
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  mem = n_elems * 2 + n_elems
+  sink = UOp.sink(fp8_out.base, x.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"quantize_fp8_scalar_{n_elems}", estimates=Estimates(ops=2*n_elems, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"quantize_fp8_scalar.cpp").read_text()
+  defines = [f"-DN_ELEMS={n_elems}", f"-DNUM_WG={NUM_WG}", 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 _quantize_fp8_delayed_bwd(gradient:UOp, kernel:UOp):
+  # NOTE: STE-equivalent backward — grad_x = grad_fp8 * scale, scale = FP8_MAX / amax_state.
+  # `gradient` is bf16 grad w.r.t. fp8 output (asm_gemm bwd already applied x_scale).
+  _, _, x, amax_state = kernel.src[1:]
+  device = x.device
+  scale = FP8_MAX / (Tensor(amax_state, device=device).float() + 1e-8)
+  grad_x = (Tensor(gradient, device=device).float() * scale).cast(dtypes.bfloat16)
+  return (None, None, grad_x.uop, None)
+
+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).
+  # 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.
+  assert x.dtype == dtypes.bfloat16, f"expected bf16, got {x.dtype}"
+  axis = x.uop.axis if isinstance(x.device, tuple) else None
+  fp8_out      = alloc_like(x.shape,  fp8_dtype,      x.device, axis)
+  amax_partial = alloc_local((NUM_WG,), dtypes.float32, x.device, axis)
+  fxn = functools.partial(_custom_quantize_fp8_with_amax, dname=dname_of(x.device))
+  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)
+  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
+
+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.
+  axis = x.uop.axis if isinstance(x.device, tuple) else None
+  fp8_out = alloc_like(x.shape, fp8_dtype, x.device, axis)
+  fxn = functools.partial(_custom_quantize_fp8_scalar, dname=dname_of(x.device))
+  fp8_out, *_ = Tensor.custom_kernel(fp8_out, x, amax_state, fxn=fxn)
+  return fp8_out
@@ -0,0 +1,48 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp8.h>
+
+// Pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation.
+
+#ifndef N_ELEMS
+#define N_ELEMS 67108864
+#endif
+#ifndef NUM_WG
+#define NUM_WG 1024
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+
+constexpr int VEC = 8;
+constexpr float FP8_MAX = 448.0f;
+
+static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+quantize_fp8_scalar(
+    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, N_ELEMS
+    const __hip_bfloat16* __restrict__ x,               // bf16, N_ELEMS
+    const float*          __restrict__ amax_state)      // fp32 scalar (delayed)
+{
+  const int tid = threadIdx.x;
+  const int wg  = blockIdx.x;
+  const int gid = wg * THREADS_PER_WG + tid;
+  const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
+
+  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
+
+  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
+    float4 x_raw = *reinterpret_cast<const float4*>(&x[base]);
+    const __hip_bfloat16 *xi = reinterpret_cast<const __hip_bfloat16*>(&x_raw);
+
+    __hip_fp8_storage_t out[VEC];
+    #pragma unroll
+    for (int i = 0; i < VEC; i++) {
+      const float v = static_cast<float>(xi[i]);
+      const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, v * scale));
+      out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3);
+    }
+    *reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
+  }
+}
@@ -0,0 +1,63 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp8.h>
+
+// One-pass bf16 -> fp8 quantize using a scalar delayed amax state,
+// AND simultaneously computes per-WG |x| max partials for the next step's amax state.
+// Saves one full HBM pass over the grad tensor vs. doing quantize + separate abs().max().
+
+#ifndef N_ELEMS
+#define N_ELEMS 67108864
+#endif
+#ifndef NUM_WG
+#define NUM_WG 1024
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+
+constexpr int VEC = 8;
+constexpr float FP8_MAX = 448.0f;
+
+static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+quantize_fp8_with_amax(
+    __hip_fp8_storage_t*  __restrict__ fp8_out,       // out: fp8, N_ELEMS
+    float*                __restrict__ amax_partial,  // out: fp32, NUM_WG per-WG partials
+    const __hip_bfloat16* __restrict__ x,             // in:  bf16, N_ELEMS
+    const float*          __restrict__ amax_state)    // in:  fp32 scalar (delayed)
+{
+  __shared__ float sdata[THREADS_PER_WG];
+
+  const int tid = threadIdx.x;
+  const int wg  = blockIdx.x;
+  const int gid = wg * THREADS_PER_WG + tid;
+  const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
+
+  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
+  float local_max = 0.0f;
+
+  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
+    float4 x_raw = *reinterpret_cast<const float4*>(&x[base]);
+    const __hip_bfloat16 *xi = reinterpret_cast<const __hip_bfloat16*>(&x_raw);
+
+    __hip_fp8_storage_t out[VEC];
+    #pragma unroll
+    for (int i = 0; i < VEC; i++) {
+      const float v = static_cast<float>(xi[i]);
+      local_max = fmaxf(local_max, fabsf(v));
+      const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, v * scale));
+      out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3);
+    }
+    *reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
+  }
+
+  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_partial[wg] = sdata[0];
+}