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

Co-authored-by: Bowen Bao <bowenbao@amd.com>
Co-authored-by: HAI <hixiao@gmail.com>

docs/diffusion/api/cli.md

@@ -101,6 +101,8 @@ For quantized transformer checkpoints, prefer:
 - `--model-path` for the base pipeline
 - `--transformer-path` for a quantized `transformers` transformer component folder
 - `--transformer-weights-path` for a quantized safetensors file, directory, or repo
+- `--quantization` for online quantization (apply quantization to unquantized models at load time, activations are quantized dynamically)
+- `--quantization-ignored-layers` layer name patterns to keep unquantized (e.g. `attention.to_`)
 
 See [Quantization](../quantization.md) for supported quantization families and examples.
 

docs/diffusion/quantization.md

@@ -10,8 +10,10 @@ Use these paths:
 - `--model-path`: the base or original model
 - `--transformer-path`: a quantized transformers-style transformer component directory that already contains its own `config.json`
 - `--transformer-weights-path`: quantized transformer weights provided as a single safetensors file, a sharded safetensors directory, a local path, or a Hugging Face repo ID
+- `--quantization`: apply online quantization to unquantized models at load time (activations are quantized dynamically)
+- `--quantization-ignored-layers` layer name patterns to keep unquantized (e.g. `attention.to_`)
 
-Recommended example:
+Recommended example for pre-quantized checkpoints:
 
 ```bash
 sglang generate \
@@ -41,13 +43,66 @@ usage and loader behavior. It is not just the numeric precision or a kernel
 backend.
 
 | quant_family     | checkpoint form                                                                            | canonical CLI                                        | supported models                                             | extra dependency                      | platform / notes                                                                                                      |
-|-------------------|--------------------------------------------------------------------------------------------|------------------------------------------------------------------------|-----------------------------------------|---------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------|
+|------------------|--------------------------------------------------------------------------------------------|------------------------------------------------------|--------------------------------------------------------------|---------------------------------------|-----------------------------------------------------------------------------------------------------------------------|
-| `fp8`             | Quantized transformer component folder, or safetensors with `quantization_config` metadata | `--transformer-path` or `--transformer-weights-path`                   | ALL                                     | None                                  | Component-folder and single-file flows are both supported                                                                              |
+| `fp8` / `mxfp4` (online quantization)  | Unquantized checkpoint (offline via AMD Quark coming soon) | `--quantization {fp8,mxfp4}` | Z-Image-Turbo (validated), others likely work. More support coming soon. | MXFP4: `aiter` on ROCm | MXFP4 requires ROCm and MI350+ (gfx95x). Weights quantized at load time, activations quantized to `fp8` / `mxfp4` dynamically. |
+| `fp8` (offline quantization) | Quantized transformer component folder, or safetensors with `quantization_config` metadata | `--transformer-path` or `--transformer-weights-path`                   | ALL                                     | None                                  | Component-folder and single-file flows are both supported                                                                              |
 | `modelopt-fp8`    | Converted ModelOpt FP8 transformer directory or repo with `config.json`                    | `--transformer-path`                                                    | FLUX.1, FLUX.2, Wan2.2, HunyuanVideo, Qwen Image, Qwen Image Edit | None                                  | Serialized config stays `quant_method=modelopt` with `quant_algo=FP8`; `dit_layerwise_offload` is supported and `dit_cpu_offload` stays disabled |
 | `modelopt-nvfp4`  | Mixed transformer directory/repo with `config.json`, or raw NVFP4 safetensors export/repo | `--transformer-path` for mixed overrides; `--transformer-weights-path` for raw exports | FLUX.1, FLUX.2, Wan2.2                  | None                                  | Mixed override repos keep the base model separate; raw exports such as `black-forest-labs/FLUX.2-dev-NVFP4` still use the weights-path flow |
 | `nunchaku-svdq`   | Pre-quantized Nunchaku transformer weights, usually named `svdq-{int4\|fp4}_r{rank}-...`   | `--transformer-weights-path`                                           | Model-specific support such as Qwen-Image, FLUX, and Z-Image | `nunchaku`                            | SGLang can infer precision and rank from the filename and supports both `int4` and `nvfp4`                                             |
 | `msmodelslim`     | Pre-quantized msmodelslim transformer weights                                              | `--model-path`                                                         | Wan2.2 family                           | None                                  | Currently only compatible with the Ascend NPU family and supports both `w8a8` and `w4a4`                                               |
 
+## Online Quantization
+
+Online quantization applies quantization to unquantized models at load time. This is useful for when pre-quantized checkpoints are not available.
+
+### FP8 Online Quantization
+
+Apply FP8 quantization to any unquantized model:
+
+```bash
+sglang generate \
+  --model-path Tongyi-MAI/Z-Image-Turbo \
+  --quantization fp8 \
+  --prompt "a beautiful sunset" \
+  --save-output
+```
+
+### MXFP4 Online Quantization
+
+MXFP4 provides aggressive 4-bit compression with online quantization. **Note: Requires ROCm and MI350+ (gfx95x) GPU.**
+
+```bash
+sglang generate \
+  --model-path Tongyi-MAI/Z-Image-Turbo \
+  --quantization mxfp4 \
+  --prompt "a beautiful sunset" \
+  --save-output
+```
+**Note:** Requires `aiter` package with MXFP4 kernel support
+
+### Skipping Layers
+
+By default, online quantization quantizes every linear layer in
+the transformer. However, `--quantization-ignored-layers` can be used to keep specific layers in their original precision:
+
+```bash
+sglang generate \
+  --model-path Tongyi-MAI/Z-Image-Turbo \
+  --quantization fp8 \
+  --quantization-ignored-layers attention.to_ \
+  --prompt "a beautiful sunset" \
+  --save-output
+
+sglang generate \
+  --model-path Tongyi-MAI/Z-Image-Turbo \
+  --quantization mxfp4 \
+  --quantization-ignored-layers attention.to_ \
+  --prompt "a beautiful sunset" \
+  --save-output
+```
+
+Each pattern is matched against the full layer prefix (e.g. `layers.0.attention.to_q`). A layer is skipped and left unquantizd if its prefix contains any of the given patterns.
+
 ## Validated ModelOpt Checkpoints
 
 This section is the canonical support matrix for the nine diffusion ModelOpt

python/sglang/jit_kernel/flash_attention_v3.py

@@ -208,8 +208,38 @@ def flash_attn_varlen_func(
 ):
 
     if not _is_fa3_supported():
-        raise NotImplementedError(
+        # Fall back to flash_attn package (FA2) on platforms without sgl-kernel FA3
-            "flash_attn at sgl-kernel is only supported on sm90 and above"
+        # (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(

python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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,
+            default=ServerArgs.quantization,
-            help='Quantization method override (e.g. "mxfp8", "fp8", "modelslim"). '
+            help=(
-            "When set, the transformer loader will use this instead of auto-detection.",
+                "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