docs: split diffusion feature guides
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
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.
|
||||
@@ -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
|
||||
|
||||
@@ -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).
|
||||
@@ -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
@@ -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
|
||||
|
||||
|
||||
Reference in new issue
Block a user