Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
094abd70c1 | ||
|
|
c27623c29c | ||
|
|
f8fa86fb1f | ||
|
|
422d8fd8f9 | ||
|
|
62e239bf1e | ||
|
|
740f45abcd | ||
|
|
5e4b9f8c6e | ||
|
|
ae811043e9 | ||
|
|
636d44e272 | ||
|
|
8e44e78038 | ||
|
|
b98b7c8529 |
No files matched your search
+110
-28
@@ -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
|
||||
@@ -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
@@ -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/*
|
||||
|
||||
@@ -7,6 +7,7 @@ owners:
|
||||
- "@yinpeiqi"
|
||||
- "@fake0fan"
|
||||
- "@Sy0307"
|
||||
- "@Gaohan123"
|
||||
primary_code_paths:
|
||||
- vllm_omni/core/**
|
||||
- vllm_omni/worker/**
|
||||
|
||||
@@ -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/**
|
||||
|
||||
@@ -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/**
|
||||
|
||||
@@ -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/**
|
||||
|
||||
@@ -6,6 +6,7 @@ architecture_state: current-plus-in-flight
|
||||
owners:
|
||||
- "@alex-jw-brooks"
|
||||
- "@linyueqian"
|
||||
- "@NickCao"
|
||||
document_stewards:
|
||||
- "@hsliuustc0106"
|
||||
- "@Gaohan123"
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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/**
|
||||
|
||||
@@ -4,6 +4,7 @@ kind: module
|
||||
status: draft
|
||||
owners:
|
||||
- "@lishunyang12"
|
||||
- "@vraiti"
|
||||
primary_code_paths:
|
||||
- vllm_omni/metrics/**
|
||||
related_code_paths:
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
```
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 "
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in new issue
Block a user