perf: fast interior loop on the big cta and fused epilogue bias

- re-enable kFastLoop on the 128x128 CTA for congruous layouts: the base-pair fragment addressing freed the registers the old offset tables spilled, and the predication-free interior loop now wins across the band (fast body 142 SASS instr with zero predicated fallback vs 719/136 generic; 128 regs, no spill)
- move the big/small CTA dispatch boundary from 3/4 to 5/8 wave: with the fast big-CTA loop the crossover sits between 49 and 63 tiles (63-tile rect +8%, 1024^3 now takes the big CTA)
- fuse the linear bias into the GEMM epilogue: FP8Params.bias_ptr adds in fp32 before the single bf16 rounding, replacing the separate out + bias elementwise pass; guarded loads keep N tails exact and batch broadcast falls out of the row-major layout
- resolve Python None bias in the pybind layer (py::object + cast) so ops/fp8.py and fp8.py pass the argument through untouched; drop the _empty_bias sentinel machinery
- add fused-bias tests covering odd N tails, no-bias parity and batched broadcast

Benchmark: L20 (sm_89), CUDA-graph e2e. Big-CTA fast loop + dispatch: 1024^3 102.6->106.3T, 1152^3 128.5->133.3T, 2048^3 173.8->178.2T, 3072^3 180.2->185.3T, 8192^3 196.2->197.7T. Bias fusion (with-bias GEMM vs unfused out + bias): 1024^3 90.5->106.1T (+17%), 2048^3 162.2->178.3T (+10%), 4096^3 178.2->191.1T (+7%). Fused bias differs from the split path by <=1 bf16 ulp and is closer to the fp64 reference. 596 tests pass.

astrai/extension/fp8.py

@@ -315,27 +315,14 @@ def _is_fp8(dtype: torch.dtype) -> bool:
     return dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
 
 
-_zero_bias: Dict[Optional[int], torch.Tensor] = {}
-
-
-def _empty_bias(x: torch.Tensor) -> torch.Tensor:
-    """Per-device cached 0-element bf16 bias (the binding only checks numel —
-    never mutated), saving a CUDA allocation per bias-less linear."""
-    key = x.device.index
-    t = _zero_bias.get(key)
-    if t is None:
-        t = torch.empty(0, device=x.device, dtype=torch.bfloat16)
-        _zero_bias[key] = t
-    return t
-
-
 def fp8_linear_forward(
     x: torch.Tensor, w: torch.Tensor, bias=None, cfg: Optional[_ActiveConfig] = None
 ):
     """Scaled fp8 linear forward (called from the aten::linear impl).
 
     Composed from the two stateless primitives: quantize x/w with the active
-    scales, run the pre-quantized GEMM, add the bias. Delayed scaling folds
+    scales, run the pre-quantized GEMM with the bias fused into its epilogue.
+    Delayed scaling folds
     the returned amax into the history ring and publishes the next scale;
     dynamic scaling measures the current amax itself. Training quantizes the
     weight every step (the optimizer bumps its version, so there is no cast
@@ -345,17 +332,18 @@ def fp8_linear_forward(
     if cfg is None:
         cfg = _current_config()
     fmt = cfg.fp8_format.fwd()
-    if bias is None:
-        bias = _empty_bias(x)
     if isinstance(cfg.recipe, DynamicScaling):
         sx = _dynamic_scale(x.reshape(-1, w.size(1)), cfg.recipe, fmt)
         sw = _dynamic_scale(w, cfg.recipe, fmt)
         x8, _ = quantize(x, sx.reciprocal(), fmt)
         w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0]
-        out = mm_fp8(x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True).reshape(
+        # Bias fuses into the GEMM epilogue (fp32 add before the single bf16
-            *x.shape[:-1], w.size(0)
+        # rounding — one rounding fewer than the separate out + bias pass);
-        )
+        # None passes through to the kernel's no-bias path.
-        return (out + bias if bias.numel() else out), sx, sw
+        out = mm_fp8(
+            x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
+        ).reshape(*x.shape[:-1], w.size(0))
+        return out, sx, sw
 
     meta = state.get_weight_meta(w)
     if not meta.w.initialized:
@@ -368,11 +356,9 @@ def fp8_linear_forward(
         w8, amax_w = w, None
     else:
         w8, amax_w = quantize(w, sw.reciprocal(), fmt)
-    out = mm_fp8(x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True).reshape(
+    out = mm_fp8(
-        *x.shape[:-1], w.size(0)
+        x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
-    )
+    ).reshape(*x.shape[:-1], w.size(0))
-    if bias.numel():
-        out = out + bias
     meta.x.update(amax_x, fmt)
     if amax_w is not None:
         meta.w.update(amax_w, fmt)

astrai/extension/ops/fp8.py

@@ -14,7 +14,7 @@ Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``;
 this module is stateless.
 """
 
-from typing import Tuple
+from typing import Optional, Tuple
 
 import torch
 from torch.library import custom_op
@@ -84,16 +84,19 @@ def fp8_gemm(
     scale: torch.Tensor,
     trans_a: int = 0,
     trans_b: int = 0,
+    bias: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
-    """FP8 GEMM: ``a @ b * scale`` with FP32 accumulation.
+    """FP8 GEMM: ``a @ b * scale (+ bias)`` with FP32 accumulation.
 
     2D or 3D (batched) operands; a size-1 batch broadcasts (matmul rules).
-    The result is always BF16; FP8 output is a separate quantize operation.
+    ``bias`` (bf16, length n) fuses into the epilogue in fp32 before the
+    single bf16 rounding. The result is always BF16; FP8 output is a
+    separate quantize operation.
     """
 
 
 @fp8_gemm.register_fake
-def _fp8_gemm_fake(a, b, scale, trans_a=0, trans_b=0):
+def _fp8_gemm_fake(a, b, scale, trans_a=0, trans_b=0, bias=None):
     dtype = torch.bfloat16
     rows = a.size(2) if trans_a else a.size(1)
     cols = b.size(1) if trans_b else b.size(2)
@@ -103,19 +106,21 @@ def _fp8_gemm_fake(a, b, scale, trans_a=0, trans_b=0):
 
 
 @fp8_gemm.register_kernel("cuda")
-def _fp8_gemm_cuda(a, b, scale, trans_a=0, trans_b=0):
+def _fp8_gemm_cuda(a, b, scale, trans_a=0, trans_b=0, bias=None):
     if a.dtype != b.dtype or a.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2):
         raise TypeError(
             f"fp8 GEMM requires matching fp8 inputs, got {a.dtype}/{b.dtype}"
         )
-    return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b)
+    return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b, bias)
 
 
 @fp8_gemm.register_kernel("cpu")
-def _fp8_gemm_cpu(a, b, scale, trans_a=0, trans_b=0):
+def _fp8_gemm_cpu(a, b, scale, trans_a=0, trans_b=0, bias=None):
     aa = a.float().transpose(-2, -1) if trans_a else a.float()
     bb = b.float().transpose(-2, -1) if trans_b else b.float()
     acc = aa @ bb * scale
+    if bias is not None and bias.numel() > 0:
+        acc = acc + bias.float()
     return acc.to(torch.bfloat16)
 
 
@@ -149,21 +154,26 @@ def mm_fp8(
     scale: torch.Tensor,
     trans_a: bool = False,
     trans_b: bool = False,
+    bias: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
-    """Pre-quantized FP8 GEMM: ``a @ b * scale``.
+    """Pre-quantized FP8 GEMM: ``a @ b * scale (+ bias)``.
 
     ``a``/``b`` must be FP8 tensors of the same format, 2D or 3D (batched,
     matmul-style broadcast on the batch dim). Inner-transposed views (e.g.
     ``x.t()``) fold into the layout at zero copy. ``scale`` is their combined
-    dequantization scale. The result is BF16; FP8 output is a separate
+    dequantization scale. ``bias`` (CUDA bf16 1D of length n) adds inside the
-    quantize operation.
+    kernel epilogue in fp32 — no separate elementwise pass. The result is
+    BF16; FP8 output is a separate quantize operation.
     """
     # Same hot-path bypass as quantize(): the binding's TORCH_CHECKs keep
-    # validation identical on the direct route.
+    # validation identical on the direct route (bias may be None — the
+    # binding resolves it to the no-bias path).
     if (
         type(a) is torch.Tensor
         and a.is_cuda
         and a.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
     ):
-        return get_module("fp8_ops").mm_fp8(a, b, scale, int(trans_a), int(trans_b))
+        return get_module("fp8_ops").mm_fp8(
-    return fp8_gemm(a, b, scale, trans_a, trans_b)
+            a, b, scale, int(trans_a), int(trans_b), bias
+        )
+    return fp8_gemm(a, b, scale, trans_a, trans_b, bias)

csrc/kernels/fp8/common.h

@@ -103,8 +103,12 @@ struct FP8QuantizeParams {
 struct FP8Params {
     // Inputs: a/b are FP8 for the pre-quantized path. Scales are
     // quantization steps (device scalars).
+    // Optional bf16 bias broadcast over output rows (fused into the epilogue
+    // before the bf16 rounding, so it adds in fp32 — one rounding fewer than
+    // the separate out + bias elementwise kernel it replaces). Null disables.
     const void* __restrict__ a_ptr = nullptr;
     const void* __restrict__ b_ptr = nullptr;
+    const void* __restrict__ bias_ptr = nullptr;
     void* __restrict__ out_ptr = nullptr;
 
     const float* __restrict__ scale = nullptr;

csrc/kernels/fp8/gemm.cuh

@@ -613,6 +613,13 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
     // scatter and the gather conflict-free: a lane quad's chunk and the 8
     // rows of one gather phase map to distinct 4-bank groups.
     const float output_scale = scale;
+    // Fused bias (idea B): added to the fp32 accumulator before the single
+    // bf16 rounding — one fewer rounding than the out + bias elementwise
+    // pass this replaces, and no extra kernel launch / m*n round-trip. The
+    // per-lane loads (2 per nt, kMt-times re-read) are L1 broadcasts; rows
+    // past the N edge skip the load (their smem slots never copy out).
+    const __nv_bfloat16* bias =
+        reinterpret_cast<const __nv_bfloat16*>(p.bias_ptr);
     __nv_bfloat16* tile_out = reinterpret_cast<__nv_bfloat16*>(fp8_gemm_smem);
     constexpr int kRowChunks = kBlockN / 8;  // 16B chunks per tile row
     static_assert(kBlockM * kBlockN * 2 <=
@@ -624,9 +631,15 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
                ((c ^ (r & (kRowChunks - 1))) * 8);
     };
     const int local_col0 = warp_n * Traits::kWarpN + thread_in_group * 2;
+    const int64_t bias_col0 = (int64_t)block_n * kBlockN;
 #pragma unroll
     for (int nt = 0; nt < kNt; ++nt) {
         const int col = local_col0 + nt * 8;
+        const int64_t gcol = bias_col0 + col;
+        const float b0 =
+            bias && gcol < n ? __bfloat162float(bias[gcol]) : 0.0f;
+        const float b1 =
+            bias && gcol + 1 < n ? __bfloat162float(bias[gcol + 1]) : 0.0f;
 #pragma unroll
         for (int mt = 0; mt < kMt; ++mt) {
             const int r0 = warp_m * Traits::kWarpM + group + mt * 16;
@@ -635,12 +648,12 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
             // m16n8 output, columns tig*2 and tig*2+1 inside one 16B chunk.
             const int off = col & 7;  // element offset within the chunk
             *reinterpret_cast<__nv_bfloat162*>(out_chunk(r0, col >> 3) + off) =
-                __floats2bfloat162_rn(tile_acc[0] * output_scale,
+                __floats2bfloat162_rn(tile_acc[0] * output_scale + b0,
-                                      tile_acc[1] * output_scale);
+                                      tile_acc[1] * output_scale + b1);
             *reinterpret_cast<__nv_bfloat162*>(out_chunk(r0 + 8, col >> 3) +
                                                off) =
-                __floats2bfloat162_rn(tile_acc[2] * output_scale,
+                __floats2bfloat162_rn(tile_acc[2] * output_scale + b0,
-                                      tile_acc[3] * output_scale);
+                                      tile_acc[3] * output_scale + b1);
         }
     }
     __syncthreads();
@@ -751,13 +764,15 @@ void launch_with_smem(int smem_bytes, dim3 grid, dim3 block,
 // congruous NT): the 128x128 CTA wins inside one full wave (81 tiles: big
 // +24%) and from ~1.5 waves up (144: +23%, 256: +39%, 2048^3 123->171 TF),
 // but loses inside the quantization dip just past one wave (100 tiles =
-// 1.09 waves: big -8%) where the finer 64x64 grid fills the tail. Below
+// 1.09 waves: big -8%) where the finer 64x64 grid fills the tail. With the
-// 3/4 wave the small CTA's extra residency wins or ties (64 tiles: tie).
+// interior fast loop on the big CTA the sub-wave boundary moved down: 63-64
-// So: big CTA iff tiles are in [3/4, 1] wave or >= 7/5 waves.
+// tiles already favor it (63-tile rect +8%, 1024^3 +2%) while 49 tiles
+// stays small-CTA territory, so the big band opens at 5/8 wave instead of
+// 3/4.
 inline bool prefer_small_cta(int64_t tiles_128, int64_t m) {
     if (m <= 64) return true;
     const int64_t waves = device_sm_count();
-    if (tiles_128 >= waves - waves / 4 && tiles_128 <= waves) return false;
+    if (tiles_128 >= waves * 5 / 8 && tiles_128 <= waves) return false;
     return tiles_128 < waves + waves * 2 / 5;
 }
 
@@ -806,8 +821,18 @@ void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
     }
     using Traits = Fp8GemmTraits<Fmt, 128, 128, kK, Stages>;
     dim3 grid((p.n + 127) / 128, (p.m + 127) / 128, p.batch);
+    // Interior-loop specialization on the big CTA as well: with the base-pair
+    // fragment addressing the doubled mainloop no longer spills, and the
+    // predication-free loads win across the band (measured, L20: 1024^3
+    // 98->103T, 2048^3 172->177T, 8192^3 201->205T, 896x1280 124->135T; the
+    // pre-base-pair attempt regressed ~3% at 131 regs). Only congruous
+    // layouts can enter fast_cta, so crosswise (TN) instantiations keep the
+    // single generic body — no dead second loop in their I-cache.
+    constexpr bool kBigFast = !std::is_same_v<LayoutA, ColMajor> &&
+                              !std::is_same_v<LayoutB, RowMajor>;
     launch_with_smem<
-        fp8_gemm_kernel<Traits, LayoutA, LayoutB, GroupRaster, false, false>>(
+        fp8_gemm_kernel<Traits, LayoutA, LayoutB, GroupRaster, false, false,
+                        kBigFast>>(
         Fp8GemmSmem<Traits, LayoutA, LayoutB, false>::kBytes, grid,
         dim3(Traits::kCtaThreads), stream, p);
 }

csrc/kernels/fp8/ops.cu

@@ -164,7 +164,7 @@ std::tuple<torch::Tensor, torch::Tensor> quantize(torch::Tensor x,
 }
 
 torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
-                     int64_t trans_a, int64_t trans_b) {
+                     int64_t trans_a, int64_t trans_b, torch::Tensor bias) {
     TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required");
     TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn ||
                     a.scalar_type() == torch::kFloat8_e5m2,
@@ -208,6 +208,16 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
     FP8Params p;
     pack_gemm(p, a_st.data_ptr(), b_st.data_ptr(), output.data_ptr(), scale,
               m, n, k, a_ld, b_ld);
+    // Fused epilogue bias (bf16, broadcast over rows and batches). An
+    // undefined or 0-element tensor keeps the plain scaled output.
+    if (bias.defined() && bias.numel() > 0) {
+        TORCH_CHECK(bias.is_cuda() && bias.scalar_type() == torch::kBFloat16,
+                    "fp8 gemm bias must be a CUDA bf16 tensor");
+        TORCH_CHECK(bias.dim() == 1 && bias.size(0) == n,
+                    "fp8 gemm bias must be 1D of length n=", n);
+        TORCH_CHECK(bias.is_contiguous(), "fp8 gemm bias must be contiguous");
+        p.bias_ptr = bias.data_ptr();
+    }
     p.batch = static_cast<int>(batch);
     p.a_batch_stride = (batch_a == 1 && batch > 1) ? 0 : a_bstride;
     p.b_batch_stride = (batch_b == 1 && batch > 1) ? 0 : b_bstride;
@@ -220,9 +230,29 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
     return output;
 }
 
+// mm_fp8 binding: Python None and an omitted argument both mean "no bias"
+// (resolved to an undefined tensor here, so every Python layer can pass its
+// bias argument through untouched instead of normalizing it host-side).
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("quantize", &quantize, py::arg("x"), py::arg("scale"),
           py::arg("fmt"));
-    m.def("mm_fp8", &mm_fp8, py::arg("a"), py::arg("b"), py::arg("scale"),
+    m.def(
-          py::arg("trans_a") = 0, py::arg("trans_b") = 0);
+        "mm_fp8",
+        [](torch::Tensor a, torch::Tensor b, torch::Tensor scale,
+           int64_t trans_a, int64_t trans_b, py::object bias) {
+            torch::Tensor t;
+            if (!bias.is_none()) {
+                // (py::isinstance<torch::Tensor> is false for real tensors
+                // here — torch's caster registers no pybind type info — so
+                // validate by attempting the cast itself.)
+                try {
+                    t = bias.cast<torch::Tensor>();
+                } catch (const py::cast_error&) {
+                    TORCH_CHECK(false, "bias must be a torch.Tensor or None");
+                }
+            }
+            return mm_fp8(a, b, scale, trans_a, trans_b, t);
+        },
+        py::arg("a"), py::arg("b"), py::arg("scale"), py::arg("trans_a") = 0,
+        py::arg("trans_b") = 0, py::arg("bias") = py::none());
 }

tests/extension/test_fp8_mma.py

@@ -116,6 +116,40 @@ def test_mm_fp8_transposed_operands(trans_a, trans_b):
     torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
 
 
+@skip_no_fp8
+@pytest.mark.parametrize("bias_on", [False, True])
+def test_mm_fp8_fused_bias(bias_on):
+    """Epilogue-fused bias matches the unfused out + bias reference (single
+    fp32 rounding vs the reference's double rounding keeps it within 1 ulp),
+    including N-tail columns and batched broadcast."""
+    torch.manual_seed(31)
+    m, n, k = 19, 13, 37  # odd n exercises the guarded bias loads
+    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
+    sa, sb = _scale(a), _scale(b)
+    a8, _ = quantize(a, sa.reciprocal(), "e4m3")
+    b8, _ = quantize(b, sb.reciprocal(), "e4m3")
+    bias = torch.randn(n, device="cuda", dtype=torch.bfloat16)
+
+    out = mm_fp8(a8, b8, sa * sb, trans_b=True, bias=bias if bias_on else None)
+    base = (_quantize(a, sa) @ _quantize(b, sb).t() * sa * sb).to(torch.bfloat16)
+    expected = base + bias if bias_on else base
+    # bias is O(1) against O(sqrt(k)) accumulators: absolute tolerance rules
+    torch.testing.assert_close(out, expected, atol=0.13, rtol=0.01)
+
+    # Batched broadcast: bias applies to every batch slice (each slice gets
+    # its own reference from its own operand values).
+    ab = torch.randn(3, m, k, device="cuda", dtype=torch.bfloat16)
+    ab8, _ = quantize(ab, sa.reciprocal(), "e4m3")
+    outb = mm_fp8(ab8, b8, sa * sb, trans_b=True, bias=bias)
+    assert outb.shape == (3, m, n)
+    for i in range(3):
+        expected_b = (_quantize(ab[i], sa) @ _quantize(b, sb).t() * sa * sb).to(
+            torch.bfloat16
+        ) + bias
+        torch.testing.assert_close(outb[i], expected_b, atol=0.13, rtol=0.01)
+
+
 @skip_no_fp8
 @pytest.mark.parametrize("trans_a", [False, True])
 @pytest.mark.parametrize("trans_b", [False, True])