feat: fused FP8 quantization output for LayerNorm

benchmarks/routines/flashinfer_benchmark_utils.py

@@ -234,6 +234,7 @@
         "gemma_fused_add_rmsnorm",
         "rmsnorm_quant",
         "fused_add_rmsnorm_quant",
+        "layernorm_quant",
         "rmsnorm_fp4quant",
         "add_rmsnorm_fp4quant",
         "fused_rmsnorm_silu",
@@ -615,6 +616,17 @@ def dtype_str_to_torch_dtype(dtype_str):
         "12.0": ["cute-dsl"],
         "12.1": ["cute-dsl"],
     },
+    "layernorm_quant": {
+        "7.5": ["cuda"],
+        "8.0": ["cuda"],
+        "8.6": ["cuda"],
+        "8.9": ["cuda"],
+        "9.0": ["cuda"],
+        "10.0": ["cuda"],
+        "10.3": ["cuda"],
+        "12.0": ["cuda"],
+        "12.1": ["cuda"],
+    },
     # NORM - FP4 Quantization (Blackwell SM100+ only, CuTe-DSL kernels)
     "rmsnorm_fp4quant": {
         "7.5": [],

benchmarks/routines/norm.py

@@ -53,6 +53,8 @@ def run_norm_test(args):
         return testRmsnormQuant(args)
     elif args.routine == "fused_add_rmsnorm_quant":
         return testFusedAddRmsnormQuant(args)
+    elif args.routine == "layernorm_quant":
+        return testLayernormQuant(args)
     elif args.routine == "rmsnorm_fp4quant":
         return testRmsnormFp4quant(args)
     elif args.routine == "add_rmsnorm_fp4quant":
@@ -1015,6 +1017,196 @@ def run_backend(backend, out_tensor, input_tensor, weight):
     return res
 
 
+def testLayernormQuant(args):
+    """
+    Test layernorm_quant API.
+
+    This test:
+    1. Generates random input tensors
+    2. Runs layernorm_quant with quantized output
+    3. Runs reference check
+    4. Measures performance metrics (memory bandwidth)
+
+    Note: LayerNorm is memory-bandwidth bound, so TB/sec is the primary metric.
+
+    Args:
+        args: Parsed command line arguments containing test configuration
+
+    Returns:
+        dict: List of dictionaries containing performance results
+    """
+    if args.verbose >= 1:
+        print("[INFO] Running testLayernormQuant")
+        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")
+
+    device = get_device(args)
+    if args.generate_repro_command:
+        print(
+            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
+        )
+
+    ## Parse input arguments
+    backends = args.backends[:]  # Make a copy to avoid modifying the original
+
+    batch_size = args.batch_size
+    hidden_size = args.hidden_size
+    # Pre-allocate scale on device once. Passing a Python float forces flashinfer's
+    # _normalize_scale_tensor to create a tensor on each call, which triggers a
+    # host->device copy that is illegal under CUDA graph capture (and emits a
+    # FutureWarning about the deprecated float form).
+    scale = torch.tensor([args.scale], dtype=torch.float32, device=device)
+    eps = args.eps
+    is_cuda_graph_compatible = not args.no_cuda_graph
+    run_refcheck = args.refcheck
+    res = []
+
+    backends = filter_backends_by_compute_capability(backends, args.routine, device)
+    if len(backends) == 0:
+        print("[ERROR] No backends to test. Exiting.")
+        return res
+
+    input_dtype = dtype_str_to_torch_dtype(args.input_dtype)
+    if input_dtype != torch.bfloat16:
+        raise ValueError(
+            f"Unsupported input dtype: {args.input_dtype}. layernorm_quant only supports bfloat16."
+        )
+
+    out_dtype = dtype_str_to_torch_dtype(args.out_dtype)
+    if out_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2]:
+        raise ValueError(
+            f"Unsupported out dtype: {args.out_dtype}. Supported dtypes are fp8_e4m3, fp8_e5m2."
+        )
+    ## Done parsing input arguments
+
+    ## Prepare input tensors (2D only for layernorm_quant)
+    input_shape = (batch_size, hidden_size)
+
+    input_tensor = torch.randn(input_shape, dtype=input_dtype, device=device)
+    gamma = torch.randn(hidden_size, dtype=torch.float32, device=device)
+    beta = torch.randn(hidden_size, dtype=torch.float32, device=device)
+    out_tensor = torch.empty(input_shape, dtype=out_dtype, device=device)
+
+    if args.verbose >= 2:
+        print(f"[VVERBOSE] {input_tensor.shape = }")
+        print(f"[VVERBOSE] {input_tensor.dtype = }")
+        print(f"[VVERBOSE] {gamma.shape = }")
+        print(f"[VVERBOSE] {out_tensor.dtype = }")
+        print(f"[VVERBOSE] {scale = }")
+
+    def run_backend(backend, out_tensor, input_tensor, gamma, beta):
+        if backend == "cuda":
+            flashinfer.norm.layernorm_quant(
+                out_tensor,
+                input_tensor,
+                gamma,
+                beta,
+                scale=scale,
+                eps=eps,
+            )
+            return out_tensor
+        else:
+            raise ValueError(f"Unsupported backend: {backend}")
+
+    # Reference: PyTorch implementation of LayerNorm + quantization
+    has_reference_output = False
+    if run_refcheck:
+        layernorm_output = torch.nn.functional.layer_norm(
+            input_tensor.float(), (hidden_size,), weight=gamma, bias=beta, eps=eps
+        )
+        # The kernel rounds the normalized value to bf16 before scaling and
+        # casting to fp8; mirror that so 1-ulp fp8 differences do not trip the
+        # tight refcheck tolerance.
+        layernorm_output = layernorm_output.to(input_dtype).float()
+        reference_output = (
+            (layernorm_output / scale)
+            .to(input_dtype)
+            .float()
+            .clamp(torch.finfo(out_dtype).min, torch.finfo(out_dtype).max)
+            .to(out_dtype)
+        )
+        has_reference_output = True
+
+    # Storage for timing results and outputs
+    backend_times = {backend: [] for backend in backends}
+    outputs = {}
+    for cur_backend in backends:
+        # Create fresh output tensor for each run
+        cur_out = torch.empty(input_shape, dtype=out_dtype, device=device)
+        if run_refcheck:
+            outputs[cur_backend] = (
+                run_backend(cur_backend, cur_out, input_tensor, gamma, beta)
+                .detach()
+                .clone()
+            )
+        backend_times[cur_backend] = bench_gpu_time(
+            fn=run_backend,
+            dry_run_iters=args.dry_run_iters,
+            repeat_iters=args.num_iters,
+            enable_cupti=args.use_cupti,
+            use_cuda_graph=is_cuda_graph_compatible,
+            input_args=(cur_backend, out_tensor, input_tensor, gamma, beta),
+        )
+
+    tested_backends = list(outputs.keys())
+    tested_outputs = list(outputs.values())
+    if len(tested_backends) > 0:
+        if run_refcheck and has_reference_output:
+            for i in range(len(tested_backends)):
+                # Compare in float for FP8 outputs
+                ref_float = reference_output.float()
+                out_float = tested_outputs[i].float()
+                (
+                    num_different_elements,
+                    num_elements,
+                    num_different_elements_percentage,
+                ) = is_close_stats(ref_float, out_float, rtol=1e-1, atol=1e-1)
+                if num_different_elements > 0:
+                    print(
+                        f"[ERROR] Output tensor mismatch from backend {tested_backends[i]}: "
+                        f"{num_different_elements}/{num_elements} ({num_different_elements_percentage:.2f}%) elements differ"
+                    )
+                    if not args.allow_output_mismatch:
+                        raise AssertionError(
+                            f"[ERROR] Backend {tested_backends[i]} output mismatch with {num_different_elements} elements"
+                        )
+
+    for backend in backends:
+        if len(backend_times[backend]) > 0:
+            median_time = np.median(backend_times[backend])
+            std_time = np.std(backend_times[backend])
+
+            # Memory bandwidth calculation for LayerNorm + Quant
+            # Read: input tensor + gamma/beta tensors
+            # Write: output tensor (quantized, smaller dtype)
+            num_elements = np.prod(input_shape)
+            problem_bytes = (
+                num_elements * input_dtype.itemsize  # input read
+                + 2 * hidden_size * torch.float32.itemsize  # gamma and beta read
+                + num_elements * out_dtype.itemsize  # output write (quantized)
+            )
+            problem_flops = num_elements * 7  # rough estimate
+            tflops = problem_flops / (10**9 * median_time)  # in TFLOPs/sec
+            tb_per_sec = problem_bytes / (10**9 * median_time)  # in TB/sec
+
+            print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec)
+
+            if args.output_path is not None:
+                cur_res = defaultdict(str)
+                cur_res["routine"] = args.routine
+                cur_res["median_time"] = median_time
+                cur_res["std_time"] = std_time
+                cur_res["tflops"] = tflops
+                cur_res["tb_per_sec"] = tb_per_sec
+                cur_res["input_dtype"] = str(input_dtype)
+                cur_res["out_dtype"] = str(out_dtype)
+                cur_res["scale"] = scale.item()
+                cur_res["eps"] = eps
+                cur_res["backend"] = backend
+                cur_res["case_tag"] = args.case_tag
+                res.append(cur_res)
+    return res
+
+
 def testFusedAddRmsnormQuant(args):
     """
     Test fused_add_rmsnorm_quant API.

csrc/flashinfer_norm_binding.cu

@@ -34,6 +34,9 @@ void gemma_fused_add_rmsnorm(TensorView input, TensorView residual, TensorView w
 
 void layernorm(Tensor out, Tensor input, Tensor gamma, Tensor beta, double eps);
 
+void layernorm_quant(TensorView out, TensorView input, TensorView gamma, TensorView beta,
+                     TensorView scale, double eps);
+
 void fused_dit_layernorm_run(TensorView input, TensorView residual, TensorView gate,
                              TensorView gate_bias, TensorView gamma, TensorView beta,
                              TensorView scale, TensorView scale_bias, TensorView shift,
@@ -58,5 +61,6 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(fused_add_rmsnorm_quant, fused_add_rmsnorm_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);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(layernorm_quant, layernorm_quant);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(fused_dit_layernorm, fused_dit_layernorm_run);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(fused_qk_rmsnorm_rope, fused_qk_rmsnorm_rope_run);

csrc/norm.cu

@@ -274,6 +274,54 @@ void layernorm(Tensor output, Tensor input, Tensor gamma, Tensor beta, double ep
   });
 }
 
+void layernorm_quant(TensorView output, TensorView input, TensorView gamma, TensorView beta,
+                     TensorView scale, double eps) {
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(output);
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(gamma);
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(beta);
+  CHECK_DEVICE(input, gamma);
+  CHECK_DEVICE(input, beta);
+  CHECK_DEVICE(input, output);
+  CHECK_DEVICE(input, scale);
+  CHECK_DIM(2, input);   // input: (batch_size, hidden_size)
+  CHECK_DIM(2, output);  // output: (batch_size, hidden_size)
+  CHECK_DIM(1, gamma);   // gamma: (hidden_size)
+  CHECK_DIM(1, beta);    // beta: (hidden_size)
+  TVM_FFI_ICHECK_EQ(input.size(1), gamma.size(0));
+  TVM_FFI_ICHECK_EQ(input.size(1), beta.size(0));
+  TVM_FFI_ICHECK_EQ(scale.numel(), 1);
+  unsigned int batch_size = input.size(0);
+  unsigned int hidden_size = input.size(1);
+  TVM_FFI_ICHECK_EQ(output.size(0), batch_size);
+  TVM_FFI_ICHECK_EQ(output.size(1), hidden_size);
+  // generalLayerNorm indexes as bidx * hidden_size + i and takes no strides
+  TVM_FFI_ICHECK_EQ(input.stride(0), hidden_size) << "input must be contiguous";
+  TVM_FFI_ICHECK_EQ(output.stride(0), hidden_size) << "output must be contiguous";
+  // Match the layernorm binding: bf16 input with fp32 gamma/beta only.
+  TVM_FFI_ICHECK_EQ(input.dtype(), dl_bfloat16) << "input must be bfloat16";
+  TVM_FFI_ICHECK_EQ(gamma.dtype(), dl_float32) << "gamma must be float32";
+  TVM_FFI_ICHECK_EQ(beta.dtype(), dl_float32) << "beta must be float32";
+  ffi::CUDADeviceGuard device_guard(input.device().device_id);
+  const cudaStream_t stream = get_stream(input.device());
+
+#ifdef ENABLE_FP8
+  DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(input.dtype(), c_type, [&] {
+    return DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP8(output.dtype(), o_type, [&] {
+      cudaError_t status = norm::LayerNormQuant(
+          static_cast<c_type*>(input.data_ptr()), static_cast<float*>(gamma.data_ptr()),
+          static_cast<float*>(beta.data_ptr()), static_cast<o_type*>(output.data_ptr()),
+          static_cast<float*>(scale.data_ptr()), batch_size, hidden_size, eps, stream);
+      TVM_FFI_ICHECK(status == cudaSuccess)
+          << "LayerNormQuant failed with error code " << cudaGetErrorString(status);
+      return true;
+    });
+  });
+#else
+  TVM_FFI_THROW(RuntimeError) << "layernorm_quant requires the norm module built with ENABLE_FP8";
+#endif
+}
+
 void fused_dit_layernorm_run(TensorView input, TensorView residual, TensorView gate,
                              TensorView gate_bias, TensorView gamma, TensorView beta,
                              TensorView scale, TensorView scale_bias, TensorView shift,

docs/api/norm.rst

@@ -17,6 +17,7 @@ Kernels for normalization layers.
     gemma_rmsnorm
     gemma_fused_add_rmsnorm
     layernorm
+    layernorm_quant
     fused_rmsnorm_silu
     fused_qk_rmsnorm_rope
     fused_dit_residual_layernorm_scale_shift

flashinfer/__init__.py

@@ -125,6 +125,7 @@
 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 layernorm as layernorm
+from .norm import layernorm_quant as layernorm_quant
 from .norm import gemma_fused_add_rmsnorm as gemma_fused_add_rmsnorm
 from .norm import gemma_rmsnorm as gemma_rmsnorm
 from .norm import rmsnorm as rmsnorm

flashinfer/norm/__init__.py

@@ -38,6 +38,7 @@
     fused_rmsnorm_silu_trace,
     gemma_fused_add_rmsnorm_trace,
     gemma_rmsnorm_trace,
+    layernorm_quant_trace,
     layernorm_trace,
     rmsnorm_quant_trace,
     rmsnorm_trace,
@@ -530,6 +531,55 @@ def _layernorm_fake(
     return input.new_empty([b, k])
 
 
+@flashinfer_api(trace=layernorm_quant_trace)
+@register_custom_op("flashinfer::layernorm_quant", mutates_args=("out",))
+def layernorm_quant(
+    out: torch.Tensor,
+    input: torch.Tensor,
+    gemma: torch.Tensor,
+    beta: torch.Tensor,
+    scale: Union[float, torch.Tensor],
+    eps: float = 1e-6,
+) -> None:
+    r"""Layer normalization + fp8 quantization.
+
+    ``out[i] = (((input[i] - E[input]) / sqrt(Var[input] + eps)) * gemma[i] + beta[i]) / scale``
+
+    Parameters
+    ----------
+    out: torch.Tensor
+        The output tensor, shape (batch_size, hidden_size). Need to be contiguous.
+        The output is quantized to the dtype of this tensor, which must be
+        float8_e4m3fn or float8_e5m2.
+    input: torch.Tensor
+        Input tensor, shape (batch_size, hidden_size). Need to be bfloat16 and contiguous.
+    gemma: torch.Tensor
+        Gemma tensor, shape (hidden_size,). Need to be float32.
+    beta: torch.Tensor
+        Beta tensor, shape (hidden_size,). Need to be float32.
+    scale: torch.Tensor
+        Scale factor for quantization, shape (1,). The normalized output is
+        divided by this scale before the fp8 cast.
+    eps: float
+        Epsilon for numerical stability.
+    """
+    scale = _normalize_scale_tensor(scale, input)
+    # No CuTe-DSL layernorm quant kernel yet; always use the CUDA JIT module.
+    get_norm_module().layernorm_quant(out, input, gemma, beta, scale, eps)
+
+
+@register_fake_op("flashinfer::layernorm_quant")
+def _layernorm_quant_fake(
+    out: torch.Tensor,
+    input: torch.Tensor,
+    gemma: torch.Tensor,
+    beta: torch.Tensor,
+    scale: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    pass
+
+
 # CuTe-DSL fused RMSNorm + FP4 Quantization kernels
 # These require SM100+ (Blackwell) GPUs and nvidia-cutlass-dsl
 try: