diff --git a/astrai/extension/fp8.py b/astrai/extension/fp8.py index b4e9c94..e778c82 100644 --- a/astrai/extension/fp8.py +++ b/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( - *x.shape[:-1], w.size(0) - ) - return (out + bias if bias.numel() else out), sx, sw + # Bias fuses into the GEMM epilogue (fp32 add before the single bf16 + # rounding — one rounding fewer than the separate out + bias pass); + # None passes through to the kernel's no-bias path. + 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( - *x.shape[:-1], w.size(0) - ) - if bias.numel(): - out = out + bias + out = mm_fp8( + x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias + ).reshape(*x.shape[:-1], w.size(0)) meta.x.update(amax_x, fmt) if amax_w is not None: meta.w.update(amax_w, fmt) diff --git a/astrai/extension/ops/fp8.py b/astrai/extension/ops/fp8.py index 992c4e7..038c54c 100644 --- a/astrai/extension/ops/fp8.py +++ b/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 - quantize operation. + dequantization scale. ``bias`` (CUDA bf16 1D of length n) adds inside the + 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 fp8_gemm(a, b, scale, trans_a, trans_b) + return get_module("fp8_ops").mm_fp8( + a, b, scale, int(trans_a), int(trans_b), bias + ) + return fp8_gemm(a, b, scale, trans_a, trans_b, bias) diff --git a/csrc/kernels/fp8/common.h b/csrc/kernels/fp8/common.h index 6d1ba0d..1bc6bf6 100644 --- a/csrc/kernels/fp8/common.h +++ b/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; diff --git a/csrc/kernels/fp8/gemm.cuh b/csrc/kernels/fp8/gemm.cuh index ce72ffd..eb4cfdd 100644 --- a/csrc/kernels/fp8/gemm.cuh +++ b/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(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, - tile_acc[1] * output_scale); + __floats2bfloat162_rn(tile_acc[0] * output_scale + b0, + tile_acc[1] * output_scale + b1); *reinterpret_cast<__nv_bfloat162*>(out_chunk(r0 + 8, col >> 3) + off) = - __floats2bfloat162_rn(tile_acc[2] * output_scale, - tile_acc[3] * output_scale); + __floats2bfloat162_rn(tile_acc[2] * output_scale + b0, + 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 -// 3/4 wave the small CTA's extra residency wins or ties (64 tiles: tie). -// So: big CTA iff tiles are in [3/4, 1] wave or >= 7/5 waves. +// 1.09 waves: big -8%) where the finer 64x64 grid fills the tail. With the +// interior fast loop on the big CTA the sub-wave boundary moved down: 63-64 +// 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; 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 && + !std::is_same_v; launch_with_smem< - fp8_gemm_kernel>( + fp8_gemm_kernel>( Fp8GemmSmem::kBytes, grid, dim3(Traits::kCtaThreads), stream, p); } diff --git a/csrc/kernels/fp8/ops.cu b/csrc/kernels/fp8/ops.cu index 7fcfc4e..05d7cee 100644 --- a/csrc/kernels/fp8/ops.cu +++ b/csrc/kernels/fp8/ops.cu @@ -164,7 +164,7 @@ std::tuple 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(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"), - py::arg("trans_a") = 0, py::arg("trans_b") = 0); + m.def( + "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 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(); + } 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()); } diff --git a/tests/extension/test_fp8_mma.py b/tests/extension/test_fp8_mma.py index afba3a1..da6458b 100644 --- a/tests/extension/test_fp8_mma.py +++ b/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])