From 2e053d6eb64a5a8a0aea12ec868bf6864abdb318 Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 24 Feb 2026 14:20:54 +0800 Subject: [PATCH] [diffusion] quant: support quant for all dits (#19156) Co-authored-by: zyzshishui --- .../multimodal_gen/configs/models/base.py | 5 -- .../multimodal_gen/docs/quantization.md | 6 +- .../component_loaders/transformer_loader.py | 32 +++---- .../loader/component_loaders/vae_loader.py | 1 - .../runtime/loader/fsdp_load.py | 26 ++++++ .../runtime/models/dits/causal_wanvideo.py | 25 ++++-- .../runtime/models/dits/flux_2.py | 64 ++++++++++--- .../runtime/models/dits/glm_image.py | 28 ++++-- .../runtime/models/dits/hunyuanvideo.py | 39 +++++++- .../runtime/models/dits/ltx_2.py | 75 +++++++++++++--- .../runtime/models/dits/mova_audio_dit.py | 29 ++++-- .../runtime/models/dits/mova_video_dit.py | 90 +++++++++++++++---- .../runtime/models/dits/qwen_image.py | 72 +++++++++++---- .../runtime/models/dits/wanvideo.py | 87 ++++++++++++++---- .../runtime/models/dits/zimage.py | 23 ++++- .../runtime/utils/hf_diffusers_utils.py | 1 - .../test/server/perf_baselines.json | 27 +++++- .../test/server/testcase_configs.py | 15 +++- .../multimodal_gen/tools/convert_hf_to_fp8.py | 11 ++- 19 files changed, 523 insertions(+), 133 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/models/base.py b/python/sglang/multimodal_gen/configs/models/base.py index 6de428ad9..6619117be 100644 --- 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__() diff --git a/python/sglang/multimodal_gen/docs/quantization.md b/python/sglang/multimodal_gen/docs/quantization.md index 542cd0292..b8588a46d 100644 --- 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 diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index 344fc6d38..3c18e16c9 100644 --- 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( diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index a237b11df..e3a9dadf4 100644 --- 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" diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 0ecb2c955..00c20138d 100644 --- 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: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py index 159f06e5d..279aca9f2 100644 --- 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) ] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index a5be6184d..bbf7cf489 100644 --- 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) ] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py index e6ab13a1c..f7f223dec 100644 --- 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) ] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py index 8ef9162d1..09a233ec9 100644 --- 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) ] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index e3383865c..7dc486320 100644 --- 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 // ( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py b/python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py index 7568bbb35..2604c0ed2 100644 --- 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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py b/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py index 31d43a458..76a5267f0 100644 --- 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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 52f1edc4d..8ae3fe74e 100644 --- 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 = { diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index f5ff18c5f..ba0074cb5 100644 --- 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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index 650245c3c..ae0e421b6 100644 --- 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: diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index 4f0ba6c5a..fd8bccd39 100644 --- 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) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 4b9ee1760..f1e062208 100644 --- 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, diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 7bf5baad3..d57c1ae34 100644 --- 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", diff --git a/python/sglang/multimodal_gen/tools/convert_hf_to_fp8.py b/python/sglang/multimodal_gen/tools/convert_hf_to_fp8.py index 6f14035f0..7c3e0051c 100644 --- 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()