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
sgl-kernel/python/sgl_kernel/elementwise.py
+4 -28
View File
@@ -104,19 +104,7 @@ def rmsnorm(
output: torch.Tensor
Normalized tensor, shape (batch_size, hidden_size).
"""
# 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()
):
if _has_flashinfer and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES:
return _flashinfer_norm.rmsnorm(input, weight, eps, out, enable_pdl)
else:
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>`_
If None, will be automatically enabled on Hopper architecture.
"""
if (
_has_flashinfer
and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
and not torch.compiler.is_dynamo_compiling()
):
if _has_flashinfer and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES:
_flashinfer_norm.fused_add_rmsnorm(input, residual, weight, eps, enable_pdl)
else:
_fused_add_rmsnorm_internal(input, residual, weight, eps, enable_pdl)
@@ -193,11 +177,7 @@ def gemma_rmsnorm(
output: torch.Tensor
Gemma Normalized tensor, shape (batch_size, hidden_size).
"""
if (
_has_flashinfer
and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
and not torch.compiler.is_dynamo_compiling()
):
if _has_flashinfer and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES:
return _flashinfer_norm.gemma_rmsnorm(input, weight, eps, out, enable_pdl)
else:
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>`_
If None, will be automatically enabled on Hopper architecture.
"""
if (
_has_flashinfer
and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
and not torch.compiler.is_dynamo_compiling()
):
if _has_flashinfer and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES:
_flashinfer_norm.gemma_fused_add_rmsnorm(
input, residual, weight, eps, enable_pdl
)