docs: split diffusion feature guides

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed 2026-08-12 08:55:41 +08:00
1 parent 8b0a6779d4
commit 7872663bf0
21 files changed
+1138 -801

No files matched your search

@@ -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.
+80
View File
@@ -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.
@@ -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
+57
View File
@@ -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).
@@ -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).
+4 -2
View File
@@ -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
+10 -3
View File
@@ -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