From f87674aa447b24fb305f3eafcab1e51b30e0d9a6 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Tue, 14 Apr 2026 20:26:27 +0800 Subject: [PATCH] [skip ci]add skills (#2710) Signed-off-by: hsliuustc0106 --- .claude/skills/add-diffusion-model/SKILL.md | 534 ++++++++++++++++ .../references/cache-dit-patterns.md | 254 ++++++++ .../references/custom-model-patterns.md | 273 +++++++++ .../references/parallelism-patterns.md | 571 ++++++++++++++++++ .../references/transformer-adaptation.md | 218 +++++++ .../references/troubleshooting.md | 178 ++++++ .claude/skills/add-tts-model/SKILL.md | 284 +++++++++ .claude/skills/readme.md | 34 ++ .gitignore | 14 +- 9 files changed, 2359 insertions(+), 1 deletion(-) create mode 100644 .claude/skills/add-diffusion-model/SKILL.md create mode 100644 .claude/skills/add-diffusion-model/references/cache-dit-patterns.md create mode 100644 .claude/skills/add-diffusion-model/references/custom-model-patterns.md create mode 100644 .claude/skills/add-diffusion-model/references/parallelism-patterns.md create mode 100644 .claude/skills/add-diffusion-model/references/transformer-adaptation.md create mode 100644 .claude/skills/add-diffusion-model/references/troubleshooting.md create mode 100644 .claude/skills/add-tts-model/SKILL.md create mode 100644 .claude/skills/readme.md diff --git a/.claude/skills/add-diffusion-model/SKILL.md b/.claude/skills/add-diffusion-model/SKILL.md new file mode 100644 index 000000000..a7e0bbf9a --- /dev/null +++ b/.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//pipeline_.py` | Always required | +| **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models//_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` | +| **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models//` 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//download_.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//` 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//download_.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 diff --git a/.claude/skills/add-diffusion-model/references/cache-dit-patterns.md b/.claude/skills/add-diffusion-model/references/cache-dit-patterns.md new file mode 100644 index 000000000..d34ce0e0f --- /dev/null +++ b/.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 | diff --git a/.claude/skills/add-diffusion-model/references/custom-model-patterns.md b/.claude/skills/add-diffusion-model/references/custom-model-patterns.md new file mode 100644 index 000000000..2434e0b5d --- /dev/null +++ b/.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//`. 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//download_.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) +``` diff --git a/.claude/skills/add-diffusion-model/references/parallelism-patterns.md b/.claude/skills/add-diffusion-model/references/parallelism-patterns.md new file mode 100644 index 000000000..933e2d232 --- /dev/null +++ b/.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//` | +| **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`. diff --git a/.claude/skills/add-diffusion-model/references/transformer-adaptation.md b/.claude/skills/add-diffusion-model/references/transformer-adaptation.md new file mode 100644 index 000000000..6e344b6a6 --- /dev/null +++ b/.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), + } +``` diff --git a/.claude/skills/add-diffusion-model/references/troubleshooting.md b/.claude/skills/add-diffusion-model/references/troubleshooting.md new file mode 100644 index 000000000..27acdd8d1 --- /dev/null +++ b/.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 diff --git a/.claude/skills/add-tts-model/SKILL.md b/.claude/skills/add-tts-model/SKILL.md new file mode 100644 index 000000000..e64e7e763 --- /dev/null +++ b/.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_.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//` +- 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) diff --git a/.claude/skills/readme.md b/.claude/skills/readme.md new file mode 100644 index 000000000..b66f2ecd1 --- /dev/null +++ b/.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. diff --git a/.gitignore b/.gitignore index 7f101a784..c0ee96806 100644 --- a/.gitignore +++ b/.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