delete review-PR

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

.claude/skills/add-diffusion-model/SKILL.md

@@ -0,0 +1,534 @@
+---
+name: add-diffusion-model
+description: Add a new diffusion model (text-to-image, text-to-video, image-to-video, text-to-audio, image editing) to vLLM-Omni, including Cache-DiT acceleration and parallelism support (TP, SP/USP, CFG-Parallel, HSDP). Use when integrating a new diffusion model, porting a diffusers pipeline or a custom model repo to vllm-omni, creating a new DiT transformer adapter, adding diffusion model support, or enabling multi-GPU parallelism and cache acceleration for an existing model.
+---
+
+# Adding a Diffusion Model to vLLM-Omni
+
+## Overview
+
+This skill guides you through adding a new diffusion model to vLLM-Omni. The model may come from HuggingFace Diffusers (structured pipeline) or from a private/custom repo. The workflow differs significantly depending on the source.
+
+## Prerequisites
+
+Before starting, determine:
+
+1. **Model category**: Text-to-Image, Text-to-Video, Image-to-Video, Image Editing, Text-to-Audio, or Omni
+2. **Reference source**: Diffusers pipeline, custom repo, or a combination
+3. **Model HuggingFace ID** or local checkpoint path
+4. **Architecture**: Scheduler, text encoder, VAE, transformer/backbone
+
+## Step 0: Classify the Migration Path
+
+Check the model's HF repo for `model_index.json`. This determines your path:
+
+| Scenario | How to identify | Migration path |
+|----------|----------------|----------------|
+| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip to Step 5 (test) and Step 7 (docs) |
+| **Diffusers-based** | Has standard `model_index.json` with `_diffusers_version`, subfolders for `transformer/`, `vae/`, etc. | Follow **Path A** below |
+| **Custom/private repo** | No diffusers `model_index.json`, weights in non-standard format, custom model code in a separate git repo | Follow **Path B** below |
+| **Hybrid** | Has some diffusers components (VAE) but custom transformer/fusion | Mix of Path A and Path B |
+
+## Path A: Diffusers-Based Model
+
+For models with a standard diffusers layout. See [references/transformer-adaptation.md](references/transformer-adaptation.md) for detailed code patterns.
+
+### A1. Analyze `model_index.json`
+
+Identify components: `transformer`, `scheduler`, `vae`, `text_encoder`, `tokenizer`.
+
+### A2. Create model directory
+
+```
+vllm_omni/diffusion/models/your_model_name/
+├── __init__.py
+├── pipeline_your_model.py
+└── your_model_transformer.py
+```
+
+### A3. Adapt transformer
+
+1. Copy from diffusers source. Remove mixins (`ModelMixin`, `ConfigMixin`, `AttentionModuleMixin`).
+2. Replace attention with `vllm_omni.diffusion.attention.layer.Attention` (QKV shape: `[B, seq, heads, head_dim]`).
+3. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`.
+4. Add `load_weights()` method mapping diffusers weight names to vllm-omni names.
+5. Add class attributes: `_repeated_blocks`, `_layerwise_offload_blocks_attr`.
+
+### A4. Adapt pipeline
+
+Inherit from `nn.Module`. The key contract:
+
+```python
+class YourPipeline(nn.Module):
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        # Load VAE, text encoder, tokenizer via from_pretrained()
+        # Instantiate transformer (weights loaded later via weights_sources)
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model, subfolder="transformer",
+                prefix="transformer.", fall_back_to_pt=True)]
+
+    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+        # Encode prompt → prepare latents → denoise loop → VAE decode
+        return DiffusionOutput(output=output)
+
+    def load_weights(self, weights):
+        return AutoWeightsLoader(self).load_weights(weights)
+```
+
+Add post/pre-process functions in the same pipeline file. Register them in `registry.py`.
+
+### A5. Register, test, docs → continue at Step 4 below.
+
+---
+
+## Path B: Custom/Private Repo Model
+
+For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3.
+
+### B1. Understand the reference repo
+
+Study the original model's code to identify:
+- **Model architecture files** (transformers, fusion modules, embeddings)
+- **Weight format** (safetensors, `.pth`, custom checkpoint structure)
+- **Weight loading helpers** (custom init functions, checkpoint loaders)
+- **Pre/post-processing** (image/audio transforms, tokenization, VAE encode/decode)
+- **External dependencies** (packages not on PyPI)
+- **Config format** (JSON config files, hardcoded dicts)
+
+### B2. Decide what lives WHERE
+
+This is the key design decision for custom models. Follow these placement rules:
+
+| Code type | Where to place | Example |
+|-----------|---------------|---------|
+| **Pipeline orchestration** (init, forward, denoise loop) | `vllm_omni/diffusion/models/<name>/pipeline_<name>.py` | Always required |
+| **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models/<name>/<name>_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` |
+| **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models/<name>/` as separate files | `autoencoder.py`, `fusion.py` |
+| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | `dreamid_omni` package via git clone |
+| **Hardcoded model configs** | Module-level dicts in pipeline file | `VIDEO_CONFIG`, `AUDIO_CONFIG` dicts |
+| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_dreamid_omni.py` |
+| **Custom `model_index.json`** | Generated by download script, placed at model root | Minimal: `{"_class_name": "YourPipeline", ...}` |
+
+### B3. Handle external dependencies
+
+If the model's code lives in a separate git repo:
+
+**Option 1: Import with graceful fallback** (recommended for models with external utils)
+
+```python
+try:
+    from external_model.utils import init_vae, load_checkpoint
+except ImportError:
+    raise ImportError(
+        "Failed to import from dependency 'external_model'. "
+        "Please run the download script first."
+    )
+```
+
+**Option 2: Port the code directly** (preferred when feasible)
+
+Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and adapt them. This avoids external dependencies. BAGEL does this — `autoencoder.py` and `bagel_transformer.py` are ported directly.
+
+**Decision criteria**: Port if the code is self-contained and won't diverge. Use external deps if the model repo is actively maintained and the code is complex.
+
+### B4. Handle custom weight loading
+
+Custom models have two patterns for weight loading:
+
+**Pattern 1: Bypass standard loader** (DreamID-Omni style)
+
+When the original model has complex custom init functions that load weights in `__init__`:
+
+```python
+class CustomPipeline(nn.Module):
+    def __init__(self, *, od_config, prefix=""):
+        super().__init__()
+        model = od_config.model
+        # Load everything eagerly in __init__ using custom helpers
+        self.vae = custom_init_vae(model, device=self.device)
+        self.text_encoder = custom_init_text_encoder(model, device=self.device)
+        self.transformer = CustomFusionModel(CONFIG)
+        load_custom_checkpoint(self.transformer,
+            checkpoint_path=os.path.join(model, "model.safetensors"))
+        # NO weights_sources defined — bypasses standard loader
+
+    def load_weights(self, weights):
+        pass  # No-op — all weights loaded in __init__
+```
+
+**Pattern 2: Use standard loader with custom `load_weights`** (BAGEL style)
+
+When weights are in safetensors format but need name remapping:
+
+```python
+class CustomPipeline(nn.Module):
+    def __init__(self, *, od_config, prefix=""):
+        super().__init__()
+        # Instantiate model architecture without weights
+        self.bagel = BagelModel(config)
+        self.vae = AutoEncoder(ae_params)
+
+        # Point loader at the safetensors in the model root
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder=None,  # weights at root, not in subfolder
+                prefix="",
+                fall_back_to_pt=False,
+            )
+        ]
+
+    def load_weights(self, weights):
+        # Custom name remapping for non-diffusers weight names
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            # Remap original weight names to vllm-omni module names
+            name = self._remap_weight_name(name)
+            if name in params:
+                default_weight_loader(params[name], tensor)
+                loaded.add(name)
+        return loaded
+```
+
+### B5. Create the `model_index.json`
+
+Custom models need a `model_index.json` at the model root for vllm-omni to discover them. For custom models, this is minimal:
+
+```json
+{
+    "_class_name": "YourModelPipeline",
+    "custom_key": "path/to/custom_weights.safetensors"
+}
+```
+
+The `_class_name` must match a key in `_DIFFUSION_MODELS` in `registry.py`. Additional keys are model-specific (accessed via `od_config.model_config`).
+
+If the model's weights come from multiple HF repos, write a **download script** that:
+1. Downloads from each repo
+2. Assembles into a single directory
+3. Generates `model_index.json`
+4. Installs any external dependencies (git clone + `.pth` file)
+
+Place at: `examples/offline_inference/<name>/download_<name>.py`
+
+### B6. Handle multi-modal inputs
+
+If the model accepts images, audio, or other multi-modal inputs, implement the protocol classes from `vllm_omni/diffusion/models/interface.py`:
+
+```python
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportAudioInput
+
+class MyPipeline(nn.Module, SupportImageInput, SupportAudioInput):
+    # Protocol markers — the engine uses these to enable proper input routing
+    pass
+```
+
+Preprocessing for custom models is typically done **inside `forward()`** rather than via registered pre-process functions, since the logic is often tightly coupled to the model.
+
+### B7. Continue at Step 4 below.
+
+---
+
+## Common Steps (Both Paths)
+
+### Step 4: Register Model in registry.py
+
+Edit `vllm_omni/diffusion/registry.py`:
+
+```python
+_DIFFUSION_MODELS = {
+    "YourModelPipeline": ("your_model_name", "pipeline_your_model", "YourModelPipeline"),
+}
+_DIFFUSION_POST_PROCESS_FUNCS = {
+    "YourModelPipeline": "get_your_model_post_process_func",  # if applicable
+}
+_DIFFUSION_PRE_PROCESS_FUNCS = {
+    "YourModelPipeline": "get_your_model_pre_process_func",  # if applicable
+}
+```
+
+The registry key is the `_class_name` from `model_index.json`. The tuple is `(folder_name, module_file, class_name)`.
+
+Create `__init__.py` exporting the pipeline class and any factory functions.
+
+### Step 5: Run, Test, Debug
+
+Use the appropriate existing example script:
+
+| Category | Script |
+|----------|--------|
+| Text-to-Image | `examples/offline_inference/text_to_image/text_to_image.py` |
+| Text-to-Video | `examples/offline_inference/text_to_video/text_to_video.py` |
+| Image-to-Video | `examples/offline_inference/image_to_video/image_to_video.py` |
+| Image-to-Image | `examples/offline_inference/image_to_image/image_edit.py` |
+| Text-to-Audio | `examples/offline_inference/text_to_audio/text_to_audio.py` |
+
+For custom/Omni models that don't fit these categories, create a dedicated example script.
+
+**Validation**: No errors, output is meaningful, quality matches reference implementation.
+
+See [references/troubleshooting.md](references/troubleshooting.md) for common errors.
+
+### Step 6: Add Example Scripts
+
+For Omni or custom models, create:
+- `examples/offline_inference/your_model_name/` — offline script + README
+- `examples/online_serving/your_model_name/` — server script + client
+- Download script if weights require assembly from multiple sources
+
+### Step 7: Update Documentation
+
+Required updates:
+1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table
+2. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
+3. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
+4. `examples/offline_inference/xxx/README.md` — offline example docs
+5. `examples/online_serve/xxx/README.md` — online serve docs
+
+### Step 8: Add E2E Tests (Recommended)
+
+Create `tests/e2e/online_serving/test_your_model_expansion.py`.
+
+### Step 9: Add Cache-DiT Acceleration
+
+Cache-DiT accelerates inference by caching intermediate computation results across denoising steps. After your model is working correctly on a single GPU, add cache-dit support.
+
+See [references/cache-dit-patterns.md](references/cache-dit-patterns.md) for detailed code patterns.
+
+#### 9a. Determine your model type
+
+| Model Type | Description | Action |
+|------------|-------------|--------|
+| **Standard single-transformer** | One transformer with one `ModuleList` of blocks | No code needed — `CacheDiTBackend` auto-detects via `enable_cache_for_dit()` |
+| **Multi-block-list** | One transformer with multiple block lists (e.g., `transformer_blocks` + `single_transformer_blocks`) | Write custom enabler with `BlockAdapter` |
+| **Dual-transformer** | Two transformers (e.g., high-noise + low-noise) | Write custom enabler with `BlockAdapter` wrapping both |
+
+#### 9b. Standard models — verify automatic support
+
+For standard single-transformer models, test directly:
+
+```python
+omni = Omni(
+    model="your-model-name",
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1,
+        "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4,
+    }
+)
+```
+
+Check logs for "Cache-dit enabled successfully on xxx". If it works, skip to Step 9e.
+
+#### 9c. Custom architectures — write a custom enabler
+
+For multi-block-list or dual-transformer models, write a custom enabler function:
+
+```python
+from cache_dit import BlockAdapter, ForwardPattern, ParamsModifier, DBCacheConfig
+
+def enable_cache_for_your_model(pipeline, cache_config):
+    db_cache_config = DBCacheConfig(
+        num_inference_steps=None,
+        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,
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=pipeline.transformer,
+            blocks=[
+                pipeline.transformer.transformer_blocks,
+                pipeline.transformer.single_transformer_blocks,
+            ],
+            forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
+            params_modifiers=[ParamsModifier(...)],
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
+        cache_dit.refresh_context(
+            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
+        )
+    return refresh_cache_context
+```
+
+#### 9d. Register the custom enabler
+
+Add your enabler to `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,
+    "YourModelPipeline": enable_cache_for_your_model,  # Add here
+}
+```
+
+#### 9e. Test Cache-DiT
+
+```python
+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) logs show cache enabled, 2) 1.5-2x speedup, 3) output quality acceptable vs baseline.
+
+If quality degrades, lower `residual_diff_threshold` (try 0.12-0.18) or increase `max_warmup_steps` (try 6-8).
+
+---
+
+### Step 10: Add Parallelism Support
+
+After the model works on a single GPU, add multi-GPU parallelism. Add each type incrementally, testing after each addition.
+
+See [references/parallelism-patterns.md](references/parallelism-patterns.md) for detailed code patterns and API reference.
+
+**Recommended order**: TP → SP/USP → CFG Parallel → HSDP
+
+#### 10a. Tensor Parallelism (TP)
+
+Shards DiT linear layers across GPUs. Requires code changes in the transformer.
+
+**What to change in the transformer**:
+1. Replace `nn.Linear` with `ColumnParallelLinear` / `RowParallelLinear` / `QKVParallelLinear`
+2. Update `load_weights()` to handle QKV fusion with `stacked_params_mapping`
+3. Use `self.to_qkv.num_heads` (local heads) instead of total heads for split sizes
+
+```python
+from vllm.model_executor.layers.linear import (
+    QKVParallelLinear, RowParallelLinear, ColumnParallelLinear,
+)
+
+# Attention: QKV → RowParallel output
+self.to_qkv = QKVParallelLinear(dim, head_dim, num_heads, num_kv_heads)
+self.to_out = RowParallelLinear(dim, dim, input_is_parallel=True)
+
+# FFN: ColumnParallel → RowParallel
+self.w1 = ColumnParallelLinear(dim, ffn_dim)
+self.w2 = RowParallelLinear(ffn_dim, dim, input_is_parallel=True)
+```
+
+**Constraints**: `num_heads % tp_size == 0` and `num_kv_heads % tp_size == 0`.
+
+**Test**: `--tensor-parallel-size 2`
+
+#### 10b. Sequence Parallelism (SP / USP)
+
+Splits sequence tokens across GPUs. Non-intrusive via `_sp_plan` on the transformer class — no changes to `forward()`.
+
+**What to change in the transformer**:
+
+Add `_sp_plan` class attribute:
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput, SequenceParallelOutput,
+)
+
+class YourTransformer(nn.Module):
+    _sp_plan = {
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+If inline tensor ops (e.g., `torch.cat`) exist between shard/gather points, extract them into `nn.Module` submodules so hooks can intercept them.
+
+For RoPE that needs splitting, add an entry for the RoPE module with `split_output=True`.
+
+**Test**: `--ulysses-degree 2` (offline) or `--usp 2` (online serving)
+
+#### 10c. CFG Parallel
+
+Distributes positive/negative CFG branches across 2 GPUs. Requires the pipeline to inherit `CFGParallelMixin`.
+
+**What to change in the pipeline**:
+
+```python
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+
+class YourPipeline(nn.Module, CFGParallelMixin):
+    def diffuse(self, ...) -> torch.Tensor:
+        for i, t in enumerate(timesteps):
+            positive_kwargs = {...}
+            negative_kwargs = {...} if do_true_cfg else None
+            noise_pred = self.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_true_cfg, true_cfg_scale=cfg_scale,
+                positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs,
+            )
+            latents = self.scheduler_step_maybe_with_cfg(
+                noise_pred, t, latents, do_true_cfg
+            )
+        return latents
+```
+
+Override `predict_noise()` if your transformer call is non-standard. Override `combine_cfg_noise()` for multi-output models (e.g., video + audio).
+
+**Constraint**: Exactly 2 GPUs. Only for models using classifier-free guidance.
+
+**Test**: `--cfg-parallel-size 2`
+
+#### 10d. HSDP (Hybrid Sharded Data Parallel)
+
+Shards transformer weights via PyTorch FSDP2 to reduce per-GPU VRAM. No code changes to the forward pass — just add a class attribute.
+
+**What to change in the transformer**:
+
+```python
+class YourTransformer(nn.Module):
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block]
+```
+
+**Constraint**: Cannot combine with TP. For standalone HSDP, set `hsdp_shard_size` explicitly.
+
+**Test**: `--use-hsdp` or `DiffusionParallelConfig(use_hsdp=True)`
+
+#### 10e. Update parallelism documentation
+
+After adding parallelism support, update:
+1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table
+2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch)
+
+---
+
+## Iterative Development Tips
+
+1. **Start minimal**: Basic generation first, no parallelism/caching
+2. **Use `--enforce-eager`**: Disable torch.compile during debugging
+3. **Use small models**: Test with smaller variants first
+4. **Check tensor shapes**: Most errors are reshape mismatches in attention
+5. **Add features incrementally**: Single GPU → TP → SP → CFG → HSDP → Cache-DiT
+6. **For custom models**: Get the model running with the original code first, then progressively replace components with vllm-omni equivalents
+7. **Cache-DiT before parallelism tuning**: Cache-DiT is lossy — verify quality at baseline before combining with parallelism
+8. **Combine lossless + lossy**: e.g., TP + SP + Cache-DiT for maximum throughput
+
+## Reference Files
+
+- [Transformer Adaptation](references/transformer-adaptation.md) — porting transformers from diffusers
+- [Custom Model Patterns](references/custom-model-patterns.md) — patterns for non-diffusers models
+- [Parallelism Patterns](references/parallelism-patterns.md) — TP, SP/USP, CFG parallel, HSDP implementation details
+- [Cache-DiT Patterns](references/cache-dit-patterns.md) — cache-dit acceleration for standard and custom architectures
+- [Troubleshooting](references/troubleshooting.md) — common errors and fixes

.claude/skills/add-diffusion-model/references/cache-dit-patterns.md

@@ -0,0 +1,254 @@
+# Cache-DiT Patterns Reference
+
+## Overview
+
+Cache-DiT accelerates Diffusion Transformers by caching intermediate computation results across denoising steps. Adjacent steps produce similar features, so redundant computations can be skipped.
+
+Three caching strategies:
+- **DBCache**: Dynamic block-level caching — selectively computes or caches transformer blocks based on residual differences
+- **TaylorSeer**: Calibration-based prediction using Taylor expansion to estimate block outputs
+- **SCM** (Step Computation Masking): Dynamic step skipping based on configurable policies
+
+**Typical speedup**: 1.5-2.5x depending on model and configuration.
+
+**Official docs**: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cache_dit
+
+## Architecture
+
+vLLM-Omni integrates cache-dit through `CacheDiTBackend`:
+
+| Component | Purpose |
+|-----------|---------|
+| `CacheDiTBackend` | Unified backend — auto-selects enabler (standard or custom) |
+| `enable_cache_for_dit()` | Default enabler for standard single-transformer models |
+| `CUSTOM_DIT_ENABLERS` dict | Registry of custom enablers keyed by pipeline class name |
+| `BlockAdapter` | Wraps complex architectures (multi-block-list or multi-transformer) |
+| `ForwardPattern` | Specifies block forward signature: `Pattern_0`, `Pattern_1`, `Pattern_2` |
+| `ParamsModifier` | Per-transformer or per-block-list config customization |
+| `DBCacheConfig` | Configuration for DBCache parameters |
+| `cache_dit.refresh_context()` | Updates cache context when `num_inference_steps` changes |
+
+**Source files**:
+- `vllm_omni/diffusion/cache/cache_dit_backend.py` — `CacheDiTBackend`, enablers, `CUSTOM_DIT_ENABLERS`
+- `vllm_omni/diffusion/cache/` — cache backend implementations
+
+## Standard Models: Automatic Support
+
+Most DiT models follow this pattern:
+- Single transformer with one `nn.ModuleList` of blocks
+- Standard forward signature
+- Compatible with cache-dit's automatic detection
+
+**Examples**: Qwen-Image, Z-Image, FLUX
+
+No code changes needed. `CacheDiTBackend` automatically uses `enable_cache_for_dit()`:
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    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, cache_config):
+    db_cache_config = DBCacheConfig(
+        num_inference_steps=None,
+        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,
+    )
+
+    cache_dit.enable_cache(pipeline.transformer, cache_config=db_cache_config)
+
+    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
+        cache_dit.refresh_context(
+            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
+        )
+    return refresh_cache_context
+```
+
+## Custom Architectures: Writing Custom Enablers
+
+### When you need a custom enabler
+
+- Model has multiple block lists in one transformer (e.g., `transformer_blocks` + `single_transformer_blocks`)
+- Model has two transformers (e.g., high-noise + low-noise like Wan2.2)
+- Model uses non-standard block forward signature
+
+### Pattern 1: Multi-Block-List (LongCat-Image style)
+
+Single transformer with two block lists:
+
+```python
+import cache_dit
+from cache_dit import BlockAdapter, ForwardPattern, ParamsModifier, DBCacheConfig
+
+def enable_cache_for_your_model(pipeline, cache_config):
+    db_cache_config = DBCacheConfig(
+        num_inference_steps=None,
+        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,
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=pipeline.transformer,
+            blocks=[
+                pipeline.transformer.transformer_blocks,
+                pipeline.transformer.single_transformer_blocks,
+            ],
+            forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
+            params_modifiers=[ParamsModifier(...)],
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
+        cache_dit.refresh_context(
+            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
+        )
+    return refresh_cache_context
+```
+
+For single transformer with multiple block lists, `refresh_context` works the same as standard models — call it once on the transformer.
+
+### Pattern 2: Dual-Transformer (Wan2.2 style)
+
+Two transformers with separate configs:
+
+```python
+def enable_cache_for_dual_transformer(pipeline, cache_config):
+    db_cache_config = DBCacheConfig(...)
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=[pipeline.transformer, pipeline.transformer_2],
+            blocks=[pipeline.transformer.blocks, pipeline.transformer_2.blocks],
+            forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2],
+            params_modifiers=[
+                ParamsModifier(...),  # Config for transformer 1
+                ParamsModifier(...),  # Config for transformer 2
+            ],
+        ),
+        cache_config=db_cache_config,
+    )
+
+    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, verbose=verbose
+        )
+        cache_dit.refresh_context(
+            pipeline.transformer_2, num_inference_steps=low_steps, verbose=verbose
+        )
+    return refresh_cache_context
+```
+
+Key difference: `refresh_context` must be called on **each transformer separately** with its own step count.
+
+### Choosing the ForwardPattern
+
+| Pattern | Block forward signature | Example models |
+|---------|------------------------|----------------|
+| `Pattern_0` | `block(hidden_states, **kwargs)` → residual added inside block | Default |
+| `Pattern_1` | `block(hidden_states, **kwargs)` → returns `(hidden_states, ...)` tuple | FLUX-style single blocks |
+| `Pattern_2` | `block(hidden_states, **kwargs)` → `(hidden_states, ...)` with different residual pattern | Wan2.2 blocks |
+
+Inspect your block's `forward()` return type and residual connection pattern to choose the right one. See [Cache-DiT API Reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/) for details.
+
+## Registering Custom Enablers
+
+Add your enabler to `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,
+    "YourModelPipeline": enable_cache_for_your_model,
+}
+```
+
+The key must match `pipeline.__class__.__name__`.
+
+## Configuration Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `Fn_compute_blocks` | 1 | Number of blocks to always compute at the front |
+| `Bn_compute_blocks` | 0 | Number of blocks to always compute at the back |
+| `max_warmup_steps` | 4 | Steps to run without caching at the beginning |
+| `max_cached_steps` | — | Max total cached steps |
+| `max_continuous_cached_steps` | — | Max consecutive cached steps |
+| `residual_diff_threshold` | 0.24 | Threshold for deciding whether to cache a block |
+
+### Tuning for quality vs speed
+
+| Goal | Adjustments |
+|------|-------------|
+| **More speed, acceptable quality loss** | Higher `residual_diff_threshold` (0.24-0.4), lower `max_warmup_steps` (2-4) |
+| **Better quality, less speed** | Lower `residual_diff_threshold` (0.12-0.18), higher `max_warmup_steps` (6-8), lower `max_continuous_cached_steps` (2) |
+
+## Testing
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+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),
+)
+```
+
+CLI (online serving):
+
+```bash
+vllm serve your-model --omni --port 8098 \
+  --cache-backend cache_dit \
+  --cache-config '{"Fn_compute_blocks": 1, "Bn_compute_blocks": 0, "max_warmup_steps": 4}'
+```
+
+**Verification checklist**:
+1. Logs show "Cache-dit enabled successfully on xxx"
+2. Performance: 1.5-2x speedup vs no cache
+3. Quality: compare output with `cache_backend=None`
+
+## Excluded Models
+
+Models listed in `_NO_CACHE_ACCELERATION` in `vllm_omni/diffusion/registry.py` do not support cache-dit (e.g., `NextStep11Pipeline`, `StableDiffusionPipeline`). Check this set before attempting to enable cache-dit.
+
+## Reference Implementations
+
+| Model | Path | Notes |
+|-------|------|-------|
+| Standard DiT | `cache_dit_backend.py::enable_cache_for_dit` | Default enabler, automatic |
+| Wan2.2 | `cache_dit_backend.py::enable_cache_for_wan22` | Dual-transformer, auto-detects mode |
+| LongCat | `cache_dit_backend.py::enable_cache_for_longcat_image` | Multi-block-list |
+| BAGEL | `cache_dit_backend.py::enable_cache_for_bagel` | Complex omni model |

.claude/skills/add-diffusion-model/references/custom-model-patterns.md

@@ -0,0 +1,273 @@
+# Custom Model Patterns Reference
+
+Patterns for adding models that don't come from the standard diffusers pipeline format.
+
+## Directory Structure Comparison
+
+### Diffusers-based model (e.g., Wan2.2)
+
+```
+vllm_omni/diffusion/models/wan2_2/
+├── __init__.py                    # Exports pipeline + transformer + helpers
+├── pipeline_wan2_2.py             # Pipeline: loads components via from_pretrained()
+├── pipeline_wan2_2_i2v.py         # Variant pipeline for image-to-video
+└── wan2_2_transformer.py          # Transformer: ported from diffusers, uses Attention layer
+```
+
+The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`.
+
+### Custom model with external deps (e.g., DreamID-Omni)
+
+```
+vllm_omni/diffusion/models/dreamid_omni/
+├── __init__.py                    # Exports pipeline only
+├── pipeline_dreamid_omni.py       # Pipeline: loads ALL weights in __init__ via custom helpers
+├── fusion.py                      # Custom fusion architecture (video + audio cross-attention)
+└── wan2_2.py                      # Re-implemented Wan backbone with split API
+
+examples/offline_inference/x_to_video_audio/
+└── download_dreamid_omni.py       # Downloads weights from 3 HF repos + clones code repo
+```
+
+All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency (`dreamid_omni` package) imported with try/except.
+
+### Custom model with ported code (e.g., BAGEL)
+
+```
+vllm_omni/diffusion/models/bagel/
+├── __init__.py
+├── pipeline_bagel.py              # Pipeline: instantiates models, uses weights_sources
+├── bagel_transformer.py           # Full LLM backbone (Qwen2-MoT) ported into vllm-omni
+└── autoencoder.py                 # Custom VAE ported from original repo
+```
+
+Model code is fully ported (no external dependency). Uses `weights_sources` and `load_weights()` with custom name remapping to handle non-diffusers safetensors format.
+
+## Weight Loading Patterns
+
+### Pattern 1: Standard diffusers flow (Wan2.2, Z-Image, FLUX)
+
+```
+init → create transformer (empty) → set weights_sources → [loader calls load_weights()]
+```
+
+- `weights_sources` points to safetensors in HF subfolder (e.g., `transformer/`)
+- `load_weights()` receives `(name, tensor)` pairs from the loader
+- Name remapping handles diffusers→vllm-omni differences (QKV fusion, Sequential index removal)
+
+### Pattern 2: Custom safetensors at root (BAGEL)
+
+```
+init → create all models (empty) → set weights_sources(subfolder=None) → [loader calls load_weights()]
+```
+
+- `weights_sources` points to **root** of model directory, not a subfolder
+- Weights have non-diffusers names (e.g., `bagel.language_model.model.layers.0.self_attn.q_proj.weight`)
+- `load_weights()` does heavy name normalization
+
+```python
+self.weights_sources = [
+    DiffusersPipelineLoader.ComponentSource(
+        model_or_path=od_config.model,
+        subfolder=None,      # root directory
+        prefix="",           # no prefix stripping
+        fall_back_to_pt=False,
+    )
+]
+```
+
+### Pattern 3: Fully custom loading (DreamID-Omni)
+
+```
+init → load ALL weights eagerly via custom helpers → load_weights() = no-op
+```
+
+- No `weights_sources` attribute — standard loader finds nothing to iterate
+- Custom init functions (e.g., `init_wan_vae_2_2()`, `load_fusion_checkpoint()`) handle downloading and loading
+- `load_weights()` is `pass`
+- Weights may come from multiple HF repos in different formats (`.pth`, `.safetensors`)
+
+Use this when:
+- The original model has complex, well-tested loading code you don't want to rewrite
+- Weights span multiple HF repos
+- Weight format is non-standard (e.g., a single `.pth` file, not sharded safetensors)
+
+## model_index.json for Custom Models
+
+Standard diffusers `model_index.json`:
+```json
+{
+    "_class_name": "WanPipeline",
+    "_diffusers_version": "0.35.0.dev0",
+    "scheduler": ["diffusers", "UniPCMultistepScheduler"],
+    "transformer": ["diffusers", "WanTransformer3DModel"],
+    "vae": ["diffusers", "AutoencoderKLWan"]
+}
+```
+
+Custom model `model_index.json` (minimal):
+```json
+{
+    "_class_name": "DreamIDOmniPipeline",
+    "fusion": "DreamID-Omni/dreamid_omni.safetensors"
+}
+```
+
+The only **required** field is `_class_name` — it must match a key in `_DIFFUSION_MODELS` in `registry.py`. Other fields are model-specific and accessible via `od_config.model_config` dict.
+
+## External Dependency Management
+
+### Git clone + .pth injection (DreamID-Omni pattern)
+
+```python
+def download_dependency():
+    CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    with open(LOCK_FILE, "w") as f:
+        fcntl.flock(f, fcntl.LOCK_EX)
+        if not DEPENDENCY_DIR.exists():
+            subprocess.run([
+                "git", "clone", "--depth", "1",
+                REPO_URL, "--branch", BRANCH,
+                str(DEPENDENCY_DIR)
+            ], check=True)
+        fcntl.flock(f, fcntl.LOCK_UN)
+
+    # Add to Python path via .pth file
+    site_packages = Path(site.getsitepackages()[0])
+    pth_file = site_packages / "vllm_omni_dependency.pth"
+    pth_file.write_text(str(DEPENDENCY_DIR))
+```
+
+### Direct port (BAGEL pattern)
+
+Copy essential files from the original repo into `vllm_omni/diffusion/models/<name>/`. Adapt imports to use vllm-omni utilities. Benefits: no external dependency, no git clone step. Drawback: must maintain the ported code.
+
+## Multi-Modal Input/Output Protocols
+
+Custom models that handle images, audio, or video I/O should implement protocol classes:
+
+```python
+from vllm_omni.diffusion.models.interface import (
+    SupportImageInput,    # Model accepts image input
+    SupportAudioInput,    # Model accepts audio input
+    SupportAudioOutput,   # Model produces audio output
+)
+
+class MyPipeline(nn.Module, SupportImageInput, SupportAudioInput, SupportAudioOutput):
+    pass  # Protocol markers enable proper engine routing
+```
+
+The engine checks `isinstance(pipeline, SupportImageInput)` at startup to configure input validation and warmup behavior.
+
+## Hardcoded Config vs Config Files
+
+Diffusers models use `config.json` in each subfolder. Custom models often use:
+
+**Module-level config dicts** (DreamID-Omni):
+```python
+VIDEO_CONFIG = {
+    "patch_size": [1, 2, 2], "model_type": "ti2v",
+    "dim": 3072, "ffn_dim": 14336, "num_heads": 24, "num_layers": 30, ...
+}
+```
+
+**Loaded from custom JSON** (BAGEL):
+```python
+cfg_path = os.path.join(model_path, "config.json")
+with open(cfg_path) as f:
+    bagel_cfg = json.load(f)
+vae_cfg = bagel_cfg.get("vae_config", {})
+```
+
+## Custom Architecture Patterns
+
+### Split forward API (DreamID-Omni)
+
+When a fusion model needs to interleave blocks from two backbones:
+
+```python
+class WanModel(nn.Module):
+    def prepare_transformer_block_kwargs(self, x, t, context, ...):
+        # Patch embed, time embed, text embed, RoPE
+        return x, e, kwargs
+
+    def post_transformer_block_out(self, x, grid_sizes, e):
+        # Output projection, unpatchify
+        return output
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError  # Fusion model handles block iteration
+```
+
+The `FusionModel` then iterates blocks in lock-step:
+```python
+for video_block, audio_block in zip(self.video_model.blocks, self.audio_model.blocks):
+    video_out = video_block(video_hidden, ...)
+    audio_out = audio_block(audio_hidden, ...)
+    # Cross-attend between modalities
+    video_out = cross_attention(video_out, audio_out)
+    audio_out = cross_attention(audio_out, video_out)
+```
+
+### LLM-as-denoiser (BAGEL)
+
+When the backbone is a language model that also does diffusion:
+
+```python
+class BagelModel(nn.Module):
+    def __init__(self):
+        self.language_model = Qwen2MoTForCausalLM(config)
+        self.vit_model = SiglipVisionModel(vit_config)
+```
+
+The LLM processes both text tokens and latent image tokens in a single forward pass, using KV caching for the text portion.
+
+## Pre/Post Processing for Custom Models
+
+Custom models typically handle pre/post processing **inside `forward()`** rather than via registered functions, because the logic is tightly coupled:
+
+```python
+def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+    # Inline preprocessing
+    image = self._load_and_resize_image(req.prompts[0].get("multi_modal_data", {}).get("image"))
+    image_latent = self._vae_encode(image)
+
+    # ... denoising loop ...
+
+    # Inline postprocessing
+    pil_image = self._decode_to_pil(latents)
+    return DiffusionOutput(output=[pil_image])
+```
+
+If pre/post functions are not registered in `_DIFFUSION_PRE_PROCESS_FUNCS` / `_DIFFUSION_POST_PROCESS_FUNCS`, the engine simply skips those steps.
+
+## Download Script Template
+
+```python
+# examples/offline_inference/<name>/download_<name>.py
+from huggingface_hub import snapshot_download
+import json, os
+
+def main(output_dir):
+    # Download model weights from HF
+    snapshot_download(repo_id="org/model-weights", local_dir=os.path.join(output_dir, "weights"))
+
+    # Download additional components if from separate repos
+    snapshot_download(repo_id="org/vae-weights", local_dir=os.path.join(output_dir, "vae"),
+        allow_patterns=["*.safetensors"])
+
+    # Generate model_index.json
+    config = {"_class_name": "YourPipeline", "custom_key": "weights/model.safetensors"}
+    with open(os.path.join(output_dir, "model_index.json"), "w") as f:
+        json.dump(config, f, indent=2)
+
+    # Install external code dependency (if needed)
+    download_dependency()
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-dir", default="./your_model")
+    args = parser.parse_args()
+    main(args.output_dir)
+```

