minke.yu/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
python/sglang/bench_one_batch_server.py
+5 -5
View File
@@ -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
+5 -5
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")
python/sglang/srt/layers/linear.py
+2 -3
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
python/sglang/srt/layers/parameter.py
+47 -2
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):
python/sglang/srt/layers/quantization/fp8.py
+65 -39
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,33 +223,29 @@ 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 create_weights(
def validate_block_quant_shapes(
self,
layer: torch.nn.Module,
input_size_per_partition: int,
output_partition_sizes: List[int],
input_size: int,
input_size_per_partition: int,
output_size: int,
params_dtype: torch.dtype,
output_size_per_partition: int,
output_partition_sizes: List[int],
skip_block_quant_check: bool = False,
**extra_weight_attrs,
):
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."
print_warning_once(
"Skipping block quantization checks for weight partition."
)
else:
# Required by row parallel
@@ -270,18 +268,40 @@ class Fp8LinearMethod(LinearMethodBase):
f"weight quantization block_n = {block_n}."
)
def create_weights(
self,
layer: torch.nn.Module,
input_size_per_partition: int,
output_partition_sizes: List[int],
input_size: int,
output_size: int,
params_dtype: torch.dtype,
skip_block_quant_check: bool = False,
**extra_weight_attrs,
):
# Copy the layer attributes
output_size_per_partition = sum(output_partition_sizes)
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
weight_dtype = (
torch.float8_e4m3fn
if self.quant_config.is_checkpoint_fp8_serialized
else params_dtype
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.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,8 +360,7 @@ class Fp8LinearMethod(LinearMethodBase):
else:
layer.register_parameter("input_scale", None)
def process_weights_after_loading(self, layer: Module) -> None:
if self.block_quant:
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
@@ -391,11 +410,15 @@ class Fp8LinearMethod(LinearMethodBase):
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:
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,13 +786,7 @@ 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:
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
@@ -847,6 +864,15 @@ class Fp8MoEMethod(FusedMoEMethodBase):
)
layer.w13_weight_scale_inv.format_ue8m0 = True
layer.w2_weight_scale_inv.format_ue8m0 = True
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:
self.process_weights_after_loading_block_quant(layer)
return
# If checkpoint is fp16 or bfloat16, quantize in place.
python/sglang/srt/layers/quantization/fp8_kernel.py
-4
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
python/sglang/srt/layers/quantization/fp8_utils.py
+5 -26
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,22 +950,7 @@ 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,
)
else:
cutlass_compatible_b = (
weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
)
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])
python/sglang/srt/metrics/collector.py
+13 -1
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)
python/sglang/srt/utils/common.py
+16 -1
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