chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] feat: support nunchaku for Z-Image-Turbo and flux.1 (int4) (#18959)

Browse Source
This commit is contained in:
Mick
2026-02-20 21:16:08 +08:00
committed by GitHub
parent 7d953440ec
commit 8d789b5c3d
9 changed files with 664 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
python/sglang/multimodal_gen/configs/models/dits/flux.py
View File

@@ -23,9 +23,36 @@ class FluxArchConfig(DiTArchConfig):
stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list)
# nunchaku checkpoint uses different weight names; map to sglang flux layout
param_names_mapping: dict = field(
default_factory=lambda: {
r"transformer\.(\w*)\.(.*)$": r"\1.\2",
# HF diffusers format
r"^transformer\.(\w*)\.(.*)$": r"\1.\2",
# transformer_blocks nunchaku format (raw export - before internal conversion)
r"^transformer_blocks\.(\d+)\.mlp_fc1\.(.*)$": r"transformer_blocks.\1.ff.net.0.proj.\2",
r"^transformer_blocks\.(\d+)\.mlp_fc2\.(.*)$": r"transformer_blocks.\1.ff.net.2.\2",
r"^transformer_blocks\.(\d+)\.mlp_context_fc1\.(.*)$": r"transformer_blocks.\1.ff_context.net.0.proj.\2",
r"^transformer_blocks\.(\d+)\.mlp_context_fc2\.(.*)$": r"transformer_blocks.\1.ff_context.net.2.\2",
r"^transformer_blocks\.(\d+)\.qkv_proj\.(.*)$": r"transformer_blocks.\1.attn.to_qkv.\2",
r"^transformer_blocks\.(\d+)\.qkv_proj_context\.(.*)$": r"transformer_blocks.\1.attn.to_added_qkv.\2",
r"^transformer_blocks\.(\d+)\.out_proj\.(.*)$": r"transformer_blocks.\1.attn.to_out.0.\2",
r"^transformer_blocks\.(\d+)\.out_proj_context\.(.*)$": r"transformer_blocks.\1.attn.to_add_out.\2",
r"^transformer_blocks\.(\d+)\.norm_q\.(.*)$": r"transformer_blocks.\1.attn.norm_q.\2",
r"^transformer_blocks\.(\d+)\.norm_k\.(.*)$": r"transformer_blocks.\1.attn.norm_k.\2",
r"^transformer_blocks\.(\d+)\.norm_added_q\.(.*)$": r"transformer_blocks.\1.attn.norm_added_q.\2",
r"^transformer_blocks\.(\d+)\.norm_added_k\.(.*)$": r"transformer_blocks.\1.attn.norm_added_k.\2",
# transformer_blocks nunchaku format (already converted with convert_flux_state_dict)
r"^transformer_blocks\.(\d+)\.attn\.add_qkv_proj\.(.*)$": r"transformer_blocks.\1.attn.to_added_qkv.\2",
# single_transformer_blocks nunchaku format (raw export - before internal conversion)
r"^single_transformer_blocks\.(\d+)\.qkv_proj\.(.*)$": r"single_transformer_blocks.\1.attn.to_qkv.\2",
r"^single_transformer_blocks\.(\d+)\.out_proj\.(.*)$": r"single_transformer_blocks.\1.attn.to_out.0.\2",
r"^single_transformer_blocks\.(\d+)\.norm_q\.(.*)$": r"single_transformer_blocks.\1.attn.norm_q.\2",
r"^single_transformer_blocks\.(\d+)\.norm_k\.(.*)$": r"single_transformer_blocks.\1.attn.norm_k.\2",
# nunchaku quantization parameter name conversions (apply to all blocks)
r"^(.*)\.smooth_orig$": r"\1.smooth_factor_orig",
r"^(.*)\.smooth$": r"\1.smooth_factor",
r"^(.*)\.lora_down$": r"\1.proj_down",
r"^(.*)\.lora_up$": r"\1.proj_up",
}
)

59
python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py
View File

