[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:
Hongsheng Liu authored and GitHub committed 2026-09-05 15:32:30 +08:00
1 parent 13c617a140
commit d4096e01fe
78 files changed
+9896 -459

No files matched your search

+90 -207
View File
@@ -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.