Migrate norm kernels to FlashInfer JIT implementation (#18871)

sgl-kernel/python/sgl_kernel/elementwise.py

@@ -1,9 +1,76 @@
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import Optional
 
 import torch
 from sgl_kernel.utils import is_arch_support_pdl
 
+try:
+    import flashinfer.norm as _flashinfer_norm
+
+    _has_flashinfer = True
+except ImportError:
+    _has_flashinfer = False
+
+_FLASHINFER_NORM_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16}
+
+
+def _rmsnorm_internal(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    out: Optional[torch.Tensor],
+    enable_pdl: Optional[bool],
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(input)
+    if enable_pdl is None:
+        enable_pdl = is_arch_support_pdl()
+    torch.ops.sgl_kernel.rmsnorm.default(out, input, weight, eps, enable_pdl)
+    return out
+
+
+def _fused_add_rmsnorm_internal(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    enable_pdl: Optional[bool],
+) -> None:
+    if enable_pdl is None:
+        enable_pdl = is_arch_support_pdl()
+    torch.ops.sgl_kernel.fused_add_rmsnorm.default(
+        input, residual, weight, eps, enable_pdl
+    )
+
+
+def _gemma_rmsnorm_internal(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    out: Optional[torch.Tensor],
+    enable_pdl: Optional[bool],
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(input)
+    if enable_pdl is None:
+        enable_pdl = is_arch_support_pdl()
+    torch.ops.sgl_kernel.gemma_rmsnorm.default(out, input, weight, eps, enable_pdl)
+    return out
+
+
+def _gemma_fused_add_rmsnorm_internal(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    enable_pdl: Optional[bool],
+) -> None:
+    if enable_pdl is None:
+        enable_pdl = is_arch_support_pdl()
+    torch.ops.sgl_kernel.gemma_fused_add_rmsnorm.default(
+        input, residual, weight, eps, enable_pdl
+    )
+
 
 # These implementations extensively draw from and build upon the FlashInfer project https://github.com/flashinfer-ai/flashinfer
 # Kudos to @yzh119
@@ -38,12 +105,23 @@ def rmsnorm(
     output: torch.Tensor
         Normalized tensor, shape (batch_size, hidden_size).
     """
-    if out is None:
-        out = torch.empty_like(input)
-    if enable_pdl is None:
-        enable_pdl = is_arch_support_pdl()
-    torch.ops.sgl_kernel.rmsnorm.default(out, input, weight, eps, enable_pdl)
-    return out
+    # torch.compiler.is_dynamo_compiling(): FlashInfer norm paths are not safe under
+    # torch.compile(..., fullgraph=True). Dynamo traces into FlashInfer's JIT module
+    # loading path, which calls Path.exists() / os.stat() — both untraceable — causing
+    # the entire compilation to fail. We fall back to the internal implementation while
+    # tracing as a temporary workaround. Once the upstream fix is merged and we upgrade
+    # FlashInfer, this check can be removed.
+    # See: https://github.com/flashinfer-ai/flashinfer/issues/2734
+    #      https://github.com/flashinfer-ai/flashinfer/pull/2733
+    if (
+        input.device.type == "musa"
+        or not _has_flashinfer
+        or input.dtype not in _FLASHINFER_NORM_SUPPORTED_DTYPES
+        or torch.compiler.is_dynamo_compiling()
+    ):
+        return _rmsnorm_internal(input, weight, eps, out, enable_pdl)
+    else:
+        return _flashinfer_norm.rmsnorm(input, weight, eps, out, enable_pdl)
 
 
 def fused_add_rmsnorm(
@@ -76,11 +154,16 @@ def fused_add_rmsnorm(
         <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_
         If None, will be automatically enabled on Hopper architecture.
     """
-    if enable_pdl is None:
-        enable_pdl = is_arch_support_pdl()
-    torch.ops.sgl_kernel.fused_add_rmsnorm.default(
-        input, residual, weight, eps, enable_pdl
-    )
+    # See is_dynamo_compiling() comment in rmsnorm() above.
+    if (
+        input.device.type == "musa"
+        or not _has_flashinfer
+        or input.dtype not in _FLASHINFER_NORM_SUPPORTED_DTYPES
+        or torch.compiler.is_dynamo_compiling()
+    ):
+        _fused_add_rmsnorm_internal(input, residual, weight, eps, enable_pdl)
+    else:
+        _flashinfer_norm.fused_add_rmsnorm(input, residual, weight, eps, enable_pdl)
 
 
 def gemma_rmsnorm(
@@ -114,12 +197,16 @@ def gemma_rmsnorm(
     output: torch.Tensor
         Gemma Normalized tensor, shape (batch_size, hidden_size).
     """
-    if out is None:
-        out = torch.empty_like(input)
-    if enable_pdl is None:
-        enable_pdl = is_arch_support_pdl()
-    torch.ops.sgl_kernel.gemma_rmsnorm.default(out, input, weight, eps, enable_pdl)
-    return out
+    # See is_dynamo_compiling() comment in rmsnorm() above.
+    if (
+        input.device.type == "musa"
+        or not _has_flashinfer
+        or input.dtype not in _FLASHINFER_NORM_SUPPORTED_DTYPES
+        or torch.compiler.is_dynamo_compiling()
+    ):
+        return _gemma_rmsnorm_internal(input, weight, eps, out, enable_pdl)
+    else:
+        return _flashinfer_norm.gemma_rmsnorm(input, weight, eps, out, enable_pdl)
 
 
 def gemma_fused_add_rmsnorm(
@@ -152,11 +239,18 @@ def gemma_fused_add_rmsnorm(
         <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_
         If None, will be automatically enabled on Hopper architecture.
     """
-    if enable_pdl is None:
-        enable_pdl = is_arch_support_pdl()
-    torch.ops.sgl_kernel.gemma_fused_add_rmsnorm.default(
-        input, residual, weight, eps, enable_pdl
-    )
+    # See is_dynamo_compiling() comment in rmsnorm() above.
+    if (
+        input.device.type == "musa"
+        or not _has_flashinfer
+        or input.dtype not in _FLASHINFER_NORM_SUPPORTED_DTYPES
+        or torch.compiler.is_dynamo_compiling()
+    ):
+        _gemma_fused_add_rmsnorm_internal(input, residual, weight, eps, enable_pdl)
+    else:
+        _flashinfer_norm.gemma_fused_add_rmsnorm(
+            input, residual, weight, eps, enable_pdl
+        )
 
 
 def _check_shape(input: torch.Tensor, output: torch.Tensor) -> None: