diff --git a/docs/.nav.yml b/docs/.nav.yml index adf7321d7..16821c595 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -54,6 +54,12 @@ nav: - contributing/model/README.md - contributing/model/adding_omni_model.md - contributing/model/adding_diffusion_model.md + - Advanced Features: + - contributing/features/cfg_parallel.md + - contributing/features/sequence_parallel.md + - contributing/features/tensor_parallel.md + - contributing/features/cache_dit.md + - contributing/features/teacache.md - CI: contributing/ci - Design Documents: - design/index.md diff --git a/docs/contributing/features/cache_dit.md b/docs/contributing/features/cache_dit.md new file mode 100644 index 000000000..f39b93578 --- /dev/null +++ b/docs/contributing/features/cache_dit.md @@ -0,0 +1,286 @@ +# Support Cache-DiT + +This section describes how to add cache-dit acceleration to a new diffusion pipeline. We use the Qwen-Image pipeline and LongCat-Image pipeline as reference implementations. + +--- + +## Table of Contents + +- [Overview](#overview) +- [Standard Models: Automatic Support](#standard-models-automatic-support) +- [Custom Architectures: Writing Custom Implementation](#custom-architectures-writing-custom-implementation) +- [Testing](#testing) +- [Troubleshooting](#troubleshooting) +- [Reference Implementations](#reference-implementations) +- [Summary](#summary) + +--- + +## Overview + +### What is Cache-DiT? + +Cache-DiT is an acceleration library for Diffusion Transformers (DiT) that caches intermediate computation results across denoising steps. The core insight is that adjacent denoising steps often produce similar intermediate features, so we can skip redundant computations by reusing cached results. + +The library supports three main caching strategies: + +- **DBCache:** Dynamic block-level caching that selectively computes or caches transformer blocks based on residual differences +- **TaylorSeer:** Calibration-based prediction that estimates block outputs using Taylor expansion +- **SCM (Step Computation Masking):** Dynamic step skipping based on configurable policies + +### Architecture + +vLLM-omni integrates cache-dit through the `CacheDiTBackend` class, which provides a unified interface for managing cache-dit acceleration on diffusion models. + +| Method/Class | Purpose | Behavior | +|--------------|---------|----------| +| [`CacheDiTBackend`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/#vllm_omni.diffusion.cache.CacheBackend) | Unified backend interface | Automatically handles enabler selection and cache refresh | +| [`enable_cache_for_dit()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Apply caching to transformer | Configures DBCache on transformer blocks | + +**Key APIs from Cache-DiT:** + +[Cache-DiT API Reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/) + +| API | Description | +|-----|-------------| +| `BlockAdapter` | Core abstraction for applying cache-dit to transformers. Specifies transformer module(s), block list(s), and forward signature pattern(s). | +| `ForwardPattern` | Defines block forward signature patterns: `Pattern_0`, `Pattern_1`, `Pattern_2` | +| `ParamsModifier` | Per-transformer or per-block-list cache configuration customization | +| `DBCacheConfig` | Configuration for DBCache parameters (warmup steps, cached steps, thresholds) | +| `refresh_context()` | Update cache context | Called when `num_inference_steps` changes | + +--- + +## Standard Models: Automatic Support + +Most DiT models follow this pattern: +- Single transformer with one `ModuleList` of blocks +- Standard forward signature +- Compatible with cache-dit's automatic detection + +**Examples:** Qwen-Image, Z-Image + +For standard single-transformer models, **no code changes are needed**. The `CacheDiTBackend` automatically uses `enable_cache_for_dit()`: + +```python +from vllm_omni import Omni + +# Works automatically for standard models +omni = Omni( + model="Qwen/Qwen-Image", # Standard single-transformer model + cache_backend="cache_dit", + cache_config={ + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + } +) +``` + +**What happens automatically:** + +```python +def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], None]: + """Default enabler for standard single-transformer DiT models.""" + + # Build cache configuration + db_cache_config = DBCacheConfig( + num_inference_steps=None, # Will be set during first inference + Fn_compute_blocks=cache_config.Fn_compute_blocks, + Bn_compute_blocks=cache_config.Bn_compute_blocks, + max_warmup_steps=cache_config.max_warmup_steps, + max_cached_steps=cache_config.max_cached_steps, + max_continuous_cached_steps=cache_config.max_continuous_cached_steps, + residual_diff_threshold=cache_config.residual_diff_threshold, + ) + + # Enable cache-dit on transformer + cache_dit.enable_cache( + pipeline.transformer, + cache_config=db_cache_config, + ) + + # Return refresh function for dynamic num_inference_steps updates + def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True): + cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) + + return refresh_cache_context +``` + +--- + +## Custom Architectures: Writing Custom Implementation + +Some models require custom handling: + +- **Dual-transformer:** Models with separate high-noise and low-noise transformers (e.g., Wan2.2) +- **Multi-block-list:** Models with multiple block lists in one transformer (e.g., LongCatImage with `transformer_blocks` + `single_transformer_blocks`) +- **Special forward patterns:** Models with non-standard block execution patterns + +### Example 1: Dual-Transformer Model (Wan2.2) + +Wan2.2 uses two transformers: one for high-noise steps and one for low-noise steps. + +**Key difference:** Use `BlockAdapter` to wrap multiple transformers with separate configurations. + +```python +# Standard: cache_dit.enable_cache(pipeline.transformer, ...) +# Custom: Use BlockAdapter to handle multiple transformers +cache_dit.enable_cache( + BlockAdapter( + transformer=[pipeline.transformer, pipeline.transformer_2], # Multiple transformers + blocks=[pipeline.transformer.blocks, pipeline.transformer_2.blocks], + forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2], + params_modifiers=[ + ParamsModifier(...), # Config for high-noise transformer + ParamsModifier(...), # Config for low-noise transformer (different params) + ], + ), + cache_config=db_cache_config, +) +``` + +**Key difference:** `refresh_context` must be called on each transformer separately. + +```python +# Standard: cache_dit.refresh_context(pipeline.transformer, num_inference_steps=N) +# Custom: Refresh each transformer with its own step count +def refresh_cache_context(pipeline, num_inference_steps, verbose=True): + high_steps, low_steps = _split_inference_steps(num_inference_steps) + cache_dit.refresh_context(pipeline.transformer, num_inference_steps=high_steps, ...) + cache_dit.refresh_context(pipeline.transformer_2, num_inference_steps=low_steps, ...) +``` + +### Example 2: Multi-Block-List Model (LongCatImage) + +LongCatImage has a single transformer with two block lists: `transformer_blocks` and `single_transformer_blocks`. + +**Key difference:** Use `BlockAdapter` to specify multiple block lists within one transformer. + +```python +# Standard: cache_dit.enable_cache(pipeline.transformer, ...) +# - Automatically detects single block list +# Custom: Use BlockAdapter to specify multiple block lists +cache_dit.enable_cache( + BlockAdapter( + transformer=pipeline.transformer, # Single transformer + blocks=[ + pipeline.transformer.transformer_blocks, # Block list 1 + pipeline.transformer.single_transformer_blocks, # Block list 2 + ], + forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1], + params_modifiers=[modifier], + ), + cache_config=db_cache_config, +) +``` + +> **Note:** For single transformer with multiple block lists, `refresh_context` works the same as standard models. + +### Registering Custom Implementations + +After writing your custom enabler, register it in `CUSTOM_DIT_ENABLERS` in `vllm_omni/diffusion/cache/cache_dit_backend.py`: + +```python +CUSTOM_DIT_ENABLERS = { + "Wan22Pipeline": enable_cache_for_wan22, + "LongCatImagePipeline": enable_cache_for_longcat_image, + "YourCustomPipeline": enable_cache_for_your_model, # Add here +} +``` + +--- + +## Testing + +After adding cache-dit support, test with: + +```python +from vllm_omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +# Test your custom model +omni = Omni( + model="your-model-name", + cache_backend="cache_dit", + cache_config={ + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + "residual_diff_threshold": 0.24, + } +) + +images = omni.generate( + "a beautiful landscape", + OmniDiffusionSamplingParams(num_inference_steps=50), +) +``` + +**Verify:** + +1. Cache is applied (check logs for "Cache-dit enabled successfully on xxx") +2. Performance improvement (should be around 1.5x-2x faster) +3. Image quality (compare with `cache_backend=None`) + +--- + +## Troubleshooting + +### Issue: Cache not applied + +**Symptoms:** No speedup observed, no cache-related log messages. + +**Causes & Solutions:** + +- **Enabler not registered:** + +**Problem:** Pipeline name not in `CUSTOM_DIT_ENABLERS` registry. + +**Solution:** Verify `pipeline.__class__.__name__` matches the registry key and add your enabler to `CUSTOM_DIT_ENABLERS`. + +### Issue: Quality degradation + +**Symptoms:** Generated images have artifacts or lower quality compared to non-cached inference. + +**Causes & Solutions:** + +- **Cache parameters too aggressive:** + +**Solution:** +```python +cache_config={ + "residual_diff_threshold": 0.12, # Lower from 0.24 (try 0.12-0.18) + "max_warmup_steps": 6, # Increase from 4 (try 6-8) + "max_continuous_cached_steps": 2, # Reduce if higher +} +``` + +Check the [user guide for cache_dit](../../user_guide/diffusion/cache_dit_acceleration.md) for more adjustable parameters. + +--- + +## Reference Implementations + +Complete examples in the codebase: + +| Model | Path | Pattern | Notes | +|-------|------|---------|-------| +| **Standard DiT** | [`cache_dit_backend.py::enable_cache_for_dit`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Default enabler | Single transformer, automatic | +| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Dual-transformer | Separate high/low noise transformers | +| **LongCat** | [`cache_dit_backend.py::enable_cache_for_longcat_image`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_longcat_image) | Multi-block-list | Two block lists in one transformer | +| **BAGEL** | [`cache_dit_backend.py::enable_cache_for_bagel`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_bagel) | Omni model | Complex architecture | + +--- + +## Summary + +Adding cache-dit support: + +1. ✅ **Check model type** - Standard models work automatically, custom architectures need enablers +2. ✅ **Write enabler** (if needed) - Use `BlockAdapter` for complex architectures +3. ✅ **Register enabler** (if needed) - Add to `CUSTOM_DIT_ENABLERS` dictionary +4. ✅ **Return refresh function** (if needed) - Handle `num_inference_steps` changes +5. ✅ **Test** - Verify with `cache_backend="cache_dit"` + +For most models, the default enabler is sufficient. Only write custom enablers for complex architectures! diff --git a/docs/contributing/features/cfg_parallel.md b/docs/contributing/features/cfg_parallel.md new file mode 100644 index 000000000..adfe447c9 --- /dev/null +++ b/docs/contributing/features/cfg_parallel.md @@ -0,0 +1,249 @@ +# Support CFG-Parallel + +This section describes how to add CFG-Parallel (Classifier-Free Guidance Parallel) to a diffusion pipeline. We use the Qwen-Image pipeline as the reference implementation. + +--- + +## Table of Contents + +- [Overview](#overview) +- [Step-by-Step Implementation](#step-by-step-implementation) +- [Customization](#customization) +- [Testing](#testing) +- [Troubleshooting](#troubleshooting) +- [Reference Implementations](#reference-implementations) +- [Summary](#summary) + +--- + +## Overview + +### What is CFG-Parallel? + +In standard Classifier-Free Guidance, each diffusion step requires two forward passes through the transformer: + +1. **Positive/Conditional**: Guided by the text prompt +2. **Negative/Unconditional**: Typically using empty or negative prompt + +CFG-Parallel eliminates this bottleneck by distributing the two forward passes across different GPU ranks, allowing them to execute simultaneously rather than sequentially. + +### Architecture + +vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic. Pipelines inherit from this mixin and implement a `diffuse()` method that orchestrates the denoising loop. + +| Method | Purpose | Automatic Behavior | +|--------|---------|-------------------| +| [`predict_noise_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with CFG | Detects parallel mode, distributes computation, gathers results | +| [`scheduler_step_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Step scheduler with sync | Rank 0 steps, broadcasts latents to all ranks | +| [`combine_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine positive/negative | Applies CFG formula with optional normalization | +| [`predict_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Forward pass wrapper | Override for custom transformer calls | +| [`cfg_normalize_function()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Normalize CFG output | Override for custom normalization | + +### How It Works + +`predict_noise_maybe_with_cfg()` automatically detects and switches between two execution modes: + +- **CFG-Parallel mode** (when `cfg_world_size > 1`): + - Rank 0 computes positive prompt prediction + - Rank 1 computes negative prompt prediction + - Results are gathered via `all_gather()` + - Combined on rank 0 using CFG formula + +- **Sequential mode** (when `cfg_world_size == 1`): + - Single rank computes both positive and negative predictions + - Directly combines them with CFG formula + +`scheduler_step_maybe_with_cfg()` ensures consistent latent states across all ranks: + +- **CFG-Parallel mode**: + - Only rank 0 performs the scheduler step (applies noise prediction to update latents) + - Updated latents are broadcast to all other ranks via `broadcast()` + - All ranks maintain synchronized latent states for the next iteration + +- **Sequential mode**: + - Single rank directly performs the scheduler step + - No synchronization needed + +--- + +## Step-by-Step Implementation + +### Step 1: Inherit `CFGParallelMixin` + +Allow your pipeline to inherit from `CFGParallelMixin` and implements the `diffuse()` method for your specific model. + +**Example (Qwen-Image):** + +```python +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +import torch.nn as nn +class YourModelPipeline(nn.Module, CFGParallelMixin): + def diffuse(self, ...) -> torch.Tensor: + for i, t in enumerate(timesteps): + # Prepare positive_kwargs (conditional) and negative_kwargs (unconditional) + positive_kwargs = {...} # hidden_states, encoder_hidden_states, etc. + negative_kwargs = {...} if do_true_cfg else None + + # Key method 1: Predict noise with automatic CFG parallel handling + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=true_cfg_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + ) + + # Key method 2: Step scheduler with automatic CFG synchronization + latents = self.scheduler_step_maybe_with_cfg( + noise_pred, t, latents, do_true_cfg + ) + + return latents +``` + +**Key Points:** + +- `positive_kwargs`: transformer arguments for conditional (text-guided) prediction +- `negative_kwargs`: transformer arguments for unconditional prediction (set to `None` if CFG disabled) +- For image editing pipelines, add `output_slice=image_seq_len` to extract the generative image portion + +### Step 2: Call `diffuse` + +Call `self.diffuse` in your pipeline's forward function: + +```python +import torch.nn as nn +class YourModelPipeline(nn.Module, CFGParallelMixin): + def forward( + self, + prompt: str, + negative_prompt: str | None = None, + guidance_scale: float = 3.5, + num_inference_steps: int = 50, + **kwargs, + ): + # Encode prompts, Initialize latents, Get timesteps + ... + # Run diffusion loop (calls the mixin's diffuse method) + latents = self.diffuse( + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + negative_prompt_embeds=negative_embeds, + negative_prompt_embeds_mask=negative_mask, + latents=latents, + timesteps=timesteps, + do_true_cfg=do_true_cfg, + true_cfg_scale=guidance_scale, + ... + ) +``` + +--- + +## Customization + +### Override `predict_noise()` for Custom Transformer Calls + +If your transformer requires custom prediction function, you can rewrite `predict_noise` function. Taking Wan2.2 as an example, which has two transformer models. The actual transformer to be called is determined by `self.transformer`. + +```python +class Wan22Pipeline(nn.Module, CFGParallelMixin): + def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor: + if current_model is None: + current_model = self.transformer + return current_model(**kwargs)[0] +``` + + +### Override `cfg_normalize_function()` for Custom Normalization + +Some models have their own normalization function. Taking LongCat Image model as an example: + +```python +class LongCatImagePipeline(nn.Module, CFGParallelMixin): + def cfg_normalize_function(self, noise_pred, comb_pred, cfg_renorm_min=0.0): + """ + Normalize the combined noise prediction. + """ + cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True) + noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True) + scale = (cond_norm / (noise_norm + 1e-8)).clamp(min=cfg_renorm_min, max=1.0) + noise_pred = comb_pred * scale + return noise_pred + + # The original cfg_normalize_function function in CFGParallelMixin + # cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True) + # noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True) + # noise_pred = comb_pred * (cond_norm / noise_norm) + # return noise_pred +``` + +--- + +## Testing + +After adding CFG-Parallel support, test with: + +```bash +cd examples/offline_inference/text_to_image +python text_to_image.py \ + --model Your-org/your-model \ + --prompt "a cup of coffee on the table" \ + --negative_prompt "ugly, unclear" \ + --cfg_scale 4.0 \ + --num_inference_steps 50 \ + --output "cfg_enabled.png" \ + --cfg_parallel_size 2 +``` + +**Verify:** + +1. Check logs for CFG parallel being activated +2. Record the `e2e_time_ms` in the log and compare with CFG-Parallel disabled +3. Compare the generated result quality with baseline +4. Record comparison results in your PR + +--- + +## Troubleshooting + +### Issue: CFG parallel not activating + +**Symptoms:** Generation still slow, logs don't show CFG parallel being used. + +**Causes & Solutions:** + +- **CFG is not enabled:** + +**Problem:** Guidance scale too low or negative prompt not provided. + +**Solution:** Ensure `guidance_scale > 1.0` and negative prompt is provided: +```python +images = pipeline( + prompt="a cat", + negative_prompt="", # Must provide (even if empty) + guidance_scale=3.5, # Must be > 1.0 +) +``` + +--- + +## Reference Implementations + +Complete examples in the codebase: + +| Model | Path | Pattern | Notes | +|-------|------|---------|-------| +| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/cfg_parallel.py` | Mixin | Dual-stream transformer | +| **Qwen-Image-Edit** | `vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py` | Mixin | Image editing with `output_slice` | +| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py` | Mixin | Dual-transformer architecture | +| **CFGParallelMixin** | `vllm_omni/diffusion/distributed/cfg_parallel.py` | Base implementation | Core mixin class | + +--- + +## Summary + +Adding CFG-Parallel support: + +1. ✅ **Create mixin** - Inherit from `CFGParallelMixin` and implement `diffuse()` method +2. ✅ **(Optional) Customize** - Override `predict_noise()` or `cfg_normalize_function()` for custom behavior +3. ✅ **Test** - Verify with `--cfg_parallel_size 2` and compare performance diff --git a/docs/contributing/features/sequence_parallel.md b/docs/contributing/features/sequence_parallel.md new file mode 100644 index 000000000..3bcee7345 --- /dev/null +++ b/docs/contributing/features/sequence_parallel.md @@ -0,0 +1,405 @@ +# Support Sequence Parallel + +This section describes how to add Sequence Parallel (SP) to a diffusion transformer model. We use the Qwen-Image transformer and Wan2.2 transformer as reference implementations. + +--- + +## Table of Contents + +- [Overview](#overview) +- [Approach 1: Non-Intrusive `_sp_plan` (Recommended)](#approach-1-non-intrusive-_sp_plan-recommended) +- [Approach 2: Intrusive Modification (For Complex Cases)](#approach-2-intrusive-modification-for-complex-cases) +- [Testing](#testing) +- [Troubleshooting](#troubleshooting) +- [Reference Implementations](#reference-implementations) +- [Summary](#summary) + +--- + +## Overview + + +### What is Sequence Parallel? + +**Terminology Note:** Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in the [diffusers library](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py). We use "Sequence Parallelism" to align with vLLM-Omni's terminology. + +Diffusion transformers process long sequences of image patches or video frames. For high-resolution generation, these sequences can become very large. Enabling SP allows each GPU to process only a portion of the sequence, with attention mechanisms (Ulysses/Ring) handling cross-GPU communication transparently. + +### Architecture + +The major APIs for Sequence Parallel: + +```python +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, # For sharding (splitting) tensors + SequenceParallelOutput, # For gathering tensors +) +from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather +``` + +| Method/Class | Purpose | Behavior | +|--------------|---------|----------| +| `SequenceParallelInput` | Declare input sharding in `_sp_plan` | Auto-shards tensors at module input | +| `SequenceParallelOutput` | Declare output gathering in `_sp_plan` | Auto-gathers tensors at module output | +| `sp_shard()` | Manual tensor sharding | Splits tensor across SP workers | +| `sp_gather()` | Manual tensor gathering | Gathers sharded tensors from all workers | + +--- + +## Approach 1: Non-Intrusive `_sp_plan` (Recommended) + +The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries. + +**When to use:** +- Standard transformer architectures +- Tensor operations happen at `nn.Module` boundaries +- Predictable sharding/gathering patterns + +**How it works:** +1. Declare `_sp_plan` dict in your transformer class +2. Framework automatically applies hooks when `sequence_parallel_size > 1` +3. Hooks shard/gather tensors at specified module boundaries +4. Attention layers handle cross-GPU communication internally + +```python +class StandardTransformer(nn.Module): + _sp_plan = { + # Shard hidden_states at first transformer block input + "blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + }, + # Gather at final output projection + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } +``` + +`StandardTransformer` has a transformer blocks list `self.blocks = nn.ModuleList([...])`, and a projection output layer `self.proj_out`. The `_sp_plan` above defines that when SP is enabled, sharding the input tensor to the first transformer block, and gathering the sharded tensor at the final output projection layer. + +**Requirements:** +- Tensor operations that need sharding/gathering must happen at **`nn.Module` boundaries** +- Inline Python operations (e.g., `torch.cat`, `pad_sequence`) **cannot be hooked** + +**Solution for inline operations:** Extract into a submodule (see Step 2 below). + +### Step 1: Understand Module Boundaries + +Identify where tensors need to be sharded or gathered in your model's forward pass: + +```python +class MyTransformer(nn.Module): + def __init__(self): + self.patch_embed = PatchEmbed() # ← Boundary 1 + self.pos_embed = RoPE() # ← Boundary 2 + self.blocks = nn.ModuleList([...]) # ← Boundary 3 + self.norm_out = LayerNorm() + self.proj_out = Linear() # ← Boundary 4 + + def forward(self, x): + x = self.patch_embed(x) # ← Shard before this? + pos = self.pos_embed(x) # ← Shard RoPE outputs? + for block in self.blocks: + x = block(x, pos) # ← Blocks process sharded x + x = self.norm_out(x) + output = self.proj_out(x) # ← Gather after this? + return output +``` + +### Step 2: Handle Inline Operations + +If your `forward()` contains inline tensor operations, **extract them into submodules**. + +**Example: Z-Image concatenates image + text features inline** + +```python +# ❌ BAD: Inline operation - hooks cannot intercept +class ZImageTransformer(nn.Module): + def forward(self, x, cap_feats): + # This concatenation happens inline - _sp_plan can't shard it! + unified = torch.cat([x, cap_feats], dim=1) + + for layer in self.layers: + unified = layer(unified) + + return unified + +# ✅ GOOD: Extract into submodule +class UnifiedPrepare(nn.Module): + """Submodule to concatenate image and text features.""" + def forward(self, x, cap_feats): + return torch.cat([x, cap_feats], dim=1) + +class ZImageTransformer(nn.Module): + def __init__(self): + super().__init__() + self.unified_prepare = UnifiedPrepare() # Now a module! + self.layers = nn.ModuleList([...]) + + def forward(self, x, cap_feats): + # Now _sp_plan can shard the output of unified_prepare! + unified = self.unified_prepare(x, cap_feats) + + for layer in self.layers: + unified = layer(unified) + + return unified +``` + +**Other common cases:** +- `pad_sequence()` → `PadSequenceModule` +- `torch.cat()` → `ConcatModule` +- `tensor.reshape()` → `ReshapeModule` +- Complex preprocessing → `PreprocessModule` + +### Step 3: Write `_sp_plan` for Your Model + +Create a class-level `_sp_plan` dictionary specifying where to shard/gather tensors. + +Typically, there are two patterns for diffusion models: + +**Pattern 1: Shard at first block, gather at output projection** + +Most common pattern for standard transformers: + +```python +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, # For sharding (splitting) tensors + SequenceParallelOutput, # For gathering tensors +) +class StandardTransformer(nn.Module): + _sp_plan = { + # Shard hidden_states at first transformer block input + "blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + }, + # Gather at final output projection + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } +``` + +**Pattern 2: Shard RoPE embeddings separately** + +When RoPE is computed in a separate module: + +```python +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, # For sharding (splitting) tensors + SequenceParallelOutput, # For gathering tensors +) +class TransformerWithRoPE(nn.Module): + _sp_plan = { + # Shard RoPE module OUTPUTS (returns tuple of cos, sin) + "rope": { + 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # cos + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # sin + }, + # Shard transformer block INPUT + "blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + }, + # Gather at output + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } +``` + +### API Reference + +**SequenceParallelInput Parameters:** + +| Parameter | Type | Description | +|-----------|------|-------------| +| `split_dim` | int | Dimension to split (usually `1` for sequence) | +| `expected_dims` | int \| None | Expected tensor rank for validation (optional) | +| `split_output` | bool | `False`: shard **input** params; `True`: shard **output** tensors | +| `auto_pad` | bool | Auto-pad if sequence not divisible by world_size (default: `False`) | + +**SequenceParallelOutput Parameters:** + +| Parameter | Type | Description | +|-----------|------|-------------| +| `gather_dim` | int | Dimension to gather (usually `1` for sequence) | +| `expected_dims` | int \| None | Expected tensor rank for validation (optional) | + +**Module Naming Conventions:** + +| Key | Meaning | Python equivalent | +|-----|---------|-------------------| +| `""` | Root model | `model` | +| `"blocks.0"` | First element of ModuleList | `model.blocks[0]` | +| `"blocks.*"` | All elements of ModuleList | `for b in model.blocks` | +| `"rope"` | Named submodule | `model.rope` | +| `"outputs.main"` | ModuleDict entry | `model.outputs["main"]` | + +**Dictionary Value Types:** + +| Key type | `split_output` | Description | +|----------|----------------|-------------| +| `"param_name"` (str) | `False` | Shard **input parameter** by name | +| `0`, `1`, ... (int) | `True` | Shard **output tuple** by index | + +--- + +## Approach 2: Intrusive Modification (For Complex Cases) + +For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls. + +**When to use:** +- Dynamic/conditional sharding logic +- Complex tensor manipulations that can't be encapsulated +- Temporary workaround during development + +```python +from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather + +def forward(self, hidden_states, ...): + if self.parallel_config.sequence_parallel_size > 1: + hidden_states = sp_shard(hidden_states, dim=1) + + # ... computation ... + + if self.parallel_config.sequence_parallel_size > 1: + output = sp_gather(output, dim=1) + + return output +``` + +--- + +## Testing + +After implementing Sequence Parallel support, thoroughly test your implementation to ensure correctness and performance across different configurations. + +**Test Different `sp_size`:** + +Test your model with various sequence parallel world sizes to verify correctness and identify optimal configurations: + +```bash +cd examples/offline_inference/text_to_image +python text_to_image.py \ + --model Your-org/your-model \ + --prompt "a cup of coffee on the table" \ + --num_inference_steps 50 \ + --ulysses_degree 2 \ + --ring_degree 2 \ + --output sp_test_image_ulysses=2_ring=2.png +``` + +**Verify:** + +1. **Correctness:** Output should be identical across all `sp_size` values +2. **Speed:** Throughput should remain stable or improve (especially for large sequences) +3. **Logs:** Check for any shape mismatch or communication errors + +**Test with Tensor Parallel:** + +Sequence Parallel can be combined with other parallelism strategies: + +```bash +cd examples/offline_inference/text_to_image +python text_to_image.py \ + --model Your-org/your-model \ + --prompt "a cup of coffee on the table" \ + --num_inference_steps 50 \ + --ulysses_degree 2 \ + --tensor_parallel_size 2 \ + --output sp_test_image_ulysses=2_tp=2.png +``` + +--- + +## Troubleshooting + +### Issue: Shape mismatch errors + +**Symptoms:** `RuntimeError: shape mismatch` during forward pass. + +**Causes & Solutions:** + +- **RoPE dimension mismatch:** + +**Problem:** RoPE embeddings not sharded, but hidden_states is sharded. + +**Solution:** Shard RoPE outputs in `_sp_plan`: +```python +_sp_plan = { + "rope": { + 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + }, + ... +} +``` + +- **Sequence Length not divisible by sp_size:** + +**Problem:** `SequenceParallelInput(auto_pad=False)` - auto_pad should be True to enable automatic sequence padding. + +**Solution:** In `SequenceParallelInput`, set `auto_pad=True`: +```python +"blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True) +} +``` + +### Issue: Inline operations not sharded + +**Symptoms:** Some tensors remain full-sized, not sharded. + +**Causes & Solutions:** + +- **Operations happen inline in `forward()`, not at module boundaries:** + +**Problem:** +```python +def forward(self, x, cap): + unified = torch.cat([x, cap], dim=1) # ← Inline operation! + # _sp_plan can't hook this +``` + +**Solution:** Extract into submodule: +```python +class ConcatModule(nn.Module): + def forward(self, x, cap): + return torch.cat([x, cap], dim=1) + +class MyModel(nn.Module): + _sp_plan = { + "concat": { + 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + }, + ... + } + def __init__(self): + self.concat = ConcatModule() # Now hookable! + + def forward(self, x, cap): + unified = self.concat(x, cap) # ← Can be sharded via _sp_plan +``` + +--- + +## Reference Implementations + +Complete examples in the codebase: + +| Model | Path | Pattern | Notes | +|-------|------|---------|-------| +| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Dual-stream + preprocessing | auto_pad, separate RoPE | +| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Dual-Transformer + RoPE | Video transformer | +| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Unified sequence | Concatenated input | +| **SP Plan Types** | `vllm_omni/diffusion/distributed/sp_plan.py` | Type definitions | SequenceParallelInput/Output | +| **Hook Implementation** | `vllm_omni/diffusion/hooks/sequence_parallel.py` | Hook mechanics | How hooks work | +| **Tests** | `tests/diffusion/distributed/test_sp_plan_hooks.py` | Test examples | Validation patterns | + +--- + +## Summary + +Adding Sequence Parallel support to a transformer: + +1. ✅ **Choose approach** - Use `_sp_plan` for standard cases, intrusive modification for complex cases +2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered? +3. ✅ **Extract inline operations** - Move `torch.cat`, `pad_sequence`, etc. to submodules +4. ✅ **Define `_sp_plan`** - Declare shard/gather points as class attribute +5. ✅ **Use `auto_pad` for variable lengths** - Support non-uniform sequences +6. ✅ **Test** - Verify with different `ulysses_degree` and `ring_degree` combinations diff --git a/docs/contributing/features/teacache.md b/docs/contributing/features/teacache.md new file mode 100644 index 000000000..6c29c0f99 --- /dev/null +++ b/docs/contributing/features/teacache.md @@ -0,0 +1,459 @@ +# Support TeaCache + +This section describes how to add TeaCache to a diffusion transformer model. We use the Qwen-Image transformer as the reference implementation. + +--- + +## Table of Contents + +- [Overview](#overview) +- [Step-by-Step Implementation](#step-by-step-implementation) +- [Customization](#customization) +- [Testing](#testing) +- [Troubleshooting](#troubleshooting) +- [Reference Implementations](#reference-implementations) +- [Summary](#summary) + +--- + +## Overview + +### What is TeaCache? + +TeaCache speeds up diffusion inference by caching transformer block computations when consecutive timesteps are similar. It provides **1.5x-2.0x speedup** with minimal quality loss. + +The core insight is that the modulated input (after normalization and timestep conditioning) changes gradually across timesteps. By measuring the L1 distance between consecutive modulated inputs and comparing it to a threshold, TeaCache decides whether to execute the full transformer blocks or reuse the cached residual from the previous step. + +vLLM-omni provides a **hook-based** TeaCache system that requires **zero changes to model code**. The hook completely intercepts the transformer's forward pass and implements adaptive caching transparently. This design allows easy integration with any transformer model by simply writing an extractor function. + +### Architecture + +The TeaCache system consists of three main components: + +| Component | Purpose | Location | +|-----------|---------|----------| +| [`CacheContext`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/#vllm_omni.diffusion.cache.CacheContext) | Dataclass containing model-specific information for caching | `vllm_omni/diffusion/cache/teacache/context.py` | +| [`EXTRACTOR_REGISTRY`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/teacache/extractors/#vllm_omni.diffusion.cache.teacache.extractors.EXTRACTOR_REGISTRY) | Maps transformer class names to extractor functions | `vllm_omni/diffusion/cache/teacache/extractors.py` | +| [`TeaCacheConfig`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/#vllm_omni.diffusion.cache.TeaCacheConfig) | Configuration including thresholds and polynomial coefficients | `vllm_omni/diffusion/cache/teacache/config.py` | + +The hook handles all caching logic automatically, including: + +- CFG-aware state management (separate states for positive/negative branches) +- CFG-parallel compatibility +- L1 distance computation with polynomial rescaling +- Residual caching and reuse + + +--- + +## Step-by-Step Implementation + +To add TeaCache support for a new model, you need to: + +1. Write an **extractor function** that returns a `CacheContext` object +2. Register the extractor in the `EXTRACTOR_REGISTRY` +3. Add model-specific polynomial coefficients to `TeaCacheConfig` + +### Step 1: Model-Specific Preprocessing + +Extract and process model inputs. This typically involves: +- Embedding image/latent inputs +- Processing text encoder outputs (if dual-stream) +- Creating timestep embeddings +- Computing positional embeddings + +**Example (Qwen-Image):** + +```python +def extract_qwen_context( + module: nn.Module, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_mask: torch.Tensor, + timestep: torch.Tensor, + img_shapes: torch.Tensor, + txt_seq_lens: torch.Tensor, + guidance: torch.Tensor | None = None, + **kwargs: Any, +) -> CacheContext: + # Validate model structure + if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0: + raise ValueError("Module must have transformer_blocks") + + # Preprocessing: embed inputs + hidden_states = module.img_in(hidden_states) + timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype) + encoder_hidden_states = module.txt_norm(encoder_hidden_states) + encoder_hidden_states = module.txt_in(encoder_hidden_states) + + # Create timestep embedding + if guidance is not None: + guidance = guidance.to(hidden_states.dtype) * 1000 + temb = ( + module.time_text_embed(timestep, hidden_states) + if guidance is None + else module.time_text_embed(timestep, guidance, hidden_states) + ) + + # Compute position embeddings + image_rotary_emb = module.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device) +``` + +### Step 2: Extract Modulated Input + +The modulated input is used for cache decisions. Extract it from the **first transformer block** after normalization and modulation. + +**Example (Qwen-Image):** + +```python + # Extract modulated input from first transformer block + block = module.transformer_blocks[0] + img_mod_params = block.img_mod(temb) + img_mod1, _ = img_mod_params.chunk(2, dim=-1) + img_modulated, _ = block.img_norm1(hidden_states, img_mod1) +``` + +**Key Points:** + +- Use the **first block** to extract modulated input early +- Apply the same normalization and modulation as the actual forward pass +- The tensor should represent the processed features that will change across timesteps + +### Step 3: Define Transformer Execution + +Create a callable that executes all transformer blocks. This encapsulates the main computation loop. + +**Example (Qwen-Image dual-stream):** + +```python + def run_transformer_blocks(): + """Execute all Qwen transformer blocks.""" + h = hidden_states + e = encoder_hidden_states + + for block in module.transformer_blocks: + e, h = block( + hidden_states=h, + encoder_hidden_states=e, + encoder_hidden_states_mask=encoder_hidden_states_mask, + temb=temb, + image_rotary_emb=image_rotary_emb, + ) + return (h, e) # Return both image and text hidden states +``` + +**Example (Single-stream model like Flux):** + +```python + def run_transformer_blocks(): + """Execute all Flux transformer blocks.""" + h = hidden_states + + for block in module.transformer_blocks: + h = block(h, temb=temb) + return (h,) # Return only image hidden states +``` + +**Key Points:** + +- Return format: +- For single-stream models: return `(hidden_states,)` +- For dual-stream models: return `(hidden_states, encoder_hidden_states)` + +### Step 4: Define Postprocessing + +Create a callable that applies final transformations to produce the model output. + +**Example (Qwen-Image):** + +```python + return_dict = kwargs.get("return_dict", True) + + def postprocess(h): + """Apply Qwen-specific output postprocessing.""" + h = module.norm_out(h, temb) + output = module.proj_out(h) + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) +``` + +### Step 5: Return CacheContext + +Package all information into a `CacheContext` object. + +```python + return CacheContext( + modulated_input=img_modulated, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, # or None for single-stream + temb=temb, + run_transformer_blocks=run_transformer_blocks, + postprocess=postprocess, + ) +``` + +**CacheContext Fields:** + +| Field | Type | Purpose | +|-------|------|---------| +| `modulated_input` | `torch.Tensor` | Tensor used for cache decision (similarity comparison) | +| `hidden_states` | `torch.Tensor` | Current hidden states (will be modified by caching) | +| `encoder_hidden_states` | `torch.Tensor | None` | Encoder states for dual-stream models, `None` for single-stream | +| `temb` | `torch.Tensor` | Timestep embedding tensor | +| `run_transformer_blocks` | `Callable[[], tuple]` | Executes transformer blocks, returns `(hidden_states, [encoder_hidden_states])` | +| `postprocess` | `Callable[[torch.Tensor], Any]` | Applies final transformations to produce model output | +| `extra_states` | `dict | None` | Optional dict for additional model-specific state | + +### Step 6: Register the Extractor + +Add your extractor to the `EXTRACTOR_REGISTRY` in `vllm_omni/diffusion/cache/teacache/extractors.py`: + +```python +EXTRACTOR_REGISTRY: dict[str, Callable] = { + "QwenImageTransformer2DModel": extract_qwen_context, + "Bagel": extract_bagel_context, + "ZImageTransformer2DModel": extract_zimage_context, + "YourModelTransformer2DModel": extract_your_model_context, # Add here +} +``` + +**Key:** Use the transformer class name (`module.__class__.__name__`) + +### Step 7: Add Model Coefficients + +Add polynomial rescaling coefficients to `vllm_omni/diffusion/cache/teacache/config.py`: + +```python +_MODEL_COEFFICIENTS = { + "QwenImageTransformer2DModel": [ + -4.50000000e02, + 2.80000000e02, + -4.50000000e01, + 3.20000000e00, + -2.00000000e-02, + ], + "YourModelTransformer2DModel": [ # Add your model's coefficients + # 5 polynomial coefficients (can reuse similar model's coefficients initially) + ], +} +``` + + +**Initial approach:** Start with coefficients from a similar model architecture, then tune empirically following [Customization](#customization) section. + +--- + +## Customization + +### Coefficient Estimation + +While you can start with coefficients from a similar model architecture, estimating custom coefficients for your specific model typically improves TeaCache performance. + +**Why Estimate Coefficients?** + +The polynomial coefficients rescale L1 distances between consecutive modulated inputs to better predict when cached residuals can be reused. Model-specific coefficients account for: + +- Architecture differences (layer count, hidden size, attention patterns) +- Training data characteristics +- Noise prediction behavior across timesteps + +| Approach | Performance | Effort | +|----------|-------------|--------| +| Using defaults from similar model | Within 5-10% of optimal | Low | +| Estimating custom coefficients | Best performance | Medium | + +#### Implement Data Collection Adapter + +Add an adapter in `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py`: + +```python +class YourModelAdapter: + """Adapter for coefficient estimation on your model.""" + + @staticmethod + def load_pipeline(model_path: str, device: str, dtype: torch.dtype) -> Any: + """Load your diffusion pipeline.""" + from your_model_package import YourModelPipeline + + pipeline = YourModelPipeline.from_pretrained( + model_path, + torch_dtype=dtype, + ) + pipeline = pipeline.to(device) + return pipeline + + @staticmethod + def get_transformer(pipeline: Any) -> tuple[Any, str]: + """Extract transformer from pipeline.""" + return pipeline.transformer, "YourTransformer2DModel" + + @staticmethod + def install_hook(transformer: Any, hook: DataCollectionHook) -> None: + """Install data collection hook on transformer.""" + from vllm_omni.diffusion.hooks import HookRegistry + + registry = HookRegistry.get_or_create(transformer) + registry.register_hook(hook._HOOK_NAME, hook) + + +# Register your adapter +_MODEL_ADAPTERS["YourModel"] = YourModelAdapter +``` + +#### Collect Data and Estimate + +```python +from vllm_omni.diffusion.cache.teacache.coefficient_estimator import ( + TeaCacheCoefficientEstimator, +) +from datasets import load_dataset +from tqdm import tqdm + +# Initialize estimator +estimator = TeaCacheCoefficientEstimator( + model_path="/path/to/your/model", + model_type="YourModel", +) + +# Load diverse prompts (paper recommends ~70 prompts) +dataset = load_dataset("nateraw/parti-prompts", split="train") +prompts = dataset["Prompt"][:70] + +# Collect data +for prompt in tqdm(prompts, desc="Collecting data"): + estimator.collect_from_prompt(prompt=prompt, num_inference_steps=50) + +# Estimate coefficients +coeffs = estimator.estimate(poly_order=4) +print(f"Estimated coefficients: {coeffs.tolist()}") +``` + +**Data Statistics Guide:** + +| Metric | Good Range | Warning Signs | +|--------|------------|---------------| +| **Count** | 2000-5000+ | < 1000: too few prompts | +| **Input Diffs (x)** | 0.01-0.10 | Very small (<0.001): model may not modulate properly | +| **Output Diffs (y)** | Should correlate with x | No correlation: check extractor | +| **Coefficient magnitude** | -1e6 to 1e6 | > 1e8: numerical instability | + +--- + +## Testing + +After adding TeaCache support, test with: + +```python +from vllm_omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +omni = Omni( + model="your-model-name", + cache_backend="tea_cache", + cache_config={ + "rel_l1_thresh": 0.2, + "coefficients": [1.33e6, -1.69e5, 7.95e3, -1.64e2, 1.26], # Your coefficients + } +) + +images = omni.generate( + "a beautiful landscape", + OmniDiffusionSamplingParams(num_inference_steps=50), +) +``` + +**Verify:** + +1. **Check logs** - Look for TeaCache initialization messages +2. **Compare performance** - Measure speedup vs baseline (expect 1.5x-2.0x) +3. **Verify output quality** - Visually compare cached vs uncached outputs (should be nearly identical) + +See more detailed examples in [user guide for teacache](../../user_guide/diffusion/teacache.md). + +--- + +## Troubleshooting + +### Issue: "Unknown model type" + +**Symptoms:** Error message indicating the model type is not recognized when enabling TeaCache. + +**Causes & Solutions:** + +- **Extractor not registered:** + +**Problem:** The transformer class name doesn't exist in `EXTRACTOR_REGISTRY`. + +**Solution:** Check the class name and add to registry: +```python +# Check transformer class name +print(pipeline.transformer.__class__.__name__) + +# Add to EXTRACTOR_REGISTRY +EXTRACTOR_REGISTRY["YourTransformer2DModel"] = extract_your_context +``` + +- **Transformer class name mismatch:** + +**Solution:** Ensure the registry key matches exactly with `module.__class__.__name__`. + +### Issue: "Cannot find coefficients" + +**Symptoms:** Error when initializing TeaCache about missing model coefficients. + +**Causes & Solutions:** + +- **Missing coefficients in config:** + +**Solution:** Add coefficients to `_MODEL_COEFFICIENTS` in `config.py`, or pass custom coefficients: +```python +omni = Omni( + model="your-model", + cache_backend="tea_cache", + cache_config={"coefficients": [1.0, -0.5, 0.1, -0.01, 0.001]} +) +``` + +### Issue: Quality Degradation + +**Symptoms:** Output images look noticeably different or have artifacts compared to baseline. + +**Causes & Solutions:** + +- **Threshold too high:** + +**Problem:** `rel_l1_thresh` is too aggressive, causing cache reuse when outputs differ significantly. + +**Solution:** Lower the threshold: +```python +cache_config={"rel_l1_thresh": 0.1} # Try 0.1-0.2 +``` + +- **Coefficients not tuned:** + +**Solution:** Estimate model-specific coefficients using the coefficient estimation process described above. + +--- + +## Reference Implementations + +Complete examples in the codebase: + +| Model | Path | Pattern | Notes | +|-------|------|---------|-------| +| **Qwen-Image** | `vllm_omni/diffusion/cache/teacache/extractors.py` | Dual-stream | `extract_qwen_context` | +| **Bagel** | `vllm_omni/diffusion/cache/teacache/extractors.py` | Omni model | `extract_bagel_context` | +| **TeaCache Core** | `vllm_omni/diffusion/cache/teacache/` | Base implementation | Hook and config | +| **Coefficient Estimator** | `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py` | Estimation tool | Adapter pattern | + +--- + +## Summary + +Adding TeaCache support: + +1. ✅ **Write extractor** - Create function returning `CacheContext` with model-specific preprocessing +2. ✅ **Register extractor** - Add to `EXTRACTOR_REGISTRY` with transformer class name +3. ✅ **Add coefficients** - Add polynomial coefficients to `_MODEL_COEFFICIENTS` +4. ✅ **Test** - Verify with `cache_backend="tea_cache"` diff --git a/docs/contributing/features/tensor_parallel.md b/docs/contributing/features/tensor_parallel.md new file mode 100644 index 000000000..aecc1aedc --- /dev/null +++ b/docs/contributing/features/tensor_parallel.md @@ -0,0 +1,279 @@ +# Support Tensor Parallel + +This section describes how to add Tensor Parallel (TP) to a diffusion transformer model. We use the Z-Image transformer as the reference implementation. + +--- + +## Table of Contents + +- [Overview](#overview) +- [Step-by-Step Implementation](#step-by-step-implementation) +- [Testing](#testing) +- [Troubleshooting](#troubleshooting) +- [Reference Implementations](#reference-implementations) +- [Summary](#summary) + +--- + +## Overview + +### What is Tensor Parallel? + +Tensor Parallel (TP) is a model parallelism technique that **shards model weights** across multiple GPUs. Each GPU holds only a portion of the model's parameters and computes only part of each layer's output. + +Diffusion transformers contain large attention and MLP layers. We can use Tensor Parallel to shard the model dimension across multiple GPUs, allowing larger models to fit in memory while achieving near-linear speedup. + +### Architecture + +The Tensor Parallel implementation relies vLLM's Parallel Layers: + +[vLLM Parallel Layers API Reference](https://docs.vllm.ai/en/latest/contributing/model/basic/?h=column#3-optional-implement-tensor-parallelism-and-quantization-support) + +**Parallel Layer Types:** + +| Layer Type | Purpose | Weight Partitioning | +|------------|---------|---------------------| +| `ColumnParallelLinear` | First FFN layer, separated QKV | Columns (output dimension) | +| `RowParallelLinear` | Second FFN layer, attention output | Rows (input dimension) | +| `QKVParallelLinear` | Multi-head/grouped-query attention QKV | Handles head replication automatically | +| `ReplicatedLinear` | Layers that shouldn't be sharded | No partitioning (replicated) | + +--- + +## Step-by-Step Implementation + + +### Step 1: Identify Linear Layers + +Find all `nn.Linear` layers in your transformer that need to be sharded. + +**Key questions:** +- Which layers should be column parallel (weight split by columns)? +- Which layers should be row parallel (weight split by rows)? + +### Step 2: Replace Linear Layers with Parallel Equivalents + +Replace `nn.Linear` with parallel layers from `vllm.model_executor.layers.linear`. + +**Example (MLP Block - Up-Down Pattern):** + +```python +class FeedForward(nn.Module): + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + # Column parallel: weight split by columns [hidden_dim/N, dim] + self.w1 = ColumnParallelLinear( + dim, + hidden_dim, + bias=False, + return_bias=False, + ) + self.act = nn.GELU() + + self.w2 = RowParallelLinear( + hidden_dim, + dim, + bias=False, + input_is_parallel=True, # Input already sharded from w1 + return_bias=False, + ) + + def forward(self, x): + # x: [batch, seq, dim] (replicated on all GPUs) + # w1 outputs sharded [batch, seq, hidden_dim/N] + x = self.w1(x) + # act operates on sharded tensors (no communication) + x = self.act(x) + # w2 outputs full dim [batch, seq, dim] via all-reduce + x = self.w2(x) + return x +``` + +**Example (Attention - QKV-Out Pattern):** + +```python +from vllm_omni.diffusion.attention.layer import Attention +class YourModelAttention(nn.Module): + def __init__(self, dim: int, num_heads: int, num_kv_heads: int): + super().__init__() + self.head_dim = dim // num_heads + + # Column parallel: QKV weight split by columns + # Each GPU gets num_heads/N heads + self.to_qkv = QKVParallelLinear( + hidden_size=dim, + head_size=self.head_dim, + total_num_heads=num_heads, + total_num_kv_heads=num_kv_heads, + bias=False, + return_bias=False, + ) + + # Row parallel: output weight split by rows + self.to_out = RowParallelLinear( + dim, + dim, + bias=False, + input_is_parallel=True, # Input sharded from attention + return_bias=False, + ) + + self.attn = Attention( + num_heads=self.to_qkv.num_heads, # Each GPU gets num_heads/N heads + head_size=self.head_dim, + softmax_scale=1.0 / (self.head_dim**0.5), + causal=False, + num_kv_heads=self.to_qkv.num_kv_heads, + ) + + def forward(self, x): + # x: [batch, seq, dim] (replicated) + # to_qkv outputs sharded [batch, seq, (q+k+v) * head_dim/N] + qkv = self.to_qkv(x) + # Split into Q, K, V (each sharded on heads) + q, k, v = qkv.split([...], dim=-1) + # Attention computed independently on each GPU + out = self.attn(q, k, v) + # to_out all-reduces to full dim + out = self.to_out(out) + return out +``` + +**Key Points:** + +- `ColumnParallelLinear` → `RowParallelLinear` is the standard pairing +- Set `input_is_parallel=True` on `RowParallelLinear` when input comes from `ColumnParallelLinear` +- Use `QKVParallelLinear` for attention projections (handles head replication automatically) + +### Step 3: Validate TP Constraints + +For correct TP operation, these dimensions **must be divisible** by `tensor_parallel_size`: + +| Dimension | Reason | Example Error | +|-----------|--------|---------------| +| `num_heads` | Heads sharded by QKVParallelLinear | `num_heads=30, tp=4` ❌ (30 % 4 ≠ 0) | +| `num_kv_heads` | KV heads sharded by QKVParallelLinear | `num_kv_heads=30, tp=4` ❌ (30 % 4 ≠ 0) | + +--- + +## Testing + +After adding Tensor Parallel support, test with: + +```python +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +parallel_config = DiffusionParallelConfig(tensor_parallel_size=2) +omni = Omni(model="your-model-name", parallel_config=parallel_config) + +output = omni.generate( + "a cup of coffee on the table", + OmniDiffusionSamplingParams(num_inference_steps=50), +) +``` + +**Or via command line:** + +```bash +cd examples/offline_inference/text_to_image +python text_to_image.py \ + --model Your-org/your-model \ + --prompt "a cup of coffee on the table" \ + --negative_prompt "ugly, unclear" \ + --cfg_scale 4.0 \ + --num_inference_steps 50 \ + --output "tp_enabled.png" \ + --tensor_parallel_size 2 +``` + +**Verify:** + +1. Check the `e2e_time_ms` in the log for speedup +2. Compare generated image quality with TP disabled +3. Verify memory usage is reduced proportionally +4. Record comparison results in your PR + +--- + +## Troubleshooting + +### Issue: TP not activating + +**Symptoms:** Model runs on single GPU, no memory savings or speedup. + +**Causes & Solutions:** + +- **Still using `nn.Linear`:** + +**Problem:** Linear layers not replaced with parallel equivalents. + +**Solution:** Replace with parallel layers: +```python +# ❌ BAD +self.proj = nn.Linear(dim, dim) + +# ✅ GOOD +self.proj = RowParallelLinear(dim, dim, input_is_parallel=True) +``` + +### Issue: Dimension mismatch errors + +**Symptoms:** `RuntimeError: shape mismatch` during forward pass. + +**Causes & Solutions:** + +- **Missing `input_is_parallel=True`:** + +**Problem:** RowParallelLinear expects sharded input but receives full tensor. + +**Solution:** Set `input_is_parallel=True` when input comes from ColumnParallelLinear: +```python +# ✅ GOOD: Correct pairing +self.w1 = ColumnParallelLinear(dim, hidden_dim, return_bias=False,) +self.w2 = RowParallelLinear( + hidden_dim, + dim, + input_is_parallel=True, # Input sharded from w1 + return_bias=False, +) +``` + +- **Incorrect split dimensions:** + +**Problem:** QKV split sizes don't match sharded dimensions. + +**Solution:** Use `self.to_qkv.num_heads` (local heads per GPU): +```python +# ❌ BAD: Uses total heads +q_size = self.total_num_heads * self.head_dim + +# ✅ GOOD: Uses local heads +q_size = self.to_qkv.num_heads * self.head_dim +``` + +--- + +## Reference Implementations + +Complete examples in the codebase: + +| Model | Path | Pattern | Notes | +|-------|------|---------|-------| +| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Standard TP | Full implementation with validation | +| **FLUX** | `vllm_omni/diffusion/models/flux/flux_transformer.py` | Dual-stream | Image + text streams | +| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Standard TP | With RoPE | +| **TP Tests** | `tests/e2e/offline_inference/test_zimage_tensor_parallel.py` | E2E testing | TP correctness and performance | +| **Constraint Tests** | `tests/diffusion/models/z_image/test_zimage_tp_constraints.py` | Unit testing | Validation logic | + +--- + +## Summary + +Adding Tensor Parallel support to a transformer: + +1. ✅ **Identify linear layers** - Which layers should be sharded? +2. ✅ **Replace with parallel layers** - Use QKVParallelLinear, ColumnParallelLinear, RowParallelLinear +3. ✅ **Validate TP constraints** - Ensure dimensions divisible by TP size +4. ✅ **Test** - Verify with `tensor_parallel_size=N`, check memory, speed, and quality diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 80212360f..529cab752 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -1,22 +1,27 @@ -# Adding a Diffusion Model -This guide walks through the process of adding a new Diffusion model to vLLM-Omni, using Qwen/Qwen-Image-Edit as a comprehensive example. +# Adding a Diffusion Model to vLLM-Omni + +This guide walks you through adding a new diffusion model to vLLM-Omni. We use **Qwen-Image** as the primary example, with references to other models (LongCat, Flux, Wan2.2) to illustrate different patterns. + + +--- + +## Table of Contents -# Table of Contents 1. [Overview](#overview) 2. [Directory Structure](#directory-structure) -3. [Step-by-Step Implementation](#step-by-step-implementation) -4. [Testing](#testing) -5. [Adding a Model Recipe](#adding-a-model-recipe) +3. [Basic Implementation](#basic-implementation) +4. [Advanced Features](#advanced-features) +5. [Troubleshooting](#troubleshooting) +6. [Pull Request Checklist](#pull-request-checklist) +7. [Reference Implementations](#reference-implementations) +8. [Summary](#summary) +--- -# Overview -When add a new diffusion model into vLLM-Omni, additional adaptation work is required due to the following reasons: +## Overview -+ New model must follow the framework’s parameter passing mechanisms and inference flow. +vLLM-Omni's diffusion inference follows this architecture: -+ Replacing the model’s default implementations with optimized modules, which is necessary to achieve the better performance. - -The diffusion execution flow as follow:

