[diffusion] quant: support quant for all dits (#19156)

Co-authored-by: zyzshishui <zyzshishui@gmail.com>
2026-02-24 14:20:54 +08:00
parent 4f25a48d7a
commit 2e053d6eb6
19 changed files with 523 additions and 133 deletions
--- a/python/sglang/multimodal_gen/configs/models/base.py
+++ b/python/sglang/multimodal_gen/configs/models/base.py
@@ -73,14 +73,9 @@ class ModelConfig:
        Update arch_config with source_model_dict
        """
        arch_config = self.arch_config
-        valid_fields = {f.name for f in fields(arch_config)}

        for key, value in source_model_dict.items():
            setattr(arch_config, key, value)
-            # else:
-            #     raise AttributeError(
-            #         f"{type(arch_config).__name__} has no field '{key}'"
-            #     )

        if hasattr(arch_config, "__post_init__"):
            arch_config.__post_init__()
--- a/python/sglang/multimodal_gen/docs/quantization.md
+++ b/python/sglang/multimodal_gen/docs/quantization.md
@@ -18,9 +18,9 @@ SVDQuant significantly reduces memory usage and accelerates inference while main

 - **Memory Optimization**: Reduces memory usage by **3.6×** compared to BF16 models.
 - **Inference Acceleration**:
-  - **3.0×** faster than the NF4 (W4A16) baseline on desktop/laptop RTX 4090 GPUs.
-  - **8.7×** speedup on laptop RTX 4090 by eliminating CPU offloading compared to 16-bit models.
-  - **3.1×** faster than BF16 and NF4 models on RTX 5090 GPUs with NVFP4.
+    - **3.0×** faster than the NF4 (W4A16) baseline on desktop/laptop RTX 4090 GPUs.
+    - **8.7×** speedup on laptop RTX 4090 by eliminating CPU offloading compared to 16-bit models.
+    - **3.1×** faster than BF16 and NF4 models on RTX 5090 GPUs with NVFP4.

 ### Supported Precisions

--- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
+++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
@@ -1,4 +1,3 @@
-import inspect
 import json
 import logging
 import os
@@ -146,12 +145,6 @@ class TransformerLoader(ComponentLoader):
        dit_config.update_model_arch(config)

        cls_name = config.pop("_class_name")
-        if cls_name is None:
-            raise ValueError(
-                "Model config does not contain a _class_name attribute. "
-                "Only diffusers format is supported."
-            )
-
        model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)

        nunchaku_config = server_args.nunchaku_config
@@ -167,20 +160,21 @@ class TransformerLoader(ComponentLoader):
            param_dtype,
        )

-        init_params: dict[str, Any] = {"config": dit_config, "hf_config": config}
        # prepare init_param
-        if "quant_config" in inspect.signature(model_cls.__init__).parameters:
-            init_params.update(
-                {
-                    "quant_config": (quant_config if quant_config else nunchaku_config),
-                }
+        init_params: dict[str, Any] = {
+            "config": dit_config,
+            "hf_config": config,
+            "quant_config": (quant_config if quant_config else nunchaku_config),
+        }
+        if (
+            init_params["quant_config"] is None
+            and server_args.transformer_weights_path is not None
+        ):
+            logger.warning(
+                f"transformer_weights_path provided, but quantization config not resolved, which is unexpected and likely to cause errors"
            )
-            if init_params["quant_config"] is None:
-                logger.warning(
-                    f"transformer_weights_path provided, but quantization config not resolved, which is unexpected and likely to cause errors"
-                )
-            else:
-                logger.debug("quantization config: %s", init_params["quant_config"])
+        else:
+            logger.debug("quantization config: %s", init_params["quant_config"])

        # Load the model using FSDP loader
        model = maybe_load_fsdp_model(
--- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
+++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
@@ -68,7 +68,6 @@ class VAELoader(ComponentLoader):

        server_args.model_paths[component_name] = component_model_path

-        logger.debug("HF model config: %s", config)
        if component_name in ("vae", "video_vae"):
            pipeline_vae_config_attr = "vae_config"
            pipeline_vae_precision = "vae_precision"
--- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
+++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
@@ -136,6 +136,14 @@ def maybe_load_fsdp_model(
        cpu_offload=cpu_offload,
        param_names_mapping=param_names_mapping_fn,
    )
+
+    for _, module in model.named_modules():
+        quant_method = getattr(module, "quant_method", None)
+        if quant_method is not None and hasattr(
+            quant_method, "process_weights_after_loading"
+        ):
+            quant_method.process_weights_after_loading(module)
+
    for n, p in chain(model.named_parameters(), model.named_buffers()):
        if p.is_meta:
            raise RuntimeError(f"Unexpected param or buffer {n} on meta device.")
@@ -252,6 +260,7 @@ def load_model_from_full_model_state_dict(
    sorted_param_names = sorted(custom_param_sd.keys())

    sharded_sd = {}
+    skipped_checkpoint_keys: list[str] = []

    # shard from loaded state_dict, custom_param_sd -> sharded_sd
    for target_param_name in sorted_param_names:
@@ -265,6 +274,7 @@ 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:
+                skipped_checkpoint_keys.append(target_param_name)
                continue

        # use meta param dtype so quantized params (e.g. FP8) keep their dtype;
@@ -329,6 +339,22 @@ def load_model_from_full_model_state_dict(
        )

    model.reverse_param_names_mapping = reverse_param_names_mapping
+
+    if skipped_checkpoint_keys:
+        logger.warning(
+            "Checkpoint keys not loaded (no matching model parameter) %s",
+            (
+                skipped_checkpoint_keys[:20]
+                if len(skipped_checkpoint_keys) > 20
+                else skipped_checkpoint_keys
+            ),
+        )
+        if len(skipped_checkpoint_keys) > 20:
+            logger.warning(
+                "... and %d more skipped keys.",
+                len(skipped_checkpoint_keys) - 20,
+            )
+
    # parameters in nn.Module that doesn't exist in safetensor files
    unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys())
    if unused_keys:
--- a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py
@@ -35,6 +35,9 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
 )
 from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.mlp import MLP
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
    _apply_rotary_emb,
    get_rotary_pos_embed,
@@ -262,16 +265,17 @@ class CausalWanTransformerBlock(nn.Module):
        added_kv_proj_dim: int | None = None,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()

        # 1. Self-attention
        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
-        self.to_q = ReplicatedLinear(dim, dim, bias=True)
-        self.to_k = ReplicatedLinear(dim, dim, bias=True)
-        self.to_v = ReplicatedLinear(dim, dim, bias=True)
+        self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config)
+        self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config)
+        self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config)

-        self.to_out = ReplicatedLinear(dim, dim, bias=True)
+        self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config)
        self.attn1 = CausalWanSelfAttention(
            dim,
            num_heads,
@@ -310,13 +314,16 @@ class CausalWanTransformerBlock(nn.Module):
            qk_norm=qk_norm,
            eps=eps,
            supported_attention_backends=cross_attn_backends,
+            quant_config=quant_config,
        )
        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
        )

        # 3. Feed-forward
-        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
+        self.ffn = MLP(
+            dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config
+        )
        self.mlp_residual = MulAdd()

        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -430,7 +437,12 @@ class CausalWanTransformer3DModel(BaseDiT, OffloadableDiTMixin):
    reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping
    lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping

-    def __init__(self, config: WanVideoConfig, hf_config: dict[str, Any]) -> None:
+    def __init__(
+        self,
+        config: WanVideoConfig,
+        hf_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
        super().__init__(config=config, hf_config=hf_config)

        inner_dim = config.num_attention_heads * config.attention_head_dim
@@ -475,6 +487,7 @@ class CausalWanTransformer3DModel(BaseDiT, OffloadableDiTMixin):
                    config.added_kv_proj_dim,
                    self._supported_attention_backends,
                    prefix=f"{config.prefix}.blocks.{i}",
+                    quant_config=quant_config,
                )
                for i in range(config.num_layers)
            ]
--- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py
@@ -24,6 +24,9 @@ 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 RMSNorm, apply_qk_norm
 from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
    NDRotaryEmbedding,
    apply_flashinfer_rope_qk_inplace,
@@ -76,6 +79,7 @@ class Flux2FeedForward(nn.Module):
        mult: float = 3.0,
        inner_dim: Optional[int] = None,
        bias: bool = False,
+        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__()
        if inner_dim is None:
@@ -84,11 +88,11 @@ class Flux2FeedForward(nn.Module):

        # Flux2SwiGLU will reduce the dimension by half
        self.linear_in = ColumnParallelLinear(
-            dim, inner_dim * 2, bias=bias, gather_output=True
+            dim, inner_dim * 2, bias=bias, gather_output=True, quant_config=quant_config
        )
        self.act_fn = Flux2SwiGLU()
        self.linear_out = ColumnParallelLinear(
-            inner_dim, dim_out, bias=bias, gather_output=True
+            inner_dim, dim_out, bias=bias, gather_output=True, quant_config=quant_config
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -112,6 +116,7 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
        eps: float = 1e-5,
        out_dim: int = None,
        elementwise_affine: bool = True,
+        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__()

@@ -128,13 +133,25 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
        self.added_proj_bias = added_proj_bias

        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,
        )

        # QK Norm
@@ -144,7 +161,11 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
        self.to_out = torch.nn.ModuleList([])
        self.to_out.append(
            ColumnParallelLinear(
-                self.inner_dim, self.out_dim, bias=out_bias, gather_output=True
+                self.inner_dim,
+                self.out_dim,
+                bias=out_bias,
+                gather_output=True,
+                quant_config=quant_config,
            )
        )
        self.to_out.append(torch.nn.Dropout(dropout))
@@ -157,21 +178,28 @@ class Flux2Attention(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, query_dim, bias=out_bias, gather_output=True
+                self.inner_dim,
+                query_dim,
+                bias=out_bias,
+                gather_output=True,
+                quant_config=quant_config,
            )

        self.attn = USPAttention(
@@ -285,6 +313,7 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
        elementwise_affine: bool = True,
        mlp_ratio: float = 4.0,
        mlp_mult_factor: int = 2,
+        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__()

@@ -307,6 +336,7 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
            self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor,
            bias=bias,
            gather_output=True,
+            quant_config=quant_config,
        )
        self.mlp_act_fn = Flux2SwiGLU()

@@ -320,6 +350,7 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
            self.out_dim,
            bias=out_bias,
            gather_output=True,
+            quant_config=quant_config,
        )

        self.attn = USPAttention(
@@ -390,6 +421,7 @@ class Flux2SingleTransformerBlock(nn.Module):
        mlp_ratio: float = 3.0,
        eps: float = 1e-6,
        bias: bool = False,
+        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__()

@@ -408,6 +440,7 @@ class Flux2SingleTransformerBlock(nn.Module):
            eps=eps,
            mlp_ratio=mlp_ratio,
            mlp_mult_factor=2,
+            quant_config=quant_config,
        )

    def forward(
@@ -461,6 +494,7 @@ class Flux2TransformerBlock(nn.Module):
        mlp_ratio: float = 3.0,
        eps: float = 1e-6,
        bias: bool = False,
+        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__()
        self.mlp_hidden_dim = int(dim * mlp_ratio)
@@ -478,14 +512,17 @@ class Flux2TransformerBlock(nn.Module):
            added_proj_bias=bias,
            out_bias=bias,
            eps=eps,
+            quant_config=quant_config,
        )

        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.ff = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias)
+        self.ff = Flux2FeedForward(
+            dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias, quant_config=quant_config
+        )

        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
        self.ff_context = Flux2FeedForward(
-            dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias
+            dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias, quant_config=quant_config
        )

    def forward(
@@ -663,7 +700,12 @@ class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin):

    param_names_mapping = FluxConfig().arch_config.param_names_mapping

-    def __init__(self, config: FluxConfig, hf_config: dict[str, Any]):
+    def __init__(
+        self,
+        config: FluxConfig,
+        hf_config: dict[str, Any],
+        quant_config: Optional[QuantizationConfig] = None,
+    ):
        super().__init__(config=config, hf_config=hf_config)
        patch_size: int = config.patch_size
        in_channels: int = config.in_channels
@@ -725,6 +767,7 @@ class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin):
                    mlp_ratio=mlp_ratio,
                    eps=eps,
                    bias=False,
+                    quant_config=quant_config,
                )
                for _ in range(num_layers)
            ]
@@ -740,6 +783,7 @@ class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin):
                    mlp_ratio=mlp_ratio,
                    eps=eps,
                    bias=False,
+                    quant_config=quant_config,
                )
                for _ in range(num_single_layers)
            ]
--- a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py
@@ -17,7 +17,6 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from diffusers.models.attention import FeedForward

 from sglang.multimodal_gen.configs.models.dits.glmimage import GlmImageDitConfig
 from sglang.multimodal_gen.runtime.distributed.parallel_state import (
@@ -29,6 +28,10 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
    ScaleResidualLayerNormScaleShift,
 )
 from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
+from sglang.multimodal_gen.runtime.layers.mlp import FeedForward
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
    _apply_rotary_emb,
    apply_flashinfer_rope_qk_inplace,
@@ -315,6 +318,7 @@ class GlmImageAttention(torch.nn.Module):
        eps,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()

@@ -329,13 +333,23 @@ class GlmImageAttention(torch.nn.Module):

        self.num_kv_heads = self.dim_head // self.inner_kv_dim

-        self.to_q = ReplicatedLinear(query_dim, self.inner_dim, bias=bias)
-        self.to_k = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias)
-        self.to_v = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias)
+        self.to_q = ReplicatedLinear(
+            query_dim, self.inner_dim, bias=bias, quant_config=quant_config
+        )
+        self.to_k = ReplicatedLinear(
+            query_dim, self.inner_kv_dim, bias=bias, quant_config=quant_config
+        )
+        self.to_v = ReplicatedLinear(
+            query_dim, self.inner_kv_dim, bias=bias, quant_config=quant_config
+        )

        # (dropout omitted)
        self.to_out = nn.ModuleList(
-            [ReplicatedLinear(self.inner_dim, self.out_dim, bias=True)]
+            [
+                ReplicatedLinear(
+                    self.inner_dim, self.out_dim, bias=True, quant_config=quant_config
+                )
+            ]
        )

        if qk_norm is None:
@@ -466,6 +480,7 @@ class GlmImageTransformerBlock(nn.Module):
        time_embed_dim: int = 512,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
    ) -> None:
        super().__init__()

@@ -483,6 +498,7 @@ class GlmImageTransformerBlock(nn.Module):
            eps=1e-5,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.attn1",
+            quant_config=quant_config,
        )

        # 2. Feedforward
@@ -687,6 +703,7 @@ class GlmImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
        self,
        config: GlmImageDitConfig,
        hf_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__(config=config, hf_config=hf_config)

@@ -747,6 +764,7 @@ class GlmImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
                    arch_config.time_embed_dim,
                    supported_attention_backends=self._supported_attention_backends,
                    prefix=f"transformer_blocks.{i}",
+                    quant_config=quant_config,
                )
                for i in range(arch_config.num_layers)
            ]
--- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py
@@ -23,6 +23,9 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
 )
 from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.mlp import MLP
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
    _apply_rotary_emb,
    get_rotary_pos_embed,
@@ -57,6 +60,7 @@ class MMDoubleStreamBlock(nn.Module):
        dtype: torch.dtype | None = None,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()

@@ -90,6 +94,7 @@ class MMDoubleStreamBlock(nn.Module):
            bias=True,
            params_dtype=dtype,
            prefix=f"{prefix}.img_attn_qkv",
+            quant_config=quant_config,
        )

        self.img_attn_q_norm = RMSNorm(head_dim, eps=1e-6, dtype=dtype)
@@ -101,6 +106,7 @@ class MMDoubleStreamBlock(nn.Module):
            bias=True,
            params_dtype=dtype,
            prefix=f"{prefix}.img_attn_proj",
+            quant_config=quant_config,
        )

        self.img_mlp = MLP(
@@ -109,6 +115,7 @@ class MMDoubleStreamBlock(nn.Module):
            bias=True,
            dtype=dtype,
            prefix=f"{prefix}.img_mlp",
+            quant_config=quant_config,
        )

        # Text modulation components
@@ -131,7 +138,11 @@ class MMDoubleStreamBlock(nn.Module):

        # Text attention components
        self.txt_attn_qkv = ReplicatedLinear(
-            hidden_size, hidden_size * 3, bias=True, params_dtype=dtype
+            hidden_size,
+            hidden_size * 3,
+            bias=True,
+            params_dtype=dtype,
+            quant_config=quant_config,
        )

        # QK norm layers for text
@@ -139,10 +150,20 @@ class MMDoubleStreamBlock(nn.Module):
        self.txt_attn_k_norm = RMSNorm(head_dim, eps=1e-6, dtype=dtype)

        self.txt_attn_proj = ReplicatedLinear(
-            hidden_size, hidden_size, bias=True, params_dtype=dtype
+            hidden_size,
+            hidden_size,
+            bias=True,
+            params_dtype=dtype,
+            quant_config=quant_config,
        )

-        self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype)
+        self.txt_mlp = MLP(
+            hidden_size,
+            mlp_hidden_dim,
+            bias=True,
+            dtype=dtype,
+            quant_config=quant_config,
+        )

        # Use UlyssesAttention to replace Distributed attention
        self.attn = UlyssesAttention(
@@ -264,6 +285,7 @@ class MMSingleStreamBlock(nn.Module):
        dtype: torch.dtype | None = None,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()

@@ -281,6 +303,7 @@ class MMSingleStreamBlock(nn.Module):
            bias=True,
            params_dtype=dtype,
            prefix=f"{prefix}.linear1",
+            quant_config=quant_config,
        )

        # Combined projection and MLP output
@@ -290,6 +313,7 @@ class MMSingleStreamBlock(nn.Module):
            bias=True,
            params_dtype=dtype,
            prefix=f"{prefix}.linear2",
+            quant_config=quant_config,
        )

        # QK norm layers
@@ -408,7 +432,12 @@ class HunyuanVideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
    reverse_param_names_mapping = HunyuanVideoConfig().reverse_param_names_mapping
    lora_param_names_mapping = HunyuanVideoConfig().lora_param_names_mapping

-    def __init__(self, config: HunyuanVideoConfig, hf_config: dict[str, Any]):
+    def __init__(
+        self,
+        config: HunyuanVideoConfig,
+        hf_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
+    ):
        super().__init__(config=config, hf_config=hf_config)

        self.patch_size = [config.patch_size_t, config.patch_size, config.patch_size]
@@ -494,6 +523,7 @@ class HunyuanVideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
                    dtype=config.dtype,
                    supported_attention_backends=self._supported_attention_backends,
                    prefix=f"{config.prefix}.double_blocks.{i}",
+                    quant_config=quant_config,
                )
                for i in range(config.num_layers)
            ]
@@ -509,6 +539,7 @@ class HunyuanVideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
                    dtype=config.dtype,
                    supported_attention_backends=self._supported_attention_backends,
                    prefix=f"{config.prefix}.single_blocks.{i + config.num_layers}",
+                    quant_config=quant_config,
                )
                for i in range(config.num_single_layers)
            ]
--- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
@@ -26,6 +26,9 @@ from sglang.multimodal_gen.runtime.layers.linear import (
    ColumnParallelLinear,
    RowParallelLinear,
 )
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
 from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
@@ -445,6 +448,7 @@ class LTX2Attention(nn.Module):
        qk_norm: bool = True,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
    ) -> None:
        super().__init__()

@@ -473,13 +477,25 @@ class LTX2Attention(nn.Module):
        self.local_heads = self.heads // tp_size

        self.to_q = ColumnParallelLinear(
-            self.query_dim, self.inner_dim, bias=True, gather_output=False
+            self.query_dim,
+            self.inner_dim,
+            bias=True,
+            gather_output=False,
+            quant_config=quant_config,
        )
        self.to_k = ColumnParallelLinear(
-            self.context_dim, self.inner_dim, bias=True, gather_output=False
+            self.context_dim,
+            self.inner_dim,
+            bias=True,
+            gather_output=False,
+            quant_config=quant_config,
        )
        self.to_v = ColumnParallelLinear(
-            self.context_dim, self.inner_dim, bias=True, gather_output=False
+            self.context_dim,
+            self.inner_dim,
+            bias=True,
+            gather_output=False,
+            quant_config=quant_config,
        )

        self.q_norm: nn.Module | None = None
@@ -502,7 +518,11 @@ class LTX2Attention(nn.Module):

        self.to_out = nn.Sequential(
            RowParallelLinear(
-                self.inner_dim, self.query_dim, bias=True, input_is_parallel=True
+                self.inner_dim,
+                self.query_dim,
+                bias=True,
+                input_is_parallel=True,
+                quant_config=quant_config,
            ),
            nn.Identity(),
        )
@@ -624,18 +644,24 @@ class LTX2Attention(nn.Module):


 class LTX2FeedForward(nn.Module):
-    def __init__(self, dim: int, dim_out: int | None = None, mult: int = 4) -> None:
+    def __init__(
+        self,
+        dim: int,
+        dim_out: int | None = None,
+        mult: int = 4,
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
        super().__init__()
        if dim_out is None:
            dim_out = dim
        inner_dim = int(dim * mult)

        self.proj_in = ColumnParallelLinear(
-            dim, inner_dim, bias=True, gather_output=True
+            dim, inner_dim, bias=True, gather_output=True, quant_config=quant_config
        )
        self.act = nn.GELU(approximate="tanh")
        self.proj_out = ColumnParallelLinear(
-            inner_dim, dim_out, bias=True, gather_output=True
+            inner_dim, dim_out, bias=True, gather_output=True, quant_config=quant_config
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -661,6 +687,7 @@ class LTX2TransformerBlock(nn.Module):
        norm_eps: float = 1e-6,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()
        self.idx = idx
@@ -675,6 +702,7 @@ class LTX2TransformerBlock(nn.Module):
            qk_norm=qk_norm,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.attn1",
+            quant_config=quant_config,
        )
        self.audio_attn1 = LTX2Attention(
            query_dim=audio_dim,
@@ -684,6 +712,7 @@ class LTX2TransformerBlock(nn.Module):
            qk_norm=qk_norm,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.audio_attn1",
+            quant_config=quant_config,
        )

        # 2. Prompt Cross-Attention
@@ -696,6 +725,7 @@ class LTX2TransformerBlock(nn.Module):
            qk_norm=qk_norm,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.attn2",
+            quant_config=quant_config,
        )
        self.audio_attn2 = LTX2Attention(
            query_dim=audio_dim,
@@ -706,6 +736,7 @@ class LTX2TransformerBlock(nn.Module):
            qk_norm=qk_norm,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.audio_attn2",
+            quant_config=quant_config,
        )

        # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
@@ -718,6 +749,7 @@ class LTX2TransformerBlock(nn.Module):
            qk_norm=qk_norm,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.audio_to_video_attn",
+            quant_config=quant_config,
        )
        self.video_to_audio_attn = LTX2Attention(
            query_dim=audio_dim,
@@ -728,11 +760,14 @@ class LTX2TransformerBlock(nn.Module):
            qk_norm=qk_norm,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.video_to_audio_attn",
+            quant_config=quant_config,
        )

        # 4. Feedforward layers
-        self.ff = LTX2FeedForward(dim, dim_out=dim)
-        self.audio_ff = LTX2FeedForward(audio_dim, dim_out=audio_dim)
+        self.ff = LTX2FeedForward(dim, dim_out=dim, quant_config=quant_config)
+        self.audio_ff = LTX2FeedForward(
+            audio_dim, dim_out=audio_dim, quant_config=quant_config
+        )

        # 5. Modulation Parameters
        self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
@@ -1001,7 +1036,12 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
                f"{arch.audio_out_channels=} {tp_size=}."
            )

-    def __init__(self, config: LTX2Config, hf_config: dict[str, Any]) -> None:
+    def __init__(
+        self,
+        config: LTX2Config,
+        hf_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
        super().__init__(config=config, hf_config=hf_config)

        arch = config.arch_config
@@ -1017,13 +1057,18 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
        # 1. Patchification input projections
        # Matches LTX2Config().param_names_mapping
        self.patchify_proj = ColumnParallelLinear(
-            arch.in_channels, self.hidden_size, bias=True, gather_output=True
+            arch.in_channels,
+            self.hidden_size,
+            bias=True,
+            gather_output=True,
+            quant_config=quant_config,
        )
        self.audio_patchify_proj = ColumnParallelLinear(
            arch.audio_in_channels,
            self.audio_hidden_size,
            bias=True,
            gather_output=True,
+            quant_config=quant_config,
        )

        # 2. Prompt embeddings
@@ -1169,6 +1214,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
                    qk_norm=True,  # Always True in LTX2
                    supported_attention_backends=self._supported_attention_backends,
                    prefix=config.prefix,
+                    quant_config=quant_config,
                )
                for idx in range(arch.num_layers)
            ]
@@ -1179,7 +1225,11 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
            self.hidden_size, eps=self.norm_eps, elementwise_affine=False
        )
        self.proj_out = ColumnParallelLinear(
-            self.hidden_size, arch.out_channels, bias=True, gather_output=True
+            self.hidden_size,
+            arch.out_channels,
+            bias=True,
+            gather_output=True,
+            quant_config=quant_config,
        )

        self.audio_norm_out = nn.LayerNorm(
@@ -1190,6 +1240,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
            arch.audio_out_channels,
            bias=True,
            gather_output=True,
+            quant_config=quant_config,
        )

        self.out_channels_raw = arch.out_channels // (
--- a/python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py
@@ -15,6 +15,9 @@ from torch.distributed.tensor import DTensor
 from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig
 from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.mlp import MLP
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin

@@ -106,7 +109,12 @@ class WanAudioModel(CachableDiT, OffloadableDiTMixin):
    reverse_param_names_mapping = MOVAAudioConfig().reverse_param_names_mapping
    lora_param_names_mapping = MOVAAudioConfig().lora_param_names_mapping

-    def __init__(self, config: MOVAAudioConfig, hf_config: dict[str, Any]) -> None:
+    def __init__(
+        self,
+        config: MOVAAudioConfig,
+        hf_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
        super().__init__(config=config, hf_config=hf_config)

        # Extract parameters from config
@@ -142,13 +150,24 @@ class WanAudioModel(CachableDiT, OffloadableDiTMixin):
            in_dim, dim, kernel_size=patch_size, stride=patch_size
        )
        self.text_embedding = MLP(
-            text_dim, dim, output_dim=dim, act_type="gelu_pytorch_tanh"
+            text_dim,
+            dim,
+            output_dim=dim,
+            act_type="gelu_pytorch_tanh",
+            quant_config=quant_config,
+        )
+        self.time_embedding = MLP(
+            freq_dim, dim, output_dim=dim, act_type="silu", quant_config=quant_config
        )
-        self.time_embedding = MLP(freq_dim, dim, output_dim=dim, act_type="silu")
        # Preserve state_dict keys (time_projection.1.weight/bias).
-        self.time_projection = nn.Sequential(nn.SiLU(), ReplicatedLinear(dim, dim * 6))
+        self.time_projection = nn.Sequential(
+            nn.SiLU(), ReplicatedLinear(dim, dim * 6, quant_config=quant_config)
+        )
        self.blocks = nn.ModuleList(
-            [DiTBlock(dim, num_heads, ffn_dim, eps) for _ in range(num_layers)]
+            [
+                DiTBlock(dim, num_heads, ffn_dim, eps, quant_config=quant_config)
+                for _ in range(num_layers)
+            ]
        )
        self.head = Head(dim, out_dim, patch_size, eps)
        self.num_heads = num_heads
--- a/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py
@@ -30,6 +30,9 @@ from sglang.multimodal_gen.runtime.layers.linear import (
    RowParallelLinear,
 )
 from sglang.multimodal_gen.runtime.layers.mlp import MLP
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -104,7 +107,13 @@ class SelfAttention(nn.Module):
    Input x should already be the local shard [B, S_local, D] when SP is enabled.
    """

-    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        eps: float = 1e-6,
+        quant_config: QuantizationConfig | None = None,
+    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
@@ -118,10 +127,18 @@ class SelfAttention(nn.Module):
        self.num_heads_per_rank = self.num_heads // self.tp_size

        # TP strategy: shard Q/K/V over heads (column-parallel), then row-parallel output.
-        self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
-        self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
-        self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
-        self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True)
+        self.q = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.k = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.v = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.o = RowParallelLinear(
+            dim, dim, bias=True, input_is_parallel=True, quant_config=quant_config
+        )
        self.norm_q = RMSNorm(dim, eps=eps)
        self.norm_k = RMSNorm(dim, eps=eps)

