Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Hongsheng Liu
|
29d2f157db | ||
|
Hongsheng Liu
|
ddcc050a94 | ||
|
Hongsheng Liu
|
50f4e675c1 | ||
|
Hongsheng Liu
|
159dceaf65 | ||
|
Hongsheng Liu
|
d10bf78c60 |
No files matched your search
+30
-112
@@ -1,11 +1,6 @@
|
||||
# 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
|
||||
@@ -26,10 +21,25 @@
|
||||
# CI
|
||||
/.buildkite/ @yenuo26 @congw729 @NickCao
|
||||
|
||||
# ComfyUI / external apps
|
||||
/apps/ @fhfuih @alex-jw-brooks
|
||||
# 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
|
||||
|
||||
# Benchmarks (out-of-tree)
|
||||
# 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/accuracy/ @david6666666
|
||||
/benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
|
||||
/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
|
||||
@@ -37,112 +47,20 @@
|
||||
/benchmarks/glm_image/ @princepride @JaredforReal
|
||||
/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
|
||||
|
||||
# =====================================================================
|
||||
# Module design ownership (docs/design/module/*.md)
|
||||
# Each section header cites the source design doc.
|
||||
# =====================================================================
|
||||
|
||||
# Autoregressive runtime — docs/design/module/ar_runtime.md
|
||||
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
|
||||
/vllm_omni/worker/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
|
||||
|
||||
# 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
|
||||
|
||||
# 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
|
||||
|
||||
# Observability — docs/design/module/observability.md
|
||||
/vllm_omni/metrics/ @lishunyang12 @vraiti
|
||||
|
||||
# Profiling — docs/design/module/profiling.md
|
||||
# Usability and Observability
|
||||
/vllm_omni/deploy/ @tzhouam @lishunyang12
|
||||
/vllm_omni/metrics/ @lishunyang12
|
||||
/vllm_omni/profiler/ @gcanlin
|
||||
|
||||
# Benchmarking (in-tree) — docs/design/module/benchmarking.md
|
||||
/vllm_omni/benchmarks/ @alex-jw-brooks @Bounty-hunter
|
||||
# ComfyUI
|
||||
/apps/ @fhfuih @alex-jw-brooks
|
||||
|
||||
# LoRA (no dedicated design doc; ownership intentionally narrowed)
|
||||
/vllm_omni/lora/ @SamitHuang @wtomin
|
||||
# Large Scale Serving
|
||||
/vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
|
||||
|
||||
# Transformers utils (related path only in design docs; preserves prior ownership)
|
||||
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
|
||||
# Quantization and configuration
|
||||
/vllm_omni/config/ @lishunyang12 @alex-jw-brooks
|
||||
/vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
|
||||
|
||||
# =====================================================================
|
||||
# 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
|
||||
# Hardware plugins and accelerators
|
||||
/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
|
||||
@@ -66,10 +66,7 @@ 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)|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
|
||||
files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)$
|
||||
|
||||
# Keep `suggestion` last
|
||||
- id: suggestion
|
||||
|
||||
+14
-4
@@ -16,21 +16,31 @@ 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
|
||||
- Speech-To-Video: user_guide/examples/offline_inference/speech_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
|
||||
- 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
|
||||
- Speech-To-Video: user_guide/examples/online_serving/speech_to_video.md
|
||||
- Text-To-Image: user_guide/examples/online_serving/text_to_image.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
|
||||
- Text-To-Image: user_guide/examples/online_serving/text_to_image.md
|
||||
- Text-To-Video: user_guide/examples/online_serving/text_to_video.md
|
||||
- General:
|
||||
- usage/*
|
||||
|
||||
@@ -7,7 +7,6 @@ owners:
|
||||
- "@yinpeiqi"
|
||||
- "@fake0fan"
|
||||
- "@Sy0307"
|
||||
- "@Gaohan123"
|
||||
primary_code_paths:
|
||||
- vllm_omni/core/**
|
||||
- vllm_omni/worker/**
|
||||
|
||||
@@ -4,7 +4,6 @@ kind: module
|
||||
status: draft
|
||||
owners:
|
||||
- "@alex-jw-brooks"
|
||||
- "@Bounty-hunter"
|
||||
primary_code_paths:
|
||||
- vllm_omni/benchmarks/**
|
||||
- vllm_omni/entrypoints/cli/benchmark/**
|
||||
|
||||
@@ -5,7 +5,6 @@ status: draft
|
||||
owners:
|
||||
- "@Isotr0py"
|
||||
- "@princepride"
|
||||
- "@SamitHuang"
|
||||
primary_code_paths:
|
||||
- vllm_omni/diffusion/cache/**
|
||||
related_code_paths:
|
||||
|
||||
@@ -3,10 +3,9 @@ title: Diffusion Model Integration
|
||||
kind: module
|
||||
status: draft
|
||||
owners:
|
||||
- "@fhfuih"
|
||||
- "@Bounty-hunter"
|
||||
- "@Isotr0py"
|
||||
- "@wtomin"
|
||||
- "@RuixiangMa"
|
||||
- "@ZJY0516"
|
||||
primary_code_paths:
|
||||
- vllm_omni/diffusion/models/**
|
||||
- vllm_omni/diffusion/model_loader/**
|
||||
|
||||
@@ -6,7 +6,6 @@ owners:
|
||||
- "@Isotr0py"
|
||||
- "@princepride"
|
||||
- "@SamitHuang"
|
||||
- "@fhfuih"
|
||||
primary_code_paths:
|
||||
- vllm_omni/diffusion/sched/**
|
||||
- vllm_omni/diffusion/executor/**
|
||||
|
||||
@@ -4,8 +4,7 @@ kind: module
|
||||
status: draft
|
||||
owners:
|
||||
- "@yuanheng-zhao"
|
||||
- "@david6666666"
|
||||
- "@lishunyang12"
|
||||
- "@Isotr0py"
|
||||
primary_code_paths:
|
||||
- vllm_omni/diffusion/offloader/**
|
||||
related_code_paths:
|
||||
|
||||
@@ -6,7 +6,6 @@ owners:
|
||||
- "@princepride"
|
||||
- "@yuanheng-zhao"
|
||||
- "@xuechendi"
|
||||
- "@Bounty-hunter"
|
||||
primary_code_paths:
|
||||
- vllm_omni/diffusion/distributed/**
|
||||
- vllm_omni/diffusion/attention/parallel/**
|
||||
|
||||
@@ -6,7 +6,6 @@ architecture_state: current-plus-in-flight
|
||||
owners:
|
||||
- "@alex-jw-brooks"
|
||||
- "@linyueqian"
|
||||
- "@NickCao"
|
||||
document_stewards:
|
||||
- "@hsliuustc0106"
|
||||
- "@Gaohan123"
|
||||
|
||||
@@ -8,15 +8,8 @@ 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,8 +6,6 @@ architecture_state: current
|
||||
owners:
|
||||
- "@Sy0307"
|
||||
- "@amy-why-3459"
|
||||
- "@Gaohan123"
|
||||
- "@alex-jw-brooks"
|
||||
document_stewards:
|
||||
- "@hsliuustc0106"
|
||||
- "@Gaohan123"
|
||||
|
||||
@@ -9,14 +9,6 @@ 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,7 +4,6 @@ kind: module
|
||||
status: draft
|
||||
owners:
|
||||
- "@lishunyang12"
|
||||
- "@vraiti"
|
||||
primary_code_paths:
|
||||
- vllm_omni/metrics/**
|
||||
related_code_paths:
|
||||
|
||||
@@ -24,22 +24,15 @@ 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",
|
||||
}
|
||||
)
|
||||
# Generated example pages keep their source-derived URLs even when their
|
||||
# navigation category changes. Quantization helpers are feature documentation,
|
||||
# not user-serving examples, so keep them under the existing Features /
|
||||
# Quantization section instead of creating a second top-level Quantization
|
||||
# category under Examples.
|
||||
EXAMPLE_NAV_TARGETS = {
|
||||
"quantization": "quantization",
|
||||
}
|
||||
GENERATED_QUANTIZATION_EXAMPLES_TITLE = "Generated Examples"
|
||||
|
||||
|
||||
def load_model_display_names() -> dict[str, str]:
|
||||
@@ -63,13 +56,6 @@ 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",
|
||||
@@ -372,9 +358,12 @@ def update_nav_file(examples: list[Example]):
|
||||
items.append({example.nav_title: rel_path.as_posix()})
|
||||
return items
|
||||
|
||||
# Build the new Examples section - start with preserved items.
|
||||
# Build the new Examples section - start with preserved items. Categories
|
||||
# not explicitly routed elsewhere retain the historical Examples behavior.
|
||||
examples_section = preserved_items.copy()
|
||||
for category in sorted(examples_by_category.keys()):
|
||||
if EXAMPLE_NAV_TARGETS.get(category, "examples") != "examples":
|
||||
continue
|
||||
items = category_items(category)
|
||||
if items:
|
||||
category_title = fix_case(category.replace("_", " ").title())
|
||||
@@ -383,6 +372,36 @@ def update_nav_file(examples: list[Example]):
|
||||
# Update the nav structure
|
||||
nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"] = examples_section
|
||||
|
||||
features_content = next(
|
||||
(item["Features"] for item in user_guide_content if isinstance(item, dict) and "Features" in item),
|
||||
None,
|
||||
)
|
||||
quantization_content = None
|
||||
if isinstance(features_content, list):
|
||||
quantization_content = next(
|
||||
(item["Quantization"] for item in features_content if isinstance(item, dict) and "Quantization" in item),
|
||||
None,
|
||||
)
|
||||
|
||||
if isinstance(quantization_content, list):
|
||||
# Remove only the group owned by this hook so repeated builds are
|
||||
# idempotent and hand-authored Quantization navigation is retained.
|
||||
quantization_content[:] = [
|
||||
item
|
||||
for item in quantization_content
|
||||
if not (isinstance(item, dict) and GENERATED_QUANTIZATION_EXAMPLES_TITLE in item)
|
||||
]
|
||||
|
||||
quantization_examples = examples_by_category.get("quantization")
|
||||
if quantization_examples:
|
||||
if not isinstance(quantization_content, list):
|
||||
logger.warning(
|
||||
"Could not find 'User Guide' -> 'Features' -> 'Quantization' section; "
|
||||
"omitting generated quantization examples from navigation"
|
||||
)
|
||||
else:
|
||||
quantization_content.append({GENERATED_QUANTIZATION_EXAMPLES_TITLE: category_items("quantization")})
|
||||
|
||||
# Write back to file
|
||||
nav_data["nav"] = nav_list
|
||||
with open(NAV_FILE, "w", encoding="utf-8") as f:
|
||||
@@ -413,8 +432,6 @@ 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"
|
||||
|
||||
@@ -283,20 +283,23 @@ blocks — use Ulysses SP (`ring_degree=1`).
|
||||
|
||||
### Which geometries run sparse
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
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.
|
||||
"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`.
|
||||
|
||||
## 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](../../diffusion/attention_backends.md)
|
||||
[vLLM-Omni standard of attention backend setting](../../../docs/user_guide/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.
|
||||
|
||||
@@ -79,8 +79,7 @@ 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 distributed
|
||||
layerwise offload:
|
||||
degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload:
|
||||
|
||||
```bash
|
||||
export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
|
||||
@@ -99,7 +98,7 @@ vllm serve "${MODEL}" \
|
||||
--usp 8 \
|
||||
--ring 1 \
|
||||
--text-encoder-tp-size 8 \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--enable-layerwise-offload \
|
||||
--vae-parallel-mode tile \
|
||||
--vae-use-tiling \
|
||||
--vae-patch-parallel-size 8 \
|
||||
@@ -113,11 +112,6 @@ 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
|
||||
@@ -140,63 +134,6 @@ 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
|
||||
@@ -217,8 +154,8 @@ configuration above:
|
||||
|
||||
| Workload | Configuration |
|
||||
|----------|---------------|
|
||||
| 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 |
|
||||
| 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 |
|
||||
|
||||
These measurements describe the validated shapes rather than a general
|
||||
throughput guarantee.
|
||||
|
||||
@@ -4,17 +4,6 @@ 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) |
|
||||
@@ -22,8 +11,7 @@ Validated on:
|
||||
| Unified memory | 128 GB LPDDR5X (~121 GiB usable) |
|
||||
| GPUs | 1 |
|
||||
| Checkpoint storage | 135 GiB per partition |
|
||||
| Measured peak allocator high-water, T2VA (FP8) | 97.7 GiB |
|
||||
| Measured peak allocator high-water, Ref2VA (FP8) | 102.8 GiB |
|
||||
| Measured peak allocator high-water (FP8) | 97.7 GiB |
|
||||
|
||||
`FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one server
|
||||
at a time.
|
||||
@@ -41,21 +29,9 @@ 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, 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.
|
||||
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.
|
||||
|
||||
## One GB10: 960x576, 8 seconds
|
||||
|
||||
@@ -118,52 +94,6 @@ 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:
|
||||
@@ -172,69 +102,32 @@ 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`. 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.
|
||||
`X-Stage-Durations` response header of `/v1/videos/sync`.
|
||||
- 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]`. 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.
|
||||
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.
|
||||
- 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.
|
||||
|
||||
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
|
||||
Repeated identical 10-step 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, and 7200 s is
|
||||
shorter than a 50-step Ref2VA run.
|
||||
default 1800 s is shorter than a 50-step run on this platform.
|
||||
|
||||
## Known limitations
|
||||
|
||||
@@ -246,17 +139,11 @@ shorter than a 50-step Ref2VA run.
|
||||
- `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 notes above.
|
||||
- Sustained load degrades throughput. See the variance note above.
|
||||
- Only one partition is resident. For Ref2VA, stop the FL2VA server and restart
|
||||
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`.
|
||||
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.
|
||||
- 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.
|
||||
```
|
||||
@@ -776,26 +776,6 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
|
||||
omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
|
||||
|
||||
|
||||
def test_from_pipeline_config_routes_ltx2_model_extra(tmp_path):
|
||||
deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml"
|
||||
deploy_path.write_text(
|
||||
"\n".join(
|
||||
[
|
||||
"pipeline: dreamzero",
|
||||
"async_chunk: false",
|
||||
"stages:",
|
||||
" - stage_id: 0",
|
||||
" extras:",
|
||||
" ltx2_use_channels_last_3d: true",
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
|
||||
|
||||
assert stage.diffusion_config.extras["ltx2_use_channels_last_3d"] is True
|
||||
|
||||
|
||||
def test_structured_diffusion_config_rejects_invalid_compile_granularity():
|
||||
with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
|
||||
omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
|
||||
|
||||
@@ -9,10 +9,6 @@ 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
|
||||
|
||||
@@ -375,264 +371,6 @@ 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,9 +5,10 @@
|
||||
|
||||
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 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.
|
||||
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.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
@@ -56,13 +57,10 @@ def test_block_aligned_video_segment_runs_sparse():
|
||||
assert plan.latent_shape == list(ALIGNED_GRID)
|
||||
|
||||
|
||||
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)
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
|
||||
@@ -157,28 +155,3 @@ 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,14 +488,6 @@ 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__()
|
||||
|
||||
@@ -11,100 +11,6 @@ import torch
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class TestLTXChannelsLast3D:
|
||||
@pytest.mark.parametrize(
|
||||
("extras", "expected"),
|
||||
[({}, False), ({"ltx2_use_channels_last_3d": True}, True)],
|
||||
)
|
||||
def test_channels_last_opt_in_is_ltx2_model_extra(self, extras, expected):
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_channels_last_3d
|
||||
|
||||
assert _ltx2_use_channels_last_3d(SimpleNamespace(extras=extras)) is expected
|
||||
|
||||
def test_channels_last_opt_in_rejects_non_boolean_model_extra(self):
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_channels_last_3d
|
||||
|
||||
with pytest.raises(TypeError, match="ltx2_use_channels_last_3d"):
|
||||
_ltx2_use_channels_last_3d(SimpleNamespace(extras={"ltx2_use_channels_last_3d": "true"}))
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("causal", "kernel_size"),
|
||||
[
|
||||
(True, 3),
|
||||
(False, 3),
|
||||
(True, (3, 1, 1)),
|
||||
(True, 1),
|
||||
],
|
||||
)
|
||||
def test_layout_preserving_causal_conv_matches_original(self, causal, kernel_size):
|
||||
from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
|
||||
_enable_channels_last_3d,
|
||||
_weight_is_channels_last_3d,
|
||||
)
|
||||
|
||||
torch.manual_seed(0)
|
||||
conv = LTX2VideoCausalConv3d(4, 5, kernel_size).to(dtype=torch.float32).eval()
|
||||
hidden_states = torch.randn(1, 4, 5, 6, 7, dtype=torch.float32)
|
||||
|
||||
# Pointwise Conv3d storage is valid in both contiguous formats because
|
||||
# all spatial/temporal kernel dimensions are singleton.
|
||||
if kernel_size != 1:
|
||||
assert not _weight_is_channels_last_3d(conv)
|
||||
with torch.no_grad():
|
||||
expected = conv(hidden_states.clone(), causal=causal)
|
||||
|
||||
converted, patched = _enable_channels_last_3d(conv)
|
||||
|
||||
assert (converted, patched) == (1, 1)
|
||||
assert _weight_is_channels_last_3d(conv)
|
||||
with torch.no_grad():
|
||||
actual = conv(hidden_states.clone(), causal=causal)
|
||||
|
||||
torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4)
|
||||
|
||||
def test_temporal_pad_replicates_edges_in_channels_last_3d(self):
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
|
||||
_temporal_pad_channels_last_3d,
|
||||
)
|
||||
|
||||
torch.manual_seed(0)
|
||||
hidden_states = torch.randn(1, 4, 3, 2, 2, dtype=torch.float32)
|
||||
|
||||
actual = _temporal_pad_channels_last_3d(hidden_states, left=2, right=1)
|
||||
expected = torch.cat(
|
||||
[
|
||||
hidden_states[:, :, :1].repeat(1, 1, 2, 1, 1),
|
||||
hidden_states,
|
||||
hidden_states[:, :, -1:],
|
||||
],
|
||||
dim=2,
|
||||
)
|
||||
|
||||
assert actual.is_contiguous(memory_format=torch.channels_last_3d)
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
|
||||
def test_contiguous_weight_keeps_original_padding_path(self, monkeypatch):
|
||||
from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d
|
||||
|
||||
import vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 as ltx2_vae
|
||||
|
||||
conv = LTX2VideoCausalConv3d(4, 5, 3).to(dtype=torch.float32).eval()
|
||||
hidden_states = torch.randn(1, 4, 3, 2, 2, dtype=torch.float32)
|
||||
|
||||
def fail_channels_last_pad(*_args, **_kwargs):
|
||||
raise AssertionError("contiguous weights must retain the original padding path")
|
||||
|
||||
monkeypatch.setattr(ltx2_vae, "_temporal_pad_channels_last_3d", fail_channels_last_pad)
|
||||
|
||||
with torch.no_grad():
|
||||
expected = conv(hidden_states, causal=True)
|
||||
actual = ltx2_vae._layout_preserving_causal_conv_forward(conv, hidden_states, causal=True)
|
||||
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
class TestLTXOutputRank:
|
||||
@pytest.mark.parametrize(
|
||||
("distributed_vae_state", "expected_decode_calls"),
|
||||
|
||||
@@ -492,10 +492,6 @@ 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
|
||||
|
||||
@@ -530,48 +526,6 @@ 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,
|
||||
@@ -580,10 +534,6 @@ 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,7 +69,6 @@ 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
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -46,9 +47,111 @@ def test_model_display_names_require_string_mapping(tmp_path, monkeypatch):
|
||||
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"))
|
||||
def test_quantization_examples_keep_source_urls_under_feature_navigation(tmp_path, monkeypatch):
|
||||
nav_file = tmp_path / "docs" / ".nav.yml"
|
||||
nav_file.parent.mkdir(parents=True)
|
||||
nav_file.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"nav": [
|
||||
{
|
||||
"User Guide": [
|
||||
{
|
||||
"Examples": [
|
||||
"examples/README.md",
|
||||
{"Offline Inference": [{"Old": "user_guide/examples/offline_inference/old.md"}]},
|
||||
]
|
||||
},
|
||||
{"Features": [{"Quantization": [{"Overview": "user_guide/quantization/overview.md"}]}]},
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
sort_keys=False,
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert generate_examples.is_general_example(general_example)
|
||||
assert not generate_examples.is_general_example(model_example)
|
||||
monkeypatch.setattr(generate_examples, "ROOT_DIR", tmp_path)
|
||||
monkeypatch.setattr(generate_examples, "NAV_FILE", nav_file)
|
||||
monkeypatch.setattr(generate_examples, "EXAMPLE_DOC_DIR", tmp_path / "docs/user_guide/examples")
|
||||
|
||||
quantization_example = SimpleNamespace(
|
||||
category="quantization",
|
||||
path=tmp_path / "examples/quantization/check_modelopt_fp8_export.py",
|
||||
nav_title="Check Modelopt FP8 Export",
|
||||
)
|
||||
offline_example = SimpleNamespace(
|
||||
category="offline_inference",
|
||||
path=tmp_path / "examples/offline_inference/demo.py",
|
||||
nav_title="Demo",
|
||||
)
|
||||
|
||||
generate_examples.update_nav_file([quantization_example, offline_example])
|
||||
generate_examples.update_nav_file([quantization_example, offline_example])
|
||||
|
||||
nav_data = yaml.safe_load(nav_file.read_text(encoding="utf-8"))
|
||||
user_guide = nav_data["nav"][0]["User Guide"]
|
||||
examples = next(item["Examples"] for item in user_guide if "Examples" in item)
|
||||
features = next(item["Features"] for item in user_guide if "Features" in item)
|
||||
quantization = next(item["Quantization"] for item in features if "Quantization" in item)
|
||||
|
||||
assert not any(isinstance(item, dict) and "Quantization" in item for item in examples)
|
||||
generated_examples = [
|
||||
item[generate_examples.GENERATED_QUANTIZATION_EXAMPLES_TITLE]
|
||||
for item in quantization
|
||||
if generate_examples.GENERATED_QUANTIZATION_EXAMPLES_TITLE in item
|
||||
]
|
||||
assert len(generated_examples) == 1
|
||||
assert generated_examples[0] == [
|
||||
{"Check Modelopt FP8 Export": "user_guide/examples/quantization/check_modelopt_fp8_export.md"}
|
||||
]
|
||||
|
||||
# A later build without quantization examples must not leave a stale
|
||||
# generated group in the hand-authored Quantization section.
|
||||
generate_examples.update_nav_file([offline_example])
|
||||
nav_data = yaml.safe_load(nav_file.read_text(encoding="utf-8"))
|
||||
user_guide = nav_data["nav"][0]["User Guide"]
|
||||
features = next(item["Features"] for item in user_guide if "Features" in item)
|
||||
quantization = next(item["Quantization"] for item in features if "Quantization" in item)
|
||||
assert not any(
|
||||
isinstance(item, dict) and generate_examples.GENERATED_QUANTIZATION_EXAMPLES_TITLE in item
|
||||
for item in quantization
|
||||
)
|
||||
|
||||
|
||||
def test_quantization_examples_are_not_added_to_examples_without_feature_target(tmp_path, monkeypatch, caplog):
|
||||
nav_file = tmp_path / "docs" / ".nav.yml"
|
||||
nav_file.parent.mkdir(parents=True)
|
||||
nav_file.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"nav": [
|
||||
{
|
||||
"User Guide": [
|
||||
{"Examples": ["examples/README.md"]},
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
sort_keys=False,
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(generate_examples, "ROOT_DIR", tmp_path)
|
||||
monkeypatch.setattr(generate_examples, "NAV_FILE", nav_file)
|
||||
monkeypatch.setattr(generate_examples, "EXAMPLE_DOC_DIR", tmp_path / "docs/user_guide/examples")
|
||||
|
||||
quantization_example = SimpleNamespace(
|
||||
category="quantization",
|
||||
path=tmp_path / "examples/quantization/check_modelopt_fp8_export.py",
|
||||
nav_title="Check Modelopt FP8 Export",
|
||||
)
|
||||
|
||||
generate_examples.update_nav_file([quantization_example])
|
||||
|
||||
nav_data = yaml.safe_load(nav_file.read_text(encoding="utf-8"))
|
||||
examples = nav_data["nav"][0]["User Guide"][0]["Examples"]
|
||||
assert not any(isinstance(item, dict) and "Quantization" in item for item in examples)
|
||||
assert "omitting generated quantization examples from navigation" in caplog.text
|
||||
@@ -20,7 +20,6 @@ def test_default_stage_config_includes_cache_backend():
|
||||
"cache_backend": "cache_dit",
|
||||
"cache_config": '{"Fn_compute_blocks": 2}',
|
||||
"vae_use_slicing": True,
|
||||
"extras": {"ltx2_use_channels_last_3d": True},
|
||||
"ulysses_degree": 2,
|
||||
}
|
||||
)[0]
|
||||
@@ -30,7 +29,6 @@ def test_default_stage_config_includes_cache_backend():
|
||||
assert engine_args["cache_backend"] == "cache_dit"
|
||||
assert engine_args["cache_config"]["Fn_compute_blocks"] == 2
|
||||
assert engine_args["vae_use_slicing"] is True
|
||||
assert engine_args["extras"]["ltx2_use_channels_last_3d"] is True
|
||||
assert engine_args["parallel_config"].ulysses_degree == 2
|
||||
assert engine_args["model_stage"] == "diffusion"
|
||||
|
||||
|
||||
@@ -123,29 +123,18 @@ 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,
|
||||
notices: list[str],
|
||||
) -> list[str]:
|
||||
def _check(label: str, path: Path, actual: int, budget: int, sample: list[str], budget_name: 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:
|
||||
# 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 [
|
||||
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.",
|
||||
]
|
||||
return []
|
||||
|
||||
|
||||
@@ -163,7 +152,6 @@ 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()))
|
||||
|
||||
@@ -174,7 +162,6 @@ 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",
|
||||
@@ -183,11 +170,8 @@ 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,17 +20,6 @@ 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.
|
||||
@@ -132,17 +121,6 @@ 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,7 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
from functools import partial
|
||||
|
||||
import torch
|
||||
@@ -17,7 +16,6 @@ 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__)
|
||||
|
||||
@@ -26,15 +24,6 @@ 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
|
||||
@@ -426,36 +415,7 @@ 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]
|
||||
@@ -473,161 +433,3 @@ 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 — 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.
|
||||
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.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -292,6 +292,25 @@ 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 "
|
||||
|
||||
@@ -1,14 +1,11 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import MethodType
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from diffusers import AutoencoderKLLTX2Video
|
||||
from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d
|
||||
from diffusers.models.autoencoders.vae import DecoderOutput
|
||||
from torch import nn
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
|
||||
@@ -22,89 +19,6 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
# Adapted from https://github.com/sgl-project/sglang/pull/27431.
|
||||
def _weight_is_channels_last_3d(module: LTX2VideoCausalConv3d) -> bool:
|
||||
weight = module.conv.weight
|
||||
return (
|
||||
weight.dim() == 5
|
||||
and hasattr(torch, "channels_last_3d")
|
||||
and weight.is_contiguous(memory_format=torch.channels_last_3d)
|
||||
)
|
||||
|
||||
|
||||
def _temporal_pad_channels_last_3d(
|
||||
hidden_states: torch.Tensor,
|
||||
*,
|
||||
left: int,
|
||||
right: int,
|
||||
) -> torch.Tensor:
|
||||
"""Replicate temporal edge frames directly into CL3D storage."""
|
||||
batch, channels, frames, height, width = hidden_states.shape
|
||||
padded = torch.empty(
|
||||
(batch, channels, frames + left + right, height, width),
|
||||
dtype=hidden_states.dtype,
|
||||
device=hidden_states.device,
|
||||
memory_format=torch.channels_last_3d,
|
||||
)
|
||||
padded[:, :, left : left + frames].copy_(hidden_states)
|
||||
if left:
|
||||
padded[:, :, :left].copy_(hidden_states[:, :, :1])
|
||||
if right:
|
||||
padded[:, :, left + frames :].copy_(hidden_states[:, :, -1:])
|
||||
return padded
|
||||
|
||||
|
||||
def _layout_preserving_causal_conv_forward(
|
||||
module: LTX2VideoCausalConv3d,
|
||||
hidden_states: torch.Tensor,
|
||||
causal: bool = True,
|
||||
) -> torch.Tensor:
|
||||
"""Diffusers LTX causal-conv math with a CL3D-preserving pad path."""
|
||||
time_kernel_size = module.kernel_size[0]
|
||||
use_channels_last_pad = hidden_states.dim() == 5 and _weight_is_channels_last_3d(module)
|
||||
|
||||
if use_channels_last_pad:
|
||||
if causal:
|
||||
left, right = time_kernel_size - 1, 0
|
||||
else:
|
||||
left = right = (time_kernel_size - 1) // 2
|
||||
if left == 0 and right == 0:
|
||||
hidden_states = hidden_states.contiguous(memory_format=torch.channels_last_3d)
|
||||
else:
|
||||
hidden_states = _temporal_pad_channels_last_3d(
|
||||
hidden_states,
|
||||
left=left,
|
||||
right=right,
|
||||
)
|
||||
elif causal:
|
||||
pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, time_kernel_size - 1, 1, 1))
|
||||
hidden_states = torch.concatenate([pad_left, hidden_states], dim=2)
|
||||
else:
|
||||
pad_size = (time_kernel_size - 1) // 2
|
||||
pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, pad_size, 1, 1))
|
||||
pad_right = hidden_states[:, :, -1:, :, :].repeat((1, 1, pad_size, 1, 1))
|
||||
hidden_states = torch.concatenate([pad_left, hidden_states, pad_right], dim=2)
|
||||
|
||||
return module.conv(hidden_states)
|
||||
|
||||
|
||||
def _enable_channels_last_3d(model: nn.Module) -> tuple[int, int]:
|
||||
"""Convert LTX Conv3d weights and install layout-preserving causal pads."""
|
||||
if not hasattr(torch, "channels_last_3d"):
|
||||
return 0, 0
|
||||
|
||||
converted_convs = 0
|
||||
patched_causal_convs = 0
|
||||
for module in model.modules():
|
||||
if isinstance(module, LTX2VideoCausalConv3d):
|
||||
module.forward = MethodType(_layout_preserving_causal_conv_forward, module)
|
||||
patched_causal_convs += 1
|
||||
if isinstance(module, nn.Conv3d):
|
||||
module.weight.data = module.weight.data.contiguous(memory_format=torch.channels_last_3d)
|
||||
converted_convs += 1
|
||||
return converted_convs, patched_causal_convs
|
||||
|
||||
|
||||
class LTX2VaeExecutor(DistributedVaeExecutor):
|
||||
def _pack_local_tiles_without_meta(self, local_results, global_padding_shape, device, dtype):
|
||||
tile_tensor = torch.zeros(global_padding_shape, device=device, dtype=dtype)
|
||||
@@ -166,16 +80,6 @@ class DistributedAutoencoderKLLTX2Video(AutoencoderKLLTX2Video, DistributedVaeMi
|
||||
model.init_distributed()
|
||||
return model
|
||||
|
||||
def enable_channels_last_3d(self) -> None:
|
||||
"""Opt into the layout-preserving CL3D path for this video VAE."""
|
||||
converted_convs, patched_causal_convs = _enable_channels_last_3d(self)
|
||||
logger.info(
|
||||
"LTX2 video VAE: converted %d Conv3d weights to channels_last_3d and "
|
||||
"installed %d layout-preserving causal conv paths",
|
||||
converted_convs,
|
||||
patched_causal_convs,
|
||||
)
|
||||
|
||||
def init_distributed(self):
|
||||
self.distributed_executor = LTX2VaeExecutor()
|
||||
|
||||
|
||||
@@ -68,7 +68,6 @@ 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:
|
||||
@@ -108,19 +107,6 @@ 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,16 +193,7 @@ class RotaryEmbedding(CustomOp):
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
# 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
|
||||
return self.forward_native(x, cos, sin)
|
||||
|
||||
def forward_native(
|
||||
self,
|
||||
|
||||
@@ -57,17 +57,6 @@ _LTX_COMPONENT_SUBFOLDERS = (
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_LTX2_CHANNELS_LAST_3D_EXTRA = "ltx2_use_channels_last_3d"
|
||||
|
||||
|
||||
def _ltx2_use_channels_last_3d(od_config: Any) -> bool:
|
||||
"""Read the LTX2-only channels-last opt-in from model extras."""
|
||||
extras = getattr(od_config, "extras", {}) or {}
|
||||
enabled = extras.get(_LTX2_CHANNELS_LAST_3D_EXTRA, False)
|
||||
if not isinstance(enabled, bool):
|
||||
raise TypeError(f"{_LTX2_CHANNELS_LAST_3D_EXTRA} must be a bool, got {type(enabled)!r}")
|
||||
return enabled
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LTXComponentProfile:
|
||||
@@ -376,8 +365,6 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
|
||||
local_files_only=local_files_only,
|
||||
dtype=dtype,
|
||||
)
|
||||
if _ltx2_use_channels_last_3d(od_config):
|
||||
pipeline.vae.enable_channels_last_3d()
|
||||
pipeline.audio_vae = _load_component(
|
||||
AutoencoderKLLTX2Audio,
|
||||
model,
|
||||
|
||||
@@ -107,15 +107,6 @@ 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:
|
||||
@@ -342,7 +333,6 @@ 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)
|
||||
@@ -375,7 +365,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 = self.rot_dim
|
||||
rot_dim = freqs.shape[-1]
|
||||
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)
|
||||
@@ -411,16 +401,11 @@ 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 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()
|
||||
# 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
|
||||
)
|
||||
if used < packed_total and not no_mask:
|
||||
if used < packed_total and not prefix_slice:
|
||||
attn_mask = torch.arange(packed_total, device=q.device)[None] < used
|
||||
metadata = AttentionMetadata(
|
||||
attn_mask=attn_mask,
|
||||
@@ -430,15 +415,6 @@ 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,
|
||||
)
|
||||
|
||||
@@ -1069,7 +1069,6 @@ class AsyncOmniEngine:
|
||||
"streaming_output": kwargs.get("diffusion_streaming_output", False),
|
||||
"enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
|
||||
"extras": {
|
||||
**(kwargs.get("extras") or {}),
|
||||
"auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None),
|
||||
"default_llama_model_id": kwargs.get("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
|
||||
},
|
||||
|
||||
@@ -629,6 +629,7 @@ class OmniServeCommand(CLISubcommand):
|
||||
action="store_true",
|
||||
help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).",
|
||||
)
|
||||
|
||||
# Parallel weight loading (faster diffusion startup)
|
||||
omni_config_group.add_argument(
|
||||
"--disable-multithread-weight-load",
|
||||
|
||||
@@ -139,18 +139,6 @@ 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