From 7872663bf044cdcbf74de72a386900efd0392b90 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Wed, 12 Aug 2026 08:55:41 +0800 Subject: [PATCH] docs: split diffusion feature guides Signed-off-by: Hongsheng Liu --- docs/.nav.yml | 21 +- .../feature/attention_backend_selection.md | 89 +++ docs/design/feature/cpu_offloading.md | 80 +++ .../feature/distributed_layerwise_offload.md | 2 +- docs/design/feature/layerwise_offload.md | 57 ++ docs/design/feature/model_level_offload.md | 56 ++ docs/design/feature/skip_softmax.md | 6 +- docs/design/index.md | 13 +- docs/features/README.md | 4 +- .../diffusion/attention_backends.md | 524 ++++-------------- .../attention_backends/dense_backends.md | 110 ++++ .../attention_backends/huggingface_hub.md | 32 ++ .../attention_backends/rainfusion.md | 72 +++ .../diffusion/attention_backends/sage.md | 63 +++ .../diffusion/attention_backends/trtllm.md | 103 ++++ docs/user_guide/diffusion/cpu_offload.md | 434 ++------------- .../cpu_offload/distributed_layerwise.md | 118 ++++ .../diffusion/cpu_offload/layerwise.md | 65 +++ .../diffusion/cpu_offload/model_level.md | 81 +++ docs/user_guide/diffusion_features.md | 7 +- docs/user_guide/feature_compatibility.md | 2 +- 21 files changed, 1138 insertions(+), 801 deletions(-) create mode 100644 docs/design/feature/attention_backend_selection.md create mode 100644 docs/design/feature/cpu_offloading.md create mode 100644 docs/design/feature/layerwise_offload.md create mode 100644 docs/design/feature/model_level_offload.md create mode 100644 docs/user_guide/diffusion/attention_backends/dense_backends.md create mode 100644 docs/user_guide/diffusion/attention_backends/huggingface_hub.md create mode 100644 docs/user_guide/diffusion/attention_backends/rainfusion.md create mode 100644 docs/user_guide/diffusion/attention_backends/sage.md create mode 100644 docs/user_guide/diffusion/attention_backends/trtllm.md create mode 100644 docs/user_guide/diffusion/cpu_offload/distributed_layerwise.md create mode 100644 docs/user_guide/diffusion/cpu_offload/layerwise.md create mode 100644 docs/user_guide/diffusion/cpu_offload/model_level.md diff --git a/docs/.nav.yml b/docs/.nav.yml index bdc11ec7c..68aa944eb 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -60,7 +60,11 @@ nav: - Overview and Compatibility: - Overview: user_guide/diffusion_features.md - Feature Compatibility: user_guide/feature_compatibility.md - - CPU Offloading: user_guide/diffusion/cpu_offload.md + - CPU Offloading: + - Overview: user_guide/diffusion/cpu_offload.md + - Model-Level: user_guide/diffusion/cpu_offload/model_level.md + - Layerwise: user_guide/diffusion/cpu_offload/layerwise.md + - Distributed Layerwise: user_guide/diffusion/cpu_offload/distributed_layerwise.md - Cache Acceleration: - TeaCache: user_guide/diffusion/cache_acceleration/teacache.md - Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md @@ -72,7 +76,13 @@ nav: - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md - - Attention Backends: user_guide/diffusion/attention_backends.md + - Attention Backends: + - Overview: user_guide/diffusion/attention_backends.md + - Dense Backends: user_guide/diffusion/attention_backends/dense_backends.md + - TRTLLM: user_guide/diffusion/attention_backends/trtllm.md + - SageAttention: user_guide/diffusion/attention_backends/sage.md + - Hugging Face Hub: user_guide/diffusion/attention_backends/huggingface_hub.md + - RainFusion: user_guide/diffusion/attention_backends/rainfusion.md - Regional Compilation: user_guide/diffusion/regional_compilation.md - Frame Interpolation: user_guide/diffusion/frame_interpolation.md - Startup and Loading: user_guide/diffusion/startup_and_loading.md @@ -125,11 +135,16 @@ nav: - Tensor Parallel: design/feature/tensor_parallel.md - VAE Patch Parallelism: design/feature/vae_parallel.md - Attention Optimization: + - Backend Selection: design/feature/attention_backend_selection.md - Skip-Softmax: design/feature/skip_softmax.md + - CPU Offloading: + - Overview: design/feature/cpu_offloading.md + - Model-Level: design/feature/model_level_offload.md + - Layerwise: design/feature/layerwise_offload.md + - Distributed Layerwise: design/feature/distributed_layerwise_offload.md - Cache-DiT: design/feature/cache_dit.md - TeaCache: design/feature/teacache.md - Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md - - Distributed Layerwise Offload: design/feature/distributed_layerwise_offload.md - Infrastructure and Performance: - Prometheus Metrics: design/metrics.md - Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md diff --git a/docs/design/feature/attention_backend_selection.md b/docs/design/feature/attention_backend_selection.md new file mode 100644 index 000000000..18e05517c --- /dev/null +++ b/docs/design/feature/attention_backend_selection.md @@ -0,0 +1,89 @@ +# Diffusion Attention Backend Selection + +This document defines the selection and extension contract for diffusion +attention backends. User-facing backend choices, installation, and tuning are +in the [attention backend guides](../../user_guide/diffusion/attention_backends.md). + +## Scope + +The contract applies to DiT and other diffusion attention layers. It is +separate from vLLM's autoregressive attention selector. + +The selection path has four responsibilities: + +1. normalize user configuration into `AttentionConfig` and `AttentionSpec`; +2. resolve a spec for an attention role; +3. ask the active platform to validate an explicit backend or choose a + hardware default; and +4. load the selected `AttentionBackend` class from the registry path. + +## Resolution contract + +`get_attn_backend_for_role()` resolves configuration in this order: + +1. exact `per_role[role]`; +2. category `per_role[role_category]`; +3. `default`; +4. platform default. + +An explicit resolution returns both the backend class and its +`AttentionSpec`. A platform-default resolution returns the class and `None`. +Layers must therefore treat the spec as optional and must not infer that a +missing spec means SDPA. + +Class resolution is cached by backend name, head size, and whether TRTLLM may +be selected automatically. Logging is cached separately by role and source so +the startup record identifies why each role received its backend. + +## Registry and platform boundary + +`DiffusionAttentionBackendEnum` maps stable configuration names to default +qualified class paths. `register_diffusion_backend()` may replace a path at +runtime without changing the public enum value. + +The active `OmniPlatform` owns compatibility policy through +`get_diffusion_attn_backend_cls()`. It must: + +- validate explicit selections and fail with an actionable error when the + requested kernel cannot run; +- choose only an available, compatible default when no backend is explicit; +- account for head size and other platform-visible constraints; and +- return a qualified class path implementing `AttentionBackend`. + +The selector must not duplicate device capability or package-availability +policy that belongs to the platform. + +## Typed backend options + +Backend-specific settings remain typed fields on `AttentionSpec` rather than +unstructured keyword dictionaries: + +- `quant` is consumed by FlashInfer and TRTLLM with backend-specific value + validation; +- `skip_softmax` is consumed by TRTLLM; and +- `block_sparse` is consumed by block-sparse backends such as RainFusion. + +A backend reads only the fields it owns and rejects incompatible values. New +options should be added to a shared typed spec only when more than one backend +shares their semantics; otherwise add a dedicated typed block. + +## Model integration contract + +Each diffusion attention call declares a stable role and, when appropriate, a +broader category. Model-specific roles permit precise overrides; categories +preserve common `self` or `cross` policies. A model also decides whether its +path is compatible with automatic TRTLLM selection, because only the model +knows whether masking and packing satisfy the kernel contract. + +## Adding or changing a backend + +An implementation change is complete only when it: + +1. implements the `AttentionBackend` interface and exposes a stable enum name; +2. adds platform validation and default routing where appropriate; +3. defines and validates any typed configuration; +4. tests explicit selection, default selection, and incompatible paths; +5. documents installation, fallback behavior, and quality implications in the + corresponding user guide; and +6. updates the overview matrix without moving algorithm details back into the + landing page. diff --git a/docs/design/feature/cpu_offloading.md b/docs/design/feature/cpu_offloading.md new file mode 100644 index 000000000..a1a7124b2 --- /dev/null +++ b/docs/design/feature/cpu_offloading.md @@ -0,0 +1,80 @@ +# CPU Offloading + +This document defines the shared architecture for diffusion CPU offloading. +The three strategies have separate user and design pages: + +| Strategy | User guide | Design contract | +| --- | --- | --- | +| Model-level | [Guide](../../user_guide/diffusion/cpu_offload/model_level.md) | [Design](model_level_offload.md) | +| Layerwise | [Guide](../../user_guide/diffusion/cpu_offload/layerwise.md) | [Design](layerwise_offload.md) | +| Distributed layerwise | [Guide](../../user_guide/diffusion/cpu_offload/distributed_layerwise.md) | [Design](distributed_layerwise_offload.md) | + +## Strategy selection + +`OffloadConfig.from_od_config()` converts diffusion configuration into one +`OffloadStrategy`. If callers set multiple flags, selection is deterministic: + +1. `DISTRIBUTED_LAYER_WISE`; +2. `LAYER_WISE`; +3. `MODEL_LEVEL`; and +4. `NONE`. + +The factory logs when a higher-priority strategy disables another. It derives +parallel and HSDP state from `DiffusionParallelConfig`; callers do not provide +a separate offload group size. + +`get_offload_backend()` then validates platform offload support, resolves the +device, and creates exactly one backend. Returning `None` means offloading is +disabled or unsupported and must not leave partially installed hooks. + +## Shared lifecycle + +Every backend implements `OffloadBackend`: + +- `enable(pipeline)` discovers modules, establishes initial residency, and + installs hooks; +- `disable()` removes owned hooks and resources; and +- `is_enabled()` reports lifecycle state. + +`disable()` does not promise to restore the pipeline's original device +placement. The caller owns any subsequent rematerialization. + +Hooks are registered through `HookRegistry` and `ModelHook`; offload backends +must use distinct hook names and remove only hooks they own. + +## Discovery and topology + +Pipeline component discovery prefers `SupportsComponentDiscovery` declarations: + +- `_dit_modules`; +- `_encoder_modules`; +- `_vae_modules`; and +- `_resident_modules`. + +Dotted paths are supported. Legacy pipelines may use the fallback scan of +well-known attribute names, but new integrations should declare components +explicitly. + +Ordinary layerwise offload obtains topology from declared +`_layerwise_offload_blocks_attrs` (with a deprecated singular-name +compatibility path). Distributed layerwise offload may instead use a pipeline +`OffloadPlan`, then falls back to declared block attributes and its documented +compatibility discovery. Discovery metadata describes structure only; the +backend remains responsible for transfer, synchronization, and storage +ownership. + +## Cross-strategy invariants + +- At most one strategy owns offload hooks for a pipeline. +- A parameter has one authoritative host representation while offloaded. +- Device storage is not freed until dependent compute or transfer work is + complete. +- Non-persistent and model-specific buffers remain correct after movement or + checkpoint rematerialization. +- Unsupported parallel or loading combinations fail before hooks mutate the + model. +- Platform streams, events, synchronization, and cache management go through + the vLLM-Omni platform abstraction. + +The diffusion [Offloader module design](../module/diffusion/offloader.md) +describes how these feature contracts fit into the larger diffusion runtime. diff --git a/docs/design/feature/distributed_layerwise_offload.md b/docs/design/feature/distributed_layerwise_offload.md index da735d60b..672b3c2e6 100644 --- a/docs/design/feature/distributed_layerwise_offload.md +++ b/docs/design/feature/distributed_layerwise_offload.md @@ -7,7 +7,7 @@ either shard those host-side weights across an existing parallel group or keep the standard loader's rank-local weights and avoid an additional collective. For user-facing commands, see the -[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md) +[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload/distributed_layerwise.md) and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md). ## Status diff --git a/docs/design/feature/layerwise_offload.md b/docs/design/feature/layerwise_offload.md new file mode 100644 index 000000000..30b52ee62 --- /dev/null +++ b/docs/design/feature/layerwise_offload.md @@ -0,0 +1,57 @@ +# Layerwise Offload + +Layerwise offload streams transformer blocks through one accelerator. For +commands and model support, see the +[layerwise user guide](../../user_guide/diffusion/cpu_offload/layerwise.md). + +## Hook ring + +`LayerWiseOffloadBackend` discovers every streamable block and registers one +`LayerwiseOffloadHook` per block. Hooks form a ring: each hook owns the pinned +host representation of its registered block and a reference to the next block +it must prefetch; the final hook prefetches the first block for the next +denoising iteration. + +During initialization, parameters and buffers are grouped by dtype and copied +into contiguous pinned CPU tensors. Original tensor storage is replaced with +an empty or meta placeholder while identity and module structure are retained. +Metadata records names, offsets, shapes, and dtypes for rematerialization. + +## Forward lifecycle + +For each block: + +1. the previous hook's completion event ensures this block is materialized; +2. the copy stream asynchronously allocates and fills the next block's device + tensors; +3. parameter and buffer storage is rebound to views of those tensors; +4. the current block computes on the normal stream; and +5. its post-forward hook replaces device storage with placeholders after all + dependent work is safe. + +The copy stream waits on the compute stream before reusing storage, and the +compute stream waits on the per-block prefetch event before forward. These +dependencies, not incidental global synchronization, define correctness. + +## Topology contract + +A DiT declares ordered block-container attributes using +`_layerwise_offload_blocks_attrs` or an `OffloadPlan`. Each resolved item must +be an executable `nn.Module`; ordering must match forward execution. Multiple +containers are concatenated in execution order. + +Non-block DiT modules, encoders, VAEs, and declared resident modules are not +streamed by this backend. + +## Invariants and limitations + +- Each offloaded tensor has one pinned host backing store. +- Parameter identity is preserved when storage is rebound. +- Blocks with heterogeneous shapes and dtypes use their recorded metadata; + implementations must not assume identical blocks. +- The ring is initialized before the first forward so block zero is available. +- This backend is single-device; distributed weight reconstruction belongs to + [Distributed Layerwise Offload](distributed_layerwise_offload.md). + +Shared selection and lifecycle behavior is defined in the +[CPU Offloading design](cpu_offloading.md). diff --git a/docs/design/feature/model_level_offload.md b/docs/design/feature/model_level_offload.md new file mode 100644 index 000000000..fd213b3aa --- /dev/null +++ b/docs/design/feature/model_level_offload.md @@ -0,0 +1,56 @@ +# Model-Level Offload + +Model-level offload enforces mutual exclusion between pipeline component +groups. For commands and model support, see the +[model-level user guide](../../user_guide/diffusion/cpu_offload/model_level.md). + +## Generic hook path + +`ModelLevelOffloadBackend.enable()` discovers DiTs, encoders, and VAEs. It +makes encoders and VAEs device resident for initialization, then installs +`SequentialOffloadHook` instances: + +- each DiT hook offloads all encoders and other DiTs before loading itself; +- each encoder hook offloads every DiT before loading itself; and +- VAEs are not part of the mutual-exclusion hook set. + +The hook moves parameters and buffers without recursively calling +`module.to()`. This avoids recursion through modules that retain references to +other transformer blocks. Host tensors are pinned when configured; DTensors +are not pinned through the ordinary tensor path. + +Before forward returns control to the model, the hook synchronizes the active +platform so compute cannot race the transfer. XPU uses blocking device-to-host +copies because its allocator cannot safely evict storage while an asynchronous +copy is pending. + +## Split-model delegation + +If a pipeline provides callable `enable_omni_model_cpu_offload`, the backend +delegates instead of installing generic encoder/DiT hooks. This supports models +whose mutually exclusive components are nested inside one transformer, such +as a reasoner and generator. + +The delegated pipeline owns its internal contexts, but must preserve the same +contract: + +- only the executing component is device resident; +- enable and disable are idempotent; +- transfers use the supplied device, pinning, and HSDP settings; and +- disabling removes all model-local offload hooks or contexts. + +The delegation is currently duck typed. If another independent implementation +is added, formalize the enable/disable pair as a protocol rather than expanding +attribute-name heuristics. + +## Extension invariants + +- New ordinary pipelines declare both DiT and encoder groups through + `SupportsComponentDiscovery`. +- Models with multiple DiTs offload the inactive DiTs as well as encoders. +- Empty or parameterless modules do not trigger transfers. +- Hook removal does not move modules; callers must not assume original + residency is restored. + +Shared selection and lifecycle behavior is defined in the +[CPU Offloading design](cpu_offloading.md). diff --git a/docs/design/feature/skip_softmax.md b/docs/design/feature/skip_softmax.md index 8e6375ee8..0527c7eeb 100644 --- a/docs/design/feature/skip_softmax.md +++ b/docs/design/feature/skip_softmax.md @@ -2,8 +2,10 @@ Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. Usage — the config keys and how to pick an operating point — is in -[Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). This page explains -the algorithm. +[TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax). +The shared selector contract is documented in +[Diffusion Attention Backend Selection](attention_backend_selection.md). This +page explains the algorithm. ## Motivation diff --git a/docs/design/index.md b/docs/design/index.md index d81ccb806..798201214 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -53,15 +53,22 @@ implementation contract; it is not, by itself, a general support claim. #### Attention optimization The [Diffusion Attention Backends](../user_guide/diffusion/attention_backends.md) -guide lists all selectable backends and their platform defaults. Skip-Softmax -is the backend-related optimization with a standalone feature-design contract: +guides list selectable backends, platform defaults, installation, and tuning. +The design contracts separate selection mechanics from backend algorithms: +- [Attention Backend Selection](feature/attention_backend_selection.md) - [Skip-Softmax](feature/skip_softmax.md) +#### CPU offloading + +- [Overview and Shared Contracts](feature/cpu_offloading.md) +- [Model-Level Offload](feature/model_level_offload.md) +- [Layerwise Offload](feature/layerwise_offload.md) +- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md) + - [Cache-DiT](feature/cache_dit.md) - [TeaCache](feature/teacache.md) - [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md) -- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md) ## Infrastructure and Performance diff --git a/docs/features/README.md b/docs/features/README.md index 0dee558a4..411a31164 100644 --- a/docs/features/README.md +++ b/docs/features/README.md @@ -48,10 +48,10 @@ defines the shared configuration and backend extension points. | Goal | User guide | Related design contract | | --- | --- | --- | | Compare acceleration methods and supported combinations | [Overview](../user_guide/diffusion_features.md), [Feature Compatibility](../user_guide/feature_compatibility.md) | [Diffusion acceleration designs](../design/index.md#diffusion-acceleration) | -| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [Distributed Layerwise Offload](../design/feature/distributed_layerwise_offload.md) | +| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [CPU Offloading](../design/feature/cpu_offloading.md) | | Reuse denoising computation | [Cache Acceleration](../user_guide/diffusion/cache_acceleration/cache_dit.md) | [Cache-DiT](../design/feature/cache_dit.md), [TeaCache](../design/feature/teacache.md) | | Distribute diffusion work across devices | [Parallelism](../user_guide/diffusion/parallelism/overview.md) | [Parallelism designs](../design/index.md#parallelism) | -| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Skip-Softmax](../design/feature/skip_softmax.md) | +| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Attention Backend Selection](../design/feature/attention_backend_selection.md) | | Compile repeated diffusion regions | [Regional Compilation](../user_guide/diffusion/regional_compilation.md) | User-facing optimization guide | | Add generated video frames | [Frame Interpolation](../user_guide/diffusion/frame_interpolation.md) | User-facing extension guide | | Reduce diffusion model startup time | [Startup and Loading](../user_guide/diffusion/startup_and_loading.md) | User-facing loading guide | diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index b6fd13b12..6897d633e 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/docs/user_guide/diffusion/attention_backends.md @@ -1,494 +1,192 @@ # Diffusion Attention Backends -This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention. +Use this page to select and configure a diffusion attention backend. Backend- +specific installation, tuning, and compatibility details live in separate +guides so that the selection contract stays easy to scan. -## Overview +Diffusion attention backend selection applies to DiT attention in image and +video generation models. It does **not** change autoregressive LLM attention, +which uses vLLM's attention selector. -Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured. +For the internal selector, registry, and platform contract, see +[Attention Backend Selection](../../design/feature/attention_backend_selection.md). -This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector. +## Choose a guide -The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default. +| Need | Guide | +| --- | --- | +| Select a conservative or platform-native dense kernel | [Dense Backends](attention_backends/dense_backends.md) | +| Use FlashInfer trtllm-gen FMHA, Skip-Softmax, or TRTLLM SAGE quantization | [TRTLLM Attention](attention_backends/trtllm.md) | +| Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) | +| Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) | +| Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) | -## Backend Options - -| Value | Notes | -|---|---| -| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. | -| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. | -| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). | -| `FLASHINFER_ATTN` | Uses FlashInfer's batch-prefill wrapper and supports mixed Q/K and V input dtypes through backend-specific configuration. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer` >= 0.6.16rc1 for mixed-dtype configurations. | -| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. | -| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. | -| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. | -| `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. | -| `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. | -| `RAINFUSION_ATTN` | MindIE-SD **RainFusion** block-sparse video attention — see [below](#rainfusion_attn-backend-and-block-sparse-video-attention). Ascend NPU only; requires `mindiesd`. Delegates to `FLASH_ATTN` for anything that is not a packed video sequence. | +## Backend options +| Value | Family | Primary use | Detail | +| --- | --- | --- | --- | +| `TORCH_SDPA` | Dense | Conservative reference; always available | [Dense Backends](attention_backends/dense_backends.md#torch_sdpa) | +| `FLASH_ATTN` | Dense | FlashAttention 4/3/2 depending on the installed package and GPU | [Dense Backends](attention_backends/dense_backends.md#flash_attn) | +| `CUDNN_ATTN` | Dense | Mask-heavy DiTs on Blackwell with cuDNN 9.5 or newer | [Dense Backends](attention_backends/dense_backends.md#cudnn_attn) | +| `FLASHINFER_ATTN` | Dense or quantized | FlashInfer batch prefill; optional mixed Q/K and V dtypes | [Dense Backends](attention_backends/dense_backends.md#flashinfer_attn) | +| `TRTLLM_ATTN` | Dense, sparse, or quantized | Datacenter Blackwell with `head_dim=128` and compatible packed paths | [TRTLLM Attention](attention_backends/trtllm.md) | +| `SAGE_ATTN` | Quantized | SageAttention 2.2 INT8 attention | [SageAttention](attention_backends/sage.md#sage_attn) | +| `SAGE_ATTN_3` | Quantized | SageAttention3 on Blackwell | [SageAttention](attention_backends/sage.md#sage_attn_3) | +| `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) | +| `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) | +| `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) | ## Configuration -Diffusion attention backends can be configured three ways, in priority order: +Backend selection is resolved in this order: -1. **`--diffusion-attention-config`** — structured per-role config (highest priority). -2. **`--diffusion-attention-backend` / `DIFFUSION_ATTENTION_BACKEND` env var** — global shorthand that sets the default backend. -3. **Platform default** — used when nothing is configured. +1. `--diffusion-attention-config` per-role configuration. +2. `--diffusion-attention-backend` or `DIFFUSION_ATTENTION_BACKEND` as a + global default. +3. The current platform's default. -`--diffusion-attention-backend` is shorthand for `--diffusion-attention-config.default.backend`. It may be combined with `--diffusion-attention-config.per_role.*` overrides, but is mutually exclusive with `--diffusion-attention-config.default.backend`. +`--diffusion-attention-backend` is shorthand for +`--diffusion-attention-config.default.backend`. Do not pass it together with an +explicit `default.backend` in the structured configuration. ### Global default -Set the default backend for every diffusion attention layer: - ```bash -# CLI flag -vllm-omni serve --diffusion-attention-backend SAGE_ATTN +vllm-omni serve --diffusion-attention-backend FLASH_ATTN -# Environment variable (also recognized for backwards compatibility) -export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN +# Backwards-compatible environment variable +export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN ``` ### Per-role configuration -Roles are free-form strings declared by each diffusion model. The two common categories are `"self"` and `"cross"`; model-specific roles (e.g. `"ltx2.audio_to_video"`) may also be declared. A role string is matched in this order: +Roles are declared by each diffusion model. Common categories are `self` and +`cross`; a model may also use a more specific role such as +`ltx2.audio_to_video`. Resolution order is: -1. Exact `per_role[role]` match -2. `per_role[role_category]` fallback (e.g. `"ltx2.audio_to_video"` → `"cross"`) -3. `default` -4. Platform default - -Use vLLM-style dotted flags or one JSON blob: +1. Exact `per_role[role]` match. +2. Category `per_role[role_category]` match. +3. `default`. +4. Platform default. ```bash # Dotted flags vllm-omni serve \ - --diffusion-attention-config.default.backend FLASH_ATTN \ - --diffusion-attention-config.per_role.cross.backend TORCH_SDPA + --diffusion-attention-config.default.backend FLASH_ATTN \ + --diffusion-attention-config.per_role.cross.backend TORCH_SDPA -# JSON +# Equivalent JSON vllm-omni serve \ - --diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}' + --diffusion-attention-config \ + '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}' ``` -A backend that needs configuration exposes it as a typed field on the spec: -`TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)) and -`RAINFUSION_ATTN`'s RainFusion block (see [below](#rainfusion_attn-backend-and-block-sparse-video-attention)): - -```bash ---diffusion-attention-config.default.backend TRTLLM_ATTN \ ---diffusion-attention-config.default.skip_softmax.target_sparsity 0.5 -``` - -### Programmatic API - -When constructing `OmniDiffusionConfig` directly: +### Python API ```python -from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec, OmniDiffusionConfig +from vllm_omni.diffusion.data import ( + AttentionConfig, + AttentionSpec, + OmniDiffusionConfig, +) config = OmniDiffusionConfig( diffusion_attention_config=AttentionConfig( default=AttentionSpec(backend="FLASH_ATTN"), - per_role={ - "cross": AttentionSpec(backend="TORCH_SDPA"), - }, + per_role={"cross": AttentionSpec(backend="TORCH_SDPA")}, ), ..., ) ``` -A plain dict is also accepted and normalized to `AttentionConfig`. +Backend-specific typed blocks are documented with their consumers: -#### FlashInfer QK16/V8 quantized attention +- `quant`: [FlashInfer](attention_backends/dense_backends.md#flashinfer-quantized-attention) + and [TRTLLM SAGE](attention_backends/trtllm.md#sage-quantization). +- `skip_softmax`: [TRTLLM Skip-Softmax](attention_backends/trtllm.md#skip-softmax). +- `block_sparse`: [RainFusion](attention_backends/rainfusion.md#configuration). -`FLASHINFER_ATTN` accepts quantization options through `AttentionSpec.quant`. For QK16/V8: - -```python -config = OmniDiffusionConfig( - diffusion_attention_config=AttentionConfig( - default=AttentionSpec( - backend="FLASHINFER_ATTN", - quant=AttnQuantSpec( - dtype_qk="bfloat16", - dtype_vo="fp8_e4m3", - ), - ), - ), - ..., -) -``` - -`dtype_qk` configures the Q and K input dtypes. `dtype_vo` configures the V input dtype. - -## Platform Defaults +## Platform defaults ### Blackwell (sm_100 / sm_103 / sm_120 / sm_121) -Auto-route preference, in order: +The CUDA auto-route preference is: -1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model declares a compatible packed/mask-free path** -2. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels) -3. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5 -4. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel -5. `TORCH_SDPA` — last resort +1. `TRTLLM_ATTN` on datacenter Blackwell (sm_100/sm_103) when FlashInfer is + available, `head_dim=128`, and the model declares a compatible packed or + mask-free path. +2. `CUDNN_ATTN` when cuDNN 9.5 or newer is available. +3. `FLASHINFER_ATTN` when FlashInfer is available but cuDNN is too old. +4. `FLASH_ATTN` when a compatible package is installed. +5. `TORCH_SDPA`. -`TRTLLM_ATTN` is skipped on workstation Blackwell (sm_120 / sm_121) and for any `head_dim != 128`, so those GPUs keep the `CUDNN_ATTN` route described below. +`TRTLLM_ATTN` is not auto-selected on workstation Blackwell +(sm_120/sm_121), for other head dimensions, or for paths that require an +unsupported mask. -`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell for compatible packed/mask-free pipelines. Workstation Blackwell (sm_120 / sm_121) and pipelines that require attention masks retain their normal fallback. +### Hopper, Ada, and Ampere -The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route. +The CUDA auto-route uses `FLASH_ATTN` when available and otherwise falls back +to `TORCH_SDPA`. `CUDNN_ATTN` and `FLASHINFER_ATTN` remain explicit options. -**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models. +Other platforms validate an explicit backend and choose their own default +through the platform implementation. Check the startup log to confirm the +resolved backend. -### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86) +## Choosing a backend manually -Auto-route preference: +Override the platform default when you need: -1. `FLASH_ATTN` — when `flash-attn` is installed -2. `TORCH_SDPA` — fallback +- a correctness reference (`TORCH_SDPA`); +- a backend-specific workaround; +- training/rollout kernel alignment (Hub backends); or +- an explicitly validated sparse or quantized speedup. -`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware. +The startup log prints the resolved backend and whether it came from explicit +configuration or platform defaulting. If no resolution message appears, check +earlier logs for diffusion-stage initialization failures. -## TRTLLM_ATTN Backend and Skip-Softmax +## Reference benchmark -`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA. Dense it is a BF16 backend like the others; on top of -that it can enable **Skip-Softmax**, a sparse-attention mode that trades a little fidelity for speed -(algorithm: [design doc](../../design/feature/skip_softmax.md)). +The following BF16 results were measured on an sm_120 RTX Pro 6000 Blackwell +with the same prompt and seed across runs. Treat them as reference results, not +portable guarantees. -Enable it through the typed `skip_softmax` block on the attention spec: - -| Key | Valid values | Meaning | -|---|---|---| -| `target_sparsity` | finite, `[0, 1]` | Operating point on a calibrated curve. Needs a calibration for the model. | -| `threshold` | finite, `≥ 0` | Direct threshold, no calibration needed. Mutually exclusive with `target_sparsity`. | -| `disabled_until_timestep` | finite, `[0, 1]` | Holds early, high-noise steps dense; skip turns on once the normalized timestep `t ≤ D`. | - -```bash -vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN", - "skip_softmax": {"target_sparsity": 0.65, "disabled_until_timestep": 0.86}}}' -``` - -Programmatically the same block is a typed `SkipSoftmaxSpec` (values validated at construction): - -```python -from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec - -AttentionConfig( - default=AttentionSpec( - backend="TRTLLM_ATTN", - skip_softmax=SkipSoftmaxSpec(target_sparsity=0.65, disabled_until_timestep=0.86), - ), -) -``` - -**Start at `target_sparsity=0.65, disabled_until_timestep=0.86`** and raise `target_sparsity` only -as far as your quality bar allows — the output diverges from dense as it climbs. The gain grows with -sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `TRTLLM_ATTN` with no -`skip_softmax` block. - -Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting -`TRTLLM_ATTN` raises rather than silently degrading. - -## TRTLLM_ATTN SAGE Quantization - -On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen -FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run -in low precision. Enable it through the typed `quant` block on the attention spec: - -| Key | Valid values | Meaning | -|---|---|---| -| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). | -| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). | -| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). | - -V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so -neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels. - -The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the -fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`), -`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`), -`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises. - -```bash -vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN", - "quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}' -``` - -```python -from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec - -AttentionConfig( - default=AttentionSpec( - backend="TRTLLM_ATTN", - quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16), - ), -) -``` - -**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel -availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103** -(B300); `int8` QK kernels are compiled for **SM100 only**. - -## RAINFUSION_ATTN Backend and Block-Sparse Video Attention - -`RAINFUSION_ATTN` runs MindIE-SD's RainFusion (`rf_v2`) kernel on Ascend NPU. It pools each -128-token block of keys, ranks them per query block, and computes attention against only the -top-scoring ones. Video tokens are rearranged into `(t, h, w)` order first, so the kept blocks are -spatiotemporal neighbours rather than arbitrary rows of the packed sequence. - -Only the video segment is sparsified. The prefix rows (text, visual conditions, audio) and the -first-frame blocks are always kept dense, which is why the realized sparsity is lower than the -nominal `sparsity` you configure. Anything the kernel cannot handle — a warmup denoise step, an -exempt layer, a sequence with no published video geometry, a video segment under 32 blocks — -delegates to `FLASH_ATTN`, so a model can select this backend unconditionally. - -Enable it through the typed `block_sparse` block on the attention spec, shared by every -block-sparse backend: - -| Key | Valid values | Meaning | -|---|---|---| -| `sparsity` | finite, `[0, 1]` | Nominal fraction of key blocks dropped per query block. `0` disables sparsity. Defaults to `0.8`. | -| `start_step` | `≥ 0` | Keep the first N denoise steps dense. Layout is decided early, so these steps dominate structural fidelity. | -| `skip_layers` | index selector, e.g. `"0-3,38"` | DiT blocks that always stay dense. | - -```bash -vllm-omni serve MiniMaxAI/MiniMax-H3 \ - --diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN", - "block_sparse": {"sparsity": 0.8, "start_step": 0}}}' -``` - -Programmatically the same block is a typed `BlockSparseSpec` (values validated at construction): - -```python -from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec - -AttentionConfig( - default=AttentionSpec( - backend="RAINFUSION_ATTN", - block_sparse=BlockSparseSpec(sparsity=0.8, start_step=0, skip_layers="0-1"), - ), -) -``` - -**Tune in the order `start_step` → `sparsity` → `skip_layers`.** Raise `start_step` first: it is the -cheapest way to recover structure, because the early high-noise steps decide global layout while the -later steps only refine texture. Then walk `sparsity` up to your quality bar. Reach for -`skip_layers` last, once an A/B against dense at the same seed points at specific blocks. - -Requires Ascend NPU with `mindiesd`; selecting it on any other platform raises. It is also -incompatible with ring sequence parallelism, since `rf_v2` needs the whole key sequence to rank -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. - -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. - -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) - -Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`. - -| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN | -|---|---|---|---|---| +| Model | Shape | `TORCH_SDPA` | `CUDNN_ATTN` | `FLASHINFER_ATTN` | +| --- | --- | ---: | ---: | ---: | | HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s | | Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** | | Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s | | FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s | -Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise. +Mask-heavy DiTs favored `CUDNN_ATTN`; lighter-mask or TP-saturated workloads +were close enough that users should benchmark their exact model and shape. -## Known Limitations +## Compatibility anchors -### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile +The following headings preserve links to sections that moved into dedicated +guides. -LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121). +## TRTLLM_ATTN Backend and Skip-Softmax -**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0: +See [TRTLLM Attention](attention_backends/trtllm.md#skip-softmax). -```bash -DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \ - --model Lightricks/LTX-2 ... -``` +## TRTLLM_ATTN SAGE Quantization -### FlashAttention-4 on Blackwell +See [TRTLLM Attention](attention_backends/trtllm.md#sage-quantization). -Install the optional CUDA 13 extra to use the CuTe-based FA4 path: +## RAINFUSION_ATTN Backend and Block-Sparse Video Attention -```bash -pip install 'vllm-omni[fa4]' -``` - -`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2 -when it is unavailable. Version `4.0.0b18` is required; earlier beta -wheels had known JIT failures on Blackwell. - -## Choosing a Backend Manually - -### When to override the default - -- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested. -- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate. -- **Workaround for known issues**: see Known Limitations above. - -### Verifying which backend is in use - -The startup log prints one of: - -``` -Using diffusion attention backend 'CUDNN_ATTN' # explicit override -Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route -Defaulting to diffusion attention backend SDPA # nothing else available -``` - -If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures. +See [RainFusion](attention_backends/rainfusion.md). ## SageAttention Installation -vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni. - -Build from source: - -```bash -git clone https://github.com/thu-ml/SageAttention.git -cd SageAttention - -export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32 -pip install . --no-build-isolation -``` - -Quick check: - -```bash -python -c "import sageattention; print(sageattention.__file__)" -``` +See [SageAttention](attention_backends/sage.md#installation). ## SageAttention3 Installation -vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni. - -Build from source: - -```bash -git clone https://github.com/thu-ml/SageAttention.git -cd SageAttention/sageattention3_blackwell -python setup.py install -``` - -Quick check: - -```bash -python -c "import sageattn3; print(sageattn3.__file__)" -``` - -Notes: - -- `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class. -- SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness. +See [SageAttention](attention_backends/sage.md#sageattention3-installation). ## HuggingFace Kernels Hub Backends -To achieve perfect numerical consistency between **training** (typically using HuggingFace Diffusers with Hub kernels) and **serving/rollout** (using vLLM-Omni), you can use the Hub-based attention backends. This eliminates numerical drift / sampling divergence caused by executing different local kernel versions during rollout. - -The following backend options are supported: -- `FLASH_ATTN_HUB` (HuggingFace `kernels-community/flash-attn2`) -- `FLASH_ATTN_3_HUB` (HuggingFace `kernels-community/flash-attn3`, Hopper sm_90+ only) - -### Installation - -To use these backends, you must install the `kernels` library: - -```bash -pip install kernels==0.14.1 -``` - -If the `kernels` library is not available in the environment, vLLM-Omni will log a warning and fall back gracefully to the corresponding local backend implementations (`FLASH_ATTN`). On CUDA GPUs below Hopper (compute capability < 9.0), `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`. - -### Usage - -Select a Hub backend using the global CLI flag or environment variables: - -```bash -# Environment variable -export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB - -# CLI flag -vllm-omni serve --diffusion-attention-backend FLASH_ATTN_HUB -``` - -## Usage Examples - -### Default (auto-route) - -```bash -python examples/offline_inference/text_to_video/text_to_video.py \ - --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \ - --prompt "A dog running across a field of golden wheat." \ - --height 480 --width 832 --num-frames 33 \ - --num-inference-steps 50 --seed 42 --guidance-scale 6.0 \ - --output hv15.mp4 -``` - -On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line. - -### Explicit backend selection - -```bash -DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \ - --model Lightricks/LTX-2 \ - --prompt "A dog running across a field of golden wheat." \ - --height 480 --width 832 --num-frames 33 \ - --num-inference-steps 40 --seed 42 --guidance-scale 4.0 \ - --output ltx2.mp4 -``` - -### SageAttention (lossy) - -```bash -DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \ - --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \ - --prompt "A dog running across a field of golden wheat." \ - --height 480 --width 832 --num-frames 33 \ - --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \ - --tensor-parallel-size 2 \ - --output hv15_sage.mp4 -``` - -Example: Wan2.2 TI2V 5B - -```bash -DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \ - --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \ - --prompt "A dog running across a field of golden wheat." \ - --height 704 --width 1280 --num-frames 49 \ - --num-inference-steps 30 --seed 42 --guidance-scale 5.0 \ - --tensor-parallel-size 2 \ - --output outputs/wan22_sage.mp4 -``` - -### Enable SageAttention3 - -Example: - -```bash -DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \ - --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \ - --prompt "A dog running across a field of golden wheat." \ - --height 480 --width 832 --num-frames 33 \ - --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \ - --tensor-parallel-size 2 \ - --output outputs/hv15_sage3.mp4 -``` +See [Hugging Face Hub Backends](attention_backends/huggingface_hub.md). diff --git a/docs/user_guide/diffusion/attention_backends/dense_backends.md b/docs/user_guide/diffusion/attention_backends/dense_backends.md new file mode 100644 index 000000000..030550dac --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends/dense_backends.md @@ -0,0 +1,110 @@ +# Dense Attention Backends + +Use these backends when you want dense attention without backend-specific +sparsity. Start with `TORCH_SDPA` for correctness comparisons, then benchmark +the fastest compatible kernel for your model, shape, and hardware. + +For selection precedence and per-role configuration, see the +[attention backend overview](../attention_backends.md). + +## `TORCH_SDPA` + +`TORCH_SDPA` calls PyTorch `scaled_dot_product_attention` and lets PyTorch's +dispatcher choose the implementation. It is always available and is the most +conservative reference when validating another backend. + +```bash +vllm-omni serve --diffusion-attention-backend TORCH_SDPA +``` + +## `FLASH_ATTN` + +`FLASH_ATTN` uses the installed FlashAttention implementation. On Blackwell it +prefers the CuTe-based FlashAttention 4 path and falls back to +FlashAttention 3 or 2 when that path is unavailable. On Hopper, Ada, and +Ampere it is the preferred automatic route when a compatible package is +installed. + +```bash +vllm-omni serve --diffusion-attention-backend FLASH_ATTN +``` + +### FlashAttention 4 on Blackwell + +Install the optional CUDA 13 extra: + +```bash +pip install 'vllm-omni[fa4]' +``` + +Version `4.0.0b18` is required; earlier beta wheels had known JIT failures on +Blackwell. If the CuTe path is unavailable, the backend falls back to the +compatible FlashAttention 3 or 2 path. + +## `CUDNN_ATTN` + +`CUDNN_ATTN` pins PyTorch SDPA to `CUDNN_ATTENTION`. It is particularly useful +for mask-heavy DiTs and is automatically preferred on Blackwell when cuDNN +9.5 or newer is available and the higher-priority TRTLLM route is not +compatible. + +```bash +vllm-omni serve --diffusion-attention-backend CUDNN_ATTN +``` + +### LTX-2.0 limitation + +LTX-2 audio attention has a symbolic head dimension during `torch.compile` +tracing. The cuDNN SDPA selector rejects that symbolic dimension and Dynamo +aborts compilation. This is tracked in +[issue #3121](https://github.com/vllm-project/vllm-omni/issues/3121). + +Use `FLASHINFER_ATTN` or `TORCH_SDPA` as a workaround: + +```bash +DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN \ + python examples/offline_inference/text_to_video/text_to_video.py \ + --model Lightricks/LTX-2 ... +``` + +## `FLASHINFER_ATTN` + +`FLASHINFER_ATTN` uses FlashInfer's batch-prefill wrapper. It is an explicit +option on CUDA platforms and an automatic Blackwell fallback when FlashInfer +is installed but cuDNN is too old for `CUDNN_ATTN`. + +```bash +vllm-omni serve --diffusion-attention-backend FLASHINFER_ATTN +``` + +### FlashInfer quantized attention + +The backend accepts an `AttentionSpec.quant` block. For QK16/V8, keep Q and K +in FP16 or BF16 and use FP8 E4M3 for V: + +```python +from vllm_omni.diffusion.data import ( + AttentionConfig, + AttentionSpec, + AttnQuantSpec, + OmniDiffusionConfig, +) + +config = OmniDiffusionConfig( + diffusion_attention_config=AttentionConfig( + default=AttentionSpec( + backend="FLASHINFER_ATTN", + quant=AttnQuantSpec( + dtype_qk="bfloat16", + dtype_vo="fp8_e4m3", + ), + ), + ), + ..., +) +``` + +`dtype_qk` controls Q and K; `dtype_vo` controls V. Mixed-dtype +configurations require FlashInfer 0.6.16rc1 or newer. The shared quantization +schema is also consumed by TRTLLM, but each backend validates its own allowed +fields and values; see [TRTLLM SAGE quantization](trtllm.md#sage-quantization). diff --git a/docs/user_guide/diffusion/attention_backends/huggingface_hub.md b/docs/user_guide/diffusion/attention_backends/huggingface_hub.md new file mode 100644 index 000000000..946786b97 --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends/huggingface_hub.md @@ -0,0 +1,32 @@ +# Hugging Face Hub Attention Backends + +Hub backends load the same published kernels commonly used with Hugging Face +Diffusers. They are useful when training and serving must use matching kernel +implementations to reduce numerical drift and sampling divergence. + +| Backend | Kernel | Platform | +| --- | --- | --- | +| `FLASH_ATTN_HUB` | `kernels-community/flash-attn2` | Compatible CUDA GPUs | +| `FLASH_ATTN_3_HUB` | `kernels-community/flash-attn3` | Hopper sm_90 or newer | + +## Installation + +```bash +pip install kernels==0.14.1 +``` + +If `kernels` is unavailable, vLLM-Omni warns and falls back to the +corresponding local FlashAttention implementation. On CUDA devices older than +Hopper, `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`. + +## Usage + +```bash +export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB + +# Equivalent CLI selection +vllm-omni serve --diffusion-attention-backend FLASH_ATTN_HUB +``` + +For per-role kernel alignment, configure the Hub backend on only the relevant +role; see the [attention backend overview](../attention_backends.md#per-role-configuration). diff --git a/docs/user_guide/diffusion/attention_backends/rainfusion.md b/docs/user_guide/diffusion/attention_backends/rainfusion.md new file mode 100644 index 000000000..10f7458be --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends/rainfusion.md @@ -0,0 +1,72 @@ +# RainFusion Attention + +`RAINFUSION_ATTN` runs MindIE-SD RainFusion (`rf_v2`) block-sparse video +attention on Ascend NPU. It pools 128-token key blocks, ranks them per query +block, and attends to the highest-scoring blocks after arranging video tokens +in `(t, h, w)` order. + +Only the video segment is sparse. Prefix rows such as text, visual conditions, +and audio, plus first-frame blocks, remain dense. Unsupported calls—including +warmup steps, skipped layers, missing video geometry, and video segments below +32 blocks—delegate to `FLASH_ATTN`, so compatible models can select RainFusion +globally. + +## Configuration + +| Key | Valid values | Meaning | +| --- | --- | --- | +| `sparsity` | finite, `[0, 1]` | Nominal dropped-key-block fraction; default `0.8`; `0` disables sparsity | +| `start_step` | integer, `>= 0` | Number of early denoise steps kept dense | +| `skip_layers` | selector such as `"0-3,38"` | DiT blocks kept dense | + +```bash +vllm-omni serve MiniMaxAI/MiniMax-H3 \ + --diffusion-attention-config '{"default":{"backend":"RAINFUSION_ATTN",\ + "block_sparse":{"sparsity":0.8,"start_step":0}}}' +``` + +```python +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec + +config = AttentionConfig( + default=AttentionSpec( + backend="RAINFUSION_ATTN", + block_sparse=BlockSparseSpec( + sparsity=0.8, + start_step=0, + skip_layers="0-1", + ), + ), +) +``` + +Tune `start_step`, then `sparsity`, then `skip_layers`. Increasing the dense +early-step window is usually the cheapest way to recover global structure; +use layer exclusions only after a same-seed dense comparison identifies +sensitive blocks. + +## Requirements and compatibility + +RainFusion requires Ascend NPU and `mindiesd`. Selecting it on another +platform raises. It is incompatible with ring sequence parallelism because +the kernel needs the complete key sequence for ranking; use Ulysses sequence +parallelism with `ring_degree=1`. + +## Geometry handling + +RainFusion handles arbitrary video grids by rearranging video spatially and, +when necessary, promoting a real-video suffix to the always-kept prefix. It +does not pad because `rf_v2` does not consume a padding attention mask. + +For MiniMax-H3 at 1344x768, grid `(62, 24, 42)` produces 62,496 video rows. +The implementation promotes 2,976 rows and leaves 59,520 sparse rows +(`465 x 128`) so the mask and kernel tiling remain aligned. + +Spatial grids aligned to 8x8 are still preferable. A latent height or width +not divisible by 8 creates an always-kept suffix and reduces realized +sparsity. At the image level, that means width and height divisible by 256. +Protected video tails require a compatible MindIE-SD release. + +For common configuration and selector behavior, see the +[attention backend overview](../attention_backends.md) and the +[backend selection design](../../../design/feature/attention_backend_selection.md). diff --git a/docs/user_guide/diffusion/attention_backends/sage.md b/docs/user_guide/diffusion/attention_backends/sage.md new file mode 100644 index 000000000..1e1fa5b60 --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends/sage.md @@ -0,0 +1,63 @@ +# SageAttention + +SageAttention backends provide lossy low-precision attention for diffusion +models. Validate output quality against `TORCH_SDPA` at the same seed before +using either backend in production. + +## `SAGE_ATTN` + +`SAGE_ATTN` uses SageAttention 2.2 with INT8-quantized attention and FP16 +accumulation. + +### Installation + +Install SageAttention into the same environment as vLLM-Omni: + +```bash +git clone https://github.com/thu-ml/SageAttention.git +cd SageAttention +export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32 +pip install . --no-build-isolation +``` + +Verify the installation: + +```bash +python -c "import sageattention; print(sageattention.__file__)" +``` + +Select it globally: + +```bash +vllm-omni serve --diffusion-attention-backend SAGE_ATTN +``` + +## `SAGE_ATTN_3` + +`SAGE_ATTN_3` uses the SageAttention3 Blackwell implementation. + +### SageAttention3 installation + +```bash +git clone https://github.com/thu-ml/SageAttention.git +cd SageAttention/sageattention3_blackwell +python setup.py install +``` + +Verify the installation: + +```bash +python -c "import sageattn3; print(sageattn3.__file__)" +``` + +```bash +vllm-omni serve --diffusion-attention-backend SAGE_ATTN_3 +``` + +`SAGE_ATTN_3` requires CUDA, an importable `sageattn3`, and a Blackwell-class +GPU. Its kernel assumes the query-head count equals the key/value-head count. +GQA and MQA diffusion calls therefore fall back to PyTorch SDPA for +correctness. + +For common configuration and platform routing, see the +[attention backend overview](../attention_backends.md). diff --git a/docs/user_guide/diffusion/attention_backends/trtllm.md b/docs/user_guide/diffusion/attention_backends/trtllm.md new file mode 100644 index 000000000..dd36ba61e --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends/trtllm.md @@ -0,0 +1,103 @@ +# TRTLLM Attention + +`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels. It supports dense +BF16 attention plus optional Skip-Softmax sparsity or SAGE quantization. + +For common selection and per-role configuration, see the +[attention backend overview](../attention_backends.md). + +## Requirements + +Selecting `TRTLLM_ATTN` requires: + +- datacenter Blackwell, sm_100 or sm_103; +- `head_dim=128`; +- FlashInfer; and +- a model-declared packed or mask-free compatible path. + +It is not supported on workstation Blackwell (sm_120/sm_121). An incompatible +explicit selection raises instead of silently changing the backend. + +## Dense mode + +Select `TRTLLM_ATTN` without `skip_softmax` or `quant` to run dense BF16: + +```bash +vllm-omni serve --diffusion-attention-backend TRTLLM_ATTN +``` + +## Skip-Softmax + +Skip-Softmax trades some fidelity for speed by skipping low-contribution +attention work. See the [feature design](../../../design/feature/skip_softmax.md) +for the algorithm and implementation contract. + +| Key | Valid values | Meaning | +| --- | --- | --- | +| `target_sparsity` | finite, `[0, 1]` | Point on a model-calibrated curve | +| `threshold` | finite, `>= 0` | Direct threshold; mutually exclusive with `target_sparsity` | +| `disabled_until_timestep` | finite, `[0, 1]` | Keep early high-noise steps dense until normalized `t <= D` | + +```bash +vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\ + "skip_softmax":{"target_sparsity":0.65,"disabled_until_timestep":0.86}}}' +``` + +```python +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec + +config = AttentionConfig( + default=AttentionSpec( + backend="TRTLLM_ATTN", + skip_softmax=SkipSoftmaxSpec( + target_sparsity=0.65, + disabled_until_timestep=0.86, + ), + ), +) +``` + +Start with `target_sparsity=0.65` and `disabled_until_timestep=0.86`. Increase +sparsity only after comparing quality against dense output at the same seed. +The benefit grows with sequence length because only attention is accelerated. + +## SAGE quantization + +TRTLLM SAGE quantizes Q/K per block and V per channel so both attention +matrix multiplications use low precision. + +| Key | Valid values | Meaning | +| --- | --- | --- | +| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype; absent means dense | +| `q_block_size` | `1`, `4`, `16` | Q scale block size; default `1` | +| `k_block_size` | `1`, `4`, `16` | K scale block size; default `16` | + +V is always quantized per channel to FP8 E4M3 and K-smoothing is internal. +Only block sizes with compiled kernels are accepted. + +```bash +vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\ + "quant":{"dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}' +``` + +```python +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec + +config = AttentionConfig( + default=AttentionSpec( + backend="TRTLLM_ATTN", + quant=AttnQuantSpec( + dtype_qk="fp8_e4m3", + q_block_size=1, + k_block_size=16, + ), + ), +) +``` + +This path requires FlashInfer 0.6.16rc1 or newer. FP8 Q/K kernels are +available on sm_100 and sm_103; INT8 Q/K kernels are compiled for sm_100 only. +The shared `AttnQuantSpec` is backend typed: values intended for +`FLASHINFER_ATTN` are rejected here and vice versa. diff --git a/docs/user_guide/diffusion/cpu_offload.md b/docs/user_guide/diffusion/cpu_offload.md index 02ab9f074..0bb798735 100644 --- a/docs/user_guide/diffusion/cpu_offload.md +++ b/docs/user_guide/diffusion/cpu_offload.md @@ -1,389 +1,75 @@ # CPU Offloading for Diffusion Models -## Overview +vLLM-Omni provides three CPU-offload strategies for diffusion models. Choose +the coarsest strategy that meets your memory target: finer-grained streaming +saves more device memory but adds more transfer and setup overhead. -vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models: +For the shared factory, discovery, and lifecycle contract, see the +[CPU Offloading design](../../design/feature/cpu_offloading.md). -1. **Model-level (Sequential) Offloading**: Mutual exclusion between DiT model and encoder - only one is on GPU at a time. -2. **Layerwise (Blockwise) Offloading**: Keeps only one transformer block on GPU at a time with compute-memory overlap. +## Choose a strategy -Both strategies use pinned memory for faster CPU-GPU transfers. The strategies are **mutually exclusive** for now - if both are enabled, layerwise takes priority. +| Strategy | Device residency | Parallel scope | Guide | +| --- | --- | --- | --- | +| Model-level (sequential) | One pipeline component group at a time | Single device | [Model-Level Offloading](cpu_offload/model_level.md) | +| Layerwise (blockwise) | One transformer block, with next-block prefetch | Single device | [Layerwise Offloading](cpu_offload/layerwise.md) | +| Distributed layerwise | Fixed two-block device buffer; optional host-weight sharding and AllGather | Multiple GPU/NPU ranks | [Distributed Layerwise Offloading](cpu_offload/distributed_layerwise.md) | +All strategies use pinned host memory for faster transfers where applicable. +Configuration priority is: + +1. Distributed layerwise offloading. +2. Layerwise offloading. +3. Model-level offloading. + +Treat the flags as mutually exclusive. If more than one is enabled, the +higher-priority strategy is selected. + +## Quick selection + +- Use [model-level offloading](cpu_offload/model_level.md) when swapping whole + encoders and DiTs is enough to fit and phase-boundary transfers are + acceptable. +- Use [layerwise offloading](cpu_offload/layerwise.md) for compute-heavy video + DiTs where block transfers can overlap computation. +- Use [distributed layerwise offloading](cpu_offload/distributed_layerwise.md) + when a multi-rank deployment also needs bounded device residency and, + optionally, sharded host weights. + +## Supported models + +| Architecture | Example models | DiT class | Model-level | Layerwise | Distributed layerwise | Layerwise block attributes | +| --- | --- | --- | :---: | :---: | :---: | --- | +| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | yes | yes | — | `transformer_blocks`, `single_transformer_blocks` | +| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | — | yes | — | `transformer_blocks`, `single_transformer_blocks` | +| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | — | yes | — | `layers` | +| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | — | yes | — | `transformer` | +| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | yes | yes | — | `transformer_blocks` | +| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | yes | yes | — | `down_blocks`, `up_blocks` | +| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | — | yes | — | `transformer_blocks` | +| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` | +| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` | +| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` | yes | yes | — | `layers` | +| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | — | yes | — | `layers`, customized modules | +| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | yes | yes | yes | `layers`, `gen_layers` | + +Model-level support requires discoverable DiT and encoder components. +Layerwise support requires transformer block topology. Distributed support +reuses that topology but still requires validation for the model, checkpoint, +and parallel configuration. + +## Compatibility anchors + +These headings preserve links to sections that moved into dedicated guides. ## Model-level (Sequential) Offloading -### How It Works - -Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks: - -- **When encoders run**: DiT transformer is offloaded to CPU -- **When DiT runs**: Encoders are offloaded to CPU, if more than one dit models, only one loaded on GPU, others get offloaded to CPU. -- **VAE**: Stays resident on GPU - -Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed. - -### Usage - -**Python API:** -```python -from vllm_omni import Omni - -m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True) -``` - -**CLI:** -```bash -vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-cpu-offload -``` - -### To Support a Model - -Implement the `SupportsComponentDiscovery` protocol to declare which -submodules serve as pipeline components (used by offloading, HSDP -sharding, and other framework features): - -```python -from typing import ClassVar -from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery - -class MyPipeline(nn.Module, SupportsComponentDiscovery): - _dit_modules: ClassVar[list[str]] = ["transformer"] - _encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"] - _vae_modules: ClassVar[list[str]] = ["vae"] - _resident_modules: ClassVar[list[str]] = [] # optional - - def __init__(self): - super().__init__() - self.transformer = ... # DiT — stays on GPU during denoising - self.text_encoder = ... # Encoder — offloaded to CPU during denoising - self.vision_model = ... # Encoder — offloaded to CPU during denoising - self.vae = ... # VAE — always on GPU -``` - -- `_dit_modules`: attribute names of denoising submodules (kept on GPU - during the diffusion loop). -- `_encoder_modules`: attribute names of encoder/vision submodules - (offloaded to CPU during the diffusion loop). -- `_vae_modules`: attribute names of VAE(s) (always kept on GPU, not - part of the mutual exclusion hooks). -- `_resident_modules`: attribute names of small submodules that must - stay on GPU during layerwise offloading (e.g. embedders, connectors). - Optional — defaults to `[]`. - -All attribute names support dotted paths for nested submodules -(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`). - -Both DiT and encoder lists are needed because the offload hooks use -mutual exclusion: when one group runs, the other moves to CPU. - -### Limitations -- Cold start latency increases -- Adds overhead from CPU-GPU transfers between encoder and denoising phases -- Support single GPU only for now - - -### Component offloading for split models (e.g. Cosmos3) - -Some models split their transformer into mutually-exclusive *components* that run -in different phases of a single forward pass rather than as separate pipeline -components -- e.g. Cosmos3's understanding (reasoner) component runs once per -generation while the generation (generator) component runs every denoising step. -Such models have no separate text encoder to swap against, so the transformer -owns a small model-local offload path and wraps each phase with -`with self._offload_context(name):` - -```python -class Cosmos3VFMTransformer(nn.Module): - def forward(self, ...): - with self._offload_context("reasoner"): - ... # understanding pass, runs once - with self._offload_context("generator"): - ... # denoising pass, runs every step -``` - -Model-level offloading then keeps exactly one component GPU-resident at a time -(the other on CPU), reusing the same `SequentialOffloadHook` `.to()` movers. The -pipeline opts in by exposing `enable_omni_model_cpu_offload` (which drives the -transformer's `enable_model_cpu_offload` and pins the VAE). Layerwise offloading -works for these models too -- each component declares its own block container via -`_layerwise_offload_blocks_attrs`. - +See [Model-Level Offloading](cpu_offload/model_level.md). ## Layerwise (Blockwise) Offloading -### How It Works - -Layerwise offloading keeps only one transformer block on GPU at a time. - -As each block completes, the next block is prefetched to GPU while the current block is freed. The pre and forward hooks utilized by layerwise offloading apply a separate CUDA stream (`copy_stream`) to overlap weight transfer with computation, and retain flattened tensors in pinned CPU memory for block parameters re-materialization. Encoders, VAE, and non-block DiT modules (embeddings, norms) always stay on GPU. - -**Execution Flow:** - -| Block | Pre-forward Hook | Forward | Post-forward Hook | -|-------|------------------|---------|-------------------| -| block-0 | Prefetch block-1 (async) | Compute block-0 | Free block-0 | -| block-1 | Prefetch block-2 (async) | Compute block-1 | Free block-1 | -| ... | ... | ... | ... | -| block-(n-1) | **Prefetch block-0** (async) | Compute block-(n-1) | Free block-(n-1) | - -Each transformer block has a `LayerwiseOffloadHook` that prefetches the next block before forward and frees the current block after forward. - -Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. For example, Wan2.2 T2V and I2V pipelines. - -### Usage - -**Python API:** -```python -from vllm_omni import Omni - -# Text-to-video -m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_layerwise_offload=True) - -# Or image-to-video -m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True) -``` - -**CLI:** -```bash -# Text-to-video -vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-layerwise-offload - -# Or image-to-video -vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --enable-layerwise-offload -``` - -### To Support a Model - -Models must define the blocks attribute name for layerwise offloading: - -```python -class WanTransformer3DModel(nn.Module): - _layerwise_offload_blocks_attrs = ["blocks"] # Attribute names containing transformer blocks - - def __init__(self): - self.blocks = nn.ModuleList([...]) # Transformer blocks -``` - -For models with multiple block types: - -```python -class Flux2Transformer2DModel(nn.Module): - _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"] -``` - -### Limitations -- Cold start latency increases because offloaded components must be moved to CPU - during setup; layerwise offload may add extra weight consolidation and pinning - work. -- Performance depends on compute cost and H2D bandwidth as well -- Support single GPU only for now - +See [Layerwise Offloading](cpu_offload/layerwise.md). ## Distributed Layerwise Offloading -### How It Works - -Distributed layerwise offloading extends single-GPU layerwise offloading to -multi-device deployments. In the default mode, each rank stores only -**1/dp_size** of the model weights in host memory; full layer weights are -reconstructed at runtime via **AllGather** on a dedicated communication stream, -overlapped with computation via a fixed double-buffer scheme. The -`--dlo-no-use-allgather` mode instead streams standard-loader rank-local -weights independently and does not shard weights across ranks. - -For the implementation design and compatibility matrix, see -[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md). - -**Key features:** -- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights, - reconstructed per-layer via `all_gather_into_tensor` -- **Fixed double-buffer**: exactly 2 transformer blocks on each device at any - time, regardless of model size -- **DP multi-concurrency**: N concurrent requests processed in parallel - (AllGather only gathers weight shards, which are request-independent) -- **mmap weight loading**: weights loaded as mmap views pointing to shared OS - page cache, eliminating O(dp_size × model_size) RSS during model creation -- **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU - (CANN/HCCL) via vLLM-Omni's platform abstraction - -**Execution Flow (per device):** - -``` -Compute Stream: [Layer N] [Layer N+1] [Layer N+2] -H2D Stream: [H2D Shard N+1] [H2D Shard N+2] -AllGather: [AG N+1] [AG N+2] -Slot Usage: Slot0: Layer N Slot1: Layer N+1 -``` - -### Usage - -**CLI:** -```bash -# 4× GPU/NPU with AllGather (recommended) -vllm serve /path/to/model --omni \ - --enable-distributed-layerwise-offload \ - --data-parallel-size 4 - -# Without DLO AllGather (each rank streams standard-loader rank-local weights) -vllm serve /path/to/model --omni \ - --enable-distributed-layerwise-offload \ - --data-parallel-size 4 \ - --dlo-no-use-allgather - -# With SP instead of DP (long sequences) -vllm serve /path/to/model --omni \ - --enable-distributed-layerwise-offload \ - --usp 4 -``` - -**Python API:** -```python -from vllm_omni import Omni - -m = Omni( - model="/path/to/model", - enable_distributed_layerwise_offload=True, - dlo_use_allgather=True, # default -) -``` - -### CLI Flags - -| Flag | Description | Default | -|------|-------------|---------| -| `--enable-distributed-layerwise-offload` | Enable DLO | `false` | -| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` | -| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` | -| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` | - -### How model weights are loaded (mmap path) - -When DLO + AllGather is active, the offloader: - -1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`) - from the normally-created transformer -2. **Converts to meta device** via `to_empty(device="meta")`, releasing random - initialization weights -3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`, - which return views into the OS page cache (shared across all ranks, 0 RSS) -4. **Calls `post_load_weights()`** to apply model-specific dtype conversions - (e.g. Cosmos3's `time_embedder` → FP32) -5. **Restores non-persistent buffers** from saved copies - -This approach requires **zero model-specific code changes** — no pipeline or -transformer modifications are needed. - -### OffloadPlan (declarative topology metadata) - -Models can optionally declare an `OffloadPlan` class variable to provide -topology metadata (block attribute names, submodules to offload) without -any offload-specific logic: - -```python -from vllm_omni.diffusion.offloader import OffloadPlan - -class MyPipeline(nn.Module): - _dit_modules = ["transformer"] - _offload_plan = OffloadPlan( - block_attrs={"transformer": ("blocks",)}, - offload_submodules={"context_encoder": "layers"}, - ) -``` - -When not declared, the offloader falls back to `_layerwise_offload_blocks_attrs` -and heuristic attribute search. This is backward-compatible — existing models -work without any changes. - -### DP Multi-concurrency - -When `--data-parallel-size > 1` and AllGather is enabled, the scheduler batches -up to `dp_size` requests per denoise step. Each DP rank processes a different -request while AllGather synchronizes only weight shards (request-independent). - -**Requirements for concurrent requests:** -- `num_inference_steps` must be specified explicitly (None is not allowed) -- All concurrent requests must have the same `num_inference_steps` value -- These constraints exist because AllGather is a collective that requires - every rank to participate at each denoise step - -### Limitations -- Online quantization (FP8) is incompatible with mmap loading — falls back to - regular loading only when the model is outside the DLO+AllGather mmap path; - online quantization with DLO+AllGather is rejected. Use - `--dlo-no-use-allgather` or disable online quantization. -- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware - weight-loader callbacks are bypassed. The no-AllGather path retains the - standard TP loader and is experimental. -- HSDP + AllGather is rejected because it would double-shard parameters. HSDP - with no-AllGather is accepted at configuration level, but has limited - end-to-end validation. -- `num_inference_steps=None` is not allowed in DP multi-concurrency mode - -**Module Discovery** - -The offloader discovers pipeline components in two ways: - -1. **Protocol-based** (preferred): If the pipeline implements - `SupportsComponentDiscovery`, its `_dit_modules`, `_encoder_modules`, - `_vae_modules`, and `_resident_modules` class variables are used - directly. All attribute names support dotted paths (e.g. - `"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules. - -2. **Fallback attribute scan**: Otherwise, the offloader scans for - well-known attribute names: - - **DiT modules**: `transformer`, `transformer_2`, `dit`, `sr_dit`, `language_model`, `transformer_blocks`, `model` - - **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder` - - **VAE**: `vae`, `audio_vae` - -**Hook System** - -Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelHook`) to register pre/post forward callbacks on modules, enabling automatic swapping without modifying model code. - -**Backend Architecture** - -``` -OffloadBackend (base class) -├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap) -│ (delegates to a pipeline's enable_omni_model_cpu_offload -│ for split models like Cosmos3) -├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook -│ (single-GPU, full weights on host) -└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook - (multi-GPU, sharded AllGather or rank-local mode) -``` - -Factory function `get_offload_backend()` selects the appropriate backend based on -configuration. - -For split models, `ModelLevelOffloadBackend.enable()` detects a pipeline's -`enable_omni_model_cpu_offload` hook and delegates to it; Cosmos3 then swaps its -reasoner/generator components inside the model forward pass. - - -## Supported Models - -| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Distributed Layerwise Offload | Blocks Attrs (Layerwise specific) | -|--------------|----------------|-----------|---------------------|-------------------|-------------------------------|-----------------------------------| -| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` | -| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` | -| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | - | `"layers"` | -| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | - | `"transformer"` | -| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"` | -| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | - | `"down_blocks"`, `"up_blocks"` | -| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` | -| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` | -| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` | -| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` | -| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` | -| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` | - -**Notes:** -- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally -- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks -- Distributed Layerwise Offloading reuses the layerwise block topology, but the - model and parallelism combination still needs validation. AllGather mode - additionally requires a compatible checkpoint/mmap layout; no-AllGather mode - retains the standard loader's rank-local tensors. See the - [DLO design](../../design/feature/distributed_layerwise_offload.md) and - [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) - for usage examples. +See [Distributed Layerwise Offloading](cpu_offload/distributed_layerwise.md). diff --git a/docs/user_guide/diffusion/cpu_offload/distributed_layerwise.md b/docs/user_guide/diffusion/cpu_offload/distributed_layerwise.md new file mode 100644 index 000000000..9f5ae3493 --- /dev/null +++ b/docs/user_guide/diffusion/cpu_offload/distributed_layerwise.md @@ -0,0 +1,118 @@ +# Distributed Layerwise Offloading + +Distributed layerwise offloading (DLO) extends block streaming to multi-device +deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of +the host weights and reconstructs each layer at runtime. Without AllGather, +each rank streams its standard-loader rank-local weights independently. + +See the [DLO feature design](../../../design/feature/distributed_layerwise_offload.md) +for the implementation contract and compatibility matrix. + +## Execution model + +DLO overlaps three operations with a fixed two-block device buffer: + +```text +Compute stream: [Layer N] [Layer N+1] [Layer N+2] +H2D stream: [H2D shard N+1] [H2D shard N+2] +AllGather: [AG N+1] [AG N+2] +Slots: slot 0: Layer N slot 1: Layer N+1 +``` + +AllGather communicates only request-independent weight shards, so data- +parallel ranks may process different requests concurrently. + +## Usage + +```bash +# Four ranks with sharded host weights and AllGather +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 4 + +# Standard-loader rank-local weights, without DLO AllGather +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 4 \ + --dlo-no-use-allgather + +# Sequence parallel deployment +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --usp 4 +``` + +```python +from vllm_omni import Omni + +omni = Omni( + model="/path/to/model", + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, +) +``` + +## Flags + +| Flag | Meaning | Default | +| --- | --- | --- | +| `--enable-distributed-layerwise-offload` | Enable DLO | `false` | +| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` | +| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` | +| `--dlo-no-use-allgather` | Keep standard-loader rank-local weights | `false` | +| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` | + +## mmap weight loading + +The DLO plus AllGather path: + +1. saves non-persistent buffers such as RoPE frequencies; +2. moves the normally created transformer to the meta device; +3. loads checkpoint tensors as mmap views backed by the shared OS page cache; +4. calls model-specific `post_load_weights()` conversions; and +5. restores the saved non-persistent buffers. + +This avoids one full checkpoint RSS copy per rank and does not require +model-specific loading code. + +## Declarative topology + +Models may declare an `OffloadPlan` instead of embedding offload logic: + +```python +from vllm_omni.diffusion.offloader import OffloadPlan + + +class MyPipeline(nn.Module): + _dit_modules = ["transformer"] + _offload_plan = OffloadPlan( + block_attrs={"transformer": ("blocks",)}, + offload_submodules={"context_encoder": "layers"}, + ) +``` + +When no plan exists, discovery falls back to +`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup. + +## Data-parallel concurrency + +With `data_parallel_size > 1` and AllGather enabled, the scheduler can process +up to `dp_size` requests per denoising step. Every concurrent request must set +the same explicit `num_inference_steps`; `None` is rejected because every rank +must enter each collective. + +## Limitations + +- Online FP8 quantization is rejected with the DLO plus AllGather mmap path. + Use `--dlo-no-use-allgather` or disable online quantization. +- Tensor parallel size greater than one is rejected in the mmap path because + it bypasses TP-aware loader callbacks. The no-AllGather path retains the + standard loader and remains experimental with TP. +- HSDP plus AllGather is rejected to avoid double sharding. HSDP without + AllGather has limited end-to-end validation. +- Resident leading layers require `--dlo-no-use-allgather` and a model + `OffloadPlan` that declares eligible `resident_dit_paths`. +- DP concurrency requires an explicit, identical inference-step count. + +See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) +for an end-to-end example. diff --git a/docs/user_guide/diffusion/cpu_offload/layerwise.md b/docs/user_guide/diffusion/cpu_offload/layerwise.md new file mode 100644 index 000000000..041278b80 --- /dev/null +++ b/docs/user_guide/diffusion/cpu_offload/layerwise.md @@ -0,0 +1,65 @@ +# Layerwise Offloading + +Layerwise, or blockwise, offloading keeps one transformer block on the +accelerator and prefetches the next block while the current block computes. +It is best suited to compute-heavy video DiTs whose block execution can hide +host-to-device transfers. + +## Execution flow + +Each block has a pre-forward and post-forward hook. Parameters are consolidated +in pinned host tensors and rematerialized for execution on a dedicated copy +stream. + +| Block | Pre-forward hook | Forward | Post-forward hook | +| --- | --- | --- | --- | +| block 0 | Prefetch block 1 | Compute block 0 | Free block 0 | +| block 1 | Prefetch block 2 | Compute block 1 | Free block 1 | +| ... | ... | ... | ... | +| last block | Prefetch block 0 | Compute last block | Free last block | + +Encoders, VAE modules, and non-block DiT modules such as embeddings and norms +remain device resident. + +## Usage + +```python +from vllm_omni import Omni + +omni = Omni( + model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", + enable_layerwise_offload=True, +) +``` + +```bash +vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --omni --enable-layerwise-offload +``` + +## Model integration + +Transformer classes declare containers of executable blocks: + +```python +class WanTransformer3DModel(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + +class Flux2Transformer2DModel(nn.Module): + _layerwise_offload_blocks_attrs = [ + "transformer_blocks", + "single_transformer_blocks", + ] +``` + +See the [layerwise design](../../../design/feature/layerwise_offload.md) for +the discovery and hook invariants. `OffloadPlan` is a separate declarative +topology path for distributed layerwise offload. + +## Limitations + +- Single device only. +- Setup consolidates and pins block parameters, increasing cold-start time. +- Performance depends on block compute time and host-to-device bandwidth; + lightweight blocks may not hide transfers. diff --git a/docs/user_guide/diffusion/cpu_offload/model_level.md b/docs/user_guide/diffusion/cpu_offload/model_level.md new file mode 100644 index 000000000..017358e89 --- /dev/null +++ b/docs/user_guide/diffusion/cpu_offload/model_level.md @@ -0,0 +1,81 @@ +# Model-Level Offloading + +Model-level, or sequential, offloading keeps only the pipeline component group +currently executing on the accelerator. It is the simplest offload strategy +and is selected with `--enable-cpu-offload`. + +## How it works + +Pre-forward hooks enforce mutual exclusion between DiT and encoder modules: + +- before an encoder runs, the DiT moves to CPU; +- before a DiT runs, encoders and other DiTs move to CPU; and +- VAE modules remain on the accelerator. + +Pinned host memory reduces transfer overhead. Transfers occur at phase +boundaries, so cold-start and encoder-to-denoiser transitions become slower. + +## Usage + +```python +from vllm_omni import Omni + +omni = Omni( + model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", + enable_cpu_offload=True, +) +``` + +```bash +vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --omni --enable-cpu-offload +``` + +## Model integration + +Pipelines should implement `SupportsComponentDiscovery`: + +```python +from typing import ClassVar + +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery + + +class MyPipeline(nn.Module, SupportsComponentDiscovery): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] +``` + +All entries may be dotted paths. DiT and encoder lists are both required for +mutual exclusion. VAE modules are pinned but not swapped; resident modules are +small modules that must stay on the accelerator for layerwise paths. + +## Split-model components + +Some models, such as Cosmos3, split one transformer into mutually exclusive +components that run in different phases. The pipeline exposes +`enable_omni_model_cpu_offload`, and the backend delegates to the model-local +contexts: + +```python +class Cosmos3VFMTransformer(nn.Module): + def forward(self, ...): + with self._offload_context("reasoner"): + ... + with self._offload_context("generator"): + ... +``` + +This preserves the same invariant—exactly one component is device resident— +while reusing sequential `.to()` movers. + +## Limitations + +- Single device only. +- Higher cold-start latency. +- Transfers between encoder and denoising phases add latency. + +See the [model-level design](../../../design/feature/model_level_offload.md) +for lifecycle and extension invariants. diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 7c664e88b..67ace0940 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -106,7 +106,10 @@ The following tables show which models support each feature: > Notes: -> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload.md). +> 1. CPU Offload has three strategies: model-level, layerwise, and distributed +> layerwise. The tables below show **layerwise support** only. Split models +> like Cosmos3 swap their reasoner/generator components for model-level +> offload; see the [CPU Offload Guide](diffusion/cpu_offload.md). > 2. The **💾Quantization** column is collapsed for readability. See [Quantization](quantization/overview.md) for per-method and per-model support details. ### ImageGen @@ -214,7 +217,7 @@ The following tables show which models support each feature: 3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible. 4. The CPU Offloading (Layerwise) row describes local layerwise offload. Multi-device Distributed Layerwise Offload has a separate topology and - compatibility matrix in the [CPU Offloading guide](diffusion/cpu_offload.md#distributed-layerwise-offloading). + compatibility matrix in the [Distributed Layerwise Offloading guide](diffusion/cpu_offload/distributed_layerwise.md). 5. The compatibility matrix uses FP8 as the representative quantization method. 6. Step Execution is not compatible with any diffusion cache backend. LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches). diff --git a/docs/user_guide/feature_compatibility.md b/docs/user_guide/feature_compatibility.md index ea05fe81e..eeb087510 100644 --- a/docs/user_guide/feature_compatibility.md +++ b/docs/user_guide/feature_compatibility.md @@ -165,7 +165,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \ - **Tensor Parallelism — Text Encoder Not Sharded**: TP currently only shards the DiT blocks. Each TP rank retains a **full copy of the text encoder weights**, leading to significant GPU memory overhead proportional to TP degree. Tracked in [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771). -- **CPU Offloading — Two Modes Are Mutually Exclusive**: Model-level offload (`enable_cpu_offload`) and layerwise offload (`enable_layerwise_offload`) cannot be used simultaneously. If both are set, layerwise takes priority and model-level is silently ignored. +- **CPU Offloading — Strategies Are Mutually Exclusive**: Distributed layerwise, layerwise, and model-level offload should not be enabled together. Selection priority is distributed layerwise, then layerwise, then model-level. - **CPU Offloading — VAE stays on GPU**: Both offloading strategies keep the VAE on GPU at all times. For high-resolution generation, VAE decode can still cause OOM. Mitigate by combining with `vae_use_tiling=True` or VAE Patch Parallelism.