[Feat] Support online quantization (fp8/int8) for DreamID-Omni (#3902)

Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -102,3 +102,5 @@ Key arguments:
 - `--audio-negative-prompt`: negative prompt for audio generation.
 - `--enable-cpu-offload`: enable CPU offload (defaults False).
 - `--cache-backend`: enable `cache_dit` for acceleration.
+- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.
+  (VAEs, the T5 text encoder, norms and modulation stay bf16.)

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -4,6 +4,7 @@
 import argparse
 import re
 import time
+from typing import Any
 
 import numpy as np
 from PIL import Image
@@ -41,6 +42,12 @@ def parse_args() -> argparse.Namespace:
         choices=[1, 2, 3, 4],
         help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).",
     )
+    parser.add_argument(
+        "--tensor-parallel-size",
+        type=int,
+        default=1,
+        help=("Number of GPUs used for tensor parallelism (TP) inside the DiT."),
+    )
     parser.add_argument(
         "--video-negative-prompt",
         default="jitter, bad hands, blur, distortion",
@@ -52,6 +59,13 @@ def parse_args() -> argparse.Namespace:
         help="Negative prompt for audio.",
     )
     parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        choices=["fp8", "int8"],
+        help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "),
+    )
     parser.add_argument(
         "--enable-cpu-offload",
         action="store_true",
@@ -100,6 +114,27 @@ def load_image_and_audio(image_paths, audio_paths):
     return image, audio
 
 
+def _extract_peak_memory_mb(result: Any) -> float:
+    """Pull worker-reported peak VRAM (MiB) generation result.
+
+    Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
+    """
+    if isinstance(result, list):
+        result = result[0] if result else None
+    if result is None:
+        return 0.0
+    val = getattr(result, "peak_memory_mb", 0.0)
+    if not val:
+        inner = getattr(result, "request_output", None)
+        if isinstance(inner, list):
+            inner = inner[0] if inner else None
+        val = getattr(inner, "peak_memory_mb", 0.0)
+    try:
+        return float(val or 0.0)
+    except (TypeError, ValueError):
+        return 0.0
+
+
 def main() -> None:
     args = parse_args()
     if args.prompt is None and args.prompt_file is None:
@@ -142,6 +177,7 @@ def main() -> None:
 
     parallel_config = DiffusionParallelConfig(
         cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
         use_hsdp=args.use_hsdp,
         hsdp_shard_size=args.hsdp_shard_size,
         hsdp_replicate_size=args.hsdp_replicate_size,
@@ -162,7 +198,7 @@ def main() -> None:
             "scm_steps_policy": "dynamic",
         }
 
-    omni = Omni(
+    omni_kwargs = dict(
         model=args.model,
         parallel_config=parallel_config,
         model_type=args.model_type,
@@ -171,6 +207,9 @@ def main() -> None:
         cache_backend=args.cache_backend,
         cache_config=cache_config,
     )
+    if args.quantization is not None:
+        omni_kwargs["quantization"] = args.quantization
+    omni = Omni(**omni_kwargs)
     start = time.perf_counter()
     outputs = omni.generate(prompt, sampling_params)
     elapsed = time.perf_counter() - start
@@ -208,6 +247,9 @@ def main() -> None:
         f.write(video_bytes)
     print(f"Saved generated video to {output_path}")
     print(f"Total time: {elapsed:.2f}s")
+    peak_mb = _extract_peak_memory_mb(outputs)
+    if peak_mb:
+        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
 
 
 if __name__ == "__main__":

vllm_omni/diffusion/models/dreamid_omni/fusion.py

@@ -1,18 +1,28 @@
 import re
+from typing import TYPE_CHECKING
 
 import torch
 import torch.nn as nn
+from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import ColumnParallelLinear
 
 from vllm_omni.diffusion.attention.layer import Attention
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
 try:
-    from dreamid_omni.modules.model import WanLayerNorm, WanRMSNorm
+    from dreamid_omni.modules.model import WanLayerNorm
 except ImportError:
     raise ImportError("Failed to import from dependency 'dreamid_omni'.")
 
 from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.models.dreamid_omni.wan2_2 import WanModel, rope_apply
+from vllm_omni.diffusion.models.dreamid_omni.wan2_2 import (
+    DistributedRMSNorm,
+    WanModel,
+    rope_apply,
+)
 
 logger = init_logger(__name__)
 
@@ -243,26 +253,36 @@ class FusedBlock(nn.Module):
 class FusionModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["fused_blocks"]
 
+    packed_modules_mapping = {
+        "to_qkv": ["q", "k", "v"],
+    }
+
     @staticmethod
     def _is_fused_block(name: str, module) -> bool:
         return "fused_blocks" in name and name.split(".")[-1].isdigit()
 
     _hsdp_shard_conditions = [_is_fused_block]
 
-    def __init__(self, video_config=None, audio_config=None):
+    def __init__(
+        self,
+        video_config=None,
+        audio_config=None,
+        *,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
         super().__init__()
         has_video = True
         has_audio = True
         self.device = get_local_device()
         if video_config is not None:
-            self.video_model = WanModel(**video_config)
+            self.video_model = WanModel(quant_config=quant_config, prefix="video_model", **video_config)
         else:
             has_video = False
             self.video_model = None
             logger.warning("No video model is provided!")
 
         if audio_config is not None:
-            self.audio_model = WanModel(**audio_config)
+            self.audio_model = WanModel(quant_config=quant_config, prefix="audio_model", **audio_config)
         else:
             has_audio = False
             self.audio_model = None
@@ -272,11 +292,12 @@ class FusionModel(nn.Module):
             assert len(self.video_model.blocks) == len(self.audio_model.blocks)
             self.num_blocks = len(self.video_model.blocks)
 
-            self.inject_cross_attention_kv_projections()
+            self.inject_cross_attention_kv_projections(quant_config=quant_config)
 
-        self.num_heads = self.video_model.num_heads
+        tp_size = get_tensor_model_parallel_world_size()
+        self.full_num_heads = self.video_model.num_heads
+        self.num_heads = self.full_num_heads // tp_size
         self.head_dim = self.video_model.dim // self.video_model.num_heads
-        # Make a single shared instance to pass in at forward time
         self.attn = Attention(
             num_heads=self.num_heads,
             head_size=self.head_dim,
@@ -315,21 +336,42 @@ class FusionModel(nn.Module):
 
         return super().load_state_dict(state_dict, strict=strict, assign=assign)
 
-    def inject_cross_attention_kv_projections(self):
-        for vid_block in self.video_model.blocks:
-            vid_block.cross_attn.k_fusion = nn.Linear(vid_block.dim, vid_block.dim)
-            vid_block.cross_attn.v_fusion = nn.Linear(vid_block.dim, vid_block.dim)
-            vid_block.cross_attn.pre_attn_norm_fusion = WanLayerNorm(vid_block.dim, elementwise_affine=True)
-            vid_block.cross_attn.norm_k_fusion = (
-                WanRMSNorm(vid_block.dim, eps=1e-6) if vid_block.qk_norm else nn.Identity()
+    def inject_cross_attention_kv_projections(
+        self,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
+        def _make_kv_proj(dim: int, prefix: str) -> ColumnParallelLinear:
+            # K/V projections for the cross-modal (a2v / v2a) attention path.
+            # Inputs are full latents, so we quantize matched to self_attn.
+            return ColumnParallelLinear(
+                dim,
+                dim,
+                bias=True,
+                gather_output=False,
+                return_bias=False,
+                quant_config=quant_config,
+                prefix=prefix,
             )
 
-        for audio_block in self.audio_model.blocks:
-            audio_block.cross_attn.k_fusion = nn.Linear(audio_block.dim, audio_block.dim)
-            audio_block.cross_attn.v_fusion = nn.Linear(audio_block.dim, audio_block.dim)
-            audio_block.cross_attn.pre_attn_norm_fusion = WanLayerNorm(audio_block.dim, elementwise_affine=True)
-            audio_block.cross_attn.norm_k_fusion = (
-                WanRMSNorm(audio_block.dim, eps=1e-6) if audio_block.qk_norm else nn.Identity()
+        tp_size = get_tensor_model_parallel_world_size()
+        for i, vid_block in enumerate(self.video_model.blocks):
+            cross = vid_block.cross_attn
+            block_prefix = f"video_model.blocks.{i}.cross_attn"
+            cross.k_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.k_fusion")
+            cross.v_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.v_fusion")
+            cross.pre_attn_norm_fusion = WanLayerNorm(vid_block.dim, elementwise_affine=True)
+            cross.norm_k_fusion = (
+                DistributedRMSNorm(vid_block.dim // tp_size, eps=1e-6) if vid_block.qk_norm else nn.Identity()
+            )
+
+        for i, audio_block in enumerate(self.audio_model.blocks):
+            cross = audio_block.cross_attn
+            block_prefix = f"audio_model.blocks.{i}.cross_attn"
+            cross.k_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.k_fusion")
+            cross.v_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.v_fusion")
+            cross.pre_attn_norm_fusion = WanLayerNorm(audio_block.dim, elementwise_affine=True)
+            cross.norm_k_fusion = (
+                DistributedRMSNorm(audio_block.dim // tp_size, eps=1e-6) if audio_block.qk_norm else nn.Identity()
             )
 
     def _detach_blocks_from_backbones(self) -> None:

vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py

@@ -13,6 +13,11 @@ from PIL import Image, ImageOps
 from torch import nn
 from torchvision.transforms import Compose, Normalize
 from tqdm import tqdm
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
@@ -124,7 +129,8 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor
         self.text_encoder = self.text_model.model
 
         # Fusion model — weights are loaded later via load_weights()
-        self.model = FusionModel(VIDEO_CONFIG, AUDIO_CONFIG)
+        quant_config = getattr(self.od_config, "quantization_config", None)
+        self.model = FusionModel(VIDEO_CONFIG, AUDIO_CONFIG, quant_config=quant_config)
         self.transformer = self.model
 
         fusion_path = self.od_config.tf_model_config.get("fusion", None)
@@ -236,10 +242,75 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor
         return ref_vae_latents, ref_audio_lengths
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        prefix = "model."
-        state_dict = {name[len(prefix) :]: tensor for name, tensor in weights if name.startswith(prefix)}
-        self.model.load_state_dict(state_dict, strict=True)
-        return {prefix + k for k in state_dict}
+        pipeline_prefix = "model."
+
+        qkv_stacked_mapping = [
+            (".self_attn.to_qkv", ".self_attn.q", "q"),
+            (".self_attn.to_qkv", ".self_attn.k", "k"),
+            (".self_attn.to_qkv", ".self_attn.v", "v"),
+        ]
+
+        params_dict = dict(self.model.named_parameters())
+        loaded_local: set[str] = set()
+
+        tp_size = get_tensor_model_parallel_world_size()
+        tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0
+
+        def _maybe_shard_weight(weight: torch.Tensor, param: torch.Tensor) -> torch.Tensor:
+            """Slice a checkpoint tensor to this rank's TP shard when the
+            param is a TP-local shape but the checkpoint stores the full one.
+            """
+            if tp_size <= 1 or weight.shape == param.shape:
+                return weight
+            if weight.ndim == 1 and weight.numel() == param.numel() * tp_size:
+                return weight.chunk(tp_size, dim=0)[tp_rank]
+            if weight.ndim == 2:
+                if weight.shape[0] == param.shape[0] * tp_size:
+                    return weight.chunk(tp_size, dim=0)[tp_rank]
+                if weight.shape[1] == param.shape[1] * tp_size:
+                    return weight.chunk(tp_size, dim=1)[tp_rank]
+            return weight
+
+        for full_name, loaded_weight in weights:
+            if not full_name.startswith(pipeline_prefix):
+                continue
+            name = full_name[len(pipeline_prefix) :]
+
+            # Fused-QKV stacking: route checkpoint's separate
+            # - ...self_attn.{q,k,v}: weights into the corresponding
+            # - ...self_attn.to_qkv : via its sharded weight_loader.
+            stacked_hit = False
+            for fused_name, raw_name, shard_id in qkv_stacked_mapping:
+                if raw_name not in name:
+                    continue
+                stacked = name.replace(raw_name, fused_name)
+                if stacked not in params_dict:
+                    continue
+                param = params_dict[stacked]
+                param.weight_loader(param, loaded_weight, shard_id)
+                loaded_local.add(stacked)
+                stacked_hit = True
+                break
+            if stacked_hit:
+                continue
+
+            if name not in params_dict:
+                logger.warning(f"Skipping DreamID-Omni weight {full_name} -- not found in model parameters")
+                continue
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader", None)
+            if weight_loader is not None:
+                # Parallel Linears handle TP sharding internally
+                weight_loader(param, loaded_weight)
+            else:
+                default_weight_loader(param, _maybe_shard_weight(loaded_weight, param))
+            loaded_local.add(name)
+
+        # Detach must happen before `DiffusersPipelineLoader._process_weights_after_loading`
+        # and before any forward pass that triggers layerwise-offload hooks.
+        self.model._detach_blocks_from_backbones()
+
+        return {pipeline_prefix + k for k in loaded_local}
 
     def get_scheduler_time_steps(self, sampling_steps, solver_name="unipc", device=0, shift=5.0):
         torch.manual_seed(4)

vllm_omni/diffusion/models/dreamid_omni/wan2_2.py

@@ -6,6 +6,16 @@ import torch.amp as amp
 import torch.nn as nn
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from vllm.distributed import (
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_reduce,
+)
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 from vllm_omni.diffusion.attention.layer import Attention
 
@@ -17,7 +27,6 @@ try:
         MLPProj,
         ModulationAdd,
         WanLayerNorm,
-        WanRMSNorm,
         rope_apply,  # diff from  wan2.2
         rope_params,
         sinusoidal_embedding_1d,
@@ -26,42 +35,99 @@ except ImportError:
     raise ImportError("Failed to import from dependency 'dreamid_omni'.")
 
 
+class DistributedRMSNorm(nn.Module):
+    """RMSNorm that computes global RMS across tensor parallel ranks.
+
+    Mirrors `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py::DistributedRMSNorm`
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        tp_size = get_tensor_model_parallel_world_size()
+        x_dtype = x.dtype
+        x_float = x.float()
+        local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True)
+        local_count = x.shape[-1]
+
+        if tp_size > 1:
+            local_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq)
+            local_count = local_count * tp_size
+
+        rms = torch.sqrt(local_sum_sq / local_count + self.eps)
+        return (x_float / rms * self.weight.float()).to(x_dtype)
+
+
 class WanSelfAttention(nn.Module):
     """Optimized self-attention module using vLLM layers."""
 
-    def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        window_size=(-1, -1),
+        qk_norm=True,
+        eps=1e-6,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         assert dim % num_heads == 0
         super().__init__()
         self.dim = dim
-        self.num_heads = num_heads
         self.head_dim = dim // num_heads
         self.window_size = window_size
         self.qk_norm = qk_norm
         self.eps = eps
 
-        # layers
-        self.q = nn.Linear(dim, dim)
-        self.k = nn.Linear(dim, dim)
-        self.v = nn.Linear(dim, dim)
-        self.o = nn.Linear(dim, dim)
-        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_qkv" if prefix else "to_qkv",
+        )
+        self.num_heads = self.to_qkv.num_heads
+        self.num_kv_heads = self.to_qkv.num_kv_heads
+        self.tp_inner_dim = self.num_heads * self.head_dim
+        self.tp_kv_dim = self.num_kv_heads * self.head_dim
+
+        self.o = RowParallelLinear(
+            dim,
+            dim,
+            bias=True,
+            input_is_parallel=True,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o" if prefix else "o",
+        )
+        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = DistributedRMSNorm(self.tp_kv_dim, eps=eps) if qk_norm else nn.Identity()
         # Unified attention layer
         self.attn = Attention(
             num_heads=self.num_heads,
             head_size=self.head_dim,
-            num_kv_heads=self.num_heads,
+            num_kv_heads=self.num_kv_heads,
             softmax_scale=1.0 / (self.head_dim**0.5),
             causal=False,
+            prefix=prefix,
         )
 
     # query, key, value function
     def qkv_fn(self, x):
-        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+        b, s = x.shape[:2]
+        qkv, _ = self.to_qkv(x)
+        q_size = self.num_heads * self.head_dim
+        kv_size = self.num_kv_heads * self.head_dim
+        q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
 
-        q = self.norm_q(self.q(x)).view(b, s, n, d)
-        k = self.norm_k(self.k(x)).view(b, s, n, d)
-        v = self.v(x).view(b, s, n, d)
+        q = self.norm_q(q).view(b, s, self.num_heads, self.head_dim)
+        k = self.norm_k(k).view(b, s, self.num_kv_heads, self.head_dim)
+        v = v.view(b, s, self.num_kv_heads, self.head_dim)
         return q, k, v
 
     def forward(self, x, seq_lens, grid_sizes, freqs, ref_lengths=None, freqs_scaling=None):
@@ -87,14 +153,97 @@ class WanSelfAttention(nn.Module):
         return x
 
 
-class WanT2VCrossAttention(WanSelfAttention):
-    def qkv_fn(self, x, context):
-        b, n, d = x.size(0), self.num_heads, self.head_dim
+class WanT2VCrossAttention(nn.Module):
+    """Text -> latent cross-attention."""
 
-        # compute query, key, value
-        q = self.norm_q(self.q(x)).view(b, -1, n, d)
-        k = self.norm_k(self.k(context)).view(b, -1, n, d)
-        v = self.v(context).view(b, -1, n, d)
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        window_size=(-1, -1),
+        qk_norm=True,
+        eps=1e-6,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        assert dim % num_heads == 0
+        self.dim = dim
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # Separate Q/K/V — Q's input differs from K/V's input (text context)
+        self.q = ColumnParallelLinear(
+            dim,
+            dim,
+            bias=True,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.q" if prefix else "q",
+        )
+        self.k = ColumnParallelLinear(
+            dim,
+            dim,
+            bias=True,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.k" if prefix else "k",
+        )
+        self.v = ColumnParallelLinear(
+            dim,
+            dim,
+            bias=True,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.v" if prefix else "v",
+        )
+
+        tp_size = get_tensor_model_parallel_world_size()
+        self.num_heads = num_heads // tp_size
+        self.num_kv_heads = self.num_heads
+        self.tp_inner_dim = self.num_heads * self.head_dim
+
+        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
+
+        self.o = RowParallelLinear(
+            dim,
+            dim,
+            bias=True,
+            input_is_parallel=True,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o" if prefix else "o",
+        )
+
+        # Kept for interface parity with WanSelfAttention;
+        # the active cross-attention path in DreamID-Omni uses the shared FusionModel.attn.
+        # Disable kv quant to prevent affecting quality.
+        self.attn = Attention(
+            num_heads=self.num_heads,
+            head_size=self.head_dim,
+            num_kv_heads=self.num_kv_heads,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+            prefix=prefix,
+            disable_kv_quant=True,
+        )
+
+    def qkv_fn(self, x, context):
+        b = x.size(0)
+
+        q = self.q(x)
+        k = self.k(context)
+        v = self.v(context)
+        q = self.norm_q(q).view(b, -1, self.num_heads, self.head_dim)
+        k = self.norm_k(k).view(b, -1, self.num_kv_heads, self.head_dim)
+        v = v.view(b, -1, self.num_kv_heads, self.head_dim)
 
         return q, k, v
 
@@ -142,6 +291,9 @@ class WanAttentionBlock(nn.Module):
         cross_attn_norm=False,
         eps=1e-6,
         additional_emb_length=None,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.dim = dim
@@ -154,7 +306,15 @@ class WanAttentionBlock(nn.Module):
 
         # layers
         self.norm1 = WanLayerNorm(dim, eps)
-        self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, eps)
+        self.self_attn = WanSelfAttention(
+            dim,
+            num_heads,
+            window_size,
+            qk_norm,
+            eps,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn" if prefix else "self_attn",
+        )
         self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
         if cross_attn_type == "i2v_cross_attn":
             assert additional_emb_length is not None, "additional_emb_length should be specified for i2v_cross_attn"
@@ -167,9 +327,32 @@ class WanAttentionBlock(nn.Module):
                 (-1, -1),
                 qk_norm,
                 eps,
+                quant_config=quant_config,
+                prefix=f"{prefix}.cross_attn" if prefix else "cross_attn",
             )
         self.norm2 = WanLayerNorm(dim, eps)
-        self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim))
+        # FFN layout: Stay compatible with upstream DreamID-Omni checkpoints
+        self.ffn = nn.Sequential(
+            ColumnParallelLinear(
+                dim,
+                ffn_dim,
+                bias=True,
+                gather_output=False,
+                return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.ffn.0" if prefix else "ffn.0",
+            ),
+            nn.GELU(approximate="tanh"),
+            RowParallelLinear(
+                ffn_dim,
+                dim,
+                bias=True,
+                input_is_parallel=True,
+                return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.ffn.2" if prefix else "ffn.2",
+            ),
+        )
 
         # modulation
         self.modulation = ModulationAdd(dim, 6)
@@ -243,6 +426,9 @@ class WanModel(ModelMixin, ConfigMixin):
         cross_attn_norm=True,
         temporal_rope_scaling_factor=1.0,
         eps=1e-6,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ):
         r"""
         Initialize the diffusion model backbone.
@@ -357,8 +543,10 @@ class WanModel(ModelMixin, ConfigMixin):
                     cross_attn_norm,
                     eps,
                     additional_emb_length,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.blocks.{layer_idx}" if prefix else f"blocks.{layer_idx}",
                 )
-                for _ in range(num_layers)
+                for layer_idx in range(num_layers)
             ]
         )