[diffusion] quant: adapt FP8 linear to sgld and support quant in flux (#17023)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/docs/quantization.md

@@ -121,3 +121,44 @@ Choose the appropriate configuration based on your hardware and requirements:
 ### Custom Model Quantization
 
 If you want to quantize your own models, you can use the [DeepCompressor](https://github.com/mit-han-lab/deepcompressor) tool. For detailed instructions, please refer to the Nunchaku official documentation.
+
+## FP8 Quantization
+
+### Usage
+
+#### Option 1: Use Pre-quantized Models (Recommended)
+
+If available, you can directly use pre-quantized FP8 models from Hugging Face or other sources. Simply load them with SGLang:
+
+```bash
+sglang generate \
+  --model-path /path/to/FLUX.1-dev-FP8/ \
+  --prompt "A Logo With Bold Large Text: SGL Diffusion" \
+  --save-output
+```
+
+#### Option 2: Convert Your Own Models
+
+If you need to convert a model to FP8 format, use the provided conversion script:
+
+**Step 1: Convert the Model**
+
+```bash
+# convert transformer to FP8 with block quantization
+python -m sglang.multimodal_gen.tools.convert_hf_to_fp8 \
+  --model-dir /path/to/FLUX.1-dev/transformer \
+  --save-dir /path/to/FLUX.1-dev/transformer-FP8 \
+  --strategy block \
+  --block-size 128 128
+```
+
+**Step 2: Run Inference**
+
+```bash
+sglang generate \
+  --model-path /path/to/FLUX.1-dev/
+  # override transformer component with path to converted model
+  --transformer-path /path/to/FLUX.1-dev/transformer-FP8
+  --prompt "A Logo With Bold Large Text: SGL Diffusion" \
+  --save-output
+```

python/sglang/multimodal_gen/runtime/layers/mlp.py

@@ -2,14 +2,24 @@
 
 # SPDX-License-Identifier: Apache-2.0
 
+from typing import Optional
+
 import torch
 import torch.nn as nn
+from diffusers.models.activations import (
+    GEGLU,
+    GELU,
+    ApproximateGELU,
+    LinearActivation,
+    SwiGLU,
+)
 
 from sglang.multimodal_gen.runtime.layers.activation import get_act_fn
 from sglang.multimodal_gen.runtime.layers.linear import (
     ColumnParallelLinear,
     RowParallelLinear,
 )
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
 
 
 class MLP(nn.Module):
@@ -26,6 +36,7 @@ class MLP(nn.Module):
         act_type: str = "gelu_pytorch_tanh",
         dtype: torch.dtype | None = None,
         prefix: str = "",
+        quant_config: QuantizationConfig = None,
     ):
         super().__init__()
         self.fc_in = ColumnParallelLinear(
@@ -33,6 +44,7 @@ class MLP(nn.Module):
             mlp_hidden_dim,
             bias=True,
             gather_output=False,
+            quant_config=quant_config,
         )
 
         self.act = get_act_fn(act_type)
@@ -43,6 +55,7 @@ class MLP(nn.Module):
             output_dim,
             bias=True,
             input_is_parallel=True,
+            quant_config=quant_config,
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -50,3 +63,56 @@ class MLP(nn.Module):
         x = self.act(x)
         x, _ = self.fc_out(x)
         return x
+
+
+class FeedForward(nn.Module):
+    r"""
+    A feed-forward layer.
+
+    Parameters:
+        dim (`int`): The number of channels in the input.
+        dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
+        mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        mult: int = 4,
+        activation_fn: str = "geglu",
+        inner_dim=None,
+        bias: bool = True,
+    ):
+        super().__init__()
+        if inner_dim is None:
+            inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        if activation_fn == "gelu":
+            act_fn = GELU(dim, inner_dim, bias=bias)
+        if activation_fn == "gelu-approximate":
+            act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias)
+        elif activation_fn == "geglu":
+            act_fn = GEGLU(dim, inner_dim, bias=bias)
+        elif activation_fn == "geglu-approximate":
+            act_fn = ApproximateGELU(dim, inner_dim, bias=bias)
+        elif activation_fn == "swiglu":
+            act_fn = SwiGLU(dim, inner_dim, bias=bias)
+        elif activation_fn == "linear-silu":
+            act_fn = LinearActivation(dim, inner_dim, bias=bias, activation="silu")
+
+        self.net = nn.ModuleList([])
+        # project in
+        self.net.append(act_fn)
+        # dummy dropout layer to match with checkpoints compatible with diffusers
+        self.net.append(nn.Dropout(0.0))
+        # project out
+        self.net.append(nn.Linear(inner_dim, dim_out, bias=bias))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for module in self.net:
+            hidden_states = module(hidden_states)
+        return hidden_states

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

@@ -5,13 +5,16 @@ from typing import Literal, get_args
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
 )
+from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config
 
-QuantizationMethods = Literal[None]
+QuantizationMethods = Literal["fp8"]
 
 QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
 
 # The customized quantization methods which will be added to this dict.
-_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {}
+_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {
+    "fp8": Fp8Config,
+}
 
 
 def register_quantization_config(quantization: str):

python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py

@@ -22,6 +22,7 @@ def is_nunchaku_available() -> bool:
     try:
         import nunchaku  # noqa
 
+        logger.debug("Nunchaku package detected")
         return True
     except Exception:
         return False

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

@@ -0,0 +1,498 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+import torch
+from torch.nn import Module
+from torch.nn.parameter import Parameter
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_tensor_model_parallel_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.multimodal_gen.runtime.models.parameter import (
+    BlockQuantScaleParameter,
+    ModelWeightParameter,
+    PerTensorScaleParameter,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.common import (
+    cpu_has_amx_support,
+    get_bool_env_var,
+    use_intel_amx_backend,
+)
+from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading
+from sglang.srt.layers.quantization.fp8_kernel import (
+    is_fp8_fnuz,
+    per_token_group_quant_fp8,
+)
+from sglang.srt.layers.quantization.fp8_utils import (
+    apply_fp8_linear,
+    can_auto_enable_marlin_fp8,
+    cutlass_fp8_supported,
+    dispatch_w8a8_block_fp8_linear,
+    input_to_float8,
+    normalize_e4m3fn_to_e4m3fnuz,
+    requant_weight_ue8m0_inplace,
+)
+from sglang.srt.layers.quantization.marlin_utils_fp8 import (
+    apply_fp8_marlin_linear,
+    prepare_fp8_layer_for_marlin,
+)
+from sglang.srt.layers.quantization.utils import (
+    convert_to_channelwise,
+    is_layer_skipped,
+    requantize_with_max_scale,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config
+
+_is_hip = current_platform.is_hip()
+_is_cuda = current_platform.is_cuda()
+_is_npu = current_platform.is_npu()
+_is_cpu_amx_available = cpu_has_amx_support()
+_is_cpu = current_platform.is_cpu()
+_is_fp8_fnuz = is_fp8_fnuz()
+_use_hip_int4 = get_bool_env_var("SGLANG_INT4_WEIGHT") and _is_hip
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
+
+if _use_aiter or _use_hip_int4:
+    pass
+
+
+ACTIVATION_SCHEMES = ["static", "dynamic"]
+
+logger = logging.getLogger(__name__)
+
+
+class Fp8Config(QuantizationConfig):
+    """Config class for FP8."""
+
+    def __init__(
+        self,
+        is_checkpoint_fp8_serialized: bool = False,
+        activation_scheme: str = "dynamic",
+        ignored_layers: Optional[List[str]] = None,
+        weight_block_size: List[int] = None,
+    ) -> None:
+        self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
+        if is_checkpoint_fp8_serialized:
+            logger.info("Detected fp8 checkpoint.")
+        if activation_scheme not in ACTIVATION_SCHEMES:
+            raise ValueError(f"Unsupported activation scheme {activation_scheme}")
+        self.activation_scheme = activation_scheme
+        self.ignored_layers = ignored_layers or []
+        if weight_block_size is not None:
+            if not is_checkpoint_fp8_serialized:
+                raise ValueError(
+                    f"The block-wise quantization only supports fp8-serialized checkpoint for now."
+                )
+            if len(weight_block_size) != 2:
+                raise ValueError(
+                    f"The quantization block size of weight must have 2 dimensions, but got {len(weight_block_size)} dimensions."
+                )
+            if activation_scheme != "dynamic":
+                raise ValueError(
+                    f"The block-wise quantization only supports dynamic activation scheme for now, but got {activation_scheme} activation scheme."
+                )
+        self.weight_block_size = weight_block_size
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "fp8"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.bfloat16, torch.half]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 80
+
+    @classmethod
+    def get_config_filenames(cls) -> List[str]:
+        return []
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> Fp8Config:
+        quant_method = cls.get_from_keys(config, ["quant_method"])
+        is_checkpoint_fp8_serialized = "fp8" in quant_method
+        activation_scheme = cls.get_from_keys(config, ["activation_scheme"])
+        ignored_layers = cls.get_from_keys_or(
+            config, ["ignored_layers", "modules_to_not_convert"], None
+        )
+        if ignored_layers:
+            # hacking 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(
+            is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
+            activation_scheme=activation_scheme,
+            ignored_layers=ignored_layers,
+            weight_block_size=weight_block_size,
+        )
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[QuantizeMethodBase]:
+        from sglang.multimodal_gen.runtime.layers.linear import LinearBase
+
+        if isinstance(layer, LinearBase):
+            if is_layer_skipped(prefix, self.ignored_layers):
+                return UnquantizedLinearMethod()
+            return Fp8LinearMethod(self)
+        return None
+
+    def get_scaled_act_names(self) -> List[str]:
+        return []
+
+
+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.
+
+    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)
+
+    Args:
+        quant_config: The quantization config.
+    """
+
+    def __init__(self, quant_config: Union[Fp8Config, W4AFp8Config]):
+        self.quant_config = quant_config
+        self.cutlass_fp8_supported = cutlass_fp8_supported()
+
+        # For GPUs that lack FP8 hardware support, we can leverage the Marlin
+        # kernel for fast weight-only FP8 quantization
+        self.use_marlin = False
+        if _is_cuda:
+            force_marlin = get_bool_env_var("SGLANG_FORCE_FP8_MARLIN")
+            auto_enable = can_auto_enable_marlin_fp8()
+            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()
+
+    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,
+    ):
+        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],
+            )
+            # Required by row parallel
+            if tp_size > 1 and input_size // input_size_per_partition == tp_size:
+                if input_size_per_partition % block_k != 0:
+                    raise ValueError(
+                        f"Weight input_size_per_partition = "
+                        f"{input_size_per_partition} is not divisible by "
+                        f"weight quantization block_k = {block_k}."
+                    )
+            # Required by column parallel or enabling merged weights
+            if (
+                tp_size > 1 and output_size // output_size_per_partition == tp_size
+            ) or len(output_partition_sizes) > 1:
+                for output_partition_size in output_partition_sizes:
+                    if output_partition_size % block_n != 0:
+                        raise ValueError(
+                            f"Weight output_partition_size = "
+                            f"{output_partition_size} is not divisible by "
+                            f"weight quantization block_n = {block_n}."
+                        )
+
+        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
+        weight_dtype = (
+            torch.float8_e4m3fn
+            if self.quant_config.is_checkpoint_fp8_serialized
+            else params_dtype
+        )
+
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition, input_size_per_partition, dtype=weight_dtype
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        # If checkpoint is serialized fp8, load them.
+        # Otherwise, wait until process_weights_after_loading.
+        if self.quant_config.is_checkpoint_fp8_serialized:
+            # WEIGHT SCALE
+            if self.block_quant:
+                if hasattr(self.quant_config, "activation_scheme"):
+                    assert self.quant_config.activation_scheme == "dynamic"
+                elif hasattr(self.quant_config, "linear_activation_scheme"):
+                    assert self.quant_config.linear_activation_scheme == "dynamic"
+                scale = BlockQuantScaleParameter(
+                    data=torch.empty(
+                        (output_size_per_partition + block_n - 1) // block_n,
+                        (input_size_per_partition + block_k - 1) // block_k,
+                        dtype=torch.float32,
+                    ),
+                    input_dim=1,
+                    output_dim=0,
+                    weight_loader=weight_loader,
+                )
+                scale.format_ue8m0 = False
+                scale[:] = torch.finfo(torch.float32).min
+                layer.register_parameter("weight_scale_inv", scale)
+            else:
+                scale = PerTensorScaleParameter(
+                    data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
+                    weight_loader=weight_loader,
+                )
+                scale[:] = torch.finfo(torch.float32).min
+                layer.register_parameter("weight_scale", scale)
+
+            # INPUT ACTIVATION SCALE
+            if (
+                hasattr(self.quant_config, "activation_scheme")
+                and self.quant_config.activation_scheme == "static"
+            ) or (
+                hasattr(self.quant_config, "linear_activation_scheme")
+                and self.quant_config.linear_activation_scheme == "static"
+            ):
+                scale = PerTensorScaleParameter(
+                    data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
+                    weight_loader=weight_loader,
+                )
+
+                scale[:] = torch.finfo(torch.float32).min
+                layer.register_parameter("input_scale", scale)
+            else:
+                layer.register_parameter("input_scale", None)
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if self.block_quant:
+            # If ROCm, normalize the weights and scales to e4m3fnuz
+            if _is_fp8_fnuz:
+                # activation_scheme: dynamic
+                weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
+                    weight=layer.weight,
+                    weight_scale=layer.weight_scale_inv,
+                    input_scale=None,
+                )
+                layer.input_scale = None
+            elif _is_cpu:
+                assert (
+                    _is_cpu_amx_available
+                ), "Fp8LinearMethod on CPU requires that CPU has AMX support"
+                _amx_process_weight_after_loading(layer, ["weight"])
+                layer.weight_scale_inv = torch.nn.Parameter(
+                    layer.weight_scale_inv.data, requires_grad=False
+                )
+                return
+            else:
+                # For fp8 linear weights run with deepgemm, the weights and scales need be requantized to ue8m0
+                from sglang.srt.layers.quantization.fp8_utils import (
+                    deepgemm_w8a8_block_fp8_linear_with_fallback,
+                )
+                from sglang.srt.model_loader.utils import (
+                    should_deepgemm_weight_requant_ue8m0,
+                )
+
+                if (
+                    should_deepgemm_weight_requant_ue8m0(
+                        weight_block_size=getattr(
+                            self.quant_config, "weight_block_size", None
+                        ),
+                    )
+                    and (
+                        self.w8a8_block_fp8_linear
+                        is deepgemm_w8a8_block_fp8_linear_with_fallback
+                    )
+                    and (not layer.weight_scale_inv.format_ue8m0)
+                ):
+                    requant_weight_ue8m0_inplace(
+                        layer.weight,
+                        layer.weight_scale_inv,
+                        self.quant_config.weight_block_size,
+                    )
+                    layer.weight_scale_inv.format_ue8m0 = True
+                weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data
+
+            layer.weight.data = weight.data
+            layer.weight_scale_inv.data = weight_scale.data
+        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 self.cutlass_fp8_supported or self.use_marlin:
+                    # apply per-channel quantization default as
+                    # cutlass sgl-kernel and marlin only support per-channel scale
+                    qweight, weight_scale = per_token_group_quant_fp8(
+                        layer.weight, layer.weight.shape[-1]
+                    )
+                    weight_scale = weight_scale.t().contiguous()
+                else:
+                    # per-tensor quantization
+                    qweight, weight_scale = input_to_float8(layer.weight)
+
+                # Update the layer with the new values.
+                layer.weight = Parameter(qweight.t(), requires_grad=False)
+                layer.weight_scale = Parameter(weight_scale, requires_grad=False)
+                layer.input_scale = None
+
+            # If checkpoint is fp8, handle that there are N scales for N
+            # shards in a fused module
+            else:
+                layer.weight_scale = Parameter(
+                    layer.weight_scale.data, requires_grad=False
+                )
+                if (
+                    hasattr(self.quant_config, "activation_scheme")
+                    and self.quant_config.activation_scheme == "static"
+                ) or (
+                    hasattr(self.quant_config, "linear_activation_scheme")
+                    and self.quant_config.linear_activation_scheme == "static"
+                ):
+                    layer.input_scale = Parameter(
+                        layer.input_scale.data, requires_grad=False
+                    )
+
+                # cutlass sgl-kernel and marlin only support per-channel scale
+                if self.cutlass_fp8_supported or self.use_marlin:
+                    weight = layer.weight
+                    weight_scale = convert_to_channelwise(
+                        layer.weight_scale, layer.logical_widths
+                    )
+                else:
+                    # Dequant -> Quant with max scale so we can run per tensor.
+                    weight = layer.weight
+                    weight_scale = layer.weight_scale
+                    # If ROCm, normalize the weights and scales to e4m3fnuz
+                    if _is_fp8_fnuz:
+                        weight, weight_scale, input_scale = (
+                            normalize_e4m3fn_to_e4m3fnuz(
+                                weight=weight,
+                                weight_scale=weight_scale,
+                                input_scale=layer.input_scale,
+                            )
+                        )
+                        if input_scale is not None:
+                            layer.input_scale = Parameter(
+                                input_scale, requires_grad=False
+                            )
+
+                    weight_scale, weight = requantize_with_max_scale(
+                        weight=weight,
+                        weight_scale=weight_scale,
+                        logical_widths=layer.logical_widths,
+                    )
+
+                # Update layer with new values.
+                layer.weight = Parameter(weight.t(), requires_grad=False)
+                layer.weight_scale = Parameter(weight_scale, requires_grad=False)
+                if (
+                    hasattr(self.quant_config, "activation_scheme")
+                    and self.quant_config.activation_scheme == "static"
+                ) or (
+                    hasattr(self.quant_config, "linear_activation_scheme")
+                    and self.quant_config.linear_activation_scheme == "static"
+                ):
+                    layer.input_scale = Parameter(
+                        layer.input_scale.max(), requires_grad=False
+                    )
+
+        if self.use_marlin:
+            if self.block_quant:
+                layer.weight_block_size = self.quant_config.weight_block_size
+            prepare_fp8_layer_for_marlin(layer, not self.block_quant)
+            # Activations not quantized for marlin.
+            del layer.input_scale
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.use_marlin:
+            return apply_fp8_marlin_linear(
+                input=x,
+                weight=layer.weight,
+                weight_scale=layer.weight_scale,
+                workspace=layer.workspace,
+                size_n=layer.output_size_per_partition,
+                size_k=layer.input_size_per_partition,
+                bias=bias,
+            )
+
+        if self.block_quant:
+            if use_intel_amx_backend(layer):
+                return torch.ops.sgl_kernel.fp8_scaled_mm_cpu(
+                    x,
+                    layer.weight,
+                    layer.weight_scale_inv,
+                    self.quant_config.weight_block_size,
+                    bias,
+                    x.dtype,
+                    True,  # is_vnni
+                )
+
+            if isinstance(x, tuple):
+                return self.w8a8_block_fp8_linear(
+                    input=x[0],
+                    weight=layer.weight,
+                    block_size=self.quant_config.weight_block_size,
+                    weight_scale=layer.weight_scale_inv,
+                    input_scale=x[1],
+                    bias=bias,
+                )
+
+            return self.w8a8_block_fp8_linear(
+                input=x,
+                weight=layer.weight,
+                block_size=self.quant_config.weight_block_size,
+                weight_scale=layer.weight_scale_inv,
+                input_scale=None,
+                bias=bias,
+            )
+
+        return apply_fp8_linear(
+            input=x,
+            weight=layer.weight,
+            weight_scale=layer.weight_scale,
+            input_scale=layer.input_scale,
+            bias=bias,
+            cutlass_fp8_supported=self.cutlass_fp8_supported,
+            use_per_token_if_dynamic=False,
+        )

python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py

@@ -24,6 +24,7 @@ from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
     get_diffusers_component_config,
     get_metadata_from_safetensors_file,
+    get_quant_config,
 )
 from sglang.multimodal_gen.runtime.utils.logging_utils import get_log_level, init_logger
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
@@ -70,6 +71,7 @@ class TransformerLoader(ComponentLoader):
     ):
         """Load the transformer based on the model path, and inference args."""
         config = get_diffusers_component_config(component_path=component_model_path)
+
         hf_config = deepcopy(config)
         cls_name = config.pop("_class_name")
         if cls_name is None:
@@ -94,6 +96,7 @@ class TransformerLoader(ComponentLoader):
         model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
 
         nunchaku_config = server_args.nunchaku_config
+
         if nunchaku_config is not None:
             nunchaku_config.model_cls = model_cls
 
@@ -120,7 +123,7 @@ class TransformerLoader(ComponentLoader):
         )
 
         logger.info(
-            "Loading %s from %s safetensors files %s, param_dtype: %s",
+            "Loading %s from %s safetensors file(s) %s, param_dtype: %s",
             cls_name,
             len(safetensors_list),
             f": {safetensors_list}" if get_log_level() == logging.DEBUG else "",
@@ -128,11 +131,12 @@ class TransformerLoader(ComponentLoader):
         )
 
         init_params: dict[str, Any] = {"config": dit_config, "hf_config": hf_config}