@@ -188,7 +205,13 @@ class CrossAttention(nn.Module):
    Uses LocalAttention instead of USPAttention for efficiency.
    """

-    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        eps: float = 1e-6,
+        quant_config: QuantizationConfig | None = None,
+    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
@@ -201,10 +224,18 @@ class CrossAttention(nn.Module):
            )
        self.num_heads_per_rank = self.num_heads // self.tp_size

-        self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
-        self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
-        self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
-        self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True)
+        self.q = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.k = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.v = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.o = RowParallelLinear(
+            dim, dim, bias=True, input_is_parallel=True, quant_config=quant_config
+        )
        self.norm_q = RMSNorm(dim, eps=eps)
        self.norm_k = RMSNorm(dim, eps=eps)

@@ -264,14 +295,15 @@ class DiTBlock(nn.Module):
        num_heads: int,
        ffn_dim: int,
        eps: float = 1e-6,
+        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.ffn_dim = ffn_dim

-        self.self_attn = SelfAttention(dim, num_heads, eps)
-        self.cross_attn = CrossAttention(dim, num_heads, eps)
+        self.self_attn = SelfAttention(dim, num_heads, eps, quant_config=quant_config)
+        self.cross_attn = CrossAttention(dim, num_heads, eps, quant_config=quant_config)
        self.norm1 = LayerNormScaleShift(
            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
        )
@@ -281,7 +313,13 @@ class DiTBlock(nn.Module):
        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
        )
-        self.ffn = MLP(dim, ffn_dim, output_dim=dim, act_type="gelu_pytorch_tanh")
+        self.ffn = MLP(
+            dim,
+            ffn_dim,
+            output_dim=dim,
+            act_type="gelu_pytorch_tanh",
+            quant_config=quant_config,
+        )
        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
        self.mlp_residual = MulAdd()

@@ -388,7 +426,12 @@ class WanModel(CachableDiT, OffloadableDiTMixin):
    reverse_param_names_mapping = MOVAVideoConfig().reverse_param_names_mapping
    lora_param_names_mapping = MOVAVideoConfig().lora_param_names_mapping

-    def __init__(self, config: MOVAVideoConfig, hf_config: dict[str, Any]) -> None:
+    def __init__(
+        self,
+        config: MOVAVideoConfig,
+        hf_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
        super().__init__(config=config, hf_config=hf_config)

        # Extract parameters from config
@@ -421,13 +464,24 @@ class WanModel(CachableDiT, OffloadableDiTMixin):
            in_dim, dim, kernel_size=patch_size, stride=patch_size
        )
        self.text_embedding = MLP(
-            text_dim, dim, output_dim=dim, act_type="gelu_pytorch_tanh"
+            text_dim,
+            dim,
+            output_dim=dim,
+            act_type="gelu_pytorch_tanh",
+            quant_config=quant_config,
+        )
+        self.time_embedding = MLP(
+            freq_dim, dim, output_dim=dim, act_type="silu", quant_config=quant_config
        )
-        self.time_embedding = MLP(freq_dim, dim, output_dim=dim, act_type="silu")
        # Preserve state_dict keys (time_projection.1.weight/bias).
-        self.time_projection = nn.Sequential(nn.SiLU(), ReplicatedLinear(dim, dim * 6))
+        self.time_projection = nn.Sequential(
+            nn.SiLU(), ReplicatedLinear(dim, dim * 6, quant_config=quant_config)
+        )
        self.blocks = nn.ModuleList(
-            [DiTBlock(dim, num_heads, ffn_dim, eps) for _ in range(num_layers)]
+            [
+                DiTBlock(dim, num_heads, ffn_dim, eps, quant_config=quant_config)
+                for _ in range(num_layers)
+            ]
        )
        self.head = Head(dim, out_dim, patch_size, eps)
        self.num_heads = num_heads
--- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
@@ -137,6 +137,8 @@ class FeedForward(nn.Module):
        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.
+        quant_config: Quantization configure.
+        prefix: The name of the layer in the state dict.
    """

    def __init__(
@@ -147,6 +149,8 @@ class FeedForward(nn.Module):
        activation_fn: str = "geglu",
        inner_dim=None,
        bias: bool = True,
+        quant_config=None,
+        prefix: str = "",
    ):
        super().__init__()
        if inner_dim is None:
@@ -154,9 +158,16 @@ class FeedForward(nn.Module):
        dim_out = dim_out if dim_out is not None else dim

        if activation_fn == "gelu":
-            act_fn = GELU(dim, inner_dim, bias=bias)
+            act_fn = GELU(dim, inner_dim, bias=bias, quant_config=None, prefix=prefix)
        if activation_fn == "gelu-approximate":
-            act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias)
+            act_fn = GELU(
+                dim,
+                inner_dim,
+                approximate="tanh",
+                bias=bias,
+                quant_config=None,
+                prefix=prefix,
+            )
        else:
            raise NotImplementedError(
                f"activation_fn '{activation_fn}' is not supported."
@@ -166,7 +177,13 @@ class FeedForward(nn.Module):
        self.net.append(act_fn)
        self.net.append(nn.Identity())
        self.net.append(
-            RowParallelLinear(inner_dim, dim_out, bias=True, input_is_parallel=True)
+            RowParallelLinear(
+                inner_dim,
+                dim_out,
+                bias=True,
+                input_is_parallel=True,
+                quant_config=None,
+            )
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -823,6 +840,11 @@ class QwenImageTransformerBlock(nn.Module):
        self.quant_config = quant_config
        self.zero_cond_t = zero_cond_t

+        mod_quant_config = (
+            quant_config
+            if (quant_config is not None and quant_config.get_name() == "svdquant")
+            else None
+        )
        # Image processing modules
        self.img_mod = nn.Sequential(
            nn.SiLU(),
@@ -831,7 +853,7 @@ class QwenImageTransformerBlock(nn.Module):
                6 * dim,
                bias=True,
                gather_output=True,
-                quant_config=quant_config,
+                quant_config=mod_quant_config,
                prefix=f"{prefix}.img_mod",
            ),  # For scale, shift, gate for norm1 and norm2
        )
@@ -860,7 +882,7 @@ class QwenImageTransformerBlock(nn.Module):
                6 * dim,
                bias=True,
                gather_output=True,
-                quant_config=quant_config,
+                quant_config=mod_quant_config,
                prefix=f"{prefix}.txt_mod",
            ),  # For scale, shift, gate for norm1 and norm2
        )
@@ -880,19 +902,33 @@ class QwenImageTransformerBlock(nn.Module):
            and quant_config.get_name() == "svdquant"
            and is_nunchaku_available()
        )
-        ff_class = (
-            diffusers.models.attention.FeedForward if nunchaku_enabled else FeedForward
-        )
-        self.img_mlp = ff_class(
-            dim=dim,
-            dim_out=dim,
-            activation_fn="gelu-approximate",
-        )
-        self.txt_mlp = ff_class(
-            dim=dim,
-            dim_out=dim,
-            activation_fn="gelu-approximate",
-        )
+        if nunchaku_enabled:
+            ff_class = diffusers.models.attention.FeedForward
+            self.img_mlp = ff_class(
+                dim=dim,
+                dim_out=dim,
+                activation_fn="gelu-approximate",
+            )
+            self.txt_mlp = ff_class(
+                dim=dim,
+                dim_out=dim,
+                activation_fn="gelu-approximate",
+            )
+        else:
+            self.img_mlp = FeedForward(
+                dim=dim,
+                dim_out=dim,
+                activation_fn="gelu-approximate",
+                quant_config=quant_config,
+                prefix=f"{prefix}.img_mlp",
+            )
+            self.txt_mlp = FeedForward(
+                dim=dim,
+                dim_out=dim,
+                activation_fn="gelu-approximate",
+                quant_config=quant_config,
+                prefix=f"{prefix}.txt_mlp",
+            )

        if nunchaku_enabled:
            nunchaku_kwargs = {
--- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
@@ -36,6 +36,9 @@ from sglang.multimodal_gen.runtime.layers.linear import (
    RowParallelLinear,
 )
 from sglang.multimodal_gen.runtime.layers.mlp import MLP
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
    NDRotaryEmbedding,
    _apply_rotary_emb,
@@ -131,6 +134,7 @@ class WanSelfAttention(nn.Module):
        eps=1e-6,
        parallel_attention=False,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        quant_config: QuantizationConfig | None = None,
    ) -> None:
        assert dim % num_heads == 0
        super().__init__()
