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

Clean up bench_one_batch warning and simplify norm dispatch (#23110)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-04-17 17:42:20 -07:00
committed by GitHub
parent 5f7aee726a
commit 9c47bbad13
2 changed files with 24 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/test/bench_one_batch_server_internal.py
+4 -1
View File
@@ -35,7 +35,10 @@ def get_cache_tokens_from_metrics(url: str) -> Optional[tuple]:
"""
try:
response = requests.get(url + "/metrics", timeout=5)
response.raise_for_status()
try:
response.raise_for_status()
except requests.exceptions.HTTPError:
return None
# Parse Prometheus text format
# Looking for: sglang:cached_tokens_total{...} <value>
sgl-kernel/python/sgl_kernel/elementwise.py
+20 -27
View File
@@ -113,14 +113,13 @@ def rmsnorm(
# 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()
_has_flashinfer
and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
and not 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)
else:
return _rmsnorm_internal(input, weight, eps, out, enable_pdl)
def fused_add_rmsnorm(
@@ -153,16 +152,14 @@ 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.
"""
# 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()
_has_flashinfer
and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
and not 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)
else:
_fused_add_rmsnorm_internal(input, residual, weight, eps, enable_pdl)
def gemma_rmsnorm(
@@ -196,16 +193,14 @@ def gemma_rmsnorm(
output: torch.Tensor
Gemma Normalized tensor, shape (batch_size, hidden_size).
"""
# 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()
_has_flashinfer
and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
and not 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)
else:
return _gemma_rmsnorm_internal(input, weight, eps, out, enable_pdl)
def gemma_fused_add_rmsnorm(
@@ -238,18 +233,16 @@ 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.
"""
# 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()
_has_flashinfer
and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
and not 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
)
else:
_gemma_fused_add_rmsnorm_internal(input, residual, weight, eps, enable_pdl)
def _check_shape(input: torch.Tensor, output: torch.Tensor) -> None: