diff --git a/python/sglang/bench_one_batch_server.py b/python/sglang/bench_one_batch_server.py index 793cbcfeb..565201b33 100644 --- a/python/sglang/bench_one_batch_server.py +++ b/python/sglang/bench_one_batch_server.py @@ -671,12 +671,12 @@ def run_benchmark(server_args: ServerArgs, bench_args: BenchArgs): profile_output_dir=bench_args.profile_output_dir, ) ) - - # Replace the profile link - for res, profile_res in zip(results, profile_results): - res.profile_link = profile_res.profile_link except Exception as e: - print(f"Error profiling, there will be no profile trace dump: {e}") + print(f"Error profiling, some profile traces may not be dumped: {e}") + + # Replace the profile link for any successful profile results + for res, profile_res in zip(results, profile_results, strict=False): + res.profile_link = profile_res.profile_link finally: if proc: kill_process_tree(proc.pid) diff --git a/python/sglang/srt/constrained/xgrammar_backend.py b/python/sglang/srt/constrained/xgrammar_backend.py index fe3f4822a..3bc306521 100644 --- a/python/sglang/srt/constrained/xgrammar_backend.py +++ b/python/sglang/srt/constrained/xgrammar_backend.py @@ -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") diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 428f3a261..ea0fdbab2 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -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 diff --git a/python/sglang/srt/layers/parameter.py b/python/sglang/srt/layers/parameter.py index 3cc1d2344..565f5b9fd 100644 --- a/python/sglang/srt/layers/parameter.py +++ b/python/sglang/srt/layers/parameter.py @@ -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): diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 25fed3e09..a2d31e9c2 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -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. diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index 27f6c86a8..7701f9757 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index e7548dbce..61219f6b0 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -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 diff --git a/python/sglang/srt/metrics/collector.py b/python/sglang/srt/metrics/collector.py index 316cc177a..5b0c47b86 100644 --- a/python/sglang/srt/metrics/collector.py +++ b/python/sglang/srt/metrics/collector.py @@ -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) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index ad8217116..f95bcf126 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -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