@@ -24,348 +29,892 @@ The diffusion execution flow as follow:

+**Key Components:** -# Directory Structure -File Structure for Adding a New Diffusion Model +1. **Request Handling:** User prompts → `OmniDiffusionRequest` +2. **Diffusion Engine:** Request → Preprocessing (Optional) → Pipeline execution -> Post-processing +3. **Pipeline Execution:** Request → Encode prompt → Diffusion steps → Vae decode + + +## Directory Structure + +Organize your model files following this structure: ``` vllm_omni/ -└── examples/ - └──offline_inference - └── example script # reuse existing if possible (e.g., image_edit.py) - └──online_serving - └── example script └── diffusion/ - └── registry.py # Registry work - ├── request.py # Request Info - └── models/your_model_name/ # Model directory (e.g., qwen_image) - └── pipeline_xxx.py # Model implementation (e.g., pipeline_qwen_image_edit.py) + ├── registry.py # ← Register your model here + ├── request.py # Request data structures + └── models/ + └── your_model_name/ # ← Create this directory + ├── __init__.py # Export pipeline and transformer + ├── pipeline_xxx.py # Pipeline implementation + └── xxx_transformer.py # Transformer implementation ``` -# Step-by-step-implementation -## Step 1: Model Implementation -The diffusion pipeline’s implementation follows **HuggingFace Diffusers**. -### 1.1 Define the Pipeline Class -Define the pipeline class, e.g., `QwenImageEditPipeline`, and initialize all required submodules, either from HuggingFace `diffusers` or custom implementations. In `QwenImageEditPipeline`, only `QwenImageTransformer2DModel` is re-implemented to support optimizations such as Ulysses-SP. When adding new models in the future, you can either reuse this re-implemented `QwenImageTransformer2DModel` or extend it as needed. +**Naming Conventions:** -### 1.2 Pre-Processing and Post-Processing Extraction -Extract the pre-processing and post-processing logic from the pipeline class to follow vLLM-Omni’s execution flow. For Qwen-Image-Edit: +- **Model directory:** `your_model_name` (lowercase, underscores), e.g., `qwen_image`, `flux`, `longcat_image`, `wan2_2` +- **Pipeline file:** `pipeline_xxx.py` where `xxx` describes the task, e.g., `pipeline_qwen_image.py`, `pipeline_qwen_image_edit.py` +- **Transformer file:** `xxx_transformer.py` matching transformer class name, e.g., `qwen_image_transformer.py`, `flux_transformer.py` + +--- + +## Basic Implementation + +This section covers the minimal steps to get a model working in vLLM-Omni with basic features (online/offline serving, batch requests). + +### Step 1: Adapt Transformer Model + +The transformer is the core denoising network. Start by copying the transformer implementation from Diffusers and making these adaptations. + + +#### 1.1: Remove Diffusers Mixins + +Diffusers' `Mixin` classes are not needed in vLLM-Omni. Remove them: + +```diff +# Before (Diffusers) +- from diffusers.models.modeling_utils import ModelMixin +- from diffusers.models.attention_processor import AttentionModuleMixin + +- class YourModelTransformer2DModel(ModelMixin, AttentionModuleMixin): ++ class YourModelTransformer2DModel(nn.Module): + """Your transformer model.""" +``` + +**Example mixins to remove:** + +- `ModelMixin` - Weight loading utilities (vLLM-Omni has its own weight loader) +- `AttentionModuleMixin` - Attention processors (using vLLM-Omni's Attention layer instead) +- `ConfigMixin` - Config management (not needed) +- `PeftAdapterMixin` - Parameter efficient finetune utilities (not needed) + +#### 1.2: Replace Attention Implementation + +**The most important adaptation:** Replace Diffusers' attention with vLLM-Omni's optimized `Attention` layer. + +**Before (Diffusers):** ```python -def get_qwen_image_edit_pre_process_func( - od_config: OmniDiffusionConfig, -): - """ - Define a pre-processing function that resizes input images and - pre-process for subsequent inference. - """ +from diffusers.models.attention_processor import dispatch_attention_fn + +class YourAttentionBlock(nn.Module): + def forward(self, hidden_states, encoder_hidden_states=None, ...): + ... + hidden_states = dispatch_attention_fn( + query, key, value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + ) ``` -```python -def get_qwen_image_edit_post_process_func( - od_config: OmniDiffusionConfig, -): - """ - Defines a post-processing function that post-process images. - """ -``` - -### 1.3 Define the forward function -The forward function of `QwenImageEditPipeline` follows the HuggingFace `diffusers` design for the most part. The key differences are: -+ As described in the overview, arguments are passed through `OnniDiffusionRequest`, so we need to get user parameters from it accordingly. -```python -prompt = req.prompt -``` -+ pre/post-processing are handled by the framework elsewhere, so skip them. - -### 1.4 Replace some ops or layers in DiT component - -vLLM-Omni provides a set of optimized operators with better performance and built-in support for parallelism, including attention, rotary embeddings (RoPE), and linear layers. - -Below is an example showing how to replace standard Transformer attention and FFN layers with vLLM-Omni implementations: - +**After (vLLM-Omni):** ```python from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + +class YourAttentionBlock(nn.Module): + def __init__(self, ...): + super().__init__() + + # Initialize vLLM-Omni's Attention layer + self.attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + softmax_scale=1.0 / (self.head_dim ** 0.5), + causal=False, # Diffusion models typically use bidirectional attention + num_kv_heads=self.num_kv_heads, + ) + + def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, ...): + ... + # Create attention metadata + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + hidden_states = self.attn(query, key, value, attn_metadata=attn_metadata) + +``` + +**Key Points:** + +- **Attention layer initialization:** Done in `__init__`, not per-forward +- **Tensor shapes:** vLLM-Omni `Attention` expects QKV to have `[B, seq, num_heads, head_dim]` shape +- **AttentionMetadata:** Wraps attention mask and other metadata + +**Attention backends:** vLLM-Omni automatically selects the attention backend given the environmental variable `DIFFUSION_ATTENTION_BACKEND`. The default attention backend is `FLASH_ATTN` for diffusion models. + +#### 1.3: Replace Imports and Utilities + +**Logger:** +```diff +- from diffusers.utils import logging +- logger = logging.get_logger(__name__) + ++ from vllm.logger import init_logger ++ logger = init_logger(__name__) +``` + +**Custom layers from vLLM and vLLM-Omni (if needed):** + +```python +from vllm.model_executor.layers.layernorm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm +``` -class MyAttention(nn.Module): - def __init__(self): +#### 1.4: Remove Training-Only Code + +Remove code that's only needed for training: + +```diff +# Remove gradient checkpointing +- if torch.is_grad_enabled() and self.gradient_checkpointing: +- hidden_states = torch.utils.checkpoint.checkpoint( +- self._forward_block, hidden_states, ... +- ) +- else: +- hidden_states = self._forward_block(hidden_states, ...) ++ hidden_states = self._forward_block(hidden_states, ...) + +# Remove training-specific attributes +- self.gradient_checkpointing = False + +# Remove dropout (set to 0 or remove) +- self.dropout = nn.Dropout(dropout_prob) ++ # Removed dropout for inference +``` + +#### 1.5: Add Configuration Support + +Add support for vLLM-Omni's `OmniDiffusionConfig`: + +```python +from vllm_omni.diffusion.data import OmniDiffusionConfig + +class YourModelTransformer2DModel(nn.Module): + def __init__( + self, + *, + od_config: OmniDiffusionConfig | None = None, # ← Add vLLM-Omni config + # ... other model-specific parameters + num_layers: int = 28, + hidden_size: int = 3072, + num_heads: int = 24, + **kwargs, + ): super().__init__() - self.attn = Attention() - self.to_qkv = QKVParallelLinear() - self.to_out = RowParallelLinear() - self.rope = RotaryEmbedding(is_neox_style=False) - def forward(self, hidden_states): - qkv, _ = self.to_qkv(hidden_states) - q, k, v = qkv.split(...) - q, k = self.rope(...) - attn_output = self.attn(q, k, v) - output = self.to_out(attn_output) + # Store config + self.od_config = od_config + self.parallel_config = od_config.parallel_config if od_config else None -class MyFFN(nn.Module): - def __init__(self): + # Model architecture + self.num_layers = num_layers + self.hidden_size = hidden_size + # ... initialize layers +``` + +### Step 2: Adapt Pipeline + +The pipeline orchestrates the full generation process (text encoding, denoising loop, VAE decoding). Adapt it from Diffusers format to vLLM-Omni's interface. + +#### 2.1: Remove Diffusers Inheritance + +**Remove Diffusers base classes:** +```diff +- from diffusers import DiffusionPipeline +- from diffusers.loaders import LoraLoaderMixin + +- class YourModelPipeline(DiffusionPipeline, LoraLoaderMixin): ++ class YourModelPipeline(nn.Module): + """Your model pipeline for vLLM-Omni.""" +``` + +#### 2.2: Adapt `__init__` Method + +**Before (Diffusers):** +```python +class YourModelPipeline(DiffusionPipeline): + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + tokenizer: CLIPTokenizer, + transformer: YourTransformer, + scheduler: FlowMatchScheduler, + ): super().__init__() - self.fc1 = ColumnParallelLinear() - self.fc2 = RowParallelLinear() - self.act = F.gelu - - def forward(self, hidden_states): - hidden, _ = self.fc1(hidden_states) - hidden = self.act(hidden) - output = self.fc2(hidden) - return output + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + ) ``` -In this example: - -+ Attention uses vLLM-Omni’s optimized attention kernel together with parallel QKV projection and RoPE. - -+ Linear layers are replaced with column- and row-parallel variants to enable tensor parallelism. - -+ The FFN follows a standard two-layer structure and can be further optimized (e.g., using fused or merged projections) if needed. - - -### 1.5 Provide a `_repeated_blocks` in DiT model -`_repeated_blocks` is the small and frequently-repeated block(s) of a model -- typically a transformer layer. - -It's used for torch compile optimizations. +**After (vLLM-Omni):** ```python -_repeated_blocks = ["QwenImageTransformerBlock"] +import os +from diffusers import AutoencoderKL +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from transformers import CLIPTextModel, CLIPTokenizer + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.models.your_model_name.your_model_transformer import ( + YourModelTransformer2DModel, +) + + +class YourModelPipeline(nn.Module): + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + self.parallel_config = od_config.parallel_config + self.device = get_local_device() + model = od_config.model + local_files_only = os.path.exists(model) + + # Load components from checkpoint + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, subfolder="scheduler", local_files_only=local_files_only) + self.text_encoder = CLIPTextModel.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only).to(self.device) + self.tokenizer = CLIPTokenizer.from_pretrained( + model, subfolder="tokenizer", local_files_only=local_files_only) + self.vae = AutoencoderKL.from_pretrained( + model, subfolder="vae", local_files_only=local_files_only).to(self.device) + + # Initialize transformer with vLLM-Omni config + transformer_kwargs = get_transformer_config_kwargs( + od_config.tf_model_config, YourModelTransformer2DModel) + self.transformer = YourModelTransformer2DModel( + od_config=od_config, **transformer_kwargs) + + # Store VAE scale factor for latent space conversions + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + self.default_sample_size = 128 # Default latent size ``` +**Key Changes:** -### 1.6 (Optional) implement sequence parallelism -vLLM-Omni has a non-intrusive `_sp_plan` that enable sequence parallel without modifying `forward()` logic. -You can refer to [How to parallelize a new model](../../user_guide/diffusion/parallelism_acceleration.md) +1. **`od_config` parameter:** All configuration through `OmniDiffusionConfig` +2. **Manual component loading:** No `register_modules()`, load each component explicitly +3. **Local files support:** Check `os.path.exists(model)` for local checkpoints +4. **Transformer with config:** Pass `od_config` to transformer constructor +#### 2.3: Adapt `__call__` → `forward` Method -### 1.7 (Optional) integrate with Cache-Dit -vLLM-Omni supports acceleration via [Cache-Dit](../../user_guide/diffusion/cache_dit_acceleration.md). Most models compatible with Diffusers can use Cache-Dit seamlessly. For new models, you can extend support by modifying`cache_dit_backend.py` +**Change signature:** -## Step 2: Extend OmniDiffusionRequest Fields -User-provided inputs are ultimately passed to the model’s forward method through OmniDiffusionRequest, so we add the required fields here to support the new model. -```python -prompt: str | list[str] | None = None -negative_prompt: str | list[str] | None = None -... +```diff +- @torch.no_grad() +- def __call__( ++ def forward( + self, ++ req: OmniDiffusionRequest, # ← Add request parameter here +- ): ++ ) -> DiffusionOutput: # ← Add return type ``` -## Step 3: Registry -+ registry diffusion model in registry.py +[`OmniDiffusionRequest`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/request/#vllm_omni.diffusion.request.OmniDiffusionRequest) is a dataclass that contains the **prompts** and **sampling parameters** [`OmniDiffusionSamplingParams`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/inputs/data/#vllm_omni.inputs.data.OmniDiffusionSamplingParams) for the diffusion pipeline execution. It also contains a request_id for other components to trace this request and its outputs. + +See some parameters in `OmniDiffusionSamplingParams` as follows: + +| parameters | type |value | function | +|:---:|:---:|:---:|:---:| +| `num_inference_steps` | `int` | 50 | The number of diffusion steps during inference| +| `guidance_scale` | `float` | 0.0 | The classifier free guidance scale | +| `width` and `height` | `int` | None | The width and height of the generated image | + +**Extract parameters from request:** + ```python +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.data import DiffusionOutput + +def forward( + self, + req: OmniDiffusionRequest, +) -> DiffusionOutput: + # Extract prompts from request + if req.prompts is not None: + prompt = [ + p if isinstance(p, str) else (p.get("prompt") or "") + for p in req.prompts + ] + + # Extract sampling parameters + sampling_params = req.sampling_params + num_inference_steps = sampling_params.num_inference_steps or 50 + guidance_scale = sampling_params.guidance_scale or 7.5 + height = sampling_params.height or (self.default_sample_size * self.vae_scale_factor) + width = sampling_params.width or (self.default_sample_size * self.vae_scale_factor) + + # For image editing pipelines, extract images from multi_modal_data + if hasattr(req, 'multi_modal_data') and req.multi_modal_data: + input_images = req.multi_modal_data.get('image', []) + + # ... rest of generation logic +``` + +For an image editing model, an example `OmniDiffusionRequest` is like: +```python +{ + "prompt": "turn this cat to a dog", + "multi_modal_data": {"image": input_image} +}, +``` + +**Wrap output:** + +```diff + # Generate images + images = self.vae.decode(latents)[0] + +- return {"images": images} ++ return DiffusionOutput(output=images) +``` + +#### 2.4: Extract Pre/Post-Processing Functions + +vLLM-Omni separates image processing from the main pipeline for better modularity. + +**Post-processing function (required):** +```python +def get_your_model_post_process_func( + od_config: OmniDiffusionConfig, +): + """ + Create post-processing function for your model. + + Returns a function that converts latents to images. + """ + from diffusers.image_processor import VaeImageProcessor + import json + + # Load VAE config to get scale factor + model_path = od_config.model + if not os.path.exists(model_path): + from vllm_omni.diffusion.model_loader.utils import download_weights_from_hf_specific + model_path = download_weights_from_hf_specific(model_path, None, ["*"]) + + vae_config_path = os.path.join(model_path, "vae/config.json") + with open(vae_config_path) as f: + vae_config = json.load(f) + vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) + + # Create image processor + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor) + + def post_process_func(images: torch.Tensor): + return image_processor.postprocess(images, output_type="pil") + + return post_process_func +``` + +**Pre-processing function (for image editing pipelines):** + +```python +def get_your_model_pre_process_func( + od_config: OmniDiffusionConfig, +): + """ + Create pre-processing function for image editing. + + Returns a function that prepares input images. + """ + from PIL import Image + from diffusers.image_processor import VaeImageProcessor + + # Load VAE config + # ... (similar to post_process_func) + + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor) + + def pre_process_func( + request: OmniDiffusionRequest, + ): + for i, prompt in enumerate(request.prompts): + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + # image pre-processing + # after pre-processing, update the request attributes + ... + return request + + return pre_process_func +``` + +#### 2.5: Add Weight Loading Support + +Add methods for automatic weight downloading and loading: + +```python +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm.model_executor.models.utils import AutoWeightsLoader + +class YourModelPipeline(nn.Module): + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + # ... initialization code + + # Define weight sources for automatic loading + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ) + ] + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """ + Customize the weight loading behavior, such as filter weights name. + + Args: + weights: Iterable of (param_name, param_tensor) tuples + + Returns: + Set of loaded parameter names + """ + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) +``` + +### Step 3: Register Model + +Register your model in `vllm_omni/diffusion/registry.py` so vLLM-Omni can discover and load it. + +#### 3.1: Register Pipeline Class + +```python +# vllm_omni/diffusion/registry.py + _DIFFUSION_MODELS = { - # arch:(mod_folder, mod_relname, cls_name) - ... - "QwenImageEditPipeline": ( - "qwen_image", - "pipeline_qwen_image_edit", - "QwenImageEditPipeline", + # Format: "PipelineClassName": (module_folder, module_file, class_name) + + # Existing models + "QwenImagePipeline": ("qwen_image", "pipeline_qwen_image", "QwenImagePipeline"), + "FluxPipeline": ("flux", "pipeline_flux", "FluxPipeline"), + + # Add your model + "YourModelPipeline": ( + "your_model_name", # Module folder name + "pipeline_your_model", # Python file name (without .py) + "YourModelPipeline", # Pipeline class name ), - ... } ``` -+ registry pre-process get function + +#### 3.2: Register Pre/Post-Processing Function + ```python +# vllm_omni/diffusion/registry.py _DIFFUSION_PRE_PROCESS_FUNCS = { # arch: pre_process_func - ... + # `pre_process_func` function must be placed in {mod_folder}/{mod_relname}.py, + # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key + "GlmImagePipeline": "get_glm_image_pre_process_func", "QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func", - ... -} -``` -+ registry post-process get function -```python + # Add your model + "YourModelPipeline": "get_your_model_pre_process_func", # Optional +} _DIFFUSION_POST_PROCESS_FUNCS = { - # arch: post_process_func - ... - "QwenImageEditPipeline": "get_qwen_image_edit_post_process_func", - ... + # Format: "PipelineClassName": "function_name" + + # Existing models + "QwenImagePipeline": "get_qwen_image_post_process_func", + "FluxPipeline": "get_flux_post_process_func", + + # Add your model + "YourModelPipeline": "get_your_model_post_process_func", } ``` -## Step 4: Add an Example Script -For each newly integrated model, we need to provide examples script under the examples/ to demonstrate how to initialize the pipeline with Omni, pass in user inputs, and generate outputs. -Key point for writing the example: -+ Use the Omni entrypoint to load the model and construct the pipeline. +#### 3.3: Export from Module -+ Show how to format user inputs and pass them via omni.generate(...). - -+ Demonstrate the common runtime arguments, such as: - - + model path or model name - - + input image(s) or prompt text - - + key diffusion parameters (e.g., inference steps, guidance scale) - - + optional acceleration backends (e.g., Cache-DiT, TeaCache) - -+ Save or display the generated results so users can validate the integration. - -## Step 5: TeaCache Coefficient Estimation (Optional) - -If your model supports TeaCache acceleration, you need to estimate the polynomial coefficients for optimal caching performance. - -### 5.1 Add Extractor Function - -First, implement an extractor function in `vllm_omni/diffusion/cache/teacache/extractors.py`. The extractor extracts the modulated input and defines how to run transformer blocks: +Create/update `__init__.py` to export your classes: ```python -def extract_your_model_context( - module: nn.Module, - hidden_states: torch.Tensor, - timestep: torch.Tensor, - **kwargs: Any, -) -> CacheContext: - # 1. Preprocessing - temb = module.time_embed(timestep) +# vllm_omni/diffusion/models/your_model_name/__init__.py - # 2. Extract modulated input (for cache decision) - modulated_input = module.transformer_blocks[0].norm1(hidden_states, temb) - - # 3. Define transformer execution - def run_transformer_blocks(): - h = hidden_states - for block in module.transformer_blocks: - h = block(h, temb=temb) - return (h,) - - # 4. Define postprocessing - def postprocess(h): - return module.proj_out(module.norm_out(h, temb)) - - return CacheContext( - modulated_input=modulated_input, - hidden_states=hidden_states, - encoder_hidden_states=None, - temb=temb, - run_transformer_blocks=run_transformer_blocks, - postprocess=postprocess, - ) -``` - -Register it in `EXTRACTOR_REGISTRY`: -```python -EXTRACTOR_REGISTRY = { - ... - "YourTransformer2DModel": extract_your_model_context, -} -``` - -### 5.2 Add Adapter for Coefficient Estimation - -Add an adapter in `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py`: - -```python -class YourModelAdapter: - @staticmethod - def load_pipeline(model_path: str, device: str, dtype: torch.dtype) -> Any: - # Load your pipeline - ... - - @staticmethod - def get_transformer(pipeline: Any) -> tuple[Any, str]: - return pipeline.transformer, "YourTransformer2DModel" - - @staticmethod - def install_hook(transformer: Any, hook: DataCollectionHook) -> None: - registry = HookRegistry.get_or_create(transformer) - registry.register_hook(hook._HOOK_NAME, hook) - -_MODEL_ADAPTERS["YourModel"] = YourModelAdapter -``` - -### 5.3 Run Coefficient Estimation - -Use the provided script to estimate coefficients: - -```python -from vllm_omni.diffusion.cache.teacache.coefficient_estimator import ( - TeaCacheCoefficientEstimator, +from .pipeline_your_model import ( + YourModelPipeline, + get_your_model_post_process_func, ) -from datasets import load_dataset -from tqdm import tqdm +from .your_model_transformer import YourModelTransformer2DModel -# Load model -estimator = TeaCacheCoefficientEstimator( - model_path="/path/to/model", - model_type="Bagel", # Your model type - device="cuda", -) - -# Load prompts (paper suggests ~70 prompts) -dataset = load_dataset("nateraw/parti-prompts", split="train") -prompts = dataset["Prompt"][:70] - -# Collect data -for prompt in tqdm(prompts): - estimator.collect_from_prompt(prompt, num_inference_steps=50) - -# Estimate coefficients -coeffs = estimator.estimate(poly_order=4) -print(f"Coefficients: {coeffs}") +__all__ = [ + "YourModelPipeline", + "YourModelTransformer2DModel", + "get_your_model_post_process_func", +] ``` -### 5.4 Interpreting Coefficient Estimation Results +--- -The estimator outputs statistics and polynomial coefficients. Here's how to interpret them: +### Step 4: Add Example Script -**Example Output:** -``` -Data statistics: -Count: 48 -Input Diffs (x): min=1.1089e-02, max=5.2555e-02, mean=2.8435e-02 -Output Diffs (y): min=2.8242e-02, max=2.9792e-01, mean=7.0312e-02 -Coefficients: [1333131.29, -168644.23, 7950.51, -163.75, 1.26] + +If your model is one of Text-to-Image, Text-to-Audio, Text-to-Video, Image-to-Image, Image-to-Video models, you can simply try one of the following offline inference scripts to run your model: + +| Model Category | Offline Inference Script | +|---|---| +| Image-to-Image | `examples/offline_inference/image_to_image/image_edit.py` | +| Image-to-Video | `examples/offline_inference/image_to_video/image_to_video.py` | +| Text-to-Image | `examples/offline_inference/text_to_image/text_to_image.py` | +| Text-to-Audio | `examples/offline_inference/text_to_audio/text_to_audio.py` | +| Text-to-Video | `examples/offline_inference/text_to_video/text_to_video.py` | + + +If new CLI arguments need to be added, please edit the offline inference script corresponding to your model category from the table above, and update the example inference script in its corresponding document file (e.g., `examples/offline_inference/text_to_video/text_to_video.md`). + +For online inference, all the supported tasks are listed in `docs/user_guide/examples/online_serving/`. If your model falls into these categories, please check the corresponding documentation in this folder and the example at `examples/online_serving/TASK_NAME`. Update them accordingly if needed. + +--- + +If your model is an Omni (understanding and generation) model, please follow the steps below. + +#### 4.1: Create Example File + +Taking **BAGEL** model as examples for both offline and online: + +- Offline: `examples/offline_inference/bagel/` +- Online: `examples/online_serving/bagel/` + +Add **two example folders** for your model: + +```bash +mkdir -p examples/offline_inference/your_model_name +mkdir -p examples/online_serving/your_model_name ``` -**What to Check:** -- **Count**: Number of timestep pairs analyzed. Should be at least 30-50 for reliable estimation. Low count suggests insufficient prompts or inference steps. -- **Input/Output Ranges**: Verify output differences correlate with input differences. If ranges seem unusual, check your prompt diversity. -- **Coefficient Magnitude**: Extremely large values (>1e8) may indicate numerical instability - try collecting more diverse data. +**Offline (recommended minimum):** create `examples/offline_inference/your_model_name/end2end.py` and a README. -**Troubleshooting:** -- If results seem unreliable, try: - - Increasing number of prompts (100+ recommended) - - Using more diverse prompts from multiple datasets - - Adjusting `num_inference_steps` (try 20, 50, 100) +- Script: `examples/offline_inference/your_model_name/end2end.py` + - Parse args like BAGEL (`--model`, `--modality`, optional `--image-path`, `--steps`, etc.) + - Use `from vllm_omni.entrypoints.omni import Omni` (or `OmniDiffusion` if your model is diffusion-only) + - Save outputs (images/audio/video/text) with deterministic filenames (e.g., `output_0_0.png`) +- Doc: `examples/offline_inference/your_model_name/README.md` + - Include at least one runnable command, e.g.: -### 5.5 Add Coefficients to Config - -Add the estimated coefficients to `vllm_omni/diffusion/cache/teacache/config.py`: - -```python -_MODEL_COEFFICIENTS = { - ... - "YourTransformer2DModel": [ - 1.04730573e+06, # a4 - -1.34150749e+05, # a3 - 6.51517806e+03, # a2 - -1.41209108e+02, # a1 - 1.17241808e+00, # a0 - ], -} +```bash +cd examples/offline_inference/your_model_name +python end2end.py --model your-org/your-model-name --modality text2img --prompts "A cute cat" ``` -## Step 6: Open a Pull Request -When submitting a pull request to add support for a new model, please include the following information in the PR description: +#### 4.2: Add Online Serving Example (OpenAI-Compatible) -+ Output verification: provide generation outputs to verify correctness and model behavior. +Mirror BAGEL’s online serving setup: -+ Inference speed: provide a comparison with the corresponding implementation in Diffusers. +- Server launcher: `examples/online_serving/your_model_name/run_server.sh` + - Wrap `vllm serve ... --omni --port ...` (and `--stage-configs-path ...` if needed) +- Client: `examples/online_serving/your_model_name/openai_chat_client.py` + - Send requests to `POST /v1/chat/completions` + - Support multimodal inputs (e.g., base64 image) if your model needs it +- Doc: `examples/online_serving/your_model_name/README.md` + - Include both “launch server” and “send request”: -+ Parallelism support: specify the supported parallel sizes and any relevant limitations. +```bash +# Terminal 1: launch server +cd examples/online_serving/your_model_name +bash run_server.sh -+ Cache acceleration: check whether the model can be accelerated using Cache-Dit or not. +# Terminal 2: send request +python openai_chat_client.py --prompt "A cute cat" --modality text2img +``` -Providing these details helps reviewers evaluate correctness, performance improvements, and parallel scalability of the new model integration. +### Step 5: Test Your Implementation -# Testing +Before submitting, thoroughly test your implementation. + +#### 5.1: Performance/Speed Check + +Manually compare **latency/throughput** and **output quality** against a Diffusers baseline. + +For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num_inference_steps**, and **guidance settings**, and run multiple trials to reduce randomness. Record the results (and your hardware / driver / CUDA versions) in your PR description. + + +#### 5.2 Functionality Check in CI + +To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes. For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md). +--- + +## Advanced Features + +Once basic implementation works, add advanced features for better performance. + +### torch.compile Support + +Enable automatic compilation for repeated blocks: + +```python +# In your_model_transformer.py + +class YourModelTransformer2DModel(nn.Module): + # Specify which blocks can be compiled + _repeated_blocks = ["YourTransformerBlock"] # List of block class names + + def __init__(self, ...): + super().__init__() + # ... initialization +``` + +vLLM-Omni automatically compiles blocks in `_repeated_blocks` when `torch.compile` is available. + +### Tensor Parallelism + +See detailed guide: [How to add Tensor Parallel support](../features/tensor_parallel.md) + +**Quick setup:** + +1. Replace Linear layers by various parallel linear layers (e.g., `ColumnParallelLinear`) in vLLM +2. Check `tp_size` validity: `hidden_dim`, `num_heads`, and `num_kv_heads` must be divisible by `tp_size` + +**Usage:** Set `tensor_parallel_size` when initializing: +```python +omni = Omni(model="your-model", tensor_parallel_size=2) +``` + +### CFG Parallelism + +See detailed guide: [How to add CFG-Parallel support](../features/cfg_parallel.md) + +**Quick setup:** + +1. Implement `diffuse()` method +2. Inherit `CFGParallelMixin` in your pipeline class + +**Usage:** Set `cfg_parallel_size` when initializing: +```python +omni = Omni(model="your-model", cfg_parallel_size=2) +``` + +### Sequence Parallelism + +See detailed guide: [How to add Sequence Parallel support](../features/sequence_parallel.md) + +**Quick setup:** + +1. Add `_sp_plan` class attribute to transformer +2. Specify where to shard/gather tensors + +**Usage:** Set `ulysses_degree` and `ring_degree` when initializing: +```python +omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2) +``` + +### Cache Acceleration + +#### TeaCache + +See detailed guide: [How to add TeaCache support](../features/teacache.md) + +**Quick setup:** + +1. Write extractor function +2. Register in `EXTRACTOR_REGISTRY` +3. Add polynomial coefficients + +**Usage:** Set `cache_backend` and `cache_config` when initializing: +```python +omni = Omni(model="your-model", + cache_backend="tea_cache", + cache_config={"rel_l1_thresh": 0.2} +) + +``` + + +#### Cache-DiT + +See detailed guide: [How to add Cache-DiT support](../features/cache_dit.md) + +**Quick setup:** + +- For standard models: Works automatically +- For complex architectures: Write custom cache config + +**Usage:** Set `cache_backend` and `cache_config` when initializing: +```python +omni = Omni(model="your-model", + cache_backend="cache_dit", + cache_config={ + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + } +) +``` + +--- + +## Troubleshooting + + +**Issue: ImportError when loading model** + +**Symptoms:** `ModuleNotFoundError` or `ImportError` when calling `Omni(model="your-model")` + +**Causes:** + +1. Model not registered in `registry.py` +2. Wrong class name in registry +3. Missing `__init__.py` exports + + +**Issue: Shape mismatch in attention** + +**Symptoms:** `RuntimeError: shape mismatch` in attention forward + +**Cause:** Incorrect tensor reshaping for vLLM-Omni's attention interface + +**Solution:** Ensure correct shapes: + +```python +# vLLM-Omni expects: [batch, seq_len, num_heads, head_dim] +query = query.view(batch_size, seq_len, self.num_heads, self.head_dim) +key = key.view(batch_size, kv_seq_len, self.num_kv_heads, self.head_dim) +value = value.view(batch_size, kv_seq_len, self.num_kv_heads, self.head_dim) + +hidden_states = self.attn(query, key, value, attn_metadata=attn_metadata) + +# Reshape back: [batch, seq_len, num_heads, head_dim] → [batch, seq_len, hidden_size] +hidden_states = hidden_states.reshape(batch_size, seq_len, -1) +``` + +**Issue: Different outputs compared to Diffusers** + +**Symptoms:** Generated images look different from Diffusers + +**Causes:** + +1. Attention backend differences (FlashAttention vs PyTorch SDPA) +2. Missing normalization or scaling + +**4. Issue: Out of memory (OOM)** + +**Symptoms:** CUDA out of memory errors + +**Solutions:** + +1. **Reduce batch size:** + ```python + omni.generate(prompts=[...], max_batch_size=2) + ``` + +2. **Use smaller image size:** + ```python + sampling_params = OmniDiffusionSamplingParams(height=512, width=512) + ``` + +3. **Enable model offloading:** + ```python + omni = Omni(model="...", enable_cpu_offload=True) + ``` + +4. **Apply vae tiling and slicing** + ```python + omni = Omni(model="...", vae_use_slicing=True, vae_use_tiling=True,) + ``` + +--- + +## Pull Request Checklist + +When submitting a PR to add your model, include: + +**1. Implementation Files** + +- ✅ Transformer model (`xxx_transformer.py`) +- ✅ Pipeline (`pipeline_xxx.py`) +- ✅ Registry entries in `registry.py` +- ✅ `__init__.py` with proper exports + +**2. Example and Tests** + +- ✅ Example script in `examples/` +- ✅ Test file in `tests/e2e/` +- ✅ Documentation (`docs/`) creation or updates + +_Note: End-to-end test files in `tests/e2e/` are optional but strongly recommended. README updates are required for all new models._ + +**3. Documentation Updates** + +- ✅ Add model to supported models table in `docs/models/supported_models.md` +- ✅ If supporting acceleration features (e.g., sequence parallelism, CFG parallel), update acceleration feature tables in: + - `docs/user_guide/diffusion_acceleration.md` + - `docs/user_guide/diffusion/parallelism_acceleration.md` + +--- + +## Model Recipe -## Adding a Model Recipe After implementing and testing your model, please add a model recipe to the [vllm-project/recipes](https://github.com/vllm-project/recipes) repository. This helps other users understand how to use your model with vLLM-Omni. + +**What to Include** + +Your recipe should include: + +1. **Model Overview**: Brief description of the model and its capabilities +2. **Installation Instructions**: Step-by-step setup instructions including: + - Installing vllm-omni and dependencies + - Installing any additional required packages (e.g., xformers, diffusers) + - Any version requirements +3. **Usage Examples**: Command-line examples demonstrating how to run the model +4. **Configuration Details**: Important configuration parameters and their meanings + +**Example** + +For reference, see the [LongCat recipe example](https://github.com/vllm-project/recipes/pull/179) which demonstrates the expected format and structure. + +**Recipe Location** + +Create your recipe file in the appropriate directory structure: +- For organization-specific models: `OrganizationName/ModelName.md` +- For general models: `ModelName.md` + +The recipe should be a Markdown file that provides clear, reproducible instructions for users to get started with your model. + +--- + +## Reference Implementations + +Study these complete examples: + +| Model | Architecture | Key Features | Files | +|-------|--------------|--------------|-------| +| **Qwen-Image** | Dual-stream transformer | CFG-Parallel, SP, TP, Cache | `vllm_omni/diffusion/models/qwen_image/` | +| **Wan2.2** | Video transformer | Dual transformers, SP, CFG-Parallel | `vllm_omni/diffusion/models/wan2_2/` | + +--- + +## Summary + +Adding a diffusion model to vLLM-Omni involves: + +1. ✅ **Adapt transformer** - Replace attention, remove mixins, add config support +2. ✅ **Adapt pipeline** - Change interface, add request handling, extract processing +3. ✅ **Register model** - Add entries to `registry.py` +4. ✅ **Add examples** - Provide runnable scripts +5. ✅ **Test thoroughly** - Verify correctness and performance +6. ✅ **Add advanced features** - Enable parallelism and acceleration (optional) +7. ✅ **Submit PR** - Include verification results and documentation + +**Need help?** Check reference implementations or ask in [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai). diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index f91a626e0..d536f4bc7 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -233,143 +233,6 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m | Hybrid Ulysses + Ring | 2 | 2 | 24.3s | 1.87x | -##### How to parallelize a new model - -NOTE: "Terminology: SP vs CP" - Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in the [diffusers library](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py). - We use "Sequence Parallelism" to align with vLLM-Omni's terminology. - ---- - -###### Non-intrusive `_sp_plan` (Recommended) - -The `_sp_plan` mechanism allows SP without modifying `forward()` logic. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries. - -**Requirements for `forward()` function:** - -- Tensor operations that need sharding/gathering must happen at **`nn.Module` boundaries** (not inline Python operations) -- If your `forward()` contains inline tensor operations (e.g., `torch.cat`, `pad_sequence`) that need sharding, **extract them into a submodule** - -**When to create a submodule:** - -```python -# ❌ BAD: Inline operations - hooks cannot intercept -def forward(self, x, cap_feats): - unified = torch.cat([x, cap_feats], dim=1) # Cannot be sharded via _sp_plan - ... - -# ✅ GOOD: Extract into a submodule -class UnifiedPrepare(nn.Module): - def forward(self, x, cap_feats): - return torch.cat([x, cap_feats], dim=1) # Now can be sharded via _sp_plan - -class MyModel(nn.Module): - def __init__(self): - self.unified_prepare = UnifiedPrepare() # Submodule - - def forward(self, x, cap_feats): - unified = self.unified_prepare(x, cap_feats) # Hook can intercept here -``` - ---- - -###### Defining `_sp_plan` - -**Type definitions** (see [diffusers `_modeling_parallel.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py) for reference): - -```python -from vllm_omni.diffusion.distributed.sp_plan import ( - SequenceParallelInput, # Corresponds to diffusers' ContextParallelInput - SequenceParallelOutput, # Corresponds to diffusers' ContextParallelOutput -) -``` - -| Parameter | Description | -|-----------|-------------| -| `split_dim` | Dimension to split/gather (usually `1` for sequence) | -| `expected_dims` | Expected tensor rank for validation (optional) | -| `split_output` | `False`: shard **input** parameters; `True`: shard **output** tensors | -| `auto_pad` | Auto-pad if sequence not divisible by world_size (Ulysses only) | - -**Key naming convention:** - -| Key | Meaning | Python equivalent | -|-----|---------|-------------------| -| `""` | Root model | `model` | -| `"blocks.0"` | First element of ModuleList | `model.blocks[0]` | -| `"blocks.*"` | All elements of ModuleList | `for b in model.blocks` | -| `"outputs.main"` | ModuleDict entry | `model.outputs["main"]` | - -**Dictionary key types:** - -| Key type | `split_output` | Description | -|----------|----------------|-------------| -| `"param_name"` (str) | `False` | Shard **input parameter** by name | -| `0`, `1` (int) | `True` | Shard **output tuple** by index | - -**Example** (similar to [diffusers `transformer_wan.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_wan.py)): - -```python -class MyTransformer(nn.Module): - _sp_plan = { - # Shard rope module OUTPUTS (returns tuple) - "rope": { - 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # cos - 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # sin - }, - # Shard transformer block INPUT parameter - "blocks.0": { - "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), - }, - # Gather at final projection - "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), - } -``` - ---- - -###### Hook flow - -``` -Input → [SequenceParallelSplitHook: pre_forward] → Module.forward() → [post_forward] → ... - ↓ -... → [SequenceParallelGatherHook: post_forward] → Output -``` - -1. **SplitHook** shards tensors before/after the target module -2. **Attention layers** handle Ulysses/Ring communication internally -3. **GatherHook** collects sharded outputs - -The framework automatically applies these hooks when `sequence_parallel_size > 1`. - ---- - -###### Method 2: Intrusive modification (For complex cases) - -For models with dynamic sharding logic that cannot be expressed via `_sp_plan`: - -```python -from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather - -def forward(self, hidden_states, ...): - if self.parallel_config.sequence_parallel_size > 1: - hidden_states = sp_shard(hidden_states, dim=1) - # ... computation ... - output = sp_gather(output, dim=1) - return output -``` - ---- - -###### Choosing the right approach - -| Scenario | Approach | -|----------|----------| -| Standard transformer | `_sp_plan` | -| Inline tensor ops need sharding | Extract to submodule + `_sp_plan` | -| Dynamic/conditional sharding | Intrusive modification | - - ### CFG-Parallel #### Offline Inference @@ -426,142 +289,3 @@ You can enable CFG-Parallel in online serving for diffusion models via `--cfg-pa ```bash vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2 ``` - -#### How to parallelize a pipeline - -This section describes how to add CFG-Parallel to a diffusion **pipeline**. We use the Qwen-Image pipeline (`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`) as the reference implementation. - -In `QwenImagePipeline`, each diffusion step runs two denoiser forward passes sequentially: - -- positive (prompt-conditioned) -- negative (negative-prompt-conditioned) - -CFG-Parallel assigns these two branches to different ranks in the **CFG group** and synchronizes the results. - -vLLM-omni provides `CFGParallelMixin` base class that encapsulates the CFG parallel logic. By inheriting from this mixin and calling its methods, pipelines can easily implement CFG parallel without writing repetitive code. - -**Key Methods in CFGParallelMixin:** -- `predict_noise_maybe_with_cfg()`: Automatically handles CFG parallel noise prediction -- `scheduler_step_maybe_with_cfg()`: Scheduler step with automatic CFG rank synchronization - -**Example Implementation:** - -```python -class QwenImageCFGParallelMixin(CFGParallelMixin): - """ - Base Mixin class for Qwen Image pipelines providing shared CFG methods. - """ - - def diffuse( - self, - prompt_embeds: torch.Tensor, - prompt_embeds_mask: torch.Tensor, - negative_prompt_embeds: torch.Tensor, - negative_prompt_embeds_mask: torch.Tensor, - latents: torch.Tensor, - img_shapes: torch.Tensor, - txt_seq_lens: torch.Tensor, - negative_txt_seq_lens: torch.Tensor, - timesteps: torch.Tensor, - do_true_cfg: bool, - guidance: torch.Tensor, - true_cfg_scale: float, - image_latents: torch.Tensor | None = None, - cfg_normalize: bool = True, - additional_transformer_kwargs: dict[str, Any] | None = None, - ) -> torch.Tensor: - self.transformer.do_true_cfg = do_true_cfg - - for i, t in enumerate(timesteps): - timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype) - - # Prepare kwargs for positive (conditional) prediction - positive_kwargs = { - "hidden_states": latents, - "timestep": timestep / 1000, - "guidance": guidance, - "encoder_hidden_states_mask": prompt_embeds_mask, - "encoder_hidden_states": prompt_embeds, - "img_shapes": img_shapes, - "txt_seq_lens": txt_seq_lens, - } - - # Prepare kwargs for negative (unconditional) prediction - if do_true_cfg: - negative_kwargs = { - "hidden_states": latents, - "timestep": timestep / 1000, - "guidance": guidance, - "encoder_hidden_states_mask": negative_prompt_embeds_mask, - "encoder_hidden_states": negative_prompt_embeds, - "img_shapes": img_shapes, - "txt_seq_lens": negative_txt_seq_lens, - } - else: - negative_kwargs = None - - # Predict noise with automatic CFG parallel handling - # - In CFG parallel mode: rank0 computes positive, rank1 computes negative - # - Automatically gathers results and combines them on rank0 - noise_pred = self.predict_noise_maybe_with_cfg( - do_true_cfg=do_true_cfg, - true_cfg_scale=true_cfg_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=cfg_normalize, - ) - - # Step scheduler with automatic CFG synchronization - # - Only rank0 computes the scheduler step - # - Automatically broadcasts updated latents to all ranks - latents = self.scheduler_step_maybe_with_cfg( - noise_pred, t, latents, do_true_cfg - ) - - return latents -``` - -**How it works:** -1. Prepare separate `positive_kwargs` and `negative_kwargs` for conditional and unconditional predictions -2. Call `predict_noise_maybe_with_cfg()` which: - - Detects if CFG parallel is enabled (`get_classifier_free_guidance_world_size() > 1`) - - Distributes computation: rank0 processes positive, rank1 processes negative - - Gathers predictions and combines them using `combine_cfg_noise()` on rank0 - - Returns combined noise prediction (only valid on rank0) -3. Call `scheduler_step_maybe_with_cfg()` which: - - Only rank0 computes the scheduler step - - Broadcasts the updated latents to all ranks for synchronization - -**How to customize** - -Some pipelines may need to customize the following functions in `CFGParallelMixin`: -1. You may need to edit `predict_noise` function for custom behaviors. -```python -def predict_noise(self, *args, **kwargs): - """ - Forward pass through transformer to predict noise. - - Subclasses should override this if they need custom behavior, - but the default implementation calls self.transformer. - """ - return self.transformer(*args, **kwargs)[0] - -``` -2. The default normalization function after combining the noise predictions from both branches is as follows. You may need to customize it. -```python -def cfg_normalize_function(self, noise_pred, comb_pred): - """ - Normalize the combined noise prediction. - - Args: - noise_pred: positive noise prediction - comb_pred: combined noise prediction after CFG - - Returns: - Normalized noise prediction tensor - """ - cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True) - noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True) - noise_pred = comb_pred * (cond_norm / noise_norm) - return noise_pred -```