feat(fp8): E4M3 fused dequant-linear kernel + FP8 quantized-linear slice (RFC 0001)

[CAICAIIs]

Summary

Adds the FP8 weight-quantization vertical slice (RFC 0001, M1) and an E4M3 fused dequant-linear CUDA kernel for the memory-bound decode path.

Decode re-reads the full weight tensor from HBM for every generated token, so it is memory-bandwidth-bound. Storing the weight as 1-byte E4M3 FP8 halves the weight bytes read, which is the memory-bound benefit from RFC 0001. On Hopper / H20 the built-in torch._scaled_mm FP8 path delivers 1.53–1.72× decode speedup vs bf16; on Ampere / A100 (which has no built-in FP8 matmul) a custom E4M3 dequant-linear kernel is provided.

quant_method is opt-in and unchanged by default ("none"); unmarked weights keep the existing bf16 path.

Motivation

On the 8×A100 node, decode is weight-read-bound and compute sits idle. Quantizing weights to FP8 (1 byte) vs bf16 (2 bytes) halves the weight-read bytes. This PR lands the M1 vertical slice: an opt-in FP8 decode path + an Ampere E4M3 kernel, both forward-only.

Changes

• areno/accel/csrc/e4m3_linear.cu — E4M3 fused dequant-linear. A branchless bit-trick decode (validated byte-for-byte vs torch.float8_e4m3fn) plus:
  • a small-M (M <= 4) memory-streaming GEMV: one warp streams one weight row with coalesced 8-byte loads, reuses the staged activation, and reduces with warp shuffles;
  • a WMMA (fp16 16×16×16) tensor-core GEMM fallback for M > 4.
    Per-tensor scalar scale applied once post-dot.
• areno/accel/kernels/e4m3_cuda.py — quantized_e4m3_linear_cuda shim + quantize/dequant helpers, with an out-buffer reuse for the hot decode path.
• areno/engine/quantization.py — FP8 (E4M3) scale/quantize/dequant reference, CPU-portable.
• areno/engine/layers/linear.py — QuantizedLinear (W8A16 dequant-forward reference) + _areno_linear_forward routes FP8-marked weights; unmarked weights unchanged.
• areno/engine/{config,modeling,worker}.py — quant_method / quant_group_size (opt-in), quantize_model_weights_fp8, and a guard that rejects quant_method='fp8' in a train worker (forward-only, no backward).
• areno/accel/kernels/fp8_linear.py — the model's E5M2 W8A16 Triton dequant-linear (the grid Triton accepts on Ampere).
• setup.py / areno/accel/csrc/extension.cpp — register the new kernel (optional out arg).
• Tests — tests/test_e4m3_decode_cpu.py, tests/test_fp8_quant_cpu.py (decode formula vs torch, quant/dequant round-trip, QuantizedLinear reference, opt-in default).
• Benchmarks — scripts/bench/{e4m3_cuda_bench,h20_fp8_scaled_mm,fp8_end_to_end_bench,triton_fp8_matmul_bench}.py.
• docs/rfcs/0001-weight-quantization-fp8-int4.md — design, gates, measured results.

Testing

• CPU: pytest tests/test_e4m3_decode_cpu.py tests/test_fp8_quant_cpu.py → 11 passed.
• GPU (A100): kernel correct vs bf16 dequant reference — max-rel err ~0.3–0.4% across M ∈ {1, 2, 4, 64}; CUDA build clean.
• Decode benchmark:
  • A100: E4M3 GEMV is memory-bound (raw read-only kernel ≈1248 GB/s ≈ 2.1× ceiling; the GEMV's practical ceiling ≈1.68× at the bare-kernel level).
  • H20 / Hopper: built-in torch._scaled_mm FP8 E4M3 → 1.53–1.72× vs bf16 (RFC §4.8).

Known limitations

• Forward-only: FP8 (E4M3/E5M2) has no backward, so it is decode-only and cannot train (guarded in worker.py).
• Ampere grid: Triton rejects E4M3 on A100, so the model's runtime FP8 decode uses E5M2 (coarse ~12 %); E4M3 is the dedicated Ampere kernel and the Hopper _scaled_mm path.
• The Ampere E4M3 GEMV is memory-bound at the bare-kernel level, but through the torch/ATen extension a per-call dispatch overhead dominates the fast kernel, so end-to-end on A100 is near parity. The >1.5× number is the Hopper _scaled_mm result.

Related

• RFC: docs/rfcs/0001-weight-quantization-fp8-int4.md
• Companion issue: [RFC] FP8 / INT4 weight quantization for AReno (CUDA train & decode) #514 (RFC, OPEN)