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

Code clean up for fp8 quantization (#16982)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-01-13 12:38:39 -08:00
committed by GitHub
parent a0b4ba9032
commit 075c5a5789
9 changed files with 309 additions and 237 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/srt/constrained/xgrammar_backend.py
View File

@@ -264,8 +264,8 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
schema=key_string, any_whitespace=self.any_whitespace
)
except (RuntimeError, json.decoder.JSONDecodeError) as e:
logging.error(f"Hit invalid json_schema: {key_string=}, {e=}")
except (RuntimeError, json.decoder.JSONDecodeError, UnicodeDecodeError) as e:
logger.error(f"Hit invalid json_schema: {key_string=}, {e=}")
return INVALID_GRAMMAR_OBJ
return self._from_context(ctx, key_string, GrammarStats(dispatch_type="json"))
@@ -273,7 +273,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
try:
ctx = self.grammar_compiler.compile_grammar(key_string)
except RuntimeError as e:
logging.error(f"Hit invalid ebnf: {key_string=}, {e=}")
logger.error(f"Hit invalid ebnf: {key_string=}, {e=}")
return INVALID_GRAMMAR_OBJ
return self._from_context(ctx, key_string, GrammarStats(dispatch_type="ebnf"))
@@ -281,7 +281,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
try:
ctx = self.grammar_compiler.compile_regex(key_string)
except RuntimeError as e:
logging.error(f"Hit invalid regex: {key_string=}, {e=}")
logger.error(f"Hit invalid regex: {key_string=}, {e=}")
return INVALID_GRAMMAR_OBJ
return self._from_context(ctx, key_string, GrammarStats(dispatch_type="regex"))
@@ -310,7 +310,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
key_string = json.dumps(structural_tag)
ctx = self.grammar_compiler.compile_structural_tag(key_string)
except (RuntimeError, json.decoder.JSONDecodeError) as e:
logging.error(f"Hit invalid structural_tag: {key_string=}, {e=}")
logger.error(f"Hit invalid structural_tag: {key_string=}, {e=}")
return INVALID_GRAMMAR_OBJ
return self._from_context(
ctx, key_string, GrammarStats(dispatch_type="structural_tag")

5
python/sglang/srt/layers/linear.py
View File

@@ -362,6 +362,7 @@ class ColumnParallelLinear(LinearBase):
def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
output_dim = getattr(param, "output_dim", None)
param_data = param.data
# Special case for GGUF
is_gguf_weight = getattr(param, "is_gguf_weight", False)
@@ -373,11 +374,9 @@ class ColumnParallelLinear(LinearBase):
if is_gguf_weight and isinstance(param, UninitializedParameter):
param.materialize(loaded_weight.shape, dtype=loaded_weight.dtype)
use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
param_data = param.data
# bitsandbytes loads the weights of the specific portion
# no need to narrow here
use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
if output_dim is not None and not use_bitsandbytes_4bit:
shard_size = param_data.shape[output_dim]
start_idx = self.tp_rank * shard_size

49
python/sglang/srt/layers/parameter.py
View File

@@ -27,6 +27,52 @@ logger = logging.getLogger(__name__)
_is_cpu = is_cpu()
def _dtype_rank(dtype: torch.dtype) -> Optional[int]:
if dtype in (
torch.float8_e4m3fn,
torch.float8_e4m3fnuz,
torch.float8_e5m2,
torch.float8_e5m2fnuz,
):
return 0
if dtype in (torch.float16, torch.bfloat16):
return 1
if dtype == torch.float32:
return 2
if dtype == torch.float64:
return 3
return None
def copy_with_check(target: torch.Tensor, loaded_weight: torch.Tensor):
"""
Copy `loaded_weight` into `target` while forbidding downcasts.
bf16/fp16 share the same rank, and all fp8 variants share the same rank.
"""
assert (
target.shape == loaded_weight.shape
), f"{target.shape=}, {loaded_weight.shape=}"
if target.dtype == loaded_weight.dtype:
target.copy_(loaded_weight)
return
target_rank = _dtype_rank(target.dtype)
loaded_rank = _dtype_rank(loaded_weight.dtype)
if target_rank is None or loaded_rank is None:
raise ValueError(
f"Unsupported copy between dtypes: {target.dtype=}, {loaded_weight.dtype=}"
)
if target_rank < loaded_rank:
raise ValueError(
f"Downcasting not allowed: {target.dtype=}, {loaded_weight.dtype=}"
)
target.copy_(loaded_weight)
class BasevLLMParameter(Parameter):
"""
Base parameter for vLLM linear layers. Extends the torch.nn.parameter
@@ -120,8 +166,7 @@ class _ColumnvLLMParameter(BasevLLMParameter):
self.output_dim, tp_rank * shard_size, shard_size
)
assert self.data.shape == loaded_weight.shape
self.data.copy_(loaded_weight)
copy_with_check(self.data, loaded_weight)
def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs):

376
python/sglang/srt/layers/quantization/fp8.py
View File

@@ -159,7 +159,7 @@ class Fp8Config(QuantizationConfig):
config, ["ignored_layers", "modules_to_not_convert"], None
)
if ignored_layers:
# hacking ministral
# hack for ministral
ignored_layers = [layer.replace("model.", "") for layer in ignored_layers]
weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None)
return cls(
@@ -192,17 +192,19 @@ class Fp8Config(QuantizationConfig):
class Fp8LinearMethod(LinearMethodBase):
"""Linear method for FP8.
Supports loading FP8 checkpoints with static weight scale and
dynamic/static activation scale.
Also supports loading quantized FP16/BF16 model checkpoints with dynamic
activation scaling. The weight scaling factor will be initialized after
the model weights are loaded.
It supports the following quantization schemes:
- Per-channel weight quantization + per-token activation quantization
- Per-tensor weight quantization + per-tensor activation quantization
- Blockwise weight quantization + blockwise activation quantization
Limitations:
1. Only support per-tensor quantization due to torch._scaled_mm support.
2. Only support float8_e4m3fn data type due to the limitation of
torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856)
It supports the following checkpoint formats:
- FP8 checkpoint
- FP16/BF16 checkpoint. In this case, the weights will be quantized to FP8 during the weight loading.
Notes:
- The activation quantization scheme can be static or dynamic. The dynamic activation quantization is more commonly used.
- On NV platforms, the per-channel weight quantization is used by default, if block quantization is not enabled.
Args:
quant_config: The quantization config.
@@ -221,8 +223,50 @@ class Fp8LinearMethod(LinearMethodBase):
self.use_marlin = force_marlin or auto_enable
self.block_quant = self.quant_config.weight_block_size is not None
self.w8a8_block_fp8_linear = dispatch_w8a8_block_fp8_linear()
self.is_checkpoint_fp8_serialized = (
self.quant_config.is_checkpoint_fp8_serialized
)
def validate_block_quant_shapes(
self,
input_size: int,
input_size_per_partition: int,
output_size: int,
output_size_per_partition: int,
output_partition_sizes: List[int],
skip_block_quant_check: bool = False,
):
tp_size = get_tensor_model_parallel_world_size()
block_n, block_k = (
self.quant_config.weight_block_size[0],
self.quant_config.weight_block_size[1],
)
if skip_block_quant_check:
print_warning_once(
"Skipping block quantization checks for weight partition."
)
else:
# Required by row parallel
if tp_size > 1 and input_size // input_size_per_partition == tp_size:
if input_size_per_partition % block_k != 0:
raise ValueError(
f"Weight input_size_per_partition = "
f"{input_size_per_partition} is not divisible by "
f"weight quantization block_k = {block_k}."
)
# Required by column parallel or enabling merged weights
if (
tp_size > 1 and output_size // output_size_per_partition == tp_size
) or len(output_partition_sizes) > 1:
for output_partition_size in output_partition_sizes:
if output_partition_size % block_n != 0:
raise ValueError(
f"Weight output_partition_size = "
f"{output_partition_size} is not divisible by "
f"weight quantization block_n = {block_n}."
)
def create_weights(
self,
@@ -235,53 +279,29 @@ class Fp8LinearMethod(LinearMethodBase):
skip_block_quant_check: bool = False,
**extra_weight_attrs,
):
# Copy the layer attributes
output_size_per_partition = sum(output_partition_sizes)
weight_loader = extra_weight_attrs.get("weight_loader")
tp_size = get_tensor_model_parallel_world_size()
if self.block_quant:
block_n, block_k = (
self.quant_config.weight_block_size[0],
self.quant_config.weight_block_size[1],
)
if skip_block_quant_check:
logger.warning_once(
f"Skipping block quantization checks for weight partition."
)
else:
# Required by row parallel
if tp_size > 1 and input_size // input_size_per_partition == tp_size:
if input_size_per_partition % block_k != 0:
raise ValueError(
f"Weight input_size_per_partition = "
f"{input_size_per_partition} is not divisible by "
f"weight quantization block_k = {block_k}."
)
# Required by column parallel or enabling merged weights
if (
tp_size > 1 and output_size // output_size_per_partition == tp_size
) or len(output_partition_sizes) > 1:
for output_partition_size in output_partition_sizes:
if output_partition_size % block_n != 0:
raise ValueError(
f"Weight output_partition_size = "
f"{output_partition_size} is not divisible by "
f"weight quantization block_n = {block_n}."
)
layer.logical_widths = output_partition_sizes
layer.input_size_per_partition = input_size_per_partition
layer.output_size_per_partition = output_size_per_partition
layer.orig_dtype = params_dtype
weight_loader = extra_weight_attrs.get("weight_loader")
# WEIGHT
if self.block_quant:
block_n, block_k = self.quant_config.weight_block_size
self.validate_block_quant_shapes(
input_size,
input_size_per_partition,
output_size,
output_size_per_partition,
output_partition_sizes,
skip_block_quant_check,
)
# Create the weight
weight_dtype = (
torch.float8_e4m3fn
if self.quant_config.is_checkpoint_fp8_serialized
else params_dtype
torch.float8_e4m3fn if self.is_checkpoint_fp8_serialized else params_dtype
)
weight = ModelWeightParameter(
data=torch.empty(
output_size_per_partition, input_size_per_partition, dtype=weight_dtype
@@ -294,7 +314,7 @@ class Fp8LinearMethod(LinearMethodBase):
# If checkpoint is serialized fp8, load them.
# Otherwise, wait until process_weights_after_loading.
if self.quant_config.is_checkpoint_fp8_serialized:
if self.is_checkpoint_fp8_serialized:
# WEIGHT SCALE
if self.block_quant:
if hasattr(self.quant_config, "activation_scheme"):
@@ -340,62 +360,65 @@ class Fp8LinearMethod(LinearMethodBase):
else:
layer.register_parameter("input_scale", None)
def process_weights_after_loading_block_quant(self, layer: Module) -> None:
# If ROCm, normalize the weights and scales to e4m3fnuz
if _is_fp8_fnuz:
# activation_scheme: dynamic
weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
weight=layer.weight,
weight_scale=layer.weight_scale_inv,
input_scale=None,
)
layer.input_scale = None
elif _is_cpu:
assert (
_is_cpu_amx_available
), "Fp8LinearMethod on CPU requires that CPU has AMX support"
_amx_process_weight_after_loading(layer, ["weight"])
layer.weight_scale_inv = torch.nn.Parameter(
layer.weight_scale_inv.data, requires_grad=False
)
return
else:
# For fp8 linear weights run with deepgemm, the weights and scales need be requantized to ue8m0
from sglang.srt.layers.quantization.fp8_utils import (
deepgemm_w8a8_block_fp8_linear_with_fallback,
)
from sglang.srt.model_loader.utils import (
should_deepgemm_weight_requant_ue8m0,
)
if (
should_deepgemm_weight_requant_ue8m0(
weight_block_size=getattr(
self.quant_config, "weight_block_size", None
),
)
and (
self.w8a8_block_fp8_linear
is deepgemm_w8a8_block_fp8_linear_with_fallback
)
and (not layer.weight_scale_inv.format_ue8m0)
):
requant_weight_ue8m0_inplace(
layer.weight,
layer.weight_scale_inv,
self.quant_config.weight_block_size,
)
layer.weight_scale_inv.format_ue8m0 = True
weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data
layer.weight.data = weight.data
layer.weight_scale_inv.data = weight_scale.data
def process_weights_after_loading(self, layer: Module) -> None:
if self.block_quant:
# If ROCm, normalize the weights and scales to e4m3fnuz
if _is_fp8_fnuz:
# activation_scheme: dynamic
weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
weight=layer.weight,
weight_scale=layer.weight_scale_inv,
input_scale=None,
)
layer.input_scale = None
elif _is_cpu:
assert (
_is_cpu_amx_available
), "Fp8LinearMethod on CPU requires that CPU has AMX support"
_amx_process_weight_after_loading(layer, ["weight"])
layer.weight_scale_inv = torch.nn.Parameter(
layer.weight_scale_inv.data, requires_grad=False
)
return
else:
# For fp8 linear weights run with deepgemm, the weights and scales need be requantized to ue8m0
from sglang.srt.layers.quantization.fp8_utils import (
deepgemm_w8a8_block_fp8_linear_with_fallback,
)
from sglang.srt.model_loader.utils import (
should_deepgemm_weight_requant_ue8m0,
)
if (
should_deepgemm_weight_requant_ue8m0(
weight_block_size=getattr(
self.quant_config, "weight_block_size", None
),
)
and (
self.w8a8_block_fp8_linear
is deepgemm_w8a8_block_fp8_linear_with_fallback
)
and (not layer.weight_scale_inv.format_ue8m0)
):
requant_weight_ue8m0_inplace(
layer.weight,
layer.weight_scale_inv,
self.quant_config.weight_block_size,
)
layer.weight_scale_inv.format_ue8m0 = True
weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data
layer.weight.data = weight.data
layer.weight_scale_inv.data = weight_scale.data
self.process_weights_after_loading_block_quant(layer)
else:
layer.weight = Parameter(layer.weight.data, requires_grad=False)
# If checkpoint not serialized fp8, quantize the weights.
if not self.quant_config.is_checkpoint_fp8_serialized:
if not self.is_checkpoint_fp8_serialized:
if self.cutlass_fp8_supported or self.use_marlin:
# apply per-channel quantization default as
# cutlass sgl-kernel and marlin only support per-channel scale
@@ -763,47 +786,33 @@ class Fp8MoEMethod(FusedMoEMethodBase):
layer.w13_input_scale = None
layer.w2_input_scale = None
def process_weights_after_loading(self, layer: Module) -> None:
if _is_hip and _use_hip_int4:
self.process_weights_hip_int4(layer)
return
# Block quant doesn't need to process weights after loading
if self.block_quant:
# If ROCm, normalize the weights and scales to e4m3fnuz
if _is_fp8_fnuz:
# activation_scheme: dynamic
w13_weight, w13_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
weight=layer.w13_weight,
weight_scale=layer.w13_weight_scale_inv,
input_scale=None,
)
w2_weight, w2_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
weight=layer.w2_weight,
weight_scale=layer.w2_weight_scale_inv,
input_scale=None,
)
# Reset the parameter
layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
layer.w13_weight_scale_inv = torch.nn.Parameter(
w13_weight_scale, requires_grad=False
)
layer.w13_input_scale = None
layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
layer.w2_weight_scale_inv = torch.nn.Parameter(
w2_weight_scale, requires_grad=False
)
layer.w2_input_scale = None
if _use_aiter:
# add this section for MI300
# Pre-shuffle weights
layer.w13_weight.data = shuffle_weight(
layer.w13_weight.contiguous(), (16, 16)
)
layer.w2_weight.data = shuffle_weight(
layer.w2_weight.contiguous(), (16, 16)
)
elif _use_aiter:
def process_weights_after_loading_block_quant(self, layer: Module) -> None:
# If ROCm, normalize the weights and scales to e4m3fnuz
if _is_fp8_fnuz:
# activation_scheme: dynamic
w13_weight, w13_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
weight=layer.w13_weight,
weight_scale=layer.w13_weight_scale_inv,
input_scale=None,
)
w2_weight, w2_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
weight=layer.w2_weight,
weight_scale=layer.w2_weight_scale_inv,
input_scale=None,
)
# Reset the parameter
layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
layer.w13_weight_scale_inv = torch.nn.Parameter(
w13_weight_scale, requires_grad=False
)
layer.w13_input_scale = None
layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
layer.w2_weight_scale_inv = torch.nn.Parameter(
w2_weight_scale, requires_grad=False
)
layer.w2_input_scale = None
if _use_aiter:
# add this section for MI300
# Pre-shuffle weights
layer.w13_weight.data = shuffle_weight(
layer.w13_weight.contiguous(), (16, 16)
@@ -811,42 +820,59 @@ class Fp8MoEMethod(FusedMoEMethodBase):
layer.w2_weight.data = shuffle_weight(
layer.w2_weight.contiguous(), (16, 16)
)
elif _is_cpu:
assert (
_is_cpu_amx_available
), "Fp8MoEMethod on CPU requires that CPU has AMX support"
_amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"])
else:
# For fp8 moe run with deepgemm, the expert weights and scales need be requantized to ue8m0
from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE
from sglang.srt.model_loader.utils import (
should_deepgemm_weight_requant_ue8m0,
elif _use_aiter:
# Pre-shuffle weights
layer.w13_weight.data = shuffle_weight(
layer.w13_weight.contiguous(), (16, 16)
)
layer.w2_weight.data = shuffle_weight(
layer.w2_weight.contiguous(), (16, 16)
)
elif _is_cpu:
assert (
_is_cpu_amx_available
), "Fp8MoEMethod on CPU requires that CPU has AMX support"
_amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"])
else:
# For fp8 moe run with deepgemm, the expert weights and scales need be requantized to ue8m0
from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE
from sglang.srt.model_loader.utils import (
should_deepgemm_weight_requant_ue8m0,
)
# Check if MoE will actually use DeepGEMM runner
will_use_deepgemm = self.is_deepgemm_moe_runner_backend_enabled()
if (
should_deepgemm_weight_requant_ue8m0(
weight_block_size=getattr(
self.quant_config, "weight_block_size", None
),
)
and will_use_deepgemm
and not layer.w13_weight_scale_inv.format_ue8m0
):
assert isinstance(
layer, DeepEPMoE
), "DeepGemm MoE is only supported with DeepEPMoE"
weight_block_size = self.quant_config.weight_block_size
requant_weight_ue8m0_inplace(
layer.w13_weight, layer.w13_weight_scale_inv, weight_block_size
)
requant_weight_ue8m0_inplace(
layer.w2_weight, layer.w2_weight_scale_inv, weight_block_size
)
layer.w13_weight_scale_inv.format_ue8m0 = True
layer.w2_weight_scale_inv.format_ue8m0 = True
# Check if MoE will actually use DeepGEMM runner
will_use_deepgemm = self.is_deepgemm_moe_runner_backend_enabled()
def process_weights_after_loading(self, layer: Module) -> None:
if _is_hip and _use_hip_int4:
self.process_weights_hip_int4(layer)
return
if (
should_deepgemm_weight_requant_ue8m0(
weight_block_size=getattr(
self.quant_config, "weight_block_size", None
),
)
and will_use_deepgemm
and not layer.w13_weight_scale_inv.format_ue8m0
):
assert isinstance(
layer, DeepEPMoE
), "DeepGemm MoE is only supported with DeepEPMoE"
weight_block_size = self.quant_config.weight_block_size
requant_weight_ue8m0_inplace(
layer.w13_weight, layer.w13_weight_scale_inv, weight_block_size
)
requant_weight_ue8m0_inplace(
layer.w2_weight, layer.w2_weight_scale_inv, weight_block_size
)
layer.w13_weight_scale_inv.format_ue8m0 = True
layer.w2_weight_scale_inv.format_ue8m0 = True
# Block quant doesn't need to process weights after loading
if self.block_quant:
self.process_weights_after_loading_block_quant(layer)
return
# If checkpoint is fp16 or bfloat16, quantize in place.

10
python/sglang/srt/layers/quantization/fp8_kernel.py
View File

@@ -1861,10 +1861,6 @@ if _is_cuda:
):
return
# FIXME: for some models, this fake registration will cause NaN outputs.
# So we gate the fake registration with an environment variable for them.
if not get_bool_env_var("SGLANG_DISABLE_SGL_KERNEL_FAKE_REGISTER"):
@torch.library.register_fake("sgl_kernel::sgl_per_token_quant_fp8")
def _(input, output_q, output_s):
return
@torch.library.register_fake("sgl_kernel::sgl_per_token_quant_fp8")
def _(input, output_q, output_s):
return

55
python/sglang/srt/layers/quantization/fp8_utils.py
View File

@@ -2,6 +2,7 @@ from __future__ import annotations
import logging
from enum import Enum
from functools import lru_cache
from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
import torch
@@ -14,13 +15,6 @@ from sglang.srt.layers.quantization.mxfp4_tensor import MXFP4QuantizeUtil
if TYPE_CHECKING:
from sglang.srt.server_args import ServerArgs
try:
from vllm import _custom_ops as ops
VLLM_AVAILABLE = True
except ImportError:
VLLM_AVAILABLE = False
from sglang.srt.layers.quantization.fp8_kernel import (
fp8_dtype,
fp8_max,
@@ -101,6 +95,7 @@ def use_rowwise_torch_scaled_mm():
USE_ROWWISE_TORCH_SCALED_MM = use_rowwise_torch_scaled_mm()
@lru_cache(maxsize=1)
def cutlass_fp8_supported():
if not _is_cuda:
return False
@@ -907,9 +902,8 @@ def apply_fp8_linear(
# We also don't pad when using torch.compile,
# as it breaks with dynamic shapes.
if pad_output is None:
pad_output = (
not get_bool_env_var("SGLANG_ENABLE_TORCH_COMPILE")
and not cutlass_fp8_supported
pad_output = not cutlass_fp8_supported and not get_bool_env_var(
"SGLANG_ENABLE_TORCH_COMPILE"
)
output_padding = 17 if pad_output else None
@@ -956,37 +950,22 @@ def apply_fp8_linear(
)
if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]:
# cutlass_scaled_mm supports per tensor/channel W and per tensor/token A
# for sgl-kernel fp8_scaled_mm, it support per channel W now
if VLLM_AVAILABLE and use_vllm_cutlass_w8a8_fp8_kernel:
# Fall back to vllm cutlass w8a8 fp8 kernel
output = ops.cutlass_scaled_mm(
qinput,
weight,
out_dtype=input.dtype,
scale_a=x_scale,
scale_b=weight_scale,
bias=bias,
cutlass_compatible_b = weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel:
# Massage the input to be 2D
qinput = qinput.view(-1, qinput.shape[-1])
output = triton_scaled_mm(
qinput, weight, x_scale, weight_scale, input.dtype, bias
)
else:
cutlass_compatible_b = (
weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
output = fp8_scaled_mm(
qinput,
weight,
x_scale,
weight_scale,
out_dtype=input.dtype,
bias=bias,
)
if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel:
# Massage the input to be 2D
qinput = qinput.view(-1, qinput.shape[-1])
output = triton_scaled_mm(
qinput, weight, x_scale, weight_scale, input.dtype, bias
)
else:
output = fp8_scaled_mm(
qinput,
weight,
x_scale,
weight_scale,
out_dtype=input.dtype,
bias=bias,
)
return output.view(*output_shape)
# torch.scaled_mm supports per tensor weights + activations only

14
python/sglang/srt/metrics/collector.py
View File

@@ -73,6 +73,8 @@ class TimeStats:
prefill_end_time_host: float = 0.0
transfer_speed_gb_s: float = 0.0
transfer_total_mb: float = 0.0
# Number of prefill retries for this request
prefill_retry_count: int = 0
# Timestamp when prefill phase finishes, obtained from `time.time()`.
# Note that this differs from the other `_time` fields tracked by the
@@ -139,7 +141,8 @@ class TimeStats:
f"forward_duration={self.format_duration(forward_duration)}, "
f"start={self.prefill_bootstrap_queue_entry_time:.3f}, "
f"transfer_speed={self.transfer_speed_gb_s:.2f}GB/s, "
f"transfer_total={self.transfer_total_mb:.2f}MB"
f"transfer_total={self.transfer_total_mb:.2f}MB, "
f"#retries={self.prefill_retry_count}"
)
elif self.disagg_mode == DisaggregationMode.DECODE:
prealloc_duration = (
@@ -455,6 +458,11 @@ class SchedulerMetricsCollector:
documentation="The number of transfer failed requests.",
labelnames=labels.keys(),
)
self.num_prefill_retries_total = Counter(
name="sglang:num_prefill_retries_total",
documentation="Total number of prefill retries.",
labelnames=labels.keys(),
)
self.kv_transfer_speed_gb_s = Gauge(
name="sglang:kv_transfer_speed_gb_s",
documentation="The transfer speed of the KV cache in GB/s.",
@@ -854,6 +862,10 @@ class SchedulerMetricsCollector:
def increment_transfer_failed_reqs(self) -> None:
self.num_transfer_failed_reqs.labels(**self.labels).inc(1)
def increment_prefill_retries(self, count: int) -> None:
if count > 0:
self.num_prefill_retries_total.labels(**self.labels).inc(count)
def observe_per_stage_req_latency(self, stage: str, latency: float) -> None:
labels_with_stage = {**self.labels, "stage": stage}
self.per_stage_req_latency_seconds.labels(**labels_with_stage).observe(latency)

17
python/sglang/srt/utils/common.py
View File

@@ -1869,9 +1869,10 @@ def crash_on_warnings():
return get_bool_env_var("SGLANG_IS_IN_CI")
@functools.lru_cache(None)
def print_warning_once(msg: str) -> None:
# Set the stacklevel to 2 to print the caller's line info
logger.warning(msg, stacklevel=2)
logger.warning(msg)
@functools.lru_cache(None)
@@ -3774,6 +3775,20 @@ def calc_diff(x, y):
return 1 - sim
@contextmanager
def temp_attr_context(obj, attr, value):
if obj is None:
yield
return
original_value = getattr(obj, attr)
setattr(obj, attr, value)
try:
yield
finally:
setattr(obj, attr, original_value)
cached_device_index = -1