Compare commits

...
Author SHA1 Message Date
Hongsheng Liu 094abd70c1 feat(diffusion): add request-gated VAE fast paths
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-11 09:31:27 +08:00
fan2956andFan c27623c29c [Bugfix] rainfusion support all shape video (#6000)
Signed-off-by: Fan <fan@FandeMacBook-Air.local>
Co-authored-by: Fan <fan@FandeMacBook-Air.local>
2026-08-11 05:10:50 +08:00
Yueqian Lin f8fa86fb1f [BugFix] Don't fail the TTS ratchet when the branch count goes down (#6008)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-08-11 05:08:31 +08:00
ooooooyeandHongsheng Liu 422d8fd8f9 [Bugfix][Diffusion] Restore supports_packed_mask_free on teacache FakeBackend (#5997)
Signed-off-by: brandneway <gyuan4892@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-10 22:20:37 +08:00
Hongsheng Liu 62e239bf1e [Docs] Fix broken attention backend link (#5998)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-10 21:55:01 +08:00
ooooooye 740f45abcd [Perf][Diffusion] MiniMax-H3: opt into packed varlen attention on NPU to eliminate quadratic mask materialization (#5891)
Signed-off-by: brandneway <gyuan4892@gmail.com>
2026-08-10 19:47:03 +08:00
5e4b9f8c6e [Misc] Align CODEOWNERS with module/feature design docs (#5958)
Signed-off-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
2026-08-10 19:25:03 +08:00
Hongsheng Liu ae811043e9 docs: keep shared task examples in navigation (#5987)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-10 19:13:52 +08:00
R0CKSTAR 636d44e272 [Perf][MUSA] Restore MiniMax-H3 dynamic RoPE fusion (#5881)
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
2026-08-10 17:24:46 +08:00
yiminghub2024 8e44e78038 Add Ref2VA measurements to the MiniMax-H3 DGX Spark (GB10) recipe (#5972)
Signed-off-by: yiminghub2024 <482890@qq.com>
2026-08-10 15:51:15 +08:00
Hongsheng Liu b98b7c8529 [Docs] Preserve generated quantization URLs and link recipes in supported models (#5969)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-10 15:45:50 +08:00
61 changed files with 2751 additions and 1384 deletions

No files matched your search

+110 -28
View File
@@ -1,6 +1,11 @@
# CODEOWNERS is used to automatically assign PR reviewers
# Format: path pattern @username or @team
# More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
#
# The vllm_omni/** source-tree rules below are derived from the `owners` and
# `primary_code_paths` front matter in docs/design/module/*.md. When updating
# ownership, update the design doc first, then regenerate the matching rule here.
# Last-match-wins: more specific rules must come AFTER broader ones.
# Default owners (fallback reviewers for all files)
* @Gaohan123 @hsliuustc0106 @ywang96
@@ -21,25 +26,10 @@
# CI
/.buildkite/ @yenuo26 @congw729 @NickCao
# Engine and model_executor
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks @NickCao
/vllm_omni/model_executor/ @tzhouam @gcanlin
/vllm_omni/worker/ @tzhouam @gcanlin
# ComfyUI / external apps
/apps/ @fhfuih @alex-jw-brooks
# Diffusion
/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
/vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
# Diffusion CPU offloading
/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao @xuechendi
# Omni and TTS support
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin @NickCao
# Benchmarks
# Benchmarks (out-of-tree)
/benchmarks/accuracy/ @david6666666
/benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
@@ -47,20 +37,112 @@
/benchmarks/glm_image/ @princepride @JaredforReal
/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
# Usability and Observability
/vllm_omni/deploy/ @tzhouam @lishunyang12
/vllm_omni/metrics/ @lishunyang12
/vllm_omni/profiler/ @gcanlin
# =====================================================================
# Module design ownership (docs/design/module/*.md)
# Each section header cites the source design doc.
# =====================================================================
# ComfyUI
/apps/ @fhfuih @alex-jw-brooks
# Autoregressive runtime — docs/design/module/ar_runtime.md
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
/vllm_omni/worker/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
# Large Scale Serving
# Engine orchestration & stage runtime — docs/design/module/engine_orchestration.md, stage_runtime.md
/vllm_omni/engine/ @tzhouam @fake0fan
# File-level overrides: input/output modality contracts own these engine files
# — docs/design/module/input_output_modality_contracts.md
/vllm_omni/engine/messages.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/engine/serialization.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/engine/mm_outputs.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/engine/output_modality.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/engine/output_processor.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
# Input/output modality contracts — docs/design/module/input_output_modality_contracts.md
/vllm_omni/inputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/outputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/request.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/data_entry_keys.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
# Entrypoints and serving boundaries — docs/design/module/entrypoints.md
/vllm_omni/entrypoints/ @alex-jw-brooks @linyueqian @NickCao
# File-level overrides within entrypoints/
/vllm_omni/entrypoints/openai/errors.py @alex-jw-brooks @NickCao
/vllm_omni/entrypoints/cli/benchmark/ @alex-jw-brooks @Bounty-hunter
# Error contracts — docs/design/module/error_contracts.md
/vllm_omni/errors.py @alex-jw-brooks @NickCao
# OmniConnector (inter-stage transport) — docs/design/module/omni_connector.md
/vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
# Quantization and configuration
# Model integration — docs/design/module/model_integration.md
/vllm_omni/model_executor/ @tzhouam @gcanlin
/vllm_omni/model_executor/models/ @tzhouam @gcanlin @Sy0307 @amy-why-3459 @linyueqian
/vllm_omni/model_extras/ @tzhouam @gcanlin @princepride
/vllm_omni/plugins/ @gcanlin @tjtanaa @xuechendi
# Config override within model_executor — docs/design/module/vllm_omni_config.md
/vllm_omni/model_executor/stage_configs/ @lishunyang12 @alex-jw-brooks
# Execution platforms (hardware backends) — docs/design/module/execution_platforms.md
/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
/vllm_omni/attention/ @gcanlin @tjtanaa @xuechendi @david6666666 @lishunyang12
# Configuration & deploy — docs/design/module/vllm_omni_config.md
/vllm_omni/config/ @lishunyang12 @alex-jw-brooks
/vllm_omni/deploy/ @lishunyang12 @alex-jw-brooks
# Quantization — docs/design/module/quantization.md
/vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
# Hardware plugins and accelerators
/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
# Observability — docs/design/module/observability.md
/vllm_omni/metrics/ @lishunyang12 @vraiti
# Profiling — docs/design/module/profiling.md
/vllm_omni/profiler/ @gcanlin
# Benchmarking (in-tree) — docs/design/module/benchmarking.md
/vllm_omni/benchmarks/ @alex-jw-brooks @Bounty-hunter
# LoRA (no dedicated design doc; ownership intentionally narrowed)
/vllm_omni/lora/ @SamitHuang @wtomin
# Transformers utils (related path only in design docs; preserves prior ownership)
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
# =====================================================================
# Diffusion module family — docs/design/module/diffusion/
# Umbrella rule first; sub-module rules after it win via last-match.
# =====================================================================
# Diffusion umbrella (top-level files: data.py, forward_context.py, etc.)
# — docs/design/module/diffusion/index.md
/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
# Cache management — docs/design/module/cache_management.md
/vllm_omni/diffusion/cache/ @Isotr0py @princepride @SamitHuang
# Diffusion quantization — docs/design/module/quantization.md
/vllm_omni/diffusion/quantization/ @david6666666 @Isotr0py @lishunyang12
# Diffusion profiling — docs/design/module/profiling.md
/vllm_omni/diffusion/profiler/ @gcanlin
# Diffusion runtime (scheduler, executor, worker)
# — docs/design/module/diffusion/diffusion_runtime.md, continuous_batching.md
/vllm_omni/diffusion/sched/ @Isotr0py @princepride @SamitHuang @fhfuih
/vllm_omni/diffusion/executor/ @Isotr0py @princepride @SamitHuang @fhfuih
/vllm_omni/diffusion/worker/ @Isotr0py @princepride @SamitHuang @fhfuih
# Diffusion model integration (pipelines, registration, checkpoint loading)
# — docs/design/module/diffusion/diffusion_model_integration.md
/vllm_omni/diffusion/models/ @fhfuih @Bounty-hunter @wtomin
/vllm_omni/diffusion/model_loader/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
/vllm_omni/diffusion/lora/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
# Diffusion parallelism (rank topology, TP/SP/PP/VAE patch)
# — docs/design/module/diffusion/parallelism.md
/vllm_omni/diffusion/distributed/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
/vllm_omni/diffusion/attention/parallel/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
# Diffusion offloader (CPU offloading, DLO)
# — docs/design/module/diffusion/offloader.md
/vllm_omni/diffusion/offloader/ @yuanheng-zhao @david6666666 @lishunyang12
+4 -1
View File
@@ -66,7 +66,10 @@ repos:
entry: python tools/pre_commit/check_tts_adapter.py
language: python
pass_filenames: false
files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)$
files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
# The budget-is-now-slack reminder is printed on a PASSING run, and
# pre-commit hides output from passing hooks unless the hook is verbose.
verbose: true
# Keep `suggestion` last
- id: suggestion
+4 -14
View File
@@ -16,31 +16,21 @@ nav:
- Examples:
- examples/README.md
- Offline Inference:
- GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
- HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
- InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
- MiMo-Audio: user_guide/examples/offline_inference/mimo_audio.md
- Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
- Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
- Speech-To-Video: user_guide/examples/offline_inference/speech_to_video.md
- Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
- Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
- Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
- X-To-Text: user_guide/examples/offline_inference/x_to_text.md
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
- Online Serving:
- BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
- vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
- Diffusers Backend Adapter: user_guide/examples/online_serving/diffusers_pipeline_adapter.md
- GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
- MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
- Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
- Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
- Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
- Speech-To-Video: user_guide/examples/online_serving/speech_to_video.md
- Text-To-Image: user_guide/examples/online_serving/text_to_image.md
- Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
- Text-To-Video: user_guide/examples/online_serving/text_to_video.md
- General:
- usage/*
+1
View File
@@ -7,6 +7,7 @@ owners:
- "@yinpeiqi"
- "@fake0fan"
- "@Sy0307"
- "@Gaohan123"
primary_code_paths:
- vllm_omni/core/**
- vllm_omni/worker/**
+1
View File
@@ -4,6 +4,7 @@ kind: module
status: draft
owners:
- "@alex-jw-brooks"
- "@Bounty-hunter"
primary_code_paths:
- vllm_omni/benchmarks/**
- vllm_omni/entrypoints/cli/benchmark/**
+1
View File
@@ -5,6 +5,7 @@ status: draft
owners:
- "@Isotr0py"
- "@princepride"
- "@SamitHuang"
primary_code_paths:
- vllm_omni/diffusion/cache/**
related_code_paths:
@@ -3,9 +3,10 @@ title: Diffusion Model Integration
kind: module
status: draft
owners:
- "@Isotr0py"
- "@fhfuih"
- "@Bounty-hunter"
- "@wtomin"
- "@ZJY0516"
- "@RuixiangMa"
primary_code_paths:
- vllm_omni/diffusion/models/**
- vllm_omni/diffusion/model_loader/**
@@ -6,6 +6,7 @@ owners:
- "@Isotr0py"
- "@princepride"
- "@SamitHuang"
- "@fhfuih"
primary_code_paths:
- vllm_omni/diffusion/sched/**
- vllm_omni/diffusion/executor/**
+2 -1
View File
@@ -4,7 +4,8 @@ kind: module
status: draft
owners:
- "@yuanheng-zhao"
- "@Isotr0py"
- "@david6666666"
- "@lishunyang12"
primary_code_paths:
- vllm_omni/diffusion/offloader/**
related_code_paths:
@@ -6,6 +6,7 @@ owners:
- "@princepride"
- "@yuanheng-zhao"
- "@xuechendi"
- "@Bounty-hunter"
primary_code_paths:
- vllm_omni/diffusion/distributed/**
- vllm_omni/diffusion/attention/parallel/**
+1
View File
@@ -6,6 +6,7 @@ architecture_state: current-plus-in-flight
owners:
- "@alex-jw-brooks"
- "@linyueqian"
- "@NickCao"
document_stewards:
- "@hsliuustc0106"
- "@Gaohan123"
+8 -1
View File
@@ -8,8 +8,15 @@ owners:
- "@xuechendi"
primary_code_paths:
- vllm_omni/platforms/**
related_code_paths:
- vllm_omni/attention/**
primary_code_path_owners:
# Additional owners scoped to attention kernels (quantized/custom attention)
- paths:
- vllm_omni/attention/**
owners:
- "@david6666666"
- "@lishunyang12"
related_code_paths: []
depends_on:
- ar_runtime.md
- diffusion/index.md
@@ -6,6 +6,8 @@ architecture_state: current
owners:
- "@Sy0307"
- "@amy-why-3459"
- "@Gaohan123"
- "@alex-jw-brooks"
document_stewards:
- "@hsliuustc0106"
- "@Gaohan123"
+8
View File
@@ -9,6 +9,14 @@ primary_code_paths:
- vllm_omni/model_executor/**
- vllm_omni/model_extras/**
- vllm_omni/plugins/**
primary_code_path_owners:
# Additional owners scoped to a sub-path of model_executor/**
- paths:
- vllm_omni/model_executor/models/**
owners:
- "@Sy0307"
- "@amy-why-3459"
- "@linyueqian"
related_code_paths:
- vllm_omni/transformers_utils/**
- vllm_omni/tokenizers/**
+1
View File
@@ -4,6 +4,7 @@ kind: module
status: draft
owners:
- "@lishunyang12"
- "@vraiti"
primary_code_paths:
- vllm_omni/metrics/**
related_code_paths:
+37 -12
View File
@@ -24,6 +24,23 @@ SERVING_MODE_TITLES = {
"online_serving": "Online serving",
}
# Keep the Examples navigation focused on reusable modality/task entry points.
# Model-specific scripts remain available under ``examples/`` and are surfaced
# from these shared pages instead of getting one page per model.
GENERAL_EXAMPLE_SLUGS = frozenset(
{
"image_to_image",
"image_to_video",
"speech_to_video",
"text_to_audio",
"text_to_image",
"text_to_speech",
"text_to_video",
"x_to_text",
"x_to_video_audio",
}
)
def load_model_display_names() -> dict[str, str]:
try:
@@ -46,6 +63,13 @@ def load_model_display_names() -> dict[str, str]:
MODEL_DISPLAY_NAMES = load_model_display_names()
def is_general_example(example: "Example") -> bool:
"""Return whether an example should get a generated serving page."""
if example.category not in SERVING_MODE_TITLES:
return True
return example.path.stem in GENERAL_EXAMPLE_SLUGS
def fix_case(text: str) -> str:
subs = {
"api": "API",
@@ -340,22 +364,21 @@ def update_nav_file(examples: list[Example]):
examples_by_category[category] = []
examples_by_category[category].append(example)
# Build the new Examples section - start with preserved items
examples_section = preserved_items.copy()
# Add examples grouped by category, sorted by category name
for category in sorted(examples_by_category.keys()):
category_examples = sorted(examples_by_category[category], key=lambda e: e.path.stem)
category_items = []
for example in category_examples:
def category_items(category: str) -> list[dict[str, str]]:
items = []
for example in sorted(examples_by_category[category], key=lambda e: e.path.stem):
doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
rel_path = doc_path.relative_to(ROOT_DIR / "docs")
category_items.append({example.nav_title: rel_path.as_posix()})
items.append({example.nav_title: rel_path.as_posix()})
return items
if category_items:
# Format category name (e.g., "offline_inference" -> "Offline Inference")
# Build the new Examples section - start with preserved items.
examples_section = preserved_items.copy()
for category in sorted(examples_by_category.keys()):
items = category_items(category)
if items:
category_title = fix_case(category.replace("_", " ").title())
examples_section.append({category_title: category_items})
examples_section.append({category_title: items})
# Update the nav structure
nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"] = examples_section
@@ -390,6 +413,8 @@ def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
for path in category.glob("*/*.md"):
examples.append(Example(path.parent, category.stem))
examples = [example for example in examples if is_general_example(example)]
# Generate the example documentation
for example in sorted(examples, key=lambda e: e.path.stem):
example_name = f"{example.path.stem}.md"
+87 -78
View File
@@ -17,84 +17,93 @@ th {
}
</style>
| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU |
|--------------|--------|-------------------|------------|---------|-----|-----------|
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | |
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | |
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | |
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | |
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | |
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | |
| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ |
| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ |
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | |
| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | |
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | |
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | |
| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | |
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | |
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ |
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU | Recipe |
|--------------|--------|-------------------|------------|---------|-----|-----------|--------|
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) |
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — |
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | — |
| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | — |
| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | — |
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — |
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — |
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | — |
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | | — |
| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — |
| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — |
| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | — |
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — |
| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — |
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | | — |
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | — |
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | — |
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | — |
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | | — |
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | | — |
| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | | [Published](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b) |
| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | | — |
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ | — |
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | | — |
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | — |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
The `Recipe` column links the corresponding deployment recipe. `Published` is
preferred when a corresponding `recipes.vllm.ai` page exists; otherwise the
link points to the repository recipe. `—` means that no recipe link has been
added for that model row yet.
For rows with an audited recipe link, the hardware checkmarks reflect hardware
explicitly documented in that recipe. Other rows retain the implementation
support metadata until their recipe is audited.
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
+12 -15
View File
@@ -283,23 +283,20 @@ blocks — use Ulysses SP (`ring_degree=1`).
### Which geometries run sparse
Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count
is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own
tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows
that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense
attention and logs `RAINFUSION_ATTN staying dense` with the row count it computed — but it gets no
speedup, so pick an aligned geometry when you want one.
MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
promoted rows remain visible to every query and key.
Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8
spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact
patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns
are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially
distant positions and selection can no longer rank them meaningfully. This is silent, and invisible
to a `sparsity=0` check, because a fully populated mask does not care how blocks are grouped.
For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
and kernel tiling aligned, at the cost of retaining more blocks.
"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned
resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`,
which you can buy back with `start_step`.
8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
that implements protected video tails for rf_v2.
## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
@@ -99,7 +99,7 @@ a native vLLM-Omni pipeline for those cases.
### Attention Backends
The diffusers backend converts
[vLLM-Omni standard of attention backend setting](../../../docs/user_guide/diffusion/attention_backends.md)
[vLLM-Omni standard of attention backend setting](../../diffusion/attention_backends.md)
to [diffusers standard](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends).
Specifically for `FLASH_ATTN`, it will first attempt to use FlashAttention-3 and then FlashAttention-2.
@@ -1,247 +0,0 @@
#!/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()
@@ -1,448 +0,0 @@
#!/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()
@@ -1,471 +0,0 @@
#!/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()
+67 -4
View File
@@ -79,7 +79,8 @@ sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`.
### Multi-NPU: 768P combined-service configuration
Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism
degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload:
degree 8, native tiled VAE patch parallelism degree 8, and distributed
layerwise offload:
```bash
export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
@@ -98,7 +99,7 @@ vllm serve "${MODEL}" \
--usp 8 \
--ring 1 \
--text-encoder-tp-size 8 \
--enable-layerwise-offload \
--enable-distributed-layerwise-offload \
--vae-parallel-mode tile \
--vae-use-tiling \
--vae-patch-parallel-size 8 \
@@ -112,6 +113,11 @@ H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no
partition restart is required. Layerwise offload applies to both DiTs.
On Atlas 800I A3 (64 GB HBM per device) the combined service does not fit at
768P without offloading or sharding: use distributed layerwise offload (as
above) or HSDP — see
[§ Memory and attention optimizations](#memory-and-attention-optimizations-a3).
### Optional optimizations
Two independent optimizations may be enabled on top of the configuration
@@ -134,6 +140,63 @@ Keep `--ring 1` when using RainFusion: the `rf_v2` kernel ranks key blocks
over the whole sequence, so ring parallelism would split away the keys it
needs. Scale with `--usp` instead.
## Memory and attention optimizations (A3)
### Fitting 768P into 64 GB HBM: distributed layerwise offload or HSDP
Each Atlas 800I A3 NPU has 64 GB of HBM, which is not enough for the
combined MiniMax-H3 service at 768P. Enable **one** of the following at
server startup:
**Distributed layerwise offload** — DiT layers are offloaded with parameters
gathered across the parallel group instead of replicated per rank:
```bash
--enable-distributed-layerwise-offload
```
Measured peak NPU memory is about 45 GB per device for Ref2VA with a 13.88 s
input video generating a 15 s 768P (1344x768) output video. Host (CPU) memory
usage is high with this option.
**HSDP** — hybrid sharded data parallelism for the DiT parameters. The
multi-stream memory-reuse knob is mandatory with this flag:
```bash
export MULTI_STREAM_MEMORY_REUSE=2
```
```bash
--use-hsdp
```
Host memory usage is small, but large shapes may still OOM; HBM usage
optimizations for this path are ongoing.
### FLASH_ATTN backend with MindIE-SD
Keep `--diffusion-attention-backend FLASH_ATTN` and install MindIE-SD (see
[§ Environment](#environment)). On NPU this backend carries most of the
memory-reduction and performance work: a mask-free packed varlen path that
never materializes the quadratic `full_qk` padding mask, and K/V prefix
slicing, both driven by the packed `cu_seqlens` metadata emitted by the H3
transformer.
### Optional: LaserAttention fused kernel
For an additional attention speedup, select the Ascend LaserAttention fused
kernel before starting the server:
```bash
export MINDIE_SD_FA_TYPE="ascend_laser_attention"
```
This requires the FLASH_ATTN backend and MindIE-SD. H3 automatically applies
exact power-of-two input pre-scaling (`laser_input_scale=256`) so the
kernel's fp16 workspace cannot overflow on outlier activations. Measured
speedup numbers will be added here.
## HTTP API examples
The request format is identical to the GPU recipe; see
@@ -154,8 +217,8 @@ configuration above:
| Workload | Configuration |
|----------|---------------|
| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
| T2VA, 209 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
These measurements describe the validated shapes rather than a general
throughput guarantee.
+128 -15
View File
@@ -4,6 +4,17 @@ This recipe uses online FP8 weight quantization, tiled VAE decode, and a single
resident partition. GB10 is a unified-memory platform, so unlike the discrete-GPU
recipes it uses **no offload of any kind** — see the capacity note below.
Validated on:
- Host: DGX Spark (GB10), aarch64
- GPUs: 1 (unified memory)
- Driver: <FILL_DRIVER_VERSION>
- vLLM: 0.26.0
- vLLM-Omni: `main` at `e1aa6eae75c460cd1893bc320546e81e66973831`
- Workloads: T2VA and Ref2VA, 960x576, `duration=8.0`, 50 steps, `flow_shift=12`,
`seed=1101`
## Capacity requirements
| Resource | DGX Spark (GB10) |
@@ -11,7 +22,8 @@ recipes it uses **no offload of any kind** — see the capacity note below.
| Unified memory | 128 GB LPDDR5X (~121 GiB usable) |
| GPUs | 1 |
| Checkpoint storage | 135 GiB per partition |
| Measured peak allocator high-water (FP8) | 97.7 GiB |
| Measured peak allocator high-water, T2VA (FP8) | 97.7 GiB |
| Measured peak allocator high-water, Ref2VA (FP8) | 102.8 GiB |
`FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one server
at a time.
@@ -29,9 +41,21 @@ math relative to the RTX and datacenter recipes:
fit in 121 GiB. Online FP8 quantizes the DiT only (62 GiB to ~31 GiB); the
Qwen3-VL text encoder and both VAEs stay BF16.
With ~97.7 GiB of the pool held by the allocator at peak, there is roughly 23 GiB
of headroom for the OS, page cache, and any other process on the box. Treat this
platform as single-tenant while a request is in flight.
With ~97.7 GiB of the pool held by the allocator at peak, T2VA leaves roughly
23 GiB of headroom for the OS, page cache, and any other process on the box.
Ref2VA is tighter on both counts. The allocator high-water rises to 102.8 GiB,
and the allocator figure understates true pool usage.
Sampling `/proc/meminfo` at 1 Hz across a 50-step Ref2VA request on a host with
`MemTotal: 127600524 kB` (121.7 GiB), `MemAvailable` started at 8.1 GiB, fell to
a trough of **7.4 GiB**, and recovered to 8.9 GiB. At the trough the pool held
roughly **114.3 GiB** — 11.4 GiB more than the 102.8 GiB the allocator reports,
which is the CUDA context, non-PyTorch allocations, and page cache the header
does not count.
Treat this platform as strictly single-tenant for Ref2VA. There is no room for a
second process, and adding reference images, reference videos, or a larger output
shape will run into the OOM killer.
## One GB10: 960x576, 8 seconds
@@ -94,6 +118,52 @@ ffprobe -v error -show_entries \
Both requests return `200 OK` with an MP4 body: H.264 video at the requested
geometry plus a 32 kHz stereo AAC track.
## One GB10: Ref2VA, 960x576, 8 seconds
`FL2VA` and `Ref2VA` are separate 135 GiB partitions and only one fits at a time.
Stop the FL2VA server before starting this one.
```bash
CUDA_VISIBLE_DEVICES=0 \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
vllm serve /path/to/MiniMax-H3/Ref2VA \
--omni --trust-remote-code --host 0.0.0.0 --port 8000 \
--init-timeout 3600 \
--num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \
--usp 1 --ring 1 --vae-patch-parallel-size 1 \
--vae-parallel-mode tile --vae-use-tiling \
--quantization fp8 --enforce-eager \
--diffusion-attention-backend CUDNN_ATTN
```
The 7200 s timeout used for T2VA is too short here: a 50-step Ref2VA request
takes over an hour on this platform. Use 14400 s.
Unlike `t2va`, `ref2va` defaults to a 16:9 output ratio, so `aspect_ratio` can be
omitted when `width` and `height` are supplied.
```bash
curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \
-D ref2va-headers.txt \
-w 'client_e2e_s=%{time_total}\n' \
--max-time 14400 \
-F 'input_reference=@/path/to/reference.jpg;type=image/jpeg' \
-F 'prompt=A hand rests on a wooden table in warm evening light, then slowly opens and turns toward the camera, with quiet room ambience.' \
-F 'width=960' \
-F 'height=576' \
-F 'fps=24' \
-F 'num_inference_steps=50' \
-F 'flow_shift=12' \
-F 'seed=1101' \
-F 'extra_params={"task":"ref2va","duration":8.0,"audio_flow_shift":3.0}' \
-o h3-gb10-ref2va.mp4
```
The response is `200 OK` with an MP4 body: H.264 at the requested geometry plus a
32 kHz stereo AAC track. `ffprobe` reports 8.032 s for `duration=8.0`.
## Measured latency
Single GB10, FP8, eager, cuDNN attention, tiled VAE, one request at a time:
@@ -102,32 +172,69 @@ Single GB10, FP8, eager, cuDNN attention, tiled VAE, one request at a time:
| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
| T2VA, 4 s, smoke shape | 10 | — | — | 8.68 s | — | 0.49 s | 95.9 s | — |
| T2VA, 8 s, 960x576 | 50 | 0.25 s | 2088 s | 42.61 s | 70 s | 4.84 s | 2169.4 s (36 min 9 s) | 97.7 GiB |
| Ref2VA, 8 s, 960x576, one image (3 runs) | 10 | 8.43-10.94 s | 665.5-748.8 s | 66.55-74.88 s | 67.9-69.8 s | 3.63 s | 770.3-861.1 s | 102.8 GiB |
| Ref2VA, 8 s, 960x576, one image | 50 | 9.66 s | 4039.6 s | 80.79 s | 69.4 s | 3.63 s | 4157.0 s (69 min 17 s) | 102.8 GiB |
How these were measured:
- Stage times come from `--enable-diffusion-pipeline-profiler`, read from the
`X-Stage-Durations` response header of `/v1/videos/sync`.
`X-Stage-Durations` response header of `/v1/videos/sync`. The keys are
`MiniMaxH3Pipeline.encode_prompt` (text encode),
`MiniMaxH3Pipeline.diffuse` (denoise), and `MiniMaxH3Pipeline.decode`
(VAE decode).
- MP4 mux is not a header field. It is the
`Video response encoding (MP4 bytes)` line in the server log.
- Peak memory is the `X-Peak-Memory-MB` header, which reports the CUDA allocator
high-water mark (`torch.cuda.max_memory_reserved`) for the request. It excludes
the CUDA context and non-PyTorch allocations, so true pool usage is somewhat
higher. `nvidia-smi --query-gpu=memory.used` is not usable on GB10 — unified
memory has no discrete VRAM to report and the field returns `[N/A]`.
- Text encode and VAE decode were measured on a separate 10-step run at the same
960x576 / 8 s geometry. Both are independent of step count, so they transfer to
the 50-step row unchanged.
memory has no discrete VRAM to report and the field returns `[N/A]`. See the
capacity section for a `/proc/meminfo` cross-check.
- For T2VA, text encode and VAE decode were measured on a separate 10-step run at
the same 960x576 / 8 s geometry. Both are independent of step count, so they
transfer to the 50-step row unchanged. The Ref2VA MP4 mux figure transfers the
same way, from the 10-step run at identical output geometry and duration.
- Ref2VA peak memory is byte-identical between the 10-step and 50-step runs
(`X-Peak-Memory-MB: 105318.000` in both), confirming that peak is set by output
geometry and reference input rather than step count.
- Stage times do not sum exactly to end-to-end. Queueing, result transfer, and
engine bookkeeping sit outside the profiled stages.
Denoise dominates at 96% of end-to-end, so step count and geometry are the only
knobs that materially change the wall clock on this platform.
Repeated identical 10-step requests varied by up to 23% in denoise time
Ref2VA costs 1.92x the wall clock of T2VA at the same geometry and step count:
4157.0 s versus 2169.4 s. Text encoding shows the most dramatic ratio — 9.66 s
versus 0.25 s, 39x — but at 0.2% of end-to-end it is not what makes Ref2VA slow.
The cost is in denoise. A single reference image expands the Qwen3-VL
presentation to 6883 tokens (logged as `MiniMax H3 ref2va Qwen presentation:
6883 tokens, 1 reference images`), and those tokens sit in the attention context
of every denoise step. Per-step time rises from 42.61 s to 80.79 s, and that
1.90x multiplies across all 50 steps. VAE decode is unchanged at ~69 s, as
expected for an identical output shape.
Repeated identical 10-step T2VA requests varied by up to 23% in denoise time
(397 s to 490 s) on a machine that had already been under sustained load. Treat
these as single-run measurements on one machine, not a throughput guarantee, and
expect the slower end of that range for back-to-back jobs.
Ref2VA per-step time is load-dependent in the same way. On an idle machine the
first 10-step request ran at 66.55 s per step. The 50-step request that followed
occupied the box for 69 consecutive minutes and averaged 80.79 s per step. Two
further 10-step requests taken immediately afterwards settled at 74.88 s and
74.31 s — 12.5% slower than the idle run, and within 0.8% of each other, so this
is a stable thermal state rather than measurement noise.
Plan for 81 s per step when sizing a full 50-step Ref2VA job; 66 s per step is
only reachable on a cold machine and does not hold for the length of one request.
Ref2VA text encode converged downward across the four runs (10.94 s, 9.66 s,
8.52 s, 8.43 s) as first-touch costs fell away; treat ~8.5 s as steady state.
VAE decode was stable at 67.9-69.8 s throughout.
Set `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` well above the expected request time — the
default 1800 s is shorter than a 50-step run on this platform.
default 1800 s is shorter than a 50-step run on this platform, and 7200 s is
shorter than a 50-step Ref2VA run.
## Known limitations
@@ -139,11 +246,17 @@ default 1800 s is shorter than a 50-step run on this platform.
- `t2va` rejects requests without an explicit `aspect_ratio`. When `width` and
`height` are both supplied they take precedence and `aspect_ratio` only has to
name one of the supported ratios.
- Sustained load degrades throughput. See the variance note above.
- Sustained load degrades throughput. See the variance notes above.
- Only one partition is resident. For Ref2VA, stop the FL2VA server and restart
the same command with `/path/to/MiniMax-H3/Ref2VA`. Ref2VA reference video
count and prompt length increase activation memory; begin with one request at
a time and keep the short edge at 576.
with `/path/to/MiniMax-H3/Ref2VA`. Ref2VA with one reference image is measured
above; reference videos and additional images push the token count higher and
were not measured. With `MemAvailable` bottoming out at 7.4 GiB during a
single-image request, there is very little room to grow — re-measure before
adding references or raising the output shape.
- Ref2VA requires a vLLM-Omni build newer than `v0.26.0`. See the version note
under **Validated on**: FP8 weight loading and image-only Ref2VA are both
broken on `release/v0.26.0`.suggestion is `v0.26.1`.
- Online FP8 is incompatible with layerwise offload — the offload path produces a
weight stride the Cutlass FP8 kernel rejects. This is not a practical
restriction here since offload is unusable on GB10 anyway.
```
+5
View File
@@ -19,6 +19,11 @@ add files for this repo. Use one Markdown file per model family by default.
> Individual pages resolve at `https://recipes.vllm.ai/<provider>/<model-id>`,
> e.g. [`bosonai/higgs-audio-v3-tts-4b`](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b).
The `Recipe` column in the Supported Models documentation links to the
published recipe when available and otherwise to the corresponding recipe in
this directory. The recipe files remain the source of truth for commands and
validation details.
Example layout:
```text
@@ -9,6 +9,10 @@ This script tests two main scenarios:
2. Case 2: Comparing FlashAttention and SDPA backends for batch_size=2 with padding
"""
import sys
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
@@ -371,6 +375,264 @@ def test_packed_varlen_metadata_must_be_complete(monkeypatch):
impl.forward_cuda(query, query, query, metadata)
# ---------------------------------------------------------------------------
# NPU mask-free packed attention (PR #5891).
#
# These tests cover the backend-side NPU branches that ``forward_fa_npu`` adds:
# - boundary resolution in ``_resolve_packed_seq_npu`` (pure Python, no device
# sync, no MindIE-SD import);
# - env dispatch in ``forward_fa_npu`` (which op a packed forward takes);
# - the laser input pre-scaling math in ``_forward_prefix_kv_slice_npu``.
#
# They run on CPU without a real NPU by injecting a fake ``mindiesd`` module
# into ``sys.modules`` (same pattern as ``test_paged_attention`` /
# ``benchmarks/conftest``). No production code is exercised for real kernels.
def _npu_impl(causal: bool = False) -> FlashAttentionImpl:
return FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=causal)
def _fake_mindiesd(monkeypatch, *, attention_forward=None, attention_forward_varlen=None):
"""Install a stub ``mindiesd`` so local ``from mindiesd import ...`` works.
Restored automatically by monkeypatch. Defaults are no-op callables so the
top-level import in ``forward_fa_npu`` succeeds even when the test mocks the
instance methods that would otherwise call them.
"""
fake = SimpleNamespace(
attention_forward=attention_forward or Mock(return_value=torch.zeros(1)),
attention_forward_varlen=attention_forward_varlen or Mock(return_value=torch.zeros(1)),
)
monkeypatch.setitem(sys.modules, "mindiesd", fake)
return fake
# --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
def test_resolve_packed_seq_accepts_real_plus_pad_contract():
impl = _npu_impl()
q = torch.randn(1, 8, 2, 4)
cu = torch.tensor([0, 5, 8], dtype=torch.int32)
extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 5, "max_seqlen_k": 5}
assert impl._resolve_packed_seq_npu(q, q, extra) == ([5, 8], [5, 8])
def test_resolve_packed_seq_single_document_no_padding():
impl = _npu_impl()
q = torch.randn(1, 8, 2, 4)
cu = torch.tensor([0, 8], dtype=torch.int32)
extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 8, "max_seqlen_k": 8}
assert impl._resolve_packed_seq_npu(q, q, extra) == ([8], [8])
@pytest.mark.parametrize(
"case",
[
"causal",
"missing_key",
"batch_ne_1",
"cu_too_long",
"cu_qk_mismatch",
"max_seqlen_not_int",
"used_gt_total",
"used_zero",
"pad_longer_than_real",
],
)
def test_resolve_packed_seq_rejects_malformed_metadata(case):
"""Any metadata outside the exact [real, pad] contract returns None so the
caller falls back to the masked path."""
cu3 = torch.tensor([0, 5, 8], dtype=torch.int32)
cu2 = torch.tensor([0, 8], dtype=torch.int32)
cu4 = torch.tensor([0, 3, 5, 8], dtype=torch.int32)
cu_short = torch.tensor([0, 3, 8], dtype=torch.int32)
base = {"cu_seqlens_q": cu3, "cu_seqlens_k": cu3, "max_seqlen_q": 5, "max_seqlen_k": 5}
q = torch.randn(1, 8, 2, 4)
impl = _npu_impl(causal=(case == "causal"))
if case == "missing_key":
extra = {"cu_seqlens_q": cu3, "max_seqlen_q": 5} # 3 of 4 keys absent
elif case == "batch_ne_1":
q = torch.randn(2, 8, 2, 4)
extra = base
elif case == "cu_too_long":
extra = {**base, "cu_seqlens_q": cu4, "cu_seqlens_k": cu4}
elif case == "cu_qk_mismatch":
extra = {**base, "cu_seqlens_q": cu3, "cu_seqlens_k": cu2}
elif case == "max_seqlen_not_int":
extra = {**base, "max_seqlen_q": 5.0}
elif case == "used_gt_total":
extra = {**base, "max_seqlen_q": 9} # 9 > total length 8
elif case == "used_zero":
extra = {**base, "max_seqlen_q": 0, "max_seqlen_k": 0}
elif case == "pad_longer_than_real":
# real=3, pad=5: the real document must be the longer one.
extra = {
"cu_seqlens_q": cu_short,
"cu_seqlens_k": cu_short,
"max_seqlen_q": 3,
"max_seqlen_k": 3,
}
else:
extra = base
assert impl._resolve_packed_seq_npu(q, q, extra) is None
# --- Test group B: env dispatch in forward_fa_npu (current behavior) --------
def _npu_impl_with_mocked_paths() -> tuple[FlashAttentionImpl, torch.Tensor]:
"""FlashAttentionImpl whose two packed paths are Mocks returning a sentinel.
Lets a test assert which branch ``forward_fa_npu`` dispatched to without
running any real kernel.
"""
impl = _npu_impl()
sentinel = torch.tensor([1.0])
impl._forward_varlen_packed_npu = Mock(return_value=sentinel)
impl._forward_prefix_kv_slice_npu = Mock(return_value=sentinel)
return impl, sentinel
def test_npu_env_unset_uses_varlen(monkeypatch):
monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
_fake_mindiesd(monkeypatch)
impl, _ = _npu_impl_with_mocked_paths()
metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
impl._forward_varlen_packed_npu.assert_called_once()
impl._forward_prefix_kv_slice_npu.assert_not_called()
def test_npu_env_laser_uses_slice(monkeypatch):
monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
_fake_mindiesd(monkeypatch)
impl, _ = _npu_impl_with_mocked_paths()
metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
impl._forward_prefix_kv_slice_npu.assert_called_once()
impl._forward_varlen_packed_npu.assert_not_called()
@pytest.mark.parametrize("fa_type", ["prompt_flash_attn", "fused_attn_score", "ascend_laser_attenton", "garbage"])
def test_npu_env_non_laser_value_falls_back_to_varlen(monkeypatch, fa_type):
"""Locks in the CURRENT dispatch behavior (PR #5891).
Any value other than ``ascend_laser_attention`` — including the valid
``prompt_flash_attn``/``fused_attn_score`` and typos — is silently routed to
the TND varlen op, which cannot honor ``MINDIE_SD_FA_TYPE``. This is review
comment #1 on the PR and is intentionally captured here as a regression
net: update this test when the dispatch starts honoring/rejecting these.
"""
monkeypatch.setenv("MINDIE_SD_FA_TYPE", fa_type)
_fake_mindiesd(monkeypatch)
impl, _ = _npu_impl_with_mocked_paths()
metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
impl._forward_varlen_packed_npu.assert_called_once()
impl._forward_prefix_kv_slice_npu.assert_not_called()
def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch):
"""Models that do not set ``npu_attn_varlen`` (e.g. Wan/Cosmos) keep using
the existing masked attention_forward path; the new branches are skipped."""
fake_forward = Mock(return_value=torch.zeros(1, 8, 2, 4))
_fake_mindiesd(monkeypatch, attention_forward=fake_forward)
impl, _ = _npu_impl_with_mocked_paths()
metadata = AttentionMetadata(
attn_mask=torch.tensor([[True, True, True, True, True, False, False, False]]),
extra={}, # no npu_attn_varlen opt-in
)
q = torch.randn(1, 8, 2, 4)
out = impl.forward_fa_npu(q, q, q, metadata)
impl._forward_varlen_packed_npu.assert_not_called()
impl._forward_prefix_kv_slice_npu.assert_not_called()
fake_forward.assert_called_once()
assert out is fake_forward.return_value
# --- Test group C: laser input pre-scaling in _forward_prefix_kv_slice_npu --
def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch):
monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
captured: dict = {}
def fake_attention_forward(query, key, value, **kwargs):
captured["query"] = query
captured["key"] = key
captured["value"] = value
captured.update(kwargs)
return torch.ones_like(query) # deterministic stand-in kernel output
_fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
impl = _npu_impl() # softmax_scale == 0.5
q = torch.randn(1, 8, 2, 4)
cu = torch.tensor([0, 5, 8], dtype=torch.int32)
extra = {
"cu_seqlens_q": cu,
"cu_seqlens_k": cu,
"max_seqlen_q": 5,
"max_seqlen_k": 5,
"laser_input_scale": 256.0,
}
out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
# K/V sliced to the valid prefix (used_k == 5); q keeps full length.
assert captured["key"].shape == (1, 5, 2, 4)
assert captured["value"].shape == (1, 5, 2, 4)
assert captured["query"].shape == (1, 8, 2, 4)
assert captured["attn_mask"] is None
# Pre-divide q/k/v by the factor and compensate the kernel scale by its
# square (exact for power-of-two); output scaled back by the factor.
torch.testing.assert_close(captured["query"], q / 256.0)
torch.testing.assert_close(captured["key"], q[:, :5] / 256.0)
assert captured["scale"] == pytest.approx(0.5 * 256.0**2)
torch.testing.assert_close(out, torch.ones(1, 8, 2, 4) * 256.0)
def test_prefix_kv_slice_no_scaling_without_factor(monkeypatch):
monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
captured: dict = {}
def fake_attention_forward(query, key, value, **kwargs):
captured["query"] = query
captured.update(kwargs)
return torch.ones_like(query)
_fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
impl = _npu_impl()
q = torch.randn(1, 8, 2, 4)
cu = torch.tensor([0, 5, 8], dtype=torch.int32)
extra = {
"cu_seqlens_q": cu,
"cu_seqlens_k": cu,
"max_seqlen_q": 5,
"max_seqlen_k": 5,
# no laser_input_scale → no pre-scaling
}
out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
torch.testing.assert_close(captured["query"], q)
assert captured["scale"] == pytest.approx(0.5)
torch.testing.assert_close(out, torch.ones(1, 8, 2, 4))
if __name__ == "__main__":
print("Running FlashAttention Padding Tests...")
print("=" * 60)
@@ -5,10 +5,9 @@
The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip:
1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row
kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the
first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
mask from the video and prefix segments separately, so the two only agree on the
block count when the video segment ends on a boundary.
kernel block, and 1344x768 gives 62x24x42 whose 62496 rows do not. Both are
handed to MindIE-SD: its rf_v2 preprocessing promotes an irregular video tail
to the always-kept prefix segment before generating the block mask.
"""
import dataclasses
@@ -57,10 +56,13 @@ def test_block_aligned_video_segment_runs_sparse():
assert plan.latent_shape == list(ALIGNED_GRID)
def test_misaligned_video_segment_stays_dense():
# Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask:
# pad keys would take a share of every real query's softmax denominator.
assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None
def test_misaligned_video_segment_runs_sparse_with_updated_mindiesd():
plan = make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID))
assert plan is not None
assert plan.prefix_len == PREFIX_ROWS
assert plan.used_len == PREFIX_ROWS + 62496
assert plan.latent_shape == list(MISALIGNED_GRID)
@pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
@@ -155,3 +157,28 @@ def test_fully_populated_mask_reproduces_dense_attention():
).transpose(1, 2)
error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.")
def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd():
"""vLLM forwards the full irregular grid; MindIE-SD owns its tail handling."""
torch.manual_seed(1)
prefix_len, grid = 384, (4, 16, 10)
video_len = grid[0] * grid[1] * grid[2]
used = prefix_len + video_len
heads, head_dim = 4, 128
q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
impl = make_impl()
impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
out = impl._forward_sparse_npu(q, k, v, plan)
reference = torch.nn.functional.scaled_dot_product_attention(
q.transpose(1, 2),
k.transpose(1, 2),
v.transpose(1, 2),
scale=impl.softmax_scale,
).transpose(1, 2)
error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
+8
View File
@@ -488,6 +488,14 @@ class _MiniMaxH3FakeAttention(nn.Module):
class FakeBackend:
supports_prefix_kv_slicing = False
# Mirror AttentionBackend.supports_packed_mask_free: H3's packed
# attention path calls attn_backend.supports_packed_mask_free() to decide
# whether to build the padding mask. Return False so the fake exercises
# the masked path (the abstract backend default).
@classmethod
def supports_packed_mask_free(cls) -> bool:
return False
def __init__(self, **kwargs):
del kwargs
super().__init__()
@@ -492,6 +492,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
class FakeBackend:
supports_prefix_kv_slicing = True
@classmethod
def supports_packed_mask_free(cls) -> bool:
return False
class FakeAttention(torch.nn.Module):
attn_backend = FakeBackend
@@ -526,6 +530,48 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
assert attention.attention.metadata.extra["valid_kv_length"] == 5
def test_packed_attention_skips_mask_for_packed_mask_free_backend():
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
MiniMaxH3Attention,
)
class FakeBackend:
supports_prefix_kv_slicing = False
@classmethod
def supports_packed_mask_free(cls) -> bool:
return True
class FakeAttention(torch.nn.Module):
attn_backend = FakeBackend
def __init__(self):
super().__init__()
self.metadata = None
def forward(self, query, key, value, metadata):
self.metadata = metadata
return query
attention = object.__new__(MiniMaxH3Attention)
torch.nn.Module.__init__(attention)
attention.attention = FakeAttention()
q = torch.randn(8, 2, 4)
attention._run_packed_attention(
q,
q,
q,
cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
max_seqlen=5,
packed_total=8,
)
assert attention.attention.metadata.attn_mask is None
assert attention.attention.metadata.extra["valid_kv_length"] == 5
assert attention.attention.metadata.extra["npu_attn_varlen"] is True
def test_packed_attention_keeps_padding_mask_for_other_backends():
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
MiniMaxH3Attention,
@@ -534,6 +580,10 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
class FakeBackend:
supports_prefix_kv_slicing = False
@classmethod
def supports_packed_mask_free(cls) -> bool:
return False
class FakeAttention(torch.nn.Module):
attn_backend = FakeBackend
@@ -69,6 +69,7 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
attention = object.__new__(MiniMaxH3Attention)
nn.Module.__init__(attention)
attention.rot_dim = 96
attention.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
attention.rope._forward_method = attention.rope.forward_native
@@ -113,6 +113,14 @@ class TestCreateDefaultDiffusion:
stages = StageConfigFactory.create_default_diffusion({"model": "x"})
assert stages[0]["engine_args"]["cache_backend"] == "none"
def test_vae_decode_precision_roundtrip(self):
od = _roundtrip_diffusion_config(model="x", vae_decode_precision="bf16")
assert od.vae_decode_precision == "bf16"
def test_invalid_vae_decode_precision_is_rejected(self):
with pytest.raises(ValueError, match="vae_decode_precision"):
OmniDiffusionConfig(model="x", vae_decode_precision="int8")
def test_single_gpu_default_devices(self):
stages = StageConfigFactory.create_default_diffusion({"model": "x"})
assert stages[0]["runtime"]["devices"] == "0"
+61 -1
View File
@@ -512,6 +512,35 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin
assert runner.pipeline.last_req.num_reqs == 1
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize(
("quality", "expected_enabled"),
[(None, False), ("lossless", False), ("high", True)],
)
def test_execute_model_scopes_vae_fast_path_to_request_quality(monkeypatch, quality, expected_enabled):
runner = _make_runner(cache_backend=None, cache_backend_name="none")
req = _make_request()
req.sampling_params.quality = quality
gate_calls = []
@contextmanager
def _record_gate(pipeline, enabled):
gate_calls.append((pipeline, enabled, "enter"))
yield
gate_calls.append((pipeline, enabled, "exit"))
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
monkeypatch.setattr(model_runner_module, "use_pipeline_vae_fast_path", _record_gate)
DiffusionModelRunner.execute_model(runner, req)
assert gate_calls == [
(runner.pipeline, expected_enabled, "enter"),
(runner.pipeline, expected_enabled, "exit"),
]
class _BatchPipeline:
"""Pipeline returning a configurable list of outputs from forward()."""
@@ -579,6 +608,18 @@ def test_execute_model_batch_rejects_output_count_mismatch(monkeypatch):
DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
@pytest.mark.core_model
@pytest.mark.cpu
def test_execute_model_batch_rejects_mixed_vae_quality() -> None:
runner = _make_batch_runner(_BatchPipeline(outputs=[DiffusionOutput(output="a"), DiffusionOutput(output="b")]))
sched = _make_scheduler_output(num_reqs=2)
sched.scheduled_new_reqs[0].req.sampling_params.quality = "lossless"
sched.scheduled_new_reqs[1].req.sampling_params.quality = "high"
with pytest.raises(RuntimeError, match="cannot mix request quality levels"):
DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
@pytest.mark.core_model
@pytest.mark.cpu
def test_execute_model_batch_routes_one_output_per_request(monkeypatch):
@@ -728,6 +769,19 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch):
@pytest.mark.core_model
@pytest.mark.cpu
def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
memory_pool_active = False
optimization_pool_states = []
@contextmanager
def _memory_pool_context(tag):
nonlocal memory_pool_active
assert tag == "weights"
memory_pool_active = True
try:
yield
finally:
memory_pool_active = False
class _DummyLoader:
def __init__(self, load_config, od_config=None):
del load_config, od_config
@@ -775,16 +829,22 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader)
monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
monkeypatch.setattr(
model_runner_module,
"optimize_pipeline_vaes",
lambda pipeline, od_config: optimization_pool_states.append(memory_pool_active),
)
monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
monkeypatch.setattr(
model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
)
DiffusionModelRunner.load_model(runner)
DiffusionModelRunner.load_model(runner, memory_pool_context_fn=_memory_pool_context)
assert runner.cache_backend is None
assert runner.od_config.cache_backend is None
assert dummy_cache_backend.enabled is False
assert optimization_pool_states == [True]
@pytest.mark.core_model
+2
View File
@@ -102,6 +102,7 @@ class TestDiffusionWorkerSleep:
def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
"""Test sleep mode level 1 (offload weights only)."""
clear_vae_caches = mocker.patch("vllm_omni.diffusion.vae_optimizations.clear_pipeline_vae_fast_path_caches")
mock_allocator_class = patch_cumem_allocator(mocker)
mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
@@ -128,6 +129,7 @@ class TestDiffusionWorkerSleep:
# Verify sleep was called with correct tags
mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
clear_vae_caches.assert_called_once_with(mock_gpu_worker.model_runner.pipeline)
assert bool(result) is True
# Verify buffers were NOT saved (level 1 doesn't save buffers)
assert len(mock_gpu_worker._sleep_saved_buffers) == 0
@@ -0,0 +1,300 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Correctness contracts for request-gated VAE decoder fast paths."""
from __future__ import annotations
import pytest
import torch
import torch.nn.functional as F
from torch import nn
from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.vae_optimizations.gate import (
VaeFastPathGate,
register_vae_fast_path_gate,
use_pipeline_vae_fast_path,
use_vae_fast_path,
)
pytestmark = [pytest.mark.diffusion]
class _PipelineWithVae(nn.Module):
def __init__(self, vae: nn.Module) -> None:
super().__init__()
self.vae = vae
class _DistributedExecutorConfig:
def __init__(self, parallel_mode: str) -> None:
self.parallel_mode = parallel_mode
@pytest.mark.core_model
@pytest.mark.cpu
def test_vae_fast_path_gate_restores_nested_state() -> None:
vae = nn.Identity()
gate = VaeFastPathGate()
register_vae_fast_path_gate(vae, gate)
with use_vae_fast_path(vae, True):
assert gate.enabled
with use_vae_fast_path(vae, False):
assert not gate.enabled
assert gate.enabled
assert not gate.enabled
@pytest.mark.core_model
@pytest.mark.cpu
def test_pipeline_vae_fast_path_controls_discovered_vae() -> None:
vae = nn.Identity()
pipeline = _PipelineWithVae(vae)
gate = VaeFastPathGate()
register_vae_fast_path_gate(vae, gate)
with use_pipeline_vae_fast_path(pipeline, True):
assert gate.enabled
assert not gate.enabled
@pytest.mark.core_model
@pytest.mark.cpu
def test_image_group_norm_fast_path_uses_compile_safe_fallback(monkeypatch) -> None:
from vllm_omni.diffusion.vae_optimizations import image
gate = VaeFastPathGate()
gate.enabled = True
norm = image.FusedGroupNormSiLU(nn.GroupNorm(2, 4), gate)
inputs = torch.randn(1, 4, 3, 3)
monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
monkeypatch.setattr(
image,
"group_norm_silu_4d",
lambda *args, **kwargs: pytest.fail("eager Triton kernel entered during compilation"),
)
actual = norm(inputs)
expected = F.silu(F.group_norm(inputs, 2, norm.weight, norm.bias, norm.eps))
torch.testing.assert_close(actual, expected)
def _small_image_vae() -> nn.Module:
from diffusers.models.autoencoders import AutoencoderKL
return AutoencoderKL(
in_channels=3,
out_channels=3,
down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D"),
up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D"),
block_out_channels=(32, 32),
layers_per_block=1,
latent_channels=4,
norm_num_groups=8,
sample_size=16,
)
def _small_wan_vae(vae_cls: type[nn.Module] | None = None) -> nn.Module:
if vae_cls is None:
from diffusers.models.autoencoders import AutoencoderKLWan
vae_cls = AutoencoderKLWan
return vae_cls(
base_dim=8,
decoder_base_dim=8,
z_dim=4,
dim_mult=[1, 2],
num_res_blocks=1,
attn_scales=[],
temperal_downsample=[False],
latents_mean=[0.0] * 4,
latents_std=[1.0] * 4,
scale_factor_temporal=1,
scale_factor_spatial=2,
)
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
def test_image_vae_install_preserves_lossless_path_and_state_dict() -> None:
from vllm_omni.diffusion.vae_optimizations import clear_pipeline_vae_fast_path_caches
from vllm_omni.diffusion.vae_optimizations.image import (
FusedGroupNormSiLU,
FusedUpsample2xConv2d,
maybe_optimize_image_vae,
)
torch.manual_seed(0)
vae = _small_image_vae().to(device="cuda", dtype=torch.bfloat16).eval()
latents = torch.randn(1, 4, 8, 8, device="cuda", dtype=torch.bfloat16)
parameter_names = {name for name, _ in vae.named_parameters()}
state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
reference = vae.decode(latents).sample
maybe_optimize_image_vae(vae)
maybe_optimize_image_vae(vae)
assert {name for name, _ in vae.named_parameters()} == parameter_names
vae.load_state_dict(state_dict, strict=True)
assert any(isinstance(module, FusedGroupNormSiLU) for module in vae.modules())
assert any(isinstance(module, FusedUpsample2xConv2d) for module in vae.modules())
assert torch.equal(vae.decode(latents).sample, reference)
with use_vae_fast_path(vae, True):
fast = vae.decode(latents).sample
torch.testing.assert_close(fast.float(), reference.float(), atol=0.1, rtol=0)
assert any(
buffer is not None
for module in vae.modules()
for name, buffer in module._buffers.items()
if name in {"_fused_weight", "_vllm_folded_value_weight", "_vllm_folded_value_bias"}
)
clear_pipeline_vae_fast_path_caches(_PipelineWithVae(vae))
assert all(
buffer is None
for module in vae.modules()
for name, buffer in module._buffers.items()
if name in {"_fused_weight", "_vllm_folded_value_weight", "_vllm_folded_value_bias"}
)
with use_vae_fast_path(vae, True):
torch.testing.assert_close(vae.decode(latents).sample.float(), reference.float(), atol=0.1, rtol=0)
assert torch.equal(vae.decode(latents).sample, reference)
assert vae.decoder.conv_in.weight.is_contiguous()
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
def test_channels_last_group_norm_silu_matches_reference(dtype: torch.dtype) -> None:
from vllm_omni.diffusion.vae_optimizations.triton_group_norm_silu import group_norm_silu_4d
torch.manual_seed(0)
x = torch.randn(1, 64, 16, 16, device="cuda", dtype=dtype).contiguous(memory_format=torch.channels_last)
weight = torch.randn(64, device="cuda", dtype=dtype)
bias = torch.randn(64, device="cuda", dtype=dtype)
actual = group_norm_silu_4d(x, weight, bias, num_groups=8, eps=1e-5)
expected = F.silu(F.group_norm(x, 8, weight, bias))
assert actual is not None
atol, rtol = (1e-5, 1e-5) if dtype == torch.float32 else (7e-2, 2e-2)
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
@pytest.mark.parametrize(
"activation_dtype,affine_dtype,atol,rtol",
[
(torch.float32, torch.float32, 1e-5, 1e-5),
(torch.bfloat16, torch.float32, 1.5e-1, 3e-2),
],
)
def test_wan_rmsnorm_silu_matches_reference(
activation_dtype: torch.dtype,
affine_dtype: torch.dtype,
atol: float,
rtol: float,
) -> None:
from vllm_omni.diffusion.vae_optimizations.triton_wan_rmsnorm_silu import wan_rmsnorm_silu
torch.manual_seed(0)
x = torch.randn(1, 96, 3, 10, 14, device="cuda", dtype=activation_dtype).contiguous(
memory_format=torch.channels_last_3d
)
gamma = torch.randn(96, 1, 1, 1, device="cuda", dtype=affine_dtype)
bias = torch.randn_like(gamma)
normalized = F.normalize(x.float() if activation_dtype != torch.float32 else x, dim=1).to(activation_dtype)
expected = F.silu(normalized * 96**0.5 * gamma + bias)
actual = wan_rmsnorm_silu(x, gamma, bias)
assert actual is not None
assert actual.stride() == x.stride()
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
def test_wan_vae_install_preserves_lossless_path_and_state_dict() -> None:
from vllm_omni.diffusion.vae_optimizations.wan import (
FusedWanRMSNormSiLU,
maybe_optimize_wan_vae,
)
torch.manual_seed(0)
vae = _small_wan_vae().cuda().eval()
latents = torch.randn(1, 4, 1, 8, 8, device="cuda")
parameter_names = {name for name, _ in vae.named_parameters()}
state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
reference = vae.decode(latents).sample
maybe_optimize_wan_vae(vae)
maybe_optimize_wan_vae(vae)
assert {name for name, _ in vae.named_parameters()} == parameter_names
vae.load_state_dict(state_dict, strict=True)
assert any(isinstance(module, FusedWanRMSNormSiLU) for module in vae.modules())
assert torch.equal(vae.decode(latents).sample, reference)
with use_vae_fast_path(vae, True):
fast = vae.decode(latents).sample
torch.testing.assert_close(fast, reference, atol=5e-4, rtol=1e-4)
assert torch.equal(vae.decode(latents).sample, reference)
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize("parallel_mode", ["auto", "spatial_shard_height", "spatial_shard_width"])
def test_wan_fast_path_skips_spatial_shard_mode(monkeypatch: pytest.MonkeyPatch, parallel_mode: str) -> None:
import vllm_omni.diffusion.vae_optimizations.wan as wan_optimizations
monkeypatch.setattr(wan_optimizations, "_HAS_TRITON", True)
vae = _small_wan_vae()
vae.distributed_executor = _DistributedExecutorConfig(parallel_mode)
wan_optimizations.maybe_optimize_wan_vae(vae)
assert not any(isinstance(module, wan_optimizations.FusedWanRMSNormSiLU) for module in vae.modules())
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.no_grad()
@pytest.mark.parametrize(
("precision", "decode_dtype"),
[("fp32", torch.float32), ("fp16", torch.float16), ("bf16", torch.bfloat16)],
)
def test_wan_decode_precision_keeps_encode_fp32(precision: str, decode_dtype: torch.dtype) -> None:
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
from vllm_omni.diffusion.vae_optimizations.precision import configure_wan_decode_precision
torch.manual_seed(0)
vae = _small_wan_vae(DistributedAutoencoderKLWan).to(device="cuda", dtype=torch.bfloat16).eval()
configure_wan_decode_precision(vae, precision)
assert next(vae.encoder.parameters()).dtype == torch.float32
assert next(vae.quant_conv.parameters()).dtype == torch.float32
assert next(vae.decoder.parameters()).dtype == decode_dtype
assert next(vae.post_quant_conv.parameters()).dtype == decode_dtype
_, grid_spec = vae.tile_split(torch.randn(1, 4, 1, 8, 8, device="cuda", dtype=decode_dtype))
assert grid_spec.output_dtype == decode_dtype
encoded = vae.encode(torch.randn(1, 3, 1, 16, 16, device="cuda", dtype=torch.bfloat16))
decoded = vae.decode(torch.randn(1, 4, 1, 8, 8, device="cuda", dtype=torch.float32))
assert encoded.latent_dist.parameters.dtype == torch.float32
assert decoded.sample.dtype == decode_dtype
+9
View File
@@ -3,6 +3,7 @@
import importlib.util
from pathlib import Path
from types import SimpleNamespace
import pytest
@@ -43,3 +44,11 @@ def test_model_display_names_require_string_mapping(tmp_path, monkeypatch):
with pytest.raises(ValueError, match="MODEL_DISPLAY_NAMES_FILE"):
generate_examples.load_model_display_names()
def test_only_general_serving_examples_get_generated_pages():
general_example = SimpleNamespace(category="offline_inference", path=Path("text_to_image"))
model_example = SimpleNamespace(category="offline_inference", path=Path("qwen3_omni"))
assert generate_examples.is_general_example(general_example)
assert not generate_examples.is_general_example(model_example)
+11
View File
@@ -76,6 +76,17 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
assert args.cfg_parallel_size == 4
def test_serve_parser_accepts_wan_decode_precision() -> None:
parser = TrackingArgumentParser()
subparsers = parser.add_subparsers(dest="subcommand")
OmniServeCommand().subparser_init(subparsers)
args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-decode-precision", "bf16"])
assert args.vae_decode_precision == "bf16"
assert args.get_explicit_kwargs_dict()["vae_decode_precision"] == "bf16"
def _make_headless_args(**kwargs) -> TrackingNamespace:
defaults = {
"model": "fake-model",
+22 -6
View File
@@ -123,18 +123,29 @@ def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]:
return found
def _check(label: str, path: Path, actual: int, budget: int, sample: list[str], budget_name: str) -> list[str]:
def _check(
label: str,
path: Path,
actual: int,
budget: int,
sample: list[str],
budget_name: str,
notices: list[str],
) -> list[str]:
if actual > budget:
listing = "\n".join(f" {line}" for line in sample)
return [
f"{path}: {actual} {label} exceeds the budget of {budget}.\n{listing}\n\n{_GUIDANCE}\n",
]
if actual < budget:
return [
f"{path}: {actual} {label}, down from a budget of {budget}. "
f"Lower {budget_name} to {actual} in {Path(__file__).relative_to(REPO_ROOT)} "
f"so the ratchet holds the new ground.",
]
# Progress is never a build failure. Removing branches is the goal, so a
# PR that does it must not hand-edit this file to stay green -- nor break
# `main` for everyone else when it merges from a stale base.
notices.append(
f"{path}: {actual} {label}, below the budget of {budget}. "
f"Lower {budget_name} to {actual} in "
f"{Path(__file__).relative_to(REPO_ROOT)} to hold the new ground."
)
return []
@@ -152,6 +163,7 @@ def main(argv: list[str] | None = None) -> int:
print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr)
return 1
notices: list[str] = []
branches = _model_type_branches(ast.parse(serving_path.read_text()))
detectors = _legacy_detectors(ast.parse(adapters_path.read_text()))
@@ -162,6 +174,7 @@ def main(argv: list[str] | None = None) -> int:
MAX_MODEL_TYPE_BRANCHES,
[f"line {lineno}: {src}" for lineno, src in branches],
"MAX_MODEL_TYPE_BRANCHES",
notices,
)
errors += _check(
"unmigrated model types in LEGACY_TTS_DETECTORS",
@@ -170,8 +183,11 @@ def main(argv: list[str] | None = None) -> int:
MAX_LEGACY_DETECTORS,
[f"line {lineno}: {name}" for lineno, name in detectors],
"MAX_LEGACY_DETECTORS",
notices,
)
for notice in notices:
print(f"check_tts_adapter: {notice}", file=sys.stderr)
for error in errors:
print(f"check_tts_adapter: {error}", file=sys.stderr)
return 1 if errors else 0
+1
View File
@@ -573,6 +573,7 @@ class _DiffusionConfigProjection:
fa_deterministic: bool = False
vae_use_slicing: bool = False
vae_use_tiling: bool = False
vae_decode_precision: Literal["fp32", "fp16", "bf16"] | None = None
mask_strategy_file_path: str | None = None
skip_time_steps: int = 15
VSA_sparsity: float = 0.0
+1
View File
@@ -401,6 +401,7 @@ class StageDeployConfig:
step_execution: bool | None = None
vae_use_slicing: bool | None = None
vae_use_tiling: bool | None = None
vae_decode_precision: str | None = None
boundary_ratio: float | None = None
flow_shift: float | None = None
diffusion_kv_cache_dtype: str | None = None
@@ -20,6 +20,17 @@ class AttentionBackend(ABC):
# avoid a slower masked-attention plan when tail padding is not semantic.
supports_prefix_kv_slicing: bool = False
@classmethod
def supports_packed_mask_free(cls) -> bool:
"""Whether packed attention never reads attn_mask on this platform.
When True, models that pack a [real, pad] two-document layout and
carry cu_seqlens/max_seqlen in ``AttentionMetadata.extra`` may skip
constructing the padding mask entirely. Backends whose mask-free
behavior is platform-dependent must check current_omni_platform.
"""
return False
# ``OmniPlatformEnum`` values this backend runs on; None means unrestricted.
# Platform resolution rejects an explicit selection outside this set, so a
# hardware-specific backend fails before the model is built.
@@ -121,6 +132,17 @@ class AttentionMetadata:
# the packed cu_seqlens tensors.
# "valid_kv_length": int — contiguous valid K/V prefix length for a
# backend that advertises supports_prefix_kv_slicing.
# "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path
# (TND npu_fusion_attention driven by cu_seqlens, mask never read).
# Requires the [real, pad] two-document packing contract; see
# FlashAttentionImpl._forward_varlen_packed_npu.
# "laser_input_scale": float — model opt-in input pre-scale for the NPU
# ascend_laser_attention path. The kernel stores unscaled QK^T in an
# fp16 workspace, so outlier activations overflow 65504 into NaN rows;
# with this set (>1), q/k/v are divided by the factor before the op,
# the kernel scale_value is multiplied by its square, and the output
# is scaled back (exact for power-of-two factors). Absent means no
# pre-scaling. See FlashAttentionImpl._forward_prefix_kv_slice_npu.
# Piecewise attention metadata (mixed causal/full masks).
# full_attn_spans: per-sample [start, end) spans in global coordinates using full attention.
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
from functools import partial
import torch
@@ -16,6 +17,7 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
piecewise_attn,
)
from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
from vllm_omni.platforms import current_omni_platform
logger = init_logger(__name__)
@@ -24,6 +26,15 @@ class FlashAttentionBackend(AttentionBackend):
accept_output_buffer: bool = True
supports_piecewise_spans: bool = True
@classmethod
def supports_packed_mask_free(cls) -> bool:
# CUDA: forward_cuda dispatches the packed cu_seqlens varlen path
# before ever reading attn_mask. NPU: forward_fa_npu honors the
# npu_attn_varlen opt-in and rebuilds the padding mask itself if the
# packed contract fails. XPU reads attn_mask, so models must keep
# constructing it there.
return current_omni_platform.is_cuda() or current_omni_platform.is_npu()
@classmethod
def supports_attention_mask(cls) -> bool:
return True
@@ -415,7 +426,36 @@ class FlashAttentionImpl(AttentionImpl):
"For installation details, see https://gitcode.com/Ascend/MindIE-SD"
"Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
)
# Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the
# model marks extra["npu_attn_varlen"] and carries packed metadata, so
# the padding document is excluded without reading or materializing
# the quadratic full_qk mask.
# - default: packed varlen via mindiesd attention_forward_varlen;
# - when MINDIE_SD_FA_TYPE=ascend_laser_attention: prefix K/V
# slicing via mindiesd attention_forward (so the env-selected op
# type is honored; the TND varlen op cannot be).
extra = attn_metadata.extra if attn_metadata else {}
if extra.get("npu_attn_varlen", False):
if os.environ.get("MINDIE_SD_FA_TYPE") == "ascend_laser_attention":
out = self._forward_prefix_kv_slice_npu(query, key, value, extra)
else:
out = self._forward_varlen_packed_npu(query, key, value, extra)
if out is not None:
return out
attention_mask = attn_metadata.attn_mask if attn_metadata else None
if attention_mask is None and extra.get("npu_attn_varlen", False):
# Models skip mask construction when this opt-in is set
# (FlashAttentionBackend.supports_packed_mask_free). The packed
# paths above declined (contract mismatch), so rebuild the padding
# mask here to keep the masked fallback correct.
used = extra.get("valid_kv_length")
if not isinstance(used, int) or not 0 < used <= query.shape[1]:
raise ValueError(
"npu_attn_varlen packed metadata is unusable and no attn_mask "
f"was constructed (valid_kv_length={used!r}, seq_len={query.shape[1]}); "
"refusing to run unmasked attention over padding rows."
)
attention_mask = torch.arange(query.shape[1], device=query.device)[None] < used
# NPU aclnnFlashAttentionScore requires mask shape to be one of:
# [B, N, Sq, Skv], [B, 1, Sq, Skv], [1, 1, Sq, Skv], or [Sq, Skv]
@@ -433,3 +473,161 @@ class FlashAttentionImpl(AttentionImpl):
op_type="fused_attn_score",
layout=layout,
)
def _resolve_packed_seq_npu(
self,
query: torch.Tensor,
key: torch.Tensor,
extra: dict,
) -> tuple[list[int], list[int]] | None:
"""Resolve packed document boundaries (cumulative end offsets per doc).
Returns (seq_q, seq_k), e.g. ([used_q, total_q], [used_k, total_k]),
or None unless the metadata matches the exact contract MiniMax-H3
produces:
- q/k/v are [1, T, N, D] (BSND, single packed batch);
- cu_seqlens describe a [real, pad] two-document packing, with the
padding document as a strict suffix;
- max_seqlen_q/k are Python ints equal to the real document length
(so boundaries are derived without any device sync).
"""
if self.causal:
return None
packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
if any(k not in extra for k in packed_keys):
return None
cu_q, cu_k = extra["cu_seqlens_q"], extra["cu_seqlens_k"]
used_q, used_k = extra["max_seqlen_q"], extra["max_seqlen_k"]
if query.shape[0] != 1 or key.shape[0] != 1:
return None
if not isinstance(used_q, int) or not isinstance(used_k, int):
return None
total_q, total_k = query.shape[1], key.shape[1]
# .shape is host-side metadata: counting documents never syncs.
if cu_q.shape[0] != cu_k.shape[0] or cu_q.shape[0] > 3:
return None
if not (0 < used_q <= total_q) or not (0 < used_k <= total_k):
return None
if used_q == total_q and used_k == total_k:
# No padding document: one full-length document.
return [total_q], [total_k]
if cu_q.shape[0] == 3 and used_q >= total_q - used_q and used_k >= total_k - used_k:
# [real, pad] packing: the real document must be the longer one
# (consistent with the max_seqlen naming).
return [used_q, total_q], [used_k, total_k]
return None
def _forward_varlen_packed_npu(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
extra: dict,
) -> torch.Tensor | None:
"""Packed varlen attention on NPU via mindiesd attention_forward_varlen.
Returns None (caller falls back to the mask path) when the packed
contract does not hold; see _resolve_packed_seq_npu.
"""
resolved = self._resolve_packed_seq_npu(query, key, extra)
if resolved is None:
return None
seq_q, seq_k = resolved
try:
from mindiesd import attention_forward_varlen
except ImportError:
raise ImportError(
"FlashAttentionBackend NPU implementation requires MindIE-SD. "
"Please install MindIE-SD to enable NPU attention support. "
"For installation details, see https://gitcode.com/Ascend/MindIE-SD"
"Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
)
q = query.squeeze(0) # [T, N, D] == TND
k = key.squeeze(0)
v = value.squeeze(0)
# attention_forward_varlen wants cu_seqlens as host lists of cumulative
# offsets and forwards cu[1:] as actual_seq_qlen/actual_seq_kvlen.
out = attention_forward_varlen(
q,
k,
v,
[0, *seq_q],
[0, *seq_k],
softmax_scale=self.softmax_scale,
)
return out.unsqueeze(0)
def _forward_prefix_kv_slice_npu(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
extra: dict,
) -> torch.Tensor | None:
"""Mask-free attention by slicing K/V to the valid prefix (zero-copy
views), then mindiesd attention_forward without a mask.
Same numerical contract as the varlen path: the padding document is a
strict suffix, so dropping it from K/V is identical to masking it out.
Query keeps full length; outputs on padding rows are never consumed
downstream. Returns None when the packed contract does not hold.
Used when MINDIE_SD_FA_TYPE=ascend_laser_attention: the TND varlen op
cannot honor the env-selected op type, while attention_forward can.
The laser kernel stores unscaled S=QK^T in an fp16 GM workspace, so
bf16 activations with large outliers overflow 65504 into ±inf and the
affected rows turn NaN. Models may opt into exact power-of-two input
pre-scaling via extra["laser_input_scale"] (see abstract.py): q/k/v
are divided by the factor, the kernel scale_value is multiplied by its
square, and the output is scaled back. Absent or invalid factor means
no pre-scaling.
"""
resolved = self._resolve_packed_seq_npu(query, key, extra)
if resolved is None:
return None
_, seq_k = resolved
used_k = seq_k[0] # real document length (first cumulative end)
try:
from mindiesd import attention_forward
except ImportError:
raise ImportError(
"FlashAttentionBackend NPU implementation requires MindIE-SD. "
"Please install MindIE-SD to enable NPU attention support. "
"For installation details, see https://gitcode.com/Ascend/MindIE-SD"
"Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
)
# mindiesd always takes BSND input regardless of `layout`; the arg
# selects the op-internal layout and laser only supports BNSD, so do
# NOT forward the model's qkv_layout ("BSND" for MiniMax-H3) here.
layout = "BNSD"
key = key[:, :used_k]
value = value[:, :used_k]
input_scale = extra.get("laser_input_scale")
preserve_input_range = isinstance(input_scale, (int, float)) and input_scale > 1
if preserve_input_range:
query = query / input_scale
key = key / input_scale
value = value / input_scale
scale = self.softmax_scale
if preserve_input_range:
scale *= input_scale**2
def _postprocess(out: torch.Tensor) -> torch.Tensor:
return out * input_scale if preserve_input_range else out
return _postprocess(
attention_forward(
query,
key,
value,
attn_mask=None,
opt_mode="manual",
op_type="fused_attn_score", # MINDIE_SD_FA_TYPE env overrides this
layout=layout,
scale=scale,
)
)
@@ -128,9 +128,9 @@ class RainFusionAttentionImpl(AttentionImpl):
whose extent the model publishes as ``AttentionMetadata.video_layout``. Every
other case — warmup denoise steps, exempt layers, a layer that does not declare
``qkv_layout="BSND"``, sequences without a published video segment, video
segments too short to pay for block selection or not aligned to the kernel's
block size — delegates to FlashAttention, so a model can select this backend
unconditionally.
segments too short to pay for block selection — delegates to FlashAttention,
so a model can select this backend unconditionally. MindIE-SD handles an
irregular video tail internally, retaining it outside the sparse blocks.
"""
def __init__(
@@ -292,25 +292,6 @@ class RainFusionAttentionImpl(AttentionImpl):
_MIN_VIDEO_BLOCKS,
)
return None
if video_len % _BLOCK_SIZE != 0:
# rf_v2 tiles the sequence into ceil((video + prefix) / 128) blocks but
# builds its block mask from the two segments pooled apart. The counts
# only agree when the video segment ends on a block boundary; otherwise
# a prefix block is forced on out of range, and the block straddling the
# seam holds real prefix rows that selection may drop. Padding the
# sequence to realign it is not an option because rf_v2 ignores
# attn_mask, so pad keys would take a share of every softmax denominator.
logger.warning_once(
"RAINFUSION_ATTN staying dense: latent grid %s gives %d video rows, which is not a "
"multiple of the %d-row kernel block. Choose a geometry whose latent_t * h * w is a "
"multiple of %d to enable sparsity.",
tuple(latent_shape),
video_len,
_BLOCK_SIZE,
_BLOCK_SIZE,
)
return None
logger.info_once(
"RAINFUSION_ATTN active: sparsity=%.2f, start_step=%d, exempt_layers=%d, "
"latent_grid=%s, prefix_rows=%d, video_rows=%d. Realized sparsity is lower than nominal "
+8
View File
@@ -711,6 +711,9 @@ class OmniDiffusionConfig:
# VAE memory optimization parameters
vae_use_slicing: bool = False
vae_use_tiling: bool = False
# Optional Wan decode-only precision. When set, the VAE encoder remains
# FP32 while decoder-side modules use the selected precision.
vae_decode_precision: str | None = None
# STA (Sliding Tile Attention) parameters
mask_strategy_file_path: str | None = None
@@ -923,6 +926,11 @@ class OmniDiffusionConfig:
)
def __post_init__(self):
if self.vae_decode_precision not in (None, "fp32", "fp16", "bf16"):
raise ValueError(
"vae_decode_precision must be one of {'fp32', 'fp16', 'bf16'} or None, "
f"got {self.vae_decode_precision!r}"
)
if self.diffusion_compile_granularity not in {"regional", "full"}:
raise ValueError(
"diffusion_compile_granularity must be 'regional' or 'full', "
@@ -106,7 +106,7 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
split_dims=(3, 4),
grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
tile_spec=tile_spec,
output_dtype=self.dtype,
output_dtype=getattr(self, "_vllm_decode_dtype", self.dtype),
)
return tiletask_list, grid_spec
+14
View File
@@ -68,6 +68,7 @@ class RMSNorm(CustomOp):
def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.float32) -> None:
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype))
self.hidden_size = hidden_size
self.variance_epsilon = eps
def _forward_fused(self, x: torch.Tensor) -> torch.Tensor:
@@ -107,6 +108,19 @@ class RMSNorm(CustomOp):
except Exception:
return self.forward_native(x)
def forward_musa(
self,
x: torch.Tensor,
) -> torch.Tensor:
# Preserve the aten::rms_norm graph so dynamic Inductor can fuse the
# H3 Q/K norm with its inline RoPE path on MUSA.
return F.rms_norm(
x,
(self.hidden_size,),
self.weight,
self.variance_epsilon,
)
def forward_npu(
self,
x: torch.Tensor,
+10 -1
View File
@@ -193,7 +193,16 @@ class RotaryEmbedding(CustomOp):
cos: torch.Tensor,
sin: torch.Tensor,
) -> torch.Tensor:
return self.forward_native(x, cos, sin)
# Keep the full-dimension NeoX form inline for MUSA compilation.
# Other layouts use the shared native implementation.
if self.half_head_dim or self.interleaved:
return self.forward_native(x, cos, sin)
if cos.dim() == 3:
cos, sin = cos[0], sin[0]
cos = cos.unsqueeze(-2)
sin = sin.unsqueeze(-2)
x1, x2 = x.chunk(2, dim=-1)
return x * cos + torch.cat((-x2, x1), dim=-1) * sin
def forward_native(
self,
@@ -107,6 +107,15 @@ MINIMAX_H3_FP32_BUFFER_NAMES = frozenset({"rope.inv_freq"})
# lookup and masked out afterwards).
MINIMAX_H3_ADALN_MODALITY_NUM = 3
# Opt-in fp16-range protection for the NPU ascend_laser_attention kernel
# (consumed only via the "laser_input_scale" extra key; other backends and
# platforms ignore it). The kernel stores unscaled QK^T in an fp16 GM
# workspace, and H3's outlier activations (per-element amax in the hundreds)
# push dot products past fp16 max 65504, turning whole 128-row blocks NaN.
# 256 is a power of two, so pre-dividing q/k/v and the compensating
# kernel-scale/output multiplies are exact in floating point.
MINIMAX_H3_LASER_INPUT_SCALE = 256.0
def _required_kwarg(kwargs: dict[str, Any], key: str) -> Any:
if key not in kwargs or kwargs[key] is None:
@@ -333,6 +342,7 @@ class MiniMaxH3Attention(nn.Module):
)
self.num_heads = self.qkv_proj.num_heads
self.num_kv_heads = self.qkv_proj.num_kv_heads
self.rot_dim = 6 * arch.rope_inv_freq_len
self.q_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
self.k_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
self.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
@@ -365,7 +375,7 @@ class MiniMaxH3Attention(nn.Module):
x: [T, heads, head_dim]; freqs: [T, rot_dim]. In the unfused path, cos/sin
are cast to the activation dtype before the elementwise math.
"""
rot_dim = freqs.shape[-1]
rot_dim = self.rot_dim
x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
cos = torch.cos(freqs).to(x.dtype) # [T, rot_dim]
sin = torch.sin(freqs).to(x.dtype)
@@ -401,11 +411,16 @@ class MiniMaxH3Attention(nn.Module):
used = min(max_seqlen, packed_total)
attn_mask = None
# Ring attention can dispatch to a different implementation from the
# configured backend, so this no-mask fast path is local-only.
prefix_slice = (
not getattr(self.attention, "use_ring", False) and self.attention.attn_backend.supports_prefix_kv_slicing
# configured backend, so the no-mask fast paths are local-only.
# supports_prefix_kv_slicing: backend slices K/V itself (cuDNN).
# supports_packed_mask_free: backend consumes the packed metadata
# without ever reading attn_mask (CUDA packed varlen, NPU
# npu_attn_varlen opt-in with its own fallback rebuild).
no_mask = not getattr(self.attention, "use_ring", False) and (
self.attention.attn_backend.supports_prefix_kv_slicing
or self.attention.attn_backend.supports_packed_mask_free()
)
if used < packed_total and not prefix_slice:
if used < packed_total and not no_mask:
attn_mask = torch.arange(packed_total, device=q.device)[None] < used
metadata = AttentionMetadata(
attn_mask=attn_mask,
@@ -415,6 +430,15 @@ class MiniMaxH3Attention(nn.Module):
"max_seqlen_q": max_seqlen,
"max_seqlen_k": max_seqlen,
"valid_kv_length": used,
# Opt the NPU flash backend into the packed varlen path so the
# quadratic full_qk mask is never materialized. Ring attention
# is excluded: it keeps the aligned padding rows for its
# fixed-size P2P buffers and still needs the mask.
"npu_attn_varlen": not getattr(self.attention, "use_ring", False),
# fp16-range protection for the ascend_laser_attention kernel
# (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other
# backend/path.
"laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE,
},
video_layout=video_layout,
)
@@ -0,0 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Optional VAE decoder optimizations."""
from vllm_omni.diffusion.vae_optimizations.gate import (
use_pipeline_vae_fast_path,
use_vae_fast_path,
)
from vllm_omni.diffusion.vae_optimizations.optimize import (
clear_pipeline_vae_fast_path_caches,
optimize_pipeline_vaes,
)
__all__ = [
"clear_pipeline_vae_fast_path_caches",
"optimize_pipeline_vaes",
"use_pipeline_vae_fast_path",
"use_vae_fast_path",
]
@@ -0,0 +1,67 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Decode-scoped gates for optional, non-bit-exact VAE fast paths."""
from __future__ import annotations
from collections.abc import Iterator
from contextlib import ExitStack, contextmanager
from weakref import WeakKeyDictionary
from torch import nn
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
class VaeFastPathGate:
"""Mutable flag shared by wrappers installed on one VAE instance."""
__slots__ = ("enabled",)
def __init__(self) -> None:
self.enabled = False
_VAE_FAST_PATH_GATES: WeakKeyDictionary[nn.Module, VaeFastPathGate] = WeakKeyDictionary()
def register_vae_fast_path_gate(vae: nn.Module, gate: VaeFastPathGate) -> None:
"""Associate an installed fast path with its owning VAE."""
_VAE_FAST_PATH_GATES[vae] = gate
@contextmanager
def use_vae_fast_path(vae: nn.Module, enabled: bool) -> Iterator[None]:
"""Enable an installed fast path for one decode and always restore it."""
gate = _VAE_FAST_PATH_GATES.get(vae)
if gate is None:
yield
return
previous_enabled = gate.enabled
gate.enabled = enabled
try:
yield
finally:
gate.enabled = previous_enabled
@contextmanager
def use_pipeline_vae_fast_path(pipeline: nn.Module, enabled: bool) -> Iterator[None]:
"""Apply one request's quality choice to every declared VAE component."""
vaes = ModuleDiscovery.discover(pipeline).vaes
with ExitStack() as stack:
for vae in vaes:
stack.enter_context(use_vae_fast_path(vae, enabled))
yield
__all__ = [
"VaeFastPathGate",
"register_vae_fast_path_gate",
"use_pipeline_vae_fast_path",
"use_vae_fast_path",
]
@@ -0,0 +1,354 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Quality-gated CUDA fast paths for Diffusers-style 2D KL VAE decoders."""
from __future__ import annotations
from types import MethodType
from typing import Any
import torch
import torch.nn.functional as F
from torch import nn
from vllm.logger import init_logger
from vllm_omni.diffusion.vae_optimizations.gate import (
VaeFastPathGate,
register_vae_fast_path_gate,
)
logger = init_logger(__name__)
try:
from vllm_omni.diffusion.vae_optimizations.triton_group_norm_silu import (
group_norm_silu_4d,
group_norm_silu_rows,
)
_HAS_TRITON = True
except ImportError: # pragma: no cover - exercised on builds without Triton
_HAS_TRITON = False
class FusedGroupNormSiLU(nn.Module):
"""Quality-gated GroupNorm + SiLU with an exact eager fallback."""
def __init__(self, norm: nn.GroupNorm, gate: VaeFastPathGate) -> None:
super().__init__()
self.num_groups = norm.num_groups
self.num_channels = norm.num_channels
self.eps = norm.eps
self.affine = norm.affine
# Register parameters directly so state-dict FQNs do not change.
self.weight = norm.weight
self.bias = norm.bias
self._vllm_gate = gate
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self._vllm_gate.enabled and x.dim() == 4 and not torch.compiler.is_compiling():
output = group_norm_silu_4d(
x,
self.weight,
self.bias,
self.num_groups,
self.eps,
apply_silu=True,
)
if output is not None:
return output
return F.silu(F.group_norm(x, self.num_groups, self.weight, self.bias, self.eps))
def _install_norm_silu(decoder: nn.Module, resnet_cls: type[nn.Module], gate: VaeFastPathGate) -> int:
def is_fusable(norm: object) -> bool:
return type(norm) is nn.GroupNorm and norm.affine and norm.weight is not None and norm.bias is not None
count = 0
for module in decoder.modules():
if (
type(module) is resnet_cls
and module.time_emb_proj is None
and module.time_embedding_norm in ("default", "group")
and module.upsample is None
and module.downsample is None
and isinstance(module.nonlinearity, nn.SiLU)
and is_fusable(module.norm1)
and is_fusable(module.norm2)
):
module.norm1 = FusedGroupNormSiLU(module.norm1, gate)
module.norm2 = FusedGroupNormSiLU(module.norm2, gate)
module.nonlinearity = nn.Identity()
count += 2
if (
type(getattr(decoder, "conv_norm_out", None)) is nn.GroupNorm
and isinstance(getattr(decoder, "conv_act", None), nn.SiLU)
and is_fusable(decoder.conv_norm_out)
):
decoder.conv_norm_out = FusedGroupNormSiLU(decoder.conv_norm_out, gate)
decoder.conv_act = nn.Identity()
count += 1
return count
_UPSAMPLE_TAP_MAP = {0: (2,), 1: (1, 2), 2: (0, 1), 3: (0,)}
def _fold_upsample2x_conv2d_weight(conv: nn.Conv2d) -> torch.Tensor:
weight = conv.weight.detach().float()
output_channels, input_channels = weight.shape[:2]
folded = weight.new_zeros(input_channels, output_channels, 4, 4)
for row in range(4):
for column in range(4):
value = weight.new_zeros(output_channels, input_channels)
for source_row in _UPSAMPLE_TAP_MAP[row]:
for source_column in _UPSAMPLE_TAP_MAP[column]:
value += weight[:, :, source_row, source_column]
folded[:, :, row, column] = value.t()
folded = folded.to(conv.weight.dtype)
if conv.weight.is_contiguous(memory_format=torch.channels_last):
folded = folded.contiguous(memory_format=torch.channels_last)
return folded.to(conv.weight.device)
class FusedUpsample2xConv2d(nn.Module):
"""Quality-gated nearest-2x + Conv2d replacement using ConvTranspose2d."""
def __init__(self, upsample: nn.Module, gate: VaeFastPathGate) -> None:
super().__init__()
# Keep only ``conv`` registered so checkpoint parameter names stay
# ``...upsamplers.N.conv.*``. The original module is an eager fallback.
object.__setattr__(self, "_original", upsample)
self.conv = upsample.conv
self.channels = upsample.channels
self._vllm_gate = gate
self.register_buffer("_fused_weight", None, persistent=False)
def forward(
self,
hidden_states: torch.Tensor,
output_size: int | None = None,
*args: Any,
**kwargs: Any,
) -> torch.Tensor:
if not self._vllm_gate.enabled or output_size is not None or hidden_states.shape[1] != self.channels:
return self._original(hidden_states, output_size=output_size, *args, **kwargs)
weight = self._fused_weight
if weight is None or weight.device != self.conv.weight.device or weight.dtype != self.conv.weight.dtype:
weight = _fold_upsample2x_conv2d_weight(self.conv)
self._fused_weight = weight
return F.conv_transpose2d(hidden_states, weight, self.conv.bias, stride=2, padding=1)
def _install_fused_upsample(decoder: nn.Module, upsample_cls: type[nn.Module], gate: VaeFastPathGate) -> int:
count = 0
for block in decoder.up_blocks:
upsamplers = getattr(block, "upsamplers", None)
if not upsamplers:
continue
for index, upsample in enumerate(upsamplers):
if type(upsample) is not upsample_cls:
continue
conv = getattr(upsample, "conv", None)
if (
upsample.use_conv
and not upsample.use_conv_transpose
and upsample.interpolate
and upsample.norm is None
and upsample.name == "conv"
and type(conv) is nn.Conv2d
and conv.kernel_size == (3, 3)
and conv.stride == (1, 1)
and conv.padding == (1, 1)
and conv.dilation == (1, 1)
and conv.groups == 1
and conv.padding_mode == "zeros"
):
upsamplers[index] = FusedUpsample2xConv2d(upsample, gate)
count += 1
return count
def _fold_attention_value_projection(module: nn.Module) -> tuple[torch.Tensor, torch.Tensor]:
value_weight = module.to_v.weight.detach().float()
value_bias = module.to_v.bias.detach().float()
output_weight = module.to_out[0].weight.detach().float()
output_bias = module.to_out[0].bias.detach().float()
dtype = module.to_v.weight.dtype
return (
(output_weight @ value_weight).to(dtype).contiguous(),
(output_weight @ value_bias + output_bias).to(dtype),
)
def _attention_fast_forward(
self: nn.Module,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
temb: torch.Tensor | None = None,
**cross_attention_kwargs: Any,
) -> torch.Tensor:
if (
not self._vllm_gate.enabled
or encoder_hidden_states is not None
or attention_mask is not None
or temb is not None
or hidden_states.ndim != 4
):
return type(self).forward(
self,
hidden_states,
encoder_hidden_states=encoder_hidden_states,
attention_mask=attention_mask,
temb=temb,
**cross_attention_kwargs,
)
residual = hidden_states
batch_size, channels, height, width = hidden_states.shape
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels)
if self.group_norm is not None:
norm = self.group_norm
normalized = None
if not torch.compiler.is_compiling():
normalized = group_norm_silu_rows(
hidden_states,
norm.weight,
norm.bias,
norm.num_groups,
norm.eps,
apply_silu=False,
)
hidden_states = normalized if normalized is not None else norm(hidden_states.transpose(1, 2)).transpose(1, 2)
query = self.to_q(hidden_states)
key = self.to_k(hidden_states)
folded_weight = self._vllm_folded_value_weight
folded_bias = self._vllm_folded_value_bias
if (
folded_weight is None
or folded_bias is None
or folded_weight.device != hidden_states.device
or folded_weight.dtype != self.to_v.weight.dtype
):
folded_weight, folded_bias = _fold_attention_value_projection(self)
self._vllm_folded_value_weight = folded_weight
self._vllm_folded_value_bias = folded_bias
value = F.linear(hidden_states, folded_weight, folded_bias)
output = F.scaled_dot_product_attention(query.unsqueeze(1), key.unsqueeze(1), value.unsqueeze(1))
output = output.squeeze(1).to(query.dtype)
output = self.to_out[1](output)
output = output.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
if self.residual_connection:
output = output + residual
return output / self.rescale_output_factor
def _attention_fast_compatible(module: nn.Module, attention_cls: type, processor_cls: type) -> bool:
return (
type(module) is attention_cls
and isinstance(module.processor, processor_cls)
and module.heads == 1
and module.scale_qk
and module.spatial_norm is None
and not module.norm_cross
and getattr(module, "norm_q", None) is None
and getattr(module, "norm_k", None) is None
and getattr(module, "add_k_proj", None) is None
and type(module.to_q) is nn.Linear
and type(module.to_k) is nn.Linear
and type(module.to_v) is nn.Linear
and type(module.to_out[0]) is nn.Linear
and isinstance(module.to_out[1], nn.Dropout)
and module.to_v.bias is not None
and module.to_out[0].bias is not None
)
def _decoder_layout_forward(self: nn.Module, *args: Any, **kwargs: Any) -> torch.Tensor:
use_channels_last = self._vllm_gate.enabled
if use_channels_last != self._vllm_channels_last:
memory_format = torch.channels_last if use_channels_last else torch.contiguous_format
self.to(memory_format=memory_format)
self._vllm_channels_last = use_channels_last
logger.debug(
"%s decoder switched to %s layout",
self._vllm_label,
"channels_last" if use_channels_last else "contiguous",
)
return type(self).forward(self, *args, **kwargs)
def maybe_optimize_image_vae(vae: nn.Module) -> nn.Module:
"""Install image VAE decoder fast paths when the structure is fully supported."""
if not _HAS_TRITON or getattr(vae, "_vllm_vae_fast_path_installed", False):
return vae
from diffusers.models.attention_processor import Attention, AttnProcessor2_0
from diffusers.models.autoencoders import AutoencoderKL
from diffusers.models.autoencoders.vae import Decoder
from diffusers.models.resnet import ResnetBlock2D
from diffusers.models.upsampling import Upsample2D
supported_types: tuple[type[nn.Module], ...] = (AutoencoderKL,)
try:
from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
supported_types += (AutoencoderKLFlux2,)
except ImportError: # pragma: no cover - older supported Diffusers builds
pass
if not isinstance(vae, supported_types):
return vae
decoder = getattr(vae, "decoder", None)
if type(decoder) is not Decoder:
return vae
attention_modules = [
module for module in decoder.modules() if _attention_fast_compatible(module, Attention, AttnProcessor2_0)
]
attention_count = sum(isinstance(module, Attention) for module in decoder.modules())
if len(attention_modules) != attention_count:
logger.info(
"%s has %d/%d attention blocks without a layout-safe rewrite; skipping VAE fast paths",
type(vae).__name__,
attention_count - len(attention_modules),
attention_count,
)
return vae
gate = VaeFastPathGate()
decoder._vllm_gate = gate
decoder._vllm_label = type(vae).__name__
decoder._vllm_channels_last = False
decoder.forward = MethodType(_decoder_layout_forward, decoder)
upsample_count = _install_fused_upsample(decoder, Upsample2D, gate)
for module in attention_modules:
module._vllm_gate = gate
module.register_buffer("_vllm_folded_value_weight", None, persistent=False)
module.register_buffer("_vllm_folded_value_bias", None, persistent=False)
module.forward = MethodType(_attention_fast_forward, module)
norm_count = _install_norm_silu(decoder, ResnetBlock2D, gate)
register_vae_fast_path_gate(vae, gate)
vae._vllm_vae_fast_path_installed = True
logger.info(
"%s installed quality-gated VAE decoder fast paths "
"(%d upsamplers, %d attention blocks, %d GroupNorm+SiLU fusions)",
type(vae).__name__,
upsample_count,
len(attention_modules),
norm_count,
)
return vae
__all__ = [
"FusedGroupNormSiLU",
"FusedUpsample2xConv2d",
"maybe_optimize_image_vae",
]
@@ -0,0 +1,70 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Pipeline-level installation for optional VAE optimizations."""
from __future__ import annotations
from torch import nn
from vllm.logger import init_logger
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
from vllm_omni.diffusion.vae_optimizations.image import maybe_optimize_image_vae
from vllm_omni.diffusion.vae_optimizations.precision import configure_wan_decode_precision
from vllm_omni.diffusion.vae_optimizations.wan import maybe_optimize_wan_vae
from vllm_omni.platforms import current_omni_platform
logger = init_logger(__name__)
_FOLDED_CACHE_NAMES = (
"_fused_weight",
"_vllm_folded_value_weight",
"_vllm_folded_value_bias",
)
def optimize_pipeline_vaes(pipeline: nn.Module, od_config: OmniDiffusionConfig) -> None:
"""Configure declared VAE components once after checkpoint loading."""
precision = getattr(od_config, "vae_decode_precision", None)
vaes = ModuleDiscovery.discover(pipeline).vaes
for vae in vaes:
# An explicit precision override is a serving contract, so failures on
# a recognized Wan VAE must abort startup instead of being hidden.
configure_wan_decode_precision(vae, precision)
if not current_omni_platform.is_cuda():
return
for vae in vaes:
for optimizer in (maybe_optimize_image_vae, maybe_optimize_wan_vae):
try:
optimizer(vae)
except Exception:
# A wrapper may already have installed its exact fallback
# before a later optional rewrite failed. With no registered
# enabled gate, those wrappers remain on the reference path.
logger.warning(
"Failed to install optional VAE fast paths on %s; using the reference path",
type(vae).__name__,
exc_info=True,
)
def clear_pipeline_vae_fast_path_caches(pipeline: nn.Module) -> None:
"""Release lazily folded buffers before the worker enters sleep mode."""
cleared = 0
for vae in ModuleDiscovery.discover(pipeline).vaes:
if not getattr(vae, "_vllm_vae_fast_path_installed", False):
continue
for module in vae.modules():
for name in _FOLDED_CACHE_NAMES:
if name in module._buffers and module._buffers[name] is not None:
setattr(module, name, None)
cleared += 1
if cleared:
logger.debug("Released %d lazy VAE fast-path buffers before sleep", cleared)
__all__ = ["clear_pipeline_vae_fast_path_caches", "optimize_pipeline_vaes"]
@@ -0,0 +1,91 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Decode-only precision configuration for structurally split Wan VAEs."""
from __future__ import annotations
from types import MethodType
from typing import Any
import torch
from torch import nn
from vllm.logger import init_logger
logger = init_logger(__name__)
PRECISION_TO_DTYPE = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}
def _decode_with_precision(
self: nn.Module,
latents: torch.Tensor,
*args: Any,
**kwargs: Any,
) -> Any:
return type(self).decode(
self,
latents.to(dtype=self._vllm_decode_dtype),
*args,
**kwargs,
)
def _encode_in_fp32(
self: nn.Module,
samples: torch.Tensor,
*args: Any,
**kwargs: Any,
) -> Any:
return type(self).encode(
self,
samples.to(dtype=torch.float32),
*args,
**kwargs,
)
def configure_wan_decode_precision(vae: nn.Module, precision: str | None) -> nn.Module:
"""Keep Wan encode in FP32 while materializing only decode modules lower."""
if precision is None:
return vae
from diffusers.models.autoencoders import AutoencoderKLWan
if not isinstance(vae, AutoencoderKLWan):
return vae
try:
decode_dtype = PRECISION_TO_DTYPE[precision]
except KeyError as exc:
raise ValueError(
f"Unsupported vae_decode_precision={precision!r}; expected one of {sorted(PRECISION_TO_DTYPE)}"
) from exc
encoder = getattr(vae, "encoder", None)
quant_conv = getattr(vae, "quant_conv", None)
decoder = getattr(vae, "decoder", None)
post_quant_conv = getattr(vae, "post_quant_conv", None)
if not all(isinstance(module, nn.Module) for module in (encoder, decoder, post_quant_conv)):
logger.info("Wan VAE lacks separate encoder/decoder modules; skipping decode-only precision")
return vae
encoder.to(dtype=torch.float32)
if isinstance(quant_conv, nn.Module):
quant_conv.to(dtype=torch.float32)
decoder.to(dtype=decode_dtype)
post_quant_conv.to(dtype=decode_dtype)
vae.encode = MethodType(_encode_in_fp32, vae)
vae._vllm_decode_dtype = decode_dtype
vae.decode = MethodType(_decode_with_precision, vae)
logger.info(
"Wan VAE configured with FP32 encode and %s decode",
precision,
)
return vae
__all__ = ["PRECISION_TO_DTYPE", "configure_wan_decode_precision"]
@@ -0,0 +1,237 @@
# ruff: noqa: N803
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Channels-last two-pass GroupNorm(+SiLU) Triton kernels.
The kernels keep the channel dimension innermost throughout a 2D VAE
decoder. Statistics are accumulated in FP32, then a separate apply pass
performs the affine transform and optional SiLU epilogue. Unsupported inputs
return ``None`` so callers can fail closed to the reference implementation.
"""
from __future__ import annotations
import torch
import triton
import triton.language as tl
_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
_MAX_CHANNELS = 2048
@triton.jit
def _gn_partial_rows_kernel(
x_ptr,
psum_ptr,
psq_ptr,
rows,
rows_per_prog,
C: tl.constexpr,
BLOCK_R: tl.constexpr,
):
chunk = tl.program_id(0).to(tl.int64)
batch = tl.program_id(1).to(tl.int64)
num_chunks = tl.num_programs(0)
row0 = chunk * rows_per_prog
cols = tl.arange(0, C)
acc_sum = tl.zeros((C,), tl.float32)
acc_sq = tl.zeros((C,), tl.float32)
x_base = x_ptr + batch * rows * C
for row_offset in range(0, rows_per_prog, BLOCK_R):
row_indices = row0 + row_offset + tl.arange(0, BLOCK_R)
mask = row_indices < rows
x = tl.load(
x_base + row_indices[:, None] * C + cols[None, :],
mask=mask[:, None],
other=0.0,
).to(tl.float32)
acc_sum += tl.sum(x, 0)
acc_sq += tl.sum(x * x, 0)
output_offsets = (batch * num_chunks + chunk) * C + cols
tl.store(psum_ptr + output_offsets, acc_sum)
tl.store(psq_ptr + output_offsets, acc_sq)
@triton.jit
def _gn_finalize_kernel(
psum_ptr,
psq_ptr,
weight_ptr,
bias_ptr,
scale_shift_ptr,
num_chunks,
group_numel,
eps,
channels,
CPG: tl.constexpr,
BLOCK_K: tl.constexpr,
):
group = tl.program_id(0).to(tl.int64)
batch = tl.program_id(1).to(tl.int64)
cols = group * CPG + tl.arange(0, CPG)
acc_sum = tl.zeros((), tl.float32)
acc_sq = tl.zeros((), tl.float32)
for chunk0 in range(0, num_chunks, BLOCK_K):
chunks = chunk0 + tl.arange(0, BLOCK_K)
mask = chunks < num_chunks
offsets = (batch * num_chunks + chunks)[:, None] * channels + cols[None, :]
acc_sum += tl.sum(tl.load(psum_ptr + offsets, mask=mask[:, None], other=0.0))
acc_sq += tl.sum(tl.load(psq_ptr + offsets, mask=mask[:, None], other=0.0))
mean = acc_sum / group_numel
variance = acc_sq / group_numel - mean * mean
variance = tl.maximum(variance, 0.0)
reciprocal_std = tl.rsqrt(variance + eps)
weight = tl.load(weight_ptr + cols).to(tl.float32)
bias = tl.load(bias_ptr + cols).to(tl.float32)
scale = weight * reciprocal_std
shift = bias - mean * scale
tl.store(scale_shift_ptr + batch * 2 * channels + cols, scale)
tl.store(scale_shift_ptr + (batch * 2 + 1) * channels + cols, shift)
@triton.jit
def _gn_apply_rows_kernel(
x_ptr,
scale_shift_ptr,
y_ptr,
rows,
C: tl.constexpr,
BLOCK_R: tl.constexpr,
SILU: tl.constexpr,
):
program = tl.program_id(0).to(tl.int64)
batch = tl.program_id(1).to(tl.int64)
cols = tl.arange(0, C)
scale = tl.load(scale_shift_ptr + batch * 2 * C + cols)
shift = tl.load(scale_shift_ptr + (batch * 2 + 1) * C + cols)
row_indices = program * BLOCK_R + tl.arange(0, BLOCK_R)
mask = row_indices < rows
offsets = batch * rows * C + row_indices[:, None] * C + cols[None, :]
x = tl.load(x_ptr + offsets, mask=mask[:, None], other=0.0).to(tl.float32)
y = x * scale[None, :] + shift[None, :]
if SILU:
y = y * tl.sigmoid(y)
tl.store(y_ptr + offsets, y.to(y_ptr.dtype.element_ty), mask=mask[:, None])
def _group_norm_silu_rows(
x: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor,
num_groups: int,
eps: float,
apply_silu: bool,
) -> torch.Tensor:
batch_size, rows, channels = x.shape
channels_per_group = channels // num_groups
block_rows = max(1, 8192 // channels)
rows_per_program = block_rows * 32
num_chunks = triton.cdiv(rows, rows_per_program)
partial_sum = torch.empty((batch_size, num_chunks, channels), device=x.device, dtype=torch.float32)
partial_sq = torch.empty_like(partial_sum)
with torch.cuda.device(x.device):
_gn_partial_rows_kernel[(num_chunks, batch_size)](
x,
partial_sum,
partial_sq,
rows,
rows_per_program,
C=channels,
BLOCK_R=block_rows,
num_warps=4,
)
scale_shift = torch.empty((batch_size, 2, channels), device=x.device, dtype=torch.float32)
block_chunks = max(1, min(4096 // channels_per_group, triton.next_power_of_2(num_chunks)))
_gn_finalize_kernel[(num_groups, batch_size)](
partial_sum,
partial_sq,
weight,
bias,
scale_shift,
num_chunks,
rows * channels_per_group,
eps,
channels,
CPG=channels_per_group,
BLOCK_K=block_chunks,
num_warps=4,
)
output = torch.empty_like(x)
_gn_apply_rows_kernel[(triton.cdiv(rows, block_rows), batch_size)](
x,
scale_shift,
output,
rows,
C=channels,
BLOCK_R=block_rows,
SILU=apply_silu,
num_warps=4,
)
return output
def _is_supported(
x: torch.Tensor,
weight: torch.Tensor | None,
bias: torch.Tensor | None,
num_groups: int,
) -> bool:
if not (x.is_cuda and x.numel() > 0 and not torch.is_grad_enabled()):
return False
if x.requires_grad or x.dtype not in _SUPPORTED_DTYPES:
return False
if weight is None or bias is None:
return False
channels = x.shape[1] if x.dim() == 4 else x.shape[-1]
if weight.shape != (channels,) or bias.shape != (channels,):
return False
if not (
weight.device == bias.device == x.device
and weight.dtype == bias.dtype == x.dtype
and weight.is_contiguous()
and bias.is_contiguous()
):
return False
if num_groups < 1 or channels % num_groups != 0:
return False
return triton.next_power_of_2(channels) == channels and channels <= _MAX_CHANNELS
def group_norm_silu_4d(
x: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor,
num_groups: int,
eps: float,
apply_silu: bool = True,
) -> torch.Tensor | None:
"""Run GroupNorm(+SiLU) on a channels-last ``(N, C, H, W)`` tensor."""
if x.dim() != 4 or not _is_supported(x, weight, bias, num_groups):
return None
batch_size, channels, height, width = x.shape
if not (channels > 1 and (height > 1 or width > 1) and x.is_contiguous(memory_format=torch.channels_last)):
return None
rows = x.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels)
output = _group_norm_silu_rows(rows, weight, bias, num_groups, eps, apply_silu)
return output.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
def group_norm_silu_rows(
x: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor,
num_groups: int,
eps: float,
apply_silu: bool = True,
) -> torch.Tensor | None:
"""Run GroupNorm(+SiLU) over contiguous ``(N, rows, C)`` input."""
if x.dim() != 3 or not x.is_contiguous():
return None
if not _is_supported(x, weight, bias, num_groups):
return None
return _group_norm_silu_rows(x, weight, bias, num_groups, eps, apply_silu)
__all__ = ["group_norm_silu_4d", "group_norm_silu_rows"]
@@ -0,0 +1,160 @@
# ruff: noqa: N803
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Channels-last-3D Wan VAE RMSNorm(+SiLU) Triton kernel."""
from __future__ import annotations
import torch
import triton
import triton.language as tl
_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
_MAX_CHANNELS = 1024
@triton.jit
def _wan_rmsnorm_silu_kernel(
x_ptr,
gamma_ptr,
bias_ptr,
output_ptr,
channels: tl.constexpr,
time_size,
height_size,
width_size,
x_stride_batch,
x_stride_channel,
x_stride_time,
x_stride_height,
x_stride_width,
output_stride_batch,
output_stride_channel,
output_stride_time,
output_stride_height,
output_stride_width,
rms_scale,
eps,
has_bias: tl.constexpr,
BLOCK_C: tl.constexpr,
):
row = tl.program_id(0).to(tl.int64)
channel_offsets = tl.arange(0, BLOCK_C)
channel_mask = channel_offsets < channels
width = row % width_size
remainder = row // width_size
height = remainder % height_size
remainder = remainder // height_size
time = remainder % time_size
batch = remainder // time_size
x_base = batch * x_stride_batch + time * x_stride_time + height * x_stride_height + width * x_stride_width
output_base = (
batch * output_stride_batch
+ time * output_stride_time
+ height * output_stride_height
+ width * output_stride_width
)
x = tl.load(
x_ptr + x_base + channel_offsets * x_stride_channel,
mask=channel_mask,
other=0.0,
).to(tl.float32)
norm = tl.sqrt(tl.sum(x * x, axis=0))
inverse_norm = 1.0 / tl.maximum(norm, eps)
# Match the eager dtype boundaries: normalize and scale in activation
# dtype, affine in the promoted dtype, then SiLU in FP32.
y = (x * inverse_norm).to(x_ptr.dtype.element_ty)
gamma = tl.load(gamma_ptr + channel_offsets, mask=channel_mask, other=1.0)
y = (y * rms_scale).to(x_ptr.dtype.element_ty)
y = (y.to(tl.float32) * gamma.to(tl.float32)).to(output_ptr.dtype.element_ty)
if has_bias:
bias = tl.load(bias_ptr + channel_offsets, mask=channel_mask, other=0.0)
y = (y.to(tl.float32) + bias.to(tl.float32)).to(output_ptr.dtype.element_ty)
y = y.to(tl.float32)
y = y * tl.sigmoid(y)
tl.store(
output_ptr + output_base + channel_offsets * output_stride_channel,
y,
mask=channel_mask,
)
def _affine_supported(x: torch.Tensor, affine: torch.Tensor) -> bool:
return (
affine.is_cuda
and affine.device == x.device
and (affine.dtype == x.dtype or affine.dtype == torch.float32)
and affine.numel() == x.shape[1]
)
def can_use_wan_rmsnorm_silu(
x: torch.Tensor,
gamma: torch.Tensor,
bias: torch.Tensor | None,
) -> bool:
"""Return whether ``wan_rmsnorm_silu`` supports this exact input."""
return (
x.is_cuda
and not torch.is_grad_enabled()
and not x.requires_grad
and x.dtype in _SUPPORTED_DTYPES
and x.ndim == 5
and x.numel() > 0
and 0 < x.shape[1] <= _MAX_CHANNELS
and x.is_contiguous(memory_format=torch.channels_last_3d)
and _affine_supported(x, gamma)
and (bias is None or _affine_supported(x, bias))
)
def wan_rmsnorm_silu(
x: torch.Tensor,
gamma: torch.Tensor,
bias: torch.Tensor | None = None,
rms_scale: float | None = None,
eps: float = 1e-12,
) -> torch.Tensor | None:
"""Fuse ``SiLU(F.normalize(x, dim=1) * scale * gamma + bias)``."""
if not can_use_wan_rmsnorm_silu(x, gamma, bias):
return None
batch_size, channels, time_size, height_size, width_size = x.shape
output_dtype = torch.promote_types(x.dtype, gamma.dtype)
output = torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=output_dtype)
block_channels = triton.next_power_of_2(channels)
num_warps = 1 if block_channels <= 64 else 4 if block_channels <= 512 else 8
gamma = gamma.reshape(channels).contiguous()
has_bias = bias is not None
bias = gamma if bias is None else bias.reshape(channels).contiguous()
if rms_scale is None:
rms_scale = channels**0.5
with torch.cuda.device(x.device):
_wan_rmsnorm_silu_kernel[(batch_size * time_size * height_size * width_size,)](
x,
gamma,
bias,
output,
channels,
time_size,
height_size,
width_size,
*x.stride(),
*output.stride(),
float(rms_scale),
eps,
has_bias,
BLOCK_C=block_channels,
num_warps=num_warps,
)
return output
__all__ = ["can_use_wan_rmsnorm_silu", "wan_rmsnorm_silu"]
@@ -0,0 +1,162 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Quality-gated CUDA fast path for Diffusers Wan VAE decoders."""
from __future__ import annotations
from types import MethodType
from typing import Any
import torch
import torch.nn.functional as F
from torch import nn
from vllm.logger import init_logger
from vllm_omni.diffusion.vae_optimizations.gate import (
VaeFastPathGate,
register_vae_fast_path_gate,
)
logger = init_logger(__name__)
try:
from vllm_omni.diffusion.vae_optimizations.triton_wan_rmsnorm_silu import (
wan_rmsnorm_silu,
)
_HAS_TRITON = True
except ImportError: # pragma: no cover - exercised on builds without Triton
_HAS_TRITON = False
class FusedWanRMSNormSiLU(nn.Module):
"""Quality-gated Wan RMSNorm + SiLU with an exact eager fallback."""
def __init__(self, norm: nn.Module, gate: VaeFastPathGate) -> None:
super().__init__()
self.gamma = norm.gamma
self.bias = norm.bias
self.scale = float(norm.scale)
self._vllm_gate = gate
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self._vllm_gate.enabled and not torch.compiler.is_compiling():
bias = self.bias if isinstance(self.bias, torch.Tensor) else None
output = wan_rmsnorm_silu(x, self.gamma, bias, rms_scale=self.scale)
if output is not None:
return output
needs_fp32_normalize = x.dtype in (torch.float16, torch.bfloat16) or any(
token in str(x.dtype) for token in ("float4_", "float8_")
)
normalized = F.normalize(x.float() if needs_fp32_normalize else x, dim=1).to(x.dtype)
return F.silu(normalized * self.scale * self.gamma + self.bias)
def _plain_silu(activation: object) -> bool:
return isinstance(activation, nn.SiLU) and not activation.inplace
def _fusable_norm(norm: object, norm_cls: type) -> bool:
return (
type(norm) is norm_cls
and getattr(norm, "channel_first", False)
and isinstance(getattr(norm, "gamma", None), torch.Tensor)
)
def _install_norm_silu(decoder: nn.Module, gate: VaeFastPathGate) -> int | None:
from diffusers.models.autoencoders.autoencoder_kl_wan import WanResidualBlock, WanRMS_norm
residual_blocks = [module for module in decoder.modules() if isinstance(module, WanResidualBlock)]
eligible = [
module
for module in residual_blocks
if type(module) is WanResidualBlock
and _plain_silu(module.nonlinearity)
and _fusable_norm(module.norm1, WanRMS_norm)
and _fusable_norm(module.norm2, WanRMS_norm)
]
if len(eligible) != len(residual_blocks):
logger.info(
"Wan VAE has %d/%d non-standard residual blocks; skipping fast path",
len(residual_blocks) - len(eligible),
len(residual_blocks),
)
return None
if not (
_fusable_norm(getattr(decoder, "norm_out", None), WanRMS_norm)
and _plain_silu(getattr(decoder, "nonlinearity", None))
):
logger.info("Wan VAE has a non-standard output head; skipping fast path")
return None
count = 0
for module in eligible:
module.norm1 = FusedWanRMSNormSiLU(module.norm1, gate)
module.norm2 = FusedWanRMSNormSiLU(module.norm2, gate)
module.nonlinearity = nn.Identity()
count += 2
decoder.norm_out = FusedWanRMSNormSiLU(decoder.norm_out, gate)
decoder.nonlinearity = nn.Identity()
return count + 1
def _set_wan_decoder_memory_format(decoder: nn.Module, *, channels_last: bool) -> None:
for module in decoder.modules():
if isinstance(module, nn.Conv3d):
module.to(memory_format=torch.channels_last_3d if channels_last else torch.contiguous_format)
elif isinstance(module, (nn.Conv2d, nn.ConvTranspose2d)):
module.to(memory_format=torch.channels_last if channels_last else torch.contiguous_format)
def _decoder_layout_forward(self: nn.Module, *args: Any, **kwargs: Any) -> torch.Tensor:
use_channels_last = self._vllm_gate.enabled
if use_channels_last != self._vllm_channels_last:
_set_wan_decoder_memory_format(self, channels_last=use_channels_last)
self._vllm_channels_last = use_channels_last
logger.debug(
"Wan VAE decoder switched to %s layout",
"channels_last_3d" if use_channels_last else "contiguous",
)
return type(self).forward(self, *args, **kwargs)
def maybe_optimize_wan_vae(vae: nn.Module) -> nn.Module:
"""Install the Wan decoder fast path when its full structure is supported."""
if not _HAS_TRITON or getattr(vae, "_vllm_vae_fast_path_installed", False):
return vae
from diffusers.models.autoencoders import AutoencoderKLWan
from diffusers.models.autoencoders.autoencoder_kl_wan import WanDecoder3d
if not isinstance(vae, AutoencoderKLWan):
return vae
decoder = getattr(vae, "decoder", None)
if type(decoder) is not WanDecoder3d:
return vae
distributed_executor = getattr(vae, "distributed_executor", None)
parallel_mode = getattr(distributed_executor, "parallel_mode", "tile")
if parallel_mode in ("auto", "spatial_shard_height", "spatial_shard_width"):
logger.info("Wan spatial-shard decode is configured; skipping quality-gated VAE fast path")
return vae
gate = VaeFastPathGate()
norm_count = _install_norm_silu(decoder, gate)
if norm_count is None:
return vae
decoder._vllm_gate = gate
decoder._vllm_channels_last = False
decoder.forward = MethodType(_decoder_layout_forward, decoder)
register_vae_fast_path_gate(vae, gate)
vae._vllm_vae_fast_path_installed = True
logger.info(
"Wan VAE installed quality-gated decoder fast path (%d RMSNorm+SiLU fusions)",
norm_count,
)
return vae
__all__ = ["FusedWanRMSNormSiLU", "maybe_optimize_wan_vae"]
@@ -43,6 +43,10 @@ from vllm_omni.diffusion.offloader import get_offload_backend
from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob
from vllm_omni.diffusion.vae_optimizations import (
optimize_pipeline_vaes,
use_pipeline_vae_fast_path,
)
from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.diffusion.worker.utils import (
@@ -243,6 +247,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
custom_pipeline_name=custom_pipeline_name,
device=self.device,
)
# Keep precision rematerialization in the tagged weights pool
# so sleep level 1 can offload every resulting parameter.
optimize_pipeline_vaes(self.pipeline, self.od_config)
time_after_load = time.perf_counter()
logger.info(
@@ -499,6 +506,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
if require_request_batch_support and not getattr(self.pipeline, "supports_request_batch", False):
raise RuntimeError(f"{type(self.pipeline).__name__} does not support request-batch forward.")
quality = getattr(reqs[0].sampling_params, "quality", None)
if any(getattr(req.sampling_params, "quality", None) != quality for req in reqs[1:]):
raise RuntimeError("A diffusion worker batch cannot mix request quality levels.")
use_fast_vae = quality == "high"
# Use no_grad() for HSDP compatibility, inference_mode() otherwise for
# better perf. HSDP2's fully_shard pre-forward hooks need tensor version
# counters, which inference tensors do not track.
@@ -523,7 +535,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=od_config):
with record_function(record_name):
raw_outputs = self.pipeline.forward(batch)
with use_pipeline_vae_fast_path(self.pipeline, use_fast_vae):
raw_outputs = self.pipeline.forward(batch)
outputs = _normalize_pipeline_outputs(
raw_outputs,
expected_count=len(reqs),
@@ -774,7 +787,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
if should_decode:
clear_pipeline_stage_durations(self.pipeline)
result = self.pipeline.post_decode(req)
with use_pipeline_vae_fast_path(
self.pipeline,
getattr(req.sampling, "quality", None) == "high",
):
result = self.pipeline.post_decode(req)
if result is not None:
self._attach_stepwise_metrics(
req,
@@ -576,6 +576,14 @@ class DiffusionWorker:
usage_before = allocator.get_current_usage()
if self.model_runner is not None and self.model_runner.pipeline is not None:
from vllm_omni.diffusion.vae_optimizations import clear_pipeline_vae_fast_path_caches
# Folded weights are allocated lazily on the request path rather
# than in the tagged weight pool. Drop them before sleep so they
# neither remain resident nor need a separate CPU backup.
clear_pipeline_vae_fast_path_caches(self.model_runner.pipeline)
if level == 2 and self.model_runner is not None:
if hasattr(self.model_runner, "graph_runners"):
self.model_runner.graph_runners.clear()
+1
View File
@@ -491,6 +491,7 @@ class OrchestratorArgs:
step_execution: bool = False
vae_use_slicing: bool = False
vae_use_tiling: bool = False
vae_decode_precision: str | None = None
enable_multithread_weight_load: bool = True
num_weight_load_threads: int = 4
enable_cpu_offload: bool = False
+1
View File
@@ -1026,6 +1026,7 @@ class AsyncOmniEngine:
"request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0),
"vae_use_slicing": kwargs.get("vae_use_slicing", False),
"vae_use_tiling": kwargs.get("vae_use_tiling", False),
"vae_decode_precision": kwargs.get("vae_decode_precision", None),
"cache_backend": cache_backend,
"cache_config": cache_config,
"enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
+9
View File
@@ -629,6 +629,15 @@ class OmniServeCommand(CLISubcommand):
action="store_true",
help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).",
)
omni_config_group.add_argument(
"--vae-decode-precision",
choices=("fp32", "fp16", "bf16"),
default=None,
help=(
"Optional Wan decode-only VAE precision. The encoder remains FP32; "
"unset preserves the pipeline's existing precision behavior."
),
)
# Parallel weight loading (faster diffusion startup)
omni_config_group.add_argument(
+12
View File
@@ -139,6 +139,18 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
)
backend_upper = "FLASH_ATTN"
if backend_upper == "FLASH_ATTN" and find_spec("mindiesd"):
# The NPU FLASH_ATTN backend imports mindiesd lazily at first
# forward, but CANN snapshots the custom-op registry at the
# first custom-op regInfo lookup in the process (e.g. a
# vllm-ascend custom op during model load/warmup). Import
# mindiesd here so its env.py prepends the mindiesd vendor
# dirs (aie_ascendc etc.) to ASCEND_CUSTOM_OPP_PATH before
# that snapshot; otherwise aclnnLaserAttention /
# FusedAttentionScore fail with EZ1001 "does not support
# opType" for the rest of the process.
import mindiesd # noqa: F401
backend = DiffusionAttentionBackendEnum[backend_upper]
logger.debug("Using diffusion attention backend '%s'", backend_upper)
return backend.get_path()