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

Fix wrong RMSNorm fallback to old Flashinfer CUDA kernel when in PCG (#29702)

Browse Source 
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
This commit is contained in:
Brayden Zhong
2026-07-03 00:49:49 +08:00
committed by GitHub
co-authored by Brayden Zhong
parent 9588cacaa1
commit 1b6d1e9752
2 changed files with 26 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/layernorm.py
+22 -5
View File
@@ -86,11 +86,28 @@ if _is_cuda or _is_xpu or _is_musa:
else: else:
_flashinfer_layernorm_available = False _flashinfer_layernorm_available = False
from sgl_kernel import ( from sgl_kernel import fused_add_rmsnorm as _sgl_fused_add_rmsnorm
fused_add_rmsnorm, from sgl_kernel import gemma_fused_add_rmsnorm as _sgl_gemma_fused_add_rmsnorm
gemma_fused_add_rmsnorm, from sgl_kernel import gemma_rmsnorm as _sgl_gemma_rmsnorm
gemma_rmsnorm, from sgl_kernel import rmsnorm as _sgl_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
+4 -28
View File
@@ -104,19 +104,7 @@ def rmsnorm(
output: torch.Tensor output: torch.Tensor
Normalized tensor, shape (batch_size, hidden_size). Normalized tensor, shape (batch_size, hidden_size).
""" """
# torch.compiler.is_dynamo_compiling(): FlashInfer norm paths are not safe under if _has_flashinfer and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES:
# 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)
@@ -152,11 +140,7 @@ 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 ( if _has_flashinfer and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES:
_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)
@@ -193,11 +177,7 @@ 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 ( if _has_flashinfer and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES:
_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)
@@ -233,11 +213,7 @@ 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 ( if _has_flashinfer and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES:
_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
) )