-        if (
-            nunchaku_config is not None
-            and "quant_config" in inspect.signature(model_cls.__init__).parameters
-        ):
-            init_params["quant_config"] = nunchaku_config
+
+        if "quant_config" in inspect.signature(model_cls.__init__).parameters:
+            quant_config = get_quant_config(config)
+            init_params["quant_config"] = (
+                quant_config if quant_config else nunchaku_config
+            )
 
         # Load the model using FSDP loader
         model = maybe_load_fsdp_model(

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -257,6 +257,7 @@ def load_model_from_full_model_state_dict(
     for target_param_name in sorted_param_names:
         full_tensor = custom_param_sd[target_param_name]
         meta_sharded_param = meta_sd.get(target_param_name)
+
         if meta_sharded_param is None:
             # For FSDP models, ensure all ranks process parameters consistently
             if strict or is_fsdp_model:
@@ -264,12 +265,16 @@ def load_model_from_full_model_state_dict(
                     f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect."
                 )
             else:
-                logger.warning(
-                    f"Parameter '{target_param_name}' from checkpoint not found in model; skipping. This is expected for optional parameters."
-                )
                 continue
 
-        target_dtype = param_dtype if param_dtype else full_tensor.dtype
+        # use meta param dtype so quantized params (e.g. FP8) keep their dtype;
+        # for non-quantized models meta dtype equals param_dtype anyway
+        if meta_sharded_param is None:
+            # for nunchaku, some scales are patched later
+            target_dtype = full_tensor.dtype
+        else:
+            target_dtype = meta_sharded_param.dtype
+
         if not hasattr(meta_sharded_param, "device_mesh"):
             full_tensor = full_tensor.to(device=device, dtype=target_dtype)
             actual_param = param_dict.get(target_param_name)
@@ -350,6 +355,7 @@ def load_model_from_full_model_state_dict(
             )
 
         meta_sharded_param = meta_sd.get(new_param_name)
+        meta_sharded_param_dtype = meta_sharded_param.dtype
 
         if "wcscales" in new_param_name or "wtscale" in new_param_name:
             init_like = torch.ones_like
@@ -358,13 +364,13 @@ def load_model_from_full_model_state_dict(
 
         if not hasattr(meta_sharded_param, "device_mesh"):
             sharded_tensor = init_like(
-                meta_sharded_param, device=device, dtype=param_dtype
+                meta_sharded_param, device=device, dtype=meta_sharded_param_dtype
             )
             if cpu_offload and not is_fsdp_model:
                 sharded_tensor = sharded_tensor.cpu()
         else:
             full_tensor = init_like(
-                meta_sharded_param, device=device, dtype=param_dtype
+                meta_sharded_param, device=device, dtype=meta_sharded_param_dtype
             )
             sharded_tensor = distribute_tensor(
                 full_tensor,

python/sglang/multimodal_gen/runtime/models/dits/flux.py

@@ -18,7 +18,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
-from diffusers.models.attention import AttentionModuleMixin, FeedForward
+from diffusers.models.attention import AttentionModuleMixin
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import (
     AdaLayerNormContinuous,
@@ -29,13 +29,12 @@ from torch.nn import LayerNorm as LayerNorm
 
 from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
 from sglang.multimodal_gen.runtime.layers.attention import USPAttention
-
-# from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm as LayerNorm
 from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm
 from sglang.multimodal_gen.runtime.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
 )
+from sglang.multimodal_gen.runtime.layers.mlp import FeedForward
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
 )
@@ -255,13 +254,25 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
             )
         else:
             self.to_q = ColumnParallelLinear(
-                query_dim, self.inner_dim, bias=bias, gather_output=True
+                query_dim,
+                self.inner_dim,
+                bias=bias,
+                gather_output=True,
+                quant_config=quant_config,
             )
             self.to_k = ColumnParallelLinear(
-                query_dim, self.inner_dim, bias=bias, gather_output=True
+                query_dim,
+                self.inner_dim,
+                bias=bias,
+                gather_output=True,
+                quant_config=quant_config,
             )
             self.to_v = ColumnParallelLinear(
-                query_dim, self.inner_dim, bias=bias, gather_output=True
+                query_dim,
+                self.inner_dim,
+                bias=bias,
+                gather_output=True,
+                quant_config=quant_config,
             )
         if not self.pre_only:
             self.to_out = torch.nn.ModuleList([])
@@ -296,18 +307,21 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
                     self.inner_dim,
                     bias=added_proj_bias,
                     gather_output=True,
+                    quant_config=quant_config,
                 )
                 self.add_k_proj = ColumnParallelLinear(
                     added_kv_proj_dim,
                     self.inner_dim,
                     bias=added_proj_bias,
                     gather_output=True,
+                    quant_config=quant_config,
                 )
                 self.add_v_proj = ColumnParallelLinear(
                     added_kv_proj_dim,
                     self.inner_dim,
                     bias=added_proj_bias,
                     gather_output=True,
+                    quant_config=quant_config,
                 )
             self.to_add_out = ColumnParallelLinear(
                 self.inner_dim,
@@ -453,7 +467,7 @@ class FluxSingleTransformerBlock(nn.Module):
                 quant_config=quant_config,
                 prefix=f"{prefix}.attn" if prefix else "attn",
             )
-            if NunchakuAdaLayerNormZeroSingle is not None:
+            if is_nunchaku_available():
                 self.norm = NunchakuAdaLayerNormZeroSingle(self.norm, scale_shift=0)
         else:
             self.proj_mlp = ColumnParallelLinear(
@@ -461,6 +475,7 @@ class FluxSingleTransformerBlock(nn.Module):
                 self.mlp_hidden_dim,
                 bias=True,
                 gather_output=True,
+                quant_config=quant_config,
             )
             self.act_mlp = nn.GELU(approximate="tanh")
             self.proj_out = ColumnParallelLinear(
@@ -468,6 +483,7 @@ class FluxSingleTransformerBlock(nn.Module):
                 dim,
                 bias=True,
                 gather_output=True,
+                quant_config=quant_config,
             )
             self.attn = FluxAttention(
                 query_dim=dim,
@@ -477,6 +493,7 @@ class FluxSingleTransformerBlock(nn.Module):
                 bias=True,
                 eps=1e-6,
                 pre_only=True,
+                quant_config=quant_config,
             )
 
     def forward(
@@ -579,12 +596,13 @@ class FluxTransformerBlock(nn.Module):
             and hasattr(quant_config, "get_name")
             and quant_config.get_name() == "svdquant"
             and is_nunchaku_available()
-            and NunchakuFeedForward is not None
         )
         self.use_nunchaku_structure = nunchaku_enabled
         self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
         self.ff_context = FeedForward(
-            dim=dim, dim_out=dim, activation_fn="gelu-approximate"
+            dim=dim,
+            dim_out=dim,
+            activation_fn="gelu-approximate",
         )
         if nunchaku_enabled:
             nunchaku_kwargs = {
@@ -594,11 +612,10 @@ class FluxTransformerBlock(nn.Module):
             }
             self.ff = NunchakuFeedForward(self.ff, **nunchaku_kwargs)
             self.ff_context = NunchakuFeedForward(self.ff_context, **nunchaku_kwargs)
-            if NunchakuAdaLayerNormZero is not None:
-                self.norm1 = NunchakuAdaLayerNormZero(self.norm1, scale_shift=0)
-                self.norm1_context = NunchakuAdaLayerNormZero(
-                    self.norm1_context, scale_shift=0
-                )
+            self.norm1 = NunchakuAdaLayerNormZero(self.norm1, scale_shift=0)
+            self.norm1_context = NunchakuAdaLayerNormZero(
+                self.norm1_context, scale_shift=0
+            )
 
     def forward(
         self,

python/sglang/multimodal_gen/runtime/utils/common.py

@@ -281,3 +281,28 @@ def get_bool_env_var(name: str, default: str = "false") -> bool:
         _warned_bool_env_var_keys.add(value)
 
     return value in truthy_values
+
+
+try:
+    import sgl_kernel  # noqa: F401
+
+    is_intel_amx_backend_available = hasattr(
+        torch.ops.sgl_kernel, "convert_weight_packed"
+    )
+except:
+    is_intel_amx_backend_available = False
+
+try:
+    # move torch._C._cpu._is_amx_tile_supported() from cpu_has_amx_support
+    # to support torch compile
+    is_amx_tile_supported = torch._C._cpu._is_amx_tile_supported()
+except:
+    is_amx_tile_supported = False
+
+
+def cpu_has_amx_support():
+    return is_amx_tile_supported and is_intel_amx_backend_available
+
+
+def use_intel_amx_backend(layer):
+    return getattr(layer, "use_intel_amx_backend", False)

python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py

@@ -26,7 +26,7 @@ import shutil
 import time
 from functools import reduce
 from pathlib import Path
-from typing import Any, Optional, Union, cast
+from typing import Any, Dict, List, Optional, Union, cast
 
 from diffusers.loaders.lora_base import (
     _best_guess_weight_name,  # watch out for potetential removal from diffusers
@@ -42,6 +42,10 @@ from safetensors import safe_open
 from transformers import AutoConfig, PretrainedConfig
 from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
 
+from sglang.multimodal_gen.runtime.layers.quantization import (
+    QuantizationConfig,
+    get_quantization_config,
+)
 from sglang.multimodal_gen.runtime.loader.weight_utils import get_lock
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -320,6 +324,80 @@ def get_diffusers_component_config(
     return combined_config
 
 
+def replace_prefix(key: str, prefix_mapping: dict[str, str]) -> str:
+    for prefix, new_prefix in prefix_mapping.items():
+        if key.startswith(prefix):
+            key = key.replace(prefix, new_prefix, 1)
+    return key
+
+
+def get_quant_config(
+    model_config,
+    packed_modules_mapping: Dict[str, List[str]] = {},
+    remap_prefix: Dict[str, str] | None = None,
+) -> QuantizationConfig:
+    if "quantization_config" not in model_config:
+        return None
+    quant_cls = get_quantization_config(
+        model_config["quantization_config"]["quant_method"]
+    )
+
+    # GGUF doesn't have config file
+    if model_config["quantization_config"]["quant_method"] == "gguf":
+        return quant_cls.from_config({})
+
+    # Read the quantization config from the HF model config, if available.
+    hf_quant_config = model_config["quantization_config"]
+    # some vision model may keep quantization_config in their text_config
+    hf_text_config = getattr(model_config, "text_config", None)
+    if hf_quant_config is None and hf_text_config is not None:
+        hf_quant_config = getattr(hf_text_config, "quantization_config", None)
+    if hf_quant_config is None:
+        # compressed-tensors uses a compressions_config
+        hf_quant_config = getattr(model_config, "compression_config", None)
+    if hf_quant_config is not None:
+        hf_quant_config["packed_modules_mapping"] = packed_modules_mapping
+        return quant_cls.from_config(hf_quant_config)
+    # In case of bitsandbytes/QLoRA, get quant config from the adapter model.
+    else:
+        model_name_or_path = model_config["model_path"]
+    is_local = os.path.isdir(model_name_or_path)
+    hf_folder = model_name_or_path
+
+    possible_config_filenames = quant_cls.get_config_filenames()
+
+    # If the quantization config is not found, use the default config.
+    if not possible_config_filenames:
+        return quant_cls()
+
+    config_files = glob.glob(os.path.join(hf_folder, "*.json"))
+
+    quant_config_files = [
+        f for f in config_files if any(f.endswith(x) for x in possible_config_filenames)
+    ]
+    if len(quant_config_files) == 0:
+        raise ValueError(
+            f"Cannot find the config file for {model_config['quantization_config']['quant_method']}"
+        )
+    if len(quant_config_files) > 1:
+        raise ValueError(
+            f"Found multiple config files for {model_config['quantization_config']['quant_method']}: "
+            f"{quant_config_files}"
+        )
+
+    quant_config_file = quant_config_files[0]
+    with open(quant_config_file) as f:
+        config = json.load(f)
+        if remap_prefix is not None:
+            exclude_modules = [
+                replace_prefix(key, remap_prefix)
+                for key in config["quantization"]["exclude_modules"]
+            ]
+            config["quantization"]["exclude_modules"] = exclude_modules
+        config["packed_modules_mapping"] = packed_modules_mapping
+        return quant_cls.from_config(config)
+
+
 # Models don't use the same configuration key for determining the maximum
 # context length.  Store them here so we can sanely check them.
 # NOTE: The ordering here is important. Some models have two of these and we

python/sglang/multimodal_gen/tools/convert_hf_to_fp8.py

@@ -0,0 +1,310 @@
+# copied and adapted from Slime
+"""
+Convert HuggingFace safetensors model to FP8 format for efficient inference.
+
+Example usage:
+    # convert FLUX.1-dev transformer to FP8
+    python -m sglang.multimodal_gen.tools.convert_hf_to_fp8 \
+        --model-dir /path/to/FLUX.1-dev/transformer \
+        --save-dir /path/to/FLUX.1-dev/transformer-FP8 \
+        --strategy block \
+        --block-size 128 128
+
+Options:
+    --model-dir MODEL_DIR
+                        path to the directory of the HF safetensors model (e.g., transformer subfolder)
+    --save-dir SAVE_DIR
+                        path to the directory to save the converted FP8 model
+    --strategy {block,channel,tensor}
+                        quantization strategy (default: block)
+    --block-size [BLOCK_SIZE ...]
+                        block size for block quantization, e.g., --block-size 128 128
+    --max-workers MAX_WORKERS
+                        number of worker threads for parallel processing (default: 1)
+"""
+
+import argparse
+import gc
+import json
+import os
+import shutil
+import threading
+from concurrent.futures import ThreadPoolExecutor
+
+import safetensors
+import safetensors.torch
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+FP8_INFO = torch.finfo(torch.float8_e4m3fn)
+FP8_MAX, FP8_MIN = FP8_INFO.max, FP8_INFO.min
+
+
+def ceildiv(a, b):
+    return -(-a // b)
+
+
+def block_fp8(weight, block_size):
+
+    # per block quant
+    block_n, block_k = block_size[0], block_size[1]
+
+    shape_0, shape_1 = weight.shape
+
+    n_tiles = ceildiv(shape_0, block_n)
+    k_tiles = ceildiv(shape_1, block_k)
+
+    q_weight = F.pad(
+        weight,
+        (0, k_tiles * block_k - shape_1, 0, n_tiles * block_n - shape_0),
+        mode="constant",
+        value=0.0,
+    )
+
+    qweight = q_weight.reshape(n_tiles, block_n, k_tiles, block_k)
+    block_max = torch.max(torch.abs(qweight), dim=1, keepdim=True)[0]
+    block_max = torch.max(block_max, dim=3, keepdim=True)[0]
+
+    scale = block_max.to(torch.float32) / FP8_MAX
+    qweight = (
+        (qweight / scale)
+        .clamp(min=FP8_MIN, max=FP8_MAX)
+        .reshape((n_tiles * block_n, k_tiles * block_k))
+        .to(torch.float8_e4m3fn)
+    )
+    qweight = qweight[:shape_0, :shape_1].clone().detach()
+    scale = scale.squeeze()
+
+    return qweight, scale
+
+
+def channel_fp8(weight):
+    channel_max = torch.max(weight.abs(), dim=-1, keepdim=True)[0]
+    scale = channel_max.clamp(min=1e-12).to(torch.float32) / FP8_MAX
+    qweight = (weight / scale).clamp(min=FP8_MIN, max=FP8_MAX)
+    qweight = qweight.to(torch.float8_e4m3fn)
+    return qweight, scale
+
+
+def tensor_fp8(weight):
+    scale = weight.abs().max().clamp(min=1e-12).to(torch.float32) / FP8_MAX
+    qweight = (weight / scale).clamp(min=FP8_MIN, max=FP8_MAX)
+    qweight = qweight.to(torch.float8_e4m3fn)
+    scale = scale.view(1)
+    return qweight, scale
+
+
+def quant_fp8(weight, strategy, block_size=None):
+    if strategy == "tensor":
+        return tensor_fp8(weight)
+    elif strategy == "channel":
+        return channel_fp8(weight)
+    else:
+        return block_fp8(weight, block_size)
+
+
+class ConversionResult:
+    def __init__(self):
+        self.lock = threading.Lock()
+        self.weight_map = {}
+        self.param_count = 0
+        self.modules_to_not_convert = []
+
+    def add_result(self, filename, q_weights, module_names):
+        with self.lock:
+            for k, v in q_weights.items():
+                self.weight_map[k] = filename
+                self.param_count += len(v)
+            self.modules_to_not_convert.extend(module_names)
+
+
+def process_file(
+    input_path, output_path, filename, strategy, block_size, result_collector
+):
+    if not filename.endswith(".safetensors"):
+        return
+
+    print(f"Processing {filename}, memory usage: {torch.cuda.memory_allocated()}")
+    weights = {}
+    q_weights = {}
+
+    with safetensors.safe_open(
+        os.path.join(input_path, filename), framework="pt", device="cuda"
+    ) as f:
+        for k in f.keys():
+            weights[k] = f.get_tensor(k)
+
+    modules_to_not_convert = []
+    for key in weights.keys():
+        if (
+            "weight" in key
+            and "layernorm" not in key
+            and "embed" not in key
+            and "router" not in key
+            and "mlp.gate." not in key
+            and "norm" not in key
+            and "lm_head" not in key
+            and "eh_proj" not in key
+            and "net" not in key
+            and "proj_out.weight" != key
+        ):
+            qw, s = quant_fp8(weights[key], strategy, block_size)
+            q_weights[key] = qw
+            if block_size:
+                scale_name = key.replace(".weight", ".weight_scale_inv")
+            else:
+                scale_name = key.replace(".weight", ".weight_scale")
+            q_weights[scale_name] = s
+        else:
+            modules_to_not_convert.append(key.replace(".weight", ""))
+            q_weights[key] = weights[key]
+
+    safetensors.torch.save_file(
+        q_weights, os.path.join(output_path, filename), metadata={"format": "pt"}
+    )
+
+    result_collector.add_result(filename, q_weights, modules_to_not_convert)
+
+
+def convert_fp8(input_path, output_path, strategy, block_size=None, max_workers=4):
+    input_path = os.path.abspath(input_path)
+    os.makedirs(output_path, exist_ok=True)
+
+    for filename in os.listdir(input_path):
+        if not filename.endswith(".safetensors") and not os.path.isdir(
+            os.path.join(input_path, filename)
+        ):
+            shutil.copyfile(
+                os.path.join(input_path, filename), os.path.join(output_path, filename)
+            )
+
+    safetensors_files = [
+        f for f in os.listdir(input_path) if f.endswith(".safetensors")
+    ]
+
+    result_collector = ConversionResult()
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = []
+        for filename in safetensors_files:
+            future = executor.submit(
+                process_file,
+                input_path,
+                output_path,
+                filename,
+                strategy,
+                block_size,
+                result_collector,
+            )
+            futures.append(future)
+
+        for future in tqdm(futures, desc="Processing files"):
+            future.result()
+
+    if strategy == "block" or strategy == "tensor":
+        quantization_config = {
+            "activation_scheme": "dynamic",
+            "fmt": "e4m3",
+            "quant_method": "fp8",
+        }
+        if block_size:
+            quantization_config["weight_block_size"] = block_size
+        if len(result_collector.modules_to_not_convert) > 0:
+            quantization_config["modules_to_not_convert"] = list(
+                set(result_collector.modules_to_not_convert)
+            )
+    else:
+        quant_group = {
+            "group_0": {
+                "input_activations": {
+                    "actorder": None,
+                    "block_structure": None,
+                    "dynamic": True,
+                    "group_size": None,
+                    "num_bits": 8,
+                    "observer": None,
+                    "observer_kwargs": {},
+                    "strategy": "token",
+                    "symmetric": True,
+                    "type": "float",
+                },
+                "output_activations": None,
+                "targets": ["Linear"],
+                "weights": {
+                    "actorder": None,
+                    "block_structure": None,
+                    "dynamic": False,
+                    "group_size": None,
+                    "num_bits": 8,
+                    "observer": "minmax",
+                    "observer_kwargs": {},
+                    "strategy": strategy,
+                    "symmetric": True,
+                    "type": "float",
+                },
+            },
+        }
+        quantization_config = {
+            "config_groups": quant_group,
+            "format": "float-quantized",
+            "ignore": list(set(result_collector.modules_to_not_convert)),
+            "quant_method": "compressed-tensors",
+            "quantization_status": "compressed",
+        }
+
+    config_path = os.path.join(input_path, "config.json")
+    if os.path.exists(config_path):
+        cfg = json.load(open(config_path))
+        cfg["quantization_config"] = quantization_config
+        json.dump(cfg, open(os.path.join(output_path, "config.json"), "w"), indent=2)
+
+    index_dict = {
+        "weight_map": result_collector.weight_map,
+        "metadata": {"total_size": result_collector.param_count},
+    }
+    json.dump(
+        index_dict,
+        open(os.path.join(output_path, "model.safetensors.index.json"), "w"),
+        indent=2,
+    )
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model-dir",
+        type=str,
+        help="Path to the directory of the HF safetensors model.",
+    )
+    parser.add_argument(
+        "--save-dir",
+        type=str,
+        help="Path to the directory to save the converted model.",
+    )
+    parser.add_argument(
+        "--strategy", type=str, default="block", choices=["block", "channel", "tensor"]
+    )
+    parser.add_argument(
+        "--block-size", type=int, nargs="*", default=None, help="eg. --block-size 32 32"
+    )
+    parser.add_argument(
+        "--max-workers",
+        type=int,
+        default=1,
+        help="Number of worker threads for parallel processing",
+    )
+    args = parser.parse_args()
+
+    if not os.path.exists(args.save_dir):
+        print(f"Creating directory {args.save_dir}")
+        os.makedirs(args.save_dir)
+    elif not os.path.isdir(args.save_dir):
+        raise ValueError("The save_dir should be a directory.")
+
+    convert_fp8(
+        args.model_dir, args.save_dir, args.strategy, args.block_size, args.max_workers
+    )