@@ -144,10 +148,18 @@ class WanSelfAttention(nn.Module):
        tp_size = get_tp_world_size()

        # layers
-        self.to_q = ColumnParallelLinear(dim, dim, gather_output=False)
-        self.to_k = ColumnParallelLinear(dim, dim, gather_output=False)
-        self.to_v = ColumnParallelLinear(dim, dim, gather_output=False)
-        self.to_out = RowParallelLinear(dim, dim, input_is_parallel=True)
+        self.to_q = ColumnParallelLinear(
+            dim, dim, gather_output=False, quant_config=quant_config
+        )
+        self.to_k = ColumnParallelLinear(
+            dim, dim, gather_output=False, quant_config=quant_config
+        )
+        self.to_v = ColumnParallelLinear(
+            dim, dim, gather_output=False, quant_config=quant_config
+        )
+        self.to_out = RowParallelLinear(
+            dim, dim, input_is_parallel=True, quant_config=quant_config
+        )
        self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
        self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
        self.tp_rmsnorm = tp_size > 1 and qk_norm
@@ -219,6 +231,7 @@ class WanI2VCrossAttention(WanSelfAttention):
        qk_norm=True,
        eps=1e-6,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        quant_config: QuantizationConfig | None = None,
    ) -> None:
        super().__init__(
            dim,
@@ -227,10 +240,15 @@ class WanI2VCrossAttention(WanSelfAttention):
            qk_norm,
            eps,
            supported_attention_backends=supported_attention_backends,
+            quant_config=quant_config,
        )

