From 34c0029f0aff4c3d1c714e7d55b2a522bbc0ff69 Mon Sep 17 00:00:00 2001 From: Colin Z <59755453+ColinZ22@users.noreply.github.com> Date: Wed, 13 May 2026 17:52:01 -0700 Subject: [PATCH] [diffusion] [AMD] feat: support online MXFP4 and fp8 quantization (#21431) Co-authored-by: Bowen Bao Co-authored-by: HAI --- docs/diffusion/api/cli.md | 2 + docs/diffusion/quantization.md | 63 ++++- .../sglang/jit_kernel/flash_attention_v3.py | 36 ++- .../runtime/layers/quantization/__init__.py | 4 +- .../runtime/layers/quantization/fp8.py | 8 +- .../runtime/layers/quantization/mxfp4.py | 237 ++++++++++++++++++ .../runtime/loader/fsdp_load.py | 10 +- .../runtime/loader/transformer_load_utils.py | 6 + .../runtime/models/dits/zimage.py | 39 ++- .../multimodal_gen/runtime/server_args.py | 29 ++- 10 files changed, 417 insertions(+), 17 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py diff --git a/docs/diffusion/api/cli.md b/docs/diffusion/api/cli.md index 6652bf964..6480687bf 100644 --- a/docs/diffusion/api/cli.md +++ b/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. diff --git a/docs/diffusion/quantization.md b/docs/diffusion/quantization.md index ccf3f8112..4f2e988cf 100644 --- a/docs/diffusion/quantization.md +++ b/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 \ @@ -40,14 +42,67 @@ Here, `quant_family` means a checkpoint and loading family with shared CLI 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 | +| quant_family | checkpoint form | canonical CLI | supported models | extra dependency | platform / notes | +|------------------|--------------------------------------------------------------------------------------------|------------------------------------------------------|--------------------------------------------------------------|---------------------------------------|-----------------------------------------------------------------------------------------------------------------------| +| `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 diff --git a/python/sglang/jit_kernel/flash_attention_v3.py b/python/sglang/jit_kernel/flash_attention_v3.py index ae4f05026..78423daaa 100644 --- a/python/sglang/jit_kernel/flash_attention_v3.py +++ b/python/sglang/jit_kernel/flash_attention_v3.py @@ -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"], diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py index ac9340697..43e8ff081 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py +++ b/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, } diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py b/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py index fcde0ab88..0ecdd52da 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py new file mode 100644 index 000000000..a296cfc99 --- /dev/null +++ b/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]) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index e50307922..14e9a64b9 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index 976ca2281..4a683d93d 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index 5c8717249..aa3792a3c 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/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 { diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index d5043fe27..9a57bf57f 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/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, - 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