Code clean up for fp8 quantization (#16982)

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)

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:
+        except (RuntimeError, json.decoder.JSONDecodeError, UnicodeDecodeError) as e:
-            logging.error(f"Hit invalid json_schema: {key_string=}, {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")

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

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
+        copy_with_check(self.data, loaded_weight)
-        self.data.copy_(loaded_weight)
 
     def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs):
 

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
+    It supports the following quantization schemes:
-    activation scaling. The weight scaling factor will be initialized after
+    - Per-channel weight quantization + per-token activation quantization
-    the model weights are loaded.
+    - Per-tensor weight quantization + per-tensor activation quantization
+    - Blockwise weight quantization + blockwise activation quantization
 
-    Limitations:
+    It supports the following checkpoint formats:
-    1. Only support per-tensor quantization due to torch._scaled_mm support.
+    - FP8 checkpoint
-    2. Only support float8_e4m3fn data type due to the limitation of
+    - FP16/BF16 checkpoint. In this case, the weights will be quantized to FP8 during the weight loading.
-       torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856)
+
+    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
+            torch.float8_e4m3fn if self.is_checkpoint_fp8_serialized else params_dtype
-            if self.quant_config.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
+            self.process_weights_after_loading_block_quant(layer)
-            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
         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:
+    def process_weights_after_loading_block_quant(self, layer: Module) -> None:
-        if _is_hip and _use_hip_int4:
+        # If ROCm, normalize the weights and scales to e4m3fnuz
-            self.process_weights_hip_int4(layer)
+        if _is_fp8_fnuz:
-            return
+            # activation_scheme: dynamic
-
+            w13_weight, w13_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
-        # Block quant doesn't need to process weights after loading
+                weight=layer.w13_weight,
-        if self.block_quant:
+                weight_scale=layer.w13_weight_scale_inv,
-            # If ROCm, normalize the weights and scales to e4m3fnuz
+                input_scale=None,
-            if _is_fp8_fnuz:
+            )
-                # activation_scheme: dynamic
+            w2_weight, w2_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
-                w13_weight, w13_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
+                weight=layer.w2_weight,
-                    weight=layer.w13_weight,
+                weight_scale=layer.w2_weight_scale_inv,
-                    weight_scale=layer.w13_weight_scale_inv,
+                input_scale=None,
-                    input_scale=None,
+            )
-                )
+            # Reset the parameter
-                w2_weight, w2_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
+            layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
-                    weight=layer.w2_weight,
+            layer.w13_weight_scale_inv = torch.nn.Parameter(
-                    weight_scale=layer.w2_weight_scale_inv,
+                w13_weight_scale, requires_grad=False
-                    input_scale=None,
+            )
-                )
+            layer.w13_input_scale = None
-                # Reset the parameter
+            layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
-                layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
+            layer.w2_weight_scale_inv = torch.nn.Parameter(
-                layer.w13_weight_scale_inv = torch.nn.Parameter(
+                w2_weight_scale, requires_grad=False
-                    w13_weight_scale, requires_grad=False
+            )
-                )
+            layer.w2_input_scale = None
-                layer.w13_input_scale = None
+            if _use_aiter:
-                layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
+                # add this section for MI300
-                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:
                 # 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:
+        elif _use_aiter:
-                assert (
+            # Pre-shuffle weights
-                    _is_cpu_amx_available
+            layer.w13_weight.data = shuffle_weight(
-                ), "Fp8MoEMethod on CPU requires that CPU has AMX support"
+                layer.w13_weight.contiguous(), (16, 16)
-                _amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"])
+            )
-            else:
+            layer.w2_weight.data = shuffle_weight(
-                # For fp8 moe run with deepgemm, the expert weights and scales need be requantized to ue8m0
+                layer.w2_weight.contiguous(), (16, 16)
-                from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE
+            )
-                from sglang.srt.model_loader.utils import (
+        elif _is_cpu:
-                    should_deepgemm_weight_requant_ue8m0,
+            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
+    def process_weights_after_loading(self, layer: Module) -> None:
-                will_use_deepgemm = self.is_deepgemm_moe_runner_backend_enabled()
+        if _is_hip and _use_hip_int4:
+            self.process_weights_hip_int4(layer)
+            return
 
-                if (
+        # Block quant doesn't need to process weights after loading
-                    should_deepgemm_weight_requant_ue8m0(
+        if self.block_quant:
-                        weight_block_size=getattr(
+            self.process_weights_after_loading_block_quant(layer)
-                            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
             return
 
         # If checkpoint is fp16 or bfloat16, quantize in place.

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.
+    @torch.library.register_fake("sgl_kernel::sgl_per_token_quant_fp8")
-    # So we gate the fake registration with an environment variable for them.
+    def _(input, output_q, output_s):
-    if not get_bool_env_var("SGLANG_DISABLE_SGL_KERNEL_FAKE_REGISTER"):
+        return
-
-        @torch.library.register_fake("sgl_kernel::sgl_per_token_quant_fp8")
-        def _(input, output_q, output_s):
-            return

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 = (
+        pad_output = not cutlass_fp8_supported and not get_bool_env_var(
-            not get_bool_env_var("SGLANG_ENABLE_TORCH_COMPILE")
+            "SGLANG_ENABLE_TORCH_COMPILE"
-            and not cutlass_fp8_supported
         )
     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
+        cutlass_compatible_b = weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
-        # for sgl-kernel fp8_scaled_mm, it support per channel W now
+        if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel:
-        if VLLM_AVAILABLE and use_vllm_cutlass_w8a8_fp8_kernel:
+            # Massage the input to be 2D
-            # Fall back to vllm cutlass w8a8 fp8 kernel
+            qinput = qinput.view(-1, qinput.shape[-1])
-            output = ops.cutlass_scaled_mm(
+            output = triton_scaled_mm(
-                qinput,
+                qinput, weight, x_scale, weight_scale, input.dtype, bias
-                weight,
-                out_dtype=input.dtype,
-                scale_a=x_scale,
-                scale_b=weight_scale,
-                bias=bias,
             )
         else:
-            cutlass_compatible_b = (
+            output = fp8_scaled_mm(
-                weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
+                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

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)

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