-        self.add_k_proj = ColumnParallelLinear(dim, dim, gather_output=False)
-        self.add_v_proj = ColumnParallelLinear(dim, dim, gather_output=False)
+        self.add_k_proj = ColumnParallelLinear(
+            dim, dim, gather_output=False, quant_config=quant_config
+        )
+        self.add_v_proj = ColumnParallelLinear(
+            dim, dim, gather_output=False, quant_config=quant_config
+        )
        self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()

    def forward(self, x, context, context_lens):
@@ -296,6 +314,7 @@ class WanTransformerBlock(nn.Module):
        prefix: str = "",
        attention_type: str = "original",
        sla_topk: float = 0.1,
+        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()

@@ -306,11 +325,19 @@ class WanTransformerBlock(nn.Module):
            elementwise_affine=False,
            dtype=torch.float32,
        )
-        self.to_q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
-        self.to_k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
-        self.to_v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
+        self.to_q = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.to_k = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.to_v = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )

-        self.to_out = RowParallelLinear(dim, dim, bias=True, reduce_results=True)
+        self.to_out = RowParallelLinear(
+            dim, dim, bias=True, reduce_results=True, quant_config=quant_config
+        )
        tp_size = get_tp_world_size()
        self.local_num_heads = divide(num_heads, tp_size)
        self_attn_backends = supported_attention_backends