@@ -16,27 +16,6 @@ from .base_config import QuantizationConfig, QuantizeMethodBase
logger = init_logger(__name__)
SVDQ_W4A4_LAYER_PATTERNS = [
"attn.to_qkv",
"attn.to_out",
"attn.add_qkv_proj",
"attn.to_add_out",
"img_mlp",
"txt_mlp",
]
AWQ_W4A16_LAYER_PATTERNS = [
"img_mod",
"txt_mod",
]
SKIP_QUANTIZATION_PATTERNS = [
"norm",
"embed",
"rotary",
"pos_embed",
]
@lru_cache(maxsize=1)
def is_nunchaku_available() -> bool:
@@ -61,13 +40,15 @@ class NunchakuConfig(QuantizationConfig):
group_size: Quantization group size (automatically set based on precision)
act_unsigned: Use unsigned activation quantization
quantized_model_path: Path to pre-quantized model weights (.safetensors)
model_cls: DiT model class that provides quantization rules via get_nunchaku_quant_rules()
"""
precision: str = "int4" # "int4" or "nvfp4"
precision: str = "int4"
rank: int = 32
group_size: Optional[int] = None
act_unsigned: bool = False
quantized_model_path: Optional[str] = None
model_cls: Optional[type] = None
@classmethod
def get_name(cls) -> str:
@@ -99,15 +80,27 @@ class NunchakuConfig(QuantizationConfig):
def get_quant_method(
self, layer: torch.nn.Module, prefix: str
) -> Optional[QuantizeMethodBase]:
if not isinstance(layer, LinearBase):
return None
for pattern in SKIP_QUANTIZATION_PATTERNS:
# get quantization rules from model class
quant_rules = self._get_quant_rules()
# priority: skip > awq_w4a16 > svdq_w4a4 > default
skip_patterns = quant_rules.get("skip", [])
for pattern in skip_patterns:
if pattern in prefix.lower():
return None
for pattern in SVDQ_W4A4_LAYER_PATTERNS:
awq_patterns = quant_rules.get("awq_w4a16", [])
for pattern in awq_patterns:
if pattern in prefix:
from ..nunchaku_linear import NunchakuAWQLinearMethod
return NunchakuAWQLinearMethod(group_size=64)
svdq_patterns = quant_rules.get("svdq_w4a4", [])
for pattern in svdq_patterns:
if pattern in prefix:
from ..nunchaku_linear import NunchakuSVDQLinearMethod
@@ -117,14 +110,7 @@ class NunchakuConfig(QuantizationConfig):
act_unsigned=self.act_unsigned,
)
for pattern in AWQ_W4A16_LAYER_PATTERNS:
if pattern in prefix:
from ..nunchaku_linear import NunchakuAWQLinearMethod
return NunchakuAWQLinearMethod(
group_size=64,
)
# default: apply svdq_w4a4 to all remaining linear layers
from ..nunchaku_linear import NunchakuSVDQLinearMethod
return NunchakuSVDQLinearMethod(
@@ -133,6 +119,13 @@ class NunchakuConfig(QuantizationConfig):
act_unsigned=self.act_unsigned,
)
def _get_quant_rules(self) -> dict[str, list[str]]:
if self.model_cls is not None and hasattr(
self.model_cls, "get_nunchaku_quant_rules"
):
return self.model_cls.get_nunchaku_quant_rules()
return {}
def __post_init__(self):
if self.group_size is None:
if self.precision == "nvfp4":

5
python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
View File

@@ -94,8 +94,9 @@ class TransformerLoader(ComponentLoader):
model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
nunchaku_config = server_args.nunchaku_config
if nunchaku_config is not None:
nunchaku_config.model_cls = model_cls
# respect dtype from checkpoint
# TODO: improve the condition
param_dtype = None
@@ -158,7 +159,7 @@ class TransformerLoader(ComponentLoader):
logger.info("Loaded model with %.2fB parameters", total_params / 1e9)
# considering the existent of mixed-precision models (e.g., nunchaku)
if next(model.parameters()).dtype != param_dtype:
if next(model.parameters()).dtype != param_dtype and param_dtype:
logger.warning(
f"Model dtype does not match expected param dtype, {next(model.parameters()).dtype} vs {param_dtype}"
)

5
python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
View File

@@ -238,7 +238,8 @@ def load_model_from_full_model_state_dict(
"""
meta_sd = model.state_dict()
param_dict = dict(model.named_parameters())
sharded_sd = {}
# map names from checkpoint to customized names
custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
full_sd_iterator, param_names_mapping
) # type: ignore
@@ -250,7 +251,7 @@ def load_model_from_full_model_state_dict(
# sort parameter names to ensure all ranks process parameters in the same order
sorted_param_names = sorted(custom_param_sd.keys())
requires_grad = False
sharded_sd = {}
# shard from loaded state_dict, custom_param_sd -> sharded_sd
for target_param_name in sorted_param_names:

79
python/sglang/multimodal_gen/runtime/loader/utils.py
View File

@@ -31,7 +31,7 @@ def set_default_torch_dtype(dtype: torch.dtype):
def get_param_names_mapping(
mapping_dict: dict[str, str],
mapping_dict: dict[str, str | tuple[str, int, int]],
) -> Callable[[str], tuple[str, Any, Any]]:
"""
Creates a mapping function that transforms parameter names using regex patterns.
@@ -44,21 +44,50 @@ def get_param_names_mapping(
"""
def mapping_fn(name: str) -> tuple[str, Any, Any]:
# Try to match and transform the name using the regex patterns in mapping_dict
for pattern, replacement in mapping_dict.items():
match = re.match(pattern, name)
if match:
merge_index = None
total_split_params = None
if isinstance(replacement, tuple):
merge_index = replacement[1]
total_split_params = replacement[2]
replacement = replacement[0]
name = re.sub(pattern, replacement, name)
return name, merge_index, total_split_params
# support chained conversions, e.g.:
# transformer.xxx.lora_down -> xxx.lora_down -> xxx.proj_down
merge_index = None
total_split_params = None
max_steps = max(8, len(mapping_dict) * 2)
applied_patterns: set[str] = set()
visited_names: set[str] = {name}
# If no pattern matches, return the original name
return name, None, None
for _ in range(max_steps):
transformed = False
for pattern, replacement in mapping_dict.items():
# avoid re-applying the same rule on its own output
if pattern in applied_patterns:
continue
if re.match(pattern, name) is None:
continue
curr_merge_index = None
curr_total_split_params = None
if isinstance(replacement, tuple):
curr_merge_index = replacement[1]
curr_total_split_params = replacement[2]
replacement = replacement[0]
new_name = re.sub(pattern, replacement, name)
if new_name != name:
if curr_merge_index is not None:
merge_index = curr_merge_index
total_split_params = curr_total_split_params
name = new_name
applied_patterns.add(pattern)
if name in visited_names:
transformed = False
break
visited_names.add(name)
transformed = True
break
if not transformed:
break
return name, merge_index, total_split_params
return mapping_fn
@@ -150,25 +179,5 @@ def _list_safetensors_files(model_path: str) -> list[str]:
BYTES_PER_GB = 1024**3
def get_memory_usage_of_component(module) -> float | None:
"""
returned value is in GB, rounded to 2 decimal digits
"""
if not isinstance(module, nn.Module):
return None
if hasattr(module, "get_memory_footprint"):
usage = module.get_memory_footprint() / BYTES_PER_GB
else:
# manually
param_size = sum(p.numel() * p.element_size() for p in module.parameters())
buffer_size = sum(b.numel() * b.element_size() for b in module.buffers())
total_size_bytes = param_size + buffer_size
usage = total_size_bytes / (1024**3)
return round(usage, 2)
# component name -> ComponentLoader class
component_name_to_loader_cls: Dict[str, Type[Any]] = {}

14
python/sglang/multimodal_gen/runtime/models/dits/base.py
View File

@@ -107,3 +107,17 @@ class CachableDiT(TeaCacheMixin, BaseDiT):
def __init__(self, config: DiTConfig, **kwargs) -> None:
super().__init__(config, **kwargs)
self._init_teacache_state()
@classmethod
def get_nunchaku_quant_rules(cls) -> dict[str, dict[str, Any]]:
"""
Get quantization rules for Nunchaku quantization.
Returns a dict mapping layer name patterns to quantization configs:
{
"skip": [list of patterns to skip quantization],
"svdq_w4a4": [list of patterns for SVDQ W4A4],
"awq_w4a16": [list of patterns for AWQ W4A16],
}
"""
return {}

490
python/sglang/multimodal_gen/runtime/models/dits/flux.py
View File

@@ -32,8 +32,17 @@ from sglang.multimodal_gen.runtime.layers.attention import USPAttention
# from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm as LayerNorm
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm
from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear
from sglang.multimodal_gen.runtime.layers.mlp import MLP
from sglang.multimodal_gen.runtime.layers.linear import (
ColumnParallelLinear,
MergedColumnParallelLinear,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
QuantizationConfig,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
NunchakuConfig,
is_nunchaku_available,
)
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
NDRotaryEmbedding,
apply_flashinfer_rope_qk_inplace,
@@ -49,19 +58,150 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__) # pylint: disable=invalid-name
try:
from nunchaku.models.attention import NunchakuFeedForward # type: ignore[import]
from nunchaku.models.normalization import ( # type: ignore[import]
NunchakuAdaLayerNormZero,
NunchakuAdaLayerNormZeroSingle,
)
from nunchaku.ops.gemm import (
svdq_gemm_w4a4_cuda as _svdq_gemm_w4a4, # type: ignore[import]
)
from nunchaku.ops.quantize import (
svdq_quantize_w4a4_act_fuse_lora_cuda as _svdq_quantize_w4a4, # type: ignore[import]
)
_nunchaku_fused_ops_available = True
except Exception:
NunchakuFeedForward = None
NunchakuAdaLayerNormZero = None
NunchakuAdaLayerNormZeroSingle = None
_svdq_gemm_w4a4 = None
_svdq_quantize_w4a4 = None
_nunchaku_fused_ops_available = False
def _fused_gelu_mlp(
x: torch.Tensor,
fc1,
fc2,
pad_size: int = 256,
) -> torch.Tensor:
"""
Fused GELU MLP matching nunchaku's fused_gelu_mlp kernel path.
nunchaku's single-block MLP checkpoint is calibrated for the fused path where:
1. fc1 GEMM + GELU + 0.171875 shift + unsigned re-quantization + fc2.lora_down
are all done in a single fused kernel call
2. fc2 GEMM then receives unsigned INT4 activations (act_unsigned=True)
Using the sequential path (fc1 → GELU → fc2 with symmetric quantization) is
fundamentally incompatible with these wscales, causing visually wrong outputs.
"""
batch_size, seq_len, channels = x.shape
x_2d = x.view(batch_size * seq_len, channels)
quantized_x, ascales, lora_act = _svdq_quantize_w4a4(
x_2d,
lora_down=fc1.proj_down,
smooth=fc1.smooth_factor,
fp4=fc1.precision == "nvfp4",
pad_size=pad_size,
)
batch_size_pad = (batch_size * seq_len + pad_size - 1) // pad_size * pad_size
is_fp4 = fc2.precision == "nvfp4"
qout_act = torch.empty(
batch_size_pad,
fc1.output_size_per_partition // 2,
dtype=torch.uint8,
device=x_2d.device,
)
if is_fp4:
qout_ascales = torch.empty(
fc1.output_size_per_partition // 16,
batch_size_pad,
dtype=torch.float8_e4m3fn,
device=x_2d.device,
)
else:
qout_ascales = torch.empty(
fc1.output_size_per_partition // 64,
batch_size_pad,
dtype=x_2d.dtype,
device=x_2d.device,
)
qout_lora_act = torch.empty(
batch_size_pad, fc2.proj_down.shape[1], dtype=torch.float32, device=x_2d.device
)
# fused: fc1 GEMM + GELU + shift + unsigned quantize + fc2.lora_down
_svdq_gemm_w4a4(
act=quantized_x,
wgt=fc1.qweight,
qout=qout_act,
ascales=ascales,
wscales=fc1.wscales,
oscales=qout_ascales,
lora_act_in=lora_act,
lora_up=fc1.proj_up,
lora_down=fc2.proj_down,
lora_act_out=qout_lora_act,
bias=fc1.bias,
smooth_factor=fc2.smooth_factor,
fp4=is_fp4,
alpha=getattr(fc1, "_nunchaku_alpha", None),
wcscales=getattr(fc1, "wcscales", None),
)
output = torch.empty(
batch_size * seq_len,
fc2.output_size_per_partition,
dtype=x_2d.dtype,
device=x_2d.device,
)
# fc2 GEMM with unsigned INT4 activations (fused kernel shifted by 0.171875)
_svdq_gemm_w4a4(
act=qout_act,
wgt=fc2.qweight,
out=output,
ascales=qout_ascales,
wscales=fc2.wscales,
lora_act_in=qout_lora_act,
lora_up=fc2.proj_up,
bias=fc2.bias,
fp4=is_fp4,
alpha=getattr(fc2, "_nunchaku_alpha", None),
wcscales=getattr(fc2, "wcscales", None),
act_unsigned=True,
)
return output.view(batch_size, seq_len, -1)
def _get_qkv_projections(
attn: "FluxAttention", hidden_states, encoder_hidden_states=None
):
query, _ = attn.to_q(hidden_states)
key, _ = attn.to_k(hidden_states)
value, _ = attn.to_v(hidden_states)
if getattr(attn, "use_fused_qkv", False):
qkv, _ = attn.to_qkv(hidden_states)
query, key, value = [x.contiguous() for x in qkv.chunk(3, dim=-1)]
else:
query, _ = attn.to_q(hidden_states)
key, _ = attn.to_k(hidden_states)
value, _ = attn.to_v(hidden_states)
encoder_query = encoder_key = encoder_value = None
if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
encoder_query, _ = attn.add_q_proj(encoder_hidden_states)
encoder_key, _ = attn.add_k_proj(encoder_hidden_states)
encoder_value, _ = attn.add_v_proj(encoder_hidden_states)
if getattr(attn, "use_fused_added_qkv", False):
added_qkv, _ = attn.to_added_qkv(encoder_hidden_states)
encoder_query, encoder_key, encoder_value = [
x.contiguous() for x in added_qkv.chunk(3, dim=-1)
]
else:
encoder_query, _ = attn.add_q_proj(encoder_hidden_states)
encoder_key, _ = attn.add_k_proj(encoder_hidden_states)
encoder_value, _ = attn.add_v_proj(encoder_hidden_states)
return query, key, value, encoder_query, encoder_key, encoder_value
@@ -81,6 +221,8 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
out_dim: int = None,
context_pre_only: Optional[bool] = None,
pre_only: bool = False,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
):
super().__init__()
@@ -96,24 +238,41 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
self.added_kv_proj_dim = added_kv_proj_dim
self.added_proj_bias = added_proj_bias
self.use_fused_qkv = isinstance(quant_config, NunchakuConfig)
self.use_fused_added_qkv = isinstance(quant_config, NunchakuConfig)
self.norm_q = RMSNorm(dim_head, eps=eps)
self.norm_k = RMSNorm(dim_head, eps=eps)
self.to_q = ColumnParallelLinear(
query_dim, self.inner_dim, bias=bias, gather_output=True
)
self.to_k = ColumnParallelLinear(
query_dim, self.inner_dim, bias=bias, gather_output=True
)
self.to_v = ColumnParallelLinear(
query_dim, self.inner_dim, bias=bias, gather_output=True
)
if self.use_fused_qkv:
self.to_qkv = MergedColumnParallelLinear(
query_dim,
[self.inner_dim] * 3,
bias=bias,
gather_output=True,
quant_config=quant_config,
prefix=f"{prefix}.to_qkv" if prefix else "to_qkv",
)
else:
self.to_q = ColumnParallelLinear(
query_dim, self.inner_dim, bias=bias, gather_output=True
)
self.to_k = ColumnParallelLinear(
query_dim, self.inner_dim, bias=bias, gather_output=True
)
self.to_v = ColumnParallelLinear(
query_dim, self.inner_dim, bias=bias, gather_output=True
)
if not self.pre_only:
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,
prefix=f"{prefix}.to_out.0" if prefix else "",
)
)
if dropout != 0.0:
@@ -122,26 +281,41 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
if added_kv_proj_dim is not None:
self.norm_added_q = RMSNorm(dim_head, eps=eps)
self.norm_added_k = RMSNorm(dim_head, eps=eps)
self.add_q_proj = ColumnParallelLinear(
added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
)
self.add_k_proj = ColumnParallelLinear(
added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
)
self.add_v_proj = ColumnParallelLinear(
added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
)
if self.use_fused_added_qkv:
self.to_added_qkv = MergedColumnParallelLinear(
added_kv_proj_dim,
[self.inner_dim] * 3,
bias=added_proj_bias,
gather_output=True,
quant_config=quant_config,
prefix=f"{prefix}.to_added_qkv" if prefix else "to_added_qkv",
)
else:
self.add_q_proj = ColumnParallelLinear(
added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
)
self.add_k_proj = ColumnParallelLinear(
added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
)
self.add_v_proj = ColumnParallelLinear(
added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
)
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,
prefix=f"{prefix}.to_add_out" if prefix else "",
)
self.attn = USPAttention(
@@ -218,13 +392,18 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
],
dim=1,
)
x, _ = self.to_out[0](x)
if len(self.to_out) == 2:
x = self.to_out[1](x)
if not self.pre_only:
x, _ = self.to_out[0](x)
if len(self.to_out) == 2:
x = self.to_out[1](x)
encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states)
return x, encoder_hidden_states
else:
if not self.pre_only:
x, _ = self.to_out[0](x)
if len(self.to_out) == 2:
x = self.to_out[1](x)
return x
@@ -235,28 +414,70 @@ class FluxSingleTransformerBlock(nn.Module):
num_attention_heads: int,
attention_head_dim: int,
mlp_ratio: float = 4.0,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
):
super().__init__()
self.mlp_hidden_dim = int(dim * mlp_ratio)
self.use_nunchaku_structure = isinstance(quant_config, NunchakuConfig)
self.norm = AdaLayerNormZeroSingle(dim)
self.proj_mlp = ColumnParallelLinear(
dim, self.mlp_hidden_dim, bias=True, gather_output=True
)
self.act_mlp = nn.GELU(approximate="tanh")
self.proj_out = ColumnParallelLinear(
dim + self.mlp_hidden_dim, dim, bias=True, gather_output=True
)
self.attn = FluxAttention(
query_dim=dim,
dim_head=attention_head_dim,
num_heads=num_attention_heads,
out_dim=dim,
bias=True,
eps=1e-6,
pre_only=True,
)
if self.use_nunchaku_structure:
self.mlp_fc1 = ColumnParallelLinear(
dim,
self.mlp_hidden_dim,
bias=True,
gather_output=True,
quant_config=quant_config,
prefix=f"{prefix}.mlp_fc1" if prefix else "mlp_fc1",
)
self.act_mlp = nn.GELU(approximate="tanh")
self.mlp_fc2 = ColumnParallelLinear(
self.mlp_hidden_dim,
dim,
bias=True,
gather_output=True,
quant_config=quant_config,
prefix=f"{prefix}.mlp_fc2" if prefix else "mlp_fc2",
)
self.attn = FluxAttention(
query_dim=dim,
dim_head=attention_head_dim,
num_heads=num_attention_heads,
out_dim=dim,
bias=True,
eps=1e-6,
pre_only=False,
quant_config=quant_config,
prefix=f"{prefix}.attn" if prefix else "attn",
)
if NunchakuAdaLayerNormZeroSingle is not None:
self.norm = NunchakuAdaLayerNormZeroSingle(self.norm, scale_shift=0)
else:
self.proj_mlp = ColumnParallelLinear(
dim,
self.mlp_hidden_dim,
bias=True,
gather_output=True,
)
self.act_mlp = nn.GELU(approximate="tanh")
self.proj_out = ColumnParallelLinear(
dim + self.mlp_hidden_dim,
dim,
bias=True,
gather_output=True,
)
self.attn = FluxAttention(
query_dim=dim,
dim_head=attention_head_dim,
num_heads=num_attention_heads,
out_dim=dim,
bias=True,
eps=1e-6,
pre_only=True,
)
def forward(
self,
@@ -271,20 +492,46 @@ class FluxSingleTransformerBlock(nn.Module):
residual = hidden_states
norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
proj_hidden_states, _ = self.proj_mlp(norm_hidden_states)
mlp_hidden_states = self.act_mlp(proj_hidden_states)
joint_attention_kwargs = joint_attention_kwargs or {}
attn_output = self.attn(
x=norm_hidden_states,
freqs_cis=freqs_cis,
**joint_attention_kwargs,
)
hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
gate = gate.unsqueeze(1)
proj_out, _ = self.proj_out(hidden_states)
hidden_states = gate * proj_out
hidden_states = residual + hidden_states
if self.use_nunchaku_structure:
if _nunchaku_fused_ops_available:
mlp_hidden_states = _fused_gelu_mlp(
norm_hidden_states, self.mlp_fc1, self.mlp_fc2
)
else:
mlp_out, _ = self.mlp_fc1(norm_hidden_states)
mlp_hidden_states = self.act_mlp(mlp_out)
mlp_hidden_states, _ = self.mlp_fc2(mlp_hidden_states)
attn_output = self.attn(
x=norm_hidden_states,
freqs_cis=freqs_cis,
**joint_attention_kwargs,
)
if isinstance(attn_output, tuple):
attn_output = attn_output[0]
hidden_states = attn_output + mlp_hidden_states
gate = gate.unsqueeze(1)
hidden_states = gate * hidden_states
hidden_states = residual + hidden_states
else:
proj_hidden_states, _ = self.proj_mlp(norm_hidden_states)
mlp_hidden_states = self.act_mlp(proj_hidden_states)
attn_output = self.attn(
x=norm_hidden_states,
freqs_cis=freqs_cis,
**joint_attention_kwargs,
)
hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
gate = gate.unsqueeze(1)
proj_out, _ = self.proj_out(hidden_states)
hidden_states = gate * proj_out
hidden_states = residual + hidden_states
if hidden_states.dtype == torch.float16:
hidden_states = hidden_states.clip(-65504, 65504)
@@ -303,6 +550,8 @@ class FluxTransformerBlock(nn.Module):
attention_head_dim: int,
qk_norm: str = "rms_norm",
eps: float = 1e-6,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
):
super().__init__()
@@ -318,22 +567,38 @@ class FluxTransformerBlock(nn.Module):
context_pre_only=False,
bias=True,
eps=eps,
quant_config=quant_config,
prefix=f"{prefix}.attn" if prefix else "attn",
)
self.norm2 = LayerNorm(dim, eps=1e-6, elementwise_affine=False)
self.ff = MLP(
input_dim=dim, mlp_hidden_dim=dim * 4, output_dim=dim, act_type="gelu"
)
self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
self.norm2_context = LayerNorm(dim, eps=1e-6, elementwise_affine=False)
self.ff_context = MLP(
input_dim=dim, mlp_hidden_dim=dim * 4, output_dim=dim, act_type="gelu"
)
nunchaku_enabled = (
quant_config is not None
and hasattr(quant_config, "get_name")
and quant_config.get_name() == "svdquant"
and is_nunchaku_available()
and NunchakuFeedForward is not None
)
self.use_nunchaku_structure = nunchaku_enabled
self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
self.ff_context = FeedForward(
dim=dim, dim_out=dim, activation_fn="gelu-approximate"
)
if nunchaku_enabled:
nunchaku_kwargs = {
"precision": quant_config.precision,
"rank": quant_config.rank,
"act_unsigned": quant_config.act_unsigned,
}
self.ff = NunchakuFeedForward(self.ff, **nunchaku_kwargs)
self.ff_context = NunchakuFeedForward(self.ff_context, **nunchaku_kwargs)
if NunchakuAdaLayerNormZero is not None:
self.norm1 = NunchakuAdaLayerNormZero(self.norm1, scale_shift=0)
self.norm1_context = NunchakuAdaLayerNormZero(
self.norm1_context, scale_shift=0
)
def forward(
self,
@@ -369,9 +634,14 @@ class FluxTransformerBlock(nn.Module):
attn_output = gate_msa.unsqueeze(1) * attn_output
hidden_states = hidden_states + attn_output
norm_hidden_states = self.norm2(hidden_states)
norm_hidden_states = (
norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
)
if self.use_nunchaku_structure:
norm_hidden_states = (
norm_hidden_states * scale_mlp[:, None] + shift_mlp[:, None]
)
else:
norm_hidden_states = (
norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
)
ff_output = self.ff(norm_hidden_states)
ff_output = gate_mlp.unsqueeze(1) * ff_output
@@ -385,10 +655,15 @@ class FluxTransformerBlock(nn.Module):
encoder_hidden_states = encoder_hidden_states + context_attn_output
norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
norm_encoder_hidden_states = (
norm_encoder_hidden_states * (1 + c_scale_mlp[:, None])
+ c_shift_mlp[:, None]
)
if self.use_nunchaku_structure:
norm_encoder_hidden_states = (
norm_encoder_hidden_states * c_scale_mlp[:, None] + c_shift_mlp[:, None]
)
else:
norm_encoder_hidden_states = (
norm_encoder_hidden_states * (1 + c_scale_mlp[:, None])
+ c_shift_mlp[:, None]
)
context_ff_output = self.ff_context(norm_encoder_hidden_states)
encoder_hidden_states = (
@@ -433,7 +708,44 @@ class FluxTransformer2DModel(CachableDiT, OffloadableDiTMixin):
param_names_mapping = FluxConfig().arch_config.param_names_mapping
def __init__(self, config: FluxConfig, hf_config: dict[str, Any]) -> None:
@classmethod
def get_nunchaku_quant_rules(cls) -> dict[str, list[str]]:
return {
"skip": [
"norm",
"embed",
"rotary",
"pos_embed",
],
"svdq_w4a4": [
"attn.to_qkv",
"attn.to_out",
"attn.add_qkv_proj",
"attn.to_added_qkv",
"attn.to_add_out",
"img_mlp",
"txt_mlp",
"attention.to_qkv",
"attention.to_out",
"proj_mlp",
"proj_out",
"mlp_fc1",
"mlp_fc2",
"ff.net",
"ff_context.net",
],
"awq_w4a16": [
"img_mod",
"txt_mod",
],
}
def __init__(
self,
config: FluxConfig,
hf_config: dict[str, Any],
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__(config=config, hf_config=hf_config)
self.config = config.arch_config
@@ -471,8 +783,10 @@ class FluxTransformer2DModel(CachableDiT, OffloadableDiTMixin):
dim=self.inner_dim,
num_attention_heads=self.config.num_attention_heads,
attention_head_dim=self.config.attention_head_dim,
quant_config=quant_config,
prefix=f"transformer_blocks.{i}",
)
for _ in range(self.config.num_layers)
for i in range(self.config.num_layers)
]
)
@@ -482,8 +796,10 @@ class FluxTransformer2DModel(CachableDiT, OffloadableDiTMixin):
dim=self.inner_dim,
num_attention_heads=self.config.num_attention_heads,
attention_head_dim=self.config.attention_head_dim,
quant_config=quant_config,
prefix=f"single_transformer_blocks.{i}",
)
for _ in range(self.config.num_single_layers)
for i in range(self.config.num_single_layers)
]
)

23
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
View File

@@ -1123,6 +1123,29 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
param_names_mapping = QwenImageDitConfig().arch_config.param_names_mapping
@classmethod
def get_nunchaku_quant_rules(cls) -> dict[str, list[str]]:
return {
"skip": [
"norm",
"embed",
"rotary",
"pos_embed",
],
"svdq_w4a4": [
"attn.to_qkv",
"attn.to_out",
"attn.add_qkv_proj",
"attn.to_add_out",
"img_mlp",
"txt_mlp",
],
"awq_w4a16": [
"img_mod",
"txt_mod",
],
}
def __init__(
self,
config: QwenImageDitConfig,

132
python/sglang/multimodal_gen/runtime/models/dits/zimage.py
View File

@@ -15,6 +15,13 @@ from sglang.multimodal_gen.runtime.layers.linear import (
ReplicatedLinear,
RowParallelLinear,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
QuantizationConfig,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
NunchakuConfig,
is_nunchaku_available,
)
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
_apply_rotary_emb,
apply_flashinfer_rope_qk_inplace,
@@ -24,6 +31,11 @@ from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
try:
from nunchaku.models.attention import NunchakuFeedForward # type: ignore[import]
except Exception:
NunchakuFeedForward = None
logger = init_logger(__name__)
_is_cuda = current_platform.is_cuda()
@@ -111,6 +123,8 @@ class ZImageAttention(nn.Module):
num_kv_heads: int,
qk_norm: bool = True,
eps: float = 1e-6,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__()
self.dim = dim
@@ -129,13 +143,26 @@ class ZImageAttention(nn.Module):
self.local_num_heads = num_heads // tp_size
self.local_num_kv_heads = num_kv_heads // tp_size
self.to_q = ColumnParallelLinear(dim, dim, bias=False, gather_output=False)
self.to_k = ColumnParallelLinear(
dim, self.head_dim * num_kv_heads, bias=False, gather_output=False
)
self.to_v = ColumnParallelLinear(
dim, self.head_dim * num_kv_heads, bias=False, gather_output=False
)
kv_dim = self.head_dim * num_kv_heads
self.use_fused_qkv = isinstance(quant_config, NunchakuConfig)
if self.use_fused_qkv:
self.to_qkv = MergedColumnParallelLinear(
dim,
[dim, kv_dim, kv_dim],
bias=False,
gather_output=False,
quant_config=quant_config,
prefix=f"{prefix}.to_qkv",
)
else:
self.to_q = ColumnParallelLinear(dim, dim, bias=False, gather_output=False)
self.to_k = ColumnParallelLinear(
dim, kv_dim, bias=False, gather_output=False
)
self.to_v = ColumnParallelLinear(
dim, kv_dim, bias=False, gather_output=False
)
if self.qk_norm:
self.norm_q = RMSNorm(self.head_dim, eps=eps)
@@ -145,7 +172,16 @@ class ZImageAttention(nn.Module):
self.norm_k = None
self.to_out = nn.ModuleList(
[RowParallelLinear(dim, dim, bias=False, input_is_parallel=True)]
[
RowParallelLinear(
dim,
dim,
bias=False,
input_is_parallel=True,
quant_config=quant_config,
prefix=f"{prefix}.to_out.0",
)
]
)
self.attn = USPAttention(
@@ -162,9 +198,23 @@ class ZImageAttention(nn.Module):
hidden_states: torch.Tensor,
freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
):
q, _ = self.to_q(hidden_states)
k, _ = self.to_k(hidden_states)
v, _ = self.to_v(hidden_states)
if self.use_fused_qkv:
qkv, _ = self.to_qkv(hidden_states)
q, k, v = qkv.split(
[
self.local_num_heads * self.head_dim,
self.local_num_kv_heads * self.head_dim,
self.local_num_kv_heads * self.head_dim,
],
dim=-1,
)
q = q.contiguous()
k = k.contiguous()
v = v.contiguous()
else:
q, _ = self.to_q(hidden_states)
k, _ = self.to_k(hidden_states)
v, _ = self.to_v(hidden_states)
q = q.view(*q.shape[:-1], self.local_num_heads, self.head_dim)
k = k.view(*k.shape[:-1], self.local_num_kv_heads, self.head_dim)
v = v.view(*v.shape[:-1], self.local_num_kv_heads, self.head_dim)
@@ -214,6 +264,8 @@ class ZImageTransformerBlock(nn.Module):
norm_eps: float,
qk_norm: bool,
modulation=True,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
):
super().__init__()
self.dim = dim
@@ -227,9 +279,37 @@ class ZImageTransformerBlock(nn.Module):
num_kv_heads=n_kv_heads,
qk_norm=qk_norm,
eps=1e-5,
quant_config=quant_config,
prefix=f"{prefix}.attention",
)
self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8))
hidden_dim = int(dim / 3 * 8)
nunchaku_enabled = (
isinstance(quant_config, NunchakuConfig) and is_nunchaku_available()
)
if nunchaku_enabled:
import diffusers
ff = diffusers.models.attention.FeedForward(
dim=dim,
dim_out=dim,
activation_fn="swiglu",
inner_dim=hidden_dim,
bias=False,
)
nunchaku_kwargs = {
"precision": quant_config.precision,
"rank": quant_config.rank,
"act_unsigned": quant_config.act_unsigned,
}
self.feed_forward = NunchakuFeedForward(ff, **nunchaku_kwargs)
# NunchakuFeedForward overrides net[2].act_unsigned=True for int4 (GELU-specific
# optimization for non-negative activations). Z-Image uses SwiGLU whose output
# can be negative, so we must restore the original act_unsigned value.
if hasattr(self.feed_forward, "net") and len(self.feed_forward.net) > 2:
self.feed_forward.net[2].act_unsigned = quant_config.act_unsigned
else:
self.feed_forward = FeedForward(dim=dim, hidden_dim=hidden_dim)
self.attention_norm1 = RMSNorm(dim, eps=norm_eps)
self.ffn_norm1 = RMSNorm(dim, eps=norm_eps)
@@ -389,10 +469,32 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
ZImageDitConfig().arch_config.reverse_param_names_mapping
)
@classmethod
def get_nunchaku_quant_rules(cls) -> dict[str, list[str]]:
return {
"skip": [
"norm",
"embed",
"rotary",
"pos_embed",
],
"svdq_w4a4": [
"attention.to_qkv",
"attention.to_out",
"img_mlp",
"txt_mlp",
],
"awq_w4a16": [
"img_mod",
"txt_mod",
],
}
def __init__(
self,
config: ZImageDitConfig,
hf_config: dict[str, Any],
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__(config=config, hf_config=hf_config)
@@ -443,6 +545,8 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
arch_config.norm_eps,
arch_config.qk_norm,
modulation=True,
quant_config=quant_config,
prefix=f"noise_refiner.{layer_id}",
)
for layer_id in range(arch_config.n_refiner_layers)
]
@@ -457,6 +561,8 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
arch_config.norm_eps,
arch_config.qk_norm,
modulation=False,
quant_config=quant_config,
prefix=f"context_refiner.{layer_id}",
)
for layer_id in range(arch_config.n_refiner_layers)
]
@@ -482,6 +588,8 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
arch_config.n_kv_heads,
arch_config.norm_eps,
arch_config.qk_norm,
quant_config=quant_config,
prefix=f"layers.{layer_id}",
)
for layer_id in range(arch_config.num_layers)
]