[Model] Add native MAGI-2 Preview diffusion support (#5918)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
This commit is contained in:
1 parent
13c617a140
commit
d4096e01fe
78 files changed
+9896
-459
No files matched your search
@@ -1,6 +1,6 @@
|
||||
---
|
||||
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.
|
||||
description: Add a new diffusion model (text-to-image, text-to-video, image-to-video, text-to-audio, image editing) to vLLM-Omni, including native non-Diffusers ports, reference-parity validation, Cache-DiT, offload, and parallelism support (TP, SP/USP, CFG-Parallel, HSDP). Use when integrating or reviewing a new diffusion model, porting a Diffusers pipeline or custom model repository, creating a DiT adapter, reusing shared examples, or qualifying multi-GPU and memory optimizations.
|
||||
---
|
||||
|
||||
# Adding a Diffusion Model to vLLM-Omni
|
||||
@@ -26,9 +26,16 @@ Check the model's HF repo for `model_index.json`. This determines your path:
|
||||
|----------|----------------|----------------|
|
||||
| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip implementation, then validate model-specific examples, tests, and docs as needed |
|
||||
| **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 |
|
||||
| **Native non-Diffusers model** | No Diffusers index, non-standard checkpoint hierarchy, or custom architecture in a separate repo | Follow **Path B** below; port the runtime natively unless an external adapter was explicitly requested |
|
||||
| **Hybrid** | Has some diffusers components (VAE) but custom transformer/fusion | Mix of Path A and Path B |
|
||||
|
||||
Before coding, write a short integration contract covering runtime ownership,
|
||||
checkpoint discovery, reference revision, I/O geometry, attention semantics,
|
||||
CFG behavior, auxiliary components, target hardware, and the default deployment.
|
||||
For Path B, hybrid models, or any optimization work, read
|
||||
[references/native-model-integration-checklist.md](references/native-model-integration-checklist.md)
|
||||
and use its phase gates.
|
||||
|
||||
## 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.
|
||||
@@ -103,9 +110,15 @@ For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py`
|
||||
|
||||
---
|
||||
|
||||
## Path B: Custom/Private Repo Model
|
||||
## Path B: Native / Non-Diffusers Model
|
||||
|
||||
For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
|
||||
For models without a Diffusers pipeline—weights in custom formats and model
|
||||
code in another public or private repository. Treat that repository as a pinned
|
||||
correctness oracle. A request for native support means the vLLM-Omni runtime
|
||||
must not import the reference implementation; an external adapter is appropriate
|
||||
only when the requested scope explicitly permits that dependency. See
|
||||
[references/custom-model-patterns.md](references/custom-model-patterns.md) for
|
||||
concrete integration patterns.
|
||||
|
||||
### B1. Understand the reference repo
|
||||
|
||||
@@ -126,16 +139,25 @@ This is the key design decision for custom models. Follow these placement rules:
|
||||
| **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 | original-repo package via git clone |
|
||||
| **Reference-only code** | Keep outside the runtime; use a pinned revision for golden outputs and architecture analysis | Reference inference script |
|
||||
| **Explicit external adapter dependency** | External package, only when the requested scope and maintainer direction allow it | Compatibility adapter, not native support |
|
||||
| **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_<name>.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:
|
||||
If the model's code lives in a separate git repo, first decide whether the
|
||||
requested deliverable is native support or an external adapter. Do not silently
|
||||
choose the adapter path.
|
||||
|
||||
**Option 1: Import with graceful fallback** (recommended for models with external utils)
|
||||
**Option 1: Port the code directly** (default for native support)
|
||||
|
||||
Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and
|
||||
adapt them to shared vLLM-Omni contracts. Keep checkpoint loading strict and
|
||||
use the pinned reference only to generate parity evidence.
|
||||
|
||||
**Option 2: Import with graceful fallback** (adapter scope only)
|
||||
|
||||
```python
|
||||
try:
|
||||
@@ -147,11 +169,9 @@ except ImportError:
|
||||
)
|
||||
```
|
||||
|
||||
**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.
|
||||
Use an external runtime dependency only when the user explicitly requested an
|
||||
adapter or a maintainer approved the exception. Document the dependency,
|
||||
pinned revision, installation path, and unsupported native features.
|
||||
|
||||
### B4. Handle custom weight loading
|
||||
|
||||
@@ -221,7 +241,8 @@ class CustomPipeline(nn.Module):
|
||||
|
||||
### 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:
|
||||
Prefer a `model_index.json` at the model root when vLLM-Omni owns an assembled
|
||||
checkpoint directory. For custom models, this is minimal:
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -230,7 +251,14 @@ Custom models need a `model_index.json` at the model root for vllm-omni to disco
|
||||
}
|
||||
```
|
||||
|
||||
The `_class_name` must match a key in `_DIFFUSION_MODELS` in `registry.py`. Additional keys are model-specific (accessed via `od_config.model_config`).
|
||||
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 released repository is immutable and has neither a root `config.json`
|
||||
nor a Diffusers index, add it to a generic native-checkpoint signature resolver:
|
||||
match the exact Hub ID or a distinctive set of local files and return the
|
||||
pipeline class. Do not use model-name substrings or add parallel one-off
|
||||
predicates in CLI, config, and serving consumers.
|
||||
|
||||
If the model's weights come from multiple HF repos, write a **download script** that:
|
||||
1. Downloads from each repo
|
||||
@@ -292,7 +320,9 @@ Use the appropriate existing example script:
|
||||
| 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.
|
||||
Reuse these shared scripts even for custom models when their request and output
|
||||
contracts fit. Create a dedicated model script only when the shared category
|
||||
cannot represent the protocol, and document that gap in the PR.
|
||||
|
||||
**Validation**: No errors, output is meaningful, quality matches reference implementation.
|
||||
|
||||
@@ -300,13 +330,17 @@ See [references/troubleshooting.md](references/troubleshooting.md) for common er
|
||||
|
||||
### Step 6: Add Example Scripts
|
||||
|
||||
For Omni or custom models, create:
|
||||
Only when the shared category scripts cannot represent the model, 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
|
||||
|
||||
Follow the [`add-recipe` skill](../add-recipe/SKILL.md) to add or update the
|
||||
model-family recipe and its `recipes/README.md` row with verified specifications,
|
||||
hardware, commands, feature links, and qualification evidence.
|
||||
|
||||
Required updates:
|
||||
1. `docs/user_guide/diffusion/parallelism/overview.md` — parallelism support overview/table
|
||||
2. `docs/user_guide/diffusion/cpu_offload.md` — if CPU offload supported (add to supported models table)
|
||||
@@ -346,107 +380,20 @@ Classify the model's **CI priority** first:
|
||||
|
||||
**Keep model-specific code inside test modules — not `tests/helpers/{slug}.py`:** deploy constants, prompts, sampling dicts, and inline `request_config` / `form_data` belong in each `test_{slug}.py` and `test_{slug}_expansion.py`. Do not add per-model files under `tests/helpers/`; reuse only repo-wide harness (`mark`, `media`, `runtime`, `stage_config`, `assertions`). **L2+ online/offline e2e:** reuse or add `send_*_request` in `tests/helpers/runtime.py` — tests call the handler, not raw `omni.generate` / HTTP. See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers**.
|
||||
|
||||
Keep the model suite proportional using the six distinct failure owners in the
|
||||
native integration checklist. Combine supported features into a few
|
||||
parametrized E2E rows instead of creating one test per optimization.
|
||||
|
||||
### 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.
|
||||
Add caching only after uncached single-device correctness. Read
|
||||
[references/cache-dit-patterns.md](references/cache-dit-patterns.md), use the
|
||||
automatic single-block-list path when possible, and add a registered
|
||||
`BlockAdapter` only for genuinely custom block topology.
|
||||
|
||||
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/cachedit/model_specific.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).
|
||||
Verify a real cache hit and compare quality with the uncached baseline. Make a
|
||||
speed claim only at a realistic step count where warmup permits hits; an
|
||||
all-warmup smoke proves integration, not acceleration.
|
||||
|
||||
---
|
||||
|
||||
@@ -460,107 +407,30 @@ See [references/parallelism-patterns.md](references/parallelism-patterns.md) for
|
||||
|
||||
#### 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`
|
||||
Replace compatible projections with vLLM parallel linears, preserve checkpoint
|
||||
fusion/loading, and use local head counts. Require query/KV head divisibility
|
||||
and compare a multi-rank forward with the one-rank oracle.
|
||||
|
||||
#### 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)
|
||||
Prefer the declarative `_sp_plan`. For packed variable-length attention or
|
||||
learned-sink LSE correction, keep model math explicit and reuse shared exchange
|
||||
utilities. Validate uneven sequence splits, RoPE coordinates, and outputs
|
||||
against the one-rank oracle.
|
||||
|
||||
#### 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`
|
||||
Confirm the model uses CFG. Then reuse `CFGParallelMixin`, overriding prediction
|
||||
or recombination only for non-standard or multi-output pipelines. Distinguish a
|
||||
packed positive/negative implementation from two independent branches and
|
||||
validate the two-rank result against the packed one-rank oracle.
|
||||
|
||||
#### 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)`
|
||||
Declare layer shard conditions and ignored rank-local modules; preserve mixed
|
||||
checkpoint dtypes. HSDP cannot combine with TP. Measure parameter loading,
|
||||
FSDP materialization, warm HBM, and host PSS rather than assuming sharding saves
|
||||
peak memory. Keep the resident layout as default if HSDP is worse.
|
||||
|
||||
#### 10e. Update parallelism documentation
|
||||
|
||||
@@ -599,6 +469,12 @@ Pipelines without `SupportsComponentDiscovery` fall back to scanning
|
||||
well-known attribute names (`transformer`, `text_encoder`, `vae`,
|
||||
etc.), which fails for non-standard names.
|
||||
|
||||
Keep model-specific checkpoint paths, nested block aliases, layout transforms,
|
||||
and component lifecycles in the model package. Change a shared offloader only
|
||||
for a general contract, demonstrate another consumer or a framework-level bug,
|
||||
and add one focused shared regression. Avoid `if ModelName` branches in shared
|
||||
backends.
|
||||
|
||||
### Step 12: Performance Profiling
|
||||
|
||||
After verifying correctness and implementing parallelism/caching, profile the model's performance to identify bottlenecks and ensure optimal execution.
|
||||
@@ -608,6 +484,12 @@ See the [Profiling Single-Stage Diffusion](../../../docs/contributing/profiling.
|
||||
2. Using Nsight Systems (`nsys`) with `profiler: "cuda"` for low-overhead CUDA traces.
|
||||
3. Controlling profiling via `omni.start_profile()` and `omni.stop_profile()`.
|
||||
|
||||
Report cold E2E, warm user latency, steady-state wave time, peak allocated and
|
||||
reserved HBM, host PSS for the full process tree, and stage boundaries. State
|
||||
whether prompt encoding and VAE/audio decoding are included. For multi-device
|
||||
layouts, plot user latency against throughput per device and keep different
|
||||
denoising-step counts on separate Pareto frontiers.
|
||||
|
||||
---
|
||||
|
||||
## Pre-commit conventions
|
||||
@@ -635,7 +517,7 @@ shellcheck). In particular:
|
||||
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
|
||||
6. **For custom models**: Run the pinned reference separately, then port the runtime natively; do not ship temporary imports from the reference implementation
|
||||
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
|
||||
|
||||
@@ -644,6 +526,7 @@ shellcheck). In particular:
|
||||
- [vllm-omni-test skill](../vllm-omni-test/SKILL.md) — L1–L4 markers, naming, Buildkite wiring, run commands
|
||||
- [Transformer Adaptation](references/transformer-adaptation.md) — porting transformers from diffusers
|
||||
- [Custom Model Patterns](references/custom-model-patterns.md) — patterns for non-diffusers models
|
||||
- [Native Model Integration Checklist](references/native-model-integration-checklist.md) — ownership boundaries, phase gates, qualification matrix, and review evidence
|
||||
- [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
|
||||
@@ -9,7 +9,9 @@ Three caching strategies:
|
||||
- **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.
|
||||
Speedup is model-, step-, and policy-dependent. Treat a cache-enabled run as
|
||||
integration evidence until logs or a focused test show real hits and repeated
|
||||
timings show an improvement over the uncached baseline.
|
||||
|
||||
**Official docs**: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cache_dit
|
||||
|
||||
@@ -243,8 +245,10 @@ vllm serve your-model --omni --port 8098 \
|
||||
|
||||
**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`
|
||||
2. A focused test or summary confirms at least one real cache hit
|
||||
3. Quality is compared with `cache_backend=None`
|
||||
4. Performance uses a realistic step count, repeated timings, and a policy
|
||||
whose warmup does not consume the whole run
|
||||
|
||||
## Excluded Models
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@ vllm_omni/diffusion/models/wan2_2/
|
||||
|
||||
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
|
||||
### External adapter (not native support)
|
||||
|
||||
```
|
||||
vllm_omni/diffusion/models/<name>/
|
||||
@@ -28,7 +28,10 @@ examples/offline_inference/<task>/
|
||||
└── download_<name>.py # Downloads weights and/or clones code repo
|
||||
```
|
||||
|
||||
All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency imported with try/except.
|
||||
All weights are loaded eagerly in `__init__`. `load_weights()` is a no-op. The
|
||||
external dependency is imported with try/except. Use this pattern only when the
|
||||
requested scope explicitly allows an adapter; it does not satisfy a request for
|
||||
a native vLLM-Omni implementation.
|
||||
|
||||
### Custom model with ported code (e.g., BAGEL)
|
||||
|
||||
@@ -86,7 +89,7 @@ init → load ALL weights eagerly via custom helpers → load_weights() = no-op
|
||||
- `load_weights()` is `pass`
|
||||
- Weights may come from multiple HF repos in different formats (`.pth`, `.safetensors`)
|
||||
|
||||
Use this when:
|
||||
Use this only for an explicitly approved external adapter 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)
|
||||
@@ -112,7 +115,14 @@ Custom model `model_index.json` (minimal):
|
||||
}
|
||||
```
|
||||
|
||||
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.
|
||||
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`.
|
||||
|
||||
When the released repository cannot be modified and contains neither this file
|
||||
nor a root `config.json`, use the generic native-checkpoint signature resolver.
|
||||
Match the exact Hub ID or a distinctive local file set; never infer support from
|
||||
a model-name substring.
|
||||
|
||||
## External Dependency Management
|
||||
|
||||
@@ -141,6 +151,10 @@ def download_dependency():
|
||||
|
||||
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.
|
||||
|
||||
This is the default for native model support. Keep the reference repository at
|
||||
a pinned revision for architecture analysis and golden outputs, but do not
|
||||
import it from production code.
|
||||
|
||||
## Multi-Modal Input/Output Protocols
|
||||
|
||||
Custom models that handle images, audio, or video I/O should implement protocol classes:
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
# Native Diffusion Model Integration Checklist
|
||||
|
||||
Use this checklist for native non-Diffusers models, hybrid pipelines, and
|
||||
model-specific optimization work.
|
||||
|
||||
## 1. Freeze the contract before coding
|
||||
|
||||
Record:
|
||||
|
||||
- Reference repository and immutable revision.
|
||||
- Official checkpoint ID, local layout, license, and required auxiliary files.
|
||||
- Pipeline stages: prompt/image encoding, denoising, video/image/audio decoding,
|
||||
and muxing.
|
||||
- Token packing, RoPE coordinates, attention sinks/LSE behavior, MoE layout,
|
||||
scheduler equations, CFG formulation, dtypes, and checkpoint transforms.
|
||||
- Output geometry, frame count/FPS, audio rate/channels, supported tasks, and
|
||||
default generation steps.
|
||||
- Available local hardware versus the official qualification hardware.
|
||||
- Intended default deployment and optional optimization paths.
|
||||
|
||||
State whether the deliverable is a native implementation or an external
|
||||
adapter. For native support, use the reference as an oracle only.
|
||||
|
||||
## 2. Use phase gates
|
||||
|
||||
| Phase | Required exit evidence |
|
||||
|---|---|
|
||||
| Architecture | Component/weight inventory; unsupported reference dependencies identified; ownership decisions documented |
|
||||
| Native correctness | Strict checkpoint load; tiny forward or one-step golden vs pinned reference; output contract validated |
|
||||
| Framework integration | Registry/discovery, shared request types, existing category example, serving route, docs, and recipe work without model-name branching in shared scripts |
|
||||
| Lossless distribution | One-rank oracle vs each supported TP/SP/CFG/HSDP/VAE layout; released-checkpoint smoke for the recommended layout |
|
||||
| Memory optimization | Ordinary/layerwise/DLO behavior, peak HBM, full-tree host PSS, and stage timeline; shared backend unchanged unless a general contract is proven |
|
||||
| Lossy optimization | Cache hit demonstrated; quality delta measured; speed claim made only when warmup permits hits |
|
||||
| Full qualification | Default-resolution E2E at official step count for every supported task; media metadata and hashes recorded |
|
||||
| Review readiness | Focused tests, Ruff/format, strict docs, DCO, current-main rebase, issue/RFC for large changes, concise PR evidence and limitations |
|
||||
|
||||
Do not start performance tuning before native one-step parity. Do not claim a
|
||||
feature from topology validation alone when released weights can exercise it.
|
||||
|
||||
## 3. Keep ownership narrow
|
||||
|
||||
Default placement:
|
||||
|
||||
| Concern | Owner |
|
||||
|---|---|
|
||||
| Architecture, checkpoint remap, packed attention math, nested block alias, auxiliary lifecycle | Model package |
|
||||
| Kernel/version dispatch already useful to multiple models | Shared attention utility |
|
||||
| Request-dependent example defaults | Typed model-extra entry consumed through a registry |
|
||||
| Exact-ID/local-file detection for repositories without standard metadata | Generic native-checkpoint signature table |
|
||||
| General offload/cache/HSDP bug | Shared layer plus one focused shared regression |
|
||||
|
||||
Before changing shared code, answer:
|
||||
|
||||
1. Is this a framework contract or only this checkpoint's layout?
|
||||
2. Can the model expose an existing declarative hook instead?
|
||||
3. Is there a second consumer or a framework-level failing test?
|
||||
4. Can the shared change avoid importing or naming the model?
|
||||
|
||||
If the answers do not justify a shared change, keep it model-local.
|
||||
|
||||
## 4. Reuse examples without leaking model logic
|
||||
|
||||
- Reuse the existing `x_to_y.py` category script.
|
||||
- Resolve the pipeline class through standard metadata or native signature
|
||||
discovery, not path substrings.
|
||||
- Put dynamic model defaults and validation in `model_extras`; let shared
|
||||
examples consume only the generic typed API.
|
||||
- Do not create image/speech/default abstractions without a real consumer.
|
||||
- Add a dedicated example only when the shared request/output protocol cannot
|
||||
express the model.
|
||||
|
||||
## 5. Build a bounded evidence matrix
|
||||
|
||||
Start with the recommended resident layout. Add only supported, useful points:
|
||||
|
||||
- Single request: TP, SP, or hybrid TP/SP.
|
||||
- Throughput: DLO/DP layouts when concurrent requests are valid.
|
||||
- CFG: compare CFG-parallel output with the packed-CFG oracle only if the model
|
||||
actually uses CFG.
|
||||
- HSDP: record loading and materialization peaks; it can use more HBM than a
|
||||
resident SP layout.
|
||||
- VAE patch parallelism: compare decoded tensors or media exactly where the
|
||||
algorithm should be lossless.
|
||||
- Cache: separate cache integration from acceleration evidence.
|
||||
|
||||
With four devices, qualify four-device layouts such as DP4, DP2SP2, SP4, or
|
||||
CFG2xSP2. Do not present CFG2xSP4 or an official eight-device runtime as locally
|
||||
validated.
|
||||
|
||||
For each measured layout record:
|
||||
|
||||
- Cold E2E and warm user latency.
|
||||
- Steady-state wave time and throughput per device.
|
||||
- Prompt, denoise, video decode, audio decode, and mux times.
|
||||
- Peak allocated/reserved HBM and full-process-tree host PSS.
|
||||
- Step count, seed, output geometry, frame/FPS/audio metadata, and parity metric.
|
||||
|
||||
Plot latency versus throughput/device. Keep one-step, four-step, and full-step
|
||||
results separate because they are not quality-equivalent.
|
||||
|
||||
## 6. Keep tests proportional
|
||||
|
||||
Prefer a small set with distinct failure ownership:
|
||||
|
||||
1. Strict weight-key/dtype/load contract.
|
||||
2. Tiny native golden or one-step reference parity.
|
||||
3. One parametrized distributed-parity test covering supported layouts.
|
||||
4. Pipeline/request/output geometry and serving smoke.
|
||||
5. Shared-example/default test when the category script changes.
|
||||
6. One shared regression only for a genuine shared-layer fix.
|
||||
|
||||
Avoid duplicating the same shape or configuration assertion across
|
||||
model-specific files. Keep expensive released-checkpoint evidence in the PR or
|
||||
recipe unless the repository's CI tier explicitly requires it.
|
||||
|
||||
## 7. Write evidence, not marketing
|
||||
|
||||
In the PR description include:
|
||||
|
||||
- A compact architecture/data-flow explanation.
|
||||
- Default deployment and why it is recommended.
|
||||
- Code-change counts by component.
|
||||
- Accuracy and distributed-parity evidence tied to commands/scripts.
|
||||
- Separate Pareto figures for different step counts.
|
||||
- Supported versus validated versus planned optimizations.
|
||||
- Hardware gaps and untested official-runtime claims.
|
||||
|
||||
Call a result a bounded local qualification unless it includes repeated timing
|
||||
and variance. State whether latency includes text/image encoding and all
|
||||
decoders. Preserve the reference revision and exact commands needed to
|
||||
reproduce the evidence.
|
||||
Reference in new issue
Block a user