@@ -370,6 +397,7 @@ class WanTransformerBlock(nn.Module):
                qk_norm=qk_norm,
                eps=eps,
                supported_attention_backends=cross_attn_backends,
+                quant_config=quant_config,
            )
        else:
            # T2V
@@ -379,6 +407,7 @@ class WanTransformerBlock(nn.Module):
                qk_norm=qk_norm,
                eps=eps,
                supported_attention_backends=cross_attn_backends,
+                quant_config=quant_config,
            )
        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
            dim,
@@ -388,7 +417,9 @@ class WanTransformerBlock(nn.Module):
        )

        # 3. Feed-forward
-        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
+        self.ffn = MLP(
+            dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config
+        )
        self.mlp_residual = MulAdd()

        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -509,6 +540,7 @@ class WanTransformerBlock_VSA(nn.Module):
        added_kv_proj_dim: int | None = None,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()

@@ -519,14 +551,22 @@ class WanTransformerBlock_VSA(nn.Module):
            elementwise_affine=False,
            dtype=torch.float32,
        )
-        self.to_q = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
-        self.to_k = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
-        self.to_v = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
+        self.to_q = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=True, quant_config=quant_config
+        )
+        self.to_k = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=True, quant_config=quant_config
+        )
+        self.to_v = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=True, quant_config=quant_config
+        )
        self.to_gate_compress = ColumnParallelLinear(
-            dim, dim, bias=True, gather_output=True
+            dim, dim, bias=True, gather_output=True, quant_config=quant_config
        )

-        self.to_out = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
+        self.to_out = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=True, quant_config=quant_config
+        )
        self.attn1 = UlyssesAttention_VSA(
            num_heads=num_heads,
            head_size=dim // num_heads,
@@ -567,6 +607,7 @@ class WanTransformerBlock_VSA(nn.Module):
                qk_norm=qk_norm,
                eps=eps,
                supported_attention_backends=cross_attn_backends,
+                quant_config=quant_config,
            )
        else:
            # T2V
@@ -576,6 +617,7 @@ class WanTransformerBlock_VSA(nn.Module):
                qk_norm=qk_norm,
                eps=eps,
                supported_attention_backends=cross_attn_backends,
+                quant_config=quant_config,
            )
        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
            dim,
@@ -585,7 +627,9 @@ class WanTransformerBlock_VSA(nn.Module):
        )

        # 3. Feed-forward
-        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
+        self.ffn = MLP(
+            dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config
+        )
        self.mlp_residual = MulAdd()

        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -685,7 +729,12 @@ class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin):
    reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping
    lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping

-    def __init__(self, config: WanVideoConfig, hf_config: dict[str, Any]) -> None:
+    def __init__(
+        self,
+        config: WanVideoConfig,
+        hf_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
        super().__init__(config=config, hf_config=hf_config)

        inner_dim = config.num_attention_heads * config.attention_head_dim
@@ -735,6 +784,7 @@ class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin):
                    prefix=f"{config.prefix}.blocks.{i}",
                    attention_type=config.attention_type,
                    sla_topk=config.sla_topk,
+                    quant_config=quant_config,
                )
                for i in range(config.num_layers)
            ]
@@ -752,6 +802,7 @@ class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin):
            config.out_channels * math.prod(config.patch_size),
            bias=True,
            gather_output=True,
+            quant_config=quant_config,
        )
        self.scale_shift_table = nn.Parameter(
            torch.randn(1, 2, inner_dim) / inner_dim**0.5
--- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py
@@ -156,12 +156,29 @@ class ZImageAttention(nn.Module):
                prefix=f"{prefix}.to_qkv",
            )
        else:
-            self.to_q = ColumnParallelLinear(dim, dim, bias=False, gather_output=False)
+            self.to_q = ColumnParallelLinear(
+                dim,
+                dim,
+                bias=False,
+                gather_output=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_q",
+            )
            self.to_k = ColumnParallelLinear(
-                dim, kv_dim, bias=False, gather_output=False
+                dim,
+                kv_dim,
+                bias=False,
+                gather_output=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_k",
            )
            self.to_v = ColumnParallelLinear(
-                dim, kv_dim, bias=False, gather_output=False
+                dim,
+                kv_dim,
+                bias=False,
+                gather_output=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_v",
            )

        if self.qk_norm:
--- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
+++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
@@ -893,7 +893,6 @@ def snapshot_download(
            "allow_patterns": allow_patterns,
            "local_files_only": local_files_only,
            "max_workers": max_workers,
-            "resume_download": True,
            "etag_timeout": 60,
        }
        hf_kwargs.update(kwargs)
--- a/python/sglang/multimodal_gen/test/server/perf_baselines.json
+++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json
@@ -316,11 +316,10 @@
            "expected_avg_denoise_ms": 38.25,
            "expected_median_denoise_ms": 35.95
        },
-        "flux_2_image_t2i_layerwise_offload": {
+        "layerwise_offload": {
            "stages_ms": {
                "InputValidationStage": 0.06,
                "TextEncodingStage": 513.58,
-                "ImageVAEEncodingStage": 0.0,
                "LatentPreparationStage": 0.46,
                "TimestepPreparationStage": 2.38,
                "DenoisingStage": 52187.62,
@@ -537,6 +536,30 @@
            "expected_avg_denoise_ms": 83.75,
            "expected_median_denoise_ms": 93.58
        },
+        "zimage_image_t2i_fp8": {
+            "stages_ms": {
+                "InputValidationStage": 0.04,
+                "TextEncodingStage": 428.59,
+                "LatentPreparationStage": 0.14,
+                "TimestepPreparationStage": 47.26,
+                "DenoisingStage": 778.56,
+                "DecodingStage": 10.39
+            },
+            "denoise_step_ms": {
+                "0": 40.9,
+                "1": 61.08,
+                "2": 95.65,
+                "3": 95.83,
+                "4": 95.65,
+                "5": 96.09,
+                "6": 96.23,
+                "7": 96.04,
+                "8": 96.29
+            },
+            "expected_e2e_ms": 1370.28,
+            "expected_avg_denoise_ms": 85.97,
+            "expected_median_denoise_ms": 95.83
+        },
        "zimage_image_t2i_multi_lora": {
            "stages_ms": {
                "InputValidationStage": 0.04,
--- a/python/sglang/multimodal_gen/test/server/testcase_configs.py
+++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py
@@ -295,6 +295,8 @@ class PerformanceSummary:
        )


+SMALL_T2I_MODEL = "Tongyi-MAI/Z-Image-Turbo"
+
 T2I_sampling_params = DiffusionSamplingParams(
    prompt="Doraemon is eating dorayaki",
    output_size="1024x1024",
@@ -397,9 +399,9 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
    # TODO: currently, we don't support sending more than one request in test, and setting `num_outputs_per_prompt` to 2 doesn't guarantee the denoising be executed twice,
    # so we do one warmup and send one request instead
    DiffusionTestCase(
-        "flux_2_image_t2i_layerwise_offload",
+        "layerwise_offload",
        DiffusionServerArgs(
-            model_path="black-forest-labs/FLUX.2-dev",
+            model_path=SMALL_T2I_MODEL,
            modality="image",
            dit_layerwise_offload=True,
            dit_offload_prefetch_size=2,
@@ -411,6 +413,15 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
        DiffusionServerArgs(model_path="Tongyi-MAI/Z-Image-Turbo", modality="image"),
        T2I_sampling_params,
    ),
+    DiffusionTestCase(
+        "zimage_image_t2i_fp8",
+        DiffusionServerArgs(
+            model_path="Tongyi-MAI/Z-Image-Turbo",
+            modality="image",
+            extras=["--transformer-path MickJ/Z-Image-Turbo-fp8"],
+        ),
+        T2I_sampling_params,
+    ),
    # Multi-LoRA test case for Z-Image-Turbo
    DiffusionTestCase(
        "zimage_image_t2i_multi_lora",
--- a/python/sglang/multimodal_gen/tools/convert_hf_to_fp8.py
+++ b/python/sglang/multimodal_gen/tools/convert_hf_to_fp8.py
@@ -147,6 +147,15 @@ def process_file(
            and "lm_head" not in key
            and "eh_proj" not in key
            and "net" not in key
+            and "txt_mod" not in key
+            and "img_mod" not in key
+            and "img_in" not in key
+            and "txt_in" not in key
+            and "time_in" not in key
+            and "vector_in" not in key
+            and "adaLN_modulation" not in key
+            and "all_final_layer" not in key
+            and "feed_forward" not in key
            and "proj_out.weight" != key
        ):
            qw, s = quant_fp8(weights[key], strategy, block_size)
@@ -294,7 +303,7 @@ if __name__ == "__main__":
    parser.add_argument(
        "--max-workers",
        type=int,
-        default=1,
+        default=8,
        help="Number of worker threads for parallel processing",
    )
    args = parser.parse_args()