feat(norm): fused Add+RMSNorm+1x128 fp8 block-quant producer

benchmarks/bench_fused_add_rmsnorm_fp8_block_quant.py

@@ -0,0 +1,163 @@
+"""Benchmark: fused Add+RMSNorm+1x128-fp8-block-quant (one kernel) vs the 2-kernel producer
+(flashinfer.fused_add_rmsnorm then a per-token-group 1x128 fp8 quant), which is what runs before
+fp8 block-scaled GEMMs in deepseek_v2-family / MoE models.
+
+Shows the kernel-level win: the fused kernel does add+norm+quant in a single memory pass, saving
+the intermediate bf16 round-trip.
+"""
+
+import numpy as np
+import torch
+
+import flashinfer
+
+
+def _round_up(x, m):
+    return (x + m - 1) // m * m
+
+
+def _run_fused(
+    out, block_scale, normed_out, input, residual, weight, eps, enable_pdl=False
+):
+    """Call the fused op via the public API if present, else the JIT module export directly."""
+    if hasattr(flashinfer, "fused_add_rmsnorm_fp8_block_quant"):
+        flashinfer.fused_add_rmsnorm_fp8_block_quant(
+            out,
+            block_scale,
+            normed_out,
+            input,
+            residual,
+            weight,
+            eps,
+            enable_pdl=enable_pdl,
+        )
+    else:
+        flashinfer.norm.get_norm_module().fused_add_rmsnorm_fp8_block_quant(
+            out, block_scale, normed_out, input, residual, weight, eps, enable_pdl
+        )
+
+
+def _ref_group_quant_fp8(normed, block=128):
+    """Reference 1x128 per-token-group fp8 quant (torch) -> (fp8, scale[M, H/128])."""
+    M, H = normed.shape
+    x = normed.float().view(M, H // block, block)
+    amax = x.abs().amax(dim=-1).clamp_min(1e-4)
+    scale = amax / 448.0
+    q = (x / scale.unsqueeze(-1)).clamp(-448, 448).to(torch.float8_e4m3fn)
+    return q.view(M, H), scale
+
+
+def bench_one(M, H, dtype=torch.bfloat16, iters=100, warmup=20):
+    dev = "cuda"
+    input = torch.randn(M, H, dtype=dtype, device=dev) * 0.1
+    residual = torch.randn(M, H, dtype=dtype, device=dev) * 0.1
+    weight = torch.randn(H, dtype=dtype, device=dev)
+    eps = 1e-6
+
+    # fused outputs
+    out = torch.empty(M, H, dtype=torch.float8_e4m3fn, device=dev)
+    block_scale = torch.empty(
+        H // 128, _round_up(M, 4), dtype=torch.float32, device=dev
+    )
+    normed_out = torch.empty(M, H, dtype=dtype, device=dev)
+
+    # try to use a production per-token-group quant for the baseline's 2nd kernel; else torch ref.
+    try:
+        from sglang.srt.layers.quantization.fp8_kernel import per_token_group_quant_fp8
+
+        def baseline_quant(n):
+            return per_token_group_quant_fp8(n, 128)
+    except Exception:
+
+        def baseline_quant(n):
+            return _ref_group_quant_fp8(n)
+
+    # Timing loops run the kernels in place (no setup clones); residual grows ~linearly (stays bounded),
+    # and memory traffic per call is identical regardless of values -> fair kernel-time comparison.
+
+    def fused():
+        _run_fused(out, block_scale, normed_out, input, residual, weight, eps)
+
+    def baseline():
+        flashinfer.fused_add_rmsnorm(
+            input, residual, weight, eps
+        )  # residual <- add+norm (bf16), in place
+        baseline_quant(residual)  # 2nd kernel: 1x128 fp8 quant of the bf16 normed
+
+    def time_it(fn):
+        for _ in range(warmup):
+            fn()
+        torch.cuda.synchronize()
+        s, e = torch.cuda.Event(True), torch.cuda.Event(True)
+        ts = []
+        for _ in range(iters):
+            s.record()
+            fn()
+            e.record()
+            torch.cuda.synchronize()
+            ts.append(s.elapsed_time(e))
+        return float(np.median(ts))  # ms
+
+    t_f = time_it(fused)
+    t_b = time_it(baseline)
+    # fused byte model: read input+residual+weight (bf16), write residual+normed(bf16)+fp8+fp32 scale
+    bytes_f = (
+        (3 * M * H * 2)
+        + (2 * M * H * 2)
+        + (M * H * 1)
+        + (H // 128 * _round_up(M, 4) * 4)
+    )
+    gbps = bytes_f / (t_f * 1e-3) / 1e9
+    print(
+        f"M={M:>5} H={H:>5} | fused {t_f * 1e3:7.1f}us  2-kernel {t_b * 1e3:7.1f}us  "
+        f"speedup {t_b / t_f:4.2f}x  | fused {gbps:6.0f} GB/s"
+    )
+    return t_b / t_f
+
+
+def check_correctness(M=99, H=4096, dtype=torch.bfloat16):
+    dev = "cuda"
+    input = torch.randn(M, H, dtype=dtype, device=dev) * 0.1
+    residual = torch.randn(M, H, dtype=dtype, device=dev) * 0.1
+    weight = torch.randn(H, dtype=dtype, device=dev)
+    eps = 1e-6
+    # reference
+    x = input.float() + residual.float()
+    var = x.pow(2).mean(-1, keepdim=True)
+    normed_ref = (x * torch.rsqrt(var + eps) * weight.float()).to(
+        dtype
+    )  # round to bf16 like kernel
+    q_ref, scale_ref = _ref_group_quant_fp8(normed_ref)  # scale_ref [M, H/128]
+    # fused
+    out = torch.empty(M, H, dtype=torch.float8_e4m3fn, device=dev)
+    block_scale = torch.empty(
+        H // 128, _round_up(M, 4), dtype=torch.float32, device=dev
+    )
+    normed_out = torch.empty(M, H, dtype=dtype, device=dev)
+    r = residual.clone()
+    _run_fused(out, block_scale, normed_out, input, r, weight, eps)
+    # checks
+    torch.testing.assert_close(r, x.to(dtype), rtol=1e-2, atol=1e-2)  # residual
+    torch.testing.assert_close(
+        normed_out.float(), normed_ref.float(), rtol=1e-2, atol=1e-2
+    )  # normed
+    scale_got = block_scale.transpose(0, 1)[:M]  # -> [M, H/128]
+    torch.testing.assert_close(
+        scale_got, scale_ref, rtol=1e-3, atol=1e-6
+    )  # block scale
+    # dequant fp8 vs reference normed
+    deq = out.float().view(M, H // 128, 128) * scale_got.unsqueeze(-1)
+    torch.testing.assert_close(deq.view(M, H), normed_ref.float(), rtol=0.1, atol=0.3)
+    print("correctness: OK (residual, normed_out, block_scale, dequant all match)")
+
+
+if __name__ == "__main__":
+    print(f"flashinfer {flashinfer.__version__} | {torch.cuda.get_device_name()}")
+    check_correctness()
+    speedups = []
+    for M in (256, 1024, 4096):
+        for H in (4096, 6144, 8192):
+            speedups.append(bench_one(M, H))
+    print(
+        f"\ngeomean speedup (fused vs 2-kernel): {np.exp(np.mean(np.log(speedups))):.2f}x"
+    )

csrc/flashinfer_norm_binding.cu

@@ -26,6 +26,10 @@ void fused_add_rmsnorm(TensorView input, TensorView residual, TensorView weight,
 void fused_add_rmsnorm_quant(TensorView output, TensorView input, TensorView residual,
                              TensorView weight, TensorView scale, double eps, bool enable_pdl);
 
+void fused_add_rmsnorm_fp8_block_quant(TensorView output, TensorView block_scale,
+                                       TensorView normed_out, TensorView input, TensorView residual,
+                                       TensorView weight, double eps, bool enable_pdl);
+
 void gemma_rmsnorm(TensorView out, TensorView input, TensorView weight, double eps,
                    bool enable_pdl);
 
@@ -58,6 +62,7 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(rmsnorm, rmsnorm);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(rmsnorm_quant, rmsnorm_quant);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(fused_add_rmsnorm, fused_add_rmsnorm);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(fused_add_rmsnorm_quant, fused_add_rmsnorm_quant);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(fused_add_rmsnorm_fp8_block_quant, fused_add_rmsnorm_fp8_block_quant);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(gemma_rmsnorm, gemma_rmsnorm);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(gemma_fused_add_rmsnorm, gemma_fused_add_rmsnorm);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(layernorm, layernorm);

csrc/norm.cu

@@ -186,6 +186,57 @@ void fused_add_rmsnorm_quant(TensorView output, TensorView input, TensorView res
   });
 }
 
+// Add-residual + RMSNorm + 1x128 fp8 block-quant. Emits fp8 `output` [M,d], fp32 `block_scale`
+// (d/128, round_up(M,4)) column-major TMA layout, bf16 `normed_out` [M,d], and updates `residual`
+// in place with (input+residual). See norm::FusedAddRMSNormFP8BlockQuant.
+void fused_add_rmsnorm_fp8_block_quant(TensorView output, TensorView block_scale,
+                                       TensorView normed_out, TensorView input, TensorView residual,
+                                       TensorView weight, double eps, bool enable_pdl) {
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(residual);
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(weight);
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(output);
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(normed_out);
+  CHECK_DEVICE(input, residual);
+  CHECK_DEVICE(input, weight);
+  CHECK_DEVICE(input, output);
+  CHECK_DEVICE(input, block_scale);
+  CHECK_DEVICE(input, normed_out);
+  CHECK_DIM(2, input);     // (batch_size, hidden_size)
+  CHECK_DIM(2, residual);  // (batch_size, hidden_size)
+  CHECK_DIM(1, weight);    // (hidden_size)
+  CHECK_DIM(2, output);    // fp8 (batch_size, hidden_size)
+  CHECK_DIM(2, normed_out);
+  CHECK_DIM(2, block_scale);  // (hidden_size/128, round_up(batch_size, 4))
+  unsigned int batch_size = input.size(0);
+  unsigned int hidden_size = input.size(1);
+  unsigned int m_pad = (batch_size + 3) & ~3u;
+  TVM_FFI_ICHECK_EQ(residual.size(0), batch_size);
+  TVM_FFI_ICHECK_EQ(residual.size(1), hidden_size);
+  TVM_FFI_ICHECK_EQ(weight.size(0), hidden_size);
+  TVM_FFI_ICHECK_EQ(normed_out.size(0), batch_size);
+  TVM_FFI_ICHECK_EQ(normed_out.size(1), hidden_size);
+  TVM_FFI_ICHECK_EQ(block_scale.size(0), hidden_size / 128);
+  TVM_FFI_ICHECK_EQ(block_scale.size(1), m_pad);
+  TVM_FFI_ICHECK_EQ(hidden_size % 128, 0);
+  ffi::CUDADeviceGuard device_guard(input.device().device_id);
+  const cudaStream_t stream = get_stream(input.device());
+
+  // 1x128 block-scale activations are e4m3 (the layout DeepGEMM/flashinfer fp8 block GEMMs
+  // consume).
+  DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(input.dtype(), c_type, [&] {
+    cudaError_t status = norm::FusedAddRMSNormFP8BlockQuant<c_type, __nv_fp8_e4m3>(
+        static_cast<c_type*>(input.data_ptr()), static_cast<c_type*>(residual.data_ptr()),
+        static_cast<c_type*>(weight.data_ptr()), static_cast<__nv_fp8_e4m3*>(output.data_ptr()),
+        static_cast<float*>(block_scale.data_ptr()), static_cast<c_type*>(normed_out.data_ptr()),
+        batch_size, hidden_size, input.stride(0), residual.stride(0), output.stride(0),
+        block_scale.stride(0), eps, enable_pdl, stream);
+    TVM_FFI_ICHECK(status == cudaSuccess)
+        << "FusedAddRMSNormFP8BlockQuant failed with error code " << cudaGetErrorString(status);
+    return true;
+  });
+}
+
 void gemma_rmsnorm(TensorView output, TensorView input, TensorView weight, double eps,
                    bool enable_pdl) {
   CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);

flashinfer/__init__.py

@@ -135,6 +135,9 @@
 from . import msa_ops as msa_ops
 from .norm import fused_add_rmsnorm as fused_add_rmsnorm
 from .norm import fused_add_rmsnorm_quant as fused_add_rmsnorm_quant
+from .norm import (
+    fused_add_rmsnorm_fp8_block_quant as fused_add_rmsnorm_fp8_block_quant,
+)
 from .norm import layernorm as layernorm
 from .norm import layernorm_quant as layernorm_quant
 from .norm import gemma_fused_add_rmsnorm as gemma_fused_add_rmsnorm

flashinfer/norm/__init__.py

@@ -389,6 +389,74 @@ def _fused_add_rmsnorm_quant_fake(
     pass
 
 
+@register_custom_op(
+    "flashinfer::fused_add_rmsnorm_fp8_block_quant",
+    mutates_args=("out", "block_scale", "normed_out", "residual"),
+)
+def fused_add_rmsnorm_fp8_block_quant(
+    out: torch.Tensor,
+    block_scale: torch.Tensor,
+    normed_out: torch.Tensor,
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    enable_pdl: Optional[bool] = None,
+) -> None:
+    r"""Fused add-residual + RMSNorm + 1x128 fp8 block quantization, in one pass.
+
+    Replaces the two-kernel producer (:func:`fused_add_rmsnorm` then a per-token-group 1x128 fp8
+    quant) used before fp8 block-scaled GEMMs. Unlike :func:`fused_add_rmsnorm_quant` (scalar
+    per-tensor scale) this emits a **dynamic per-1x128-block** fp32 scale in the column-major
+    TMA-aligned layout that :mod:`flashinfer.deep_gemm` consumes.
+
+    Step 1: ``residual = input + residual`` (pre-norm, written in place).
+    Step 2: ``normed = RMSNorm(residual) * weight`` (bf16, written to ``normed_out``).
+    Step 3: ``out, block_scale = quant_1x128_fp8(normed)`` (fp8 from the bf16-rounded normed, so
+    ``out`` matches re-quantizing ``normed_out``).
+
+    Parameters
+    ----------
+    out : torch.Tensor
+        fp8 output, shape ``(batch_size, hidden_size)``, dtype float8_e4m3fn / float8_e5m2.
+    block_scale : torch.Tensor
+        fp32 block scales, shape ``(hidden_size // 128, round_up(batch_size, 4))``, contiguous.
+        This is the column-major (MN-major, TMA-aligned) buffer deep_gemm expects: the logical
+        ``(batch_size, hidden_size // 128)`` scale is ``block_scale.transpose(0, 1)[:batch_size]``.
+    normed_out : torch.Tensor
+        bf16/fp16 normed output, shape ``(batch_size, hidden_size)`` (for consumers that need the
+        pre-quant activation, e.g. an MoE router).
+    input, residual : torch.Tensor
+        Shape ``(batch_size, hidden_size)``. ``residual`` is updated in place with ``input+residual``.
+    weight : torch.Tensor
+        RMSNorm weight, shape ``(hidden_size,)``. ``hidden_size`` must be a multiple of 128
+        (and of 256 for hidden_size<=8192, 512 for 8192<hidden_size<=16384).
+    eps : float
+        Epsilon for numerical stability.
+    enable_pdl : bool
+        Whether to enable programmatic dependent launch.
+    """
+    if enable_pdl is None or enable_pdl:
+        enable_pdl = device_support_pdl(input.device)
+    get_norm_module().fused_add_rmsnorm_fp8_block_quant(
+        out, block_scale, normed_out, input, residual, weight, eps, enable_pdl
+    )
+
+
+@register_fake_op("flashinfer::fused_add_rmsnorm_fp8_block_quant")
+def _fused_add_rmsnorm_fp8_block_quant_fake(
+    out: torch.Tensor,
+    block_scale: torch.Tensor,
+    normed_out: torch.Tensor,
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    enable_pdl: Optional[bool] = None,
+) -> None:
+    pass
+
+
 @flashinfer_api(trace=gemma_rmsnorm_trace)
 def gemma_rmsnorm(
     input: torch.Tensor,