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

[AMD][MXFP4] Online MXFP4 quantization 1/N - dense and MOE models w. original BF16 weight (#18005)

Browse Source 
Co-authored-by: Bowen Bao <bowenbao@amd.com>
Co-authored-by: Colin Zeng <Colin.Zeng@amd.com>
This commit is contained in:
fxmarty-amd
2026-06-03 12:55:24 -07:00
committed by GitHub
co-authored by Bowen Bao Colin Zeng
parent e0b692600f
commit 293816ab14
12 changed files with 509 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/model_config.py
+2
View File
@@ -1145,6 +1145,7 @@ class ModelConfig:
"mxfp4",
"auto-round",
"quark_int4fp8_moe",
"quark_mxfp4",
]
optimized_quantization_methods = [
"fp8",
@@ -1167,6 +1168,7 @@ class ModelConfig:
"petit_nvfp4",
"quark",
"modelslim",
"quark_mxfp4",
]
compatible_quantization_methods = {
"modelopt_fp8": ["modelopt"],
python/sglang/srt/constants.py
+2
View File
@@ -10,3 +10,5 @@ GPU_MEMORY_ALL_TYPES = [
]
HEALTH_CHECK_RID_PREFIX = "HEALTH_CHECK"
GIB_BYTES = 1073741824 # 1024**3
python/sglang/srt/layers/quantization/__init__.py
+1
View File
@@ -90,6 +90,7 @@ BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = {
"petit_nvfp4": PetitNvFp4Config,
"fbgemm_fp8": FBGEMMFp8Config,
"quark": QuarkConfig,
"quark_mxfp4": QuarkConfig,
"auto-round": AutoRoundConfig,
"modelslim": ModelSlimConfig,
"quark_int4fp8_moe": QuarkInt4Fp8Config,
python/sglang/srt/layers/quantization/quark/quark.py
+103 -3
View File
@@ -29,6 +29,8 @@ from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.utils import get_device_capability
if TYPE_CHECKING:
from transformers import PretrainedConfig
from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
__all__ = ["QuarkLinearMethod", "QuarkFusedMoEMethod"]
@@ -40,21 +42,47 @@ class QuarkConfig(QuantizationConfig):
def __init__(
self,
quant_config: dict[str, Any],
quant_config: Optional[dict[str, Any]] = None,
hf_config: "PretrainedConfig | None" = None,
kv_cache_group: Optional[list[str]] = None,
kv_cache_config: Optional[dict[str, Any]] = None,
pack_method: str = "reorder",
is_prequantized: bool = False,
online_scheme: Optional[str] = None,
):
super().__init__()
if kv_cache_group is None:
kv_cache_group = []
if online_scheme is not None:
assert not is_prequantized
if online_scheme == "quark_mxfp4":
quant_config = self._create_online_mxfp4_config(
model_type=hf_config.model_type
)
else:
raise ValueError(f"Unsupported online_scheme: {online_scheme}")
if quant_config is None:
raise ValueError("Either quant_config or online_scheme must be provided")
self.quant_config = quant_config
self.kv_cache_group = kv_cache_group
self.kv_cache_config = kv_cache_config
self.pack_method = pack_method
self.exclude_layers = cast(list[str], self.quant_config.get("exclude", []))
self.is_prequantized = is_prequantized
self.packed_modules_mapping = self.quant_config["packed_modules_mapping"]
self._quantized_layers = set()
@property
def quantized_layers(self) -> tuple[list[str], int]:
# Extract unique layer types (last part after ".")
layer_types = sorted(
set(name.split(".")[-1] for name in self._quantized_layers)
)
return layer_types, len(self._quantized_layers)
def get_linear_method(self) -> "QuarkLinearMethod":
return QuarkLinearMethod(self)
@@ -83,6 +111,7 @@ class QuarkConfig(QuantizationConfig):
self, layer: torch.nn.Module, prefix: str
) -> Optional["QuantizeMethodBase"]:
# Check if the layer is skipped for quantization.
if should_ignore_layer(
prefix,
ignore=self.exclude_layers,
@@ -97,14 +126,17 @@ class QuarkConfig(QuantizationConfig):
if isinstance(layer, LinearBase):
scheme = self.get_linear_scheme(layer=layer, layer_name=prefix)
layer.scheme = scheme
self._quantized_layers.add(prefix)
return QuarkLinearMethod(self)
if isinstance(layer, RadixAttention):
self._quantized_layers.add(prefix)
return QuarkKVCacheMethod(self)
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
if isinstance(layer, FusedMoE):
self._quantized_layers.add(prefix)
layer.scheme = self.get_moe_scheme(layer, prefix)
return QuarkFusedMoEMethod(self)
@@ -175,12 +207,72 @@ class QuarkConfig(QuantizationConfig):
kv_cache_group=kv_cache_group,
kv_cache_config=kv_cache_config,
pack_method=pack_method,
is_prequantized=True,
)
@classmethod
def get_config_filenames(cls) -> list[str]:
return []
@staticmethod
def _create_online_mxfp4_config(model_type: str) -> dict[str, Any]:
"""
Create a synthetic quant_config for online MXFP4 quantization.
"""
# MOE gate/router is typically implemented as a ReplicatedLinear, and skipped for quantization for accuracy reasons.
# lm_head/embed_tokens is also skipped for accuracy reasons, normally not handled by `QuarkConfig` in any case, but adding them here for safety.
exclude = [
"re:.*gate$",
"re:.*router",
"re:.*lm_head",
"re:.*embed_tokens",
]
if model_type == "qwen3_5_moe":
# Exclusion for accuracy adapted from
# https://huggingface.co/amd/Qwen3.5-397B-A17B-MXFP4/blob/main/config.json
exclude.extend(
[
"re:.*n_proj_a",
"re:.*in_proj_b",
"re:.*in_proj_qkv",
"re:.*in_proj_z",
"re:.*o_proj",
"re:.*out_proj",
"re:.*qkv_proj",
"re:.*shared_expert",
]
)
return {
"packed_modules_mapping": {},
"exclude": exclude,
"global_quant_config": {
"weight": {
"dtype": "fp4",
"qscheme": "per_group",
"group_size": 32,
"is_dynamic": False,
"scale_format": "e8m0",
},
"input_tensors": {
"dtype": "fp4",
"qscheme": "per_group",
"group_size": 32,
"is_dynamic": True,
"scale_format": "e8m0",
},
"output_tensors": None,
"bias": None,
},
"layer_quant_config": {},
"layer_type_quant_config": {},
"export": {
"kv_cache_group": [],
"pack_method": "reorder",
},
}
def _check_scheme_supported(self, min_capability: int, error: bool = True) -> bool:
capability_tuple = get_device_capability()
@@ -337,7 +429,11 @@ class QuarkConfig(QuantizationConfig):
input_config = cast(dict[str, Any], config.get("input_tensors"))
if self._is_mx_fp4(weight_config, input_config):
return QuarkW4A4MXFP4(weight_config, input_config)
return QuarkW4A4MXFP4(
weight_config,
input_config,
is_checkpoint_mxfp4_serialized=self.is_prequantized,
)
if self._is_fp8_w8a8(weight_config, input_config):
is_fp8_w8a8_supported = self._check_scheme_supported(
QuarkW8A8Fp8.get_min_capability(), error=False
@@ -383,7 +479,11 @@ class QuarkConfig(QuantizationConfig):
input_config = layer_quant_config.get("input_tensors")
if self._is_mx_fp4(weight_config, input_config):
return QuarkW4A4MXFp4MoE(weight_config, input_config)
return QuarkW4A4MXFp4MoE(
weight_config,
input_config,
is_checkpoint_mxfp4_serialized=self.is_prequantized,
)
elif self._is_fp8_w8a8(weight_config, input_config):
return QuarkW8A8FP8MoE(weight_config, input_config)
else:
python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py
+68 -6
View File
@@ -1,5 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
import logging
from typing import Any, Callable, Optional
import torch
@@ -7,6 +8,7 @@ import torch
from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter
from sglang.srt.layers.quantization.quark.schemes import QuarkLinearScheme
from sglang.srt.utils import is_hip
from sglang.srt.utils.common import mxfp_supported
_is_hip = is_hip()
if _is_hip:
@@ -19,6 +21,7 @@ if _is_hip:
__all__ = ["QuarkW4A4MXFP4"]
logger = logging.getLogger(__name__)
OCP_MX_BLOCK_SIZE = 32
@@ -26,19 +29,35 @@ OCP_MX_BLOCK_SIZE = 32
class QuarkW4A4MXFP4(QuarkLinearScheme):
def __init__(
self, weight_quant_spec: dict[str, Any], input_quant_spec: dict[str, Any]
self,
weight_quant_spec: dict[str, Any],
input_quant_spec: dict[str, Any],
is_checkpoint_mxfp4_serialized: bool = True,
):
self.out_dtype = torch.get_default_dtype()
self.qscheme = "per_group"
self.weight_quant_spec = weight_quant_spec
self.input_quant_spec = input_quant_spec
self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized
if not self.is_checkpoint_mxfp4_serialized:
if not mxfp_supported():
raise NotImplementedError(
"Online MXFP4 quantization requires an AMD ROCm device with "
"FP4 hardware support (gfx95x, e.g. MI355x)."
)
logger.info_once(
"Using online MXFP4 quantization from a higher precision checkpoint. Beware that this optimization may degrade prediction quality - please validate your model accuracy. More details at https://docs.sglang.io/advanced_features/quantization.html#online-quantization."
)
@classmethod
def get_min_capability(cls) -> int:
return 70
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
return
if not self.is_checkpoint_mxfp4_serialized:
assert layer.weight.dtype == torch.uint8
assert layer.weight_scale.dtype == torch.uint8
def create_weights(
self,
@@ -49,10 +68,19 @@ class QuarkW4A4MXFP4(QuarkLinearScheme):
weight_loader: Callable,
**kwargs,
):
self.input_size_per_partition = input_size_per_partition
output_size_per_partition = sum(output_partition_sizes)
self.output_size_per_partition = output_size_per_partition
layer.logical_widths = output_partition_sizes
original_weight_loader = weight_loader
if not self.is_checkpoint_mxfp4_serialized:
weight_loader = self.get_online_mxfp4_weight_loader(layer, weight_loader)
# WEIGHT
# Both serialized and online quantization use packed uint8 format
weight = PackedvLLMParameter(
data=torch.empty(
output_size_per_partition,
@@ -76,19 +104,50 @@ class QuarkW4A4MXFP4(QuarkLinearScheme):
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
weight_loader=original_weight_loader,
)
layer.register_parameter("weight_scale", weight_scale)
def get_online_mxfp4_weight_loader(
self,
layer,
original_weight_loader: Callable,
) -> Callable:
"""
Wrap the original weight loader to perform online MXFP4 quantization.
"""
def online_mxfp4_weight_loader(
param: torch.nn.Parameter,
loaded_weight: torch.Tensor,
shard_id: int | str | None = None,
):
# Materialize on device the loaded weight.
loaded_weight = loaded_weight.to(param.device)
# Quantize the loaded weight shard immediately. Since MXFP4 uses per-group quantization, there is no need to load all shards (e.g. q_proj, k_proj, v_proj) before doing online quantization.
qweight, weight_scale = dynamic_mxfp4_quant(loaded_weight)
# Required e.g. for q_proj, k_proj, v_proj.
kwargs = {}
if shard_id is not None:
kwargs["loaded_shard_id"] = shard_id
# Use the original weight loader to handle the loading logic
# (e.g. qkv sharding, etc.)
original_weight_loader(param, qweight, **kwargs)
layer.weight_scale.weight_loader(layer.weight_scale, weight_scale, **kwargs)
return online_mxfp4_weight_loader
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
# This path does not have support for bias currently
assert bias is None, "bias is not supported"
# Bias will be added after the GEMM if provided
three_d = False
fused_gemm_split_cat = False
x_s = None
@@ -152,6 +211,9 @@ class QuarkW4A4MXFP4(QuarkLinearScheme):
else:
gemm_afp4wfp4(x_q, layer.weight, x_s, layer.weight_scale, self.out_dtype, y)
if bias is not None:
y = y + bias
if fused_gemm_split_cat:
return k, v
elif three_d:
python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py
+77 -10
View File
@@ -16,6 +16,7 @@ from sglang.srt.utils import (
is_hip,
set_weight_attrs,
)
from sglang.srt.utils.common import mxfp_supported
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher import (
@@ -35,14 +36,25 @@ if _use_aiter:
from aiter.ops.shuffle import shuffle_weight
from aiter.utility.fp4_utils import e8m0_shuffle
if _is_hip:
from aiter.ops.triton.quant import dynamic_mxfp4_quant
else:
dynamic_mxfp4_quant = None
OCP_MX_BLOCK_SIZE = 32
class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
def __init__(self, weight_config: dict[str, Any], input_config: dict[str, Any]):
def __init__(
self,
weight_config: dict[str, Any],
input_config: dict[str, Any],
is_checkpoint_mxfp4_serialized: bool = True,
):
self.weight_quant = weight_config
self.input_quant = input_config
self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized
weight_qscheme = self.weight_quant.get("qscheme")
input_qscheme = self.input_quant.get("qscheme")
@@ -56,6 +68,18 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
self.static_input_scales = not self.input_quant.get("is_dynamic")
self.with_bias = False
if not self.is_checkpoint_mxfp4_serialized:
if not mxfp_supported():
raise NotImplementedError(
"Online MXFP4 quantization for MoE layers requires an AMD ROCm "
"device with FP4 hardware support (gfx95x, e.g. MI355x)."
)
logger.info_once(
"Using online MXFP4 quantization for MoE layers from a higher precision checkpoint. "
"Beware that this optimization may degrade prediction quality - please validate your model accuracy. "
"More details at https://docs.sglang.io/advanced_features/quantization.html#online-quantization."
)
@classmethod
def get_min_capability(cls) -> int:
return 70
@@ -90,7 +114,16 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
params_dtype = torch.uint8
# WEIGHTS
original_weight_loader = extra_weight_attrs.get("weight_loader")
if self.is_checkpoint_mxfp4_serialized:
weight_loader = original_weight_loader
else:
weight_loader = self.get_online_weight_loader(layer, original_weight_loader)
extra_weight_attrs["weight_loader"] = weight_loader
# WEIGHTS — always uint8 (packed mxfp4), always on device
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
@@ -101,7 +134,6 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
@@ -114,10 +146,11 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# WEIGHT_SCALES
extra_weight_attrs["weight_loader"] = original_weight_loader
w13_weight_scale = torch.nn.Parameter(
torch.ones(
num_experts,
@@ -140,26 +173,60 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
),
requires_grad=False,
)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
float_dtype = torch.get_default_dtype()
def get_online_weight_loader(self, layer, original_weight_loader):
"""
Wrap the original weight loader to perform online MXFP4 quantization.
"""
def online_mxfp4_moe_weight_loader(
param: torch.nn.Parameter,
loaded_weight: torch.Tensor,
weight_name: str,
shard_id: str,
expert_id: int,
):
if dynamic_mxfp4_quant is None:
raise NotImplementedError(
"Online MXFP4 quantization for MoE is only supported on AMD GPUs."
)
# Materialize on device the loaded weight.
loaded_weight = loaded_weight.to(param.device)
# Quantize the high-precision shard loaded_weight to MXFP4.
qweight, weight_scale = dynamic_mxfp4_quant(loaded_weight)
original_weight_loader(param, qweight, weight_name, shard_id, expert_id)
if "w13" in weight_name:
scale_param = layer.w13_weight_scale
scale_weight_name = "w13_weight_scale"
else:
# w2.
scale_param = layer.w2_weight_scale
scale_weight_name = "w2_weight_scale"
scale_param.weight_loader(
scale_param, weight_scale, scale_weight_name, shard_id, expert_id
)
return online_mxfp4_moe_weight_loader
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
# Pre-shuffle weight scales
s0, s1, _ = layer.w13_weight_scale.shape
w13_weight_scale = layer.w13_weight_scale.view(s0 * s1, -1)
w13_weight_scale = e8m0_shuffle(w13_weight_scale)
# layer.w13_weight_scale = torch.nn.Parameter(w13_weight_scale, requires_grad=False)
layer.w13_weight_scale.data = w13_weight_scale.view(s0, s1, -1)
s0, s1, _ = layer.w2_weight_scale.shape
w2_weight_scale = layer.w2_weight_scale.view(s0 * s1, -1)
w2_weight_scale = e8m0_shuffle(w2_weight_scale)
# layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, requires_grad=False)
layer.w2_weight_scale.data = w2_weight_scale.view(s0, s1, -1)
# Pre-shuffle weight
python/sglang/srt/model_executor/model_runner.py
+14
View File
@@ -1417,6 +1417,20 @@ class ModelRunner(ModelRunnerKVCacheMixin):
f"avail mem={after_avail_memory:.2f} GB, "
f"mem usage={self.weight_load_mem_usage:.2f} GB."
)
# TODO: Make sure all models have `quant_config` attribute, and all online quantization methods register which layers they actually quantize.
if (
hasattr(self.model, "quant_config")
and hasattr(self.model.quant_config, "quantized_layers")
and self.server_args.quantization is not None
):
layer_types, quantized_layers_count = (
self.model.quant_config.quantized_layers
)
logger.info(
f"Online {self.server_args.quantization} quantization: quantized {quantized_layers_count} layers of types: {layer_types}"
)
if self.server_args.debug_tensor_dump_output_folder is not None:
dump_folder = self.server_args.debug_tensor_dump_output_folder
if self.spec_algorithm.is_eagle():
python/sglang/srt/model_loader/loader.py
+24
View File
@@ -37,12 +37,14 @@ import huggingface_hub
import numpy as np
import torch
from sglang.srt.constants import GIB_BYTES
from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
RemoteInstanceWeightLoaderBackend,
get_remote_instance_transfer_engine_info_per_rank,
register_memory_region,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import get_available_gpu_memory
# Try to import accelerate (optional dependency)
try:
@@ -82,6 +84,7 @@ from sglang.srt.model_loader.utils import (
get_model_architecture,
set_default_torch_dtype,
)
from sglang.srt.utils.common import is_cuda_alike
# Constants for memory management
DEFAULT_GPU_MEMORY_FRACTION_FOR_CALIBRATION = (
@@ -746,8 +749,29 @@ class DefaultModelLoader(BaseModelLoader):
@staticmethod
def load_weights_and_postprocess(model, weights, target_device):
# Used in tests to verify memory savings when using online quantization.
if is_cuda_alike():
peak_memory = torch.cuda.max_memory_allocated()
logger.debug(
"Peak GPU memory before loading weights: %s GiB",
f"{peak_memory / GIB_BYTES:.3f}",
)
memory_start = get_available_gpu_memory(
target_device.type, gpu_id=torch.cuda.current_device()
)
model.load_weights(weights)
# Used in tests to verify memory savings when using online quantization.
if is_cuda_alike():
memory_end = get_available_gpu_memory(
target_device.type, gpu_id=torch.cuda.current_device()
)
logger.debug(
"Memory increase during load_weights: %s GiB",
f"{memory_start - memory_end:.3f}",
)
for _, module in model.named_modules():
quant_method = getattr(module, "quant_method", None)
if quant_method is not None:
python/sglang/srt/model_loader/weight_utils.py
+6
View File
@@ -294,9 +294,15 @@ def get_quant_config(
possible_config_filenames = quant_cls.get_config_filenames()
# If the quantization config is not found, use the default config.
# TODO: standardize the handling of online quantization with custom handlenames (mxfp8, quark_mxfp4, etc.)
if not possible_config_filenames:
if model_config.quantization == "mxfp8":
return Fp8Config(use_mxfp8=True, is_checkpoint_fp8_serialized=False)
if model_config.quantization == "quark_mxfp4":
return quant_cls(
online_scheme=model_config.quantization,
hf_config=model_config.hf_config,
)
return quant_cls()
config_files = glob.glob(os.path.join(hf_folder, "*.json"))
python/sglang/srt/server_args.py
+11 -6
View File
@@ -125,32 +125,36 @@ LOAD_FORMAT_CHOICES = [
"runai_streamer",
]
# TODO: this list should likely contain only methods that support online quantization, or that support using custom quantization classes compatible with a given `quant_method` in config.json.
# Some of the choices here do NOT support online quantization.
QUANTIZATION_CHOICES = [
"awq",
"fp8",
"mxfp8",
"fp8", # MOE + linear online quantization.
"mxfp8", # MOE + linear online quantization.
"gptq",
"marlin",
"gptq_marlin",
"awq_marlin",
"bitsandbytes",
"gguf",
# Modelopt has some online quantization support through ModelOptModelLoader.
"modelopt",
"modelopt_fp8",
"modelopt_fp4",
"modelopt_mixed",
"petit_nvfp4",
"w8a8_int8",
"w8a8_fp8",
"moe_wna16",
"w8a8_int8", # mentioned in quantization.md documentation, supporting compressed-tensors quant_method.
"w8a8_fp8", # mentioned in quantization.md documentation, supporting compressed-tensors quant_method.
"moe_wna16", # custom loading logic for gptq/awq checkpoints (likely untested/unused)
"qoq",
"w4afp8",
"mxfp4",
"mxfp4", # MOE-only.
"auto-round",
"compressed-tensors", # for Ktransformers
"modelslim", # for NPU
"quark", # AMD Quark quantizer (FP8 / MXFP4 / Int4FP8 etc.)
"quark_int4fp8_moe",
"quark_mxfp4", # Online MOE + linear quantization.
# Apple Silicon MLX backend — on-the-fly quantization of fp16 weights at load
# time via mlx.nn.quantize. Only takes effect when SGLANG_USE_MLX=1.
"mlx_q4", # 4 bits, group_size=64 (mlx-community default)
@@ -158,6 +162,7 @@ QUANTIZATION_CHOICES = [
"unquant",
]
SPECULATIVE_DRAFT_MODEL_QUANTIZATION_CHOICES = QUANTIZATION_CHOICES
ATTENTION_BACKEND_CHOICES = [