Compare commits
9
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a3a3fbb781 | ||
|
|
dd6f32cb47 | ||
|
|
c27623c29c | ||
|
|
f8fa86fb1f | ||
|
|
422d8fd8f9 | ||
|
|
62e239bf1e | ||
|
|
740f45abcd | ||
|
|
5e4b9f8c6e | ||
|
|
ae811043e9 |
No files matched your search
+110
-28
@@ -1,6 +1,11 @@
|
||||
# CODEOWNERS is used to automatically assign PR reviewers
|
||||
# Format: path pattern @username or @team
|
||||
# More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
|
||||
#
|
||||
# The vllm_omni/** source-tree rules below are derived from the `owners` and
|
||||
# `primary_code_paths` front matter in docs/design/module/*.md. When updating
|
||||
# ownership, update the design doc first, then regenerate the matching rule here.
|
||||
# Last-match-wins: more specific rules must come AFTER broader ones.
|
||||
|
||||
# Default owners (fallback reviewers for all files)
|
||||
* @Gaohan123 @hsliuustc0106 @ywang96
|
||||
@@ -21,25 +26,10 @@
|
||||
# CI
|
||||
/.buildkite/ @yenuo26 @congw729 @NickCao
|
||||
|
||||
# Engine and model_executor
|
||||
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
|
||||
/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
|
||||
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks @NickCao
|
||||
/vllm_omni/model_executor/ @tzhouam @gcanlin
|
||||
/vllm_omni/worker/ @tzhouam @gcanlin
|
||||
# ComfyUI / external apps
|
||||
/apps/ @fhfuih @alex-jw-brooks
|
||||
|
||||
# Diffusion
|
||||
/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
|
||||
/vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
|
||||
|
||||
# Diffusion CPU offloading
|
||||
/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao @xuechendi
|
||||
|
||||
# Omni and TTS support
|
||||
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
|
||||
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin @NickCao
|
||||
|
||||
# Benchmarks
|
||||
# Benchmarks (out-of-tree)
|
||||
/benchmarks/accuracy/ @david6666666
|
||||
/benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
|
||||
/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
|
||||
@@ -47,20 +37,112 @@
|
||||
/benchmarks/glm_image/ @princepride @JaredforReal
|
||||
/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
|
||||
|
||||
# Usability and Observability
|
||||
/vllm_omni/deploy/ @tzhouam @lishunyang12
|
||||
/vllm_omni/metrics/ @lishunyang12
|
||||
/vllm_omni/profiler/ @gcanlin
|
||||
# =====================================================================
|
||||
# Module design ownership (docs/design/module/*.md)
|
||||
# Each section header cites the source design doc.
|
||||
# =====================================================================
|
||||
|
||||
# ComfyUI
|
||||
/apps/ @fhfuih @alex-jw-brooks
|
||||
# Autoregressive runtime — docs/design/module/ar_runtime.md
|
||||
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
|
||||
/vllm_omni/worker/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
|
||||
|
||||
# Large Scale Serving
|
||||
# Engine orchestration & stage runtime — docs/design/module/engine_orchestration.md, stage_runtime.md
|
||||
/vllm_omni/engine/ @tzhouam @fake0fan
|
||||
# File-level overrides: input/output modality contracts own these engine files
|
||||
# — docs/design/module/input_output_modality_contracts.md
|
||||
/vllm_omni/engine/messages.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
|
||||
/vllm_omni/engine/serialization.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
|
||||
/vllm_omni/engine/mm_outputs.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
|
||||
/vllm_omni/engine/output_modality.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
|
||||
/vllm_omni/engine/output_processor.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
|
||||
|
||||
# Input/output modality contracts — docs/design/module/input_output_modality_contracts.md
|
||||
/vllm_omni/inputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
|
||||
/vllm_omni/outputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
|
||||
/vllm_omni/request.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
|
||||
/vllm_omni/data_entry_keys.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
|
||||
|
||||
# Entrypoints and serving boundaries — docs/design/module/entrypoints.md
|
||||
/vllm_omni/entrypoints/ @alex-jw-brooks @linyueqian @NickCao
|
||||
# File-level overrides within entrypoints/
|
||||
/vllm_omni/entrypoints/openai/errors.py @alex-jw-brooks @NickCao
|
||||
/vllm_omni/entrypoints/cli/benchmark/ @alex-jw-brooks @Bounty-hunter
|
||||
|
||||
# Error contracts — docs/design/module/error_contracts.md
|
||||
/vllm_omni/errors.py @alex-jw-brooks @NickCao
|
||||
|
||||
# OmniConnector (inter-stage transport) — docs/design/module/omni_connector.md
|
||||
/vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
|
||||
|
||||
# Quantization and configuration
|
||||
# Model integration — docs/design/module/model_integration.md
|
||||
/vllm_omni/model_executor/ @tzhouam @gcanlin
|
||||
/vllm_omni/model_executor/models/ @tzhouam @gcanlin @Sy0307 @amy-why-3459 @linyueqian
|
||||
/vllm_omni/model_extras/ @tzhouam @gcanlin @princepride
|
||||
/vllm_omni/plugins/ @gcanlin @tjtanaa @xuechendi
|
||||
# Config override within model_executor — docs/design/module/vllm_omni_config.md
|
||||
/vllm_omni/model_executor/stage_configs/ @lishunyang12 @alex-jw-brooks
|
||||
|
||||
# Execution platforms (hardware backends) — docs/design/module/execution_platforms.md
|
||||
/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
|
||||
/vllm_omni/attention/ @gcanlin @tjtanaa @xuechendi @david6666666 @lishunyang12
|
||||
|
||||
# Configuration & deploy — docs/design/module/vllm_omni_config.md
|
||||
/vllm_omni/config/ @lishunyang12 @alex-jw-brooks
|
||||
/vllm_omni/deploy/ @lishunyang12 @alex-jw-brooks
|
||||
|
||||
# Quantization — docs/design/module/quantization.md
|
||||
/vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
|
||||
|
||||
# Hardware plugins and accelerators
|
||||
/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
|
||||
# Observability — docs/design/module/observability.md
|
||||
/vllm_omni/metrics/ @lishunyang12 @vraiti
|
||||
|
||||
# Profiling — docs/design/module/profiling.md
|
||||
/vllm_omni/profiler/ @gcanlin
|
||||
|
||||
# Benchmarking (in-tree) — docs/design/module/benchmarking.md
|
||||
/vllm_omni/benchmarks/ @alex-jw-brooks @Bounty-hunter
|
||||
|
||||
# LoRA (no dedicated design doc; ownership intentionally narrowed)
|
||||
/vllm_omni/lora/ @SamitHuang @wtomin
|
||||
|
||||
# Transformers utils (related path only in design docs; preserves prior ownership)
|
||||
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
|
||||
|
||||
# =====================================================================
|
||||
# Diffusion module family — docs/design/module/diffusion/
|
||||
# Umbrella rule first; sub-module rules after it win via last-match.
|
||||
# =====================================================================
|
||||
|
||||
# Diffusion umbrella (top-level files: data.py, forward_context.py, etc.)
|
||||
# — docs/design/module/diffusion/index.md
|
||||
/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
|
||||
|
||||
# Cache management — docs/design/module/cache_management.md
|
||||
/vllm_omni/diffusion/cache/ @Isotr0py @princepride @SamitHuang
|
||||
|
||||
# Diffusion quantization — docs/design/module/quantization.md
|
||||
/vllm_omni/diffusion/quantization/ @david6666666 @Isotr0py @lishunyang12
|
||||
|
||||
# Diffusion profiling — docs/design/module/profiling.md
|
||||
/vllm_omni/diffusion/profiler/ @gcanlin
|
||||
|
||||
# Diffusion runtime (scheduler, executor, worker)
|
||||
# — docs/design/module/diffusion/diffusion_runtime.md, continuous_batching.md
|
||||
/vllm_omni/diffusion/sched/ @Isotr0py @princepride @SamitHuang @fhfuih
|
||||
/vllm_omni/diffusion/executor/ @Isotr0py @princepride @SamitHuang @fhfuih
|
||||
/vllm_omni/diffusion/worker/ @Isotr0py @princepride @SamitHuang @fhfuih
|
||||
|
||||
# Diffusion model integration (pipelines, registration, checkpoint loading)
|
||||
# — docs/design/module/diffusion/diffusion_model_integration.md
|
||||
/vllm_omni/diffusion/models/ @fhfuih @Bounty-hunter @wtomin
|
||||
/vllm_omni/diffusion/model_loader/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
|
||||
/vllm_omni/diffusion/lora/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
|
||||
|
||||
# Diffusion parallelism (rank topology, TP/SP/PP/VAE patch)
|
||||
# — docs/design/module/diffusion/parallelism.md
|
||||
/vllm_omni/diffusion/distributed/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
|
||||
/vllm_omni/diffusion/attention/parallel/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
|
||||
|
||||
# Diffusion offloader (CPU offloading, DLO)
|
||||
# — docs/design/module/diffusion/offloader.md
|
||||
/vllm_omni/diffusion/offloader/ @yuanheng-zhao @david6666666 @lishunyang12
|
||||
@@ -66,7 +66,10 @@ repos:
|
||||
entry: python tools/pre_commit/check_tts_adapter.py
|
||||
language: python
|
||||
pass_filenames: false
|
||||
files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)$
|
||||
files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
|
||||
# The budget-is-now-slack reminder is printed on a PASSING run, and
|
||||
# pre-commit hides output from passing hooks unless the hook is verbose.
|
||||
verbose: true
|
||||
|
||||
# Keep `suggestion` last
|
||||
- id: suggestion
|
||||
|
||||
@@ -7,6 +7,7 @@ owners:
|
||||
- "@yinpeiqi"
|
||||
- "@fake0fan"
|
||||
- "@Sy0307"
|
||||
- "@Gaohan123"
|
||||
primary_code_paths:
|
||||
- vllm_omni/core/**
|
||||
- vllm_omni/worker/**
|
||||
|
||||
@@ -4,6 +4,7 @@ kind: module
|
||||
status: draft
|
||||
owners:
|
||||
- "@alex-jw-brooks"
|
||||
- "@Bounty-hunter"
|
||||
primary_code_paths:
|
||||
- vllm_omni/benchmarks/**
|
||||
- vllm_omni/entrypoints/cli/benchmark/**
|
||||
|
||||
@@ -5,6 +5,7 @@ status: draft
|
||||
owners:
|
||||
- "@Isotr0py"
|
||||
- "@princepride"
|
||||
- "@SamitHuang"
|
||||
primary_code_paths:
|
||||
- vllm_omni/diffusion/cache/**
|
||||
related_code_paths:
|
||||
|
||||
@@ -3,9 +3,10 @@ title: Diffusion Model Integration
|
||||
kind: module
|
||||
status: draft
|
||||
owners:
|
||||
- "@Isotr0py"
|
||||
- "@fhfuih"
|
||||
- "@Bounty-hunter"
|
||||
- "@wtomin"
|
||||
- "@ZJY0516"
|
||||
- "@RuixiangMa"
|
||||
primary_code_paths:
|
||||
- vllm_omni/diffusion/models/**
|
||||
- vllm_omni/diffusion/model_loader/**
|
||||
|
||||
@@ -6,6 +6,7 @@ owners:
|
||||
- "@Isotr0py"
|
||||
- "@princepride"
|
||||
- "@SamitHuang"
|
||||
- "@fhfuih"
|
||||
primary_code_paths:
|
||||
- vllm_omni/diffusion/sched/**
|
||||
- vllm_omni/diffusion/executor/**
|
||||
|
||||
@@ -4,7 +4,8 @@ kind: module
|
||||
status: draft
|
||||
owners:
|
||||
- "@yuanheng-zhao"
|
||||
- "@Isotr0py"
|
||||
- "@david6666666"
|
||||
- "@lishunyang12"
|
||||
primary_code_paths:
|
||||
- vllm_omni/diffusion/offloader/**
|
||||
related_code_paths:
|
||||
|
||||
@@ -6,6 +6,7 @@ owners:
|
||||
- "@princepride"
|
||||
- "@yuanheng-zhao"
|
||||
- "@xuechendi"
|
||||
- "@Bounty-hunter"
|
||||
primary_code_paths:
|
||||
- vllm_omni/diffusion/distributed/**
|
||||
- vllm_omni/diffusion/attention/parallel/**
|
||||
|
||||
@@ -6,6 +6,7 @@ architecture_state: current-plus-in-flight
|
||||
owners:
|
||||
- "@alex-jw-brooks"
|
||||
- "@linyueqian"
|
||||
- "@NickCao"
|
||||
document_stewards:
|
||||
- "@hsliuustc0106"
|
||||
- "@Gaohan123"
|
||||
|
||||
@@ -8,8 +8,15 @@ owners:
|
||||
- "@xuechendi"
|
||||
primary_code_paths:
|
||||
- vllm_omni/platforms/**
|
||||
related_code_paths:
|
||||
- vllm_omni/attention/**
|
||||
primary_code_path_owners:
|
||||
# Additional owners scoped to attention kernels (quantized/custom attention)
|
||||
- paths:
|
||||
- vllm_omni/attention/**
|
||||
owners:
|
||||
- "@david6666666"
|
||||
- "@lishunyang12"
|
||||
related_code_paths: []
|
||||
depends_on:
|
||||
- ar_runtime.md
|
||||
- diffusion/index.md
|
||||
|
||||
@@ -6,6 +6,8 @@ architecture_state: current
|
||||
owners:
|
||||
- "@Sy0307"
|
||||
- "@amy-why-3459"
|
||||
- "@Gaohan123"
|
||||
- "@alex-jw-brooks"
|
||||
document_stewards:
|
||||
- "@hsliuustc0106"
|
||||
- "@Gaohan123"
|
||||
|
||||
@@ -9,6 +9,14 @@ primary_code_paths:
|
||||
- vllm_omni/model_executor/**
|
||||
- vllm_omni/model_extras/**
|
||||
- vllm_omni/plugins/**
|
||||
primary_code_path_owners:
|
||||
# Additional owners scoped to a sub-path of model_executor/**
|
||||
- paths:
|
||||
- vllm_omni/model_executor/models/**
|
||||
owners:
|
||||
- "@Sy0307"
|
||||
- "@amy-why-3459"
|
||||
- "@linyueqian"
|
||||
related_code_paths:
|
||||
- vllm_omni/transformers_utils/**
|
||||
- vllm_omni/tokenizers/**
|
||||
|
||||
@@ -4,6 +4,7 @@ kind: module
|
||||
status: draft
|
||||
owners:
|
||||
- "@lishunyang12"
|
||||
- "@vraiti"
|
||||
primary_code_paths:
|
||||
- vllm_omni/metrics/**
|
||||
related_code_paths:
|
||||
|
||||
@@ -283,23 +283,20 @@ blocks — use Ulysses SP (`ring_degree=1`).
|
||||
|
||||
### Which geometries run sparse
|
||||
|
||||
Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count
|
||||
is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own
|
||||
tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows
|
||||
that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense
|
||||
attention and logs `RAINFUSION_ATTN staying dense` with the row count it computed — but it gets no
|
||||
speedup, so pick an aligned geometry when you want one.
|
||||
MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
|
||||
necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
|
||||
does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
|
||||
promoted rows remain visible to every query and key.
|
||||
|
||||
Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8
|
||||
spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact
|
||||
patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns
|
||||
are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially
|
||||
distant positions and selection can no longer rank them meaningfully. This is silent, and invisible
|
||||
to a `sparsity=0` check, because a fully populated mask does not care how blocks are grouped.
|
||||
For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
|
||||
128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
|
||||
MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
|
||||
and kernel tiling aligned, at the cost of retaining more blocks.
|
||||
|
||||
"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned
|
||||
resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`,
|
||||
which you can buy back with `start_step`.
|
||||
8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
|
||||
creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
|
||||
grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
|
||||
that implements protected video tails for rf_v2.
|
||||
|
||||
## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
|
||||
|
||||
|
||||
@@ -99,7 +99,7 @@ a native vLLM-Omni pipeline for those cases.
|
||||
### Attention Backends
|
||||
|
||||
The diffusers backend converts
|
||||
[vLLM-Omni standard of attention backend setting](../../../docs/user_guide/diffusion/attention_backends.md)
|
||||
[vLLM-Omni standard of attention backend setting](../../diffusion/attention_backends.md)
|
||||
to [diffusers standard](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends).
|
||||
|
||||
Specifically for `FLASH_ATTN`, it will first attempt to use FlashAttention-3 and then FlashAttention-2.
|
||||
|
||||
@@ -79,7 +79,8 @@ sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`.
|
||||
### Multi-NPU: 768P combined-service configuration
|
||||
|
||||
Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism
|
||||
degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload:
|
||||
degree 8, native tiled VAE patch parallelism degree 8, and distributed
|
||||
layerwise offload:
|
||||
|
||||
```bash
|
||||
export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
|
||||
@@ -98,7 +99,7 @@ vllm serve "${MODEL}" \
|
||||
--usp 8 \
|
||||
--ring 1 \
|
||||
--text-encoder-tp-size 8 \
|
||||
--enable-layerwise-offload \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--vae-parallel-mode tile \
|
||||
--vae-use-tiling \
|
||||
--vae-patch-parallel-size 8 \
|
||||
@@ -112,6 +113,11 @@ H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
|
||||
The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no
|
||||
partition restart is required. Layerwise offload applies to both DiTs.
|
||||
|
||||
On Atlas 800I A3 (64 GB HBM per device) the combined service does not fit at
|
||||
768P without offloading or sharding: use distributed layerwise offload (as
|
||||
above) or HSDP — see
|
||||
[§ Memory and attention optimizations](#memory-and-attention-optimizations-a3).
|
||||
|
||||
### Optional optimizations
|
||||
|
||||
Two independent optimizations may be enabled on top of the configuration
|
||||
@@ -134,6 +140,63 @@ Keep `--ring 1` when using RainFusion: the `rf_v2` kernel ranks key blocks
|
||||
over the whole sequence, so ring parallelism would split away the keys it
|
||||
needs. Scale with `--usp` instead.
|
||||
|
||||
## Memory and attention optimizations (A3)
|
||||
|
||||
### Fitting 768P into 64 GB HBM: distributed layerwise offload or HSDP
|
||||
|
||||
Each Atlas 800I A3 NPU has 64 GB of HBM, which is not enough for the
|
||||
combined MiniMax-H3 service at 768P. Enable **one** of the following at
|
||||
server startup:
|
||||
|
||||
**Distributed layerwise offload** — DiT layers are offloaded with parameters
|
||||
gathered across the parallel group instead of replicated per rank:
|
||||
|
||||
```bash
|
||||
--enable-distributed-layerwise-offload
|
||||
```
|
||||
|
||||
Measured peak NPU memory is about 45 GB per device for Ref2VA with a 13.88 s
|
||||
input video generating a 15 s 768P (1344x768) output video. Host (CPU) memory
|
||||
usage is high with this option.
|
||||
|
||||
**HSDP** — hybrid sharded data parallelism for the DiT parameters. The
|
||||
multi-stream memory-reuse knob is mandatory with this flag:
|
||||
|
||||
```bash
|
||||
export MULTI_STREAM_MEMORY_REUSE=2
|
||||
```
|
||||
|
||||
```bash
|
||||
--use-hsdp
|
||||
```
|
||||
|
||||
Host memory usage is small, but large shapes may still OOM; HBM usage
|
||||
optimizations for this path are ongoing.
|
||||
|
||||
### FLASH_ATTN backend with MindIE-SD
|
||||
|
||||
Keep `--diffusion-attention-backend FLASH_ATTN` and install MindIE-SD (see
|
||||
[§ Environment](#environment)). On NPU this backend carries most of the
|
||||
memory-reduction and performance work: a mask-free packed varlen path that
|
||||
never materializes the quadratic `full_qk` padding mask, and K/V prefix
|
||||
slicing, both driven by the packed `cu_seqlens` metadata emitted by the H3
|
||||
transformer.
|
||||
|
||||
### Optional: LaserAttention fused kernel
|
||||
|
||||
For an additional attention speedup, select the Ascend LaserAttention fused
|
||||
kernel before starting the server:
|
||||
|
||||
```bash
|
||||
export MINDIE_SD_FA_TYPE="ascend_laser_attention"
|
||||
```
|
||||
|
||||
This requires the FLASH_ATTN backend and MindIE-SD. H3 automatically applies
|
||||
exact power-of-two input pre-scaling (`laser_input_scale=256`) so the
|
||||
kernel's fp16 workspace cannot overflow on outlier activations. Measured
|
||||
speedup numbers will be added here.
|
||||
|
||||
|
||||
## HTTP API examples
|
||||
|
||||
The request format is identical to the GPU recipe; see
|
||||
@@ -154,8 +217,8 @@ configuration above:
|
||||
|
||||
| Workload | Configuration |
|
||||
|----------|---------------|
|
||||
| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
|
||||
| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
|
||||
| T2VA, 209 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
|
||||
| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
|
||||
|
||||
These measurements describe the validated shapes rather than a general
|
||||
throughput guarantee.
|
||||
|
||||
@@ -776,6 +776,26 @@ 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,6 +9,10 @@ This script tests two main scenarios:
|
||||
2. Case 2: Comparing FlashAttention and SDPA backends for batch_size=2 with padding
|
||||
"""
|
||||
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@@ -371,6 +375,264 @@ def test_packed_varlen_metadata_must_be_complete(monkeypatch):
|
||||
impl.forward_cuda(query, query, query, metadata)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# NPU mask-free packed attention (PR #5891).
|
||||
#
|
||||
# These tests cover the backend-side NPU branches that ``forward_fa_npu`` adds:
|
||||
# - boundary resolution in ``_resolve_packed_seq_npu`` (pure Python, no device
|
||||
# sync, no MindIE-SD import);
|
||||
# - env dispatch in ``forward_fa_npu`` (which op a packed forward takes);
|
||||
# - the laser input pre-scaling math in ``_forward_prefix_kv_slice_npu``.
|
||||
#
|
||||
# They run on CPU without a real NPU by injecting a fake ``mindiesd`` module
|
||||
# into ``sys.modules`` (same pattern as ``test_paged_attention`` /
|
||||
# ``benchmarks/conftest``). No production code is exercised for real kernels.
|
||||
|
||||
|
||||
def _npu_impl(causal: bool = False) -> FlashAttentionImpl:
|
||||
return FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=causal)
|
||||
|
||||
|
||||
def _fake_mindiesd(monkeypatch, *, attention_forward=None, attention_forward_varlen=None):
|
||||
"""Install a stub ``mindiesd`` so local ``from mindiesd import ...`` works.
|
||||
|
||||
Restored automatically by monkeypatch. Defaults are no-op callables so the
|
||||
top-level import in ``forward_fa_npu`` succeeds even when the test mocks the
|
||||
instance methods that would otherwise call them.
|
||||
"""
|
||||
fake = SimpleNamespace(
|
||||
attention_forward=attention_forward or Mock(return_value=torch.zeros(1)),
|
||||
attention_forward_varlen=attention_forward_varlen or Mock(return_value=torch.zeros(1)),
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "mindiesd", fake)
|
||||
return fake
|
||||
|
||||
|
||||
# --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
|
||||
|
||||
|
||||
def test_resolve_packed_seq_accepts_real_plus_pad_contract():
|
||||
impl = _npu_impl()
|
||||
q = torch.randn(1, 8, 2, 4)
|
||||
cu = torch.tensor([0, 5, 8], dtype=torch.int32)
|
||||
extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 5, "max_seqlen_k": 5}
|
||||
|
||||
assert impl._resolve_packed_seq_npu(q, q, extra) == ([5, 8], [5, 8])
|
||||
|
||||
|
||||
def test_resolve_packed_seq_single_document_no_padding():
|
||||
impl = _npu_impl()
|
||||
q = torch.randn(1, 8, 2, 4)
|
||||
cu = torch.tensor([0, 8], dtype=torch.int32)
|
||||
extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 8, "max_seqlen_k": 8}
|
||||
|
||||
assert impl._resolve_packed_seq_npu(q, q, extra) == ([8], [8])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"case",
|
||||
[
|
||||
"causal",
|
||||
"missing_key",
|
||||
"batch_ne_1",
|
||||
"cu_too_long",
|
||||
"cu_qk_mismatch",
|
||||
"max_seqlen_not_int",
|
||||
"used_gt_total",
|
||||
"used_zero",
|
||||
"pad_longer_than_real",
|
||||
],
|
||||
)
|
||||
def test_resolve_packed_seq_rejects_malformed_metadata(case):
|
||||
"""Any metadata outside the exact [real, pad] contract returns None so the
|
||||
caller falls back to the masked path."""
|
||||
cu3 = torch.tensor([0, 5, 8], dtype=torch.int32)
|
||||
cu2 = torch.tensor([0, 8], dtype=torch.int32)
|
||||
cu4 = torch.tensor([0, 3, 5, 8], dtype=torch.int32)
|
||||
cu_short = torch.tensor([0, 3, 8], dtype=torch.int32)
|
||||
base = {"cu_seqlens_q": cu3, "cu_seqlens_k": cu3, "max_seqlen_q": 5, "max_seqlen_k": 5}
|
||||
q = torch.randn(1, 8, 2, 4)
|
||||
|
||||
impl = _npu_impl(causal=(case == "causal"))
|
||||
if case == "missing_key":
|
||||
extra = {"cu_seqlens_q": cu3, "max_seqlen_q": 5} # 3 of 4 keys absent
|
||||
elif case == "batch_ne_1":
|
||||
q = torch.randn(2, 8, 2, 4)
|
||||
extra = base
|
||||
elif case == "cu_too_long":
|
||||
extra = {**base, "cu_seqlens_q": cu4, "cu_seqlens_k": cu4}
|
||||
elif case == "cu_qk_mismatch":
|
||||
extra = {**base, "cu_seqlens_q": cu3, "cu_seqlens_k": cu2}
|
||||
elif case == "max_seqlen_not_int":
|
||||
extra = {**base, "max_seqlen_q": 5.0}
|
||||
elif case == "used_gt_total":
|
||||
extra = {**base, "max_seqlen_q": 9} # 9 > total length 8
|
||||
elif case == "used_zero":
|
||||
extra = {**base, "max_seqlen_q": 0, "max_seqlen_k": 0}
|
||||
elif case == "pad_longer_than_real":
|
||||
# real=3, pad=5: the real document must be the longer one.
|
||||
extra = {
|
||||
"cu_seqlens_q": cu_short,
|
||||
"cu_seqlens_k": cu_short,
|
||||
"max_seqlen_q": 3,
|
||||
"max_seqlen_k": 3,
|
||||
}
|
||||
else:
|
||||
extra = base
|
||||
|
||||
assert impl._resolve_packed_seq_npu(q, q, extra) is None
|
||||
|
||||
|
||||
# --- Test group B: env dispatch in forward_fa_npu (current behavior) --------
|
||||
|
||||
|
||||
def _npu_impl_with_mocked_paths() -> tuple[FlashAttentionImpl, torch.Tensor]:
|
||||
"""FlashAttentionImpl whose two packed paths are Mocks returning a sentinel.
|
||||
|
||||
Lets a test assert which branch ``forward_fa_npu`` dispatched to without
|
||||
running any real kernel.
|
||||
"""
|
||||
impl = _npu_impl()
|
||||
sentinel = torch.tensor([1.0])
|
||||
impl._forward_varlen_packed_npu = Mock(return_value=sentinel)
|
||||
impl._forward_prefix_kv_slice_npu = Mock(return_value=sentinel)
|
||||
return impl, sentinel
|
||||
|
||||
|
||||
def test_npu_env_unset_uses_varlen(monkeypatch):
|
||||
monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
|
||||
_fake_mindiesd(monkeypatch)
|
||||
impl, _ = _npu_impl_with_mocked_paths()
|
||||
metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
|
||||
|
||||
impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
|
||||
|
||||
impl._forward_varlen_packed_npu.assert_called_once()
|
||||
impl._forward_prefix_kv_slice_npu.assert_not_called()
|
||||
|
||||
|
||||
def test_npu_env_laser_uses_slice(monkeypatch):
|
||||
monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
|
||||
_fake_mindiesd(monkeypatch)
|
||||
impl, _ = _npu_impl_with_mocked_paths()
|
||||
metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
|
||||
|
||||
impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
|
||||
|
||||
impl._forward_prefix_kv_slice_npu.assert_called_once()
|
||||
impl._forward_varlen_packed_npu.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fa_type", ["prompt_flash_attn", "fused_attn_score", "ascend_laser_attenton", "garbage"])
|
||||
def test_npu_env_non_laser_value_falls_back_to_varlen(monkeypatch, fa_type):
|
||||
"""Locks in the CURRENT dispatch behavior (PR #5891).
|
||||
|
||||
Any value other than ``ascend_laser_attention`` — including the valid
|
||||
``prompt_flash_attn``/``fused_attn_score`` and typos — is silently routed to
|
||||
the TND varlen op, which cannot honor ``MINDIE_SD_FA_TYPE``. This is review
|
||||
comment #1 on the PR and is intentionally captured here as a regression
|
||||
net: update this test when the dispatch starts honoring/rejecting these.
|
||||
"""
|
||||
monkeypatch.setenv("MINDIE_SD_FA_TYPE", fa_type)
|
||||
_fake_mindiesd(monkeypatch)
|
||||
impl, _ = _npu_impl_with_mocked_paths()
|
||||
metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
|
||||
|
||||
impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
|
||||
|
||||
impl._forward_varlen_packed_npu.assert_called_once()
|
||||
impl._forward_prefix_kv_slice_npu.assert_not_called()
|
||||
|
||||
|
||||
def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch):
|
||||
"""Models that do not set ``npu_attn_varlen`` (e.g. Wan/Cosmos) keep using
|
||||
the existing masked attention_forward path; the new branches are skipped."""
|
||||
fake_forward = Mock(return_value=torch.zeros(1, 8, 2, 4))
|
||||
_fake_mindiesd(monkeypatch, attention_forward=fake_forward)
|
||||
impl, _ = _npu_impl_with_mocked_paths()
|
||||
metadata = AttentionMetadata(
|
||||
attn_mask=torch.tensor([[True, True, True, True, True, False, False, False]]),
|
||||
extra={}, # no npu_attn_varlen opt-in
|
||||
)
|
||||
q = torch.randn(1, 8, 2, 4)
|
||||
|
||||
out = impl.forward_fa_npu(q, q, q, metadata)
|
||||
|
||||
impl._forward_varlen_packed_npu.assert_not_called()
|
||||
impl._forward_prefix_kv_slice_npu.assert_not_called()
|
||||
fake_forward.assert_called_once()
|
||||
assert out is fake_forward.return_value
|
||||
|
||||
|
||||
# --- Test group C: laser input pre-scaling in _forward_prefix_kv_slice_npu --
|
||||
|
||||
|
||||
def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch):
|
||||
monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
|
||||
captured: dict = {}
|
||||
|
||||
def fake_attention_forward(query, key, value, **kwargs):
|
||||
captured["query"] = query
|
||||
captured["key"] = key
|
||||
captured["value"] = value
|
||||
captured.update(kwargs)
|
||||
return torch.ones_like(query) # deterministic stand-in kernel output
|
||||
|
||||
_fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
|
||||
impl = _npu_impl() # softmax_scale == 0.5
|
||||
q = torch.randn(1, 8, 2, 4)
|
||||
cu = torch.tensor([0, 5, 8], dtype=torch.int32)
|
||||
extra = {
|
||||
"cu_seqlens_q": cu,
|
||||
"cu_seqlens_k": cu,
|
||||
"max_seqlen_q": 5,
|
||||
"max_seqlen_k": 5,
|
||||
"laser_input_scale": 256.0,
|
||||
}
|
||||
|
||||
out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
|
||||
|
||||
# K/V sliced to the valid prefix (used_k == 5); q keeps full length.
|
||||
assert captured["key"].shape == (1, 5, 2, 4)
|
||||
assert captured["value"].shape == (1, 5, 2, 4)
|
||||
assert captured["query"].shape == (1, 8, 2, 4)
|
||||
assert captured["attn_mask"] is None
|
||||
# Pre-divide q/k/v by the factor and compensate the kernel scale by its
|
||||
# square (exact for power-of-two); output scaled back by the factor.
|
||||
torch.testing.assert_close(captured["query"], q / 256.0)
|
||||
torch.testing.assert_close(captured["key"], q[:, :5] / 256.0)
|
||||
assert captured["scale"] == pytest.approx(0.5 * 256.0**2)
|
||||
torch.testing.assert_close(out, torch.ones(1, 8, 2, 4) * 256.0)
|
||||
|
||||
|
||||
def test_prefix_kv_slice_no_scaling_without_factor(monkeypatch):
|
||||
monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
|
||||
captured: dict = {}
|
||||
|
||||
def fake_attention_forward(query, key, value, **kwargs):
|
||||
captured["query"] = query
|
||||
captured.update(kwargs)
|
||||
return torch.ones_like(query)
|
||||
|
||||
_fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
|
||||
impl = _npu_impl()
|
||||
q = torch.randn(1, 8, 2, 4)
|
||||
cu = torch.tensor([0, 5, 8], dtype=torch.int32)
|
||||
extra = {
|
||||
"cu_seqlens_q": cu,
|
||||
"cu_seqlens_k": cu,
|
||||
"max_seqlen_q": 5,
|
||||
"max_seqlen_k": 5,
|
||||
# no laser_input_scale → no pre-scaling
|
||||
}
|
||||
|
||||
out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
|
||||
|
||||
torch.testing.assert_close(captured["query"], q)
|
||||
assert captured["scale"] == pytest.approx(0.5)
|
||||
torch.testing.assert_close(out, torch.ones(1, 8, 2, 4))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Running FlashAttention Padding Tests...")
|
||||
print("=" * 60)
|
||||
|
||||
@@ -5,10 +5,9 @@
|
||||
|
||||
The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip:
|
||||
1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row
|
||||
kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the
|
||||
first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
|
||||
mask from the video and prefix segments separately, so the two only agree on the
|
||||
block count when the video segment ends on a boundary.
|
||||
kernel block, and 1344x768 gives 62x24x42 whose 62496 rows do not. Both are
|
||||
handed to MindIE-SD: its rf_v2 preprocessing promotes an irregular video tail
|
||||
to the always-kept prefix segment before generating the block mask.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
@@ -57,10 +56,13 @@ def test_block_aligned_video_segment_runs_sparse():
|
||||
assert plan.latent_shape == list(ALIGNED_GRID)
|
||||
|
||||
|
||||
def test_misaligned_video_segment_stays_dense():
|
||||
# Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask:
|
||||
# pad keys would take a share of every real query's softmax denominator.
|
||||
assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None
|
||||
def test_misaligned_video_segment_runs_sparse_with_updated_mindiesd():
|
||||
plan = make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID))
|
||||
|
||||
assert plan is not None
|
||||
assert plan.prefix_len == PREFIX_ROWS
|
||||
assert plan.used_len == PREFIX_ROWS + 62496
|
||||
assert plan.latent_shape == list(MISALIGNED_GRID)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
|
||||
@@ -155,3 +157,28 @@ def test_fully_populated_mask_reproduces_dense_attention():
|
||||
).transpose(1, 2)
|
||||
error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
|
||||
assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.")
|
||||
def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd():
|
||||
"""vLLM forwards the full irregular grid; MindIE-SD owns its tail handling."""
|
||||
torch.manual_seed(1)
|
||||
prefix_len, grid = 384, (4, 16, 10)
|
||||
video_len = grid[0] * grid[1] * grid[2]
|
||||
used = prefix_len + video_len
|
||||
heads, head_dim = 4, 128
|
||||
q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
|
||||
|
||||
impl = make_impl()
|
||||
impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
|
||||
plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
|
||||
out = impl._forward_sparse_npu(q, k, v, plan)
|
||||
|
||||
reference = torch.nn.functional.scaled_dot_product_attention(
|
||||
q.transpose(1, 2),
|
||||
k.transpose(1, 2),
|
||||
v.transpose(1, 2),
|
||||
scale=impl.softmax_scale,
|
||||
).transpose(1, 2)
|
||||
error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
|
||||
assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
|
||||
@@ -488,6 +488,14 @@ class _MiniMaxH3FakeAttention(nn.Module):
|
||||
class FakeBackend:
|
||||
supports_prefix_kv_slicing = False
|
||||
|
||||
# Mirror AttentionBackend.supports_packed_mask_free: H3's packed
|
||||
# attention path calls attn_backend.supports_packed_mask_free() to decide
|
||||
# whether to build the padding mask. Return False so the fake exercises
|
||||
# the masked path (the abstract backend default).
|
||||
@classmethod
|
||||
def supports_packed_mask_free(cls) -> bool:
|
||||
return False
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
del kwargs
|
||||
super().__init__()
|
||||
|
||||
@@ -11,6 +11,100 @@ 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,6 +492,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
|
||||
class FakeBackend:
|
||||
supports_prefix_kv_slicing = True
|
||||
|
||||
@classmethod
|
||||
def supports_packed_mask_free(cls) -> bool:
|
||||
return False
|
||||
|
||||
class FakeAttention(torch.nn.Module):
|
||||
attn_backend = FakeBackend
|
||||
|
||||
@@ -526,6 +530,48 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
|
||||
assert attention.attention.metadata.extra["valid_kv_length"] == 5
|
||||
|
||||
|
||||
def test_packed_attention_skips_mask_for_packed_mask_free_backend():
|
||||
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
||||
MiniMaxH3Attention,
|
||||
)
|
||||
|
||||
class FakeBackend:
|
||||
supports_prefix_kv_slicing = False
|
||||
|
||||
@classmethod
|
||||
def supports_packed_mask_free(cls) -> bool:
|
||||
return True
|
||||
|
||||
class FakeAttention(torch.nn.Module):
|
||||
attn_backend = FakeBackend
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.metadata = None
|
||||
|
||||
def forward(self, query, key, value, metadata):
|
||||
self.metadata = metadata
|
||||
return query
|
||||
|
||||
attention = object.__new__(MiniMaxH3Attention)
|
||||
torch.nn.Module.__init__(attention)
|
||||
attention.attention = FakeAttention()
|
||||
q = torch.randn(8, 2, 4)
|
||||
|
||||
attention._run_packed_attention(
|
||||
q,
|
||||
q,
|
||||
q,
|
||||
cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
|
||||
max_seqlen=5,
|
||||
packed_total=8,
|
||||
)
|
||||
|
||||
assert attention.attention.metadata.attn_mask is None
|
||||
assert attention.attention.metadata.extra["valid_kv_length"] == 5
|
||||
assert attention.attention.metadata.extra["npu_attn_varlen"] is True
|
||||
|
||||
|
||||
def test_packed_attention_keeps_padding_mask_for_other_backends():
|
||||
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
||||
MiniMaxH3Attention,
|
||||
@@ -534,6 +580,10 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
|
||||
class FakeBackend:
|
||||
supports_prefix_kv_slicing = False
|
||||
|
||||
@classmethod
|
||||
def supports_packed_mask_free(cls) -> bool:
|
||||
return False
|
||||
|
||||
class FakeAttention(torch.nn.Module):
|
||||
attn_backend = FakeBackend
|
||||
|
||||
|
||||
@@ -20,6 +20,7 @@ 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]
|
||||
@@ -29,6 +30,7 @@ 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,18 +123,29 @@ def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]:
|
||||
return found
|
||||
|
||||
|
||||
def _check(label: str, path: Path, actual: int, budget: int, sample: list[str], budget_name: str) -> list[str]:
|
||||
def _check(
|
||||
label: str,
|
||||
path: Path,
|
||||
actual: int,
|
||||
budget: int,
|
||||
sample: list[str],
|
||||
budget_name: str,
|
||||
notices: list[str],
|
||||
) -> list[str]:
|
||||
if actual > budget:
|
||||
listing = "\n".join(f" {line}" for line in sample)
|
||||
return [
|
||||
f"{path}: {actual} {label} exceeds the budget of {budget}.\n{listing}\n\n{_GUIDANCE}\n",
|
||||
]
|
||||
if actual < budget:
|
||||
return [
|
||||
f"{path}: {actual} {label}, down from a budget of {budget}. "
|
||||
f"Lower {budget_name} to {actual} in {Path(__file__).relative_to(REPO_ROOT)} "
|
||||
f"so the ratchet holds the new ground.",
|
||||
]
|
||||
# Progress is never a build failure. Removing branches is the goal, so a
|
||||
# PR that does it must not hand-edit this file to stay green -- nor break
|
||||
# `main` for everyone else when it merges from a stale base.
|
||||
notices.append(
|
||||
f"{path}: {actual} {label}, below the budget of {budget}. "
|
||||
f"Lower {budget_name} to {actual} in "
|
||||
f"{Path(__file__).relative_to(REPO_ROOT)} to hold the new ground."
|
||||
)
|
||||
return []
|
||||
|
||||
|
||||
@@ -152,6 +163,7 @@ def main(argv: list[str] | None = None) -> int:
|
||||
print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
notices: list[str] = []
|
||||
branches = _model_type_branches(ast.parse(serving_path.read_text()))
|
||||
detectors = _legacy_detectors(ast.parse(adapters_path.read_text()))
|
||||
|
||||
@@ -162,6 +174,7 @@ def main(argv: list[str] | None = None) -> int:
|
||||
MAX_MODEL_TYPE_BRANCHES,
|
||||
[f"line {lineno}: {src}" for lineno, src in branches],
|
||||
"MAX_MODEL_TYPE_BRANCHES",
|
||||
notices,
|
||||
)
|
||||
errors += _check(
|
||||
"unmigrated model types in LEGACY_TTS_DETECTORS",
|
||||
@@ -170,8 +183,11 @@ def main(argv: list[str] | None = None) -> int:
|
||||
MAX_LEGACY_DETECTORS,
|
||||
[f"line {lineno}: {name}" for lineno, name in detectors],
|
||||
"MAX_LEGACY_DETECTORS",
|
||||
notices,
|
||||
)
|
||||
|
||||
for notice in notices:
|
||||
print(f"check_tts_adapter: {notice}", file=sys.stderr)
|
||||
for error in errors:
|
||||
print(f"check_tts_adapter: {error}", file=sys.stderr)
|
||||
return 1 if errors else 0
|
||||
|
||||
@@ -20,6 +20,17 @@ class AttentionBackend(ABC):
|
||||
# avoid a slower masked-attention plan when tail padding is not semantic.
|
||||
supports_prefix_kv_slicing: bool = False
|
||||
|
||||
@classmethod
|
||||
def supports_packed_mask_free(cls) -> bool:
|
||||
"""Whether packed attention never reads attn_mask on this platform.
|
||||
|
||||
When True, models that pack a [real, pad] two-document layout and
|
||||
carry cu_seqlens/max_seqlen in ``AttentionMetadata.extra`` may skip
|
||||
constructing the padding mask entirely. Backends whose mask-free
|
||||
behavior is platform-dependent must check current_omni_platform.
|
||||
"""
|
||||
return False
|
||||
|
||||
# ``OmniPlatformEnum`` values this backend runs on; None means unrestricted.
|
||||
# Platform resolution rejects an explicit selection outside this set, so a
|
||||
# hardware-specific backend fails before the model is built.
|
||||
@@ -121,6 +132,17 @@ class AttentionMetadata:
|
||||
# the packed cu_seqlens tensors.
|
||||
# "valid_kv_length": int — contiguous valid K/V prefix length for a
|
||||
# backend that advertises supports_prefix_kv_slicing.
|
||||
# "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path
|
||||
# (TND npu_fusion_attention driven by cu_seqlens, mask never read).
|
||||
# Requires the [real, pad] two-document packing contract; see
|
||||
# FlashAttentionImpl._forward_varlen_packed_npu.
|
||||
# "laser_input_scale": float — model opt-in input pre-scale for the NPU
|
||||
# ascend_laser_attention path. The kernel stores unscaled QK^T in an
|
||||
# fp16 workspace, so outlier activations overflow 65504 into NaN rows;
|
||||
# with this set (>1), q/k/v are divided by the factor before the op,
|
||||
# the kernel scale_value is multiplied by its square, and the output
|
||||
# is scaled back (exact for power-of-two factors). Absent means no
|
||||
# pre-scaling. See FlashAttentionImpl._forward_prefix_kv_slice_npu.
|
||||
|
||||
# Piecewise attention metadata (mixed causal/full masks).
|
||||
# full_attn_spans: per-sample [start, end) spans in global coordinates using full attention.
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
from functools import partial
|
||||
|
||||
import torch
|
||||
@@ -16,6 +17,7 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
|
||||
piecewise_attn,
|
||||
)
|
||||
from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -24,6 +26,15 @@ class FlashAttentionBackend(AttentionBackend):
|
||||
accept_output_buffer: bool = True
|
||||
supports_piecewise_spans: bool = True
|
||||
|
||||
@classmethod
|
||||
def supports_packed_mask_free(cls) -> bool:
|
||||
# CUDA: forward_cuda dispatches the packed cu_seqlens varlen path
|
||||
# before ever reading attn_mask. NPU: forward_fa_npu honors the
|
||||
# npu_attn_varlen opt-in and rebuilds the padding mask itself if the
|
||||
# packed contract fails. XPU reads attn_mask, so models must keep
|
||||
# constructing it there.
|
||||
return current_omni_platform.is_cuda() or current_omni_platform.is_npu()
|
||||
|
||||
@classmethod
|
||||
def supports_attention_mask(cls) -> bool:
|
||||
return True
|
||||
@@ -415,7 +426,36 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
"For installation details, see https://gitcode.com/Ascend/MindIE-SD"
|
||||
"Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
|
||||
)
|
||||
# Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the
|
||||
# model marks extra["npu_attn_varlen"] and carries packed metadata, so
|
||||
# the padding document is excluded without reading or materializing
|
||||
# the quadratic full_qk mask.
|
||||
# - default: packed varlen via mindiesd attention_forward_varlen;
|
||||
# - when MINDIE_SD_FA_TYPE=ascend_laser_attention: prefix K/V
|
||||
# slicing via mindiesd attention_forward (so the env-selected op
|
||||
# type is honored; the TND varlen op cannot be).
|
||||
extra = attn_metadata.extra if attn_metadata else {}
|
||||
if extra.get("npu_attn_varlen", False):
|
||||
if os.environ.get("MINDIE_SD_FA_TYPE") == "ascend_laser_attention":
|
||||
out = self._forward_prefix_kv_slice_npu(query, key, value, extra)
|
||||
else:
|
||||
out = self._forward_varlen_packed_npu(query, key, value, extra)
|
||||
if out is not None:
|
||||
return out
|
||||
attention_mask = attn_metadata.attn_mask if attn_metadata else None
|
||||
if attention_mask is None and extra.get("npu_attn_varlen", False):
|
||||
# Models skip mask construction when this opt-in is set
|
||||
# (FlashAttentionBackend.supports_packed_mask_free). The packed
|
||||
# paths above declined (contract mismatch), so rebuild the padding
|
||||
# mask here to keep the masked fallback correct.
|
||||
used = extra.get("valid_kv_length")
|
||||
if not isinstance(used, int) or not 0 < used <= query.shape[1]:
|
||||
raise ValueError(
|
||||
"npu_attn_varlen packed metadata is unusable and no attn_mask "
|
||||
f"was constructed (valid_kv_length={used!r}, seq_len={query.shape[1]}); "
|
||||
"refusing to run unmasked attention over padding rows."
|
||||
)
|
||||
attention_mask = torch.arange(query.shape[1], device=query.device)[None] < used
|
||||
|
||||
# NPU aclnnFlashAttentionScore requires mask shape to be one of:
|
||||
# [B, N, Sq, Skv], [B, 1, Sq, Skv], [1, 1, Sq, Skv], or [Sq, Skv]
|
||||
@@ -433,3 +473,161 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
op_type="fused_attn_score",
|
||||
layout=layout,
|
||||
)
|
||||
|
||||
def _resolve_packed_seq_npu(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
extra: dict,
|
||||
) -> tuple[list[int], list[int]] | None:
|
||||
"""Resolve packed document boundaries (cumulative end offsets per doc).
|
||||
|
||||
Returns (seq_q, seq_k), e.g. ([used_q, total_q], [used_k, total_k]),
|
||||
or None unless the metadata matches the exact contract MiniMax-H3
|
||||
produces:
|
||||
- q/k/v are [1, T, N, D] (BSND, single packed batch);
|
||||
- cu_seqlens describe a [real, pad] two-document packing, with the
|
||||
padding document as a strict suffix;
|
||||
- max_seqlen_q/k are Python ints equal to the real document length
|
||||
(so boundaries are derived without any device sync).
|
||||
"""
|
||||
if self.causal:
|
||||
return None
|
||||
packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
|
||||
if any(k not in extra for k in packed_keys):
|
||||
return None
|
||||
cu_q, cu_k = extra["cu_seqlens_q"], extra["cu_seqlens_k"]
|
||||
used_q, used_k = extra["max_seqlen_q"], extra["max_seqlen_k"]
|
||||
if query.shape[0] != 1 or key.shape[0] != 1:
|
||||
return None
|
||||
if not isinstance(used_q, int) or not isinstance(used_k, int):
|
||||
return None
|
||||
total_q, total_k = query.shape[1], key.shape[1]
|
||||
# .shape is host-side metadata: counting documents never syncs.
|
||||
if cu_q.shape[0] != cu_k.shape[0] or cu_q.shape[0] > 3:
|
||||
return None
|
||||
if not (0 < used_q <= total_q) or not (0 < used_k <= total_k):
|
||||
return None
|
||||
if used_q == total_q and used_k == total_k:
|
||||
# No padding document: one full-length document.
|
||||
return [total_q], [total_k]
|
||||
if cu_q.shape[0] == 3 and used_q >= total_q - used_q and used_k >= total_k - used_k:
|
||||
# [real, pad] packing: the real document must be the longer one
|
||||
# (consistent with the max_seqlen naming).
|
||||
return [used_q, total_q], [used_k, total_k]
|
||||
return None
|
||||
|
||||
def _forward_varlen_packed_npu(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
extra: dict,
|
||||
) -> torch.Tensor | None:
|
||||
"""Packed varlen attention on NPU via mindiesd attention_forward_varlen.
|
||||
|
||||
Returns None (caller falls back to the mask path) when the packed
|
||||
contract does not hold; see _resolve_packed_seq_npu.
|
||||
"""
|
||||
resolved = self._resolve_packed_seq_npu(query, key, extra)
|
||||
if resolved is None:
|
||||
return None
|
||||
seq_q, seq_k = resolved
|
||||
|
||||
try:
|
||||
from mindiesd import attention_forward_varlen
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"FlashAttentionBackend NPU implementation requires MindIE-SD. "
|
||||
"Please install MindIE-SD to enable NPU attention support. "
|
||||
"For installation details, see https://gitcode.com/Ascend/MindIE-SD"
|
||||
"Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
|
||||
)
|
||||
q = query.squeeze(0) # [T, N, D] == TND
|
||||
k = key.squeeze(0)
|
||||
v = value.squeeze(0)
|
||||
# attention_forward_varlen wants cu_seqlens as host lists of cumulative
|
||||
# offsets and forwards cu[1:] as actual_seq_qlen/actual_seq_kvlen.
|
||||
out = attention_forward_varlen(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
[0, *seq_q],
|
||||
[0, *seq_k],
|
||||
softmax_scale=self.softmax_scale,
|
||||
)
|
||||
return out.unsqueeze(0)
|
||||
|
||||
def _forward_prefix_kv_slice_npu(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
extra: dict,
|
||||
) -> torch.Tensor | None:
|
||||
"""Mask-free attention by slicing K/V to the valid prefix (zero-copy
|
||||
views), then mindiesd attention_forward without a mask.
|
||||
|
||||
Same numerical contract as the varlen path: the padding document is a
|
||||
strict suffix, so dropping it from K/V is identical to masking it out.
|
||||
Query keeps full length; outputs on padding rows are never consumed
|
||||
downstream. Returns None when the packed contract does not hold.
|
||||
|
||||
Used when MINDIE_SD_FA_TYPE=ascend_laser_attention: the TND varlen op
|
||||
cannot honor the env-selected op type, while attention_forward can.
|
||||
|
||||
The laser kernel stores unscaled S=QK^T in an fp16 GM workspace, so
|
||||
bf16 activations with large outliers overflow 65504 into ±inf and the
|
||||
affected rows turn NaN. Models may opt into exact power-of-two input
|
||||
pre-scaling via extra["laser_input_scale"] (see abstract.py): q/k/v
|
||||
are divided by the factor, the kernel scale_value is multiplied by its
|
||||
square, and the output is scaled back. Absent or invalid factor means
|
||||
no pre-scaling.
|
||||
"""
|
||||
resolved = self._resolve_packed_seq_npu(query, key, extra)
|
||||
if resolved is None:
|
||||
return None
|
||||
_, seq_k = resolved
|
||||
used_k = seq_k[0] # real document length (first cumulative end)
|
||||
|
||||
try:
|
||||
from mindiesd import attention_forward
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"FlashAttentionBackend NPU implementation requires MindIE-SD. "
|
||||
"Please install MindIE-SD to enable NPU attention support. "
|
||||
"For installation details, see https://gitcode.com/Ascend/MindIE-SD"
|
||||
"Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
|
||||
)
|
||||
# mindiesd always takes BSND input regardless of `layout`; the arg
|
||||
# selects the op-internal layout and laser only supports BNSD, so do
|
||||
# NOT forward the model's qkv_layout ("BSND" for MiniMax-H3) here.
|
||||
layout = "BNSD"
|
||||
key = key[:, :used_k]
|
||||
value = value[:, :used_k]
|
||||
|
||||
input_scale = extra.get("laser_input_scale")
|
||||
preserve_input_range = isinstance(input_scale, (int, float)) and input_scale > 1
|
||||
if preserve_input_range:
|
||||
query = query / input_scale
|
||||
key = key / input_scale
|
||||
value = value / input_scale
|
||||
scale = self.softmax_scale
|
||||
if preserve_input_range:
|
||||
scale *= input_scale**2
|
||||
|
||||
def _postprocess(out: torch.Tensor) -> torch.Tensor:
|
||||
return out * input_scale if preserve_input_range else out
|
||||
|
||||
return _postprocess(
|
||||
attention_forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_mask=None,
|
||||
opt_mode="manual",
|
||||
op_type="fused_attn_score", # MINDIE_SD_FA_TYPE env overrides this
|
||||
layout=layout,
|
||||
scale=scale,
|
||||
)
|
||||
)
|
||||
@@ -128,9 +128,9 @@ class RainFusionAttentionImpl(AttentionImpl):
|
||||
whose extent the model publishes as ``AttentionMetadata.video_layout``. Every
|
||||
other case — warmup denoise steps, exempt layers, a layer that does not declare
|
||||
``qkv_layout="BSND"``, sequences without a published video segment, video
|
||||
segments too short to pay for block selection or not aligned to the kernel's
|
||||
block size — delegates to FlashAttention, so a model can select this backend
|
||||
unconditionally.
|
||||
segments too short to pay for block selection — delegates to FlashAttention,
|
||||
so a model can select this backend unconditionally. MindIE-SD handles an
|
||||
irregular video tail internally, retaining it outside the sparse blocks.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -292,25 +292,6 @@ class RainFusionAttentionImpl(AttentionImpl):
|
||||
_MIN_VIDEO_BLOCKS,
|
||||
)
|
||||
return None
|
||||
if video_len % _BLOCK_SIZE != 0:
|
||||
# rf_v2 tiles the sequence into ceil((video + prefix) / 128) blocks but
|
||||
# builds its block mask from the two segments pooled apart. The counts
|
||||
# only agree when the video segment ends on a block boundary; otherwise
|
||||
# a prefix block is forced on out of range, and the block straddling the
|
||||
# seam holds real prefix rows that selection may drop. Padding the
|
||||
# sequence to realign it is not an option because rf_v2 ignores
|
||||
# attn_mask, so pad keys would take a share of every softmax denominator.
|
||||
logger.warning_once(
|
||||
"RAINFUSION_ATTN staying dense: latent grid %s gives %d video rows, which is not a "
|
||||
"multiple of the %d-row kernel block. Choose a geometry whose latent_t * h * w is a "
|
||||
"multiple of %d to enable sparsity.",
|
||||
tuple(latent_shape),
|
||||
video_len,
|
||||
_BLOCK_SIZE,
|
||||
_BLOCK_SIZE,
|
||||
)
|
||||
return None
|
||||
|
||||
logger.info_once(
|
||||
"RAINFUSION_ATTN active: sparsity=%.2f, start_step=%d, exempt_layers=%d, "
|
||||
"latent_grid=%s, prefix_rows=%d, video_rows=%d. Realized sparsity is lower than nominal "
|
||||
|
||||
@@ -1,11 +1,14 @@
|
||||
# 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 (
|
||||
@@ -19,6 +22,89 @@ 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)
|
||||
@@ -80,6 +166,16 @@ 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()
|
||||
|
||||
|
||||
@@ -57,6 +57,17 @@ _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:
|
||||
@@ -365,6 +376,8 @@ 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,6 +107,15 @@ MINIMAX_H3_FP32_BUFFER_NAMES = frozenset({"rope.inv_freq"})
|
||||
# lookup and masked out afterwards).
|
||||
MINIMAX_H3_ADALN_MODALITY_NUM = 3
|
||||
|
||||
# Opt-in fp16-range protection for the NPU ascend_laser_attention kernel
|
||||
# (consumed only via the "laser_input_scale" extra key; other backends and
|
||||
# platforms ignore it). The kernel stores unscaled QK^T in an fp16 GM
|
||||
# workspace, and H3's outlier activations (per-element amax in the hundreds)
|
||||
# push dot products past fp16 max 65504, turning whole 128-row blocks NaN.
|
||||
# 256 is a power of two, so pre-dividing q/k/v and the compensating
|
||||
# kernel-scale/output multiplies are exact in floating point.
|
||||
MINIMAX_H3_LASER_INPUT_SCALE = 256.0
|
||||
|
||||
|
||||
def _required_kwarg(kwargs: dict[str, Any], key: str) -> Any:
|
||||
if key not in kwargs or kwargs[key] is None:
|
||||
@@ -402,11 +411,16 @@ class MiniMaxH3Attention(nn.Module):
|
||||
used = min(max_seqlen, packed_total)
|
||||
attn_mask = None
|
||||
# Ring attention can dispatch to a different implementation from the
|
||||
# configured backend, so this no-mask fast path is local-only.
|
||||
prefix_slice = (
|
||||
not getattr(self.attention, "use_ring", False) and self.attention.attn_backend.supports_prefix_kv_slicing
|
||||
# configured backend, so the no-mask fast paths are local-only.
|
||||
# supports_prefix_kv_slicing: backend slices K/V itself (cuDNN).
|
||||
# supports_packed_mask_free: backend consumes the packed metadata
|
||||
# without ever reading attn_mask (CUDA packed varlen, NPU
|
||||
# npu_attn_varlen opt-in with its own fallback rebuild).
|
||||
no_mask = not getattr(self.attention, "use_ring", False) and (
|
||||
self.attention.attn_backend.supports_prefix_kv_slicing
|
||||
or self.attention.attn_backend.supports_packed_mask_free()
|
||||
)
|
||||
if used < packed_total and not prefix_slice:
|
||||
if used < packed_total and not no_mask:
|
||||
attn_mask = torch.arange(packed_total, device=q.device)[None] < used
|
||||
metadata = AttentionMetadata(
|
||||
attn_mask=attn_mask,
|
||||
@@ -416,6 +430,15 @@ class MiniMaxH3Attention(nn.Module):
|
||||
"max_seqlen_q": max_seqlen,
|
||||
"max_seqlen_k": max_seqlen,
|
||||
"valid_kv_length": used,
|
||||
# Opt the NPU flash backend into the packed varlen path so the
|
||||
# quadratic full_qk mask is never materialized. Ring attention
|
||||
# is excluded: it keeps the aligned padding rows for its
|
||||
# fixed-size P2P buffers and still needs the mask.
|
||||
"npu_attn_varlen": not getattr(self.attention, "use_ring", False),
|
||||
# fp16-range protection for the ascend_laser_attention kernel
|
||||
# (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other
|
||||
# backend/path.
|
||||
"laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE,
|
||||
},
|
||||
video_layout=video_layout,
|
||||
)
|
||||
|
||||
@@ -1069,6 +1069,7 @@ 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,7 +629,6 @@ 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,6 +139,18 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
|
||||
)
|
||||
backend_upper = "FLASH_ATTN"
|
||||
|
||||
if backend_upper == "FLASH_ATTN" and find_spec("mindiesd"):
|
||||
# The NPU FLASH_ATTN backend imports mindiesd lazily at first
|
||||
# forward, but CANN snapshots the custom-op registry at the
|
||||
# first custom-op regInfo lookup in the process (e.g. a
|
||||
# vllm-ascend custom op during model load/warmup). Import
|
||||
# mindiesd here so its env.py prepends the mindiesd vendor
|
||||
# dirs (aie_ascendc etc.) to ASCEND_CUSTOM_OPP_PATH before
|
||||
# that snapshot; otherwise aclnnLaserAttention /
|
||||
# FusedAttentionScore fail with EZ1001 "does not support
|
||||
# opType" for the rest of the process.
|
||||
import mindiesd # noqa: F401
|
||||
|
||||
backend = DiffusionAttentionBackendEnum[backend_upper]
|
||||
logger.debug("Using diffusion attention backend '%s'", backend_upper)
|
||||
return backend.get_path()
|
||||
|
||||
Reference in new issue
Block a user