.claude/skills/add-diffusion-model/references/parallelism-patterns.md

@@ -0,0 +1,571 @@
+# Parallelism Patterns Reference
+
+## Overview
+
+vLLM-Omni supports multiple parallelism strategies for diffusion models. Each targets a different bottleneck:
+
+| Strategy | Splits | Best For | Constraint |
+|----------|--------|----------|------------|
+| Tensor Parallel (TP) | Model layers across GPUs | Latency reduction, large models | Requires fast GPU interconnect, `num_heads % tp == 0` |
+| Sequence Parallel (SP/USP) | Sequence tokens across GPUs | Long sequences (video, high-res) | Near-linear scaling |
+| CFG Parallel | Positive/negative CFG branches | Models using classifier-free guidance | Exactly 2 GPUs |
+| HSDP | Weight shards via FSDP2 | VRAM reduction | Cannot combine with TP |
+| VAE Patch Parallel | VAE decode spatial tiles | Large VAE outputs | Auto-enables tiling |
+
+**Recommended integration order**: TP → SP → CFG Parallel → HSDP
+
+**Official design docs**:
+- TP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/tensor_parallel
+- SP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/sequence_parallel
+- CFG: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cfg_parallel
+- HSDP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/hsdp
+
+---
+
+## Tensor Parallelism (TP)
+
+Replace standard `nn.Linear` with vLLM's parallel linear layers. This is the most invasive change but provides direct VRAM savings and compute speedup.
+
+### Layer replacement rules
+
+| Pattern | vLLM Layer | When to Use |
+|---------|-----------|-------------|
+| Fan-out (first in FFN) | `ColumnParallelLinear` | Projection that splits output across ranks |
+| Fan-in (second in FFN) | `RowParallelLinear` | Projection that gathers across ranks |
+| QKV projection | `QKVParallelLinear` | Fused Q/K/V for self-attention |
+| Single Q or K or V | `ColumnParallelLinear` | Separate projections (cross-attention) |
+| Attention output | `RowParallelLinear` | Output projection after attention |
+| Must not shard | `ReplicatedLinear` | Layers that must stay replicated |
+
+### MLP Block (Up-Down Pattern)
+
+```python
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear, RowParallelLinear,
+)
+
+class TPFeedForward(nn.Module):
+    def __init__(self, dim, ffn_dim):
+        super().__init__()
+        self.fc1 = ColumnParallelLinear(dim, ffn_dim, bias=False, return_bias=False)
+        self.fc2 = RowParallelLinear(
+            ffn_dim, dim, bias=False,
+            input_is_parallel=True,  # Input already sharded from fc1
+            return_bias=False,
+        )
+
+    def forward(self, x):
+        x, _ = self.fc1(x)
+        x = torch.nn.functional.gelu(x)
+        x, _ = self.fc2(x)
+        return x
+```
+
+### Attention Block (QKV-Out Pattern)
+
+```python
+from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
+from vllm_omni.diffusion.attention.layer import Attention
+
+class TPSelfAttention(nn.Module):
+    def __init__(self, dim, num_heads, num_kv_heads=None):
+        super().__init__()
+        num_kv_heads = num_kv_heads or num_heads
+        self.head_dim = dim // num_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,
+        )
+        self.to_out = RowParallelLinear(
+            dim, dim, bias=False,
+            input_is_parallel=True,
+            return_bias=False,
+        )
+        self.attn = Attention(
+            num_heads=self.to_qkv.num_heads,      # Local heads per GPU
+            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,  # Local KV heads per GPU
+        )
+
+    def forward(self, x):
+        qkv, _ = self.to_qkv(x)
+        q, k, v = qkv.split(
+            [self.to_qkv.num_heads * self.head_dim,
+             self.to_qkv.num_kv_heads * self.head_dim,
+             self.to_qkv.num_kv_heads * self.head_dim],
+            dim=-1,
+        )
+        B, S, _ = x.shape
+        q = q.view(B, S, self.to_qkv.num_heads, self.head_dim)
+        k = k.view(B, S, self.to_qkv.num_kv_heads, self.head_dim)
+        v = v.view(B, S, self.to_qkv.num_kv_heads, self.head_dim)
+        out = self.attn(q, k, v)
+        out = out.reshape(B, S, -1)
+        out, _ = self.to_out(out)
+        return out
+```
+
+### QKV Fusion in load_weights
+
+When you fuse separate Q/K/V into `QKVParallelLinear`, map diffusers' separate weight names:
+
+```python
+stacked_params_mapping = [
+    ("to_qkv", "to_q", "q"),
+    ("to_qkv", "to_k", "k"),
+    ("to_qkv", "to_v", "v"),
+]
+
+def load_weights(self, weights):
+    params = dict(self.named_parameters())
+    loaded = set()
+    for name, tensor in weights:
+        for fused_name, orig_name, shard_id in stacked_params_mapping:
+            if orig_name in name:
+                name = name.replace(orig_name, fused_name)
+                param = params[name]
+                param.weight_loader(param, tensor, shard_id)
+                loaded.add(name)
+                break
+        else:
+            if name in params:
+                param = params[name]
+                if hasattr(param, "weight_loader"):
+                    param.weight_loader(param, tensor)
+                else:
+                    default_weight_loader(param, tensor)
+                loaded.add(name)
+    return loaded
+```
+
+### RMSNorm with TP
+
+When RMSNorm sits between TP-sharded dimensions, use `DistributedRMSNorm` — it computes global RMS via all-reduce across TP ranks. See the Wan2.2 implementation for the pattern.
+
+### TP Constraints
+
+- `num_heads % tp_size == 0`
+- `num_kv_heads % tp_size == 0`
+- Use `self.to_qkv.num_heads` (local per-GPU count), not total heads, for split sizes
+
+### Testing TP
+
+```bash
+python text_to_image.py --model Your-org/your-model \
+  --tensor-parallel-size 2 --output "tp_test.png"
+```
+
+**Verify**: speedup, memory reduction proportional to TP size, quality matches single-GPU.
+
+### Reference implementations
+
+| Model | Path |
+|-------|------|
+| Z-Image | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` |
+| FLUX | `vllm_omni/diffusion/models/flux/flux_transformer.py` |
+| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` |
+
+---
+
+## Sequence Parallelism (SP / USP)
+
+SP splits sequence tokens across GPUs using Ulysses (all-to-all) or Ring (P2P) communication. It is applied non-intrusively via the `_sp_plan` dict — no changes to `forward()` logic.
+
+### Approach 1: Non-Intrusive `_sp_plan` (Recommended)
+
+The framework automatically registers hooks to shard inputs and gather outputs at `nn.Module` boundaries.
+
+#### Step 1: Identify module boundaries
+
+Find where tensors need sharding/gathering:
+
+```python
+class MyTransformer(nn.Module):
+    def __init__(self):
+        self.patch_embed = PatchEmbed()    # Before blocks
+        self.pos_embed = RoPE()            # RoPE may need splitting
+        self.blocks = nn.ModuleList([...]) # Blocks process sharded x
+        self.norm_out = LayerNorm()
+        self.proj_out = Linear()           # Gather after this
+
+    def forward(self, x):
+        x = self.patch_embed(x)
+        pos = self.pos_embed(x)
+        for block in self.blocks:
+            x = block(x, pos)
+        x = self.norm_out(x)
+        return self.proj_out(x)
+```
+
+#### Step 2: Handle inline operations
+
+`_sp_plan` hooks only work at `nn.Module` boundaries. Inline ops like `torch.cat()` must be extracted into submodules:
+
+```python
+# BAD: Inline — hooks can't intercept
+unified = torch.cat([x, cap_feats], dim=1)
+
+# GOOD: Extract into submodule
+class UnifiedPrepare(nn.Module):
+    def forward(self, x, cap_feats):
+        return torch.cat([x, cap_feats], dim=1)
+
+self.unified_prepare = UnifiedPrepare()
+unified = self.unified_prepare(x, cap_feats)
+```
+
+Common cases: `torch.cat()`, `pad_sequence()`, `tensor.reshape()`, complex preprocessing.
+
+#### Step 3: Write `_sp_plan`
+
+**Pattern 1: Shard at first block, gather at output** (most common)
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput, SequenceParallelOutput,
+)
+
+class StandardTransformer(nn.Module):
+    _sp_plan = {
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+**Pattern 2: Shard RoPE outputs separately**
+
+```python
+class TransformerWithRoPE(nn.Module):
+    _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),
+        },
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+**Pattern 3: Dual-stream (shard image, replicate text)**
+
+```python
+class DualStreamTransformer(nn.Module):
+    _sp_plan = {
+        "rope_preparer": {
+            2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
+            3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
+        },
+        "transformer_blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+### API Reference
+
+**SequenceParallelInput**:
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `split_dim` | int | Dimension to split (usually 1 for sequence) |
+| `expected_dims` | int/None | Expected tensor rank for validation |
+| `split_output` | bool | `False`: shard input params; `True`: shard output tensors |
+| `auto_pad` | bool | Auto-pad if sequence not divisible by world_size |
+
+**SequenceParallelOutput**:
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `gather_dim` | int | Dimension to gather (usually 1 for sequence) |
+| `expected_dims` | int/None | Expected tensor rank for validation |
+
+**Module naming**:
+
+| Key | Meaning |
+|-----|---------|
+| `"blocks.0"` | First element of ModuleList |
+| `"blocks.*"` | All elements of ModuleList |
+| `"rope"` | Named submodule |
+
+**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 (Complex Cases)
+
+For dynamic sharding logic that can't 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)
+    for block in self.blocks:
+        hidden_states = block(hidden_states)
+    if self.parallel_config.sequence_parallel_size > 1:
+        hidden_states = sp_gather(hidden_states, dim=1)
+    return hidden_states
+```
+
+Use intrusive modification as a last resort — `_sp_plan` is preferred for maintainability.
+
+### UAA Mode (Experimental)
+
+`ulysses_mode="advanced_uaa"` handles arbitrary sequence lengths and head counts that aren't divisible by `ulysses_degree`. Uses variable all-to-all split sizes and temporary head padding.
+
+### Combining SP methods
+
+Ulysses and Ring can be combined: `ulysses_degree × ring_degree = total SP GPUs`.
+
+```python
+DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)  # 4 GPUs total
+```
+
+### Testing SP
+
+```bash
+# Offline
+python text_to_image.py --model Your-model --ulysses-degree 2
+
+# Online serving
+vllm serve Your-model --omni --usp 2
+```
+
+### Reference implementations
+
+| Model | Path |
+|-------|------|
+| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` |
+| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` |
+| Z-Image | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` |
+
+---
+
+## CFG Parallelism
+
+Distributes positive/negative Classifier-Free Guidance branches across 2 GPUs.
+
+### Implementation
+
+Inherit `CFGParallelMixin` and implement `diffuse()`:
+
+```python
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+
+class YourPipeline(nn.Module, CFGParallelMixin):
+    def diffuse(self, latents, timesteps, prompt_embeds, negative_embeds,
+                do_true_cfg, true_cfg_scale, **kwargs):
+        for i, t in enumerate(timesteps):
+            positive_kwargs = {
+                "hidden_states": latents,
+                "encoder_hidden_states": prompt_embeds,
+                "timestep": t,
+            }
+            negative_kwargs = {
+                "hidden_states": latents,
+                "encoder_hidden_states": negative_embeds,
+                "timestep": t,
+            } if do_true_cfg else None
+
+            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,
+            )
+            latents = self.scheduler_step_maybe_with_cfg(
+                noise_pred, t, latents, do_true_cfg
+            )
+        return latents
+```
+
+### Customization hooks
+
+| Method | Override when |
+|--------|-------------|
+| `predict_noise()` | Non-standard transformer call (e.g., dual-transformer like Wan2.2) |
+| `cfg_normalize_function()` | Custom normalization (e.g., LongCat with clamping) |
+| `combine_cfg_noise()` | Multi-output models (e.g., video + audio: CFG on video, positive-only on audio) |
+
+**Custom predict_noise** (Wan2.2 — selects active transformer):
+
+```python
+def predict_noise(self, current_model=None, **kwargs):
+    if current_model is None:
+        current_model = self.transformer
+    return current_model(**kwargs)[0]
+```
+
+**Custom combine_cfg_noise** (multi-output):
+
+```python
+def combine_cfg_noise(self, positive_pred, negative_pred, scale, normalize):
+    video_pos, audio_pos = positive_pred
+    video_neg, audio_neg = negative_pred
+    video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize)
+    return (video_combined, audio_pos)
+```
+
+### Composite scheduler for multi-output
+
+When each output has its own schedule:
+
+```python
+class VideoAudioScheduler:
+    def __init__(self, video_scheduler, audio_scheduler):
+        self.video_scheduler = video_scheduler
+        self.audio_scheduler = audio_scheduler
+
+    def step(self, noise_pred, t, latents, return_dict=False, generator=None):
+        video_out = self.video_scheduler.step(
+            noise_pred[0], t[0], latents[0], return_dict=False, generator=generator
+        )[0]
+        audio_out = self.audio_scheduler.step(
+            noise_pred[1], t[1], latents[1], return_dict=False, generator=generator
+        )[0]
+        return ((video_out, audio_out),)
+```
+
+### Testing CFG Parallel
+
+```bash
+python text_to_image.py --model Your-model \
+  --cfg-parallel-size 2 --cfg-scale 4.0 \
+  --negative-prompt "ugly, unclear"
+```
+
+**Constraint**: `guidance_scale > 1.0` and negative prompt must be provided.
+
+### Reference implementations
+
+| Model | Path |
+|-------|------|
+| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/cfg_parallel.py` |
+| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py` |
+| Mixin base | `vllm_omni/diffusion/distributed/cfg_parallel.py` |
+
+---
+
+## HSDP (Hybrid Sharded Data Parallel)
+
+Shards model weights across GPUs using PyTorch FSDP2. Reduces per-GPU VRAM without changing computation.
+
+### Implementation
+
+Add `_hsdp_shard_conditions` to the transformer class:
+
+```python
+class YourTransformer(nn.Module):
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block]
+```
+
+For MoE models, add additional conditions:
+
+```python
+class MoETransformer(nn.Module):
+    @staticmethod
+    def _is_transformer_block(name, module):
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    @staticmethod
+    def _is_moe_expert(name, module):
+        return "experts" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block, _is_moe_expert]
+```
+
+A module is sharded if **any** condition returns `True`.
+
+### Constraints
+
+- Cannot combine with Tensor Parallelism
+- For standalone HSDP (no other parallelism), `hsdp_shard_size` must be specified explicitly
+- Can combine with SP: HSDP reduces memory while SP distributes sequence
+
+### Testing HSDP
+
+```python
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+parallel_config = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=8)
+omni = Omni(model="your-model", parallel_config=parallel_config)
+```
+
+Or CLI:
+
+```bash
+vllm serve Your-model --omni --use-hsdp
+```
+
+**Verify**: logs show "HSDP Inference: replicate_size=..., shard_size=..." and "Sharded N modules + root". Check VRAM reduction.
+
+### Reference implementations
+
+| Model | Path |
+|-------|------|
+| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` |
+| HSDP Core | `vllm_omni/diffusion/distributed/hsdp.py` |
+
+---
+
+## VAE Patch Parallelism
+
+Shards VAE decode spatially across ranks using tiling:
+
+```bash
+python text_to_image.py --model Your-model --vae-patch-parallel-size 4
+```
+
+Auto-enables `--vae-use-tiling`. Uses `DistributedAutoencoderKLWan` or similar distributed VAE. Set `vae_patch_parallel_size` in `DiffusionParallelConfig`.
+
+---
+
+## Combining Parallelism Methods
+
+Common multi-GPU recipes:
+
+```bash
+# 4 GPUs: CFG (2) × Ulysses (2)
+python text_to_image.py --model Qwen/Qwen-Image \
+  --cfg-parallel-size 2 --ulysses-degree 2
+
+# 8 GPUs: Ulysses (4) × Ring (2) + VAE patch (8)
+python text_to_video.py --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --ulysses-degree 4 --ring-degree 2 --vae-patch-parallel-size 8
+
+# 2 GPUs: HSDP + Ulysses (cannot combine HSDP with TP)
+vllm serve Your-model --omni --use-hsdp --usp 2
+```
+
+## Discovering Parallelism Support
+
+Check which parallelism methods a model supports:
+
+| Check | How |
+|-------|-----|
+| **Ulysses / Ring SP** | Transformer defines `_sp_plan`. Search: `grep -r '_sp_plan' vllm_omni/diffusion/models/` |
+| **CFG Parallel** | Pipeline inherits `CFGParallelMixin`. Search: `grep -r 'CFGParallelMixin' vllm_omni/diffusion/models/` |
+| **TP** | Uses `ColumnParallelLinear` / `QKVParallelLinear`. Search: `grep -r 'ParallelLinear\|QKVParallel' vllm_omni/diffusion/models/<model>/` |
+| **HSDP** | Transformer defines `_hsdp_shard_conditions`. Search: `grep -r '_hsdp_shard_conditions' vllm_omni/diffusion/models/` |
+
+The canonical per-model support table is in `docs/user_guide/diffusion/parallelism_acceleration.md`.

.claude/skills/add-diffusion-model/references/transformer-adaptation.md

@@ -0,0 +1,218 @@
+# Transformer Adaptation Reference
+
+## Adapting a Diffusers Transformer to vLLM-Omni
+
+### Step-by-step Checklist
+
+1. Copy the transformer class from diffusers source
+2. Remove all mixin classes — inherit only from `nn.Module`
+3. Replace attention dispatch with `vllm_omni.diffusion.attention.layer.Attention`
+4. Replace logger with `vllm.logger.init_logger`
+5. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`
+6. Remove training-only code (gradient checkpointing, dropout)
+7. Add `load_weights()` method for weight loading from safetensors
+8. Add class-level attributes for acceleration features
+
+### Mixin Removal
+
+Remove these diffusers mixins (and their imports):
+
+```python
+# Remove all of these:
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention_processor import AttentionModuleMixin
+from diffusers.loaders import PeftAdapterMixin, FromOriginalModelMixin
+
+# Replace:
+class MyTransformer(ModelMixin, ConfigMixin, AttentionModuleMixin):
+# With:
+class MyTransformer(nn.Module):
+```
+
+Also remove `@register_to_config` decorators from `__init__`.
+
+### Attention Replacement
+
+The vLLM-Omni `Attention` layer wraps backend selection (FlashAttention, SDPA, SageAttn, etc.) and supports sequence parallelism hooks.
+
+**QKV tensor shape must be `[batch, seq_len, num_heads, head_dim]`.**
+
+#### Self-Attention Pattern
+
+```python
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+
+class SelfAttentionBlock(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.to_q = nn.Linear(dim, dim)
+        self.to_k = nn.Linear(dim, dim)
+        self.to_v = nn.Linear(dim, dim)
+        self.to_out = nn.Linear(dim, dim)
+
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim ** 0.5),
+            causal=False,
+            num_kv_heads=num_heads,
+        )
+
+    def forward(self, x, attn_mask=None):
+        B, S, _ = x.shape
+        q = self.to_q(x).view(B, S, self.num_heads, self.head_dim)
+        k = self.to_k(x).view(B, S, self.num_heads, self.head_dim)
+        v = self.to_v(x).view(B, S, self.num_heads, self.head_dim)
+
+        attn_metadata = AttentionMetadata(attn_mask=attn_mask)
+        out = self.attn(q, k, v, attn_metadata=attn_metadata)
+        out = out.reshape(B, S, -1)
+        return self.to_out(out)
+```
+
+#### Fused QKV with TP (Advanced)
+
+For tensor parallelism, use vLLM's parallel linear layers:
+
+```python
+from vllm.model_executor.layers.linear import (
+    QKVParallelLinear, RowParallelLinear
+)
+
+class TPSelfAttention(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            total_num_kv_heads=num_heads,
+        )
+        self.to_out = RowParallelLinear(dim, dim)
+
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim ** 0.5),
+            causal=False,
+            num_kv_heads=num_heads,
+        )
+```
+
+### Logger Replacement
+
+```python
+# Replace:
+from diffusers.utils import logging
+logger = logging.get_logger(__name__)
+
+# With:
+from vllm.logger import init_logger
+logger = init_logger(__name__)
+```
+
+### Custom Layers from vLLM-Omni
+
+Available utility layers:
+
+```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
+```
+
+### Config Support
+
+```python
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+class MyTransformer(nn.Module):
+    def __init__(self, *, od_config=None, num_layers=28, hidden_size=3072, **kwargs):
+        super().__init__()
+        self.od_config = od_config
+        self.parallel_config = od_config.parallel_config if od_config else None
+        # ... build layers
+```
+
+The transformer config values come from `model_index.json` → `config.json` in the transformer subfolder. The pipeline uses `get_transformer_config_kwargs(od_config.tf_model_config, TransformerClass)` to filter config keys to match the `__init__` signature.
+
+### Weight Loading
+
+The `load_weights` method receives an iterable of `(name, tensor)` from safetensors files, with the prefix (e.g., `"transformer."`) already stripped by the loader.
+
+```python
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+class MyTransformer(nn.Module):
+    def load_weights(self, weights):
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            # Optional: remap names from diffusers to vllm-omni naming
+            # e.g., "ff.net.0.proj" -> "ff.net_0.proj"
+
+            if name in params:
+                param = params[name]
+                if hasattr(param, "weight_loader"):
+                    param.weight_loader(param, tensor)
+                else:
+                    default_weight_loader(param, tensor)
+                loaded.add(name)
+        return loaded
+```
+
+#### QKV Fusion in load_weights
+
+If you fused separate Q/K/V into a `QKVParallelLinear`, you need to map diffusers' separate weight names:
+
+```python
+stacked_params_mapping = [
+    ("to_qkv", "to_q", "q"),
+    ("to_qkv", "to_k", "k"),
+    ("to_qkv", "to_v", "v"),
+]
+
+def load_weights(self, weights):
+    params = dict(self.named_parameters())
+    loaded = set()
+    for name, tensor in weights:
+        for fused_name, orig_name, shard_id in stacked_params_mapping:
+            if orig_name in name:
+                name = name.replace(orig_name, fused_name)
+                param = params[name]
+                param.weight_loader(param, tensor, shard_id)
+                loaded.add(name)
+                break
+        else:
+            # Normal loading
+            ...
+    return loaded
+```
+
+### Class-Level Attributes for Features
+
+```python
+class MyTransformer(nn.Module):
+    # torch.compile: list block class names that repeat and can be compiled
+    _repeated_blocks = ["MyTransformerBlock"]
+
+    # CPU offload: attribute name of the nn.ModuleList containing blocks
+    _layerwise_offload_blocks_attr = "blocks"
+
+    # LoRA: mapping of fused param names to original param names
+    packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
+
+    # Sequence parallelism plan (advanced — add after basic impl works)
+    _sp_plan = {
+        "blocks.0": SequenceParallelInput(split_dim=1),
+        "proj_out": SequenceParallelOutput(gather_dim=1),
+    }
+```

.claude/skills/add-diffusion-model/references/troubleshooting.md

@@ -0,0 +1,178 @@
+# Troubleshooting Reference
+
+## Common Errors When Adding a Diffusion Model
+
+### ImportError / ModuleNotFoundError
+
+**Cause**: Missing or incorrect registration.
+
+**Fix checklist**:
+1. Model registered in `vllm_omni/diffusion/registry.py` `_DIFFUSION_MODELS` dict
+2. `__init__.py` exports the pipeline class
+3. Pipeline file exists at the correct path: `vllm_omni/diffusion/models/{folder}/{file}.py`
+4. Class name in registry matches the actual class name in the file
+
+### Shape Mismatch in Attention
+
+**Symptom**: `RuntimeError: shape mismatch` or `expected 4D tensor`
+
+**Cause**: QKV tensors not reshaped to `[batch, seq_len, num_heads, head_dim]`.
+
+**Fix**: Before calling `self.attn(q, k, v, ...)`, ensure:
+```python
+q = q.view(batch, seq_len, self.num_heads, self.head_dim)
+k = k.view(batch, kv_seq_len, self.num_kv_heads, self.head_dim)
+v = v.view(batch, kv_seq_len, self.num_kv_heads, self.head_dim)
+```
+
+After attention, reshape back:
+```python
+out = out.reshape(batch, seq_len, -1)
+```
+
+### Weight Loading Failures
+
+**Symptom**: `RuntimeError: size mismatch for parameter ...` or missing keys
+
+**Debugging**:
+1. Print diffusers weight names: `safetensors.safe_open(path, "pt").keys()`
+2. Print model parameter names: `dict(model.named_parameters()).keys()`
+3. Compare and add name remappings in `load_weights()`
+
+**Common remappings needed**:
+- `ff.net.0.proj` → `ff.net_0.proj` (PyTorch Sequential indexing)
+- `.to_out.0.` → `.to_out.` (Sequential unwrapping)
+- `scale_shift_table` → moved to a wrapper module
+
+### Black/Blank/Noisy Output
+
+**Possible causes**:
+1. **Wrong latent normalization**: Check VAE expects latents scaled by `vae.config.scaling_factor`
+2. **Wrong scheduler**: Using the wrong scheduler class or wrong `flow_shift`
+3. **Missing CFG**: Some models require `guidance_scale > 1.0` with negative prompt
+4. **Wrong timestep format**: Some schedulers expect float, others expect int/long
+5. **Missing post-processing**: Raw VAE output may need denormalization
+
+**Quick test**: Run with diffusers directly using the same seed and compare latents at each step.
+
+### OOM (Out of Memory)
+
+**Solutions** (in order of preference):
+1. `--enforce-eager` to disable torch.compile (saves compile memory)
+2. `--enable-cpu-offload` for model-level offload
+3. `--enable-layerwise-offload` for block-level offload (better for large models)
+4. `--vae-use-slicing --vae-use-tiling` for VAE memory reduction
+5. Reduce resolution: `--height 480 --width 832`
+6. Use TP: `--tensor-parallel-size 2`
+
+### Different Output vs Diffusers Reference
+
+**Common causes**:
+1. **Attention backend difference**: FlashAttention vs SDPA may produce slightly different results. Set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` to match diffusers
+2. **Float precision**: vLLM-Omni may use bfloat16 where diffusers uses float32 for some operations
+3. **Missing normalization**: Check all LayerNorm/RMSNorm are preserved
+4. **Scheduler rounding**: Some schedulers have numerical sensitivity
+
+### Tensor Parallel Errors
+
+**Symptom**: `AssertionError: not divisible` or incorrect output with TP>1
+
+**Fix**:
+1. Verify `num_heads % tp_size == 0` and `num_kv_heads % tp_size == 0`
+2. Ensure `ColumnParallelLinear` / `RowParallelLinear` are used correctly
+3. Check that norms between parallel layers use distributed norm if needed
+4. Verify `load_weights` handles TP sharding for norm weights
+5. Use `self.to_qkv.num_heads` (local heads per GPU) for QKV split sizes, not total heads
+
+**Missing `input_is_parallel=True`**:
+
+`RowParallelLinear` expects sharded input from `ColumnParallelLinear`:
+```python
+self.w1 = ColumnParallelLinear(dim, hidden_dim, return_bias=False)
+self.w2 = RowParallelLinear(hidden_dim, dim, input_is_parallel=True, return_bias=False)
+```
+
+### Sequence Parallel Errors
+
+**Symptom**: Incorrect output or crashes with `--ulysses-degree N` or `--usp N`
+
+**Possible causes**:
+1. **Inline operations between shard/gather points**: `torch.cat()`, `pad_sequence()` etc. not at `nn.Module` boundaries. Fix: extract into submodule.
+2. **Wrong `split_dim`**: Check the tensor shape at the shard point. Sequence dimension is typically `dim=1` for `[B, S, D]` tensors.
+3. **RoPE not sharded**: If RoPE is computed separately, add it to `_sp_plan` with `split_output=True`.
+4. **Sequence not divisible by SP degree**: Use `auto_pad=True` in `SequenceParallelInput` or switch to `ulysses_mode="advanced_uaa"`.
+
+**Debugging**: Add `expected_dims=N` to `SequenceParallelInput`/`Output` for shape validation at runtime.
+
+### CFG Parallel Errors
+
+**Symptom**: CFG parallel not activating, no speedup
+
+**Fix checklist**:
+1. Pipeline inherits `CFGParallelMixin`
+2. `guidance_scale > 1.0`
+3. Negative prompt provided (even if empty string)
+4. `--cfg-parallel-size 2` specified
+5. `diffuse()` method calls `predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()`
+
+**Symptom**: Different output with CFG parallel vs sequential
+
+**Possible cause**: Non-deterministic scheduler. Fix: pass `generator=torch.Generator(device).manual_seed(seed)` to `scheduler_step_maybe_with_cfg()`.
+
+### HSDP Errors
+
+**Symptom**: HSDP not activating or errors during weight loading
+
+**Fix checklist**:
+1. Transformer defines `_hsdp_shard_conditions` class attribute
+2. Shard condition functions return `True` for correct modules (test with `model.named_modules()`)
+3. Not combining with TP (HSDP and TP are incompatible)
+4. For standalone HSDP, `hsdp_shard_size` is specified explicitly
+
+**Verify**: Check logs for "HSDP Inference: replicate_size=..., shard_size=..." and "Sharded N modules + root".
+
+### Cache-DiT Not Applied
+
+**Symptom**: No speedup, no cache-related log messages
+
+**Fix checklist**:
+1. Model not in `_NO_CACHE_ACCELERATION` in `registry.py`
+2. Pipeline class name matches `CUSTOM_DIT_ENABLERS` key (if using custom enabler)
+3. `cache_backend="cache_dit"` specified
+4. Check logs for "Cache-dit enabled successfully on xxx"
+
+**Verify pipeline name**: `print(pipeline.__class__.__name__)` — must match registry key.
+
+### Cache-DiT Quality Degradation
+
+**Symptom**: Artifacts or lower quality with cache-dit
+
+**Fix**: Reduce aggressiveness:
+```python
+cache_config={
+    "residual_diff_threshold": 0.12,      # Lower from 0.24
+    "max_warmup_steps": 6,                # Increase from 4
+    "max_continuous_cached_steps": 2,      # Reduce if higher
+}
+```
+
+If quality is still poor, the model may need a custom enabler with per-block-list `ParamsModifier` tuning.
+
+### Model Not Detected / Wrong Pipeline Class
+
+**Symptom**: `ValueError: Model class ... not found in diffusion model registry`
+
+**Cause**: The model's `model_index.json` has a `_class_name` for the pipeline that doesn't match registry keys.
+
+**Fix**: The registry key must match the diffusers pipeline class name from `model_index.json`. If using a different name, map it in the registry:
+```python
+"DiffusersPipelineClassName": ("your_folder", "your_file", "YourVllmClassName"),
+```
+
+## Debugging Workflow
+
+1. **Add verbose logging**: Use `logger.info()` to print tensor shapes at each stage
+2. **Compare step-by-step**: Run diffusers and vllm-omni side by side, comparing tensors after each major operation
+3. **Use small configs**: Reduce `num_inference_steps=2`, small resolution for fast iteration
+4. **Test transformer isolation**: Feed the same input to both diffusers and vllm-omni transformers, compare outputs
+5. **Binary search for bugs**: Comment out blocks/layers to isolate where divergence starts

.claude/skills/add-tts-model/SKILL.md

@@ -0,0 +1,284 @@
+---
+name: add-tts-model
+description: "Integrate a new text-to-speech model into vLLM-Omni from HuggingFace reference implementation through production-ready serving with streaming and CUDA graph acceleration. Use when adding a new TTS model, wiring stage separation for speech synthesis, enabling online voice generation serving, debugging TTS integration behavior, or building audio output pipelines."
+---
+
+# TTS Model Integration Workflow
+
+## Overview
+
+```
+HF Reference -> Stage Separation -> Online Serving -> Async Chunk -> CUDA Graph
+   (Phase 1)      (Phase 2)          (Phase 3)        (Phase 4)     (Phase 5)
+```
+
+## Phase 1: HuggingFace Reference
+
+**Goal**: Understand the reference implementation and verify it produces correct audio.
+
+### Steps
+
+1. **Run the reference model** end-to-end using the official HuggingFace / GitHub code
+2. **Document the architecture**:
+   - What are the sub-models? (AR decoder, codec decoder, vocoder, etc.)
+   - What is the token vocabulary? (semantic codes, RVQ codebooks, special tokens)
+   - What is the output format? (sample rate, channels, codec type)
+3. **Capture reference outputs** for comparison during integration
+4. **Identify the config structure**: `config.json` fields, `model_type`, sub-model configs
+
+### Key Questions
+
+- How many codebooks? What are the codebook sizes?
+- What special tokens exist? (`<|voice|>`, `<|audio_start|>`, `<|im_end|>`, etc.)
+- What is the token-to-ID mapping for codec codes?
+- What is the hop length / frame rate of the codec?
+- Does the model support voice cloning? How? (reference audio encoding, speaker embeddings, etc.)
+
+### Deliverables
+
+- Working reference script that produces audio
+- Architecture diagram / notes
+- Token vocabulary mapping
+- Reference audio samples for regression testing
+
+## Phase 2: Stage Separation (Offline Inference)
+
+**Goal**: Split the model into vLLM-Omni stages and get offline inference working.
+
+### Steps
+
+1. **Register the model** in `vllm_omni/model_executor/models/registry.py`
+2. **Create config classes** (`configuration_<model>.py`) with `model_type` registration
+3. **Implement Stage 0** (AR model):
+   - Subclass appropriate base (e.g., wrap Qwen3 decoder layers)
+   - Implement `forward()` for autoregressive token generation
+   - Handle special token logic (start/stop tokens, codec token mapping)
+   - If dual-AR (like Fish Speech), implement Fast AR as a nested module
+4. **Implement Stage 1** (Decoder):
+   - Load codec weights (may need lazy loading from separate checkpoint)
+   - Implement `forward()`: codec codes -> audio waveform
+   - Return `OmniOutput` with `multimodal_outputs`
+5. **Create stage config YAML** defining both stages, memory allocation, and model paths
+6. **Create stage input processor** for prompt building
+7. **Write end2end.py** test script
+
+### Critical Parameters to Get Right
+
+| Parameter | Impact if Wrong |
+|-----------|----------------|
+| Hop length | Audio duration wrong, streaming noise |
+| Token ID mapping | Garbage codes -> noise output |
+| Codebook count/size | Shape mismatch crashes |
+| Stop token | Generation never stops or stops too early |
+| dtype / autocast | Numerical issues, silent quality degradation |
+| Repetition penalty | Must match reference (often 1.0 for TTS) |
+
+### Debugging Priority (from experience)
+
+When audio output is wrong, check in this order:
+
+1. **RoPE / attention**: Are position encodings correct? Is the attention mask right?
+2. **Normalization**: RMSNorm epsilon, layer norm placement (pre vs post)
+3. **Hop length**: Product of all upsample rates in the codec decoder
+4. **Token mapping**: Are codec IDs correctly offset from the vocabulary base?
+5. **Sampling parameters**: Temperature, top_k, top_p, repetition_penalty
+6. **Tensor layout**: Codebook-major vs frame-major ordering
+7. **dtype**: Float32 for codec decoders (autocast can corrupt audio)
+
+### Deliverables
+
+- Model files in `vllm_omni/model_executor/models/<model_name>/`
+- Stage config YAML
+- Working `end2end.py` with correct audio output
+- README.md in the example directory
+
+## Phase 3: Online Serving
+
+**Goal**: Expose the model via `/v1/audio/speech` API endpoint.
+
+### Steps
+
+1. **Register in `serving_speech.py`**:
+   - Add model stage name to `_TTS_MODEL_STAGES` set
+   - Add model detection flag (e.g., `_is_fish_speech`)
+   - Implement prompt builder method (e.g., `_build_fish_speech_prompt()`)
+2. **Handle model-specific parameters**:
+   - Voice cloning: `ref_audio` encoding and prompt injection
+   - `max_new_tokens` override in sampling params
+   - Model-specific default values
+3. **Create client scripts**: `speech_client.py`, `run_server.sh`
+4. **Test all response formats**: wav, mp3, flac, pcm
+5. **Add Gradio demo**: Interactive web UI with streaming support
+
+### Voice Cloning Pattern
+
+```python
+import base64
+from pathlib import Path
+
+def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list:
+    """Build prompt with reference audio for voice cloning in serving_speech.py."""
+    audio_bytes = Path(ref_audio_path).read_bytes()
+    codes = codec.encode(audio_bytes)  # Encode on CPU using model's codec (e.g., DAC)
+    token_ids = [code + codec.vocab_offset for code in codes.flatten().tolist()]
+    return [
+        {"role": "system", "content": f"<|voice|>{''.join(chr(t) for t in token_ids)}"},
+        {"role": "user", "content": text},
+    ]
+```
+
+### Deliverables
+
+- Updated `serving_speech.py` with model-specific prompt builder
+- Client scripts and server launcher
+- Gradio demo with streaming and voice cloning UI
+- Documentation (offline + online serving docs)
+
+## Phase 4: Async Chunk (Streaming)
+
+**Goal**: Enable inter-stage streaming so audio chunks are produced while AR generation continues.
+
+### Steps
+
+1. **Update stage config YAML**:
+   ```yaml
+   async_chunk: true
+   codec_chunk_frames: 25      # frames per chunk
+   codec_left_context_frames: 25  # overlap for smooth boundaries
+   ```
+2. **Implement chunk handling in Stage 1**:
+   - Accept partial input (chunk of codec codes)
+   - Handle left context for smooth audio boundaries
+   - Return partial audio in `OmniOutput`
+3. **Test streaming**:
+   - Verify audio quality matches non-streaming output
+   - Check for artifacts at chunk boundaries
+   - Measure TTFA (time to first audio)
+4. **Update online serving** to support `stream=true` with PCM output
+
+### Streaming Architecture
+
+```
+Stage 0 (AR)                    Stage 1 (Decoder)
+  |                                |
+  |-- chunk 0 (25 frames) ------> decode -> audio chunk 0 -> client
+  |-- chunk 1 (25 frames) ------> decode -> audio chunk 1 -> client
+  |-- chunk 2 (25 frames) ------> decode -> audio chunk 2 -> client
+  ...
+```
+
+### Key Considerations
+
+- **Left context overlap**: Prevents audible artifacts at chunk boundaries
+- **Hop length matters**: `context_audio_samples = context_frames * hop_length`
+- **First chunk latency**: Can use larger initial chunk for better quality, then smaller chunks
+
+### Deliverables
+
+- Updated stage config with async_chunk enabled
+- Smooth streaming audio without boundary artifacts
+- TTFA metrics
+
+## Phase 5: CUDA Graph Acceleration
+
+**Goal**: Capture the AR loop as a CUDA graph for significant speedup.
+
+### Steps
+
+1. **Identify the hot loop**: The AR decoding loop that runs N steps per token
+2. **Create static buffers**:
+   - KV caches with fixed max sequence length
+   - Pre-built causal masks and position tensors per step
+   - Static input/output tensors
+3. **Implement graph capture**:
+   - Warm up with real data
+   - Capture the forward pass
+   - Replay with updated inputs
+4. **Handle constraints**:
+   - Use `torch.argmax` instead of `torch.multinomial` (graph-safe)
+   - Fixed batch size (fall back to eager for other sizes)
+   - No dynamic control flow inside the graph
+
+### Example: Code Predictor CUDA Graph (Qwen3-TTS)
+
+```python
+import torch
+
+class CodePredictorGraph:
+    """Captures the 16-step code predictor AR loop as a single CUDA graph."""
+
+    def setup_graph(self, device: torch.device, kv_heads: int = 4, head_dim: int = 64):
+        self.num_steps = 16
+        self.kv_cache = torch.zeros(1, kv_heads, self.num_steps, head_dim, device=device)
+        self.positions = torch.arange(self.num_steps, device=device)
+        self.causal_mask = torch.tril(torch.ones(self.num_steps, self.num_steps, device=device))
+        self.input_buf = torch.zeros(1, 1, kv_heads * head_dim, device=device)
+        self.output_buf = torch.zeros(1, self.num_steps, device=device, dtype=torch.long)
+        # Warm up, then: self.graph = torch.cuda.CUDAGraph(); self.graph.capture(...)
+
+    def run_graph(self, initial_input: torch.Tensor) -> torch.Tensor:
+        self.input_buf.copy_(initial_input)
+        self.graph.replay()
+        return self.output_buf.clone()
+```
+
+### Performance Expectations
+
+Based on Qwen3-TTS code predictor experience:
+- **3-5x speedup** for the graphed component
+- Only effective for fixed batch sizes (typically batch_size=1)
+- Falls back to eager mode for unsupported configurations
+
+### Deliverables
+
+- CUDA graph implementation for the AR hot loop
+- Benchmark script comparing eager vs graph performance
+- Documentation of constraints and fallback behavior
+
+## Integration Checklist
+
+Use this checklist when integrating a new TTS model:
+
+### Phase 1: HF Reference
+- [ ] Reference model runs and produces correct audio
+- [ ] Architecture documented (stages, codebooks, tokens, sample rate)
+- [ ] Reference audio samples saved for comparison
+
+### Phase 2: Stage Separation
+- [ ] Model registered in `registry.py`
+- [ ] Config classes created with `model_type` registration
+- [ ] Stage 0 (AR) implemented and generates correct tokens
+- [ ] Stage 1 (Decoder) produces correct audio from tokens
+- [ ] Stage config YAML created
+- [ ] `end2end.py` produces audio matching reference quality
+- [ ] README.md written
+
+### Phase 3: Online Serving
+- [ ] Model added to `serving_speech.py`
+- [ ] Prompt builder handles text input correctly
+- [ ] Voice cloning works (if supported)
+- [ ] All response formats work (wav, mp3, flac, pcm)
+- [ ] Client scripts and server launcher created
+- [ ] Gradio demo working
+- [ ] Documentation added (offline + online docs, nav, supported models)
+
+### Phase 4: Async Chunk
+- [ ] Stage config updated with `async_chunk: true`
+- [ ] Stage 1 handles partial chunks correctly
+- [ ] No audio artifacts at chunk boundaries
+- [ ] Streaming via API (`stream=true`) works
+- [ ] TTFA measured and acceptable
+
+### Phase 5: CUDA Graph
+- [ ] Hot loop identified and profiled
+- [ ] Static buffers allocated
+- [ ] Graph captured and replays correctly
+- [ ] Benchmark shows meaningful speedup
+- [ ] Fallback to eager works for unsupported configs
+
+## References
+
+- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) -- supported models and usage
+- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) -- complete example of Phases 1-3
+- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) -- complete example of all 5 phases
+- [Adding a TTS model (developer guide)](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/model/adding_tts_model.md)

.claude/skills/readme.md

@@ -0,0 +1,34 @@
+# Claude Skills for vLLM-Omni
+
+This directory contains Claude Code skills maintained for the `vllm-omni`
+repository. These skills capture repeatable workflows for common contributor
+tasks such as model integration, pull request review, and release note
+generation.
+
+## Directory Structure
+
+Each skill lives in its own directory under `.claude/skills/`. A skill may
+include:
+
+- `SKILL.md`: the main workflow and operating instructions
+- `references/`: focused reference material used by the skill
+- `scripts/`: small helper scripts used by the skill
+
+## Available Skills
+
+- `add-diffusion-model`: guides integration of a new diffusion model into
+  `vllm-omni`
+- `add-omni-model`: covers addition of new omni-modality model support
+- `add-tts-model`: covers integration of new TTS models and related serving
+  workflows
+- `generate-release-note`: helps prepare release notes for repository changes
+- `review-pr`: provides a structured workflow for reviewing pull requests
+
+## Maintenance Guidelines
+
+- Keep skill names short and task-oriented.
+- Prefer repository-local paths, commands, and examples.
+- Avoid hardcoding fast-changing support matrices unless the skill is actively
+  maintained alongside those changes.
+- Treat skills as contributor tooling: optimize for clarity, actionability, and
+  low maintenance overhead.

.gitignore

@@ -158,7 +158,19 @@ cython_debug/
 
 # Claude
 CLAUDE.md
-.claude/
+/.claude/*
+!.claude/skills/
+!.claude/skills/readme.md
+!.claude/skills/add-diffusion-model/
+!.claude/skills/add-diffusion-model/SKILL.md
+!.claude/skills/add-diffusion-model/references/
+!.claude/skills/add-diffusion-model/references/*.md
+!.claude/skills/add-tts-model/
+!.claude/skills/add-tts-model/SKILL.md
+!.claude/skills/review-pr/
+!.claude/skills/review-pr/SKILL.md
+!.claude/skills/review-pr/references/
+!.claude/skills/review-pr/references/*.md
 
 # Codex
 AGENTS.md