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

[diffusion] [AMD] feat: support online MXFP4 and fp8 quantization (#21431)

Browse Source 
Co-authored-by: Bowen Bao <bowenbao@amd.com>
Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
Colin Z
2026-05-14 08:52:01 +08:00
committed by GitHub
co-authored by Bowen Bao HAI
parent af7511e0e8
commit 34c0029f0a
10 changed files with 417 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/flash_attention_v3.py
+33 -3
View File
@@ -208,9 +208,39 @@ def flash_attn_varlen_func(
):
if not _is_fa3_supported():
raise NotImplementedError(
"flash_attn at sgl-kernel is only supported on sm90 and above"
)
# Fall back to flash_attn package (FA2) on platforms without sgl-kernel FA3
# (e.g. ROCm, or CUDA < sm90)
if cu_seqlens_q is not None:
from flash_attn import flash_attn_varlen_func as fa2_flash_attn_varlen_func
return fa2_flash_attn_varlen_func(
q,
k,
v,
cu_seqlens_q,
cu_seqlens_k,
max_seqlen_q,
max_seqlen_k,
softmax_scale=softmax_scale,
causal=causal,
window_size=window_size,
softcap=softcap,
return_attn_probs=return_softmax_lse,
)
else:
# 4D inputs (batch, seqlen, nheads, headdim) without cu_seqlens
from flash_attn import flash_attn_func as fa2_flash_attn_func
return fa2_flash_attn_func(
q,
k,
v,
softmax_scale=softmax_scale,
causal=causal,
window_size=window_size,
softcap=softcap,
return_attn_probs=return_softmax_lse,
)
return _call_fa3_kernel(
_load_fa3_kernels()["flash_attn_varlen_func"],
python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py
+3 -1
View File
@@ -14,10 +14,11 @@ from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
ModelOptFp8Config,
)
from sglang.multimodal_gen.runtime.layers.quantization.modelslim import ModelSlimConfig
from sglang.multimodal_gen.runtime.layers.quantization.mxfp4 import Mxfp4Config
from sglang.multimodal_gen.runtime.layers.quantization.mxfp8_npu import MXFP8Config
QuantizationMethods = Literal[
"fp8", "modelopt", "modelopt_fp8", "modelopt_fp4", "modelslim", "mxfp8"
"fp8", "modelopt", "modelopt_fp8", "modelopt_fp4", "modelslim", "mxfp4"
]
QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
@@ -29,6 +30,7 @@ _CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {
"modelopt_fp4": ModelOptFp4Config,
"modelslim": ModelSlimConfig,
"fp8": Fp8Config,
"mxfp4": Mxfp4Config,
"mxfp8": MXFP8Config,
}
python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py
+7 -1
View File
@@ -83,6 +83,7 @@ class Fp8Config(QuantizationConfig):
activation_scheme: str = "dynamic",
ignored_layers: Optional[List[str]] = None,
weight_block_size: List[int] = None,
packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
) -> None:
self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
if is_checkpoint_fp8_serialized:
@@ -91,6 +92,7 @@ class Fp8Config(QuantizationConfig):
raise ValueError(f"Unsupported activation scheme {activation_scheme}")
self.activation_scheme = activation_scheme
self.ignored_layers = ignored_layers or []
self.packed_modules_mapping = packed_modules_mapping or {}
if weight_block_size is not None:
if not is_checkpoint_fp8_serialized:
raise ValueError(
@@ -147,7 +149,11 @@ class Fp8Config(QuantizationConfig):
from sglang.multimodal_gen.runtime.layers.linear import LinearBase
if isinstance(layer, LinearBase):
if is_layer_skipped(prefix, self.ignored_layers):
if is_layer_skipped(
prefix,
self.ignored_layers,
fused_mapping=self.packed_modules_mapping,
):
return UnquantizedLinearMethod()
return Fp8LinearMethod(self)
return None
python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py
+237
View File
@@ -0,0 +1,237 @@
import logging
from typing import Dict, List, Optional
import torch
from torch.nn.parameter import Parameter
from sglang.multimodal_gen.runtime.layers.linear import (
LinearMethodBase,
UnquantizedLinearMethod,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
QuantizationConfig,
)
from sglang.multimodal_gen.runtime.models.parameter import (
ModelWeightParameter,
PerTensorScaleParameter,
)
from sglang.srt.layers.quantization.utils import is_layer_skipped
from sglang.srt.utils import is_hip, mxfp_supported
logger = logging.getLogger(__name__)
_is_hip = is_hip()
if _is_hip:
try:
import aiter
from aiter.ops.gemm_op_a4w4 import gemm_a4w4
from aiter.ops.shuffle import shuffle_weight
from aiter.utility.fp4_utils import dynamic_mxfp4_quant
except ImportError as e:
logger.warning(f"aiter MXFP4 kernels not available: {e}")
aiter = None
shuffle_weight = None
dynamic_mxfp4_quant = None
gemm_a4w4 = None
# The gemm_a4w4 ASM kernel has degraded precision when the output
# dimension (N) is smaller than its minimum tile size.
# Layers with output_size falls below this threshold will stay unquantized
_MXFP4_MIN_OUTPUT_DIM = 256
class Mxfp4Config(QuantizationConfig):
"""
MXFP4 quantization config for diffusion models.
Supports online quantization from unquantized BF16/FP16 checkpoints.
Note: MXFP4 requires ROCm and MI350+ (gfx95x).
"""
def __init__(
self,
is_checkpoint_mxfp4_serialized: bool = False,
ignored_layers: Optional[List[str]] = None,
packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
):
super().__init__()
self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized
self.ignored_layers = ignored_layers or []
self.packed_modules_mapping = packed_modules_mapping or {}
@classmethod
def get_name(cls) -> str:
return "mxfp4"
@classmethod
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
return [torch.bfloat16, torch.float16]
@classmethod
def get_min_capability(cls) -> int:
return 95 # gfx95x, Note: mxfp_supported() is a better check
@classmethod
def get_config_filenames(cls) -> list[str]:
return [] # No config file needed for online quantization
@classmethod
def from_config(cls, config: dict) -> "Mxfp4Config":
"""Create from model config (for pre-quantized checkpoints)."""
is_serialized = config.get("quant_method") == "mxfp4"
return cls(is_checkpoint_mxfp4_serialized=is_serialized)
def get_quant_method(self, layer, prefix: str):
from sglang.multimodal_gen.runtime.layers.linear import LinearBase
if isinstance(layer, LinearBase):
if is_layer_skipped(
prefix,
self.ignored_layers,
fused_mapping=self.packed_modules_mapping,
):
logger.debug(
f"MXFP4: Keeping layer {prefix} unquantized (in ignored_layers)"
)
return UnquantizedLinearMethod()
# Skip layers whose output dims are too small, see ASM kernel comment above
output_size = getattr(layer, "output_size", None)
if output_size is not None and output_size < _MXFP4_MIN_OUTPUT_DIM:
logger.info(
f"MXFP4: Keeping layer {prefix} unquantized "
f"(output_size={output_size} < {_MXFP4_MIN_OUTPUT_DIM})"
)
return UnquantizedLinearMethod()
logger.debug(f"MXFP4: Replacing layer {prefix} with MXFP4 linear method")
return Mxfp4LinearMethod(self)
else:
logger.debug(f"MXFP4: Skipping layer {prefix} (not a LinearBase)")
return None
class Mxfp4LinearMethod(LinearMethodBase):
"""
MXFP4 online quantization method for linear layers.
Quantizes unquantized BF16/FP16 weights to MXFP4 format during
process_weights_after_loading().
"""
def __init__(self, quant_config: Mxfp4Config):
self.quant_config = quant_config
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,
**extra_weight_attrs,
):
"""
Creates BF16/FP16 parameters that will be
quantized to MXFP4 in process_weights_after_loading().
"""
output_size_per_partition = sum(output_partition_sizes)
weight_loader = extra_weight_attrs.get("weight_loader")
weight = ModelWeightParameter(
data=torch.empty(
output_size_per_partition,
input_size_per_partition,
dtype=params_dtype,
),
weight_loader=weight_loader,
input_dim=1,
output_dim=0,
)
layer.register_parameter("weight", weight)
# Placeholder scale (will be created during quantization)
weight_scale = PerTensorScaleParameter(
data=torch.empty(1, dtype=torch.float32),
weight_loader=weight_loader,
)
layer.register_parameter("weight_scale", weight_scale)
def process_weights_after_loading(self, layer: torch.nn.Module):
"""
Quantize BF16/FP16 weights to MXFP4 after loading from checkpoint.
Converts weights from unquantized format to:
- Packed uint8 (2 FP4 values per byte)
- E8M0 scales (one per 32-element block)
"""
if not mxfp_supported():
platform = "unknown"
if _is_hip:
try:
platform = torch.cuda.get_device_properties(0).gcnArchName
except:
platform = "ROCm (unknown arch)"
raise RuntimeError(
f"MXFP4 quantization requires ROCm and MI350+ (gfx95x). "
f"Current platform: {platform}."
)
# Check if weights are already quantized
if layer.weight.dtype not in [torch.bfloat16, torch.float16]:
# Already quantized or unexpected dtype
logger.info("Weights are quantized or unexpected dtype")
return
if any(fn is None for fn in (dynamic_mxfp4_quant, shuffle_weight, gemm_a4w4)):
raise RuntimeError(
"aiter MXFP4 kernels not available. "
"Install aiter with MXFP4 support."
)
weight_data = layer.weight.data
was_on_cpu = weight_data.device.type == "cpu"
if was_on_cpu:
weight_data = weight_data.cuda()
w_quant, mx_scales = dynamic_mxfp4_quant(weight_data, shuffle=True)
w_quant_shuffled = shuffle_weight(w_quant)
if was_on_cpu:
w_quant_shuffled = w_quant_shuffled.cpu()
mx_scales = mx_scales.cpu()
layer.weight = Parameter(w_quant_shuffled, requires_grad=False)
layer.weight_scale = Parameter(mx_scales, requires_grad=False)
logger.debug(
f"MXFP4: Quantized layer weights - weight {layer.weight.shape} {layer.weight.dtype}, "
f"scale {layer.weight_scale.shape}"
)
def apply(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
if not mxfp_supported():
raise RuntimeError(
"MXFP4 inference requires ROCm and MI350+ (gfx95x). "
"Current platform not supported."
)
# Handle 3D input tensors [batch, seq, hidden]
original_shape = x.shape
if x.dim() == 3:
x = x.view(-1, x.shape[-1])
x_fp4, x_scale = dynamic_mxfp4_quant(x, shuffle=True)
y = gemm_a4w4(x_fp4, layer.weight, x_scale, layer.weight_scale)
if bias is not None:
y = y + bias
return y.view(*original_shape[:-1], layer.weight.shape[0])
python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
+9 -1
View File
@@ -610,6 +610,7 @@ def load_model_from_full_model_state_dict(
"wcscales",
"wtscale",
"input_scale",
"weight_scale",
"bias",
"norm_q",
"norm_k",
@@ -641,7 +642,14 @@ def load_model_from_full_model_state_dict(
if missing_param_init == "ones" or any(
p in new_param_name
for p in ("wcscales", "wtscale", "input_scale", "norm_q", "norm_k")
for p in (
"wcscales",
"wtscale",
"input_scale",
"weight_scale",
"norm_q",
"norm_k",
)
):
init_like = torch.ones_like
elif missing_param_init == "zeros" or missing_param_init is None:
python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py
+6
View File
@@ -368,6 +368,12 @@ def resolve_transformer_quant_load_spec(
safetensors_list=safetensors_list,
component_model_path=component_model_path,
)
if quant_config is not None:
packed = getattr(model_cls, "packed_modules_mapping", None)
if packed and hasattr(quant_config, "packed_modules_mapping"):
quant_config.packed_modules_mapping = packed
nunchaku_config = server_args.nunchaku_config
# resolve target param dtype
python/sglang/multimodal_gen/runtime/models/dits/zimage.py
+35 -4
View File
@@ -114,13 +114,31 @@ class TimestepEmbedder(nn.Module):
class FeedForward(nn.Module):
def __init__(self, dim: int, hidden_dim: int):
def __init__(
self,
dim: int,
hidden_dim: int,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
):
super().__init__()
# Use MergedColumnParallelLinear for gate and up projection (fused)
self.w13 = MergedColumnParallelLinear(
dim, [hidden_dim, hidden_dim], bias=False, gather_output=False
dim,
[hidden_dim, hidden_dim],
bias=False,
gather_output=False,
quant_config=quant_config,
prefix=f"{prefix}.w13",
)
self.w2 = RowParallelLinear(
hidden_dim,
dim,
bias=False,
input_is_parallel=True,
quant_config=quant_config,
prefix=f"{prefix}.w2",
)
self.w2 = RowParallelLinear(hidden_dim, dim, bias=False, input_is_parallel=True)
self.act = SiluAndMul()
def forward(self, x):
@@ -409,7 +427,12 @@ class ZImageTransformerBlock(nn.Module):
if hasattr(self.feed_forward, "net") and len(self.feed_forward.net) > 2:
self.feed_forward.net[2].act_unsigned = quant_config.act_unsigned
else:
self.feed_forward = FeedForward(dim=dim, hidden_dim=hidden_dim)
self.feed_forward = FeedForward(
dim=dim,
hidden_dim=hidden_dim,
quant_config=quant_config,
prefix=f"{prefix}.feed_forward",
)
self.attention_norm1 = RMSNorm(dim, eps=norm_eps)
self.ffn_norm1 = RMSNorm(dim, eps=norm_eps)
@@ -600,6 +623,14 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
ZImageDitConfig().arch_config.reverse_param_names_mapping
)
# Maps fused runtime layer names to their checkpoint shard names.
# Used by is_layer_skipped() to correctly handle --quantization-ignored-layers
# Only list fusions that are unconditional. Conditional fusions (e.g. to_qkv for
# Nunchaku) are handled by their own quant path.
packed_modules_mapping = {
"w13": ["w1", "w3"],
}
@classmethod
def get_nunchaku_quant_rules(cls) -> dict[str, list[str]]:
return {
python/sglang/multimodal_gen/runtime/server_args.py
+26 -3
View File
@@ -182,6 +182,12 @@ class ServerArgs(DisaggArgsMixin):
# path to pre-quantized transformer weights (single .safetensors or directory).
transformer_weights_path: str | None = None
# Quantization method for online quantization
quantization: str | None = None
# Layer name patterns to skip during online quantization
quantization_ignored_layers: list[str] | None = None
# can restrict layers to adapt, e.g. ["q_proj"]
# Will adapt only q, k, v, o by default.
lora_target_modules: list[str] | None = None
@@ -1144,9 +1150,26 @@ class ServerArgs(DisaggArgsMixin):
parser.add_argument(
"--quantization",
type=str,
default=None,
help='Quantization method override (e.g. "mxfp8", "fp8", "modelslim"). '
"When set, the transformer loader will use this instead of auto-detection.",
default=ServerArgs.quantization,
help=(
"Quantization method for the transformer. If omitted, the method is "
"auto-detected from the checkpoint config or safetensors metadata when "
"possible. Applies to both pre-quantized checkpoints and online "
"quantization. Use this flag to override auto-detection. "
"Options: 'fp8', 'mxfp8', 'mxfp4', 'modelslim'. "
"Note: MXFP4 requires ROCm and MI350+ (gfx95x)."
),
)
parser.add_argument(
"--quantization-ignored-layers",
type=str,
nargs="+",
default=ServerArgs.quantization_ignored_layers,
help=(
"Layer name patterns to keep unquantized during online quantization "
"(fp8/mxfp4). Each pattern is matched against the layer prefix. "
"Example: --quantization-ignored-layers img_mod txt_mod to_out"
),
)
# Nunchaku SVDQuant quantization parameters