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

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<torch::Tensor> out);
 std::vector<torch::Tensor> 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");

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)

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

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

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)

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."""