[Quant] ModelOpt FP8 for Wan2.2 & HunyuanVideo-1.5 video-gen (#3305)

Signed-off-by: WeiQing Chen <david6666666@users.noreply.github.com>

examples/quantization/check_modelopt_fp8_export.py

@@ -0,0 +1,247 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Verify a ModelOpt FP8 diffusers checkpoint exported by
+quantize_wan2_2_modelopt_fp8.py / quantize_hunyuanvideo_15_modelopt_fp8.py.
+
+Three checks:
+  A. transformer/config.json has a sane quantization_config block.
+  B. transformer/*.safetensors contains FP8 (float8_e4m3fn) quantized tensors.
+  C. transformer disk size is materially smaller than a BF16 baseline.
+
+Example:
+    python examples/quantization/check_modelopt_fp8_export.py \\
+        --output ./hv15-480p-modelopt-fp8
+
+    # Optional: compare disk size against a local or HF BF16 baseline.
+    python examples/quantization/check_modelopt_fp8_export.py \\
+        --output ./hv15-480p-modelopt-fp8 \\
+        --baseline hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from collections import Counter
+from pathlib import Path
+
+SUPPORTED_ALGOS = ("FP8",)
+
+
+def _check_config(transformer_dir: Path) -> tuple[int, str | None]:
+    """Returns (status, quant_algo). status: 0 pass, 1 fail, 2 warn."""
+    cfg_path = transformer_dir / "config.json"
+    if not cfg_path.exists():
+        print(f"[FAIL] {cfg_path} missing.")
+        return 1, None
+
+    with cfg_path.open(encoding="utf-8") as f:
+        cfg = json.load(f)
+
+    qc = cfg.get("quantization_config")
+    if not isinstance(qc, dict):
+        print(f"[FAIL] No `quantization_config` block in {cfg_path}.")
+        return 1, None
+
+    print(f"[A] quantization_config from {cfg_path}:")
+    print(json.dumps(qc, indent=2))
+
+    quant_algo = qc.get("quant_algo")
+    issues = []
+    if qc.get("quant_method") != "modelopt":
+        issues.append(f"quant_method={qc.get('quant_method')!r} (expected 'modelopt')")
+    if quant_algo not in SUPPORTED_ALGOS:
+        issues.append(
+            f"quant_algo={quant_algo!r} (expected one of {SUPPORTED_ALGOS} — vllm-omni adapter may not auto-detect)"
+        )
+
+    if issues:
+        print("[A] WARN — config looks incomplete:")
+        for issue in issues:
+            print(f"    - {issue}")
+        return 2, quant_algo
+    print(f"[A] PASS — config looks correct (quant_algo={quant_algo}).")
+    return 0, quant_algo
+
+
+def _read_safetensors_header(path: Path) -> dict:
+    """Read the JSON header of a safetensors file. Bypass-safe — doesn't materialize tensors.
+
+    Returns {tensor_name: {'dtype': 'F8_E4M3', 'shape': [...], 'data_offsets': [...]}}.
+    Header dtype strings: F8_E4M3, F8_E5M2, BF16, F16, F32, F64, I8, I16, I32, I64, BOOL, U8, ...
+    """
+    import struct
+
+    with open(path, "rb") as f:
+        header_len = struct.unpack("<Q", f.read(8))[0]
+        header = json.loads(f.read(header_len))
+    header.pop("__metadata__", None)
+    return header
+
+
+def _classify_weight_scale_granularity(weight_scale_shapes: list[list[int]]) -> str:
+    """Infer per-tensor vs per-channel vs per-block from sample weight_scale shapes.
+
+    ModelOpt block-wise produces shapes like `[16, 1, 16, 1]` (broadcasting dims of 1
+    interleaved with block-count dims). We count "meaningful" dims — ones with size > 1 —
+    and classify: 0 meaningful dims = per-tensor (scalar), 1 = per-channel, 2+ = per-block.
+    """
+    if not weight_scale_shapes:
+        return "no weight_scale tensors found"
+
+    def meaningful_dims(shape: list[int]) -> int:
+        return sum(1 for d in shape if d > 1)
+
+    per_tensor = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 0)
+    per_channel = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 1)
+    per_block = sum(1 for s in weight_scale_shapes if meaningful_dims(s) >= 2)
+    total = len(weight_scale_shapes)
+    if per_tensor == total:
+        return "per-tensor (all scalar scales)"
+    if per_channel == total:
+        return "per-channel (1 meaningful dim)"
+    if per_block == total:
+        return "per-block (2+ meaningful dims — e.g. [M//bm, 1, N//bn, 1] for tiles)"
+    return f"mixed: per-tensor={per_tensor}, per-channel={per_channel}, per-block={per_block} of {total}"
+
+
+def _check_safetensors(transformer_dir: Path) -> int:
+    """Returns 0 on pass, 1 on fail. Reads on-disk dtype from the safetensors header."""
+    files = sorted(transformer_dir.glob("*.safetensors"))
+    if not files:
+        print(f"[FAIL] No *.safetensors in {transformer_dir}.")
+        return 1
+
+    header_dtype_counts: Counter[str] = Counter()
+    sample_quant_weight_keys: list[str] = []
+    sample_scale_keys: list[str] = []
+    weight_scale_shapes: list[list[int]] = []
+    sample_weight_scale_entries: list[tuple[str, list[int]]] = []
+
+    for f in files:
+        try:
+            header = _read_safetensors_header(f)
+        except Exception as exc:
+            print(f"[B] WARN — could not parse header of {f}: {exc}")
+            continue
+        for k, info in header.items():
+            dtype = info.get("dtype", "?")
+            header_dtype_counts[dtype] += 1
+            # FP8 stores weights as F8_E4M3 directly.
+            if dtype.startswith("F8") and k.endswith(".weight") and len(sample_quant_weight_keys) < 5:
+                sample_quant_weight_keys.append(k)
+            if k.endswith(("_scale", ".weight_scale", ".input_scale", "_scale_inv")) and len(sample_scale_keys) < 5:
+                sample_scale_keys.append(k)
+            if k.endswith(".weight_scale"):
+                weight_scale_shapes.append(info.get("shape", []))
+                if len(sample_weight_scale_entries) < 5:
+                    sample_weight_scale_entries.append((k, info.get("shape", [])))
+
+    print(f"\n[B] On-disk dtype counts across {len(files)} safetensors file(s) (from header, not get_tensor):")
+    for dtype, count in sorted(header_dtype_counts.items(), key=lambda kv: -kv[1]):
+        marker = "  <-- FP8" if dtype.startswith("F8") else ""
+        print(f"    {dtype:10s} {count:>6d}{marker}")
+
+    quant_count = sum(c for d, c in header_dtype_counts.items() if d.startswith("F8"))
+    if quant_count == 0:
+        print("[B] FAIL — no FP8 tensors on disk. Calibration likely did not actually quantize the weights.")
+        return 1
+
+    print(f"[B] PASS — {quant_count} FP8 tensors stored on disk.")
+    if sample_quant_weight_keys:
+        print(f"    sample quantized weight tensors: {sample_quant_weight_keys[:3]}")
+    if sample_scale_keys:
+        print(f"    sample scale tensors:            {sample_scale_keys[:3]}")
+    print("    (Note: torch's get_tensor() may return these as bf16 views on some versions —")
+    print("     irrelevant; vLLM's loader uses native FP8 ops.)")
+
+    # Weight-scale granularity — per-tensor (scalar) vs per-channel (1-D) vs per-block (N-D).
+    print(f"\n    weight_scale granularity: {_classify_weight_scale_granularity(weight_scale_shapes)}")
+    for key, shape in sample_weight_scale_entries[:3]:
+        print(f"      {key}: shape {shape}")
+    return 0
+
+
+def _disk_size_gib(p: Path) -> float:
+    return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) / (1024**3)
+
+
+def _check_size_vs_baseline(transformer_dir: Path, baseline: str | None) -> int:
+    """Returns 0 always (informational only)."""
+    quant_size = _disk_size_gib(transformer_dir)
+    print(f"\n[C] FP8 transformer disk size: {quant_size:.2f} GiB")
+
+    if baseline is None:
+        print("[C] SKIP — pass --baseline <path or HF id> to compare against BF16.")
+        return 0
+
+    baseline_path = Path(baseline)
+    if not baseline_path.exists():
+        # Try HF download.
+        try:
+            from huggingface_hub import snapshot_download
+        except ImportError:
+            print("[C] SKIP — huggingface_hub not installed and baseline not a local path.")
+            return 0
+        print(f"    Downloading baseline transformer from HF: {baseline}")
+        baseline_path = Path(snapshot_download(baseline, allow_patterns=["transformer/*"]))
+
+    bf16_dir = baseline_path / "transformer" if (baseline_path / "transformer").exists() else baseline_path
+    bf16_size = _disk_size_gib(bf16_dir)
+    if bf16_size == 0:
+        print(f"[C] WARN — baseline transformer dir empty: {bf16_dir}")
+        return 0
+
+    # Expected reduction for FP8: ~50%.
+    min_reduction = 30
+    reduction = (1 - quant_size / bf16_size) * 100
+    print(f"[C] BF16 baseline transformer disk size: {bf16_size:.2f} GiB ({bf16_dir})")
+    print(f"[C] Disk reduction: {reduction:.1f}%  (FP8 transformer is {quant_size / bf16_size:.0%} of BF16)")
+    if reduction < min_reduction:
+        print(
+            f"[C] WARN — FP8 should typically reduce disk by ~40-50%; <{min_reduction}% suggests partial quantization."
+        )
+    return 0
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--output", required=True, help="Path to the exported ModelOpt FP8 checkpoint root.")
+    p.add_argument(
+        "--baseline",
+        default=None,
+        help="Optional BF16 baseline (local diffusers dir or HF id) for disk-size comparison.",
+    )
+    args = p.parse_args()
+
+    out_root = Path(args.output).expanduser().resolve()
+    transformer_dir = out_root / "transformer"
+    if not transformer_dir.exists():
+        print(f"[FAIL] {transformer_dir} does not exist.")
+        sys.exit(1)
+
+    print(f"Checking: {out_root}\n")
+
+    fail = 0
+    config_status, _quant_algo = _check_config(transformer_dir)
+    fail |= config_status
+    fail |= _check_safetensors(transformer_dir)
+    _check_size_vs_baseline(transformer_dir, args.baseline)
+
+    print()
+    if fail == 0:
+        print("=" * 60)
+        print("ALL CHECKS PASSED — checkpoint looks ready for vllm-omni serving.")
+    elif fail == 1:
+        print("=" * 60)
+        print("FAILURES detected — calibration may need to be re-run.")
+        sys.exit(1)
+    else:
+        print("=" * 60)
+        print("WARNINGS only — checkpoint may serve but with caveats. See [A] above.")
+
+
+if __name__ == "__main__":
+    main()

examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py

@@ -0,0 +1,448 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quantize HunyuanVideo-1.5 (480p T2V) to a ModelOpt FP8 Hugging Face checkpoint.
+
+Calibrates the DiT transformer using a small video prompt set and exports a
+diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
+The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
+(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
+
+Layers kept full precision match the #2728 / #2795 pattern: modulation,
+AdaLayerNorm, entry/exit projections, embeddings, the token refiner path,
+and final proj_out. MHA quantizers are off by default; HV-1.5 self-attention
+empirically degrades under FP8 (see #2920 ablation).
+
+Example:
+    python examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py \\
+        --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
+        --output ./hv15-480p-modelopt-fp8 \\
+        --overwrite
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import re
+import shutil
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+from diffusers import DiffusionPipeline
+
+DEFAULT_PROMPTS = [
+    "A dog running across a field of golden wheat.",
+    "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
+    "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
+    "A crackling campfire at night under a starry sky, sparks rising into the dark.",
+    "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
+    "A close-up of a blooming rose covered in morning dew, soft natural light.",
+    "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
+    "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
+]
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model", required=True, help="Input HV-1.5 diffusers directory or HF id.")
+    p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
+    p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
+    p.add_argument("--height", type=int, default=480)
+    p.add_argument("--width", type=int, default=832)
+    p.add_argument(
+        "--num-frames",
+        type=int,
+        default=33,
+        help="Frames per calibration sample. 33 matches the typical short benchmark.",
+    )
+    p.add_argument("--guidance-scale", type=float, default=6.0)
+    p.add_argument(
+        "--calib-steps",
+        type=int,
+        default=10,
+        help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
+    )
+    p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument(
+        "--prompt",
+        action="append",
+        default=[],
+        help="Custom calibration prompt. Repeat to provide multiple.",
+    )
+    p.add_argument(
+        "--quantize-mha",
+        action="store_true",
+        help="Enable FP8 attention K/V/softmax quantizers. Off by default — empirically degrades HV-1.5 video output.",
+    )
+    p.add_argument(
+        "--weight-block-size",
+        type=str,
+        default=None,
+        help="Per-block weight quantization as 'M,N' (e.g. '128,128' for 128x128 tiles). "
+        "Default: per-tensor (one scale per linear). Block-wise typically gives tighter quality at "
+        "negligible memory cost. Static FP8 is exempt from upstream vLLM's online block-wise gate.",
+    )
+    p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
+    return p
+
+
+def _parse_block_size(spec: str | None) -> list[int] | None:
+    if spec is None:
+        return None
+    parts = [int(x) for x in spec.split(",") if x.strip()]
+    if len(parts) != 2:
+        raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
+    return parts
+
+
+def _require_modelopt() -> Any:
+    try:
+        import modelopt.torch.quantization as mtq
+    except ModuleNotFoundError as exc:
+        raise SystemExit(
+            "NVIDIA ModelOpt is not installed. Install with:\n"
+            "  pip install 'nvidia-modelopt[all]'\n"
+            f"Original error: {exc}"
+        ) from exc
+    return mtq
+
+
+def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
+    model_path = args.model
+    output_dir = Path(args.output).expanduser().resolve()
+    if output_dir.exists():
+        if not args.overwrite:
+            raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
+        shutil.rmtree(output_dir)
+    return model_path, output_dir
+
+
+def _select_dtype(name: str) -> torch.dtype:
+    return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
+
+
+def _build_prompts(args: argparse.Namespace) -> list[str]:
+    prompts = args.prompt or DEFAULT_PROMPTS
+    if args.calib_size <= 0:
+        raise SystemExit("--calib-size must be positive.")
+    if len(prompts) < args.calib_size:
+        repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
+        prompts = (prompts * repeats)[: args.calib_size]
+    return prompts[: args.calib_size]
+
+
+# Layers to KEEP at full precision (mirror of the #2920 wiring + #2728/#2795 skip pattern).
+# - x_embedder, image_embedder, context_embedder*, time_embed*, cond_type_embed: entry/embedding
+# - norm_out, norm1*.linear, norm1_context*.linear, norm2*, norm2_context*: AdaLayerNorm modulation
+# - proj_out: final output projection
+# - token_refiner*: text-encoder refinement uses diffusers raw nn.Linear
+def _filter_func_hv15(name: str) -> bool:
+    pattern = re.compile(
+        r"(proj_out.*|"
+        r".*(x_embedder|image_embedder|context_embedder|context_embedder_2|"
+        r"time_embed|cond_type_embed|"
+        r"norm_out|norm1\.linear|norm1_context\.linear|norm2|norm2_context|"
+        r"token_refiner).*)"
+    )
+    return pattern.match(name) is not None
+
+
+def _mha_filter_func(name: str) -> bool:
+    pattern = re.compile(
+        r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
+    )
+    return pattern.match(name) is not None
+
+
+def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
+    if not hasattr(mtq, "disable_quantizer"):
+        return
+    mtq.disable_quantizer(backbone, _filter_func_hv15)
+    if not quantize_mha:
+        mtq.disable_quantizer(backbone, _mha_filter_func)
+
+
+def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
+    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
+    if hasattr(pipe, "set_progress_bar_config"):
+        pipe.set_progress_bar_config(disable=True)
+    pipe.to("cuda")
+    return pipe
+
+
+def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
+    generator = torch.Generator(device="cuda")
+
+    # Try to set guidance on the pipeline's guider object up front (modern
+    # diffusers HV-1.5 uses a Guider abstraction, not a per-call kwarg). Falls
+    # back silently — calibration uses whatever default the pipeline ships with.
+    guider = getattr(pipe, "guider", None)
+    if guider is not None and hasattr(guider, "guidance_scale"):
+        try:
+            guider.guidance_scale = args.guidance_scale
+        except Exception:
+            pass
+
+    base_kwargs = dict(
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        num_inference_steps=args.calib_steps,
+        output_type="latent",
+    )
+
+    def forward_loop(*_unused_args, **_unused_kwargs) -> None:
+        with torch.inference_mode():
+            for idx, prompt in enumerate(prompts):
+                generator.manual_seed(args.seed + idx)
+                # Try with guidance_scale first; fall back without on TypeError
+                # for pipelines (like HV-1.5) that take CFG via guider config.
+                try:
+                    pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
+                except TypeError as exc:
+                    if "guidance_scale" not in str(exc):
+                        raise
+                    pipe(prompt=prompt, generator=generator, **base_kwargs)
+
+    return forward_loop
+
+
+def _summarize_export(output_dir: Path) -> None:
+    cfg_path = output_dir / "transformer" / "config.json"
+    if not cfg_path.exists():
+        print(f"[warn] {cfg_path} missing.", file=sys.stderr)
+        return
+    with cfg_path.open(encoding="utf-8") as f:
+        cfg = json.load(f)
+    qc = cfg.get("quantization_config")
+    if not isinstance(qc, dict):
+        print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
+        return
+    print("Export summary:")
+    print(f"  quant_method: {qc.get('quant_method')}")
+    print(f"  quant_algo:   {qc.get('quant_algo')}")
+    producer = qc.get("producer")
+    if isinstance(producer, dict):
+        print(f"  producer:     {producer.get('name')} {producer.get('version')}")
+    print(f"  config path:  {cfg_path}")
+
+
+def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
+    """Convert in-memory weights of quantized modules to actual FP8 storage.
+
+    `export_hf_checkpoint` skips this step for unknown model types (HV-1.5 isn't
+    in ModelOpt's recognized-model registry), so we must call the per-weight
+    export helper ourselves. Same workaround as the HunyuanImage-3 calibration
+    helper.
+    """
+    from modelopt.torch.export.quant_utils import (
+        QUANTIZATION_NONE,
+        get_quantization_format,
+        quantizer_attr_names,
+        weight_attr_names,
+    )
+    from modelopt.torch.export.unified_export_hf import _export_quantized_weight
+
+    exported = 0
+    for name, module in backbone.named_modules():
+        try:
+            quantization_format = get_quantization_format(module)
+        except Exception as exc:
+            print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
+            continue
+        if quantization_format == QUANTIZATION_NONE:
+            continue
+        for weight_name in weight_attr_names(module):
+            quantizer_attrs = quantizer_attr_names(weight_name)
+            weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
+            if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
+                continue
+            _export_quantized_weight(module, dtype, weight_name)
+            exported += 1
+    return exported
+
+
+def _hv15_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
+    """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913).
+
+    Same shape as the HunyuanImage-3 author's _hunyuan_quant_config(). When
+    `weight_block_size` is given, advertise block-wise weight quantization in
+    the saved metadata (so consumers know to expect multi-element scale tensors).
+    """
+    weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
+    if weight_block_size is not None:
+        weights_cfg["strategy"] = "block"
+        weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
+    return {
+        "config_groups": {
+            "group_0": {
+                "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
+                "weights": weights_cfg,
+                "targets": ["Linear"],
+            }
+        },
+        "ignore": [
+            "context_embedder*",
+            "context_embedder_2*",
+            "cond_type_embed*",
+            "image_embedder*",
+            "norm1.linear*",
+            "norm1_context.linear*",
+            "norm2*",
+            "norm2_context*",
+            "norm_out*",
+            "proj_out*",
+            "time_embed*",
+            "token_refiner*",
+            "x_embedder*",
+        ],
+        "producer": {"name": "modelopt"},
+        "quant_algo": "FP8",
+        "quant_method": "modelopt",
+    }
+
+
+def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
+    """Inject quant_algo: FP8 + config_groups into transformer/config.json so
+    vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8."""
+    cfg_path = output_dir / "transformer" / "config.json"
+    with cfg_path.open(encoding="utf-8") as f:
+        cfg = json.load(f)
+
+    new_qc = _hv15_quant_config_block(weight_block_size=weight_block_size)
+    existing = cfg.get("quantization_config")
+    if isinstance(existing, dict):
+        producer = existing.get("producer")
+        if isinstance(producer, dict):
+            new_qc["producer"] = producer
+
+    cfg["quantization_config"] = new_qc
+    with cfg_path.open("w", encoding="utf-8") as f:
+        json.dump(cfg, f, indent=2)
+
+
+def _save_pipeline_with_fp8_transformer(
+    pipe: DiffusionPipeline,
+    model_path: str,
+    output_dir: Path,
+    max_shard_size: str = "5GB",
+) -> None:
+    """Save the pipeline with the (now FP8) transformer.
+
+    Copies the source directory verbatim except for `transformer/`, then
+    saves the transformer with quantizers hidden so the state dict contains
+    only the FP8 weights + scale tensors.
+    """
+    from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
+
+    src = Path(model_path)
+    if not src.exists():
+        from huggingface_hub import snapshot_download
+
+        src = Path(snapshot_download(model_path))
+
+    if output_dir.exists():
+        shutil.rmtree(output_dir)
+    shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer"))
+
+    transformer_out = output_dir / "transformer"
+    # `hide_quantizers_from_state_dict` walks named_modules(); pass the actual
+    # nn.Module (transformer), not the diffusers Pipeline wrapper.
+    with hide_quantizers_from_state_dict(pipe.transformer):
+        pipe.transformer.save_pretrained(
+            str(transformer_out),
+            safe_serialization=True,
+            max_shard_size=max_shard_size,
+        )
+
+
+def main() -> None:
+    args = _build_parser().parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
+
+    mtq = _require_modelopt()
+    model_path, output_dir = _ensure_paths(args)
+    dtype = _select_dtype(args.dtype)
+    prompts = _build_prompts(args)
+
+    print("Quantization plan:")
+    weight_block_size = _parse_block_size(args.weight_block_size)
+
+    print(f"  input:           {args.model}")
+    print(f"  output:          {output_dir}")
+    print(f"  dtype:           {dtype}")
+    print(f"  height/width:    {args.height}x{args.width}")
+    print(f"  num_frames:      {args.num_frames}")
+    print(f"  calib_size:      {len(prompts)}")
+    print(f"  calib_steps:     {args.calib_steps}")
+    print(f"  quantize_mha:    {args.quantize_mha}")
+    print(
+        f"  weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
+    )
+
+    pipe = _load_pipeline(model_path, dtype)
+    backbone = pipe.transformer
+
+    quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
+    if weight_block_size is not None:
+        # Switch from per-tensor (default) to block-wise weight quantization.
+        # ModelOpt's wildcard "*weight_quantizer" matches every linear's weight quantizer.
+        quant_config["quant_cfg"]["*weight_quantizer"] = {
+            "num_bits": (4, 3),  # E4M3 (FP8 weights, same as default)
+            "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
+        }
+        print(
+            f"  -> overriding weight quantizer with block_sizes={weight_block_size} "
+            f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
+        )
+
+    forward_loop = _build_forward_loop(pipe, args, prompts)
+    quantized = mtq.quantize(backbone, quant_config, forward_loop)
+    if quantized is not None:
+        pipe.transformer = quantized
+        backbone = quantized
+
+    _disable_known_problematic_quantizers(mtq, backbone, quantize_mha=args.quantize_mha)
+
+    print("\nForcing FP8 weight serialization (HV-1.5 isn't in ModelOpt's recognized-model registry,")
+    print("so we have to call the per-weight export helper ourselves)...")
+    exported = _force_export_quantized_weights(backbone, dtype)
+    print(f"  -> {exported} weights converted to FP8 in memory")
+    if exported == 0:
+        raise SystemExit(
+            "No quantized weights were exported. Calibration may have skipped every layer "
+            "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
+            "weight quantizers."
+        )
+
+    print("\nSaving pipeline with FP8 transformer...")
+    _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
+    _patch_quant_config(output_dir, weight_block_size=weight_block_size)
+    print(f"Saved to: {output_dir}")
+    _summarize_export(output_dir)
+
+    print("\nNext: validate the checkpoint with vllm-omni:")
+    print(
+        "  python examples/offline_inference/text_to_video/text_to_video.py \\\n"
+        f"    --model {output_dir} \\\n"
+        "    --quantization fp8 \\\n"
+        "    --prompt 'A dog running across a field of golden wheat.' \\\n"
+        f"    --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
+        "    --num-inference-steps 30 --guidance-scale 6.0 --seed 42 \\\n"
+        "    --output outputs/hv15_modelopt_fp8.mp4 \\\n"
+        "    --enforce-eager"
+    )
+    print(
+        "\n  (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
+        "checkpoint's config.json has modelopt metadata.)"
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/quantization/quantize_wan2_2_modelopt_fp8.py

@@ -0,0 +1,471 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quantize Wan2.2 (TI2V-5B, 704x1280 T2V) to a ModelOpt FP8 Hugging Face checkpoint.
+
+Calibrates the DiT transformer using a small video prompt set and exports a
+diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
+The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
+(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
+
+Layers kept full precision match the #2728 / #2795 pattern: condition embedder
+(time/text/image), patch embedding, modulation (scale_shift_table), final
+norm + proj_out, and sequence-parallel helpers. All attention + FFN linears
+are quantized — static calibration handles the numerics that online FP8
+couldn't (see #2920 ablation).
+
+Default target is `Wan-AI/Wan2.2-TI2V-5B-Diffusers`, the dense 5B variant that
+fits 80GB BF16. The A14B MoE variants need 2+ GPUs and are out of scope here.
+
+Example:
+    python examples/quantization/quantize_wan2_2_modelopt_fp8.py \\
+        --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\
+        --output ./wan22-ti2v-modelopt-fp8 \\
+        --overwrite
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import re
+import shutil
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+from diffusers import DiffusionPipeline
+
+DEFAULT_PROMPTS = [
+    "A dog running across a field of golden wheat.",
+    "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
+    "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
+    "A crackling campfire at night under a starry sky, sparks rising into the dark.",
+    "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
+    "A close-up of a blooming rose covered in morning dew, soft natural light.",
+    "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
+    "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
+]
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model", required=True, help="Input Wan2.2 diffusers directory or HF id.")
+    p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
+    p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
+    p.add_argument("--height", type=int, default=704, help="Calibration video height (Wan2.2 TI2V-5B native: 704).")
+    p.add_argument("--width", type=int, default=1280, help="Calibration video width (Wan2.2 TI2V-5B native: 1280).")
+    p.add_argument(
+        "--num-frames",
+        type=int,
+        default=49,
+        help="Frames per calibration sample. 49 matches the typical short benchmark; "
+        "use 17 to reduce memory pressure during calibration.",
+    )
+    p.add_argument("--guidance-scale", type=float, default=5.0)
+    p.add_argument(
+        "--calib-steps",
+        type=int,
+        default=10,
+        help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
+    )
+    p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument(
+        "--prompt",
+        action="append",
+        default=[],
+        help="Custom calibration prompt. Repeat to provide multiple.",
+    )
+    p.add_argument(
+        "--quantize-mha",
+        action="store_true",
+        help="Enable FP8 attention K/V/softmax quantizers. Off by default — Wan2.2's long attention "
+        "sequences amplified FP8 drift in the online ablation (see #2920).",
+    )
+    p.add_argument(
+        "--weight-block-size",
+        type=str,
+        default=None,
+        help="Per-block weight quantization as 'M,N' (e.g. '128,128'). Default per-tensor. "
+        "Note: vllm-omni's ModelOpt adapter may not yet dispatch block-wise scales — check #2924 "
+        "for the HV-1.5 investigation status before relying on this.",
+    )
+    p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
+    return p
+
+
+def _parse_block_size(spec: str | None) -> list[int] | None:
+    if spec is None:
+        return None
+    parts = [int(x) for x in spec.split(",") if x.strip()]
+    if len(parts) != 2:
+        raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
+    return parts
+
+
+def _require_modelopt() -> Any:
+    try:
+        import modelopt.torch.quantization as mtq
+    except ModuleNotFoundError as exc:
+        raise SystemExit(
+            "NVIDIA ModelOpt is not installed. Install with:\n"
+            "  pip install 'nvidia-modelopt[all]'\n"
+            f"Original error: {exc}"
+        ) from exc
+    return mtq
+
+
+def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
+    model_path = args.model
+    output_dir = Path(args.output).expanduser().resolve()
+    if output_dir.exists():
+        if not args.overwrite:
+            raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
+        shutil.rmtree(output_dir)
+    return model_path, output_dir
+
+
+def _select_dtype(name: str) -> torch.dtype:
+    return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
+
+
+def _build_prompts(args: argparse.Namespace) -> list[str]:
+    prompts = args.prompt or DEFAULT_PROMPTS
+    if args.calib_size <= 0:
+        raise SystemExit("--calib-size must be positive.")
+    if len(prompts) < args.calib_size:
+        repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
+        prompts = (prompts * repeats)[: args.calib_size]
+    return prompts[: args.calib_size]
+
+
+# Layers to KEEP at full precision. Wan2.2's module naming:
+# - condition_embedder: time_embedder, time_proj, text_embedder, image_embedder (I2V)
+# - patch_embedding: Conv3dLayer (already not Linear, belt-and-suspenders skip)
+# - scale_shift_table: nn.Parameter modulation (not Linear, but pattern guard)
+# - norm_out: AdaLayerNorm final
+# - proj_out: final nn.Linear
+# - timestep_proj_prepare / output_scale_shift_prepare: SP helpers
+def _filter_func_wan22(name: str) -> bool:
+    pattern = re.compile(
+        r"(proj_out.*|"
+        r".*(condition_embedder|patch_embedding|"
+        r"norm_out|scale_shift_table|"
+        r"timestep_proj_prepare|output_scale_shift_prepare).*)"
+    )
+    return pattern.match(name) is not None
+
+
+def _mha_filter_func(name: str) -> bool:
+    pattern = re.compile(
+        r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
+    )
+    return pattern.match(name) is not None
+
+
+def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
+    if not hasattr(mtq, "disable_quantizer"):
+        return
+    mtq.disable_quantizer(backbone, _filter_func_wan22)
+    if not quantize_mha:
+        mtq.disable_quantizer(backbone, _mha_filter_func)
+
+
+def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
+    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
+    if hasattr(pipe, "set_progress_bar_config"):
+        pipe.set_progress_bar_config(disable=True)
+    pipe.to("cuda")
+    return pipe
+
+
+def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
+    generator = torch.Generator(device="cuda")
+
+    # Try setting guidance on the pipeline's guider if present (newer diffusers APIs).
+    guider = getattr(pipe, "guider", None)
+    if guider is not None and hasattr(guider, "guidance_scale"):
+        try:
+            guider.guidance_scale = args.guidance_scale
+        except Exception:
+            pass
+
+    base_kwargs = dict(
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        num_inference_steps=args.calib_steps,
+        output_type="latent",
+    )
+
+    def forward_loop(*_unused_args, **_unused_kwargs) -> None:
+        with torch.inference_mode():
+            for idx, prompt in enumerate(prompts):
+                generator.manual_seed(args.seed + idx)
+                # Try with guidance_scale first; fall back without on TypeError
+                # for pipelines that take CFG via guider config only.
+                try:
+                    pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
+                except TypeError as exc:
+                    if "guidance_scale" not in str(exc):
+                        raise
+                    pipe(prompt=prompt, generator=generator, **base_kwargs)
+
+    return forward_loop
+
+
+def _summarize_export(output_dir: Path) -> None:
+    cfg_path = output_dir / "transformer" / "config.json"
+    if not cfg_path.exists():
+        print(f"[warn] {cfg_path} missing.", file=sys.stderr)
+        return
+    with cfg_path.open(encoding="utf-8") as f:
+        cfg = json.load(f)
+    qc = cfg.get("quantization_config")
+    if not isinstance(qc, dict):
+        print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
+        return
+    print("Export summary:")
+    print(f"  quant_method: {qc.get('quant_method')}")
+    print(f"  quant_algo:   {qc.get('quant_algo')}")
+    producer = qc.get("producer")
+    if isinstance(producer, dict):
+        print(f"  producer:     {producer.get('name')} {producer.get('version')}")
+    print(f"  config path:  {cfg_path}")
+
+
+def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
+    """Convert in-memory weights of quantized modules to actual FP8 storage.
+
+    `export_hf_checkpoint` skips this step for unknown model types (Wan2.2 isn't
+    in ModelOpt's recognized-model registry), so we must call the per-weight
+    export helper ourselves. Same workaround as the HunyuanVideo-1.5 / HunyuanImage-3
+    calibration helpers.
+    """
+    from modelopt.torch.export.quant_utils import (
+        QUANTIZATION_NONE,
+        get_quantization_format,
+        quantizer_attr_names,
+        weight_attr_names,
+    )
+    from modelopt.torch.export.unified_export_hf import _export_quantized_weight
+
+    exported = 0
+    for name, module in backbone.named_modules():
+        try:
+            quantization_format = get_quantization_format(module)
+        except Exception as exc:
+            print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
+            continue
+        if quantization_format == QUANTIZATION_NONE:
+            continue
+        for weight_name in weight_attr_names(module):
+            quantizer_attrs = quantizer_attr_names(weight_name)
+            weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
+            if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
+                continue
+            _export_quantized_weight(module, dtype, weight_name)
+            exported += 1
+    return exported
+
+
+def _wan22_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
+    """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913)."""
+    weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
+    if weight_block_size is not None:
+        weights_cfg["strategy"] = "block"
+        weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
+    return {
+        "config_groups": {
+            "group_0": {
+                "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
+                "weights": weights_cfg,
+                "targets": ["Linear"],
+            }
+        },
+        "ignore": [
+            "condition_embedder*",
+            "norm_out*",
+            "output_scale_shift_prepare*",
+            "patch_embedding*",
+            "proj_out*",
+            "scale_shift_table*",
+            "timestep_proj_prepare*",
+        ],
+        "producer": {"name": "modelopt"},
+        "quant_algo": "FP8",
+        "quant_method": "modelopt",
+    }
+
+
+def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
+    """Inject quant_algo: FP8 + config_groups into each transformer's config.json
+    so vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8.
+
+    Patches both ``transformer`` and ``transformer_2`` (A14B) when present."""
+    for sub in ("transformer", "transformer_2"):
+        cfg_path = output_dir / sub / "config.json"
+        if not cfg_path.exists():
+            continue
+        with cfg_path.open(encoding="utf-8") as f:
+            cfg = json.load(f)
+
+        new_qc = _wan22_quant_config_block(weight_block_size=weight_block_size)
+        existing = cfg.get("quantization_config")
+        if isinstance(existing, dict):
+            producer = existing.get("producer")
+            if isinstance(producer, dict):
+                new_qc["producer"] = producer
+
+        cfg["quantization_config"] = new_qc
+        with cfg_path.open("w", encoding="utf-8") as f:
+            json.dump(cfg, f, indent=2)
+
+
+def _list_transformers(pipe: DiffusionPipeline) -> list[tuple[str, torch.nn.Module]]:
+    """Return [(attr_name, module)] for every DiT on the pipeline.
+
+    Wan2.2 A14B has both ``transformer`` (low-noise) and ``transformer_2``
+    (high-noise). TI2V-5B has only ``transformer``.
+    """
+    out = []
+    for attr in ("transformer", "transformer_2"):
+        mod = getattr(pipe, attr, None)
+        if mod is not None:
+            out.append((attr, mod))
+    return out
+
+
+def _save_pipeline_with_fp8_transformer(
+    pipe: DiffusionPipeline,
+    model_path: str,
+    output_dir: Path,
+    max_shard_size: str = "5GB",
+) -> None:
+    """Copy source dir verbatim minus transformer dirs, then save each quantized DiT.
+
+    Handles A14B's dual transformer (``transformer`` + ``transformer_2``) as well
+    as the single-DiT TI2V-5B.
+    """
+    from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
+
+    src = Path(model_path)
+    if not src.exists():
+        from huggingface_hub import snapshot_download
+
+        src = Path(snapshot_download(model_path))
+
+    if output_dir.exists():
+        shutil.rmtree(output_dir)
+    shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer", "transformer_2"))
+
+    for attr, backbone in _list_transformers(pipe):
+        transformer_out = output_dir / attr
+        # Pass the nn.Module (transformer), not the Pipeline wrapper.
+        with hide_quantizers_from_state_dict(backbone):
+            backbone.save_pretrained(
+                str(transformer_out),
+                safe_serialization=True,
+                max_shard_size=max_shard_size,
+            )
+
+
+def main() -> None:
+    args = _build_parser().parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
+
+    mtq = _require_modelopt()
+    model_path, output_dir = _ensure_paths(args)
+    dtype = _select_dtype(args.dtype)
+    prompts = _build_prompts(args)
+    weight_block_size = _parse_block_size(args.weight_block_size)
+
+    print("Quantization plan:")
+    print(f"  input:           {args.model}")
+    print(f"  output:          {output_dir}")
+    print(f"  dtype:           {dtype}")
+    print(f"  height/width:    {args.height}x{args.width}")
+    print(f"  num_frames:      {args.num_frames}")
+    print(f"  calib_size:      {len(prompts)}")
+    print(f"  calib_steps:     {args.calib_steps}")
+    print(f"  quantize_mha:    {args.quantize_mha}")
+    print(
+        f"  weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
+    )
+
+    pipe = _load_pipeline(model_path, dtype)
+    transformers = _list_transformers(pipe)
+    if not transformers:
+        raise SystemExit("Pipeline has no transformer or transformer_2 attribute.")
+    print(f"  found {len(transformers)} transformer(s): {', '.join(a for a, _ in transformers)}")
+
+    quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
+    if weight_block_size is not None:
+        quant_config["quant_cfg"]["*weight_quantizer"] = {
+            "num_bits": (4, 3),
+            "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
+        }
+        print(
+            f"  -> overriding weight quantizer with block_sizes={weight_block_size} "
+            f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
+        )
+
+    forward_loop = _build_forward_loop(pipe, args, prompts)
+
+    # Quantize each DiT. The forward_loop runs the full pipeline, so quantizers
+    # on every attached transformer (A14B: transformer + transformer_2) get
+    # calibrated during the same passes.
+    for attr, backbone in transformers:
+        print(f"\nQuantizing {attr}...")
+        quantized = mtq.quantize(backbone, quant_config, forward_loop)
+        if quantized is not None:
+            setattr(pipe, attr, quantized)
+        _disable_known_problematic_quantizers(mtq, getattr(pipe, attr), quantize_mha=args.quantize_mha)
+
+    # Export weights to FP8 only AFTER all transformers are calibrated. Exporting
+    # inside the loop would convert transformer's weights to real FP8, and the
+    # next transformer's forward_loop (full pipeline) would then fake-quantize an
+    # already-FP8 weight -> "fake_e4m3fy not implemented for Float8_e4m3fn".
+    print("\nForcing FP8 weight serialization (Wan2.2 isn't in ModelOpt's recognized-model registry,")
+    print("so we have to call the per-weight export helper ourselves)...")
+    total_exported = 0
+    for attr, backbone in _list_transformers(pipe):
+        exported = _force_export_quantized_weights(backbone, dtype)
+        print(f"  -> {exported} weights converted to FP8 in {attr}")
+        total_exported += exported
+
+    if total_exported == 0:
+        raise SystemExit(
+            "No quantized weights were exported. Calibration may have skipped every layer "
+            "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
+            "weight quantizers."
+        )
+
+    print("\nSaving pipeline with FP8 transformer(s)...")
+    _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
+    _patch_quant_config(output_dir, weight_block_size=weight_block_size)
+    print(f"Saved to: {output_dir}")
+    _summarize_export(output_dir)
+
+    print("\nNext: validate the checkpoint with vllm-omni:")
+    print(
+        "  python examples/offline_inference/text_to_video/text_to_video.py \\\n"
+        f"    --model {output_dir} \\\n"
+        "    --quantization fp8 \\\n"
+        "    --prompt 'A dog running across a field of golden wheat.' \\\n"
+        f"    --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
+        "    --num-inference-steps 30 --guidance-scale 5.0 --seed 42 \\\n"
+        "    --output outputs/wan22_modelopt_fp8.mp4"
+    )
+    print(
+        "\n  (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
+        "checkpoint's config.json has modelopt metadata.)"
+    )
+
+
+if __name__ == "__main__":
+    main()

tests/e2e/accuracy/test_videogen_modelopt_quant.py

@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end ModelOpt FP8 quality accuracy for video-gen DiTs.
+
+Modeled on ``test_hunyuan_image3.py``'s ``test_quantized_dit_matches_bf16_accuracy``:
+for each calibrated FP8 checkpoint the test generates one video with BF16 and one
+with the FP8 DiT under the same seed, scores prompt-faithfulness with CLIP on the
+middle frame, then asserts both ``CLIP >= absolute_floor`` and
+``CLIP_drop <= drop_threshold``. The test is gated by an opt-in env var (model
+paths are local) so it never runs in generic CI.
+
+To run::
+
+    VIDEOGEN_RUN_QUANT_ACCURACY=1 \\
+    WAN22_A14B_BF16_MODEL=Wan-AI/Wan2.2-T2V-A14B-Diffusers \\
+    WAN22_A14B_FP8_MODEL=/path/to/wan22-a14b-modelopt-fp8 \\
+    HV15_BF16_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
+    HV15_FP8_MODEL=/path/to/hv15-modelopt-fp8 \\
+        pytest -s -v tests/e2e/accuracy/test_videogen_modelopt_quant.py
+"""
+
+from __future__ import annotations
+
+import gc
+import json
+import os
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from tests.e2e.accuracy.helpers import CLIPScorer, model_output_dir
+
+pytestmark = [pytest.mark.local_model, pytest.mark.diffusion]
+
+# ---------------------------------------------------------------------------
+# Tunables (override via env to suit local hardware / calibration budget)
+# ---------------------------------------------------------------------------
+SEED = 42
+HEIGHT = int(os.environ.get("VIDEOGEN_QUANT_HEIGHT", "480"))
+WIDTH = int(os.environ.get("VIDEOGEN_QUANT_WIDTH", "832"))
+NUM_FRAMES = int(os.environ.get("VIDEOGEN_QUANT_NUM_FRAMES", "17"))
+NUM_INFERENCE_STEPS = int(os.environ.get("VIDEOGEN_QUANT_STEPS", "20"))
+
+WAN_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic."
+HV_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic."
+
+# CLIP gates: absolute floor + max drop from the BF16 reference.
+CLIP_ABSOLUTE_FLOOR = float(os.environ.get("VIDEOGEN_QUANT_CLIP_FLOOR", "20.0"))
+CLIP_DROP_FP8 = float(os.environ.get("VIDEOGEN_QUANT_CLIP_DROP_FP8", "7.0"))
+
+RUN_ENV = "VIDEOGEN_RUN_QUANT_ACCURACY"
+_TRUE = {"1", "true", "yes", "on"}
+
+
+# ---------------------------------------------------------------------------
+@dataclass(frozen=True)
+class _Spec:
+    """A model under test (BF16 + its FP8 variant)."""
+
+    family: str
+    bf16_env: str
+    fp8_env: str
+    prompt: str
+    guidance_scale: float
+    guidance_scale_2: float | None = None
+    boundary_ratio: float | None = None
+
+
+WAN_SPEC = _Spec(
+    family="wan22_a14b_t2v",
+    bf16_env="WAN22_A14B_BF16_MODEL",
+    fp8_env="WAN22_A14B_FP8_MODEL",
+    prompt=WAN_PROMPT,
+    guidance_scale=3.5,
+    guidance_scale_2=4.0,
+    boundary_ratio=0.875,
+)
+
+HV_SPEC = _Spec(
+    family="hunyuanvideo15_t2v",
+    bf16_env="HV15_BF16_MODEL",
+    fp8_env="HV15_FP8_MODEL",
+    prompt=HV_PROMPT,
+    guidance_scale=6.0,
+)
+
+
+@dataclass(frozen=True)
+class _QuantCase:
+    name: str
+    spec: _Spec
+    clip_drop: float
+
+
+def _cases() -> list[pytest.ParameterSet]:
+    cases = [
+        _QuantCase("wan_fp8", WAN_SPEC, CLIP_DROP_FP8),
+        _QuantCase("hv_fp8", HV_SPEC, CLIP_DROP_FP8),
+    ]
+    params: list[pytest.ParameterSet] = []
+    opted_in = os.environ.get(RUN_ENV, "").lower() in _TRUE
+    for case in cases:
+        marks: list[Any] = []
+        if not opted_in:
+            marks.append(pytest.mark.skip(reason=f"Set {RUN_ENV}=1 to run video-gen quant accuracy."))
+        if not os.environ.get(case.spec.bf16_env):
+            marks.append(pytest.mark.skip(reason=f"Set {case.spec.bf16_env} to run {case.name}."))
+        if not os.environ.get(case.spec.fp8_env):
+            marks.append(pytest.mark.skip(reason=f"Set {case.spec.fp8_env} to run {case.name}."))
+        params.append(pytest.param(case, id=case.name, marks=marks))
+    return params
+
+
+# ---------------------------------------------------------------------------
+def _middle_frame(frames: list[Image.Image] | list[np.ndarray] | np.ndarray) -> Image.Image:
+    if isinstance(frames, np.ndarray):
+        idx = frames.shape[0] // 2
+        arr = frames[idx]
+        if arr.dtype != np.uint8:
+            arr = (arr * 255).clip(0, 255).astype(np.uint8)
+        return Image.fromarray(arr).convert("RGB")
+    if not frames:
+        raise AssertionError("Empty video output")
+    mid = frames[len(frames) // 2]
+    if isinstance(mid, np.ndarray):
+        if mid.dtype != np.uint8:
+            mid = (mid * 255).clip(0, 255).astype(np.uint8)
+        return Image.fromarray(mid).convert("RGB")
+    return mid.convert("RGB")
+
+
+def _extract_frames(result) -> list[Image.Image]:
+    """Pull a frame list out of an Omni text-to-video result."""
+    from vllm_omni.entrypoints.omni import OmniRequestOutput
+
+    if isinstance(result, list):
+        result = result[0] if result else None
+    if isinstance(result, OmniRequestOutput) and result.is_pipeline_output and result.request_output is not None:
+        result = result.request_output
+    images = getattr(result, "images", None)
+    if not images:
+        raise AssertionError("Pipeline output had no images")
+    if isinstance(images, list) and images and isinstance(images[0], tuple):
+        # (frames, fps)
+        images = images[0][0]
+    return list(images)
+
+
+def _generate_video(
+    *,
+    model: str,
+    spec: _Spec,
+    quantization: str | None,
+    output_path: Path,
+) -> tuple[Image.Image, float]:
+    """Generate one video and return its middle frame + elapsed wall time."""
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+    from vllm_omni.entrypoints.omni import Omni
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.platforms import current_omni_platform
+
+    kwargs: dict[str, Any] = dict(
+        model=model,
+        parallel_config=DiffusionParallelConfig(),
+        enforce_eager=True,
+        vae_use_tiling=True,
+    )
+    if quantization:
+        kwargs["quantization"] = quantization
+    if spec.boundary_ratio is not None:
+        kwargs["boundary_ratio"] = spec.boundary_ratio
+    omni = Omni(**kwargs)
+
+    try:
+        gen = torch.Generator(device=current_omni_platform.device_type).manual_seed(SEED)
+        sampling = dict(
+            height=HEIGHT,
+            width=WIDTH,
+            generator=gen,
+            guidance_scale=spec.guidance_scale,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            num_frames=NUM_FRAMES,
+            seed=SEED,
+        )
+        if spec.guidance_scale_2 is not None:
+            sampling["guidance_scale_2"] = spec.guidance_scale_2
+
+        import time as _time
+
+        t0 = _time.perf_counter()
+        out = omni.generate({"prompt": spec.prompt}, OmniDiffusionSamplingParams(**sampling))
+        elapsed = _time.perf_counter() - t0
+
+        frames = _extract_frames(out)
+        mid = _middle_frame(frames)
+        mid.save(output_path)
+        return mid, elapsed
+    finally:
+        del omni
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.accelerator.empty_cache()
+
+
+@pytest.mark.parametrize("case", _cases())
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.accelerator.device_count() < 1,
+    reason="Needs at least 1 GPU.",
+)
+def test_quantized_videogen_dit_matches_bf16_clip(case: _QuantCase, accuracy_artifact_root: Path) -> None:
+    """FP8 DiT should preserve prompt-aligned video quality vs BF16."""
+    out_dir = model_output_dir(accuracy_artifact_root, case.spec.family + "-quant")
+
+    with tempfile.TemporaryDirectory():
+        bf16_model = os.environ[case.spec.bf16_env]
+        quant_model = os.environ[case.spec.fp8_env]
+
+        bf16_frame, bf16_time = _generate_video(
+            model=bf16_model,
+            spec=case.spec,
+            quantization=None,
+            output_path=out_dir / "bf16.png",
+        )
+        quant_frame, quant_time = _generate_video(
+            model=quant_model,
+            spec=case.spec,
+            quantization="fp8",
+            output_path=out_dir / f"{case.name}.png",
+        )
+
+    clip = CLIPScorer()
+    bf16_clip = clip.score(bf16_frame, case.spec.prompt)
+    quant_clip = clip.score(quant_frame, case.spec.prompt)
+    clip_drop = bf16_clip - quant_clip
+    speedup = (bf16_time / quant_time) if quant_time else float("nan")
+
+    metrics = {
+        "case": case.name,
+        "family": case.spec.family,
+        "quantization": "fp8",
+        "bf16_model": bf16_model,
+        "quant_model": quant_model,
+        "prompt": case.spec.prompt,
+        "height": HEIGHT,
+        "width": WIDTH,
+        "num_frames": NUM_FRAMES,
+        "num_inference_steps": NUM_INFERENCE_STEPS,
+        "bf16_elapsed_s": bf16_time,
+        "quant_elapsed_s": quant_time,
+        "speedup_vs_bf16": speedup,
+        "bf16_clip_score": bf16_clip,
+        "quant_clip_score": quant_clip,
+        "clip_score_drop": clip_drop,
+    }
+    metrics_path = out_dir / f"{case.name}_metrics.json"
+    metrics_path.write_text(json.dumps(metrics, indent=2, sort_keys=True), encoding="utf-8")
+
+    print(f"\n[{case.name}] BF16 {bf16_time:.2f}s | FP8 {quant_time:.2f}s (x{speedup:.2f})")
+    print(f"           CLIP BF16={bf16_clip:.2f}  FP8={quant_clip:.2f}  drop={clip_drop:.2f}")
+    print(f"           metrics={metrics_path}")
+
+    assert quant_clip >= CLIP_ABSOLUTE_FLOOR, (
+        f"{case.name} CLIP below floor: got {quant_clip:.2f}, expected >= {CLIP_ABSOLUTE_FLOOR:.2f}"
+    )
+    assert clip_drop <= case.clip_drop, (
+        f"{case.name} CLIP drop too large: got {clip_drop:.2f}, expected <= {case.clip_drop:.2f}"
+    )

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -433,8 +433,13 @@ class DiffusersPipelineLoader:
         # that have loaded weights tracking currently.
         if loaded_weights is not None:
             weights_not_loaded = weights_to_load - loaded_weights
-            # NOTE: if the model is quantized, ignore not_loaded check for scale weights
-            weights_scale_not_loaded = {name for name in weights_not_loaded if name.endswith("weight_scale")}
+            # NOTE: if the model is quantized, ignore not_loaded check for scale
+            # weights. ModelOpt FP8 carries a per-tensor `weight_scale` and a
+            # static activation `input_scale`, which the quant method may
+            # fold/track differently than plain parameters.
+            weights_scale_not_loaded = {
+                name for name in weights_not_loaded if name.endswith(("weight_scale", "input_scale"))
+            }
             weights_not_loaded = weights_not_loaded - weights_scale_not_loaded
             if weights_not_loaded:
                 self._check_unloaded_weights(weights_not_loaded)

vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from collections.abc import Iterable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn.functional as F
@@ -31,6 +31,9 @@ from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
 logger = init_logger(__name__)
 
 
@@ -332,6 +335,8 @@ class HunyuanVideo15Attention(nn.Module):
         out_bias: bool = True,
         eps: float = 1e-6,
         out_dim: int | None = None,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
 
@@ -350,6 +355,8 @@ class HunyuanVideo15Attention(nn.Module):
             head_size=self.head_dim,
             total_num_heads=self.heads,
             bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_qkv",
         )
 
         self.to_out = nn.ModuleList(
@@ -360,6 +367,8 @@ class HunyuanVideo15Attention(nn.Module):
                     bias=out_bias,
                     input_is_parallel=True,
                     return_bias=False,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.to_out.0",
                 ),
                 nn.Identity(),  # placeholder for dropout (none used)
             ]
@@ -374,6 +383,8 @@ class HunyuanVideo15Attention(nn.Module):
                 head_size=self.head_dim,
                 total_num_heads=self.heads,
                 bias=added_proj_bias,
+                quant_config=quant_config,
+                prefix=f"{prefix}.add_kv_proj",
             )
 
             self.to_add_out = RowParallelLinear(
@@ -382,6 +393,8 @@ class HunyuanVideo15Attention(nn.Module):
                 bias=out_bias,
                 input_is_parallel=True,
                 return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_add_out",
             )
 
         self.rope = RotaryEmbedding(is_neox_style=False)
@@ -401,6 +414,8 @@ class HunyuanVideo15Attention(nn.Module):
         image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
         hidden_states_mask: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        # Ensure contiguous for FP8 quantized linear layers
+        hidden_states = hidden_states.contiguous()
         qkv, _ = self.to_qkv(hidden_states)
         q_size = self.to_qkv.num_heads * self.head_dim
         kv_size = self.to_qkv.num_kv_heads * self.head_dim
@@ -421,6 +436,7 @@ class HunyuanVideo15Attention(nn.Module):
             key = self.rope(key, cos, sin)
 
         if encoder_hidden_states is not None:
+            encoder_hidden_states = encoder_hidden_states.contiguous()
             encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states)
             add_q_size = self.add_kv_proj.num_heads * self.head_dim
             add_kv_size = self.add_kv_proj.num_kv_heads * self.head_dim
@@ -492,6 +508,8 @@ class HunyuanVideo15TransformerBlock(nn.Module):
         attention_head_dim: int,
         mlp_ratio: float = 4.0,
         qk_norm: str = "rms_norm",
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         hidden_size = num_attention_heads * attention_head_dim
@@ -507,13 +525,23 @@ class HunyuanVideo15TransformerBlock(nn.Module):
             out_dim=hidden_size,
             bias=True,
             eps=1e-6,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
         )
 
         self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.ff = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio)
+        self.ff = FeedForward(
+            dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio, quant_config=quant_config, prefix=f"{prefix}.ff"
+        )
 
         self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.ff_context = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio)
+        self.ff_context = FeedForward(
+            dim=hidden_size,
+            dim_out=hidden_size,
+            mult=mlp_ratio,
+            quant_config=quant_config,
+            prefix=f"{prefix}.ff_context",
+        )
 
     def forward(
         self,
@@ -601,6 +629,7 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
         target_size: int = 640,
         task_type: str = "i2v",
         use_meanflow: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
     ):
         super().__init__()
 
@@ -632,9 +661,14 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
         self.transformer_blocks = nn.ModuleList(
             [
                 HunyuanVideo15TransformerBlock(
-                    num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm
+                    num_attention_heads,
+                    attention_head_dim,
+                    mlp_ratio=mlp_ratio,
+                    qk_norm=qk_norm,
+                    quant_config=quant_config,
+                    prefix=f"transformer_blocks.{i}",
                 )
-                for _ in range(num_layers)
+                for i in range(num_layers)
             ]
         )
 

vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py

@@ -143,7 +143,9 @@ class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, Diff
             self.scheduler._shift = od_config.flow_shift
 
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HunyuanVideo15Transformer3DModel)
-        self.transformer = HunyuanVideo15Transformer3DModel(od_config=od_config, **transformer_kwargs)
+        self.transformer = HunyuanVideo15Transformer3DModel(
+            od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
+        )
 
         # Check if model uses meanflow (distilled variants)
         self.use_meanflow = getattr(od_config.tf_model_config, "use_meanflow", False)

vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml

@@ -0,0 +1,38 @@
+# Stage config for running HunyuanVideo-1.5 DiT with ModelOpt FP8 auto-detect.
+# Single GPU. Bump `tensor_parallel_size` and `devices` for multi-GPU TP.
+#
+# Use with a ModelOpt FP8 checkpoint (e.g. produced by
+# examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py).
+
+stage_args:
+  - stage_id: 0
+    stage_type: diffusion
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: dit
+      model_class_name: HunyuanVideo15Pipeline
+      max_num_seqs: 1
+      enforce_eager: true
+      trust_remote_code: true
+      distributed_executor_backend: "mp"
+      # HunyuanVideo-1.5 peak VRAM is dominated by the single-shot VAE decode of
+      # the full video (the DiT contributes little), so tiling the decode is the
+      # main memory lever — far more than FP8 (e.g. 480x832x33: ~117 -> ~47 GiB).
+      # Default it on here; override with --vae-use-tiling/false if you have headroom.
+      vae_use_tiling: true
+      parallel_config:
+        tensor_parallel_size: 1
+
+    final_output: true
+    final_output_type: video
+    is_comprehension: false
+    default_sampling_params:
+      seed: 42
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1

vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml

@@ -0,0 +1,41 @@
+# Stage config for running Wan2.2 TI2V-5B DiT with ModelOpt FP8 auto-detect.
+# Single GPU (TI2V-5B fits 80GB BF16; FP8 drops by ~half).
+# For the A14B MoE variants, bump `tensor_parallel_size` and `devices`.
+#
+# Use with a ModelOpt FP8 checkpoint (e.g. produced by
+# examples/quantization/quantize_wan2_2_modelopt_fp8.py).
+
+stage_args:
+  - stage_id: 0
+    stage_type: diffusion
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: dit
+      # Wan2.2 TI2V-5B is served by the base Wan22Pipeline, which auto-detects
+      # expand_timesteps (TI2V-5B mode) from model_index.json. (The standalone
+      # Wan22TI2VPipeline was consolidated into Wan22Pipeline upstream.)
+      model_class_name: Wan22Pipeline
+      max_num_seqs: 1
+      enforce_eager: true
+      trust_remote_code: true
+      distributed_executor_backend: "mp"
+      # Tile the VAE decode by default. Wan2.2's 480p peak is DiT-dominated (so
+      # FP8 is the main lever here), but tiling adds headroom for higher
+      # resolutions / longer clips at negligible cost. Override with --vae-use-tiling/false.
+      vae_use_tiling: true
+      parallel_config:
+        tensor_parallel_size: 1
+
+    final_output: true
+    final_output_type: video
+    is_comprehension: false
+    default_sampling_params:
+      seed: 42
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1

vllm_omni/patch.py

@@ -217,3 +217,43 @@ def _patch_scaled_mm_fp8_contiguous_activation():
 
 
 _patch_scaled_mm_fp8_contiguous_activation()
+
+
+def _patch_flashinfer_fp8_scaled_mm_output_shape():
+    """Restore the N-D output shape for the FlashInfer FP8 ScaledMM kernel.
+
+    ``FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm`` returns the raw 2-D
+    GEMM result and ignores ``output_shape``, unlike the CUTLASS / PyTorch
+    ScaledMM kernels which reshape to it. A 3-D activation ``(B, S, D)`` thus
+    collapses to ``(B*S, D)``, breaking diffusion DiTs that reshape the linear
+    output by absolute dim (e.g. Wan2.2 ``qkv.unflatten(2, ...)``) with
+    ``IndexError: Dimension out of range``. It only bites >2-D inputs, so LLM
+    (token-flattened, 2-D) paths are unaffected.
+
+    Carried here because the upstream fix may not have landed yet; this override
+    becomes a harmless no-op once vLLM honors ``output_shape`` itself.
+    """
+    try:
+        from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
+            FlashInferFP8ScaledMMLinearKernel,
+        )
+    except ImportError:
+        return
+
+    _original_apply_scaled_mm = FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm
+    if getattr(_original_apply_scaled_mm, "_omni_output_shape_patched", False):
+        return
+
+    def _patched_apply_scaled_mm(self, *, A, B, out_dtype, As, Bs, bias, output_shape):  # noqa: N803
+        out = _original_apply_scaled_mm(
+            self, A=A, B=B, out_dtype=out_dtype, As=As, Bs=Bs, bias=bias, output_shape=output_shape
+        )
+        if tuple(out.shape) != tuple(output_shape):
+            out = out.view(*output_shape)
+        return out
+
+    _patched_apply_scaled_mm._omni_output_shape_patched = True
+    FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm = _patched_apply_scaled_mm
+
+
+_patch_flashinfer_fp8_scaled_mm_output_shape()