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

[diffusion] quant: update Wan2.2 modelOpt CI checkpoints (#25483)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-05-20 09:05:39 +08:00
committed by GitHub
parent 7f154ba449
commit 80fc524809
9 changed files with 87 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/registry.py
+4 -1
View File
@@ -739,7 +739,10 @@ def _register_configs():
register_configs(
sampling_param_cls=Wan2_2_T2V_A14B_SamplingParam,
pipeline_config_cls=Wan2_2_T2V_A14B_Config,
hf_model_paths=["Wan-AI/Wan2.2-T2V-A14B-Diffusers"],
hf_model_paths=[
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
"nvidia/Wan2.2-T2V-A14B-Diffusers-NVFP4",
],
)
register_configs(
sampling_param_cls=Wan2_2_I2V_A14B_SamplingParam,
python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py
+20 -7
View File
@@ -202,7 +202,7 @@ class ModelOptFp4Config(ModelOptQuantConfig):
exclude_modules: List[str] = None,
packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
checkpoint_uses_packed_qkv: bool = False,
swap_weight_nibbles: bool = True,
swap_weight_nibbles: bool = False,
) -> None:
super().__init__(exclude_modules, packed_modules_mapping)
self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
@@ -261,7 +261,7 @@ class ModelOptFp4Config(ModelOptQuantConfig):
def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config:
group_size = None
exclude_modules = []
swap_weight_nibbles = True
swap_weight_nibbles = False
# Flat format (config.json quantization_config)
quant_method = config.get("quant_algo")
@@ -273,7 +273,10 @@ class ModelOptFp4Config(ModelOptQuantConfig):
first_group = next(iter(config_groups.values()), {})
group_size = first_group.get("weights", {}).get("group_size")
exclude_modules = config.get("ignore", [])
swap_weight_nibbles = config.get("swap_weight_nibbles", True)
swap_weight_nibbles = config.get(
"swap_weight_nibbles",
config.get("checkpoint_uses_packed_qkv", False),
)
else:
# Nested format (hf_quant_config.json)
try:
@@ -283,7 +286,10 @@ class ModelOptFp4Config(ModelOptQuantConfig):
exclude_modules = quant_config.get("exclude_modules", [])
swap_weight_nibbles = quant_config.get(
"swap_weight_nibbles",
config.get("swap_weight_nibbles", True),
config.get(
"swap_weight_nibbles",
config.get("checkpoint_uses_packed_qkv", False),
),
)
except (ValueError, KeyError):
raise ValueError("Cannot find 'quant_algo' in quantization config.")
@@ -494,7 +500,9 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
w = layer.weight.data
w_swapped = _prepare_nvfp4_weight_bytes(
w,
swap_weight_nibbles=getattr(self.quant_config, "swap_weight_nibbles", True),
swap_weight_nibbles=getattr(
self.quant_config, "swap_weight_nibbles", False
),
)
_, flashinfer_backend = _get_fp4_gemm_op()
@@ -554,8 +562,13 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
padded_scales[:B, :M, :K] = scales
_, flashinfer_backend = _get_fp4_gemm_op()
if flashinfer_backend is None:
# CUTLASS (sgl_kernel) path: blockwise interleave to TMA layout
uses_flux1_scale_layout = not getattr(
self.quant_config, "checkpoint_uses_packed_qkv", False
) and getattr(layer, "prefix", "").startswith(
("transformer_blocks.", "single_transformer_blocks.")
)
if flashinfer_backend is None or uses_flux1_scale_layout:
# CUTLASS and FLUX.1 CUDNN paths need the TMA scale layout.
padded_scales = padded_scales.reshape(
B, M_padded // 128, 4, 32, K_padded // 4, 4
)
python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py
+3 -3
View File
@@ -88,12 +88,12 @@ def _merge_modelopt_fp4_configs(
inferred_config.packed_modules_mapping = getattr(
existing_config, "packed_modules_mapping", {}
)
inferred_config.swap_weight_nibbles = getattr(
existing_config, "swap_weight_nibbles", True
)
inferred_config.checkpoint_uses_packed_qkv = getattr(
inferred_config, "checkpoint_uses_packed_qkv", False
) or getattr(existing_config, "checkpoint_uses_packed_qkv", False)
inferred_config.swap_weight_nibbles = getattr(
inferred_config, "swap_weight_nibbles", False
) or getattr(existing_config, "swap_weight_nibbles", False)
if getattr(inferred_config, "group_size", None) is None:
inferred_config.group_size = getattr(existing_config, "group_size", None)
python/sglang/multimodal_gen/test/server/consistency_threshold.json
+6
View File
@@ -13,6 +13,12 @@
"psnr_threshold": 24.0,
"mean_abs_diff_threshold": 8.0
},
"flux1_modelopt_fp8_t2i": {
"clip_threshold": 0.92,
"ssim_threshold": 0.94,
"psnr_threshold": 28.0,
"mean_abs_diff_threshold": 8.0
},
"flux_2_klein_image_t2i": {
"clip_threshold": 0.94,
"ssim_threshold": 0.78,
python/sglang/multimodal_gen/test/server/gpu_cases.py
+14 -7
View File
@@ -11,8 +11,9 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
MODELOPT_NVFP4_B200_ENV_VARS,
MODELOPT_QWEN_IMAGE_EDIT_FP8_TRANSFORMER,
MODELOPT_QWEN_IMAGE_FP8_TRANSFORMER,
MODELOPT_WAN22_FP8_TRANSFORMER,
MODELOPT_WAN22_NVFP4_TRANSFORMER,
MODELOPT_WAN22_FP8_MODEL,
MODELOPT_WAN22_NVFP4_B200_ENV_VARS,
MODELOPT_WAN22_NVFP4_MODEL,
T2V_PROMPT,
DiffusionSamplingParams,
DiffusionServerArgs,
@@ -419,6 +420,7 @@ else:
modality="image",
sampling_params=MODELOPT_T2I_CI_sampling_params,
extras=["--transformer-path", MODELOPT_FLUX1_FP8_TRANSFORMER],
run_consistency_check=True,
),
_make_modelopt_ci_case(
"flux2_modelopt_fp8_t2i",
@@ -426,13 +428,15 @@ else:
modality="image",
sampling_params=MODELOPT_T2I_CI_sampling_params,
extras=["--transformer-path", MODELOPT_FLUX2_FP8_TRANSFORMER],
run_consistency_check=True,
),
_make_modelopt_ci_case(
"wan22_modelopt_fp8_t2v",
model_path=DEFAULT_WAN_2_2_T2V_A14B_MODEL_NAME_FOR_TEST,
model_path=MODELOPT_WAN22_FP8_MODEL,
modality="video",
sampling_params=MODELOPT_T2V_CI_sampling_params,
extras=["--transformer-path", MODELOPT_WAN22_FP8_TRANSFORMER],
extras=[],
run_consistency_check=True,
),
_make_modelopt_ci_case(
"hunyuanvideo_modelopt_fp8_t2v",
@@ -469,6 +473,7 @@ else:
sampling_params=MODELOPT_T2I_CI_sampling_params,
extras=["--transformer-path", MODELOPT_FLUX1_NVFP4_TRANSFORMER],
env_vars=MODELOPT_NVFP4_B200_ENV_VARS,
run_consistency_check=True,
),
_make_modelopt_ci_case(
"flux2_modelopt_nvfp4_t2i",
@@ -477,14 +482,16 @@ else:
sampling_params=MODELOPT_T2I_CI_sampling_params,
extras=["--transformer-weights-path", MODELOPT_FLUX2_NVFP4_WEIGHTS],
env_vars=MODELOPT_NVFP4_B200_ENV_VARS,
run_consistency_check=True,
),
_make_modelopt_ci_case(
"wan22_modelopt_nvfp4_t2v",
model_path=DEFAULT_WAN_2_2_T2V_A14B_MODEL_NAME_FOR_TEST,
model_path=MODELOPT_WAN22_NVFP4_MODEL,
modality="video",
sampling_params=MODELOPT_T2V_CI_sampling_params,
extras=["--transformer-path", MODELOPT_WAN22_NVFP4_TRANSFORMER],
env_vars=MODELOPT_NVFP4_B200_ENV_VARS,
extras=[],
env_vars=MODELOPT_WAN22_NVFP4_B200_ENV_VARS,
run_consistency_check=True,
),
]
python/sglang/multimodal_gen/test/server/testcase_configs.py
+7 -5
View File
@@ -441,7 +441,7 @@ 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_WAN22_FP8_MODEL = "nvidia/Wan2.2-T2V-A14B-Diffusers-FP8"
MODELOPT_HUNYUANVIDEO_FP8_TRANSFORMER = (
"lmsys/hunyuanvideo-modelopt-fp8-sglang-transformer"
)
@@ -451,10 +451,11 @@ MODELOPT_QWEN_IMAGE_EDIT_FP8_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 = (
"lmsys/wan22-t2v-a14b-modelopt-nvfp4-sglang-transformer"
)
MODELOPT_WAN22_NVFP4_MODEL = "nvidia/Wan2.2-T2V-A14B-Diffusers-NVFP4"
MODELOPT_NVFP4_B200_ENV_VARS = {"SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND": "cudnn"}
MODELOPT_WAN22_NVFP4_B200_ENV_VARS = {
"SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND": "trtllm"
}
def _make_modelopt_ci_case(
@@ -465,6 +466,7 @@ def _make_modelopt_ci_case(
sampling_params: DiffusionSamplingParams,
extras: list[str],
env_vars: dict[str, str] | None = None,
run_consistency_check: bool = False,
) -> DiffusionTestCase:
return DiffusionTestCase(
case_id,
@@ -477,7 +479,7 @@ def _make_modelopt_ci_case(
),
sampling_params,
run_perf_check=False,
run_consistency_check=False,
run_consistency_check=run_consistency_check,
run_component_accuracy_check=False,
)
python/sglang/multimodal_gen/test/test_utils.py
+1 -1
View File
@@ -33,7 +33,7 @@ if TYPE_CHECKING:
logger = init_logger(__name__)
SGL_TEST_FILES_CI_DATA_REVISION = "c8305f1dd8cc82197f36c17d0f503adc94016cc7"
SGL_TEST_FILES_CI_DATA_REVISION = "94eab4fcca6d4ddc77cdb3622f13033b61e81002"
SGL_TEST_FILES_CONSISTENCY_GT_ROOT = (
"https://raw.githubusercontent.com/"
f"sgl-project/ci-data/{SGL_TEST_FILES_CI_DATA_REVISION}/"
python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py
+2 -4
View File
@@ -218,9 +218,7 @@ def build_modelopt_nvfp4_transformer(
patterns.extend(keep_bf16_patterns)
resolved_swap_weight_nibbles = (
swap_weight_nibbles
if swap_weight_nibbles is not None
else (False if pattern_preset == "flux1-nvfp4" else True)
swap_weight_nibbles if swap_weight_nibbles is not None else False
)
output_config = _updated_quant_config(
_load_config(source_dir),
@@ -373,7 +371,7 @@ def _parse_args() -> argparse.Namespace:
default=None,
help=(
"Whether the runtime should swap packed FP4 nibbles before padding. "
"Defaults to false for --pattern-preset flux1-nvfp4 and true otherwise."
"Defaults to false."
),
)
parser.add_argument(