minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 15 Packages Projects Releases Wiki Activity

[Diffusion] Add Qwen Image ModelOpt FP8 support (#23155)

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Xiaoyu Zhang
2026-05-04 00:24:22 +08:00
committed by GitHub
co-authored by Mick
parent 5925572c95
commit f2d1390909
7 changed files with 273 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md
+29 -2
View File
@@ -61,9 +61,9 @@ This repo now contains:
- trajectory similarity validation:
[`python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py`](../../../tools/compare_diffusion_trajectory_similarity.py)
Validated documentation and CI coverage currently center on six ModelOpt diffusion transformer override families:
Validated documentation and CI coverage currently center on these ModelOpt diffusion transformer override families:
- FP8: FLUX.1-dev, FLUX.2-dev, Wan2.2
- FP8: FLUX.1-dev, FLUX.2-dev, Wan2.2, Qwen Image, Qwen Image Edit
- NVFP4: FLUX.1-dev, FLUX.2-dev, Wan2.2
Treat a new family, a new precision, or a new checkpoint layout as unsupported until it has a documented matrix row and a matching validation story.
@@ -194,6 +194,33 @@ For `FLUX.1-dev`, the validated fallback set currently keeps these modules in BF
Use `--model-type flux1` to force that profile, or rely on `--model-type auto` when the export config identifies `FluxTransformer2DModel`.
Qwen Image and Qwen Image Edit share `QwenImageTransformer2DModel`, so one
ModelOpt FP8 fallback preset covers both. The validated Qwen Image fallback set
keeps these modules in BF16:
- `img_in`
- `txt_in`
- `time_text_embed.timestep_embedder.linear_1`
- `time_text_embed.timestep_embedder.linear_2`
- `norm_out.linear`
- `proj_out`
- `transformer_blocks.*.img_mlp.net.2`
- `transformer_blocks.*.img_mod`
- `transformer_blocks.*.txt_mod`
Use `--model-type qwen-image` to force that profile, or rely on
`--model-type auto` when the export config identifies
`QwenImageTransformer2DModel`.
Qwen modulation weights can appear in safetensors as `.img_mod.1.weight` and
`.txt_mod.1.weight`. Canonicalize those module names to `.img_mod` and
`.txt_mod` before fallback matching.
For Qwen Image FP8, explicit BF16 fallback tensors must be written before
honoring ModelOpt ignored weights. Otherwise converter stats can report a
fallback while the output checkpoint still retains the source FP8 tensor, which
causes severe image-quality regressions.
For FLUX.1-dev NVFP4 model families that need a mixed BF16+NVFP4 checkpoint, build the merged transformer explicitly:
```bash
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+124 -20
View File
@@ -10,7 +10,6 @@ import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.models.attention import FeedForward
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.normalization import AdaLayerNormContinuous
@@ -532,10 +531,27 @@ class QwenImageCrossAttention(nn.Module):
prefix=f"{prefix}.to_qkv",
)
else:
# Use separate Q/K/V projections for non-quantized models
self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
self.to_q = ReplicatedLinear(
dim,
self.inner_dim,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.to_q",
)
self.to_k = ReplicatedLinear(
dim,
self.inner_dim,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.to_k",
)
self.to_v = ReplicatedLinear(
dim,
self.inner_dim,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.to_v",
)
if self.qk_norm:
self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
@@ -552,15 +568,26 @@ class QwenImageCrossAttention(nn.Module):
prefix=f"{prefix}.to_added_qkv",
)
else:
# Use separate Q/K/V projections for non-quantized models
self.add_q_proj = ReplicatedLinear(
added_kv_proj_dim, self.inner_dim, bias=True
added_kv_proj_dim,
self.inner_dim,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.add_q_proj",
)
self.add_k_proj = ReplicatedLinear(
added_kv_proj_dim, self.inner_dim, bias=True
added_kv_proj_dim,
self.inner_dim,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.add_k_proj",
)
self.add_v_proj = ReplicatedLinear(
added_kv_proj_dim, self.inner_dim, bias=True
added_kv_proj_dim,
self.inner_dim,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.add_v_proj",
)
if context_pre_only is not None and not context_pre_only:
@@ -705,6 +732,65 @@ class QwenImageCrossAttention(nn.Module):
return img_attn_output, txt_attn_output
class QwenImageGELU(nn.Module):
def __init__(
self,
dim: int,
inner_dim: int,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__()
self.proj = ReplicatedLinear(
dim,
inner_dim,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.proj",
)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
hidden_states, _ = self.proj(hidden_states)
return F.gelu(hidden_states, approximate="tanh")
class QwenImageFeedForward(nn.Module):
def __init__(
self,
dim: int,
dim_out: int,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
mult: int = 4,
) -> None:
super().__init__()
inner_dim = dim * mult
self.net = nn.ModuleList(
[
QwenImageGELU(
dim,
inner_dim,
quant_config=quant_config,
prefix=f"{prefix}.net.0",
),
nn.Dropout(0.0),
ReplicatedLinear(
inner_dim,
dim_out,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.net.2",
),
]
)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
hidden_states = self.net[0](hidden_states)
hidden_states = self.net[1](hidden_states)
hidden_states, _ = self.net[2](hidden_states)
return hidden_states
class QwenImageTransformerBlock(nn.Module):
def __init__(
self,
@@ -794,15 +880,17 @@ class QwenImageTransformerBlock(nn.Module):
activation_fn="gelu-approximate",
)
else:
self.img_mlp = FeedForward(
self.img_mlp = QwenImageFeedForward(
dim=dim,
dim_out=dim,
activation_fn="gelu-approximate",
quant_config=quant_config,
prefix=f"{prefix}.img_mlp",
)
self.txt_mlp = FeedForward(
self.txt_mlp = QwenImageFeedForward(
dim=dim,
dim_out=dim,
activation_fn="gelu-approximate",
quant_config=quant_config,
prefix=f"{prefix}.txt_mlp",
)
if nunchaku_enabled:
@@ -1124,8 +1212,20 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
self.img_in = nn.Linear(in_channels, self.inner_dim)
self.txt_in = nn.Linear(joint_attention_dim, self.inner_dim)
self.img_in = ReplicatedLinear(
in_channels,
self.inner_dim,
bias=True,
quant_config=quant_config,
prefix="img_in",
)
self.txt_in = ReplicatedLinear(
joint_attention_dim,
self.inner_dim,
bias=True,
quant_config=quant_config,
prefix="txt_in",
)
self.transformer_blocks = nn.ModuleList(
[
@@ -1144,8 +1244,12 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
self.norm_out = AdaLayerNormContinuous(
self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6
)
self.proj_out = nn.Linear(
self.inner_dim, patch_size * patch_size * self.out_channels, bias=True
self.proj_out = ReplicatedLinear(
self.inner_dim,
patch_size * patch_size * self.out_channels,
bias=True,
quant_config=quant_config,
prefix="proj_out",
)
self.timestep_zero = torch.zeros(
@@ -1230,7 +1334,7 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
if isinstance(encoder_hidden_states_mask, list):
encoder_hidden_states_mask = encoder_hidden_states_mask[0]
hidden_states = self.img_in(hidden_states)
hidden_states, _ = self.img_in(hidden_states)
timestep = (timestep / 1000).to(hidden_states.dtype)
@@ -1242,7 +1346,7 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
modulate_index = None
encoder_hidden_states = self.txt_norm(encoder_hidden_states)
encoder_hidden_states = self.txt_in(encoder_hidden_states)
encoder_hidden_states, _ = self.txt_in(encoder_hidden_states)
block_attention_kwargs = attention_kwargs.copy() if attention_kwargs else {}
if encoder_hidden_states_mask is not None:
@@ -1295,7 +1399,7 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
# Use only the image part (hidden_states) from the dual-stream blocks
hidden_states = self.norm_out(hidden_states, temb_txt)
output = self.proj_out(hidden_states)
output, _ = self.proj_out(hidden_states)
return output
python/sglang/multimodal_gen/test/server/gpu_cases.py
+17
View File
@@ -5,6 +5,8 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
MODELOPT_FLUX2_FP8_TRANSFORMER,
MODELOPT_FLUX2_NVFP4_WEIGHTS,
MODELOPT_NVFP4_B200_ENV_VARS,
MODELOPT_QWEN_IMAGE_EDIT_FP8_TRANSFORMER,
MODELOPT_QWEN_IMAGE_FP8_TRANSFORMER,
MODELOPT_WAN22_FP8_TRANSFORMER,
MODELOPT_WAN22_NVFP4_TRANSFORMER,
T2V_PROMPT,
@@ -14,6 +16,7 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
HUNYUAN3D_SHAPE_sampling_params,
MODELOPT_T2I_CI_sampling_params,
MODELOPT_T2V_CI_sampling_params,
MODELOPT_TI2I_CI_sampling_params,
MULTI_FRAME_I2I_sampling_params,
MULTI_IMAGE_TI2I_sampling_params,
MULTI_IMAGE_TI2I_UPLOAD_sampling_params,
@@ -400,6 +403,20 @@ else:
sampling_params=MODELOPT_T2V_CI_sampling_params,
extras=["--transformer-path", MODELOPT_WAN22_FP8_TRANSFORMER],
),
_make_modelopt_ci_case(
"qwen_image_modelopt_fp8_t2i",
model_path=DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST,
modality="image",
sampling_params=MODELOPT_T2I_CI_sampling_params,
extras=["--transformer-path", MODELOPT_QWEN_IMAGE_FP8_TRANSFORMER],
),
_make_modelopt_ci_case(
"qwen_image_edit_modelopt_fp8_ti2i",
model_path=DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST,
modality="image",
sampling_params=MODELOPT_TI2I_CI_sampling_params,
extras=["--transformer-path", MODELOPT_QWEN_IMAGE_EDIT_FP8_TRANSFORMER],
),
_make_modelopt_ci_case(
"flux1_modelopt_nvfp4_t2i",
model_path=DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST,
python/sglang/multimodal_gen/test/server/testcase_configs.py
+11
View File
@@ -368,6 +368,13 @@ MODELOPT_T2I_CI_sampling_params = DiffusionSamplingParams(
extras={"num_inference_steps": 12, "seed": 0},
)
MODELOPT_TI2I_CI_sampling_params = DiffusionSamplingParams(
prompt="Convert 2D style to 3D style",
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
output_size="512x512",
extras={"num_inference_steps": 8, "seed": 0},
)
TI2I_sampling_params = DiffusionSamplingParams(
prompt="Convert 2D style to 3D style",
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
@@ -435,6 +442,10 @@ HUNYUAN3D_SHAPE_sampling_params = DiffusionSamplingParams(
MODELOPT_FLUX1_FP8_TRANSFORMER = "lmsys/flux1-dev-modelopt-fp8-sglang-transformer"
MODELOPT_FLUX2_FP8_TRANSFORMER = "lmsys/flux2-dev-modelopt-fp8-sglang-transformer"
MODELOPT_WAN22_FP8_TRANSFORMER = "lmsys/wan22-t2v-a14b-modelopt-fp8-sglang-transformer"
MODELOPT_QWEN_IMAGE_FP8_TRANSFORMER = "lmsys/qwen-image-modelopt-fp8-sglang-transformer"
MODELOPT_QWEN_IMAGE_EDIT_FP8_TRANSFORMER = (
"lmsys/qwen-image-edit-modelopt-fp8-sglang-transformer"
)
MODELOPT_FLUX1_NVFP4_TRANSFORMER = "lmsys/flux1-dev-modelopt-nvfp4-sglang-transformer"
MODELOPT_FLUX2_NVFP4_WEIGHTS = "black-forest-labs/FLUX.2-dev-NVFP4"
MODELOPT_WAN22_NVFP4_TRANSFORMER = (
python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py
+21 -5
View File
@@ -76,6 +76,15 @@ DEFAULT_LTX2_KEEP_BF16_PATTERNS = [
r"^transformer_blocks\.(0|43|44|45|46|47)\.(attn1|attn2|audio_attn1|audio_attn2|audio_to_video_attn|video_to_audio_attn)\.to_out\.0$",
r"^transformer_blocks\.(0|43|44|45|46|47)\.(ff|audio_ff)\.proj_(in|out)$",
]
DEFAULT_QWEN_IMAGE_KEEP_BF16_PATTERNS = [
r"^img_in$",
r"^txt_in$",
r"^time_text_embed\.timestep_embedder\.linear_[12]$",
r"^norm_out\.linear$",
r"^proj_out$",
r"^transformer_blocks\.\d+\.img_mlp\.net\.2$",
r"^transformer_blocks\.\d+\.(img_mod|txt_mod)$",
]
def _resolve_transformer_dir(path: str) -> str:
@@ -173,6 +182,7 @@ def _module_name_variants(weight_name: str) -> list[str]:
canonicalized.append(
re.sub(r"(\.audio_ff|\.ff)\.net\.2$", r"\1.proj_out", variant)
)
canonicalized.append(re.sub(r"(\.(img_mod|txt_mod))\.1$", r"\1", variant))
variants.extend(canonicalized)
deduped: list[str] = []
@@ -259,12 +269,16 @@ def get_default_keep_bf16_patterns(
return list(DEFAULT_FLUX1_KEEP_BF16_PATTERNS)
if model_type == "flux2":
return list(DEFAULT_FLUX2_KEEP_BF16_PATTERNS)
if model_type == "qwen-image":
return list(DEFAULT_QWEN_IMAGE_KEEP_BF16_PATTERNS)
if model_type == "none":
return []
if class_name == "FluxTransformer2DModel":
return list(DEFAULT_FLUX1_KEEP_BF16_PATTERNS)
if class_name == "Flux2Transformer2DModel":
return list(DEFAULT_FLUX2_KEEP_BF16_PATTERNS)
if class_name == "QwenImageTransformer2DModel":
return list(DEFAULT_QWEN_IMAGE_KEEP_BF16_PATTERNS)
return []
@@ -552,13 +566,14 @@ def build_modelopt_fp8_transformer(
if name in fallback_scale_names:
del shard_tensors[name]
continue
if name in fallback_tensors:
shard_tensors[name] = fallback_tensors[name]
continue
if name.endswith(".weight") and is_ignored_by_modelopt(
name, ignore_patterns
):
preserved_ignored_weight_count += 1
continue
if name in fallback_tensors:
shard_tensors[name] = fallback_tensors[name]
scale_key = _resolve_scale_key(name, fp8_scale_map)
if (
name.endswith(".weight")
@@ -645,12 +660,13 @@ def _parse_args() -> argparse.Namespace:
)
parser.add_argument(
"--model-type",
choices=["auto", "flux1", "flux2", "ltx2", "none"],
choices=["auto", "flux1", "flux2", "ltx2", "qwen-image", "none"],
default="auto",
help=(
"Optional model-family BF16 fallback profile. 'none' uses the generic "
"conversion path. 'auto' enables the validated FLUX.1 / FLUX.2 / LTX-2 "
"fallback set when the export config matches those transformer classes."
"conversion path. 'auto' enables the validated FLUX.1 / FLUX.2 / LTX-2 / "
"Qwen Image fallback set when the export config matches those transformer "
"classes."
),
)
parser.add_argument(