minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Revert "Fix wrong RMSNorm fallback to old Flashinfer CUDA kernel when in PCG" (#30083)

Browse Source
This commit is contained in:
Cheng Wan
2026-07-03 19:05:20 -07:00
committed by GitHub
parent c21f6f19cf
commit 5f623ad24e
2 changed files with 33 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/layernorm.py
+5 -22
View File
@@ -86,28 +86,11 @@ if _is_cuda or _is_xpu or _is_musa:
else: else:
_flashinfer_layernorm_available = False _flashinfer_layernorm_available = False
from sgl_kernel import fused_add_rmsnorm as _sgl_fused_add_rmsnorm from sgl_kernel import (
from sgl_kernel import gemma_fused_add_rmsnorm as _sgl_gemma_fused_add_rmsnorm fused_add_rmsnorm,
from sgl_kernel import gemma_rmsnorm as _sgl_gemma_rmsnorm gemma_fused_add_rmsnorm,
from sgl_kernel import rmsnorm as _sgl_rmsnorm gemma_rmsnorm,
rmsnorm,
from sglang.srt.utils.custom_op import register_custom_op_from_extern
rmsnorm = register_custom_op_from_extern(
_sgl_rmsnorm, op_name="sgl_rmsnorm", out_shape="input"
)
fused_add_rmsnorm = register_custom_op_from_extern(
_sgl_fused_add_rmsnorm,
op_name="sgl_fused_add_rmsnorm",
mutates_args=["input", "residual"],
)
gemma_rmsnorm = register_custom_op_from_extern(
_sgl_gemma_rmsnorm, op_name="sgl_gemma_rmsnorm", out_shape="input"
)
gemma_fused_add_rmsnorm = register_custom_op_from_extern(
_sgl_gemma_fused_add_rmsnorm,
op_name="sgl_gemma_fused_add_rmsnorm",
mutates_args=["input", "residual"],
) )
_has_aiter_layer_norm = False _has_aiter_layer_norm = False
_has_vllm_rms_norm = False _has_vllm_rms_norm = False
sgl-kernel/python/sgl_kernel/elementwise.py
+28 -4
View File
@@ -104,7 +104,19 @@ def rmsnorm(
output: torch.Tensor output: torch.Tensor
Normalized tensor, shape (batch_size, hidden_size). Normalized tensor, shape (batch_size, hidden_size).
""" """
if _has_flashinfer and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES: # 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 (
_has_flashinfer
and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
and not torch.compiler.is_dynamo_compiling()
):
return _flashinfer_norm.rmsnorm(input, weight, eps, out, enable_pdl) return _flashinfer_norm.rmsnorm(input, weight, eps, out, enable_pdl)
else: else:
return _rmsnorm_internal(input, weight, eps, out, enable_pdl) return _rmsnorm_internal(input, weight, eps, out, enable_pdl)
@@ -140,7 +152,11 @@ def fused_add_rmsnorm(
<https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_ <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 None, will be automatically enabled on Hopper architecture.
""" """
if _has_flashinfer and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES: if (
_has_flashinfer
and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
and not torch.compiler.is_dynamo_compiling()
):
_flashinfer_norm.fused_add_rmsnorm(input, residual, weight, eps, enable_pdl) _flashinfer_norm.fused_add_rmsnorm(input, residual, weight, eps, enable_pdl)
else: else:
_fused_add_rmsnorm_internal(input, residual, weight, eps, enable_pdl) _fused_add_rmsnorm_internal(input, residual, weight, eps, enable_pdl)
@@ -177,7 +193,11 @@ def gemma_rmsnorm(
output: torch.Tensor output: torch.Tensor
Gemma Normalized tensor, shape (batch_size, hidden_size). Gemma Normalized tensor, shape (batch_size, hidden_size).
""" """
if _has_flashinfer and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES: if (
_has_flashinfer
and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
and not torch.compiler.is_dynamo_compiling()
):
return _flashinfer_norm.gemma_rmsnorm(input, weight, eps, out, enable_pdl) return _flashinfer_norm.gemma_rmsnorm(input, weight, eps, out, enable_pdl)
else: else:
return _gemma_rmsnorm_internal(input, weight, eps, out, enable_pdl) return _gemma_rmsnorm_internal(input, weight, eps, out, enable_pdl)
@@ -213,7 +233,11 @@ def gemma_fused_add_rmsnorm(
<https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_ <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 None, will be automatically enabled on Hopper architecture.
""" """
if _has_flashinfer and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES: if (
_has_flashinfer
and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
and not torch.compiler.is_dynamo_compiling()
):
_flashinfer_norm.gemma_fused_add_rmsnorm( _flashinfer_norm.gemma_fused_add_rmsnorm(
input, residual, weight, eps, enable_pdl input, residual, weight, eps, enable_pdl
) )