Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5308222abb | ||
|
|
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__()
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@@ -41,3 +43,21 @@ def test_wan_vae_execution_context_uses_platform_autocast(mocker):
|
||||
dtype=torch.bfloat16,
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
|
||||
def test_prepare_pipeline_installs_data_movement_before_compile(monkeypatch):
|
||||
model = _DummyOmniAutoencoderKLWan(dtype=torch.float32)
|
||||
installed = []
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover",
|
||||
lambda pipeline: SimpleNamespace(vaes=[model]),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
wan_vae_module,
|
||||
"install_wan_vae_data_movement",
|
||||
lambda vae: installed.append(vae),
|
||||
)
|
||||
|
||||
wan_vae_module.prepare_pipeline_wan_vae_data_movement(torch.nn.Identity())
|
||||
|
||||
assert installed == [model]
|
||||
@@ -0,0 +1,274 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from diffusers.models.autoencoders import AutoencoderKLWan
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders import wan_vae_data_movement
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
|
||||
from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import (
|
||||
_cache_payload,
|
||||
_causal_padding,
|
||||
_is_dynamic_spatial_conv,
|
||||
_run_cached_causal_conv,
|
||||
install_wan_vae_data_movement,
|
||||
)
|
||||
from vllm_omni.diffusion.kernels.wan_vae_data_movement import cat_pad_5d, dup_up3d_add
|
||||
|
||||
pytestmark = pytest.mark.core_model
|
||||
cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
|
||||
|
||||
def _dense_5d(shape: tuple[int, ...], dtype: torch.dtype, channels_last: bool) -> torch.Tensor:
|
||||
tensor = torch.randn(shape, device="cuda", dtype=dtype)
|
||||
if channels_last:
|
||||
return tensor.contiguous(memory_format=torch.channels_last_3d)
|
||||
return tensor.contiguous()
|
||||
|
||||
|
||||
def _reference_cat_pad(
|
||||
x: torch.Tensor,
|
||||
cache: torch.Tensor | None,
|
||||
padding: tuple[int, ...],
|
||||
) -> torch.Tensor:
|
||||
reference_padding = list(padding)
|
||||
if cache is not None:
|
||||
x = torch.cat([cache, x], dim=2)
|
||||
reference_padding[4] -= cache.shape[2]
|
||||
return F.pad(x, reference_padding)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@cuda_only
|
||||
@torch.no_grad()
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
|
||||
@pytest.mark.parametrize("channels_last", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"channels,frames,height,width,cache_frames,padding",
|
||||
[
|
||||
(96, 1, 10, 14, 0, (1, 1, 1, 1, 2, 0)),
|
||||
(96, 1, 10, 14, 1, (1, 1, 1, 1, 2, 0)),
|
||||
(96, 1, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
|
||||
(64, 1, 10, 14, 2, (0, 0, 0, 0, 2, 0)),
|
||||
(48, 4, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
|
||||
],
|
||||
)
|
||||
def test_cat_pad_5d_is_bitwise_and_layout_exact(
|
||||
dtype: torch.dtype,
|
||||
channels_last: bool,
|
||||
channels: int,
|
||||
frames: int,
|
||||
height: int,
|
||||
width: int,
|
||||
cache_frames: int,
|
||||
padding: tuple[int, ...],
|
||||
) -> None:
|
||||
torch.cuda.manual_seed(0)
|
||||
x = _dense_5d((1, channels, frames, height, width), dtype, channels_last)
|
||||
cache = None
|
||||
if cache_frames:
|
||||
pad_height, pad_width = padding[2], padding[0]
|
||||
buffer = _dense_5d(
|
||||
(1, channels, cache_frames, height + 2 * pad_height, width + 2 * pad_width),
|
||||
dtype,
|
||||
channels_last,
|
||||
)
|
||||
cache = buffer[:, :, :, pad_height : pad_height + height, pad_width : pad_width + width]
|
||||
|
||||
reference = _reference_cat_pad(x, cache, padding)
|
||||
output = cat_pad_5d(x, cache, padding)
|
||||
assert output is not None
|
||||
assert output.shape == reference.shape
|
||||
assert output.stride() == reference.stride()
|
||||
assert torch.equal(output, reference)
|
||||
|
||||
pair = cat_pad_5d(x, cache, padding, keep_cache_frames=2)
|
||||
assert pair is not None
|
||||
output_with_cache, next_cache = pair
|
||||
assert torch.equal(output_with_cache, reference)
|
||||
pad_height, pad_width = padding[2], padding[0]
|
||||
keep_frames = min(2, reference.shape[2])
|
||||
expected_cache = reference[
|
||||
:,
|
||||
:,
|
||||
-keep_frames:,
|
||||
pad_height : pad_height + height,
|
||||
pad_width : pad_width + width,
|
||||
]
|
||||
assert next_cache.shape == expected_cache.shape
|
||||
assert torch.equal(next_cache, expected_cache)
|
||||
assert next_cache.is_contiguous(memory_format=torch.channels_last_3d) is channels_last
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@cuda_only
|
||||
@torch.no_grad()
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
|
||||
@pytest.mark.parametrize("main_layout", ["contiguous", "channels_last", "permuted"])
|
||||
@pytest.mark.parametrize(
|
||||
"in_channels,out_channels,frames,height,width,factor_temporal,factor_spatial,drop_first",
|
||||
[
|
||||
(128, 64, 1, 10, 14, 2, 2, False),
|
||||
(128, 64, 1, 10, 14, 2, 2, True),
|
||||
(64, 32, 2, 10, 14, 1, 2, False),
|
||||
],
|
||||
)
|
||||
def test_dup_up3d_add_is_bitwise_and_layout_exact(
|
||||
dtype: torch.dtype,
|
||||
main_layout: str,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
frames: int,
|
||||
height: int,
|
||||
width: int,
|
||||
factor_temporal: int,
|
||||
factor_spatial: int,
|
||||
drop_first: bool,
|
||||
) -> None:
|
||||
torch.cuda.manual_seed(0)
|
||||
repeats = out_channels * factor_temporal * factor_spatial * factor_spatial // in_channels
|
||||
source = _dense_5d((1, in_channels, frames, height, width), dtype, main_layout == "channels_last")
|
||||
out_frames = frames * factor_temporal - (factor_temporal - 1 if drop_first else 0)
|
||||
shape = (1, out_channels, out_frames, height * factor_spatial, width * factor_spatial)
|
||||
if main_layout == "channels_last":
|
||||
main = _dense_5d(shape, dtype, True)
|
||||
elif main_layout == "permuted":
|
||||
main = torch.randn(
|
||||
(shape[0], shape[2], shape[1], shape[3], shape[4]),
|
||||
device="cuda",
|
||||
dtype=dtype,
|
||||
).permute(0, 2, 1, 3, 4)
|
||||
else:
|
||||
main = _dense_5d(shape, dtype, False)
|
||||
|
||||
shortcut = source.repeat_interleave(repeats, dim=1)
|
||||
shortcut = shortcut.view(
|
||||
1,
|
||||
out_channels,
|
||||
factor_temporal,
|
||||
factor_spatial,
|
||||
factor_spatial,
|
||||
frames,
|
||||
height,
|
||||
width,
|
||||
)
|
||||
shortcut = shortcut.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
|
||||
shortcut = shortcut.view(shape[0], shape[1], frames * factor_temporal, shape[3], shape[4])
|
||||
if drop_first:
|
||||
shortcut = shortcut[:, :, factor_temporal - 1 :]
|
||||
reference = main + shortcut
|
||||
|
||||
output = dup_up3d_add(
|
||||
main,
|
||||
source,
|
||||
factor_temporal,
|
||||
factor_spatial,
|
||||
repeats,
|
||||
drop_first,
|
||||
)
|
||||
assert output is not None
|
||||
assert output.shape == reference.shape
|
||||
assert output.stride() == reference.stride()
|
||||
assert torch.equal(output, reference)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@cuda_only
|
||||
@torch.no_grad()
|
||||
@pytest.mark.parametrize("channels_last", [False, True])
|
||||
@pytest.mark.parametrize("temporal_only", [False, True])
|
||||
def test_cached_conv_chunk_loop_matches_reference(monkeypatch, channels_last: bool, temporal_only: bool) -> None:
|
||||
from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T, WanCausalConv3d
|
||||
|
||||
torch.cuda.manual_seed(0)
|
||||
channels = 64
|
||||
if temporal_only:
|
||||
conv = WanCausalConv3d(channels, 2 * channels, (3, 1, 1), padding=(1, 0, 0))
|
||||
else:
|
||||
conv = WanCausalConv3d(channels, channels, 3, padding=1)
|
||||
conv = conv.to(device="cuda", dtype=torch.float32)
|
||||
chunks = [_dense_5d((1, channels, 1, 10, 14), torch.float32, channels_last) for _ in range(4)]
|
||||
|
||||
def run(force_reference: bool, start):
|
||||
cache = [start]
|
||||
outputs = []
|
||||
original = wan_vae_data_movement.cat_pad_5d
|
||||
if force_reference:
|
||||
monkeypatch.setattr(wan_vae_data_movement, "cat_pad_5d", None)
|
||||
try:
|
||||
for chunk in chunks:
|
||||
outputs.append(_run_cached_causal_conv(conv, chunk, cache, 0))
|
||||
finally:
|
||||
monkeypatch.setattr(wan_vae_data_movement, "cat_pad_5d", original)
|
||||
return outputs, cache[0]
|
||||
|
||||
for start in (None, "Rep"):
|
||||
fused_outputs, fused_cache = run(False, start)
|
||||
reference_outputs, reference_cache = run(True, start)
|
||||
for output, reference in zip(fused_outputs, reference_outputs, strict=True):
|
||||
assert torch.equal(output, reference)
|
||||
payload = _cache_payload(fused_cache)
|
||||
assert payload is not None and payload.shape[2] == CACHE_T
|
||||
assert torch.equal(payload, reference_cache[:, :, -CACHE_T:])
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_installer_is_idempotent_and_preserves_state_dict_keys() -> None:
|
||||
vae = OmniAutoencoderKLWan(
|
||||
base_dim=8,
|
||||
decoder_base_dim=8,
|
||||
z_dim=4,
|
||||
dim_mult=[1, 1],
|
||||
num_res_blocks=1,
|
||||
temperal_downsample=[False, False],
|
||||
is_residual=True,
|
||||
)
|
||||
keys_before = set(vae.state_dict())
|
||||
assert install_wan_vae_data_movement(vae)
|
||||
assert install_wan_vae_data_movement(vae)
|
||||
assert set(vae.state_dict()) == keys_before
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_dynamic_spatial_conv_exposes_direct_path_compatibility_contract() -> None:
|
||||
conv = torch.nn.Conv3d(2, 2, 3)
|
||||
conv._vllm_omni_dynamic_spatial_shard_conv = True
|
||||
conv._source_padding = (1, 1, 1, 1, 2, 0)
|
||||
|
||||
assert _is_dynamic_spatial_conv(conv)
|
||||
assert _causal_padding(conv) == conv._source_padding
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@cuda_only
|
||||
@torch.no_grad()
|
||||
@pytest.mark.parametrize("channels_last", [False, True])
|
||||
def test_tiny_wan_decoder_is_bitwise_exact(channels_last: bool) -> None:
|
||||
config = dict(
|
||||
base_dim=8,
|
||||
decoder_base_dim=8,
|
||||
z_dim=4,
|
||||
dim_mult=[1, 1],
|
||||
num_res_blocks=1,
|
||||
temperal_downsample=[False, True],
|
||||
is_residual=True,
|
||||
)
|
||||
torch.manual_seed(0)
|
||||
reference = AutoencoderKLWan(**config).eval().to("cuda")
|
||||
candidate = OmniAutoencoderKLWan(**config).eval().to("cuda")
|
||||
candidate.load_state_dict(reference.state_dict())
|
||||
if channels_last:
|
||||
for model in (reference, candidate):
|
||||
for module in model.decoder.modules():
|
||||
if isinstance(module, torch.nn.Conv3d):
|
||||
module.weight.data = module.weight.data.contiguous(memory_format=torch.channels_last_3d)
|
||||
|
||||
latents = torch.randn(1, 4, 3, 4, 4, device="cuda")
|
||||
if channels_last:
|
||||
latents = latents.contiguous(memory_format=torch.channels_last_3d)
|
||||
expected = reference.decode(latents, return_dict=False)[0]
|
||||
output = candidate.decode(latents, return_dict=False)[0]
|
||||
assert output.stride() == expected.stride()
|
||||
assert torch.equal(output, expected)
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -728,6 +728,8 @@ 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):
|
||||
prepared_pipelines = []
|
||||
|
||||
class _DummyLoader:
|
||||
def __init__(self, load_config, od_config=None):
|
||||
del load_config, od_config
|
||||
@@ -775,6 +777,11 @@ 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,
|
||||
"prepare_pipeline_wan_vae_data_movement",
|
||||
lambda pipeline: prepared_pipelines.append(pipeline),
|
||||
)
|
||||
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
|
||||
@@ -785,6 +792,7 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
|
||||
assert runner.cache_backend is None
|
||||
assert runner.od_config.cache_backend is None
|
||||
assert dummy_cache_backend.enabled is False
|
||||
assert prepared_pipelines == [runner.pipeline]
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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 "
|
||||
|
||||
@@ -18,13 +18,26 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
|
||||
GridSpec,
|
||||
TileTask,
|
||||
)
|
||||
from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import (
|
||||
install_wan_vae_data_movement,
|
||||
)
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def prepare_pipeline_wan_vae_data_movement(pipeline: torch.nn.Module) -> None:
|
||||
"""Install Wan data-movement methods before offload or compilation setup."""
|
||||
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
|
||||
|
||||
for vae in ModuleDiscovery.discover(pipeline).vaes:
|
||||
if isinstance(vae, OmniAutoencoderKLWan):
|
||||
install_wan_vae_data_movement(vae)
|
||||
|
||||
|
||||
class OmniAutoencoderKLWan(AutoencoderKLWan):
|
||||
def _execution_context(self):
|
||||
install_wan_vae_data_movement(self)
|
||||
try:
|
||||
first_param = next(self.parameters())
|
||||
except StopIteration:
|
||||
|
||||
@@ -0,0 +1,302 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Guarded Wan VAE decoder data-movement fast paths.
|
||||
|
||||
This module adapts the lossless Wan decoder optimizations from SGLang while
|
||||
keeping Diffusers' parameter names and module topology unchanged. Every
|
||||
optimized method falls back to the Diffusers 0.38.0 operation order when the
|
||||
kernel, device, layout, dtype, autograd state, or module type is unsupported.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import MethodType
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from diffusers.models.autoencoders.autoencoder_kl_wan import (
|
||||
CACHE_T,
|
||||
DupUp3D,
|
||||
WanCausalConv3d,
|
||||
WanDecoder3d,
|
||||
WanResample,
|
||||
WanResidualBlock,
|
||||
WanResidualUpBlock,
|
||||
)
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
|
||||
|
||||
try:
|
||||
from vllm_omni.diffusion.kernels.wan_vae_data_movement import (
|
||||
cat_pad_5d,
|
||||
dup_up3d_add,
|
||||
)
|
||||
except ImportError: # pragma: no cover - Triton is optional on non-GPU installs.
|
||||
cat_pad_5d = None
|
||||
dup_up3d_add = None
|
||||
|
||||
_INSTALL_ATTR = "_vllm_omni_wan_data_movement_installed"
|
||||
_DYNAMIC_SPATIAL_CONV_ATTR = "_vllm_omni_dynamic_spatial_shard_conv"
|
||||
|
||||
|
||||
def _cache_payload(cache: Any) -> torch.Tensor | None:
|
||||
return cache if isinstance(cache, torch.Tensor) else None
|
||||
|
||||
|
||||
def _is_dynamic_spatial_conv(conv: nn.Module) -> bool:
|
||||
return bool(getattr(conv, _DYNAMIC_SPATIAL_CONV_ATTR, False))
|
||||
|
||||
|
||||
def _spatial_shard_is_active() -> bool:
|
||||
return wan_spatial_shard._SPATIAL_SHARD_CONTEXT.get() is not None
|
||||
|
||||
|
||||
def _causal_padding(conv: nn.Module) -> tuple[int, ...]:
|
||||
if _is_dynamic_spatial_conv(conv):
|
||||
return tuple(conv._source_padding)
|
||||
return tuple(conv._padding)
|
||||
|
||||
|
||||
def _can_use_fused_cache_path(conv: nn.Module, x: torch.Tensor) -> bool:
|
||||
return bool(
|
||||
cat_pad_5d is not None
|
||||
and (type(conv) is WanCausalConv3d or _is_dynamic_spatial_conv(conv))
|
||||
and x.dim() == 5
|
||||
and x.is_cuda
|
||||
and not torch.is_grad_enabled()
|
||||
and not torch.compiler.is_compiling()
|
||||
and not _spatial_shard_is_active()
|
||||
)
|
||||
|
||||
|
||||
def _run_cached_causal_conv(
|
||||
conv: nn.Module,
|
||||
x: torch.Tensor,
|
||||
cache_list: list[Any],
|
||||
index: int,
|
||||
) -> torch.Tensor:
|
||||
"""Run a causal conv and refresh its feature-cache slot."""
|
||||
cache = cache_list[index]
|
||||
is_repeat_marker = isinstance(cache, str)
|
||||
payload = None if is_repeat_marker else _cache_payload(cache)
|
||||
if _can_use_fused_cache_path(conv, x) and (
|
||||
payload is None or (payload.device == x.device and payload.dtype == x.dtype)
|
||||
):
|
||||
pair = cat_pad_5d(x, payload, _causal_padding(conv), keep_cache_frames=CACHE_T)
|
||||
if pair is not None:
|
||||
conv_input, cache_list[index] = pair
|
||||
return nn.Conv3d.forward(conv, conv_input)
|
||||
|
||||
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
||||
if cache_x.shape[2] < CACHE_T and payload is not None:
|
||||
cache_x = torch.cat(
|
||||
[payload[:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
|
||||
dim=2,
|
||||
)
|
||||
elif cache_x.shape[2] < CACHE_T and is_repeat_marker:
|
||||
cache_x = torch.cat([torch.zeros_like(cache_x), cache_x], dim=2)
|
||||
|
||||
out = conv(x) if payload is None else conv(x, payload)
|
||||
cache_list[index] = cache_x
|
||||
return out
|
||||
|
||||
|
||||
def _causal_conv_forward(self: WanCausalConv3d, x: torch.Tensor, cache_x: torch.Tensor | None = None):
|
||||
padding = list(self._padding)
|
||||
if (
|
||||
any(padding)
|
||||
and _can_use_fused_cache_path(self, x)
|
||||
and (cache_x is None or (cache_x.device == x.device and cache_x.dtype == x.dtype))
|
||||
):
|
||||
conv_input = cat_pad_5d(x, cache_x, padding)
|
||||
if conv_input is not None:
|
||||
return nn.Conv3d.forward(self, conv_input)
|
||||
|
||||
if cache_x is not None and padding[4] > 0:
|
||||
cache_x = cache_x.to(x.device)
|
||||
x = torch.cat([cache_x, x], dim=2)
|
||||
padding[4] -= cache_x.shape[2]
|
||||
x = F.pad(x, padding)
|
||||
return nn.Conv3d.forward(self, x)
|
||||
|
||||
|
||||
def _resample_forward(
|
||||
self: WanResample,
|
||||
x: torch.Tensor,
|
||||
feat_cache: list[Any] | None = None,
|
||||
feat_idx: list[int] | None = None,
|
||||
):
|
||||
if feat_idx is None:
|
||||
feat_idx = [0]
|
||||
batch, channels, frames, height, width = x.size()
|
||||
if self.mode == "upsample3d" and feat_cache is not None:
|
||||
index = feat_idx[0]
|
||||
if feat_cache[index] is None:
|
||||
feat_cache[index] = "Rep"
|
||||
feat_idx[0] += 1
|
||||
else:
|
||||
x = _run_cached_causal_conv(self.time_conv, x, feat_cache, index)
|
||||
feat_idx[0] += 1
|
||||
x = x.reshape(batch, 2, channels, frames, height, width)
|
||||
x = torch.stack((x[:, 0], x[:, 1]), 3)
|
||||
x = x.reshape(batch, channels, frames * 2, height, width)
|
||||
|
||||
frames = x.shape[2]
|
||||
x = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width)
|
||||
x = self.resample(x)
|
||||
x = x.view(batch, frames, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4)
|
||||
|
||||
if self.mode == "downsample3d" and feat_cache is not None:
|
||||
index = feat_idx[0]
|
||||
if feat_cache[index] is None:
|
||||
feat_cache[index] = x.clone()
|
||||
feat_idx[0] += 1
|
||||
else:
|
||||
cache_x = x[:, :, -1:, :, :].clone()
|
||||
x = self.time_conv(torch.cat([feat_cache[index][:, :, -1:, :, :], x], dim=2))
|
||||
feat_cache[index] = cache_x
|
||||
feat_idx[0] += 1
|
||||
return x
|
||||
|
||||
|
||||
def _residual_block_forward(
|
||||
self: WanResidualBlock,
|
||||
x: torch.Tensor,
|
||||
feat_cache: list[Any] | None = None,
|
||||
feat_idx: list[int] | None = None,
|
||||
):
|
||||
if feat_idx is None:
|
||||
feat_idx = [0]
|
||||
residual = self.conv_shortcut(x)
|
||||
x = self.nonlinearity(self.norm1(x))
|
||||
if feat_cache is not None:
|
||||
index = feat_idx[0]
|
||||
x = _run_cached_causal_conv(self.conv1, x, feat_cache, index)
|
||||
feat_idx[0] += 1
|
||||
else:
|
||||
x = self.conv1(x)
|
||||
|
||||
x = self.dropout(self.nonlinearity(self.norm2(x)))
|
||||
if feat_cache is not None:
|
||||
index = feat_idx[0]
|
||||
x = _run_cached_causal_conv(self.conv2, x, feat_cache, index)
|
||||
feat_idx[0] += 1
|
||||
else:
|
||||
x = self.conv2(x)
|
||||
return x + residual
|
||||
|
||||
|
||||
def _residual_up_block_forward(
|
||||
self: WanResidualUpBlock,
|
||||
x: torch.Tensor,
|
||||
feat_cache: list[Any] | None = None,
|
||||
feat_idx: list[int] | None = None,
|
||||
first_chunk: bool = False,
|
||||
):
|
||||
if feat_idx is None:
|
||||
feat_idx = [0]
|
||||
shortcut_source = x.clone()
|
||||
for resnet in self.resnets:
|
||||
if feat_cache is None:
|
||||
x = resnet(x)
|
||||
else:
|
||||
x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx)
|
||||
|
||||
if self.upsampler is not None:
|
||||
if feat_cache is None:
|
||||
x = self.upsampler(x)
|
||||
else:
|
||||
x = self.upsampler(x, feat_cache=feat_cache, feat_idx=feat_idx)
|
||||
|
||||
shortcut = self.avg_shortcut
|
||||
if shortcut is None:
|
||||
return x
|
||||
if (
|
||||
dup_up3d_add is not None
|
||||
and type(shortcut) is DupUp3D
|
||||
and x.is_cuda
|
||||
and shortcut_source.is_cuda
|
||||
and x.dtype == shortcut_source.dtype
|
||||
and not torch.is_grad_enabled()
|
||||
and not torch.compiler.is_compiling()
|
||||
and not _spatial_shard_is_active()
|
||||
):
|
||||
fused = dup_up3d_add(
|
||||
x,
|
||||
shortcut_source,
|
||||
shortcut.factor_t,
|
||||
shortcut.factor_s,
|
||||
shortcut.repeats,
|
||||
first_chunk,
|
||||
)
|
||||
if fused is not None:
|
||||
return fused
|
||||
return x + shortcut(shortcut_source, first_chunk=first_chunk)
|
||||
|
||||
|
||||
def _decoder_forward(
|
||||
self: WanDecoder3d,
|
||||
x: torch.Tensor,
|
||||
feat_cache: list[Any] | None = None,
|
||||
feat_idx: list[int] | None = None,
|
||||
first_chunk: bool = False,
|
||||
):
|
||||
if feat_idx is None:
|
||||
feat_idx = [0]
|
||||
if feat_cache is not None:
|
||||
index = feat_idx[0]
|
||||
x = _run_cached_causal_conv(self.conv_in, x, feat_cache, index)
|
||||
feat_idx[0] += 1
|
||||
else:
|
||||
x = self.conv_in(x)
|
||||
|
||||
x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx)
|
||||
for up_block in self.up_blocks:
|
||||
x = up_block(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
|
||||
|
||||
x = self.nonlinearity(self.norm_out(x))
|
||||
if feat_cache is not None:
|
||||
index = feat_idx[0]
|
||||
x = _run_cached_causal_conv(self.conv_out, x, feat_cache, index)
|
||||
feat_idx[0] += 1
|
||||
else:
|
||||
x = self.conv_out(x)
|
||||
return x
|
||||
|
||||
|
||||
def install_wan_vae_data_movement(vae: nn.Module) -> bool:
|
||||
"""Install guarded fast paths on a Diffusers Wan decoder once."""
|
||||
if getattr(vae, _INSTALL_ATTR, False):
|
||||
return True
|
||||
if cat_pad_5d is None or dup_up3d_add is None:
|
||||
return False
|
||||
|
||||
decoder = getattr(vae, "decoder", None)
|
||||
if type(decoder) is not WanDecoder3d:
|
||||
return False
|
||||
|
||||
decoder.forward = MethodType(_decoder_forward, decoder)
|
||||
patched = 0
|
||||
for module in decoder.modules():
|
||||
module_type = type(module)
|
||||
if module_type is WanCausalConv3d:
|
||||
module.forward = MethodType(_causal_conv_forward, module)
|
||||
patched += 1
|
||||
elif _is_dynamic_spatial_conv(module):
|
||||
# The dynamic wrapper owns halo exchange while a spatial request is
|
||||
# active. Its direct path can still use the lossless fused cache
|
||||
# movement through ``_run_cached_causal_conv``.
|
||||
patched += 1
|
||||
elif module_type is WanResidualBlock:
|
||||
module.forward = MethodType(_residual_block_forward, module)
|
||||
elif module_type is WanResidualUpBlock:
|
||||
module.forward = MethodType(_residual_up_block_forward, module)
|
||||
elif module_type is WanResample:
|
||||
module.forward = MethodType(_resample_forward, module)
|
||||
|
||||
if patched == 0:
|
||||
return False
|
||||
setattr(vae, _INSTALL_ATTR, True)
|
||||
return True
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -0,0 +1,405 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Bit-exact data-movement kernels for the Wan causal VAE decoder.
|
||||
|
||||
Adapted from SGLang PR #34125:
|
||||
https://github.com/sgl-project/sglang/pull/34125
|
||||
|
||||
The kernels in this module only move values, fill zeros, or perform the same
|
||||
two-operand addition as the PyTorch reference path. They preserve the input
|
||||
layout instead of selecting a new Conv3d layout, so enabling them does not
|
||||
change the decoder's cuDNN accumulation order.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
_MAX_INT32 = 2**31 - 1
|
||||
_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _cat_pad_5d_kernel(
|
||||
x_ptr,
|
||||
cache_ptr,
|
||||
out_ptr,
|
||||
keep_ptr,
|
||||
total,
|
||||
channels,
|
||||
frames,
|
||||
height,
|
||||
width,
|
||||
cache_frames,
|
||||
out_frames,
|
||||
out_height,
|
||||
out_width,
|
||||
zero_front_frames,
|
||||
pad_height,
|
||||
pad_width,
|
||||
stride_x_batch,
|
||||
stride_x_channel,
|
||||
stride_x_frame,
|
||||
stride_x_height,
|
||||
stride_x_width,
|
||||
stride_cache_batch,
|
||||
stride_cache_channel,
|
||||
stride_cache_frame,
|
||||
stride_cache_height,
|
||||
stride_cache_width,
|
||||
has_cache: tl.constexpr,
|
||||
keep_frames_const: tl.constexpr,
|
||||
channels_inner: tl.constexpr,
|
||||
index_64bit: tl.constexpr,
|
||||
block_size_const: tl.constexpr,
|
||||
):
|
||||
if index_64bit:
|
||||
offsets = tl.program_id(0).to(tl.int64) * block_size_const + tl.arange(0, block_size_const).to(tl.int64)
|
||||
else:
|
||||
offsets = tl.program_id(0) * block_size_const + tl.arange(0, block_size_const)
|
||||
mask = offsets < total
|
||||
|
||||
if channels_inner:
|
||||
out_channel = offsets % channels
|
||||
remaining = offsets // channels
|
||||
out_width_idx = remaining % out_width
|
||||
remaining = remaining // out_width
|
||||
out_height_idx = remaining % out_height
|
||||
remaining = remaining // out_height
|
||||
out_frame = remaining % out_frames
|
||||
out_batch = remaining // out_frames
|
||||
else:
|
||||
out_width_idx = offsets % out_width
|
||||
remaining = offsets // out_width
|
||||
out_height_idx = remaining % out_height
|
||||
remaining = remaining // out_height
|
||||
out_frame = remaining % out_frames
|
||||
remaining = remaining // out_frames
|
||||
out_channel = remaining % channels
|
||||
out_batch = remaining // channels
|
||||
|
||||
input_width = out_width_idx - pad_width
|
||||
input_height = out_height_idx - pad_height
|
||||
input_frame = out_frame - zero_front_frames
|
||||
|
||||
spatial_valid = (input_width >= 0) & (input_width < width) & (input_height >= 0) & (input_height < height)
|
||||
from_cache = spatial_valid & (input_frame >= 0) & (input_frame < cache_frames)
|
||||
from_x = spatial_valid & (input_frame >= cache_frames) & (input_frame < cache_frames + frames)
|
||||
|
||||
x_frame = input_frame - cache_frames
|
||||
x_offset = (
|
||||
out_batch * stride_x_batch
|
||||
+ out_channel * stride_x_channel
|
||||
+ x_frame * stride_x_frame
|
||||
+ input_height * stride_x_height
|
||||
+ input_width * stride_x_width
|
||||
)
|
||||
values = tl.load(x_ptr + x_offset, mask=mask & from_x, other=0.0)
|
||||
if has_cache:
|
||||
cache_offset = (
|
||||
out_batch * stride_cache_batch
|
||||
+ out_channel * stride_cache_channel
|
||||
+ input_frame * stride_cache_frame
|
||||
+ input_height * stride_cache_height
|
||||
+ input_width * stride_cache_width
|
||||
)
|
||||
cache_values = tl.load(cache_ptr + cache_offset, mask=mask & from_cache, other=0.0)
|
||||
values = tl.where(from_cache, cache_values, values)
|
||||
tl.store(out_ptr + offsets, values, mask=mask)
|
||||
|
||||
if keep_frames_const > 0:
|
||||
keep_frame = out_frame - (out_frames - keep_frames_const)
|
||||
keep_mask = mask & spatial_valid & (keep_frame >= 0)
|
||||
if channels_inner:
|
||||
keep_offset = (
|
||||
((out_batch * keep_frames_const + keep_frame) * height + input_height) * width + input_width
|
||||
) * channels + out_channel
|
||||
else:
|
||||
keep_offset = (
|
||||
((out_batch * channels + out_channel) * keep_frames_const + keep_frame) * height + input_height
|
||||
) * width + input_width
|
||||
tl.store(keep_ptr + keep_offset, values, mask=keep_mask)
|
||||
|
||||
|
||||
def _uses_channels_last_3d(tensor: torch.Tensor) -> bool:
|
||||
return tensor.shape[1] > 1 and tensor.stride(1) == 1
|
||||
|
||||
|
||||
def cat_pad_5d(
|
||||
x: torch.Tensor,
|
||||
cache_x: torch.Tensor | None,
|
||||
padding: list[int] | tuple[int, ...],
|
||||
keep_cache_frames: int = 0,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None:
|
||||
"""Fuse temporal concat, constant padding, and optional cache refresh.
|
||||
|
||||
``padding`` follows the Wan causal Conv3d convention
|
||||
``(w_left, w_right, h_top, h_bottom, t_front, t_back)``. The output uses
|
||||
the same dense layout that the reference ``cat``/``pad`` path selects.
|
||||
Unsupported inputs return ``None`` so callers can execute the reference
|
||||
implementation.
|
||||
"""
|
||||
if len(padding) != 6:
|
||||
return None
|
||||
pad_width_left, pad_width_right, pad_height_top, pad_height_bottom, pad_front, pad_back = padding
|
||||
if any(value < 0 for value in padding):
|
||||
return None
|
||||
if pad_width_left != pad_width_right or pad_height_top != pad_height_bottom or pad_back != 0:
|
||||
return None
|
||||
if x.dim() != 5 or not x.is_cuda or x.dtype not in _SUPPORTED_DTYPES:
|
||||
return None
|
||||
if keep_cache_frames < 0:
|
||||
return None
|
||||
|
||||
channels_inner = _uses_channels_last_3d(x)
|
||||
cache_frames = 0
|
||||
if cache_x is not None:
|
||||
if (
|
||||
cache_x.dim() != 5
|
||||
or cache_x.dtype != x.dtype
|
||||
or cache_x.device != x.device
|
||||
or cache_x.shape[0] != x.shape[0]
|
||||
or cache_x.shape[1] != x.shape[1]
|
||||
or cache_x.shape[3:] != x.shape[3:]
|
||||
or _uses_channels_last_3d(cache_x) != channels_inner
|
||||
):
|
||||
return None
|
||||
cache_frames = cache_x.shape[2]
|
||||
|
||||
zero_front_frames = pad_front - cache_frames
|
||||
if zero_front_frames < 0:
|
||||
return None
|
||||
|
||||
batch, channels, frames, height, width = x.shape
|
||||
out_frames = pad_front + frames
|
||||
out_height = height + 2 * pad_height_top
|
||||
out_width = width + 2 * pad_width_left
|
||||
keep_frames = min(keep_cache_frames, out_frames)
|
||||
memory_format = torch.channels_last_3d if channels_inner else torch.contiguous_format
|
||||
out = torch.empty(
|
||||
(batch, channels, out_frames, out_height, out_width),
|
||||
device=x.device,
|
||||
dtype=x.dtype,
|
||||
memory_format=memory_format,
|
||||
)
|
||||
total = out.numel()
|
||||
if total == 0 or total > _MAX_INT32 * 4:
|
||||
return None
|
||||
|
||||
if keep_frames:
|
||||
keep_arg = torch.empty(
|
||||
(batch, channels, keep_frames, height, width),
|
||||
device=x.device,
|
||||
dtype=x.dtype,
|
||||
memory_format=memory_format,
|
||||
)
|
||||
else:
|
||||
keep_arg = out
|
||||
|
||||
if cache_x is None:
|
||||
cache_arg = x
|
||||
cache_strides = (0, 0, 0, 0, 0)
|
||||
else:
|
||||
cache_arg = cache_x
|
||||
cache_strides = cache_x.stride()
|
||||
|
||||
block_size = 512
|
||||
grid = (triton.cdiv(total, block_size),)
|
||||
with torch.get_device_module().device(x.device):
|
||||
_cat_pad_5d_kernel[grid](
|
||||
x,
|
||||
cache_arg,
|
||||
out,
|
||||
keep_arg,
|
||||
total,
|
||||
channels,
|
||||
frames,
|
||||
height,
|
||||
width,
|
||||
cache_frames,
|
||||
out_frames,
|
||||
out_height,
|
||||
out_width,
|
||||
zero_front_frames,
|
||||
pad_height_top,
|
||||
pad_width_left,
|
||||
*x.stride(),
|
||||
*cache_strides,
|
||||
has_cache=cache_x is not None,
|
||||
keep_frames_const=keep_frames,
|
||||
channels_inner=channels_inner,
|
||||
index_64bit=total >= _MAX_INT32,
|
||||
block_size_const=block_size,
|
||||
)
|
||||
|
||||
if keep_cache_frames:
|
||||
return out, keep_arg
|
||||
return out
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _dup_up3d_add_kernel(
|
||||
main_ptr,
|
||||
source_ptr,
|
||||
out_ptr,
|
||||
total,
|
||||
out_channels,
|
||||
out_frames,
|
||||
out_height,
|
||||
out_width,
|
||||
frame_offset,
|
||||
stride_main_batch,
|
||||
stride_main_channel,
|
||||
stride_main_frame,
|
||||
stride_main_height,
|
||||
stride_main_width,
|
||||
stride_source_batch,
|
||||
stride_source_channel,
|
||||
stride_source_frame,
|
||||
stride_source_height,
|
||||
stride_source_width,
|
||||
stride_out_batch,
|
||||
stride_out_channel,
|
||||
stride_out_frame,
|
||||
stride_out_height,
|
||||
stride_out_width,
|
||||
factor_temporal_const: tl.constexpr,
|
||||
factor_spatial_const: tl.constexpr,
|
||||
repeats_const: tl.constexpr,
|
||||
channels_inner: tl.constexpr,
|
||||
index_64bit: tl.constexpr,
|
||||
block_size_const: tl.constexpr,
|
||||
):
|
||||
if index_64bit:
|
||||
offsets = tl.program_id(0).to(tl.int64) * block_size_const + tl.arange(0, block_size_const).to(tl.int64)
|
||||
else:
|
||||
offsets = tl.program_id(0) * block_size_const + tl.arange(0, block_size_const)
|
||||
mask = offsets < total
|
||||
|
||||
if channels_inner:
|
||||
out_channel = offsets % out_channels
|
||||
remaining = offsets // out_channels
|
||||
out_width_idx = remaining % out_width
|
||||
remaining = remaining // out_width
|
||||
out_height_idx = remaining % out_height
|
||||
remaining = remaining // out_height
|
||||
out_frame = remaining % out_frames
|
||||
out_batch = remaining // out_frames
|
||||
else:
|
||||
out_width_idx = offsets % out_width
|
||||
remaining = offsets // out_width
|
||||
out_height_idx = remaining % out_height
|
||||
remaining = remaining // out_height
|
||||
out_frame = remaining % out_frames
|
||||
remaining = remaining // out_frames
|
||||
out_channel = remaining % out_channels
|
||||
out_batch = remaining // out_channels
|
||||
|
||||
unsliced_frame = out_frame + frame_offset
|
||||
source_frame = unsliced_frame // factor_temporal_const
|
||||
temporal_remainder = unsliced_frame % factor_temporal_const
|
||||
source_height = out_height_idx // factor_spatial_const
|
||||
height_remainder = out_height_idx % factor_spatial_const
|
||||
source_width = out_width_idx // factor_spatial_const
|
||||
width_remainder = out_width_idx % factor_spatial_const
|
||||
repeated_channel = (
|
||||
(out_channel * factor_temporal_const + temporal_remainder) * factor_spatial_const + height_remainder
|
||||
) * factor_spatial_const + width_remainder
|
||||
source_channel = repeated_channel // repeats_const
|
||||
|
||||
main_offset = (
|
||||
out_batch * stride_main_batch
|
||||
+ out_channel * stride_main_channel
|
||||
+ out_frame * stride_main_frame
|
||||
+ out_height_idx * stride_main_height
|
||||
+ out_width_idx * stride_main_width
|
||||
)
|
||||
source_offset = (
|
||||
out_batch * stride_source_batch
|
||||
+ source_channel * stride_source_channel
|
||||
+ source_frame * stride_source_frame
|
||||
+ source_height * stride_source_height
|
||||
+ source_width * stride_source_width
|
||||
)
|
||||
out_offset = (
|
||||
out_batch * stride_out_batch
|
||||
+ out_channel * stride_out_channel
|
||||
+ out_frame * stride_out_frame
|
||||
+ out_height_idx * stride_out_height
|
||||
+ out_width_idx * stride_out_width
|
||||
)
|
||||
main_value = tl.load(main_ptr + main_offset, mask=mask, other=0.0)
|
||||
source_value = tl.load(source_ptr + source_offset, mask=mask, other=0.0)
|
||||
value = main_value.to(tl.float32) + source_value.to(tl.float32)
|
||||
tl.store(out_ptr + out_offset, value, mask=mask)
|
||||
|
||||
|
||||
def dup_up3d_add(
|
||||
main: torch.Tensor,
|
||||
source: torch.Tensor,
|
||||
factor_temporal: int,
|
||||
factor_spatial: int,
|
||||
repeats: int,
|
||||
drop_first_frames: bool,
|
||||
) -> torch.Tensor | None:
|
||||
"""Evaluate ``main + DupUp3D(source)`` without materializing the shortcut."""
|
||||
if main.dim() != 5 or source.dim() != 5:
|
||||
return None
|
||||
if factor_temporal <= 0 or factor_spatial <= 0 or repeats <= 0:
|
||||
return None
|
||||
if factor_temporal & (factor_temporal - 1) or factor_spatial & (factor_spatial - 1):
|
||||
return None
|
||||
if repeats & (repeats - 1):
|
||||
return None
|
||||
if not main.is_cuda or not source.is_cuda:
|
||||
return None
|
||||
if main.dtype != source.dtype or main.dtype not in _SUPPORTED_DTYPES or main.device != source.device:
|
||||
return None
|
||||
|
||||
batch, source_channels, frames, height, width = source.shape
|
||||
frame_offset = factor_temporal - 1 if drop_first_frames else 0
|
||||
numerator = source_channels * repeats
|
||||
denominator = factor_temporal * factor_spatial * factor_spatial
|
||||
if numerator % denominator:
|
||||
return None
|
||||
expected_shape = (
|
||||
batch,
|
||||
numerator // denominator,
|
||||
frames * factor_temporal - frame_offset,
|
||||
height * factor_spatial,
|
||||
width * factor_spatial,
|
||||
)
|
||||
if tuple(main.shape) != expected_shape:
|
||||
return None
|
||||
|
||||
out = torch.empty_like(main)
|
||||
total = out.numel()
|
||||
if total == 0 or total > _MAX_INT32 * 4:
|
||||
return None
|
||||
|
||||
block_size = 512
|
||||
grid = (triton.cdiv(total, block_size),)
|
||||
with torch.get_device_module().device(main.device):
|
||||
_dup_up3d_add_kernel[grid](
|
||||
main,
|
||||
source,
|
||||
out,
|
||||
total,
|
||||
expected_shape[1],
|
||||
expected_shape[2],
|
||||
expected_shape[3],
|
||||
expected_shape[4],
|
||||
frame_offset,
|
||||
*main.stride(),
|
||||
*source.stride(),
|
||||
*out.stride(),
|
||||
factor_temporal_const=factor_temporal,
|
||||
factor_spatial_const=factor_spatial,
|
||||
repeats_const=repeats,
|
||||
channels_inner=out.stride(1) == 1 and expected_shape[1] > 1,
|
||||
index_64bit=total >= _MAX_INT32,
|
||||
block_size_const=block_size,
|
||||
)
|
||||
return out
|
||||
@@ -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,
|
||||
|
||||
@@ -21,6 +21,7 @@ from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import install_wan_vae_data_movement
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
@@ -215,6 +216,7 @@ class HeliosPipeline(
|
||||
self.vae = AutoencoderKLWan.from_pretrained(
|
||||
model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
install_wan_vae_data_movement(self.vae)
|
||||
|
||||
transformer_config = load_transformer_config(model, "transformer", local_files_only)
|
||||
self.transformer = create_transformer_from_config(
|
||||
|
||||
@@ -19,6 +19,7 @@ from torch import nn
|
||||
from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import install_wan_vae_data_movement
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.models.interface import (
|
||||
SupportImageInput,
|
||||
@@ -374,6 +375,7 @@ class LingBotVideoPipeline(
|
||||
torch_dtype=vae_dtype,
|
||||
local_files_only=local_files_only,
|
||||
).to(self.device)
|
||||
install_wan_vae_data_movement(self.vae)
|
||||
self.scheduler = FlowUniPCMultistepScheduler.from_pretrained(
|
||||
model,
|
||||
subfolder=scheduler_subfolder,
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -30,6 +30,9 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import (
|
||||
from vllm_omni.diffusion.cache.selector import get_cache_backend
|
||||
from vllm_omni.diffusion.compile import regionally_compile
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
|
||||
prepare_pipeline_wan_vae_data_movement,
|
||||
)
|
||||
from vllm_omni.diffusion.forward_context import set_forward_context
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.models.interface import (
|
||||
@@ -243,6 +246,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
custom_pipeline_name=custom_pipeline_name,
|
||||
device=self.device,
|
||||
)
|
||||
prepare_pipeline_wan_vae_data_movement(self.pipeline)
|
||||
time_after_load = time.perf_counter()
|
||||
|
||||
logger.info(
|
||||
|
||||
@@ -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