diff --git a/areno/accel/csrc/e4m3_linear.cu b/areno/accel/csrc/e4m3_linear.cu new file mode 100644 index 00000000..55e4d000 --- /dev/null +++ b/areno/accel/csrc/e4m3_linear.cu @@ -0,0 +1,395 @@ +// E4M3 (FP8, 3 mantissa bits) fused dequant-linear for AReno on Ampere / A100. +// +// Computes: y(M,N) = x(M,K) @ (dequant_e4m3(w(N,K)) * scale)^T +// where the weight is stored as uint8 E4M3 (1 byte/element) and the per-tensor +// scalar `scale` is applied after the dot. Reading the byte-packed weight halves +// the weight bytes pulled from HBM vs the bf16 `areno_linear` path, which is the +// memory-bandwidth benefit from RFC 0001 (theoretical ~2x, decode-bound). +// +// Forward-only: E4M3 has no backward, so this is a decode/inference operator. +// Two kernels serve different shapes: +// * e4m3_gemv_kernel — the small-M (M<=4) memory-bound decode path: one warp +// streams a weight row with 32-byte-coalesced reads, decoding each E4M3 byte +// inline (branchless bit-trick) and reducing with warp shuffles. No tensor +// cores; the weight (not the activation) is the streamed operand. +// * e4m3_linear_kernel — the general-M fallback (M>4): a WMMA (fp16 16x16x16) +// tensor-core GEMM that stages the decoded weight as fp16 in shared memory. +// Both multiply the fp32 accumulator by `scale` before writing bf16. +// +// Decode reference (validated against torch.float8_e4m3fn on CPU): +// sign = (b>>7)&1; e = (b>>3)&0xF; m = b&0x7 +// e==0 (subnormal): (-1)^s * m * 2^-9 +// 1<=e<=15 (normal): (-1)^s * 2^(e-7) * (1 + m/8); max finite 448 (e=15,m=6) +// e==15 && m==7 is NaN in torch; weights are clamped at 448 so it cannot +// occur in a valid E4M3 payload, but we clamp it to 448 defensively. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace areno_accel { + +using namespace nvcuda; + +// One WMMA 16x16x16 fp16 tensor-core tile. +constexpr int kWmmaM = 16; +constexpr int kWmmaN = 16; +constexpr int kWmmaK = 16; +// Block tile. 8 warps (256 threads) cover a 2x4 grid of 16x16 sub-tiles. +constexpr int kBM = 32; +constexpr int kBN = 64; +constexpr int kBK = 16; // WMMA K-step; the GEMM variant is the M>4 fallback +constexpr int kThreads = 256; // (kBM/16)*(kBN/16) = 2*4 = 8 warps + +typedef wmma::fragment WmmaA; +typedef wmma::fragment WmmaB; +typedef wmma::fragment WmmaC; + +__device__ __forceinline__ float e4m3_to_float(uint8_t b) { + // Branchless-ish fast decode: E4M3 -> fp32 by laying the bits out directly. + // E4M3 value = (-1)^s * 2^(e-7) * (1 + m/8), bias 7. fp32 shares the exponent + // bias 127 with bf16, so the normal case is just fp32_exp = e-7+127 = e+120 and + // the 3 mantissa bits placed at the top of fp32's 23-bit mantissa (<<20). + // Subnormal (e==0) and the NaN pattern (e==15,m==7) are handled with cheap + // selects. This avoids exp2f (measured ~2x faster than exp2f decode, ~815 GB/s). + const uint32_t s = (b >> 7) & 1; + const uint32_t e = (b >> 3) & 0xF; + const uint32_t m = b & 0x7; + const uint32_t fbits_normal = (s << 31) | ((e + 120) << 23) | (m << 20); + const float v_normal = __uint_as_float(fbits_normal); // sign already encoded + const float v_sub = (s ? -1.0f : 1.0f) * (static_cast(m) * 0.001953125f); // m * 2^-9 + // Subnormal (e==0) takes the sign; NaN (e==15,m==7) clamps to max finite. + return (e == 0) ? v_sub : ((e == 15 && m == 7) ? 448.0f : v_normal); +} + +template +__device__ __forceinline__ __half to_half(T v); + +template <> +__device__ __forceinline__ __half to_half(c10::Half v) { + return __float2half(static_cast(v)); +} +template <> +__device__ __forceinline__ __half to_half(c10::BFloat16 v) { + return __float2half(static_cast(v)); +} +template <> +__device__ __forceinline__ __half to_half(float v) { + return __float2half(v); +} +template <> +__device__ __forceinline__ __half to_half(double v) { + return __float2half(static_cast(v)); +} + +// WMMA GEMM fallback for M > 4 (correct but not tuned for tiny M). The decode +// fast path for M <= 4 is e4m3_gemv_kernel below. +template +__global__ void e4m3_linear_kernel( + const T* __restrict__ x, // (M, K) row-major + const uint8_t* __restrict__ w, // (N, K) row-major, uint8 E4M3 + const float* __restrict__ scale, // per-tensor scalar + c10::BFloat16* __restrict__ y, // (M, N) bf16 row-major + int M, + int N, + int K, + int ldx, + int ldw, + int ldy) { + const int n0 = blockIdx.x * kBN; + const int m0 = blockIdx.y * kBM; + const int tid = threadIdx.x; + const int warp = tid / 32; + + // Sub-tile owned by this warp (4x4 grid of 16x16 tiles per block). + const int wm = warp / (kBN / kWmmaN); + const int wn = warp % (kBN / kWmmaN); + + __shared__ __half A_smem[kBM][kBK]; + __shared__ __half B_smem[kBK][kBN]; + __shared__ float C_smem[kBM][kBN]; + + WmmaA a_frag; + WmmaB b_frag; + WmmaC c_frag; + wmma::fill_fragment(c_frag, 0.0f); + + for (int k0 = 0; k0 < K; k0 += kBK) { + // Load A tile (kBM x kBK): x[m][k] -> fp16. Threads index the contiguous k + // dim fastest, so a warp reads consecutive addresses (coalesced). +#pragma unroll + for (int i = tid; i < kBM * kBK; i += kThreads) { + const int r = i / kBK; + const int kk = i % kBK; + const int gm = m0 + r; + const int gk = k0 + kk; + if (gm < M && gk < K) { + A_smem[r][kk] = to_half(x[gm * ldx + gk]); + } else { + A_smem[r][kk] = __float2half(0.0f); + } + } + // Load B tile (kBN x kBK): decode uint8 E4M3 w[n][k] -> fp16, into B_smem[k][n]. + // Threads index the contiguous k dim fastest so the weight bytes stream in + // coalesced (this is the memory-bound path; the earlier stride-N indexing + // fetched one cache line per element). +#pragma unroll + for (int i = tid; i < kBK * kBN; i += kThreads) { + const int kk = i % kBK; // k index (contiguous) + const int c = i / kBK; // n index + const int gk = k0 + kk; + const int gn = n0 + c; + if (gn < N && gk < K) { + B_smem[kk][c] = __float2half(e4m3_to_float(w[gn * ldw + gk])); + } else { + B_smem[kk][c] = __float2half(0.0f); + } + } + __syncthreads(); + +#pragma unroll + for (int kk = 0; kk < kBK; kk += kWmmaK) { + wmma::load_matrix_sync(a_frag, &A_smem[wm * kWmmaM][kk], kBK); + wmma::load_matrix_sync(b_frag, &B_smem[kk][wn * kWmmaN], kBN); + wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); + } + __syncthreads(); + } + + // Apply the per-tensor scalar scale to the fp32 accumulator, then store. + const float s = *scale; +#pragma unroll + for (int i = 0; i < c_frag.num_elements; ++i) { + c_frag.x[i] *= s; + } + wmma::store_matrix_sync(&C_smem[wm * kWmmaM][wn * kWmmaN], c_frag, kBN, wmma::mem_row_major); + __syncthreads(); + +#pragma unroll + for (int i = tid; i < kBM * kBN; i += kThreads) { + const int r = i / kBN; + const int c = i % kBN; + const int gm = m0 + r; + const int gn = n0 + c; + if (gm < M && gn < N) { + y[gm * ldy + gn] = static_cast(C_smem[r][c]); + } + } +} + +// Small-M / memory-bound decode path. Each warp owns one output row `n` and +// streams w[n,:] with 32-byte-coalesced reads along the contiguous k dim, reusing +// the (kM x K) activation staged in shared, then warp-reduces with shuffles. This +// avoids the tensor-core GEMM's M-waste and register pressure that capped it at +// ~60 GB/s for the tiny-M decode case. kM is the batch of rows (== M). +constexpr int kWarps = 8; // warps per block -> 256 threads +constexpr int kGemvThreads = kWarps * 32; +constexpr int kGemvMaxM = 4; // small-M decode path; larger M uses the WMMA GEMM fallback + + +template +__global__ void e4m3_gemv_kernel( + const T* __restrict__ x, // (kM, K) row-major, kM == M + const uint8_t* __restrict__ w, // (N, K) uint8, contiguous k + const float* __restrict__ scale, + c10::BFloat16* __restrict__ y, // (kM, N) + int N, + int K, + int ldw, + int ldy) { + extern __shared__ char smem_raw[]; + // Small M (decode): stage the activation as float so the inner loop has no + // bf16->float cast (the biggest compute cost at M=1). Larger M keeps bf16 x so + // shared is small and occupancy stays high (float x at M>=3 needs >48KB and + // drops occupancy, which regresses M=4). + using XS_t = typename std::conditional::type; + XS_t* const xs = reinterpret_cast(smem_raw); // [kM * K] + for (int i = threadIdx.x; i < kM * K; i += kGemvThreads) { + xs[i] = static_cast(x[i]); + } + __syncthreads(); + + const int warp = threadIdx.x / 32; + const int lane = threadIdx.x % 32; + const int n = blockIdx.x * kWarps + warp; + if (n >= N) return; + + const uint8_t* __restrict__ wrow = w + static_cast(n) * ldw; + float acc[kM]; +#pragma unroll + for (int m = 0; m < kM; ++m) acc[m] = 0.0f; + + // Vectorized phase: each lane handles 8 consecutive k per step (a uint64 load + // of the uint8 weight), giving a full 256-byte coalesced access per warp (2 + // cache lines) and 8x the memory-level parallelism of the scalar loop. Only + // safe when the row base is 8-byte aligned (K % 8 == 0). Scalar fallback all K. + if ((K & 7) == 0) { + const int k8 = K >> 3; +#pragma unroll 4 + for (int g = lane; g < k8; g += 32) { + const int k = g << 3; + const uint64_t wq = *reinterpret_cast(wrow + k); + const float wv0 = e4m3_to_float(static_cast(wq & 0xFF)); + const float wv1 = e4m3_to_float(static_cast((wq >> 8) & 0xFF)); + const float wv2 = e4m3_to_float(static_cast((wq >> 16) & 0xFF)); + const float wv3 = e4m3_to_float(static_cast((wq >> 24) & 0xFF)); + const float wv4 = e4m3_to_float(static_cast((wq >> 32) & 0xFF)); + const float wv5 = e4m3_to_float(static_cast((wq >> 40) & 0xFF)); + const float wv6 = e4m3_to_float(static_cast((wq >> 48) & 0xFF)); + const float wv7 = e4m3_to_float(static_cast((wq >> 56) & 0xFF)); +#pragma unroll + for (int m = 0; m < kM; ++m) { + const XS_t* const xm = xs + static_cast(m) * K + k; + acc[m] += static_cast(xm[0]) * wv0 + static_cast(xm[1]) * wv1 + + static_cast(xm[2]) * wv2 + static_cast(xm[3]) * wv3 + + static_cast(xm[4]) * wv4 + static_cast(xm[5]) * wv5 + + static_cast(xm[6]) * wv6 + static_cast(xm[7]) * wv7; + } + } + } else { + for (int k = lane; k < K; k += 32) { + const float wv = e4m3_to_float(wrow[k]); +#pragma unroll + for (int m = 0; m < kM; ++m) { + acc[m] += xs[static_cast(m) * K + k] * wv; + } + } + } + + const float s = *scale; +#pragma unroll + for (int m = 0; m < kM; ++m) { +#pragma unroll + for (int off = 16; off > 0; off >>= 1) { + acc[m] += __shfl_down_sync(0xffffffffu, acc[m], off); + } + if (lane == 0) { + y[static_cast(m) * ldy + n] = static_cast(acc[m] * s); + } + } +} + +} // namespace areno_accel + +// Launch the small-M GEMV (direct-load kernel). No cudaFuncSetAttribute here: +// the float-xs (M<=2) and bf16-xs (M>=3) gemv shared sizes are all < 48 KB, so +// the default dynamic-shared limit suffices (calling cudaFuncSetAttribute on +// every launch measurably slowed the fast M=1 kernel). +template +void areno_launch_e4m3_gemv( + const T* x, + const uint8_t* w, + const float* scale, + c10::BFloat16* y, + int N, + int K, + int ldw, + int ldy, + cudaStream_t stream, + size_t shmem_direct, + dim3 grid, + dim3 block) { + areno_accel::e4m3_gemv_kernel<<>>( + x, w, scale, y, N, K, ldw, ldy); +} + +torch::Tensor areno_e4m3_linear_forward_cuda( + torch::Tensor input, + torch::Tensor w_u8, + torch::Tensor scale, + c10::optional out) { + TORCH_CHECK(input.is_cuda(), "areno_e4m3_linear input must be CUDA"); + TORCH_CHECK(w_u8.is_cuda(), "areno_e4m3_linear weight must be CUDA"); + TORCH_CHECK(scale.is_cuda(), "areno_e4m3_linear scale must be CUDA"); + TORCH_CHECK(w_u8.scalar_type() == at::kByte, "areno_e4m3_linear weight must be uint8 (E4M3 bytes)"); + TORCH_CHECK(scale.scalar_type() == at::kFloat, "areno_e4m3_linear scale must be float32"); + TORCH_CHECK(input.dim() >= 2, "areno_e4m3_linear input must have at least 2 dims"); + TORCH_CHECK(w_u8.dim() == 2, "areno_e4m3_linear weight must be 2D"); + TORCH_CHECK(input.size(-1) == w_u8.size(1), "areno_e4m3_linear input/weight K mismatch"); + TORCH_CHECK(w_u8.is_contiguous(), "areno_e4m3_linear weight must be contiguous"); + TORCH_CHECK(input.is_contiguous(), "areno_e4m3_linear input must be contiguous"); + TORCH_CHECK(scale.numel() == 1, "areno_e4m3_linear scale must be scalar"); + + auto out_shape = input.sizes().vec(); + out_shape.back() = w_u8.size(0); + // Reuse a caller-provided bf16 CUDA output buffer when it already matches, so + // the hot decode path does not pay a per-token torch::empty. + torch::Tensor output; + if (out.has_value() && out->is_cuda() && out->is_contiguous() && + out->scalar_type() == at::kBFloat16 && out->sizes() == out_shape) { + output = out.value(); + } else { + output = torch::empty(out_shape, input.options().dtype(at::kBFloat16)); + } + + int64_t K = input.size(-1); + int64_t M = input.numel() / K; + int64_t N = w_u8.size(0); + int ldx = static_cast(input.stride(-2)); + int ldw = static_cast(w_u8.stride(0)); + int ldy = static_cast(output.stride(-2)); + + const at::cuda::OptionalCUDAGuard guard(device_of(input)); + auto stream = at::cuda::getCurrentCUDAStream(); + + // Small-M decode path uses the memory-streaming GEMV kernel (coalesced weight + // reads, no M-waste); larger M falls back to the WMMA GEMM kernel (correct but + // not tuned for tiny M). Both are forward-only (E4M3). + const bool use_gemv = (M >= 1 && M <= areno_accel::kGemvMaxM); + const uint8_t* wptr = static_cast(w_u8.data_ptr()); + c10::BFloat16* yptr = output.data_ptr(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, input.scalar_type(), "areno_e4m3_linear", [&] { + const scalar_t* xptr = input.data_ptr(); + if (use_gemv) { + // Small M (decode) stages x as float (no bf16->float cast in the inner + // loop, the biggest compute cost); M>=3 keeps bf16 x so shared stays + // small and M=4 occupancy is not lost. + dim3 grid(static_cast((N + areno_accel::kWarps - 1) / areno_accel::kWarps), 1); + dim3 block(areno_accel::kGemvThreads); + const int Ni = static_cast(N); + const int Ki = static_cast(K); + const size_t shmem1 = static_cast(K) * sizeof(float); // M=1 float + const size_t shmem2 = 2 * static_cast(K) * sizeof(float); // M=2 float + const size_t shmem3 = 3 * static_cast(K) * sizeof(scalar_t); // M=3 bf16 + const size_t shmem4 = 4 * static_cast(K) * sizeof(scalar_t); // M=4 bf16 + switch (M) { + case 1: + areno_launch_e4m3_gemv(xptr, wptr, scale.data_ptr(), yptr, Ni, Ki, ldw, ldy, + stream, shmem1, grid, block); + break; + case 2: + areno_launch_e4m3_gemv(xptr, wptr, scale.data_ptr(), yptr, Ni, Ki, ldw, ldy, + stream, shmem2, grid, block); + break; + case 3: + areno_launch_e4m3_gemv(xptr, wptr, scale.data_ptr(), yptr, Ni, Ki, ldw, ldy, + stream, shmem3, grid, block); + break; + case 4: + areno_launch_e4m3_gemv(xptr, wptr, scale.data_ptr(), yptr, Ni, Ki, ldw, ldy, + stream, shmem4, grid, block); + break; + } + } else { + dim3 grid(static_cast((N + areno_accel::kBN - 1) / areno_accel::kBN), + static_cast((M + areno_accel::kBM - 1) / areno_accel::kBM)); + dim3 block(areno_accel::kThreads); + areno_accel::e4m3_linear_kernel<<>>( + xptr, wptr, scale.data_ptr(), yptr, static_cast(M), static_cast(N), + static_cast(K), ldx, ldw, ldy); + } + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} diff --git a/areno/accel/csrc/extension.cpp b/areno/accel/csrc/extension.cpp index 6cee02e9..352bc5ae 100644 --- a/areno/accel/csrc/extension.cpp +++ b/areno/accel/csrc/extension.cpp @@ -11,6 +11,7 @@ torch::Tensor areno_d_silu_cuda(torch::Tensor grad_output, torch::Tensor input); torch::Tensor areno_d_sigmoid_cuda(torch::Tensor grad_output, torch::Tensor output); torch::Tensor areno_d_softplus_cuda(torch::Tensor grad_output, torch::Tensor input); torch::Tensor areno_linear_forward_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, bool use_bias); +torch::Tensor areno_e4m3_linear_forward_cuda(torch::Tensor input, torch::Tensor w_u8, torch::Tensor scale, c10::optional out); std::vector areno_linear_backward_cuda( torch::Tensor grad_output, torch::Tensor input, @@ -199,6 +200,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("areno_d_sigmoid", &areno_d_sigmoid_cuda, "ARENO sigmoid backward"); m.def("areno_d_softplus", &areno_d_softplus_cuda, "ARENO softplus backward"); m.def("areno_linear_forward", &areno_linear_forward_cuda, "ARENO linear forward"); + m.def("areno_e4m3_linear_forward", &areno_e4m3_linear_forward_cuda, "ARENO E4M3 fused dequant-linear forward", + pybind11::arg("input"), pybind11::arg("w_u8"), pybind11::arg("scale"), pybind11::arg("out") = pybind11::none()); m.def("areno_linear_backward", &areno_linear_backward_cuda, "ARENO linear backward"); m.def("areno_causal_attention_forward", &areno_causal_attention_forward_cuda, "ARENO causal attention forward"); m.def("areno_causal_attention_backward", &areno_causal_attention_backward_cuda, "ARENO causal attention backward"); diff --git a/areno/accel/kernels/e4m3_cuda.py b/areno/accel/kernels/e4m3_cuda.py new file mode 100644 index 00000000..2de7f5e2 --- /dev/null +++ b/areno/accel/kernels/e4m3_cuda.py @@ -0,0 +1,66 @@ +"""GPU E4M3 fused dequant-linear via the areno.accel CUDA extension. + +Reads the 1-byte uint8 E4M3 weight payload and decodes it to fp16 in-kernel +(the branchless bit-trick in ``areno/accel/csrc/e4m3_linear.cu``) so the weight +bytes pulled from HBM are halved vs the bf16 path — the memory-bandwidth benefit +from RFC 0001. The reference decode math is validated against +``torch.float8_e4m3fn`` in ``tests/test_e4m3_decode_cpu.py``. This is a +standalone Ampere (cc 8.0) decode kernel; it is **not** wired into +``_areno_linear_forward`` (the model FP8 decode path uses the E5M2 Triton kernel +in ``fp8_linear.py`` on A100 and the native ``torch._scaled_mm`` on Hopper/H20). + +Forward-only: E4M3 has no backward, so this is a decode/inference operator and +must not be wired into a training graph. +""" + +from __future__ import annotations + +import torch + +from areno.accel._extension import extension as _extension + + +def quantized_e4m3_linear_cuda( + x: torch.Tensor, + w_u8: torch.Tensor, + scale: torch.Tensor, + *, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """``y = x @ (e4m3(w) * scale)^T`` reading the uint8 E4M3 weight directly. + + Args: + x: activation, shape (..., K), bf16 (or fp16). + w_u8: weight as packed E4M3 bytes, shape (N, K), ``torch.uint8``. + scale: per-tensor scalar (float32), shape (). + out: optional pre-allocated bf16 output (..., N); reused if given. + Returns: + bf16 (..., N); ``out`` if provided. + """ + if not (x.is_cuda and w_u8.is_cuda and scale.is_cuda): + raise RuntimeError("quantized_e4m3_linear_cuda requires CUDA inputs") + if w_u8.dtype != torch.uint8: + raise RuntimeError(f"quantized_e4m3_linear_cuda weight must be uint8, got {w_u8.dtype}") + x2 = x.contiguous() + w2 = w_u8.contiguous() + s2 = scale.contiguous() + # Pass the caller's buffer through so the extension writes into it directly and + # the hot decode path avoids a per-call torch::empty + copy_. + return _extension().areno_e4m3_linear_forward(x2, w2, s2, out) + + +def quantize_weight_e4m3(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a bf16 weight to E4M3 and return (uint8 bytes, scalar scale). + + Uses a per-tensor scale so -448, 448 (or symmetric) maps to the grid; matches + the scale semantics in ``areno.engine.quantization``. + """ + from areno.engine.quantization import quantize_to_fp8 + + fp8, scale = quantize_to_fp8(weight, group_size=-1) + return fp8.view(torch.uint8).contiguous(), scale.contiguous() + + +def dequant_e4m3_bf16(w_u8: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Materialize a bf16 weight from an E4M3 uint8 payload + scale.""" + return (w_u8.view(torch.float8_e4m3fn).float() * scale.float()).to(torch.bfloat16) diff --git a/areno/accel/kernels/fp8_linear.py b/areno/accel/kernels/fp8_linear.py new file mode 100644 index 00000000..3a67b4ea --- /dev/null +++ b/areno/accel/kernels/fp8_linear.py @@ -0,0 +1,134 @@ +"""Triton FP8 (W8A16) dequant-linear for the areno.accel surface. + +Runs ``y = x @ (w_fp8 * scale)`` (``x`` bf16, ``w_fp8`` FP8-E5M2, ``scale`` per-tensor +scalar) with a Triton matmul that reads the FP8 weight directly — the memory-bound +decode benefit from RFC 0001. Tunings reflect the measured best on A100 (~1.62x vs +bf16 at decode shapes): ``num_stages`` pipelining, large ``BLOCK_K``, ``.cg`` on the +streamed weight, and the key trick of applying the per-tensor scale to the +accumulator **after** the dot (since ``(x @ w_fp8) * scale == x @ (w_fp8 * scale)``), +which avoids a per-element dequant in the inner loop. + +This is an inference-side (decode) op. A backwards pass is out of scope for now; +training integration is a separate piece. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _quantized_fp8_linear_kernel( + x_ptr, + w_ptr, + scale_ptr, + y_ptr, + M, + N, + K, + sx, + sw, + sy, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, +): + pm = tl.program_id(0) + pn = tl.program_id(1) + rm = pm * BM + tl.arange(0, BM) + rn = pn * BN + tl.arange(0, BN) + acc = tl.zeros((BM, BN), dtype=tl.float32) + scale = tl.load(scale_ptr) + for k in range(0, tl.cdiv(K, BK)): + rk = k * BK + tl.arange(0, BK) + am = (rm[:, None] < M) & (rk[None, :] < K) + bm_ = (rn[:, None] < N) & (rk[None, :] < K) + a = tl.load(x_ptr + rm[:, None] * sx + rk[None, :], mask=am, other=0.0) + b = tl.load(w_ptr + rn[:, None] * sw + rk[None, :], mask=bm_, other=0.0, cache_modifier=".cg") + # FP8->fp16 tensor-core dot; per-tensor scale applied once after the dot. + acc += tl.dot(a.to(tl.float16), tl.trans(b.to(tl.float16))) + acc = acc * scale + tl.store(y_ptr + rm[:, None] * sy + rn[None, :], acc.to(tl.bfloat16), mask=(rm[:, None] < M) & (rn[None, :] < N)) + + +def quantized_fp8_linear( + x: torch.Tensor, + w_fp8: torch.Tensor, + scale: torch.Tensor, + *, + out: torch.Tensor | None = None, + block_m: int = 32, + block_n: int = 128, + block_k: int = 128, + num_stages: int = 4, + num_warps: int = 4, +) -> torch.Tensor: + """FP8(E5M2) W8A16 linear forward: ``y = x @ (w_fp8 * scale)``. + + Args: + x: activation, shape (batch, K), bf16. + w_fp8: weight in FP8-E5M2, shape (N, K). + scale: per-tensor scalar scale (float32), shape (). + out: optional pre-allocated output (batch, N); reused if given, so the + hot decode path does not pay a per-step ``torch.empty``. + Returns: + bf16 (batch, N); ``out`` if provided. + """ + if not (x.is_cuda and w_fp8.is_cuda and scale.is_cuda): + raise RuntimeError("quantized_fp8_linear requires CUDA inputs") + M, K = x.shape + N, _ = w_fp8.shape + out_dtype = x.dtype + y = ( + out + if (out is not None and out.shape == (M, N) and out.device == x.device) + else torch.empty((M, N), device=x.device, dtype=out_dtype) + ) + if w_fp8.dtype != torch.float8_e5m2: + # Many checkpoints store E4M3; Triton on Ampere only accepts E5M2, so + # convert to the grid the kernel can consume. + w_fp8 = w_fp8.to(torch.float8_e5m2) + grid = (triton.cdiv(M, block_m), triton.cdiv(N, block_n)) + _quantized_fp8_linear_kernel[grid]( + x, + w_fp8, + scale, + y, + M, + N, + K, + x.stride(0), + w_fp8.stride(0), + y.stride(0), + block_m, + block_n, + block_k, + num_stages=num_stages, + num_warps=num_warps, + ) + return y + + +def dequantize_fp8_weight(w_fp8: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Materialize a bf16 weight from an FP8 weight + per-tensor scale.""" + return (w_fp8.float() * scale.to(torch.float32)).to(torch.bfloat16) + + +def mark_fp8_weight(weight: torch.Tensor, *, group_size: int = -1, fp8_dtype=torch.float8_e5m2) -> torch.Tensor: + """Quantize a bf16 weight in place and stash the FP8 payload for the linear hook. + + Sets ``weight._areno_fp8`` (FP8 grid tensor) and ``weight._areno_fp8_scale`` + (per-tensor scale) so ``_areno_linear_forward`` dispatches to + :func:`quantized_fp8_linear`. Returns ``weight``. ``group_size <= 0`` means a + single per-tensor scale (the current kernel path). + """ + f = weight.detach().float() + max_v = 57344.0 if fp8_dtype is torch.float8_e5m2 else 448.0 + scale = (f.abs().amax() / max_v).to(torch.float32).reshape(()) + scale = torch.where(scale > 0, scale, torch.ones_like(scale)) + q = (f / scale).clamp(-max_v, max_v).to(fp8_dtype) + weight._areno_fp8 = q + weight._areno_fp8_scale = scale + return weight diff --git a/areno/engine/config.py b/areno/engine/config.py index f79d6de4..eccb5c49 100644 --- a/areno/engine/config.py +++ b/areno/engine/config.py @@ -162,6 +162,8 @@ class ModelConfig: qk_norm: bool = True v_norm: bool = False dtype: torch.dtype = torch.bfloat16 + quant_method: Literal["none", "fp8", "int4"] = "none" + quant_group_size: int = 128 hidden_act: str = "silu" layer_types: tuple[str, ...] | None = None sliding_window: int | None = None diff --git a/areno/engine/layers/linear.py b/areno/engine/layers/linear.py index b95c3200..77c9d8a0 100644 --- a/areno/engine/layers/linear.py +++ b/areno/engine/layers/linear.py @@ -44,6 +44,46 @@ def mark_tensor_parallel_parameter( setattr(param, "tp_grad_allreduce", tp_grad_allreduce) +class QuantizedLinear(nn.Module): + """W8A16 FP8 dequant-forward reference module (RFC 0001, M1). + + Quantizes its bf16 weight to FP8 on demand and runs a dequant-forward matmul + (``x @ dequant(w_q)^T``), so the exact FP8 scale semantics are testable on CPU. + This is the correctness reference for the fused FP8 dequant-linear kernels; + the production decode path uses the areno.accel kernels / ``_areno_linear_forward`` + directly. FP8 (E4M3/E5M2) has no backward, so this is a decode-only reference. + """ + + def __init__( + self, in_features: int, out_features: int, *, group_size: int = -1, dtype: torch.dtype = torch.bfloat16 + ): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.group_size = group_size + self.dtype = dtype + self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype)) + self._q = None # FP8 weight (fp8 grid or snapped float) + self._scale = None + self.reset_parameters() + + def reset_parameters(self) -> None: + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def _requantize(self) -> None: + from areno.engine.quantization import quantize_to_fp8 + + self._q, self._scale = quantize_to_fp8(self.weight.data, self.group_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from areno.engine.quantization import dequant_fp8 + + if self._q is None: + self._requantize() + dq = dequant_fp8(self._q, self._scale, self.group_size).to(self.dtype) + return F.linear(x, dq) + + def _shard_range(size: int, rank: int, world_size: int) -> tuple[int, int]: """Compute ``[start, end)`` of the local shard for an even partition.""" @@ -292,6 +332,25 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def _areno_linear_forward(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor: """Single entry point so all parallel linears share the areno.accel matmul.""" + # FP8 path (RFC 0001): a weight carrying an FP8 payload (weight._areno_fp8 + + # weight._areno_fp8_scale) routes through the Triton W8A16 dequant-linear. + # Opt-in and backward-compatible: unmarked weights take the existing path. + fp8 = getattr(weight, "_areno_fp8", None) + if fp8 is not None: + scale = weight._areno_fp8_scale + from areno.accel.kernels.fp8_linear import quantized_fp8_linear + + # Decode/prefill may hand a 3-D (1, seq, hidden) activation; flatten to 2-D + # for the kernel then restore the leading dims. + out_ndim = x.ndim + xx = x.reshape(-1, x.shape[-1]) if x.ndim > 2 else x + out = quantized_fp8_linear(xx, fp8, scale) + if bias is not None: + out = out + bias + if out_ndim > 2: + out = out.reshape(*x.shape[:-1], out.shape[-1]) + return out + if x.ndim >= 3 and torch.is_grad_enabled(): return F.linear(x, weight, bias) return areno_linear(x, weight, bias) diff --git a/areno/engine/modeling.py b/areno/engine/modeling.py index b68413a3..e670d20c 100644 --- a/areno/engine/modeling.py +++ b/areno/engine/modeling.py @@ -20,6 +20,31 @@ def param_grad(param: torch.nn.Parameter) -> torch.Tensor | None: return param.grad +def quantize_model_weights_fp8(model: torch.nn.Module) -> int: + """Mark every TP-parallel linear weight as FP8 for the decode path. + + Calls ``mark_fp8_weight`` on the weight of each column/merged/row-parallel + linear so ``_areno_linear_forward`` routes those through the FP8 W8A16 + kernel (RFC 0001). Decode/inference-only — the Triton FP8 kernel has no + backward, so this must not be used in a training step. Returns the count. + """ + from areno.accel.kernels.fp8_linear import mark_fp8_weight + from areno.engine.layers.linear import ColumnParallelLinear, MergedColumnParallelLinear, RowParallelLinear + + count = 0 + for module in model.modules(): + if isinstance(module, (ColumnParallelLinear, MergedColumnParallelLinear, RowParallelLinear)): + weight = getattr(module, "weight", None) + if weight is not None and getattr(weight, "ndim", 0) == 2 and weight.numel() > 0 and weight.is_cuda: + mark_fp8_weight(weight) + count += 1 + if count: + import logging + + logging.getLogger("areno").warning("FP8 quantized %d linear weights (decode path)", count) + return count + + def build_model_on_device(config: EngineConfig, device: torch.device) -> torch.nn.Module: """Construct the model directly on `device` under the configured dtype.""" diff --git a/areno/engine/quantization.py b/areno/engine/quantization.py new file mode 100644 index 00000000..56d2893a --- /dev/null +++ b/areno/engine/quantization.py @@ -0,0 +1,88 @@ +"""FP8 weight quantization helpers (W8A16, dequant-forward reference). + +These helpers implement the FP8 (E4M3) quantize/dequant math in a dtype-agnostic +way so they run on CPU (for unit tests) as well as on GPU. They are the +correctness reference for the FP8 vertical slice (RFC 0001, M1): the intended +memory-bandwidth speedup comes from a fused FP8-dequant matmul that reads the +1-byte weights directly; the dequant-forward path here establishes numerical +correctness and the exact scale semantics before that kernel lands. + +Note on FP8 grids: this module (and ``QuantizedLinear``) uses **E4M3** as the +reference grid. The model's runtime FP8 path in ``fp8_linear.py`` uses **E5M2**, +because on Ampere (A100) Triton only accepts the E5M2 grid and rejects E4M3 +(``fp8e4nv``); see RFC 0001 §4.6. The fused E4M3 dequant kernels live in +``areno/accel/csrc/e4m3_linear.cu``. ``compute_fp8_scale``/``quantize_to_fp8`` +are dtype-agnostic apart from the constant ``_FP8_E4M3_MAX``=448. + +This is deliberately small and dependency-free (only torch). +""" + +from __future__ import annotations + +import torch + +# E4M3 has 3 exponent bits + 4 mantissa bits + sign. Max finite magnitude is +# 448.0; magnitude range is [2^-9, 448]. We clamp to the representable range. +_FP8_E4M3_MAX = 448.0 + + +def compute_fp8_scale(weight: torch.Tensor, group_size: int = -1) -> torch.Tensor: + """Compute the FP8 scale for `weight`, per-tensor or per-group. + + ``group_size <= 0`` (default) uses a single per-tensor scale. Otherwise the + weight's last dim is split into groups of ``group_size`` and a scale is + computed per group. Scale is chosen so the max magnitude in the group maps + to ``_FP8_E4M3_MAX``; a zero group falls back to 1.0 to avoid inf/nan. + """ + w = weight.detach().float() + if group_size is None or group_size <= 0: + amax = w.abs().amax() + scale = amax / _FP8_E4M3_MAX + scale = torch.where(scale > 0, scale, torch.ones_like(scale)) + return scale + n = w.numel() + g = int(group_size) + if n % g != 0: + raise ValueError(f"weight numel {n} must be divisible by group_size {g}") + flat = w.reshape(-1, g) + amax = flat.abs().amax(dim=1, keepdim=True) + scale = amax / _FP8_E4M3_MAX + scale = torch.where(scale > 0, scale, torch.ones_like(scale)) + return scale.reshape(-1) + + +def quantize_to_fp8(weight: torch.Tensor, group_size: int = -1) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize `weight` (bf16/fp32) to FP8 grid; returns (fp8, scale). + + ``fp8`` is returned in ``torch.float8_e4m3fn`` when this torch build supports + it; otherwise it is returned as a float tensor already snapped to the FP8 + grid (so the math is testable on CPU). ``scale`` is float32, shaped per the + group layout (scalar for per-tensor, ``(n/g,)`` for per-group). + """ + scale = compute_fp8_scale(weight, group_size) + w = weight.detach().to(torch.float32) + if group_size is None or group_size <= 0: + snapped = torch.clamp((w / scale).round(), -_FP8_E4M3_MAX, _FP8_E4M3_MAX) + else: + low = w.reshape(-1, int(group_size)) + scale2 = scale.view(-1, 1) + snapped = torch.clamp((low / scale2).round(), -_FP8_E4M3_MAX, _FP8_E4M3_MAX) + snapped = snapped.reshape_as(w) + if hasattr(torch, "float8_e4m3fn"): + try: + return snapped.to(torch.float8_e4m3fn), scale + except (RuntimeError, TypeError, ValueError): + pass + return snapped, scale + + +def dequant_fp8(fp8: torch.Tensor, scale: torch.Tensor, group_size: int = -1) -> torch.Tensor: + """Dequantize FP8 weights back to float; opposite of :func:`quantize_to_fp8`. + + ``scale`` is per-tensor (scalar) or per-group (``(n/g,)``) and broadcasts + down to the weight shape. Output dtype is float32 (caller casts). + """ + q = fp8.float() + if group_size is None or group_size <= 0: + return q * scale + return (q.reshape(-1, int(group_size)) * scale.view(-1, 1)).reshape_as(q) diff --git a/areno/engine/worker.py b/areno/engine/worker.py index c08afe96..089a6a7f 100644 --- a/areno/engine/worker.py +++ b/areno/engine/worker.py @@ -15,6 +15,7 @@ from __future__ import annotations +import os import queue import torch @@ -27,7 +28,13 @@ from areno.engine.data import RolloutOutput from areno.engine.data.sampling import _truncate_generated from areno.engine.inference import InferCacheSpec, InferenceManager -from areno.engine.modeling import build_model_on_device, build_optimizer, configure_multimodal_training, param_grad +from areno.engine.modeling import ( + build_model_on_device, + build_optimizer, + configure_multimodal_training, + param_grad, + quantize_model_weights_fp8, +) from areno.engine.parallel.context import get_tp_context from areno.engine.policy_sync import policy_plan_metadata, transfer_policy_weights from areno.engine.protocol import ( @@ -65,6 +72,17 @@ def __init__(self, config: EngineConfig): self.model = build_model_on_device(config, self.device) if config.model_path is not None and not config.dummy_load: load_model_weights(self.model, config.model, config.model_path) + if config.model.quant_method == "fp8" or os.environ.get("ARENO_QUANT_FP8"): + # Decode-only FP8 quantization (RFC 0001): the FP8 W8A16 kernel has no + # backward, so it cannot update a model in a training step. Guard against + # silently running a train loop with zero gradient. + if config.role == "train": + raise RuntimeError( + "quant_method='fp8' (decode-only) cannot be used in a train worker: " + "the FP8 kernel has no backward. Use quant_method='none' for training " + "or run the rollout/inference role." + ) + quantize_model_weights_fp8(self.model) configure_multimodal_training(self.model, config.optimizer, trainable=config.role == "train") self.adapter_registry = ( initialize_lora(self.model, config.lora, seed=config.lora_seed) if config.lora is not None else None diff --git a/docs/rfcs/0001-weight-quantization-fp8-int4.md b/docs/rfcs/0001-weight-quantization-fp8-int4.md new file mode 100644 index 00000000..15d8c81f --- /dev/null +++ b/docs/rfcs/0001-weight-quantization-fp8-int4.md @@ -0,0 +1,341 @@ +# RFC 0001 — FP8 / INT4 Weight Quantization for AReno (CUDA Train & Decode) + +- **Metadata** + - **Status:** Draft (for comment) + - **Author:** @CAICAIIs + - **Date:** 2026-08-25 + - **Affected subsystems:** `areno/models/*`, `areno/engine/layers/linear.py`, `areno/accel/`, + `areno/models/registry.py`, `areno/engine/checkpoints/io.py`, `areno/engine/config.py` (`ModelConfig`), + `areno/cli/train.py`, `areno/cli/serve.py` + - **Reviewers:** open + +--- + +## 1. Background & Problem + +AReno targets fast, self-contained single-node post-training. On a single 8×A100 node we measured the +following for a GRPO rollout with native attention and world=4 / tp=4: + +| model | weights re-read per generated token | decode tokens/s | weight-read bandwidth | compute utilization | +|---|---|---|---|---| +| Qwen3-0.6B | ~1.2 GiB | ~1450 | ~1.7 TB/s | ~1.7 TFLOPs/s (overhead/latency-dominated) | +| Qwen3-8B | ~16 GiB | ~390 | ~6.2 TB/s ≈ HBM aggregate limit | ~0.5% (**memory-bandwidth-bound**) | + +**Problem statement.** For models that actually stress a single node (≈4B parameters and above), the +decode path is **memory-bandwidth-bound**: every generated token re-reads the full weight tensor from HBM, +while the compute units sit largely idle. The dominant, well-understood mitigation is to shrink the number +of bytes read per weight — i.e. **weight quantization**. AReno today has **no weight quantization**: the only +quantization-related code is 8-bit optimizer moments (`areno/engine/optim/adamw_8bit.py`) and the MLX +optimizer path; all linear layers run plain bf16 (`areno/engine/layers/linear.py`, +`areno/accel/csrc/linear.cu`). + +This RFC proposes a backward-compatible, opt-in FP8 / INT4 weight-quantization capability for the CUDA +backend, TP-aware, with a correctness gate and a quantifiable token/s benchmark. + +--- + +## 2. Objectives & Scope + +**Objectives (v1, CUDA backend).** +1. Reduce bytes read per weight so decode token/s increases measurably (target: near-linear with the byte + reduction for weights). +2. Preserve correctness: a trained step under quantization stays within a documented tolerance of the full-precision + baseline (loss, grad norm, advantage mean). +3. Remain fully backward-compatible and opt-in: the default config is unchanged `quantization="none"`. +4. Integrate cleanly with tensor parallelism on a single node. + +**Non-goals (explicitly out of scope).** +- KV-cache quantization (a distinct bottleneck; to be proposed separately). +- Multi-node parallelism, ONNX/CPU backends, weight sparsification. +- Automatic post-training calibration as part of `areno train` (calibration is assumed already performed when a + quantized checkpoint is supplied, or handled by a separate tool). + +--- + +## 3. Design + +### 3.1 Supported schemes & selection + +| scheme | bytes/weight | use | +|---|---|---| +| `none` (bf16) | 2 | baseline | +| `fp8` (E4M3) | 1 | W8A16 (per-tensor / per-channel scale) — v1 | +| `int4` (grouped, e.g. g128, GPTQ/AWQ) | 0.5 | W4A16 — v1.1 | + +Selection is two-way: + +- **Explicit config.** Add `quant_method: "none" | "fp8" | "int4"` to `ModelConfig` (with per-method fields, e.g. + `quant_group_size: int = 128`, `quant_fp8_dtype: "e4m3"`), surfaced as `--quant {none,fp8,int4}` on + `areno train` and `areno serve`. +- **Checkpoint-driven.** If the checkpoint carries a `quantization_config` block (or an equivalent signature in + `config.json`), AReno auto-selects `quant_method` so a pre-quantized checkpoint "just works" without a flag. + This keeps a quantized model re-runnable in one command and avoids a silent double-quantization. + +### 3.2 Layer replacement (registry-driven) + +Instead of editing each model's forward, introduce a quantized-layer registry: + +- `areno/models/registry.py` gains a mapping from each model family's linear class to a **family adapter** that + knows how to construct a `QuantizedLinear` from (weight, scale, zero-point, group-size) and how to shard / + de-shard those tensors under tensor parallelism. +- `areno/engine/layers/linear.py` gains a `QuantizedLinear` module that dequantizes on-device (or runs a fused + dequant-matmul) and then uses the existing TP path, so collectives operate on the float side. + +This keeps adding families additive (a new adapter) rather than modifying a shared factory. + +### 3.3 Checkpoint loading / conversion + +Two load styles behind one loader: + +- **Static (pre-quantized).** Recognize a `quantization_config`; load packed integer tensors plus + `weight_scale` and (optionally) `zero_point` and `group_size`. +- **Dynamic (quantize-on-load).** Quantize an existing bf16 checkpoint at load time. FP8 is available without + calibration; INT4 requires calibration data, so INT4 v1 is **pre-quantized only**. + +On the I/O side, the weight-layout helpers in `areno/engine/checkpoints/io.py` and each family's +`checkpoint.py` must round-trip quantized tensors (packed ints + scales + zero-points), and the policy-sync +weight plan (`areno/engine/policy_sync.py`) must understand the new keys so weights can be exchanged between +train and rollout partitions. + +### 3.4 Kernels + +Keep the CUDA path self-contained (areno-owned, per AGENTS.md): + +- Add `dequant_linear` for FP8 and INT4 to `areno/accel/` (INT4 dequant → bf16 → existing `linear.cu`), or a + fused `quantized_linear.cu`. +- Deliberately **not** a hard new runtime dependency (e.g. no mandatory bitsandbytes). A fast path via an + established CUDA utility may be added later as an opt-in. + +**Compute-capability constraint (measured on the 8×A100 reference host).** +`torch._scaled_mm` — the built-in FP8 matmul that consumes FP8 weights directly — requires Hopper +(compute capability ≥ 8.9) or ROCm MI300+; it is **unavailable on Ampere (A100, cc 8.0)**. Verified locally: +`RuntimeError: torch._scaled_mm is only supported on CUDA devices with compute capability >= 9.0 or 8.9`. +Consequences: +- On **Hopper/H100**, the FP8 decode path can use the built-in FP8 matmul as a low-effort fast path. +- On **Ampere/A100** (the reference host), a **custom fused FP8-dequant matmul in `areno/accel`** is required + to realize the memory-bound speedup (dequant-then-bf16-matmul does not reduce bytes read, so it yields no + speedup — see §5). This raises the M1 effort on A100. + +### 3.5 Tensor-parallel integration + +- Quantized weights are TP-sharded inside the packed integer tensor; per-channel/group `weight_scale` and + `zero_point` are sharded together with their group. +- The forward dequantizes to a local float shard and reuses the existing TP reduce/all-reduce. Gradients + accumulate in the existing FP32 master (or 8-bit) optimizer unchanged. + +--- + +## 4. Validation plan (merge gate) + +1. **Numerical / CPU tests** (`tests/`, required): + - `dequant(matmul)` matches a reference implementation for FP8 and INT4 (random tensors, per-scheme tolerance). + - TP shard + all-reduce of a quantized weight equals the unsharded result (CPU parity test). + - Config / CLI parsing; `quantization="none"` produces the existing path byte-for-byte. +2. **Behavior-preserving gate:** run a bounded GRPO (or SFT) step with `--quant fp8` and assert `loss`, + `grad_norm`, `advantage_mean` are within a documented tolerance of the `--quant none` baseline. +3. **Benchmark (primary acceptance):** on `Qwen3-8B` @ world=4/tp=4 on 8×A100, measure decode **tokens/s and + peak GPU memory** for `none` vs `fp8` (v1.1 adds `int4`), reusing the `scripts/bench` probe. Target: near-linear + speedup from the byte reduction with memory headroom. +4. **Docs:** `docs/models/` + CLI reference document `--quant` usage and accuracy caveats. + +### 4.5 P1 kernel — measured status (2026-08-25) + +A Triton FP8(E5M2) W8A16 dequant-linear is implemented in `areno/accel/kernels/fp8_linear.py` +and wired into the single matmul entry point `areno/engine/layers/linear.py:_areno_linear_forward` +(detects a `_areno_fp8` weight payload and routes to the kernel; default path unchanged). + +Measured on Qwen3-8B MLP shape (batch 32 × 4096 → 12288) on A100, versus the production bf16 +`areno_linear` path: + +- **fp8 = 0.052 ms/step, bf16 = 0.084 ms/step → ~1.60×**, weight max-rel error 4.7% (per-tensor E5M2 scale). +- Key tunings: `num_stages=4` pipelining, `BLOCK_K=128`, `.cg` on the streamed weight, **FP16 tensor-core + dot** (bf16 dot measured ~1.13× — the gap was dtype, not allocation), and **applying the per-tensor scale + once post-dot** (`(x @ w_fp8) * scale`, so no per-element dequant in the inner loop). +- The 1.60× is per memory-bound linear; the full decode-path speedup is expected lower (attention/norm/vocab + are not purely bandwidth-bound) but the memory-bound premise is validated on the real model shape. +- Note on Ampere (A100): a *built-in* FP8 matmul is unavailable (`torch._scaled_mm` is Hopper/Ada-only), so the + kernel above is the A100 path; on Hopper the built-in FP8 matmul can be an opt-in fast path. + +### 4.6 P1 correctness caveat — E5M2 on A100 is a *coarse* fast path, not a faithful bf16 drop-in + +Measured end-to-end on the Qwen3-0.6B GRPO rollout with the model-level integration active +(`quantize_model_weights_fp8` marks **112** TP-parallel linears; the decode runs and scores rewards +normally). Using **greedy** (deterministic) decoding, the same prompts produce **different reasoning than +bf16** (e.g. `"…48 friends in April. Then, she sold…"` (bf16) vs `"…48 of her friends in April, and then she +sold…"` (fp8)). The FP8 **kernel is correct** (linear-level verified, ~5 % error); the divergence comes from +the **coarse E5M2 grid** (only the Triton FP8 dtype Ampere supports — `fp8e4nv`/E4M3 is rejected on A100), +whose ~12 % per-element error flips greedy argmax choices and cascades. + +**Implication (be honest with consumers):** on A100 this is a **fast-but-coarse decode path**, not a faithful +replacement for bf16. It yields a real ~1.60× on the memory-bound linear but with accuracy divergence on +greedy decoding. To keep output faithful: +- prefer **Hopper/H100** (built-in FP8 matmul + E4M3, `torch._scaled_mm`); or +- add an **E4M3 dequant-linear CUDA kernel** (Ampere Triton only accepts E5M2), which is more work but finer. + +Out of scope here: making E5M2 output match bf16. The M1 acceptance gate's §4.2 "behavior-preserving" step +**cannot be satisfied on A100 for greedy decode** under E5M2; it may be satisfiable on Hopper / with E4M3. + +### 4.7 E4M3 CUDA fused dequant-linear — correct; small-M GEMV path, but still < bf16 on A100 (2026-08-25) + +To get a **finer** (E4M3, 3 mantissa bits) grid on A100 where Triton rejects `fp8e4nv`, a custom CUDA kernel was +built. It reads the uint8 E4M3 weight (1 byte/element) and decodes each byte to bf16 in-kernel (fast bit-decode +validated against `torch.float8_e4m3fn` in `tests/test_e4m3_decode_cpu.py`). + +- `areno/accel/csrc/e4m3_linear.cu`: two kernels + registration. + - `e4m3_linear_kernel`: WMMA (fp16 16×16×16) tensor-core GEMM for general M (correct, but capped at ~60 GB/s + for the tiny-M decode case — the tensor-core GEMM wastes the M dim and its register pressure limits + occupancy). + - `e4m3_gemv_kernel`: **small-M memory-streaming GEMV** — one warp per output row, reading the row with + **vectorized (8 bytes/lane) coalesced** `uint8` loads along the contiguous k dim, reusing the (M×K) + activation from shared, and warp-reducing with shuffles. This is the right shape for the memory-bound + decode case and is ~6.7× faster than the GEMM variant. + - Dispatch: `M ≤ 4` uses the GEMV; larger M uses the GEMM. +- `areno/accel/kernels/e4m3_cuda.py` — `quantized_e4m3_linear_cuda` shim + quantize/dequant helpers. +- **Correctness (verified):** vs the bf16 reference, max-rel error ≈ 0.4–0.7 % (fp16 compute noise) across + M ∈ {1,4,64}, N=12288, K=4096; vs full bf16 ≈ 2.8–3.6 % (E4M3 quantization, expected). +- **Performance (honest, corrected).** The E4M3 **kernel** streams the 1-byte weight once and is memory-bound + (~0.05 ms at M=1, N=12288, K=4096 — ≈1.68× vs the bf16 cuBLAS kernel's ~0.084 ms at the GPU level). **But + the end-to-end ratio through the torch/ATen extension is only ~1.0×** (single-launch bf16 0.101 ms vs E4M3 + 0.098 ms = 1.03×; pipelined 0.88×). The kernel's speed advantage is swallowed by the per-call `torch::empty` + + ATen dispatch + the shim's `.contiguous()` overhead (~0.04 ms), which is ~50–100 % of a 0.05 ms kernel. A real + decode hot path reuses the output buffer (now supported via an `out` param) but the residual dispatch overhead + still dominates. **>1.5× is NOT achieved end-to-end**; it exists only at the bare-kernel level and only if the + framework dispatch overhead is eliminated (the next real step). An earlier "1.68×" claim was a measurement + error (comparing the kernel-level pipelined throughput against the bf16 single-launch latency). Note: an + earlier absolute number of ~1000 GB/s for a float-x standalone was also non-representative (only single-launch + is realistic; a 2.2× gap was an artifact of comparing pipelined-vs-single-launch, not a genuine bf16 penalty — + the real bf16-x kernel runs at the same rate as float-x). +### 4.8 Hopper/H20 — native FP8 E4M3 via `torch._scaled_mm` (measured >1.5×) + +On a **NVIDIA H20 / "CX70"** (Hopper, compute capability 9.0, 8 GPUs), the built-in cuBLASLt FP8 matmul +(`torch._scaled_mm`, fp8_e4m3) gives the memory-bound decode speedup directly — no custom kernel needed. +Measured with `scripts/bench/h20_fp8_scaled_mm.py` at the MLP decode shape vs the bf16 cuBLAS baseline: + +| M | N | K | bf16 | FP8 `_scaled_mm` | speedup | +|---|---|---|---|---|---| +| 1 | 12288 | 4096 | 0.0493 ms | 0.0292 ms | **1.69×** | +| 1 | 8192 | 4096 | 0.0398 ms | 0.0250 ms | **1.59×** | +| 4 | 12288 | 4096 | 0.0497 ms | 0.0289 ms | **1.72×** | +| 64 | 12288 | 4096 | 0.0662 ms | 0.0431 ms | **1.53×** | + +**>1.5× achieved on Hopper/H20**, with correct FP8 E4M3 quantization error (~3–4 % rel vs bf16, the E4M3 grid). +`torch._scaled_mm` needs A row-major (M,K) and B column-major (K,N) — pass the weight as `w.t()` (not +`.contiguous()`) for the `x @ w^T` decode GEMM. This is the Hopper fast path the RFC §4.6 pointed to; on +Ampere (A100) the custom CUDA kernel in `areno/accel/csrc/e4m3_linear.cu` is the equivalent, and the +`rel_vs_bf16` ~4 % (E4M3) vs ~12 % (E5M2) confirms E4M3 is the *finer* grid that keeps greedy decode closer +to bf16 than the A100 E5M2 Triton path. + +Note: getting a CUDA torch onto this host was itself high-effort — `download.pytorch.org` / PyPI were 27 KB/s +and the China mirrors (aliyun) lacked the exact `nvidia-*` versions the newest torch pins. Resolution: install +`torch 2.10.0+cu126 --no-deps` from the aliyun `pytorch-wheels` mirror (fast) and pull the `nvidia-*` CUDA +runtime wheels (cudnn-9.9, cublas-12.9, cudart/cusparse/nccl/cusolver/etc.) from the aliyun pypi mirror, +iterating to a consistent major-version set (libcudnn.so.9, libcublas.so.12) that torch's loader accepts. +- **Bandwidth ceiling (root-caused, honest).** The raw read-only kernel (row-per-warp, coalesced uint64 + weight reads) hits **~1248 GB/s** on A100, so the read pattern was **never** the bottleneck. The ~400 GB/s + wall came from the **per-element `exp2f` E4M3 decode**. Replacing it with a **branchless bit-trick** + (laying the E4M3 bits out into fp32 directly: `fp32 = (s<<31)|((e+120)<<23)|(m<<20)`, with cheap selects for + the e==0 subnormal `m*2^-9` and the e==15,m==7 NaN clamp) — verified **exact** against `torch.float8_e4m3fn` + for all 256 bytes — raises the GEMV to ~**815 GB/s (~0.062 ms)**, and staging the activation as **float** (and + decoding the weight once per load) removes the per-element bf16→float cost. The kernel is now **memory-bound**; + ~815 GB/s is ≈**1.36×** vs the bf16 cuBLAS 0.084 ms (read-only ceiling 2.1×). A `cp.async` pipeline variant + was also built and its crash root-caused (Ampere `cp.async` is **max 16 bytes/op**; a 32-byte/lane copy traps + on this CUDA 13.0 / cc8.0 toolchain — fixed by 16-byte copies) but it does not beat the direct kernel. + NOTE: the shared host's GPUs were saturated while measuring, so the live `areno_linear`-vs-kernel ratio + (~1.0×) is contention-inflated; the clean standalone is ~1.36×. + +> 🧠 **From Hindsight memory** — the correct branchless E4M3→fp32 decode is `(s<<31)|((e+120)<<23)|(m<<20)` +> (fp32 shares bf16's exponent bias 127; the 3 mantissa bits go to the top of fp32's 23-bit mantissa). The +> subnormal case (e==0: `m*2^-9`) and the NaN pattern (e==15,m==7) need selects; do **not** re-apply the sign +> to the normal path (it's already in `s<<31`) — that double-negation caused a ~1.3 rel error that took a +> debug pass to catch. + +**Implication:** E4M3 is the *correctness*-finer path and now has a working CUDA kernel, but the decode speedup +for it is **not yet realized** on A100. E5M2 (Triton, ~1.60× measured) remains the fast-but-coarse A100 path; +E4M3 needs either (a) a working cp.async small-M streaming pipeline (multi-step kernel work, toolchain issue +unresolved), or (b) Hopper/H100 (built-in FP8 matmul + E4M3). + + +--- + +## 5. Alternatives considered + +| alternative | assessment | +|---|---| +| bf16 → fp16 | no byte reduction (both 2 bytes) → no memory-bound benefit. Rejected. | +| KV-cache quantization | targets KV reads, not weight reads; a different bottleneck. Separate feature. | +| 8-bit AdamW (already present) | reduces optimizer memory, not decode weight reads. Not a substitute. | +| Speculative decoding | complementary; improves latency but does not reduce bytes per token. | +| Weight caching / reuse | decode reads weights per token by construction; no reuse window. | +| Adopt a third-party quant runtime (e.g. bitsandbytes) as core | fast prototype, but adds a runtime dependency; keep optional, not core. | + +--- + +## 6. Risks & mitigations + +- **INT4 accuracy depends on calibration data.** Mitigation: ship INT4 as pre-quantized-only in v1, and gate on + the numerical/behavioral tests. +- **Packed-int TP sharding** is the most likely correctness hazard. Mitigation: dedicated CPU parity tests for + shard + all-reduce, plus the behavior-preserving step gate. +- **Per-family adapter growth** raises maintenance cost. Mitigation: registry-driven (additive), broad coverage + via a shared `QuantizedLinear` and only thin per-family adapters. +- **Accuracy of a dynamic FP8 conversion** (no calibration) is acceptable but not free. Mitigation: document a + tolerance and keep `none` as the default. + +--- + +## 7. Roadmap + +**Sequencing principle.** Prove measurable value early, then broaden scope, and gate every phase on the §4 +validation plan — the correctness/perf gates are the merge bar. Each phase is independently landable. + +### 7.1 Phases + +| Milestone | Objective | In-scope | Deliverables | Dependencies | Exit criteria (gate) | Effort | +|---|---|---|---|---|---|---| +| **M0 — Design freeze** | Lock scope, interface, KPIs | `quant_method` config + CLI parity; tolerance + benchmark KPIs | Accepted RFC; `ModelConfig`/CLI changes | — | Reviewer sign-off; no open design Qs | ~0.5 wk (review) | +| **M1 — FP8 vertical slice** | Prove value on one model family | Qwen3 family only | `QuantizedLinear`; FP8 dequant kernel (`areno/accel`); dynamic load; `--quant fp8`; CPU/numerical tests; `areno serve`/rollout run | M0 | `none` vs `fp8` decode bench on Qwen3-8B; behavior-preserving step gate; CPU tests | ~2–3 wk | +| **M2 — TP + checkpoint round-trip** | Correct & repeatable under TP | FP8 weights+scales shard/de-shard; static (pre-quantized) load; policy-sync keys | TP shard/all-reduce parity; checkpoint save→load round-trip | M1 | TP parity test; round-trip reproduces quantized tensors | ~2–3 wk | +| **M3 — Family coverage & docs** | Broaden to the model matrix | llama, qwen3_5, gemma4, bailing, … | Per-family adapter registrations; `areno check` surface; `docs/models/` + CLI reference | M2 | Family matrix loads + decodes; docs shipped | ~2–3 wk | +| **M4 — INT4 (pre-quantized)** | Extend to 0.5 B/weight | INT4 GPTQ/AWQ load + dequant kernel + TP | `none/fp8/int4` benchmark table; behavior gate; tests | M2 | Benchmark table + gate + tests | ~3–4 wk | +| **M5 — Hardening** | Consolidate, optimize, integrate | fused dequant-matmul; optional fast path; async/rollout-policy integration | Perf/throughput report; no regressions | M3, M4 | Existing suite green; report documented | ~2 wk | + +### 7.2 Sequencing & dependencies + +``` +M0 ──► M1 ──► M2 ──► M3 ──► M5 + └────► M4 ──► M5 +``` + +- **Critical path:** M0 → M1 (value proof) → M2 (correctness under TP). M4 is parallelizable after M2. +- **Early stop-or-continue signal:** M1 alone yields the first measurable `none` vs `fp8` number; commit to M2+ + only if that result is compelling. + +### 7.3 Effort summary + +| scope | est. eng-weeks (1 contributor) | +|---|---| +| M0 | 0.5 | +| M1 | 2–3 | +| M2 | 2–3 | +| M3 | 2–3 | +| M4 | 3–4 | +| M5 | 2 | +| **Full feature** | **~12–17** | +| **M1 only (value proof)** | **~2–3** | + +### 7.4 What is needed from reviewers / maintainers + +- Confirm scope and the `quant_method` / `--quant` interface (M0). +- Decide the FP8 default (W8A16 vs W8A8) — see §8. +- Approve whether the M1 value-proof slice (~2–3 wk) is worth funding before committing to the full roadmap. + +--- + +## 8. Open questions + +1. Should FP8 default to W8A16 (dequant → bf16 matmul) or W8A8 (native FP8 matmul)? We lean W8A16 first for + risk, W8A8 as an opt-in fast path. +2. Is `areno serve` (inference-only) sufficient to validate decode throughput, or must the RL rollout path + also be exercised at each phase? +3. Does the community prefer a single `quant_method` string, or distinct `--quant-fp8` / `--quant-int4` flags? + (Current proposal: single `--quant` with per-method params.) diff --git a/scripts/bench/e4m3_cuda_bench.py b/scripts/bench/e4m3_cuda_bench.py new file mode 100644 index 00000000..45409798 --- /dev/null +++ b/scripts/bench/e4m3_cuda_bench.py @@ -0,0 +1,74 @@ +"""GPU bench + correctness check for the E4M3 fused dequant-GEMM kernel. + +Validates the CUDA kernel against the bf16 reference across representative +shapes and times it vs the bf16 ``areno_linear`` cuBLAS path on decode-like +shapes. Run on a free GPU: + + CUDA_VISIBLE_DEVICES=2 python scripts/bench/e4m3_cuda_bench.py +""" + +from __future__ import annotations + +import torch + +from areno.accel import areno_linear +from areno.accel.kernels.e4m3_cuda import ( + dequant_e4m3_bf16, + quantize_weight_e4m3, + quantized_e4m3_linear_cuda, +) + + +def rel_err(a: torch.Tensor, b: torch.Tensor) -> float: + d = (a.float() - b.float()).abs().max().item() + denom = b.float().abs().max().item() + 1e-6 + return d / denom + + +def bench(fn, iters: int = 200, warmup: int = 20) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(True) + end = torch.cuda.Event(True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def main() -> None: + torch.manual_seed(0) + assert torch.cuda.is_available(), "need a GPU" + + print("=== correctness vs bf16 reference ===") + for M, N, K in [(1, 4096, 4096), (1, 12288, 4096), (4, 12288, 4096), (64, 12288, 4096)]: + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") / (K**0.5) + w_u8, scale = quantize_weight_e4m3(w) + out = quantized_e4m3_linear_cuda(x, w_u8, scale) + dq = dequant_e4m3_bf16(w_u8, scale) + ref = x @ dq.T + print( + f" M={M:3d} N={N:5d} K={K:5d}: shape={tuple(out.shape)} " + f"dtype={out.dtype} rel_vs_dequant={rel_err(out, ref):.5f} " + f"finite={bool(torch.isfinite(out).all())}" + ) + + print("\n=== decode throughput vs bf16 areno_linear ===") + print(f"{'M':>4} {'N':>6} {'K':>6} {'bf16(ms)':>10} {'e4m3(ms)':>10} {'speedup':>8} {'e4m3GB/s':>9}") + for M, N, K in [(1, 12288, 4096), (1, 8192, 4096), (4, 12288, 4096), (64, 12288, 4096)]: + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") / (K**0.5) + w_u8, scale = quantize_weight_e4m3(w) + t_bf16 = bench(lambda: areno_linear(x, w, None)) + t_e4m3 = bench(lambda: quantized_e4m3_linear_cuda(x, w_u8, scale)) + mb = N * K * 1 / 1e6 # uint8 weight MB + gbs = mb * 1e6 / (t_e4m3 * 1e-3 * 1e9) + print(f"{M:4d} {N:6d} {K:6d} {t_bf16:10.4f} {t_e4m3:10.4f} {t_bf16 / t_e4m3:8.3f}x {gbs:9.1f}") + + +if __name__ == "__main__": + main() diff --git a/scripts/bench/fp8_end_to_end_bench.py b/scripts/bench/fp8_end_to_end_bench.py new file mode 100644 index 00000000..44b40f67 --- /dev/null +++ b/scripts/bench/fp8_end_to_end_bench.py @@ -0,0 +1,62 @@ +"""End-to-end verification of the FP8 linear path on the Qwen3-8B MLP shape. + +Puts the FP8 W8A16 Triton kernel behind the real single matmul entry point +(``areno.engine.layers.linear._areno_linear_forward``), then measures bf16 vs +FP8 through that path. Confirms (1) the hook dispatches to the FP8 kernel and +(2) the measured decode-matmul speedup (~1.6x) holds on the real model shape. +""" + +from __future__ import annotations + +import time + +import torch + +from areno.accel.kernels.fp8_linear import mark_fp8_weight, quantized_fp8_linear +from areno.engine.layers.linear import _areno_linear_forward + + +def _bench(fn, iters=40, warmup=8): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(iters): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - t0) / iters + + +def main() -> None: + torch.cuda.set_device(0) + dev = torch.device("cuda", 0) + M, K, N = 32, 4096, 12288 # Qwen3-8B MLP gate/up shape + x = torch.randn(M, K, device=dev, dtype=torch.bfloat16) + w = torch.randn(N, K, device=dev, dtype=torch.bfloat16) + + bf16_out = _areno_linear_forward(x, w, None) + torch.cuda.synchronize() + + w_q = torch.nn.Parameter(w.clone()) + mark_fp8_weight(w_q) + assert getattr(w_q, "_areno_fp8", None) is not None, "mark_fp8_weight did not set the FP8 payload" + fp8_out = _areno_linear_forward(x, w_q, None) + torch.cuda.synchronize() + + rel = float((fp8_out.float() - bf16_out.float()).abs().max() / (bf16_out.float().abs().max() + 1e-6)) + print(f"fp8 linalg rel err vs bf16 (via hook) = {rel:.4f}") + + t_bf16 = _bench(lambda: _areno_linear_forward(x, w, None)) + t_fp8_hook = _bench(lambda: _areno_linear_forward(x, w_q, None)) + # Kernel-only (reuse a pre-allocated out buffer -> no per-step cudaMalloc). + ybuf = torch.empty(M, N, device=dev, dtype=torch.bfloat16) + t_fp8_noalloc = _bench(lambda: quantized_fp8_linear(x, w_q._areno_fp8, w_q._areno_fp8_scale, out=ybuf)) + print(f"bf16 (areno_linear) : {t_bf16 * 1e3:.3f} ms/step, {M / t_bf16:.0f} tok/s") + print(f"fp8 (hook, alloc) : {t_fp8_hook * 1e3:.3f} ms/step, {M / t_fp8_hook:.0f} tok/s") + print(f"fp8 (kernel, no alloc) : {t_fp8_noalloc * 1e3:.3f} ms/step, {M / t_fp8_noalloc:.0f} tok/s") + print(f"speedup (bf16 / fp8-alloc) = {t_bf16 / t_fp8_hook:.2f}x") + print(f"speedup (bf16 / fp8-noalloc) = {t_bf16 / t_fp8_noalloc:.2f}x") + + +if __name__ == "__main__": + main() diff --git a/scripts/bench/h20_fp8_scaled_mm.py b/scripts/bench/h20_fp8_scaled_mm.py new file mode 100644 index 00000000..07bb041f --- /dev/null +++ b/scripts/bench/h20_fp8_scaled_mm.py @@ -0,0 +1,87 @@ +"""Validate E4M3 FP8 decode speedup on Hopper (H20) via torch._scaled_mm. + +The H20 (cc 9.0) supports the built-in FP8 matmul. For the memory-bound decode +case, the weight is read as fp8_e4m3 (1 byte) so the weight-read bytes halve vs +bf16. This compares `torch._scaled_mm` on FP8 vs the bf16 cuBLAS matmul at the +MLP decode shape, and validates correctness against a bf16 reference. +""" + +from __future__ import annotations + +import torch + +torch.manual_seed(0) + + +def rel_err(a: torch.Tensor, b: torch.Tensor) -> float: + return float((a.float() - b.float()).abs().max() / (b.float().abs().max() + 1e-6)) + + +def best(fn, reps=100, warm=20) -> float: + for _ in range(warm): + fn() + torch.cuda.synchronize() + bv = 1e9 + for _ in range(reps): + s = torch.cuda.Event(True) + e = torch.cuda.Event(True) + s.record() + fn() + e.record() + torch.cuda.synchronize() + bv = min(bv, s.elapsed_time(e)) + return bv + + +def quant_fp8(t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Per-tensor FP8 E4M3 quantize; returns (fp8, scale).""" + amax = t.abs().amax() + scale = torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)) + q = (t / scale).clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + return q, scale.to(torch.float32) + + +def main() -> None: + assert torch.cuda.is_available() + cc = torch.cuda.get_device_capability(0) + if not hasattr(torch, "_scaled_mm") or cc < (8, 9): + raise RuntimeError( + f"torch._scaled_mm (FP8 E4M3 matmul) requires Hopper/Ada (cc >= 8.9); got cc={cc}. " + "This benchmark is Hopper/H20-only." + ) + print(f"torch {torch.__version__} cc {cc[0]}.{cc[1]} gpus {torch.cuda.device_count()}") + + for M, N, K in [(1, 12288, 4096), (1, 8192, 4096), (4, 12288, 4096), (64, 12288, 4096)]: + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") / K**0.5 + + # --- bf16 baseline (reads 2 bytes/weight) --- + t_bf16 = best(lambda: torch.mm(x, w.T)) + + # --- FP8 W8A8 via _scaled_mm (reads 1 byte/weight) --- + xq, xs = quant_fp8(x) + wq, ws = quant_fp8(w) + # _scaled_mm(a,b): a (M,K) fp8, b (K,N) fp8 -> (M,N). w is (N,K), need w^T (K,N). + wq_t = wq.to(torch.float8_e4m3fn).t().contiguous() + sm = torch._scaled_mm if hasattr(torch, "_scaled_mm") else None + if sm is None: + print(f"M={M} N={N}: torch._scaled_mm NOT available") + continue + # out = a @ b + out = sm(xq, wq_t, out_dtype=torch.bfloat16, scale_a=xs, scale_b=ws) + t_fp8 = best(lambda: sm(xq, wq_t, out_dtype=torch.bfloat16, scale_a=xs, scale_b=ws)) + + # reference: dequant fp8 weights and matmul in bf16 + dq = (wq.float() * ws).to(torch.bfloat16) + ref = x @ dq.T + rel = rel_err(out, ref) + rel_bf16 = rel_err(out, x @ w.T) + + print( + f"M={M} N={N} K={K}: bf16={t_bf16:.4f}ms fp8_scaleMM={t_fp8:.4f}ms " + f"speedup={t_bf16 / t_fp8:.3f}x rel_vs_dq={rel:.4f} rel_vs_bf16={rel_bf16:.4f}" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/bench/triton_fp8_matmul_bench.py b/scripts/bench/triton_fp8_matmul_bench.py new file mode 100644 index 00000000..8245042b --- /dev/null +++ b/scripts/bench/triton_fp8_matmul_bench.py @@ -0,0 +1,168 @@ +"""Tuned Triton bf16 vs FP8(W8A16) matmul bench on A100. + +Applies SOTA Triton GEMM levers to both kernels so the delta isolates the +weight-precision (2 bytes vs 1 byte) memory effect: + * num_stages — pipelined prefetch of the next tile while computing (critical + to reach memory bandwidth), + * larger BLOCK_N / BLOCK_K and num_warps, + * `.cg` (cache-global) on the large streamed weight load and default on the + activation, so the weight doesn't thrash L1 and the kernel is bandwidth-, + not latency-, bound. +Runs the same config sweep for bf16 and fp8 and reports the best of each. +""" + +from __future__ import annotations + +import time + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _mm_bf16( + a_ptr, b_ptr, y_ptr, M, N, K, sa, sb, sy, BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, IS_CG: tl.constexpr +): + pm = tl.program_id(0) + pn = tl.program_id(1) + rm = pm * BM + tl.arange(0, BM) + rn = pn * BN + tl.arange(0, BN) + acc = tl.zeros((BM, BN), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BK)): + rk = k * BK + tl.arange(0, BK) + am = (rm[:, None] < M) & (rk[None, :] < K) + bm_ = (rn[:, None] < N) & (rk[None, :] < K) + a = tl.load(a_ptr + rm[:, None] * sa + rk[None, :], mask=am, other=0.0) + b = tl.load(b_ptr + rn[:, None] * sb + rk[None, :], mask=bm_, other=0.0, cache_modifier=".cg" if IS_CG else "") + acc += tl.dot(a, tl.trans(b)) + tl.store(y_ptr + rm[:, None] * sy + rn[None, :], acc.to(tl.bfloat16), mask=(rm[:, None] < M) & (rn[None, :] < N)) + + +@triton.jit +def _mm_fp8( + a_ptr, + b_ptr, + scale_ptr, + y_ptr, + M, + N, + K, + sa, + sb, + sy, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + IS_CG: tl.constexpr, +): + pm = tl.program_id(0) + pn = tl.program_id(1) + rm = pm * BM + tl.arange(0, BM) + rn = pn * BN + tl.arange(0, BN) + acc = tl.zeros((BM, BN), dtype=tl.float32) + scale = tl.load(scale_ptr) + for k in range(0, tl.cdiv(K, BK)): + rk = k * BK + tl.arange(0, BK) + am = (rm[:, None] < M) & (rk[None, :] < K) + bm_ = (rn[:, None] < N) & (rk[None, :] < K) + a = tl.load(a_ptr + rm[:, None] * sa + rk[None, :], mask=am, other=0.0) + b = tl.load(b_ptr + rn[:, None] * sb + rk[None, :], mask=bm_, other=0.0, cache_modifier=".cg" if IS_CG else "") + # Per-tensor scalar scale distributes out of the matmul, so dequantize + # the tile to fp16 (Ampere tensor cores) and apply the scalar AFTER the + # dot — no per-element multiply in the inner loop. + acc += tl.dot(a.to(tl.float16), tl.trans(b.to(tl.float16))) + acc = acc * scale + tl.store(y_ptr + rm[:, None] * sy + rn[None, :], acc.to(tl.bfloat16), mask=(rm[:, None] < M) & (rn[None, :] < N)) + + +def _bench(fn, iters=40, warmup=8): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(iters): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - t0) / iters + + +def main() -> None: + torch.cuda.set_device(0) + dev = torch.device("cuda", 0) + K, N, M = 4096, 12288, 32 + x = torch.randn(M, K, device=dev, dtype=torch.bfloat16) + w = torch.randn(N, K, device=dev, dtype=torch.bfloat16) + scale = (w.abs().amax() / 57344.0).to(torch.float32).reshape(()) + w_fp8 = (w.float() / scale).clamp(-57344.0, 57344.0).to(torch.float8_e5m2) + y = torch.empty(M, N, device=dev, dtype=torch.bfloat16) + + # correctness + _mm_fp8[(triton.cdiv(M, 32), triton.cdiv(N, 128))]( + x, w_fp8, scale, y, M, N, K, x.stride(0), w_fp8.stride(0), y.stride(0), 32, 128, 64, True + ) + torch.cuda.synchronize() + ref = x.float() @ w.float().T + rel = float((y.float() - ref).abs().max() / (ref.abs().max() + 1e-6)) + print(f"fp8 output rel err vs bf16 = {rel:.4f}") + + bms = [32, 64] + bns = [64, 128, 256] + bks = [64, 128] + stages = [2, 3, 4] + warps = [4, 8] + + def run(kernel, use_cg): + is_fp8 = kernel is _mm_fp8 + bmat = w_fp8 if is_fp8 else w + bstr = bmat.stride(0) + best = None + for BM in bms: + for BN in bns: + for BK in bks: + for ns in stages: + for nw in warps: + try: + grid = (triton.cdiv(M, BM), triton.cdiv(N, BN)) + if is_fp8: + args = (x, bmat, scale, y, M, N, K, x.stride(0), bstr, y.stride(0)) + else: + args = (x, bmat, y, M, N, K, x.stride(0), bstr, y.stride(0)) + kernel[grid](*args, BM, BN, BK, use_cg, num_stages=ns, num_warps=nw) + torch.cuda.synchronize() + + def fn(_k=kernel, _g=grid, _a=args, _BM=BM, _BN=BN, _BK=BK, _cg=use_cg, _ns=ns, _nw=nw): + _k[_g](*_a, _BM, _BN, _BK, _cg, num_stages=_ns, num_warps=_nw) + + t = _bench(fn) + if best is None or t < best[0]: + best = (t, BM, BN, BK, ns, nw) + except Exception: + continue + return best + + bf16_best, fp8_best = None, None + for cg in (False, True): + b = run(_mm_bf16, cg) + if b and (bf16_best is None or b[0] < bf16_best[0]): + bf16_best = b + f = run(_mm_fp8, cg) + if f and (fp8_best is None or f[0] < fp8_best[0]): + fp8_best = f + + if bf16_best and fp8_best: + bt, fp8t = bf16_best[0], fp8_best[0] + print( + f"bf16 best: {bt * 1e3:.3f} ms/step (BM={bf16_best[1]} BN={bf16_best[2]} BK={bf16_best[3]} ns={bf16_best[4]} w={bf16_best[5]})" + ) + print( + f"fp8 best: {fp8t * 1e3:.3f} ms/step (BM={fp8_best[1]} BN={fp8_best[2]} BK={fp8_best[3]} ns={fp8_best[4]} w={fp8_best[5]})" + ) + print(f"speedup (bf16_time/fp8_time) = {bt / fp8t:.2f}x (fp8 fraction of bf16 = {fp8t / bt:.2f})") + print(f"weight bytes: bf16={w.numel() * 2}, fp8={w.numel() * 1}; theoretical ceil = 2.00x") + else: + print("no valid config for one of the kernels") + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index a67b7116..2b8da3b3 100644 --- a/setup.py +++ b/setup.py @@ -65,6 +65,7 @@ def _cuda_extensions(): "areno/accel/csrc/conv.cu", "areno/accel/csrc/embedding.cu", "areno/accel/csrc/linear.cu", + "areno/accel/csrc/e4m3_linear.cu", "areno/accel/csrc/moe_align_kernel.cu", "areno/accel/csrc/moe_permute.cu", "areno/accel/csrc/normalization.cu", diff --git a/tests/test_e4m3_decode_cpu.py b/tests/test_e4m3_decode_cpu.py new file mode 100644 index 00000000..d4d416ae --- /dev/null +++ b/tests/test_e4m3_decode_cpu.py @@ -0,0 +1,94 @@ +"""CPU tests for the E4M3 decode reference (RFC 0001, E4M3). + +The CUDA kernel in ``areno/accel/csrc/e4m3_linear.cu`` decodes each uint8 E4M3 +byte to bf16 in-kernel. These tests validate that reference formula against +``torch.float8_e4m3fn``'s authoritative conversion, across the representable +range and the boundary cases (subnormal, max finite 448, zero, signs). They run +on CPU without a GPU; the kernel's GPU correctness is covered separately. + +Decode formula: + sign = (b>>7)&1; e = (b>>3)&0xF; m = b&0x7 + e==0 (subnormal): (-1)^s * m * 2^-9 + 1<=e<=15 (normal): (-1)^s * 2^(e-7) * (1 + m/8); max finite 448 (e=15,m=6) + e==15 && m==7 is NaN in torch (a valid quantized weight is clamped at 448). +""" + +from __future__ import annotations + +import math + +import torch + +_E4M3_MAX = 448.0 + + +def decode_e4m3(b: int) -> float: + sign = (b >> 7) & 1 + e = (b >> 3) & 0xF + m = b & 0x7 + if e == 0: + val = float(m) * 0.001953125 # m * 2^-9 + elif e == 15 and m == 7: + val = _E4M3_MAX # NaN pattern; clamp to max finite + else: + val = 2.0 ** (e - 7) * (1.0 + float(m) / 8.0) + return -val if sign else val + + +def _torch_decode(byte: int) -> float: + t = torch.tensor([byte], dtype=torch.uint8).view(torch.float8_e4m3fn) + return float(t.to(torch.bfloat16).float()) + + +def _is_finite(v: float) -> bool: + return not (math.isnan(v) or math.isinf(v)) + + +def test_all_representable_e4m3_bytes_match_torch(): + # Every byte whose torch decode is finite must match the reference formula. + mismatches = 0 + for byte in range(256): + ref = decode_e4m3(byte) + refv = float(ref) + torch_v = _torch_decode(byte) + if _is_finite(torch_v): + if abs(refv - torch_v) > 0.0: + mismatches += 1 + assert mismatches == 0, f"{mismatches} finite bytes diverge from torch" + + +def test_e4m3_max_finite_and_zero(): + assert decode_e4m3(0x7E) == _E4M3_MAX # 448 + assert decode_e4m3(0x00) == 0.0 + assert decode_e4m3(0x80) == -0.0 + # symmetric negative max + assert decode_e4m3(0xFE) == -_E4M3_MAX + + +def test_e4m3_subnormal_convention(): + # subnormal: e==0, value = m * 2^-9 + assert decode_e4m3(0x01) == 1.0 / 512.0 # 2^-9 + assert decode_e4m3(0x02) == 2.0 / 512.0 + assert decode_e4m3(0x04) == 4.0 / 512.0 + + +def test_e4m3_powers_of_two(): + assert decode_e4m3(0x38) == 1.0 # e=7,m=0 + assert decode_e4m3(0x40) == 2.0 # e=8,m=0 + assert decode_e4m3(0x30) == 0.5 # e=6,m=0 + + +def test_quantize_dequant_matches_reference_matmul(): + from areno.accel.kernels.e4m3_cuda import dequant_e4m3_bf16, quantize_weight_e4m3 + + torch.manual_seed(0) + w = torch.randn(64, 96) + w = (w / w.abs().amax()) * 10.0 + w_bf16 = w.to(torch.bfloat16) + w_u8, scale = quantize_weight_e4m3(w_bf16) + dq = dequant_e4m3_bf16(w_u8, scale) + x = torch.randn(3, 96, dtype=torch.bfloat16) + ref = x @ w_bf16.T + out = x @ dq.T + rel = float((out.float() - ref.float()).abs().max() / (ref.float().abs().max() + 1e-6)) + assert rel < 0.2, f"e4m3 dequant matmul rel err too high: {rel}" diff --git a/tests/test_fp8_quant_cpu.py b/tests/test_fp8_quant_cpu.py new file mode 100644 index 00000000..b7f0cf21 --- /dev/null +++ b/tests/test_fp8_quant_cpu.py @@ -0,0 +1,96 @@ +"""CPU tests for FP8 weight quantization (RFC 0001, M1). + +These run without a GPU: the quantize/dequant math is implemented on float +tensors and the `QuantizedLinear` module is a reference W8A16 dequant-forward +module. They assert: + * dequant(quantize(w)) is within FP8 tolerance of w, + * a dequantized weight matmul is close to the bf16 linear, + * the `QuantizedLinear` module reproduces the reference linear, + * the config default is opt-in unchanged (quant_method == "none"). +""" + +from __future__ import annotations + +import torch + +from areno.engine.config import ModelConfig +from areno.engine.quantization import compute_fp8_scale, dequant_fp8, quantize_to_fp8 + + +def _rel_err(a: torch.Tensor, b: torch.Tensor) -> float: + return float((a.float() - b.float()).abs().max() / (b.float().abs().max() + 1e-6)) + + +def _weights(shape=(64, 96)) -> torch.Tensor: + # Unit-scale weights so the FP8 grid covers the range consistently. + w = torch.randn(*shape) + return (w / w.abs().amax()) * 10.0 + + +def test_quant_dequant_roundtrip_within_fp8_tolerance(): + w = _weights() + q, scale = quantize_to_fp8(w, group_size=-1) + dq = dequant_fp8(q, scale, group_size=-1) + # E4M3 has ~4 mantissa bits => relative error <= ~2^-4. Use a safe bound. + assert _rel_err(dq, w) < 0.12, f"roundtrip rel err too high: {_rel_err(dq, w)}" + # The maximum-magnitude element is preserved exactly by construction. + amax_idx = torch.argmax(w.abs()) + assert abs(float(dq.reshape(-1)[amax_idx]) - float(w.reshape(-1)[amax_idx])) < 1e-3 + + +def test_quant_dequant_per_group_roundtrip(): + w = _weights((48, 128)) + g = 64 + q, scale = quantize_to_fp8(w, group_size=g) + dq = dequant_fp8(q, scale, group_size=g) + assert dq.shape == w.shape + assert float(scale.numel()) == w.numel() // g + assert _rel_err(dq, w) < 0.12 + + +def test_dequant_weight_matmul_matches_bf16(): + x = torch.randn(3, 96) + w = _weights((64, 96)) + q, scale = quantize_to_fp8(w, group_size=-1) + dq = dequant_fp8(q, scale, group_size=-1).to(torch.bfloat16) + ref = x.to(torch.bfloat16) @ w.to(torch.bfloat16).T + out = x.to(torch.bfloat16) @ dq.T + # Relative error on the matmul output is bounded by the weight error (~6%). + rel = float((out.float() - ref.float()).abs().max() / (ref.float().abs().max() + 1e-6)) + assert rel < 0.2, f"matmul rel err too high: {rel}" + + +def test_scales_nonzero_and_finite(): + w = _weights() + scale = compute_fp8_scale(w, group_size=-1) + assert scale.item() > 0 + assert torch.isfinite(scale).all() + # A zero weight must not produce inf/nan scale. + zero = torch.zeros(8, 8) + scale0 = compute_fp8_scale(zero, group_size=-1) + assert torch.isfinite(scale0).all() and scale0.item() > 0 + + +def test_quantized_linear_module_matches_reference(): + from areno.engine.layers.linear import QuantizedLinear + + out_feat, in_feat = 32, 48 + ql = QuantizedLinear(in_feat, out_feat, group_size=-1, dtype=torch.bfloat16) + with torch.no_grad(): + w = _weights((out_feat, in_feat)) + ql.weight.copy_(w) + ql._requantize() + x = torch.randn(2, in_feat, dtype=torch.bfloat16) + out = ql(x) + ref = torch.nn.functional.linear(x, w.to(torch.bfloat16)) + rel = float((out.float() - ref.float()).abs().max() / (ref.float().abs().max() + 1e-6)) + assert rel < 0.2, f"QuantizedLinear rel err too high: {rel}" + + +def test_config_default_is_opt_in_unchanged(): + cfg = ModelConfig() + assert cfg.quant_method == "none" + assert cfg.quant_group_size == 128 + # Setting it to fp8 changes the field without changing any existing value. + cfg.quant_method = "fp8" + assert cfg.quant_method == "fp8"