Add layernorm op for inputs of mixed dtype (#1926)

<!-- .github/pull_request_template.md -->

## 📌 Description

<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

## Reviewer Notes

Continued from #1914
with small review fixes included, as the original author will be out for
the next week.

---------

Co-authored-by: kaixih <[email protected]>

Author: akhilg-nv

csrc/flashinfer_norm_binding.cu

@@ -26,7 +26,10 @@ void gemma_rmsnorm(TensorView out, TensorView input, TensorView weight, double e
 void gemma_fused_add_rmsnorm(TensorView input, TensorView residual, TensorView weight, double eps,
                              bool enable_pdl);
 
+void layernorm(Tensor out, Tensor input, Tensor gamma, Tensor beta, double eps);
+
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(rmsnorm, rmsnorm);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(fused_add_rmsnorm, fused_add_rmsnorm);
 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

@@ -160,3 +160,36 @@ void gemma_fused_add_rmsnorm(TensorView input, TensorView residual, TensorView w
     return true;
   });
 }
+
+void layernorm(Tensor output, Tensor input, Tensor gamma, Tensor beta, double eps) {
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(gamma);
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(beta);
+  CHECK_DEVICE(input, gamma);
+  CHECK_DEVICE(input, beta);
+  CHECK_DIM(2, input);  // input: (batch_size, hidden_size)
+  CHECK_DIM(1, gamma);  // gamma: (hidden_size)
+  CHECK_DIM(1, beta);   // beta: (hidden_size)
+  TVM_FFI_ICHECK_EQ(input->shape[1], gamma->shape[0]);
+  TVM_FFI_ICHECK_EQ(input->shape[1], beta->shape[0]);
+  unsigned int batch_size = input->shape[0];
+  unsigned int hidden_size = input->shape[1];
+  TVM_FFI_ICHECK_EQ(output->shape[0], batch_size);
+  TVM_FFI_ICHECK_EQ(output->shape[1], hidden_size);
+  cudaSetDevice(input->device.device_id);
+  const cudaStream_t stream = get_stream(input->device);
+  // TODO(kaixih): This is currently our only use case; Add more if needed.
+  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";
+
+  DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(input->dtype, c_type, [&] {
+    cudaError_t status =
+        norm::LayerNorm(static_cast<c_type*>(input->data), static_cast<float*>(gamma->data),
+                        static_cast<float*>(beta->data), static_cast<c_type*>(output->data),
+                        batch_size, hidden_size, eps, stream);
+    TVM_FFI_ICHECK(status == cudaSuccess)
+        << "LayerNorm failed with error code " << cudaGetErrorString(status);
+    return true;
+  });
+}

docs/api/norm.rst

@@ -14,3 +14,4 @@ Kernels for normalization layers.
     fused_add_rmsnorm
     gemma_rmsnorm
     gemma_fused_add_rmsnorm
+    layernorm

flashinfer/__init__.py

@@ -90,6 +90,7 @@
 from .gemm import tgv_gemm_sm100 as tgv_gemm_sm100
 from .mla import BatchMLAPagedAttentionWrapper as BatchMLAPagedAttentionWrapper
 from .norm import fused_add_rmsnorm as fused_add_rmsnorm
+from .norm import layernorm as layernorm
 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/jit/norm.py

@@ -19,10 +19,15 @@
 
 
 def gen_norm_module() -> JitSpec:
+    nvcc_flags = [
+        "-DENABLE_BF16",
+        "-DENABLE_FP8",
+    ]
     return gen_jit_spec(
         "norm",
         [
             jit_env.FLASHINFER_CSRC_DIR / "norm.cu",
             jit_env.FLASHINFER_CSRC_DIR / "flashinfer_norm_binding.cu",
         ],
+        extra_cuda_cflags=nvcc_flags,
     )

flashinfer/norm.py

@@ -244,3 +244,43 @@ def _gemma_fused_add_rmsnorm_fake(
     enable_pdl: Optional[bool] = None,
 ) -> None:
     pass
+
+
+@register_custom_op("flashinfer::layernorm", mutates_args=())
+def layernorm(
+    input: torch.Tensor,
+    gemma: torch.Tensor,
+    beta: torch.Tensor,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    r"""Layer normalization.
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (batch_size, hidden_size). Need to be bfloat16.
+    gemma: torch.Tensor
+        Gemma tensor, shape (hidden_size,). Need to be float32.
+    beta: torch.Tensor
+        Beta tensor, shape (hidden_size,). Need to be float32.
+    eps: float
+        Epsilon for numerical stability.
+
+    Returns
+    -------
+    output: torch.Tensor
+        Layer Normalized tensor, shape (batch_size, hidden_size). Same dtype as input.
+    """
+    out = torch.empty_like(input)
+    get_norm_module().layernorm(out, input, gemma, beta, eps)
+    return out
+
+
+@register_fake_op("flashinfer::layernorm")
+def _layernorm_fake(
+    input: torch.Tensor,
+    gemma: torch.Tensor,
+    beta: torch.Tensor,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    b, k = input.shape
+    return input.new_empty([b, k])