chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py

@@ -0,0 +1,883 @@
+"""Build an SGLang-loadable ModelOpt FP8 diffusion transformer.
+
+The core conversion path is model-agnostic:
+- read the ModelOpt diffusers transformer export
+- rebuild per-layer `weight_scale` / `input_scale` tensors from `backbone.pt`
+- materialize SGLang-native `float8_e4m3fn` weights
+- preserve ModelOpt `ignore` layers in their original dtype
+
+Some models still benefit from a small validated BF16 fallback set. Those
+fallback profiles are intentionally isolated so the generic FP8 conversion path
+remains reusable across future diffusion backbones.
+
+Example:
+
+    python -m sglang.multimodal_gen.tools.build_modelopt_fp8_transformer \
+        --modelopt-hf-dir /tmp/modelopt_flux2_fp8/hf \
+        --modelopt-backbone-ckpt /tmp/modelopt_flux2_fp8/ckpt/backbone.pt \
+        --base-transformer-dir /path/to/FLUX.2-dev/transformer \
+        --output-dir /tmp/modelopt_flux2_fp8/sglang_transformer
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import os
+import re
+import shutil
+from collections import defaultdict
+from pathlib import Path
+from typing import Callable, Iterable, Mapping, Sequence
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import load_file, save_file
+
+from sglang.multimodal_gen.runtime.utils.quantization_utils import (
+    normalize_flat_modelopt_quant_config,
+)
+
+INDEX_FILENAMES = [
+    "model.safetensors.index.json",
+    "diffusion_pytorch_model.safetensors.index.json",
+]
+FP8_E4M3_MAXBOUND = 448.0
+DEFAULT_FLUX2_KEEP_BF16_PATTERNS = [
+    r"^time_guidance_embed\.(timestep_embedder|guidance_embedder)\.linear_[12]$",
+    r"^double_stream_modulation_(img|txt)\.linear$",
+    r"^single_stream_modulation\.linear$",
+    r"^x_embedder$",
+    r"^context_embedder$",
+    r"^norm_out\.linear$",
+]
+DEFAULT_FLUX1_KEEP_BF16_PATTERNS = [
+    r"^transformer_blocks\.\d+\.norm1\.linear$",
+    r"^transformer_blocks\.\d+\.norm1_context\.linear$",
+    r"^transformer_blocks\.\d+\.ff\.net\.0\.proj$",
+    r"^transformer_blocks\.\d+\.ff\.net\.2$",
+    r"^transformer_blocks\.\d+\.ff_context\.net\.0\.proj$",
+    r"^transformer_blocks\.\d+\.ff_context\.net\.2$",
+    r"^single_transformer_blocks\.\d+\.norm\.linear$",
+]
+DEFAULT_LTX2_KEEP_BF16_PATTERNS = [
+    r"^(audio_)?adaln_single\.emb\.timestep_embedder\.linear_[12]$",
+    r"^(audio_)?adaln_single\.linear$",
+    r"^audio_caption_projection\.linear_[12]$",
+    r"^audio_patchify_proj$",
+    r"^audio_proj_out$",
+    r"^av_ca_(a2v_gate|audio_scale_shift|v2a_gate|video_scale_shift)_adaln_single\.emb\.timestep_embedder\.linear_[12]$",
+    r"^av_ca_(a2v_gate|audio_scale_shift|v2a_gate|video_scale_shift)_adaln_single\.linear$",
+    r"^caption_projection\.linear_[12]$",
+    r"^patchify_proj$",
+    r"^proj_out$",
+    r"^transformer_blocks\.(0|43|44|45|46|47)\.(attn1|attn2|audio_attn1|audio_attn2|audio_to_video_attn|video_to_audio_attn)\.to_(q|k|v)$",
+    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_HUNYUANVIDEO_KEEP_BF16_PATTERNS = [
+    r"^context_embedder\.",
+    r"^x_embedder\.proj$",
+    r"^time_text_embed\.(timestep_embedder|guidance_embedder|text_embedder)\.linear_[12]$",
+    r"^norm_out\.linear$",
+    r"^proj_out$",
+    r"^transformer_blocks\.\d+\.norm1\.linear$",
+    r"^transformer_blocks\.\d+\.norm1_context\.linear$",
+    r"^single_transformer_blocks\.\d+\.norm\.linear$",
+]
+HUNYUANVIDEO_RUNTIME_NAME_REPLACEMENTS = [
+    (
+        r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1$",
+        r"txt_in.t_embedder.mlp.fc_in",
+    ),
+    (
+        r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_2$",
+        r"txt_in.t_embedder.mlp.fc_out",
+    ),
+    (r"^context_embedder\.proj_in$", r"txt_in.input_embedder"),
+    (
+        r"^context_embedder\.time_text_embed\.text_embedder\.linear_1$",
+        r"txt_in.c_embedder.fc_in",
+    ),
+    (
+        r"^context_embedder\.time_text_embed\.text_embedder\.linear_2$",
+        r"txt_in.c_embedder.fc_out",
+    ),
+    (
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm1$",
+        r"txt_in.refiner_blocks.\1.norm1",
+    ),
+    (
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm2$",
+        r"txt_in.refiner_blocks.\1.norm2",
+    ),
+    (
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_[qkv]$",
+        r"txt_in.refiner_blocks.\1.self_attn_qkv",
+    ),
+    (
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_out\.0$",
+        r"txt_in.refiner_blocks.\1.self_attn_proj",
+    ),
+    (
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?$",
+        r"txt_in.refiner_blocks.\1.mlp.fc_in",
+    ),
+    (
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?$",
+        r"txt_in.refiner_blocks.\1.mlp.fc_out",
+    ),
+    (
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear$",
+        r"txt_in.refiner_blocks.\1.adaLN_modulation.linear",
+    ),
+    (r"^x_embedder\.proj$", r"img_in.proj"),
+    (r"^time_text_embed\.timestep_embedder\.linear_1$", r"time_in.mlp.fc_in"),
+    (r"^time_text_embed\.timestep_embedder\.linear_2$", r"time_in.mlp.fc_out"),
+    (r"^time_text_embed\.guidance_embedder\.linear_1$", r"guidance_in.mlp.fc_in"),
+    (r"^time_text_embed\.guidance_embedder\.linear_2$", r"guidance_in.mlp.fc_out"),
+    (r"^time_text_embed\.text_embedder\.linear_1$", r"vector_in.fc_in"),
+    (r"^time_text_embed\.text_embedder\.linear_2$", r"vector_in.fc_out"),
+    (r"^transformer_blocks\.(\d+)\.norm1\.linear$", r"double_blocks.\1.img_mod.linear"),
+    (
+        r"^transformer_blocks\.(\d+)\.norm1_context\.linear$",
+        r"double_blocks.\1.txt_mod.linear",
+    ),
+    (r"^transformer_blocks\.(\d+)\.attn\.norm_q$", r"double_blocks.\1.img_attn_q_norm"),
+    (r"^transformer_blocks\.(\d+)\.attn\.norm_k$", r"double_blocks.\1.img_attn_k_norm"),
+    (r"^transformer_blocks\.(\d+)\.attn\.to_[qkv]$", r"double_blocks.\1.img_attn_qkv"),
+    (
+        r"^transformer_blocks\.(\d+)\.attn\.add_[qkv]_proj$",
+        r"double_blocks.\1.txt_attn_qkv",
+    ),
+    (
+        r"^transformer_blocks\.(\d+)\.attn\.to_out\.0$",
+        r"double_blocks.\1.img_attn_proj",
+    ),
+    (
+        r"^transformer_blocks\.(\d+)\.attn\.to_add_out$",
+        r"double_blocks.\1.txt_attn_proj",
+    ),
+    (
+        r"^transformer_blocks\.(\d+)\.attn\.norm_added_q$",
+        r"double_blocks.\1.txt_attn_q_norm",
+    ),
+    (
+        r"^transformer_blocks\.(\d+)\.attn\.norm_added_k$",
+        r"double_blocks.\1.txt_attn_k_norm",
+    ),
+    (
+        r"^transformer_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?$",
+        r"double_blocks.\1.img_mlp.fc_in",
+    ),
+    (
+        r"^transformer_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?$",
+        r"double_blocks.\1.img_mlp.fc_out",
+    ),
+    (
+        r"^transformer_blocks\.(\d+)\.ff_context\.net\.0(?:\.proj)?$",
+        r"double_blocks.\1.txt_mlp.fc_in",
+    ),
+    (
+        r"^transformer_blocks\.(\d+)\.ff_context\.net\.2(?:\.proj)?$",
+        r"double_blocks.\1.txt_mlp.fc_out",
+    ),
+    (r"^single_transformer_blocks\.(\d+)\.attn\.norm_q$", r"single_blocks.\1.q_norm"),
+    (r"^single_transformer_blocks\.(\d+)\.attn\.norm_k$", r"single_blocks.\1.k_norm"),
+    (
+        r"^single_transformer_blocks\.(\d+)\.attn\.to_[qkv]$",
+        r"single_blocks.\1.linear1",
+    ),
+    (r"^single_transformer_blocks\.(\d+)\.proj_mlp$", r"single_blocks.\1.linear1"),
+    (r"^single_transformer_blocks\.(\d+)\.proj_out$", r"single_blocks.\1.linear2"),
+    (
+        r"^single_transformer_blocks\.(\d+)\.norm\.linear$",
+        r"single_blocks.\1.modulation.linear",
+    ),
+    (r"^norm_out\.linear$", r"final_layer.adaLN_modulation.linear"),
+    (r"^proj_out$", r"final_layer.linear"),
+]
+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:
+    candidate = Path(path).expanduser().resolve()
+    if (candidate / "config.json").is_file():
+        return str(candidate)
+    transformer_dir = candidate / "transformer"
+    if (transformer_dir / "config.json").is_file():
+        return str(transformer_dir)
+    raise FileNotFoundError(f"Could not resolve a transformer directory from: {path}")
+
+
+def _resolve_backbone_ckpt(path: str) -> str:
+    candidate = Path(path).expanduser().resolve()
+    if candidate.is_file():
+        return str(candidate)
+    backbone_path = candidate / "backbone.pt"
+    if backbone_path.is_file():
+        return str(backbone_path)
+    raise FileNotFoundError(f"Could not resolve backbone.pt from: {path}")
+
+
+def _find_index_file(model_dir: str) -> str | None:
+    for filename in INDEX_FILENAMES:
+        candidate = os.path.join(model_dir, filename)
+        if os.path.isfile(candidate):
+            return filename
+
+    matches = sorted(
+        filename
+        for filename in os.listdir(model_dir)
+        if filename.endswith(".safetensors.index.json")
+    )
+    return matches[0] if matches else None
+
+
+def _load_weight_map(model_dir: str) -> tuple[dict[str, str], str | None]:
+    index_filename = _find_index_file(model_dir)
+    if index_filename is not None:
+        with open(os.path.join(model_dir, index_filename), encoding="utf-8") as f:
+            index_data = json.load(f)
+        return dict(index_data["weight_map"]), index_filename
+
+    safetensors_files = sorted(
+        filename
+        for filename in os.listdir(model_dir)
+        if filename.endswith(".safetensors")
+    )
+    if len(safetensors_files) != 1:
+        raise ValueError(
+            f"Expected an index file or a single safetensors shard in {model_dir}, "
+            f"found {len(safetensors_files)} shard(s)."
+        )
+
+    shard_name = safetensors_files[0]
+    with safe_open(
+        os.path.join(model_dir, shard_name), framework="pt", device="cpu"
+    ) as f:
+        weight_map = {key: shard_name for key in f.keys()}
+    index_filename = f"{Path(shard_name).stem}.safetensors.index.json"
+    return weight_map, index_filename
+
+
+def _load_config(model_dir: str) -> dict:
+    config_path = os.path.join(model_dir, "config.json")
+    with open(config_path, encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _load_first_shard_metadata(
+    model_dir: str, weight_map: Mapping[str, str]
+) -> dict[str, str]:
+    if not weight_map:
+        return {}
+    first_shard = next(iter(weight_map.values()))
+    with safe_open(
+        os.path.join(model_dir, first_shard), framework="pt", device="cpu"
+    ) as f:
+        return dict(f.metadata() or {})
+
+
+def _map_hunyuanvideo_runtime_module_name(module_name: str) -> list[str]:
+    mapped_names: list[str] = []
+    for pattern, replacement in HUNYUANVIDEO_RUNTIME_NAME_REPLACEMENTS:
+        mapped = re.sub(pattern, replacement, module_name)
+        if mapped != module_name:
+            mapped_names.append(mapped)
+    return mapped_names
+
+
+def _get_runtime_module_name_mapper(
+    *, model_type: str, class_name: str | None
+) -> Callable[[str], list[str]] | None:
+    if model_type == "hunyuan-video" or class_name == "HunyuanVideoTransformer3DModel":
+        return _map_hunyuanvideo_runtime_module_name
+    return None
+
+
+def _module_name_variants(
+    weight_name: str,
+    runtime_name_mapper: Callable[[str], list[str]] | None = None,
+) -> list[str]:
+    module_name = weight_name[:-7] if weight_name.endswith(".weight") else weight_name
+    variants = [module_name]
+
+    for prefix in ("model.diffusion_model.", "velocity_model."):
+        if module_name.startswith(prefix):
+            variants.append(module_name[len(prefix) :])
+
+    canonicalized: list[str] = []
+    for variant in variants:
+        canonicalized.append(
+            re.sub(r"(\.audio_ff|\.ff)\.net\.0\.proj$", r"\1.proj_in", variant)
+        )
+        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)
+    if runtime_name_mapper is not None:
+        runtime_variants: list[str] = []
+        for variant in variants:
+            runtime_variants.extend(runtime_name_mapper(variant))
+        variants.extend(runtime_variants)
+
+    deduped: list[str] = []
+    for variant in variants:
+        if variant not in deduped:
+            deduped.append(variant)
+    return deduped
+
+
+def _preferred_module_name(
+    weight_name: str,
+    runtime_name_mapper: Callable[[str], list[str]] | None = None,
+) -> str:
+    return _module_name_variants(weight_name, runtime_name_mapper)[-1]
+
+
+def _scale_key_candidates(weight_name: str) -> list[str]:
+    candidates = [weight_name]
+    if weight_name.startswith("model.diffusion_model."):
+        candidates.append(
+            "velocity_model." + weight_name[len("model.diffusion_model.") :]
+        )
+    return candidates
+
+
+def _resolve_scale_key(
+    weight_name: str,
+    scale_map: Mapping[str, Mapping[str, torch.Tensor]],
+) -> str | None:
+    for candidate in _scale_key_candidates(weight_name):
+        if candidate in scale_map:
+            return candidate
+    return None
+
+
+def _is_ltx2_x0_export(
+    *,
+    config: Mapping[str, object],
+    source_metadata: Mapping[str, str],
+    source_weight_map: Mapping[str, str],
+) -> bool:
+    if config.get("_class_name") != "X0Model":
+        return False
+    if not any(name.startswith("model.diffusion_model.") for name in source_weight_map):
+        return False
+    try:
+        metadata_config = json.loads(str(source_metadata.get("config", "")))
+    except json.JSONDecodeError:
+        return False
+    return isinstance(metadata_config.get("transformer"), dict)
+
+
+def _build_output_config(
+    *,
+    source_config: Mapping[str, object],
+    source_metadata: Mapping[str, str],
+    quant_config: Mapping[str, object],
+    is_ltx2_x0_export: bool,
+) -> dict[str, object]:
+    if is_ltx2_x0_export:
+        metadata_config = json.loads(str(source_metadata["config"]))
+        output_config = dict(metadata_config["transformer"])
+        output_config["_class_name"] = "LTX2VideoTransformer3DModel"
+    else:
+        output_config = dict(source_config)
+
+    output_config["quantization_config"] = dict(quant_config)
+    return output_config
+
+
+def _should_keep_ltx2_transformer_key(weight_name: str) -> bool:
+    if not weight_name.startswith("model.diffusion_model."):
+        return False
+    connector_prefixes = (
+        "model.diffusion_model.audio_embeddings_connector.",
+        "model.diffusion_model.video_embeddings_connector.",
+    )
+    return not weight_name.startswith(connector_prefixes)
+
+
+def get_default_keep_bf16_patterns(
+    *, model_type: str, class_name: str | None
+) -> list[str]:
+    if model_type == "ltx2":
+        return list(DEFAULT_LTX2_KEEP_BF16_PATTERNS)
+    if model_type == "flux1":
+        return list(DEFAULT_FLUX1_KEEP_BF16_PATTERNS)
+    if model_type == "flux2":
+        return list(DEFAULT_FLUX2_KEEP_BF16_PATTERNS)
+    if model_type == "hunyuan-video":
+        return list(DEFAULT_HUNYUANVIDEO_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 == "HunyuanVideoTransformer3DModel":
+        return list(DEFAULT_HUNYUANVIDEO_KEEP_BF16_PATTERNS)
+    if class_name == "QwenImageTransformer2DModel":
+        return list(DEFAULT_QWEN_IMAGE_KEEP_BF16_PATTERNS)
+    return []
+
+
+def should_keep_bf16(
+    weight_name: str,
+    keep_bf16_patterns: Sequence[str],
+    runtime_name_mapper: Callable[[str], list[str]] | None = None,
+) -> bool:
+    if not keep_bf16_patterns:
+        return False
+
+    return any(
+        re.search(pattern, module_name)
+        for pattern in keep_bf16_patterns
+        for module_name in _module_name_variants(weight_name, runtime_name_mapper)
+    )
+
+
+def is_ignored_by_modelopt(
+    weight_name: str,
+    ignore_patterns: Sequence[str],
+    runtime_name_mapper: Callable[[str], list[str]] | None = None,
+) -> bool:
+    if not ignore_patterns:
+        return False
+
+    for pattern in ignore_patterns:
+        regex_str = pattern.replace(".", r"\.").replace("*", r".*")
+        if any(
+            re.fullmatch(regex_str, module_name)
+            for module_name in _module_name_variants(weight_name, runtime_name_mapper)
+        ):
+            return True
+    return False
+
+
+def build_fp8_scale_map(
+    model_state_dict: Mapping[str, torch.Tensor],
+    *,
+    maxbound: float = FP8_E4M3_MAXBOUND,
+) -> dict[str, dict[str, torch.Tensor]]:
+    scale_map: dict[str, dict[str, torch.Tensor]] = {}
+    for key, value in model_state_dict.items():
+        if key.endswith(".weight_quantizer._amax"):
+            layer_name = key[: -len(".weight_quantizer._amax")]
+            scale_map.setdefault(f"{layer_name}.weight", {})["weight_scale"] = (
+                value.detach().to(torch.float32).reshape(1).cpu() / maxbound
+            )
+        elif key.endswith(".input_quantizer._amax"):
+            layer_name = key[: -len(".input_quantizer._amax")]
+            scale_map.setdefault(f"{layer_name}.weight", {})["input_scale"] = (
+                value.detach().to(torch.float32).reshape(1).cpu() / maxbound
+            )
+
+    return {
+        weight_name: scale_tensors
+        for weight_name, scale_tensors in scale_map.items()
+        if {"weight_scale", "input_scale"} <= set(scale_tensors)
+    }
+
+
+def quantize_fp8_weight(
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+) -> torch.Tensor:
+    if weight.dtype == torch.float8_e4m3fn:
+        return weight.contiguous()
+
+    scale = weight_scale.to(weight.device, dtype=torch.float32)
+    if scale.numel() != 1:
+        raise ValueError(
+            "Only per-tensor FP8 scales are supported for diffusion checkpoints, "
+            f"got shape {tuple(scale.shape)}."
+        )
+
+    quantized = (weight.to(torch.float32) / scale.reshape(1)).to(torch.float8_e4m3fn)
+    return quantized.cpu().contiguous()
+
+
+def _copy_non_shard_files(source_dir: str, output_dir: str) -> None:
+    ignored = set(INDEX_FILENAMES)
+    for entry in os.listdir(source_dir):
+        if entry.endswith(".safetensors") or entry in ignored:
+            continue
+        source_path = os.path.join(source_dir, entry)
+        output_path = os.path.join(output_dir, entry)
+        if os.path.isdir(source_path):
+            shutil.copytree(source_path, output_path, dirs_exist_ok=True)
+        else:
+            shutil.copy2(source_path, output_path)
+
+
+def _load_selected_tensors(
+    model_dir: str,
+    weight_map: Mapping[str, str],
+    tensor_names: Iterable[str],
+) -> dict[str, torch.Tensor]:
+    tensors: dict[str, torch.Tensor] = {}
+    names_by_file: dict[str, list[str]] = defaultdict(list)
+    for name in tensor_names:
+        names_by_file[weight_map[name]].append(name)
+
+    for filename, names in names_by_file.items():
+        shard_path = os.path.join(model_dir, filename)
+        with safe_open(shard_path, framework="pt", device="cpu") as f:
+            for name in names:
+                tensors[name] = f.get_tensor(name).contiguous()
+    return tensors
+
+
+def build_modelopt_fp8_transformer(
+    *,
+    modelopt_hf_dir: str,
+    modelopt_backbone_ckpt: str,
+    output_dir: str,
+    base_transformer_dir: str | None = None,
+    model_type: str = "auto",
+    keep_bf16_patterns: Sequence[str] | None = None,
+    maxbound: float = FP8_E4M3_MAXBOUND,
+    overwrite: bool = False,
+) -> dict[str, int]:
+    source_dir = _resolve_transformer_dir(modelopt_hf_dir)
+    backbone_ckpt_path = _resolve_backbone_ckpt(modelopt_backbone_ckpt)
+    base_dir = (
+        _resolve_transformer_dir(base_transformer_dir) if base_transformer_dir else None
+    )
+
+    config = _load_config(source_dir)
+    quant_config = config.get("quantization_config")
+    if not isinstance(quant_config, dict):
+        raise ValueError(
+            "Expected a flat quantization_config dict in the ModelOpt export."
+        )
+    if quant_config.get("quant_method") != "modelopt":
+        raise ValueError(
+            "This tool only supports ModelOpt diffusers FP8 exports "
+            "(quant_method=modelopt)."
+        )
+
+    source_weight_map_all, index_filename = _load_weight_map(source_dir)
+    source_metadata = _load_first_shard_metadata(source_dir, source_weight_map_all)
+    is_ltx2_export = _is_ltx2_x0_export(
+        config=config,
+        source_metadata=source_metadata,
+        source_weight_map=source_weight_map_all,
+    )
+    class_name = config.get("_class_name")
+    runtime_name_mapper = _get_runtime_module_name_mapper(
+        model_type=model_type, class_name=class_name
+    )
+    ignore_patterns = list(quant_config.get("ignore", []) or [])
+    patterns = list(
+        get_default_keep_bf16_patterns(model_type=model_type, class_name=class_name)
+    )
+    if is_ltx2_export and model_type == "auto":
+        patterns.extend(DEFAULT_LTX2_KEEP_BF16_PATTERNS)
+    if keep_bf16_patterns:
+        patterns.extend(keep_bf16_patterns)
+    if patterns and base_dir is None and not is_ltx2_export:
+        raise ValueError(
+            "BF16 fallback patterns are enabled, but --base-transformer-dir was not provided."
+        )
+
+    output_path = Path(output_dir).expanduser().resolve()
+    if output_path.exists():
+        if not overwrite:
+            raise FileExistsError(
+                f"Output directory already exists: {output_path}. "
+                "Use --overwrite to replace it."
+            )
+        shutil.rmtree(output_path)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    _copy_non_shard_files(source_dir, str(output_path))
+
+    if is_ltx2_export:
+        source_weight_map = {
+            name: filename
+            for name, filename in source_weight_map_all.items()
+            if _should_keep_ltx2_transformer_key(name)
+        }
+    else:
+        source_weight_map = source_weight_map_all
+    base_weight_map: dict[str, str] = {}
+    if base_dir is not None:
+        base_weight_map, _ = _load_weight_map(base_dir)
+    fallback_weight_names = sorted(
+        weight_name
+        for weight_name in source_weight_map
+        if weight_name.endswith(".weight")
+        and should_keep_bf16(weight_name, patterns, runtime_name_mapper)
+    )
+    fallback_weight_names_set = set(fallback_weight_names)
+
+    backbone_state = torch.load(backbone_ckpt_path, map_location="cpu")[
+        "model_state_dict"
+    ]
+    fp8_scale_map = build_fp8_scale_map(backbone_state, maxbound=maxbound)
+    quant_algo = str(quant_config.get("quant_algo", "")).upper()
+    if quant_algo and "FP8" not in quant_algo:
+        raise ValueError(
+            "This tool only supports ModelOpt diffusers FP8 exports, "
+            f"got quant_algo={quant_config.get('quant_algo')!r}."
+        )
+    if not quant_algo and not fp8_scale_map:
+        raise ValueError(
+            "Could not infer an FP8 ModelOpt export: quantization_config.quant_algo "
+            "is missing and backbone.pt does not contain FP8 scale tensors."
+        )
+    effective_quant_config = json.loads(json.dumps(quant_config))
+    if not quant_algo:
+        effective_quant_config["quant_algo"] = "FP8"
+    effective_quant_config = (
+        normalize_flat_modelopt_quant_config(effective_quant_config)
+        or effective_quant_config
+    )
+
+    auto_ignore_modules = sorted(
+        {
+            _preferred_module_name(weight_name, runtime_name_mapper)
+            for weight_name in source_weight_map
+            if weight_name.endswith(".weight")
+            and _resolve_scale_key(weight_name, fp8_scale_map) is None
+        }
+    )
+    fallback_ignore_modules = sorted(
+        {
+            _preferred_module_name(weight_name, runtime_name_mapper)
+            for weight_name in fallback_weight_names
+        }
+    )
+    ignore_patterns = sorted(
+        {
+            *ignore_patterns,
+            *auto_ignore_modules,
+            *fallback_ignore_modules,
+        }
+    )
+    effective_quant_config["ignore"] = ignore_patterns
+    serialized_quant_config = json.dumps(effective_quant_config, sort_keys=True)
+    output_config = _build_output_config(
+        source_config=config,
+        source_metadata=source_metadata,
+        quant_config=effective_quant_config,
+        is_ltx2_x0_export=is_ltx2_export,
+    )
+
+    fallback_tensors = (
+        _load_selected_tensors(base_dir, base_weight_map, fallback_weight_names)
+        if fallback_weight_names and base_dir is not None
+        else {}
+    )
+    fallback_scale_names = {
+        scale_name
+        for weight_name in fallback_weight_names
+        for scale_name in (
+            weight_name[:-7] + ".weight_scale",
+            weight_name[:-7] + ".input_scale",
+        )
+    }
+
+    weights_by_file: dict[str, list[str]] = defaultdict(list)
+    for weight_name, filename in source_weight_map.items():
+        weights_by_file[filename].append(weight_name)
+
+    updated_weight_map: dict[str, str] = {}
+    total_size = 0
+    added_scale_count = 0
+    preserved_ignored_weight_count = 0
+
+    for filename, names in sorted(weights_by_file.items()):
+        shard_path = os.path.join(source_dir, filename)
+        shard_tensors = load_file(shard_path, device="cpu")
+        selected_names = set(names)
+
+        with safe_open(shard_path, framework="pt", device="cpu") as f:
+            metadata = dict(f.metadata() or {})
+
+        metadata.setdefault("format", "pt")
+        metadata["_class_name"] = str(
+            output_config.get("_class_name", metadata.get("_class_name", ""))
+        )
+        metadata["config"] = json.dumps(output_config, sort_keys=True)
+        metadata["quantization_config"] = serialized_quant_config
+        metadata["_quantization_metadata"] = serialized_quant_config
+
+        for name in list(shard_tensors.keys()):
+            if name not in selected_names:
+                del shard_tensors[name]
+                continue
+            if "_quantizer." in name:
+                del shard_tensors[name]
+                continue
+            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, runtime_name_mapper
+            ):
+                preserved_ignored_weight_count += 1
+                continue
+            scale_key = _resolve_scale_key(name, fp8_scale_map)
+            if (
+                name.endswith(".weight")
+                and scale_key is not None
+                and name not in fallback_tensors
+                and name not in fallback_weight_names_set
+            ):
+                scale_tensors = fp8_scale_map[scale_key]
+                shard_tensors[name] = quantize_fp8_weight(
+                    shard_tensors[name], scale_tensors["weight_scale"]
+                )
+                shard_tensors[name[:-7] + ".weight_scale"] = scale_tensors[
+                    "weight_scale"
+                ]
+                shard_tensors[name[:-7] + ".input_scale"] = scale_tensors["input_scale"]
+                added_scale_count += 2
+
+        save_file(shard_tensors, os.path.join(output_path, filename), metadata=metadata)
+
+        for name, tensor in shard_tensors.items():
+            updated_weight_map[name] = filename
+            total_size += tensor.element_size() * tensor.numel()
+
+        del shard_tensors
+        gc.collect()
+
+    with open(output_path / index_filename, "w", encoding="utf-8") as f:
+        json.dump(
+            {
+                "metadata": {"total_size": total_size},
+                "weight_map": updated_weight_map,
+            },
+            f,
+            indent=2,
+            sort_keys=True,
+        )
+
+    with open(output_path / "config.json", "w", encoding="utf-8") as f:
+        json.dump(output_config, f, indent=2, sort_keys=True)
+
+    return {
+        "quantized_weights": sum(
+            1
+            for name in source_weight_map
+            if name.endswith(".weight")
+            and _resolve_scale_key(name, fp8_scale_map) is not None
+            and not is_ignored_by_modelopt(name, ignore_patterns, runtime_name_mapper)
+        ),
+        "bf16_fallback_weights": len(fallback_weight_names),
+        "preserved_ignored_weights": preserved_ignored_weight_count,
+        "added_scale_tensors": added_scale_count,
+        "output_shards": len(weights_by_file),
+    }
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Build an SGLang-loadable ModelOpt FP8 diffusion transformer from a "
+            "ModelOpt diffusers export."
+        )
+    )
+    parser.add_argument(
+        "--modelopt-hf-dir",
+        required=True,
+        help="ModelOpt --hf-ckpt-dir output, or its transformer subdirectory.",
+    )
+    parser.add_argument(
+        "--modelopt-backbone-ckpt",
+        required=True,
+        help="Path to backbone.pt, or the directory that contains it.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        required=True,
+        help="Directory to write the converted SGLang transformer checkpoint.",
+    )
+    parser.add_argument(
+        "--base-transformer-dir",
+        help=(
+            "Original BF16 transformer directory (or parent model dir). Required when "
+            "BF16 fallback layers are enabled."
+        ),
+    )
+    parser.add_argument(
+        "--model-type",
+        choices=[
+            "auto",
+            "flux1",
+            "flux2",
+            "ltx2",
+            "hunyuan-video",
+            "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 / "
+            "HunyuanVideo / Qwen Image fallback sets when the export config matches "
+            "those transformer classes."
+        ),
+    )
+    parser.add_argument(
+        "--keep-bf16-pattern",
+        action="append",
+        default=[],
+        help=(
+            "Regex matched against module names without the trailing .weight. "
+            "Matching weights are copied from --base-transformer-dir instead of "
+            "staying in FP8."
+        ),
+    )
+    parser.add_argument(
+        "--maxbound",
+        type=float,
+        default=FP8_E4M3_MAXBOUND,
+        help="FP8 maxbound used to turn ModelOpt amax into a scale. E4M3 uses 448.",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Replace --output-dir if it already exists.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+    stats = build_modelopt_fp8_transformer(
+        modelopt_hf_dir=args.modelopt_hf_dir,
+        modelopt_backbone_ckpt=args.modelopt_backbone_ckpt,
+        output_dir=args.output_dir,
+        base_transformer_dir=args.base_transformer_dir,
+        model_type=args.model_type,
+        keep_bf16_patterns=args.keep_bf16_pattern,
+        maxbound=args.maxbound,
+        overwrite=args.overwrite,
+    )
+    print(json.dumps(stats, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py

@@ -0,0 +1,400 @@
+"""Build an SGLang-loadable ModelOpt NVFP4 diffusion transformer.
+
+This tool keeps the ModelOpt-exported NVFP4 tensors for most transformer
+modules, but can replace a validated subset of numerically sensitive modules
+with their original BF16 tensors from the base transformer checkpoint.
+
+It is primarily intended for FLUX.1-dev style ModelOpt NVFP4 exports where:
+- the base pipeline should remain separate from the quantized transformer
+- fallback BF16 modules are model-family specific
+- the serialized FP4 weight byte order may already match the runtime kernel
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import shutil
+from collections import defaultdict
+from pathlib import Path
+from typing import Iterable, Mapping, Sequence
+
+from safetensors import safe_open
+from safetensors.torch import load_file, save_file
+
+INDEX_FILENAMES = [
+    "model.safetensors.index.json",
+    "diffusion_pytorch_model.safetensors.index.json",
+]
+
+DEFAULT_FLUX1_NVFP4_FALLBACK_PATTERNS = [
+    "transformer_blocks.*.norm1.linear*",
+    "transformer_blocks.*.norm1_context.linear*",
+    "transformer_blocks.*.ff.net.0.proj*",
+    "transformer_blocks.*.ff.net.2*",
+    "transformer_blocks.*.ff_context.net.0.proj*",
+    "transformer_blocks.*.ff_context.net.2*",
+    "single_transformer_blocks.*.norm.linear*",
+    "single_transformer_blocks.*.proj_mlp*",
+]
+
+_TENSOR_MODULE_SUFFIXES = (
+    ".weight_scale_2",
+    ".weight_scale",
+    ".input_scale",
+    ".weight",
+    ".bias",
+)
+
+
+def _resolve_transformer_dir(path: str) -> str:
+    candidate = Path(path).expanduser().resolve()
+    if (candidate / "config.json").is_file():
+        return str(candidate)
+    transformer_dir = candidate / "transformer"
+    if (transformer_dir / "config.json").is_file():
+        return str(transformer_dir)
+    raise FileNotFoundError(f"Could not resolve a transformer directory from: {path}")
+
+
+def _find_index_file(model_dir: str) -> str | None:
+    for filename in INDEX_FILENAMES:
+        candidate = os.path.join(model_dir, filename)
+        if os.path.isfile(candidate):
+            return filename
+
+    matches = sorted(
+        filename
+        for filename in os.listdir(model_dir)
+        if filename.endswith(".safetensors.index.json")
+    )
+    return matches[0] if matches else None
+
+
+def _load_weight_map(model_dir: str) -> tuple[dict[str, str], str | None]:
+    index_filename = _find_index_file(model_dir)
+    if index_filename is not None:
+        with open(os.path.join(model_dir, index_filename), encoding="utf-8") as f:
+            index_data = json.load(f)
+        return dict(index_data["weight_map"]), index_filename
+
+    safetensors_files = sorted(
+        filename
+        for filename in os.listdir(model_dir)
+        if filename.endswith(".safetensors")
+    )
+    if len(safetensors_files) != 1:
+        raise ValueError(
+            f"Expected an index file or a single safetensors shard in {model_dir}, "
+            f"found {len(safetensors_files)} shard(s)."
+        )
+
+    shard_name = safetensors_files[0]
+    with safe_open(
+        os.path.join(model_dir, shard_name), framework="pt", device="cpu"
+    ) as f:
+        weight_map = {key: shard_name for key in f.keys()}
+    index_filename = f"{Path(shard_name).stem}.safetensors.index.json"
+    return weight_map, index_filename
+
+
+def _load_config(model_dir: str) -> dict:
+    config_path = os.path.join(model_dir, "config.json")
+    with open(config_path, encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _write_config(model_dir: Path, config: Mapping[str, object]) -> None:
+    with open(model_dir / "config.json", "w", encoding="utf-8") as f:
+        json.dump(config, f, indent=2, sort_keys=True)
+        f.write("\n")
+
+
+def _copy_non_shard_files(source_dir: str, output_dir: str) -> None:
+    ignored = set(INDEX_FILENAMES)
+    for entry in os.listdir(source_dir):
+        if entry.endswith(".safetensors") or entry in ignored:
+            continue
+        source_path = os.path.join(source_dir, entry)
+        output_path = os.path.join(output_dir, entry)
+        if os.path.isdir(source_path):
+            shutil.copytree(source_path, output_path, dirs_exist_ok=True)
+        else:
+            shutil.copy2(source_path, output_path)
+
+
+def _load_selected_tensors(
+    model_dir: str,
+    weight_map: Mapping[str, str],
+    tensor_names: Iterable[str],
+):
+    tensors = {}
+    names_by_file: dict[str, list[str]] = defaultdict(list)
+    for name in tensor_names:
+        names_by_file[weight_map[name]].append(name)
+
+    for filename, names in names_by_file.items():
+        shard_path = os.path.join(model_dir, filename)
+        with safe_open(shard_path, framework="pt", device="cpu") as f:
+            for name in names:
+                tensors[name] = f.get_tensor(name).contiguous()
+    return tensors
+
+
+def _module_name_for_tensor(tensor_name: str) -> str:
+    for suffix in _TENSOR_MODULE_SUFFIXES:
+        if tensor_name.endswith(suffix):
+            return tensor_name[: -len(suffix)]
+    return tensor_name
+
+
+def _matches_any_pattern(module_name: str, patterns: Sequence[str]) -> bool:
+    if not patterns:
+        return False
+    for pattern in patterns:
+        regex_str = pattern.replace(".", r"\.").replace("*", r".*")
+        if re.fullmatch(regex_str, module_name):
+            return True
+    return False
+
+
+def _preset_patterns(pattern_preset: str) -> list[str]:
+    if pattern_preset == "none":
+        return []
+    if pattern_preset == "flux1-nvfp4":
+        return list(DEFAULT_FLUX1_NVFP4_FALLBACK_PATTERNS)
+    raise ValueError(f"Unsupported pattern preset: {pattern_preset}")
+
+
+def _updated_quant_config(
+    source_config: Mapping[str, object],
+    *,
+    fallback_patterns: Sequence[str],
+    swap_weight_nibbles: bool,
+) -> dict[str, object]:
+    output_config = json.loads(json.dumps(source_config))
+    quant_config = output_config.get("quantization_config")
+    if not isinstance(quant_config, dict):
+        raise ValueError("Expected a flat quantization_config dict in config.json.")
+    if (
+        quant_config.get("quant_method") != "modelopt"
+        or "FP4" not in str(quant_config.get("quant_algo", "")).upper()
+    ):
+        raise ValueError(
+            "This tool only supports ModelOpt diffusion NVFP4 exports "
+            "(quant_method=modelopt, quant_algo=FP4/NVFP4)."
+        )
+
+    ignore_patterns = list(quant_config.get("ignore", []) or [])
+    for pattern in fallback_patterns:
+        if pattern not in ignore_patterns:
+            ignore_patterns.append(pattern)
+
+    quant_config["ignore"] = ignore_patterns
+    quant_config.setdefault(
+        "quant_type", str(quant_config.get("quant_algo", "")).upper()
+    )
+    quant_config["swap_weight_nibbles"] = swap_weight_nibbles
+    return output_config
+
+
+def build_modelopt_nvfp4_transformer(
+    *,
+    base_transformer_dir: str,
+    modelopt_hf_dir: str,
+    output_dir: str,
+    pattern_preset: str = "none",
+    keep_bf16_patterns: Sequence[str] | None = None,
+    swap_weight_nibbles: bool | None = None,
+    overwrite: bool = False,
+) -> dict[str, int | bool]:
+    source_dir = _resolve_transformer_dir(modelopt_hf_dir)
+    base_dir = _resolve_transformer_dir(base_transformer_dir)
+
+    patterns = _preset_patterns(pattern_preset)
+    if keep_bf16_patterns:
+        patterns.extend(keep_bf16_patterns)
+
+    resolved_swap_weight_nibbles = (
+        swap_weight_nibbles if swap_weight_nibbles is not None else False
+    )
+    output_config = _updated_quant_config(
+        _load_config(source_dir),
+        fallback_patterns=patterns,
+        swap_weight_nibbles=resolved_swap_weight_nibbles,
+    )
+    quant_config = output_config["quantization_config"]
+    serialized_quant_config = json.dumps(quant_config, sort_keys=True)
+
+    output_path = Path(output_dir).expanduser().resolve()
+    if output_path.exists():
+        if not overwrite:
+            raise FileExistsError(
+                f"Output directory already exists: {output_path}. "
+                "Use --overwrite to replace it."
+            )
+        shutil.rmtree(output_path)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    _copy_non_shard_files(source_dir, str(output_path))
+    _write_config(output_path, output_config)
+
+    source_weight_map, index_filename = _load_weight_map(source_dir)
+    base_weight_map, _ = _load_weight_map(base_dir)
+
+    fallback_tensor_names = sorted(
+        name
+        for name in base_weight_map
+        if name in source_weight_map
+        and _matches_any_pattern(_module_name_for_tensor(name), patterns)
+    )
+    fallback_tensors = _load_selected_tensors(
+        base_dir,
+        base_weight_map,
+        fallback_tensor_names,
+    )
+    fallback_modules = {
+        _module_name_for_tensor(tensor_name) for tensor_name in fallback_tensor_names
+    }
+
+    weights_by_file: dict[str, list[str]] = defaultdict(list)
+    for tensor_name, filename in source_weight_map.items():
+        weights_by_file[filename].append(tensor_name)
+
+    updated_weight_map: dict[str, str] = {}
+    total_size = 0
+    replaced_tensor_count = 0
+    removed_aux_tensor_count = 0
+
+    for filename, tensor_names in sorted(weights_by_file.items()):
+        shard_path = os.path.join(source_dir, filename)
+        shard_tensors = load_file(shard_path, device="cpu")
+
+        with safe_open(shard_path, framework="pt", device="cpu") as f:
+            metadata = dict(f.metadata() or {})
+
+        metadata.setdefault("format", "pt")
+        metadata["quantization_config"] = serialized_quant_config
+        metadata["_quantization_metadata"] = serialized_quant_config
+
+        for name in list(shard_tensors.keys()):
+            if "_quantizer." in name:
+                del shard_tensors[name]
+                removed_aux_tensor_count += 1
+                continue
+
+            module_name = _module_name_for_tensor(name)
+            if module_name not in fallback_modules:
+                continue
+
+            if name in fallback_tensors:
+                shard_tensors[name] = fallback_tensors[name]
+                replaced_tensor_count += 1
+            else:
+                del shard_tensors[name]
+                removed_aux_tensor_count += 1
+
+        save_file(shard_tensors, os.path.join(output_path, filename), metadata=metadata)
+
+        for name, tensor in shard_tensors.items():
+            updated_weight_map[name] = filename
+            total_size += tensor.element_size() * tensor.numel()
+
+    if index_filename is None:
+        raise ValueError(
+            "Expected a sharded or indexed ModelOpt HF export, but no index file was found."
+        )
+
+    with open(output_path / index_filename, "w", encoding="utf-8") as f:
+        json.dump(
+            {
+                "metadata": {"total_size": total_size},
+                "weight_map": updated_weight_map,
+            },
+            f,
+            indent=2,
+            sort_keys=True,
+        )
+        f.write("\n")
+
+    return {
+        "fallback_modules": len(fallback_modules),
+        "replaced_tensors": replaced_tensor_count,
+        "removed_aux_tensors": removed_aux_tensor_count,
+        "output_shards": len(weights_by_file),
+        "swap_weight_nibbles": resolved_swap_weight_nibbles,
+    }
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Build an SGLang-loadable ModelOpt NVFP4 diffusion transformer and "
+            "optionally keep selected modules in BF16."
+        )
+    )
+    parser.add_argument(
+        "--base-transformer-dir",
+        required=True,
+        help="Original BF16 transformer directory, or a parent model directory.",
+    )
+    parser.add_argument(
+        "--modelopt-hf-dir",
+        required=True,
+        help="ModelOpt --hf-ckpt-dir output, or its transformer subdirectory.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        required=True,
+        help="Directory to write the mixed transformer checkpoint.",
+    )
+    parser.add_argument(
+        "--pattern-preset",
+        choices=["none", "flux1-nvfp4"],
+        default="none",
+        help="Optional model-family BF16 fallback preset.",
+    )
+    parser.add_argument(
+        "--keep-bf16-pattern",
+        action="append",
+        default=[],
+        help=(
+            "Glob-style pattern matched against module names without trailing tensor "
+            "suffixes such as .weight or .bias."
+        ),
+    )
+    parser.add_argument(
+        "--swap-weight-nibbles",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+        help=(
+            "Whether the runtime should swap packed FP4 nibbles before padding. "
+            "Defaults to false."
+        ),
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Replace --output-dir if it already exists.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+    stats = build_modelopt_nvfp4_transformer(
+        base_transformer_dir=args.base_transformer_dir,
+        modelopt_hf_dir=args.modelopt_hf_dir,
+        output_dir=args.output_dir,
+        pattern_preset=args.pattern_preset,
+        keep_bf16_patterns=args.keep_bf16_pattern,
+        swap_weight_nibbles=args.swap_weight_nibbles,
+        overwrite=args.overwrite,
+    )
+    print(json.dumps(stats, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py

@@ -0,0 +1,631 @@
+"""Compare diffusion BF16 and quantized runs via trajectory-latent similarity.
+
+This tool runs two SGLang diffusion variants with the same prompt and seed,
+captures intermediate denoising latents via `return_trajectory_latents`, and
+reports cosine / error metrics for each timestep plus final frame metrics.
+
+The intended use is quant validation with reduced deterministic settings:
+- same prompt / seed / resolution / step count for both variants
+- BF16 reference on the base model
+- FP8 candidate via `--candidate-transformer-path` and/or component overrides
+
+Example:
+
+    python -m sglang.multimodal_gen.tools.compare_diffusion_trajectory_similarity \
+        --model-path /path/to/model \
+        --prompt "A futuristic cyberpunk city at night" \
+        --width 512 --height 512 --num-inference-steps 8 --seed 42 \
+        --text-encoder-cpu-offload \
+        --candidate-transformer-path /tmp/modelopt_flux2_fp8/sglang_transformer \
+        --output-json /tmp/flux2_similarity.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import json
+import math
+import os
+from pathlib import Path
+from typing import Any, Sequence
+
+import imageio.v3 as iio
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+def parse_component_overrides(entries: Sequence[str] | None) -> dict[str, str]:
+    overrides: dict[str, str] = {}
+    for entry in entries or []:
+        if "=" not in entry:
+            raise ValueError(
+                f"Invalid component override '{entry}'. Expected format component=path."
+            )
+        component, path = entry.split("=", 1)
+        component = component.strip().replace("-", "_")
+        path = path.strip()
+        if not component or not path:
+            raise ValueError(
+                f"Invalid component override '{entry}'. Expected format component=path."
+            )
+        overrides[component] = path
+    return overrides
+
+
+def _cosine_similarity(flat_a: torch.Tensor, flat_b: torch.Tensor) -> float:
+    norm_a = torch.linalg.vector_norm(flat_a).item()
+    norm_b = torch.linalg.vector_norm(flat_b).item()
+    if norm_a == 0.0 and norm_b == 0.0:
+        return 1.0
+    if norm_a == 0.0 or norm_b == 0.0:
+        return 0.0
+    return float(F.cosine_similarity(flat_a, flat_b, dim=0).item())
+
+
+def compute_tensor_metrics(lhs: Any, rhs: Any) -> dict[str, float]:
+    lhs_tensor = torch.as_tensor(lhs).detach().cpu().float()
+    rhs_tensor = torch.as_tensor(rhs).detach().cpu().float()
+    if lhs_tensor.shape != rhs_tensor.shape:
+        raise ValueError(
+            f"Metric shape mismatch: {tuple(lhs_tensor.shape)} vs {tuple(rhs_tensor.shape)}"
+        )
+
+    diff = lhs_tensor - rhs_tensor
+    mse = float(diff.square().mean().item())
+    rmse = float(math.sqrt(mse))
+    mae = float(diff.abs().mean().item())
+    max_abs = float(diff.abs().max().item())
+    l2 = float(torch.linalg.vector_norm(diff).item())
+    cosine = _cosine_similarity(lhs_tensor.reshape(-1), rhs_tensor.reshape(-1))
+    return {
+        "cosine_similarity": cosine,
+        "mae": mae,
+        "mse": mse,
+        "rmse": rmse,
+        "max_abs": max_abs,
+        "l2": l2,
+    }
+
+
+def compute_uint8_frame_metrics(lhs: Any, rhs: Any) -> dict[str, float]:
+    metrics = compute_tensor_metrics(lhs, rhs)
+    mse = metrics["mse"]
+    metrics["psnr_db"] = (
+        float("inf") if mse == 0.0 else 20 * math.log10(255.0) - 10 * math.log10(mse)
+    )
+    return metrics
+
+
+def _normalize_step_index(step_index: int, num_steps: int) -> int:
+    if num_steps <= 0:
+        raise ValueError("num_steps must be positive.")
+    if step_index < 0:
+        step_index += num_steps
+    if step_index < 0 or step_index >= num_steps:
+        raise IndexError(
+            f"Requested step index {step_index} is outside the valid range [0, {num_steps})."
+        )
+    return step_index
+
+
+def _maybe_scalar(timestep: torch.Tensor | None, index: int) -> float | None:
+    if timestep is None:
+        return None
+    value = timestep[index]
+    if isinstance(value, torch.Tensor):
+        value = value.detach().cpu()
+        if value.numel() == 1:
+            return float(value.item())
+    return float(value)
+
+
+def summarize_trajectory_metrics(
+    reference_latents: Any,
+    candidate_latents: Any,
+    *,
+    reference_timesteps: Any = None,
+    candidate_timesteps: Any = None,
+    step_index: int = -1,
+) -> dict[str, Any]:
+    ref = torch.as_tensor(reference_latents).detach().cpu().float()
+    cand = torch.as_tensor(candidate_latents).detach().cpu().float()
+    if ref.shape != cand.shape:
+        raise ValueError(
+            f"Trajectory shape mismatch: {tuple(ref.shape)} vs {tuple(cand.shape)}"
+        )
+    if ref.ndim < 2:
+        raise ValueError(
+            f"Expected trajectory latents with an explicit timestep dimension, got {tuple(ref.shape)}"
+        )
+
+    num_steps = ref.shape[1]
+    selected_step = _normalize_step_index(step_index, num_steps)
+    ref_t = (
+        torch.as_tensor(reference_timesteps).detach().cpu()
+        if reference_timesteps is not None
+        else None
+    )
+    cand_t = (
+        torch.as_tensor(candidate_timesteps).detach().cpu()
+        if candidate_timesteps is not None
+        else None
+    )
+
+    per_step: list[dict[str, Any]] = []
+    for idx in range(num_steps):
+        metrics = compute_tensor_metrics(ref[:, idx], cand[:, idx])
+        metrics["step_index"] = idx
+        metrics["reference_timestep"] = _maybe_scalar(ref_t, idx)
+        metrics["candidate_timestep"] = _maybe_scalar(cand_t, idx)
+        per_step.append(metrics)
+
+    return {
+        "trajectory_shape": list(ref.shape),
+        "num_steps": num_steps,
+        "selected_step_index": selected_step,
+        "selected_step_metrics": per_step[selected_step],
+        "per_step_metrics": per_step,
+    }
+
+
+def summarize_output_frame_metrics(
+    reference_frames: Sequence[Any],
+    candidate_frames: Sequence[Any],
+) -> dict[str, Any]:
+    if len(reference_frames) != len(candidate_frames):
+        raise ValueError(
+            f"Output frame count mismatch: {len(reference_frames)} vs {len(candidate_frames)}"
+        )
+    if not reference_frames:
+        raise ValueError("No output frames available for comparison.")
+
+    ref_stack = np.stack([np.asarray(frame) for frame in reference_frames], axis=0)
+    cand_stack = np.stack([np.asarray(frame) for frame in candidate_frames], axis=0)
+
+    frame0_metrics = compute_uint8_frame_metrics(ref_stack[0], cand_stack[0])
+    mid_index = len(reference_frames) // 2
+    mid_metrics = compute_uint8_frame_metrics(
+        ref_stack[mid_index], cand_stack[mid_index]
+    )
+    all_metrics = compute_uint8_frame_metrics(ref_stack, cand_stack)
+
+    return {
+        "num_frames": len(reference_frames),
+        "frame0_metrics": frame0_metrics,
+        "mid_frame_index": mid_index,
+        "mid_frame_metrics": mid_metrics,
+        "all_frames_metrics": all_metrics,
+    }
+
+
+def extract_result_frames(result: Any) -> list[np.ndarray]:
+    if result.frames is not None:
+        return [np.asarray(frame) for frame in result.frames]
+
+    sample = result.samples
+    if sample is None:
+        if result.output_file_path:
+            output_path = Path(result.output_file_path)
+            if not output_path.exists():
+                raise ValueError(
+                    "GenerationResult did not contain frames or samples, and its "
+                    f"output_file_path does not exist: {output_path}"
+                )
+            if output_path.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}:
+                return [np.asarray(iio.imread(output_path))]
+            return [np.asarray(frame) for frame in iio.imiter(output_path)]
+        raise ValueError(
+            "GenerationResult did not contain frames, samples, or a readable output_file_path."
+        )
+
+    if isinstance(sample, torch.Tensor):
+        tensor = sample.detach().cpu().float()
+        if tensor.ndim == 3:
+            tensor = tensor.unsqueeze(1)
+        if tensor.ndim != 4:
+            raise ValueError(
+                f"Unsupported tensor sample shape for frame extraction: {tuple(tensor.shape)}"
+            )
+        tensor = (tensor * 255).clamp(0, 255).to(torch.uint8)
+        frames = tensor.permute(1, 2, 3, 0).contiguous().numpy()
+        return [frame for frame in frames]
+
+    array = np.asarray(sample)
+    if array.ndim == 2:
+        array = array[..., None]
+    if array.ndim == 3:
+        if array.shape[-1] in (1, 3, 4):
+            array = array[None, ...]
+        else:
+            array = array[..., None]
+    if array.ndim != 4:
+        raise ValueError(
+            f"Unsupported numpy sample shape for frame extraction: {tuple(array.shape)}"
+        )
+    if array.dtype != np.uint8:
+        array = (np.clip(array, 0.0, 1.0) * 255.0).astype(np.uint8)
+    return [frame for frame in array]
+
+
+def build_server_kwargs(args: argparse.Namespace, *, variant: str) -> dict[str, Any]:
+    component_paths = parse_component_overrides(
+        getattr(args, f"{variant}_component_path") or []
+    )
+    transformer_path = getattr(args, f"{variant}_transformer_path")
+
+    kwargs: dict[str, Any] = {
+        "model_path": args.model_path,
+        "model_id": args.model_id,
+        "backend": args.backend,
+        "num_gpus": args.num_gpus,
+        "dit_cpu_offload": args.dit_cpu_offload,
+        "dit_layerwise_offload": args.dit_layerwise_offload,
+        "text_encoder_cpu_offload": args.text_encoder_cpu_offload,
+        "vae_cpu_offload": args.vae_cpu_offload,
+        "pin_cpu_memory": args.pin_cpu_memory,
+        "enable_cfg_parallel": args.enable_cfg_parallel,
+        "ulysses_degree": args.ulysses_degree,
+    }
+    if args.sp_degree is not None:
+        kwargs["sp_degree"] = args.sp_degree
+    if transformer_path is not None:
+        kwargs["transformer_weights_path"] = transformer_path
+    if component_paths:
+        kwargs["component_paths"] = component_paths
+    return kwargs
+
+
+def build_sampling_kwargs(
+    args: argparse.Namespace, *, output_dir: str | None = None
+) -> dict[str, Any]:
+    kwargs: dict[str, Any] = {
+        "prompt": args.prompt,
+        "width": args.width,
+        "height": args.height,
+        "num_inference_steps": args.num_inference_steps,
+        "guidance_scale": args.guidance_scale,
+        "seed": args.seed,
+        "return_frames": True,
+        "return_trajectory_latents": True,
+        "return_trajectory_decoded": args.return_trajectory_decoded,
+        "save_output": output_dir is not None,
+    }
+    if output_dir is not None:
+        kwargs["output_path"] = output_dir
+    if args.num_frames is not None:
+        kwargs["num_frames"] = args.num_frames
+    if args.guidance_scale_2 is not None:
+        kwargs["guidance_scale_2"] = args.guidance_scale_2
+    return kwargs
+
+
+def _normalize_single_result(result: Any):
+    if isinstance(result, list):
+        if len(result) != 1:
+            raise ValueError(
+                f"Expected a single generation result, got {len(result)} results."
+            )
+        result = result[0]
+    if result is None:
+        raise RuntimeError("Generation returned no result.")
+    return result
+
+
+def _clear_diffusion_fp4_backend_caches() -> None:
+    from sglang.multimodal_gen.runtime.layers.quantization import (
+        modelopt_quant as diffusion_modelopt_quant,
+    )
+    from sglang.multimodal_gen.runtime.platforms import current_platform
+
+    diffusion_modelopt_quant._get_fp4_gemm_op.cache_clear()
+    current_platform.__class__.get_modelopt_fp4_gemm_op.cache_clear()
+    current_platform.__class__.get_modelopt_flashinfer_fp4_backend.cache_clear()
+
+
+@contextlib.contextmanager
+def override_diffusion_fp4_backend(backend: str | None):
+    env_name = "SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND"
+    previous = os.environ.get(env_name)
+
+    if backend is None:
+        os.environ.pop(env_name, None)
+    else:
+        os.environ[env_name] = backend
+
+    _clear_diffusion_fp4_backend_caches()
+    try:
+        yield
+    finally:
+        if previous is None:
+            os.environ.pop(env_name, None)
+        else:
+            os.environ[env_name] = previous
+        _clear_diffusion_fp4_backend_caches()
+
+
+def _extract_total_duration_ms(result: Any) -> float | None:
+    metrics = getattr(result, "metrics", None)
+    if not isinstance(metrics, dict):
+        return None
+    total_duration_ms = metrics.get("total_duration_ms")
+    if total_duration_ms is None:
+        return None
+    return float(total_duration_ms)
+
+
+def run_variant(
+    *,
+    server_kwargs: dict[str, Any],
+    sampling_kwargs: dict[str, Any],
+    fp4_gemm_backend: str | None,
+    warmup_runs: int,
+    measure_runs: int,
+):
+    from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import (
+        DiffGenerator,
+    )
+
+    if warmup_runs < 0:
+        raise ValueError("warmup_runs must be >= 0.")
+    if measure_runs <= 0:
+        raise ValueError("measure_runs must be >= 1.")
+
+    with override_diffusion_fp4_backend(fp4_gemm_backend):
+        with DiffGenerator.from_pretrained(
+            local_mode=True, **server_kwargs
+        ) as generator:
+            for _ in range(warmup_runs):
+                _normalize_single_result(
+                    generator.generate(sampling_params_kwargs=sampling_kwargs)
+                )
+
+            measured_results = []
+            for _ in range(measure_runs):
+                measured_results.append(
+                    _normalize_single_result(
+                        generator.generate(sampling_params_kwargs=sampling_kwargs)
+                    )
+                )
+
+    final_result = measured_results[-1]
+    generation_times = [float(result.generation_time) for result in measured_results]
+    peak_memories = [float(result.peak_memory_mb) for result in measured_results]
+    total_duration_ms = [
+        duration
+        for duration in (
+            _extract_total_duration_ms(result) for result in measured_results
+        )
+        if duration is not None
+    ]
+
+    return {
+        "result": final_result,
+        "fp4_gemm_backend": fp4_gemm_backend or "default",
+        "warmup_runs": warmup_runs,
+        "measure_runs": measure_runs,
+        "generation_time_s": generation_times[-1],
+        "avg_generation_time_s": sum(generation_times) / len(generation_times),
+        "per_run_generation_time_s": generation_times,
+        "peak_memory_mb": peak_memories[-1],
+        "max_peak_memory_mb": max(peak_memories) if peak_memories else 0.0,
+        "per_run_peak_memory_mb": peak_memories,
+        "total_duration_ms": total_duration_ms[-1] if total_duration_ms else None,
+        "avg_total_duration_ms": (
+            sum(total_duration_ms) / len(total_duration_ms)
+            if total_duration_ms
+            else None
+        ),
+        "per_run_total_duration_ms": total_duration_ms,
+    }
+
+
+def _to_jsonable(result: dict[str, Any]) -> dict[str, Any]:
+    return json.loads(json.dumps(result, allow_nan=True))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model-path", required=True)
+    parser.add_argument(
+        "--model-id",
+        help=(
+            "Optional model ID override passed to DiffGenerator.from_pretrained. "
+            "Use this when --model-path points to a local directory whose name "
+            "does not match a registered native SGLang model."
+        ),
+    )
+    parser.add_argument("--backend", default="sglang")
+    parser.add_argument("--prompt", required=True)
+    parser.add_argument("--output-json", required=True)
+    parser.add_argument("--width", type=int, required=True)
+    parser.add_argument("--height", type=int, required=True)
+    parser.add_argument("--num-frames", type=int)
+    parser.add_argument("--num-inference-steps", type=int, required=True)
+    parser.add_argument("--guidance-scale", type=float, required=True)
+    parser.add_argument("--guidance-scale-2", type=float)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--num-gpus", type=int, default=1)
+    parser.add_argument("--ulysses-degree", type=int, default=1)
+    parser.add_argument("--sp-degree", type=int)
+    parser.add_argument("--trajectory-step-index", type=int, default=-1)
+    parser.add_argument("--reference-transformer-path")
+    parser.add_argument("--candidate-transformer-path")
+    parser.add_argument(
+        "--reference-fp4-gemm-backend",
+        help=(
+            "Optional NVFP4 GEMM backend override for the reference run, e.g. "
+            "'flashinfer_trtllm'."
+        ),
+    )
+    parser.add_argument(
+        "--candidate-fp4-gemm-backend",
+        help=(
+            "Optional NVFP4 GEMM backend override for the candidate run, e.g. "
+            "'flashinfer_trtllm'."
+        ),
+    )
+    parser.add_argument("--warmup-runs", type=int, default=0)
+    parser.add_argument("--measure-runs", type=int, default=1)
+    parser.add_argument(
+        "--reference-component-path",
+        action="append",
+        default=[],
+        help="Repeatable component override in the form component=path.",
+    )
+    parser.add_argument(
+        "--candidate-component-path",
+        action="append",
+        default=[],
+        help="Repeatable component override in the form component=path.",
+    )
+    parser.add_argument("--save-output-dir")
+    parser.add_argument(
+        "--return-trajectory-decoded",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
+    parser.add_argument(
+        "--enable-cfg-parallel",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
+    parser.add_argument(
+        "--text-encoder-cpu-offload",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
+    parser.add_argument(
+        "--vae-cpu-offload",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
+    parser.add_argument(
+        "--dit-cpu-offload",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
+    parser.add_argument(
+        "--dit-layerwise-offload",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
+    parser.add_argument(
+        "--pin-cpu-memory",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
+    args = parser.parse_args()
+
+    output_json = Path(args.output_json).expanduser().resolve()
+    output_json.parent.mkdir(parents=True, exist_ok=True)
+
+    save_root: Path | None = None
+    if args.save_output_dir:
+        save_root = Path(args.save_output_dir).expanduser().resolve()
+        save_root.mkdir(parents=True, exist_ok=True)
+
+    ref_server_kwargs = build_server_kwargs(args, variant="reference")
+    cand_server_kwargs = build_server_kwargs(args, variant="candidate")
+
+    ref_sampling_kwargs = build_sampling_kwargs(
+        args,
+        output_dir=str(save_root / "reference") if save_root else None,
+    )
+    cand_sampling_kwargs = build_sampling_kwargs(
+        args,
+        output_dir=str(save_root / "candidate") if save_root else None,
+    )
+
+    reference_run = run_variant(
+        server_kwargs=ref_server_kwargs,
+        sampling_kwargs=ref_sampling_kwargs,
+        fp4_gemm_backend=args.reference_fp4_gemm_backend,
+        warmup_runs=args.warmup_runs,
+        measure_runs=args.measure_runs,
+    )
+    candidate_run = run_variant(
+        server_kwargs=cand_server_kwargs,
+        sampling_kwargs=cand_sampling_kwargs,
+        fp4_gemm_backend=args.candidate_fp4_gemm_backend,
+        warmup_runs=args.warmup_runs,
+        measure_runs=args.measure_runs,
+    )
+    reference = reference_run["result"]
+    candidate = candidate_run["result"]
+
+    result = {
+        "model_path": args.model_path,
+        "prompt": args.prompt,
+        "seed": args.seed,
+        "warmup_runs": args.warmup_runs,
+        "measure_runs": args.measure_runs,
+        "server_kwargs": {
+            "reference": ref_server_kwargs,
+            "candidate": cand_server_kwargs,
+        },
+        "backend_overrides": {
+            "reference_fp4_gemm_backend": reference_run["fp4_gemm_backend"],
+            "candidate_fp4_gemm_backend": candidate_run["fp4_gemm_backend"],
+        },
+        "sampling_kwargs": {
+            "width": args.width,
+            "height": args.height,
+            "num_frames": args.num_frames,
+            "num_inference_steps": args.num_inference_steps,
+            "guidance_scale": args.guidance_scale,
+            "guidance_scale_2": args.guidance_scale_2,
+        },
+        "reference_generation": {
+            key: value for key, value in reference_run.items() if key != "result"
+        }
+        | {"output_file_path": reference.output_file_path},
+        "candidate_generation": {
+            key: value for key, value in candidate_run.items() if key != "result"
+        }
+        | {"output_file_path": candidate.output_file_path},
+        "trajectory_metrics": summarize_trajectory_metrics(
+            reference.trajectory_latents,
+            candidate.trajectory_latents,
+            reference_timesteps=reference.trajectory_timesteps,
+            candidate_timesteps=candidate.trajectory_timesteps,
+            step_index=args.trajectory_step_index,
+        ),
+        "output_metrics": summarize_output_frame_metrics(
+            extract_result_frames(reference),
+            extract_result_frames(candidate),
+        ),
+    }
+
+    output_json.write_text(
+        json.dumps(_to_jsonable(result), indent=2, sort_keys=True), encoding="utf-8"
+    )
+
+    selected = result["trajectory_metrics"]["selected_step_metrics"]
+    frame0 = result["output_metrics"]["frame0_metrics"]
+    print(
+        json.dumps(
+            {
+                "output_json": str(output_json),
+                "trajectory_selected_step": result["trajectory_metrics"][
+                    "selected_step_index"
+                ],
+                "reference_avg_generation_time_s": result["reference_generation"][
+                    "avg_generation_time_s"
+                ],
+                "candidate_avg_generation_time_s": result["candidate_generation"][
+                    "avg_generation_time_s"
+                ],
+                "trajectory_cosine": selected["cosine_similarity"],
+                "trajectory_mae": selected["mae"],
+                "frame0_psnr_db": frame0["psnr_db"],
+                "frame0_mae": frame0["mae"],
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

python/sglang/multimodal_gen/tools/convert_hf_to_fp8.py

@@ -0,0 +1,320 @@
+# copied and adapted from Slime
+"""
+Convert HuggingFace safetensors model to FP8 format for efficient inference.
+
+Example usage:
+    # convert FLUX.1-dev transformer to FP8
+    python -m sglang.multimodal_gen.tools.convert_hf_to_fp8 \
+        --model-dir /path/to/FLUX.1-dev/transformer \
+        --save-dir /path/to/FLUX.1-dev/transformer-FP8 \
+        --strategy block \
+        --block-size 128 128
+
+Options:
+    --model-dir MODEL_DIR
+                        path to the directory of the HF safetensors model (e.g., transformer subfolder)
+    --save-dir SAVE_DIR
+                        path to the directory to save the converted FP8 model
+    --strategy {block,channel,tensor}
+                        quantization strategy (default: block)
+    --block-size [BLOCK_SIZE ...]
+                        block size for block quantization, e.g., --block-size 128 128
+    --max-workers MAX_WORKERS
+                        number of worker threads for parallel processing (default: 1)
+"""
+
+import argparse
+import gc
+import json
+import os
+import shutil
+import threading
+from concurrent.futures import ThreadPoolExecutor
+
+import safetensors
+import safetensors.torch
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+FP8_INFO = torch.finfo(torch.float8_e4m3fn)
+FP8_MAX, FP8_MIN = FP8_INFO.max, FP8_INFO.min
+
+
+def ceildiv(a, b):
+    return -(-a // b)
+
+
+def block_fp8(weight, block_size):
+
+    # per block quant
+    block_n, block_k = block_size[0], block_size[1]
+
+    shape_0, shape_1 = weight.shape
+
+    n_tiles = ceildiv(shape_0, block_n)
+    k_tiles = ceildiv(shape_1, block_k)
+
+    q_weight = F.pad(
+        weight,
+        (0, k_tiles * block_k - shape_1, 0, n_tiles * block_n - shape_0),
+        mode="constant",
+        value=0.0,
+    )
+
+    qweight = q_weight.reshape(n_tiles, block_n, k_tiles, block_k)
+    block_max = torch.max(torch.abs(qweight), dim=1, keepdim=True)[0]
+    block_max = torch.max(block_max, dim=3, keepdim=True)[0]
+
+    scale = block_max.to(torch.float32) / FP8_MAX
+    qweight = (
+        (qweight / scale)
+        .clamp(min=FP8_MIN, max=FP8_MAX)
+        .reshape((n_tiles * block_n, k_tiles * block_k))
+        .to(torch.float8_e4m3fn)
+    )
+    qweight = qweight[:shape_0, :shape_1].clone().detach()
+    scale = scale.squeeze()
+
+    return qweight, scale
+
+
+def channel_fp8(weight):
+    channel_max = torch.max(weight.abs(), dim=-1, keepdim=True)[0]
+    scale = channel_max.clamp(min=1e-12).to(torch.float32) / FP8_MAX
+    qweight = (weight / scale).clamp(min=FP8_MIN, max=FP8_MAX)
+    qweight = qweight.to(torch.float8_e4m3fn)
+    return qweight, scale
+
+
+def tensor_fp8(weight):
+    scale = weight.abs().max().clamp(min=1e-12).to(torch.float32) / FP8_MAX
+    qweight = (weight / scale).clamp(min=FP8_MIN, max=FP8_MAX)
+    qweight = qweight.to(torch.float8_e4m3fn)
+    scale = scale.view(1)
+    return qweight, scale
+
+
+def quant_fp8(weight, strategy, block_size=None):
+    if strategy == "tensor":
+        return tensor_fp8(weight)
+    elif strategy == "channel":
+        return channel_fp8(weight)
+    else:
+        return block_fp8(weight, block_size)
+
+
+class ConversionResult:
+    def __init__(self):
+        self.lock = threading.Lock()
+        self.weight_map = {}
+        self.param_count = 0
+        self.modules_to_not_convert = []
+
+    def add_result(self, filename, q_weights, module_names):
+        with self.lock:
+            for k, v in q_weights.items():
+                self.weight_map[k] = filename
+                self.param_count += v.numel()
+            self.modules_to_not_convert.extend(module_names)
+
+
+def process_file(
+    input_path, output_path, filename, strategy, block_size, result_collector
+):
+    if not filename.endswith(".safetensors"):
+        return
+
+    print(f"Processing {filename}, memory usage: {torch.cuda.memory_allocated()}")
+    weights = {}
+    q_weights = {}
+
+    with safetensors.safe_open(
+        os.path.join(input_path, filename), framework="pt", device="cuda"
+    ) as f:
+        for k in f.keys():
+            weights[k] = f.get_tensor(k)
+
+    modules_to_not_convert = []
+    for key in weights.keys():
+        if (
+            "weight" in key
+            and "layernorm" not in key
+            and "embed" not in key
+            and "router" not in key
+            and "mlp.gate." not in key
+            and "norm" not in key
+            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 "modulation" 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)
+            q_weights[key] = qw
+            if block_size:
+                scale_name = key.replace(".weight", ".weight_scale_inv")
+            else:
+                scale_name = key.replace(".weight", ".weight_scale")
+            q_weights[scale_name] = s
+        else:
+            modules_to_not_convert.append(key.replace(".weight", ""))
+            q_weights[key] = weights[key]
+
+    safetensors.torch.save_file(
+        q_weights, os.path.join(output_path, filename), metadata={"format": "pt"}
+    )
+
+    result_collector.add_result(filename, q_weights, modules_to_not_convert)
+
+
+def convert_fp8(input_path, output_path, strategy, block_size=None, max_workers=4):
+    input_path = os.path.abspath(input_path)
+    os.makedirs(output_path, exist_ok=True)
+
+    for filename in os.listdir(input_path):
+        if not filename.endswith(".safetensors") and not os.path.isdir(
+            os.path.join(input_path, filename)
+        ):
+            shutil.copyfile(
+                os.path.join(input_path, filename), os.path.join(output_path, filename)
+            )
+
+    safetensors_files = [
+        f for f in os.listdir(input_path) if f.endswith(".safetensors")
+    ]
+
+    result_collector = ConversionResult()
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = []
+        for filename in safetensors_files:
+            future = executor.submit(
+                process_file,
+                input_path,
+                output_path,
+                filename,
+                strategy,
+                block_size,
+                result_collector,
+            )
+            futures.append(future)
+
+        for future in tqdm(futures, desc="Processing files"):
+            future.result()
+
+    if strategy == "block" or strategy == "tensor":
+        quantization_config = {
+            "activation_scheme": "dynamic",
+            "fmt": "e4m3",
+            "quant_method": "fp8",
+        }
+        if block_size:
+            quantization_config["weight_block_size"] = block_size
+        if len(result_collector.modules_to_not_convert) > 0:
+            quantization_config["modules_to_not_convert"] = list(
+                set(result_collector.modules_to_not_convert)
+            )
+    else:
+        quant_group = {
+            "group_0": {
+                "input_activations": {
+                    "actorder": None,
+                    "block_structure": None,
+                    "dynamic": True,
+                    "group_size": None,
+                    "num_bits": 8,
+                    "observer": None,
+                    "observer_kwargs": {},
+                    "strategy": "token",
+                    "symmetric": True,
+                    "type": "float",
+                },
+                "output_activations": None,
+                "targets": ["Linear"],
+                "weights": {
+                    "actorder": None,
+                    "block_structure": None,
+                    "dynamic": False,
+                    "group_size": None,
+                    "num_bits": 8,
+                    "observer": "minmax",
+                    "observer_kwargs": {},
+                    "strategy": strategy,
+                    "symmetric": True,
+                    "type": "float",
+                },
+            },
+        }
+        quantization_config = {
+            "config_groups": quant_group,
+            "format": "float-quantized",
+            "ignore": list(set(result_collector.modules_to_not_convert)),
+            "quant_method": "compressed-tensors",
+            "quantization_status": "compressed",
+        }
+
+    config_path = os.path.join(input_path, "config.json")
+    if os.path.exists(config_path):
+        cfg = json.load(open(config_path))
+        cfg["quantization_config"] = quantization_config
+        json.dump(cfg, open(os.path.join(output_path, "config.json"), "w"), indent=2)
+
+    index_dict = {
+        "weight_map": result_collector.weight_map,
+        "metadata": {"total_size": result_collector.param_count},
+    }
+    json.dump(
+        index_dict,
+        open(os.path.join(output_path, "model.safetensors.index.json"), "w"),
+        indent=2,
+    )
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model-dir",
+        type=str,
+        help="Path to the directory of the HF safetensors model.",
+    )
+    parser.add_argument(
+        "--save-dir",
+        type=str,
+        help="Path to the directory to save the converted model.",
+    )
+    parser.add_argument(
+        "--strategy", type=str, default="block", choices=["block", "channel", "tensor"]
+    )
+    parser.add_argument(
+        "--block-size", type=int, nargs="*", default=None, help="eg. --block-size 32 32"
+    )
+    parser.add_argument(
+        "--max-workers",
+        type=int,
+        default=8,
+        help="Number of worker threads for parallel processing",
+    )
+    args = parser.parse_args()
+
+    if not os.path.exists(args.save_dir):
+        print(f"Creating directory {args.save_dir}")
+        os.makedirs(args.save_dir)
+    elif not os.path.isdir(args.save_dir):
+        raise ValueError("The save_dir should be a directory.")
+
+    convert_fp8(
+        args.model_dir, args.save_dir, args.strategy, args.block_size, args.max_workers
+    )

python/sglang/multimodal_gen/tools/wan_repack.py

@@ -0,0 +1,228 @@
+### Based on https://github.com/huggingface/diffusers/blob/main/scripts/convert_wan_to_diffusers.py
+
+import argparse
+import json
+import pathlib
+import shutil
+from typing import Any, Dict, List
+
+from safetensors.torch import load_file, save_file
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+TRANSFORMER_KEYS_RENAME_DICT = {
+    "time_embedding.0": "condition_embedder.time_embedder.linear_1",
+    "time_embedding.2": "condition_embedder.time_embedder.linear_2",
+    "text_embedding.0": "condition_embedder.text_embedder.linear_1",
+    "text_embedding.2": "condition_embedder.text_embedder.linear_2",
+    "time_projection.1": "condition_embedder.time_proj",
+    "head.modulation": "scale_shift_table",
+    "head.head": "proj_out",
+    "modulation": "scale_shift_table",
+    "ffn.0": "ffn.net.0.proj",
+    "ffn.2": "ffn.net.2",
+    # Hack to swap the layer names
+    # The original model calls the norms in following order: norm1, norm3, norm2
+    # We convert it to: norm1, norm2, norm3
+    "norm2": "norm__placeholder",
+    "norm3": "norm2",
+    "norm__placeholder": "norm3",
+    # For the I2V model
+    "img_emb.proj.0": "condition_embedder.image_embedder.norm1",
+    "img_emb.proj.1": "condition_embedder.image_embedder.ff.net.0.proj",
+    "img_emb.proj.3": "condition_embedder.image_embedder.ff.net.2",
+    "img_emb.proj.4": "condition_embedder.image_embedder.norm2",
+    # for the FLF2V model
+    "img_emb.emb_pos": "condition_embedder.image_embedder.pos_embed",
+    # Add attention component mappings
+    "self_attn.q": "attn1.to_q",
+    "self_attn.k": "attn1.to_k",
+    "self_attn.v": "attn1.to_v",
+    "self_attn.o": "attn1.to_out.0",
+    "self_attn.norm_q": "attn1.norm_q",
+    "self_attn.norm_k": "attn1.norm_k",
+    "cross_attn.q": "attn2.to_q",
+    "cross_attn.k": "attn2.to_k",
+    "cross_attn.v": "attn2.to_v",
+    "cross_attn.o": "attn2.to_out.0",
+    "cross_attn.norm_q": "attn2.norm_q",
+    "cross_attn.norm_k": "attn2.norm_k",
+    "attn2.to_k_img": "attn2.add_k_proj",
+    "attn2.to_v_img": "attn2.add_v_proj",
+    "attn2.norm_k_img": "attn2.norm_added_k",
+    # MXFP4 msmodelslim wraps Linear layers with a `.linear.` subpath;
+    # strip it so keys match the SGLang model parameters.
+    ".linear.": ".",
+    # NonFusionSmoothQuantWrapper exports smooth quant scale as `.div.mul_scale`;
+    # strip `.div.` so it loads as a direct parameter `mul_scale` on the linear layer.
+    ".div.": ".",
+}
+
+SUPPORTED_MODEL_TYPES = ["Wan2.2-T2V-A14B", "Wan2.2-I2V-A14B", "Wan2.2-TI2V-5B"]
+
+# Cascade models have two transformers (high_noise + low_noise)
+CASCADE_MODEL_TYPES = {"Wan2.2-T2V-A14B", "Wan2.2-I2V-A14B"}
+
+
+def get_transformer_config(model_type: str) -> Dict[str, Any]:
+    if model_type in SUPPORTED_MODEL_TYPES:
+        return TRANSFORMER_KEYS_RENAME_DICT
+    else:
+        raise ValueError(
+            f"Unsupported model_type: {model_type}. Supported: {SUPPORTED_MODEL_TYPES}"
+        )
+
+
+def get_transformer_dirs(model_type: str) -> List[str]:
+    """Return the list of transformer directory names for a given model type."""
+    if model_type in CASCADE_MODEL_TYPES:
+        return ["transformer", "transformer_2"]
+    return ["transformer"]
+
+
+def get_quant_subpath(
+    model_type: str, quant_path: pathlib.Path, transformer_dir: str
+) -> pathlib.Path:
+    """Return the quant weights subdirectory for a given transformer."""
+    if model_type in CASCADE_MODEL_TYPES:
+        sub = (
+            "high_noise_model"
+            if transformer_dir == "transformer"
+            else "low_noise_model"
+        )
+        return quant_path / sub
+    return quant_path
+
+
+def update_dict_(d: Dict[str, Any], old_key: str, new_key: str) -> None:
+    d[new_key] = d.pop(old_key)
+
+
+def load_sharded_safetensors(directory: pathlib.Path, pattern: str) -> dict:
+    candidates = sorted(directory.glob(pattern))
+    if not candidates:
+        raise FileNotFoundError(f"No file matching '{pattern}' found in {directory}")
+
+    state_dict = {}
+    for f in candidates:
+        state_dict.update(load_file(f))
+    return state_dict
+
+
+def convert_transformer(
+    model_type: str, model_dir: pathlib.Path, output_dir: pathlib.Path
+) -> None:
+    """Convert a single quantized transformer directory into Diffusers format."""
+    model_path = pathlib.Path(model_dir)
+    out_path = pathlib.Path(output_dir)
+    out_path.mkdir(parents=True, exist_ok=True)
+    RENAME_DICT = get_transformer_config(model_type)
+
+    state_dict = load_sharded_safetensors(model_path, "quant_model_weight*.safetensors")
+
+    json_candidates = sorted(model_path.glob("quant_model_description*.json"))
+    if not json_candidates:
+        raise FileNotFoundError(
+            f"No quant_model_description*.json found in {model_path}"
+        )
+    with open(json_candidates[0]) as f:
+        quant_config = json.load(f)
+
+    for key in list(state_dict.keys()):
+        new_key = key[:]
+        for replace_key, rename_key in RENAME_DICT.items():
+            new_key = new_key.replace(replace_key, rename_key)
+        if new_key != key:
+            update_dict_(state_dict, key, new_key)
+            # The quant JSON only covers quantized layers, not all model keys
+            if key in quant_config:
+                update_dict_(quant_config, key, new_key)
+
+    save_file(state_dict, out_path / "diffusion_pytorch_model.safetensors")
+
+    with open(out_path / "quant_model_description.json", "w") as f:
+        json.dump(quant_config, f, indent=2)
+
+
+def repack(
+    model_type: str,
+    original_model_path: pathlib.Path,
+    quant_path: pathlib.Path,
+    output_path: pathlib.Path,
+) -> None:
+    """
+    Full one-step repack workflow:
+      1. Copy the original HF Diffusers model to output_path, excluding transformer dir(s).
+      2. For each transformer: convert quant weights and copy config.json from original.
+    """
+    transformer_dirs = get_transformer_dirs(model_type)
+
+    # Step 1: Copy original model, skipping transformer dirs (they will be replaced)
+    logger.debug(f"Step 1: Copying original model to {output_path}")
+    logger.debug(f"        (skipping: {transformer_dirs})")
+    shutil.copytree(
+        str(original_model_path),
+        str(output_path),
+        ignore=shutil.ignore_patterns(*transformer_dirs),
+    )
+
+    # Step 2+: Convert each transformer
+    for i, tdir in enumerate(transformer_dirs):
+        q_path = get_quant_subpath(model_type, quant_path, tdir)
+        out_tdir = output_path / tdir
+        logger.debug(
+            f"\nStep {i + 2}: Converting {tdir} (quant source: {q_path.name})..."
+        )
+        convert_transformer(model_type, q_path, out_tdir)
+
+        # Copy config.json from the original transformer dir
+        src_config = original_model_path / tdir / "config.json"
+        if src_config.is_file():
+            shutil.copy2(str(src_config), str(out_tdir / "config.json"))
+            logger.debug(f"  Copied config.json from original {tdir}/")
+
+    logger.info(f"\nDone! Repacked model saved to: {output_path}")
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        description="Repack msmodelslim quantized Wan2.2 weights into HF Diffusers format"
+    )
+    parser.add_argument(
+        "--model-type",
+        type=str,
+        required=True,
+        choices=SUPPORTED_MODEL_TYPES,
+        help="Model type to convert",
+    )
+    parser.add_argument(
+        "--original-model-path",
+        type=str,
+        required=True,
+        help="Path to the original HF Diffusers model (e.g., /weights/Wan2.2-TI2V-5B-Diffusers)",
+    )
+    parser.add_argument(
+        "--quant-path",
+        type=str,
+        required=True,
+        help="Path to msmodelslim quantized weights directory",
+    )
+    parser.add_argument(
+        "--output-path",
+        type=str,
+        required=True,
+        help="Output path for the repacked model (e.g., /weights/Wan2.2-TI2V-5B-Diffusers-MXFP8)",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = get_args()
+    repack(
+        model_type=args.model_type,
+        original_model_path=pathlib.Path(args.original_model_path),
+        quant_path=pathlib.Path(args.quant_path),
+        output_path=pathlib.Path(args.output_path),
+    )