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: