anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Model] Add native MAGI-2 Preview diffusion support (#5918)

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.buildkite/cuda/test-nightly.yml
+6
View File
@@ -343,6 +343,12 @@ steps:
- pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · MAGI-2 Preview Serving Smoke"
timeout_in_minutes: 180
commands:
- pytest -s -v tests/e2e/online_serving/test_magi2.py -m "advanced_model and diffusion and cuda and H100 and cards_4" --run-level "advanced_model"
mirror_hardwares: h100_4
- label: ":full_moon: Diffusion X2V · Doc Test"
timeout_in_minutes: 60
commands:
.claude/skills/add-diffusion-model/SKILL.md
+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
.claude/skills/add-diffusion-model/references/cache-dit-patterns.md
+7 -3
View File
@@ -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
.claude/skills/add-diffusion-model/references/custom-model-patterns.md
+18 -4
View File
@@ -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:
.claude/skills/add-diffusion-model/references/native-model-integration-checklist.md
+131
View File
@@ -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.
.claude/skills/add-recipe/SKILL.md
+126
View File
@@ -0,0 +1,126 @@
---
name: add-recipe
description: Add or update an in-repository vLLM-Omni model recipe with verified task, input, output, hardware, command, feature, and validation contracts. Use when creating files under recipes/, restructuring a recipe after review, documenting a newly supported model, or synchronizing recipe claims with support tables and shared feature guides.
---
# Add a vLLM-Omni Recipe
## Source the contract
Read [`recipes/TEMPLATE.md`](../../../recipes/TEMPLATE.md) and
[`recipes/README.md`](../../../recipes/README.md) before editing. Inspect the
canonical model card/repository, implementation, shared examples, tests, and
available local qualification evidence.
Do not infer support from a model name, sibling recipe, registry entry, or
topology validator. Mark unexecuted configurations as configuration-only.
## Route by model family
Read only the contributor and user documentation relevant to the model:
| Family | Contributor guide | User-facing documentation to synchronize |
|---|---|---|
| Diffusion | [`adding a diffusion model`](../add-diffusion-model/SKILL.md) | `docs/user_guide/diffusion_features.md`, applicable diffusion feature guides, and shared image/video/audio examples |
| TTS | [`adding a TTS model`](../add-tts-model/SKILL.md) | `examples/offline_inference/text_to_speech/README.md`, `examples/online_serving/text_to_speech/README.md`, and `docs/serving/speech_api.md` |
| Omni | [`adding an omni model`](../../../docs/contributing/model/adding_omni_model.md) | model-family offline/online example docs, `docs/serving/chat_completions_api.md`, and `docs/user_guide/feature_compatibility.md` |
All families must update `docs/models/supported_models.md` and the matching row
in `recipes/README.md`. Do not add a modality-specific support document when
the repository has no such table; update the closest shared user contract.
## Structure the recipe
Keep the recipe task-oriented and use one hardware-specific file named
`recipes/<vendor>/<model>-<hardware>.md`. A recipe may contain multiple
deployment topologies on that hardware, but each additional accelerator model
or platform requires its own suffixed file.
1. **Summary:** vendor, exact model ID, runtime, modes, named hardware,
recommended deployment, and maintainer.
2. **Supported model contract:** put task, input, output, and provided-profile
tables before setup or commands.
3. **References:** link the canonical model source, shared examples, supported
model table, and feature matrix.
4. **Checkpoint/setup:** pin revisions when possible and state required assets.
5. **Hardware:** document the single accelerator named by the file suffix,
including every validated device-count/topology profile on it.
6. **Commands:** reuse shared `examples/` entrypoints unless the model contract
truly requires a dedicated script. Include offline and online paths that
were validated.
7. **Supported features:** use a compact model-specific topology/status table
with links to shared feature guides.
8. **Verification:** provide a quick command and exact expected output
contract.
9. **Qualification evidence:** report correctness, memory, and timing evidence
with its measurement scope and caveats.
## Separate hardware from software
For every locally validated hardware profile, record:
- accelerator vendor/model and per-device memory;
- number of devices;
- interconnect (`NVLink`, `PCIe`, or the platform equivalent);
- host memory when CPU staging/offload is material;
- whether the profile is runtime-qualified or configuration-only.
Record OS, Python, driver/runtime, framework versions, and the vLLM-Omni
revision in a separate software-environment table. Keep precision, worker
count, and DP/TP/SP/PP sizes with the deployment profile or exact command.
Never generalize a result from one accelerator family to another or combine
NVIDIA, AMD, NPU, or distinct accelerator models in one recipe. Create a
separate hardware-suffixed recipe and `recipes/README.md` row. Distinguish
upstream requirements from hardware exercised by the PR.
## Keep shared documentation concise
Link the applicable diffusion, TTS, omni, serving, or design guide for generic
feature semantics and launch instructions. The recipe should contain only the
model-specific status, valid topology, required flag difference, and evidence
boundary.
Put unsupported combinations in the same feature table. Avoid repeating the
same model-specific prose below the global feature matrix; the recipe is the
detailed source of truth.
## Report evidence precisely
For every performance or memory value, state:
- checkpoint/revision and workload;
- accelerator and device count;
- precision and topology;
- cold/warm scope, step count, and concurrency;
- peak HBM and host-memory metric when measured;
- output/quality guard and sample count;
- monitoring overhead or other known variance.
Call a single bounded run qualification evidence, not a benchmark. Do not add
a benchmark script unless the contribution explicitly requires one.
## Synchronize repository documentation
Update all applicable locations:
- the row in `recipes/README.md`;
- `docs/models/supported_models.md`;
- the family-specific documentation selected in **Route by model family**;
- shared example documentation when commands or flags change.
Keep global support tables compact. Link to the recipe instead of duplicating
its deployment explanation.
## Validate
Run, at minimum:
```bash
pre-commit run --files <changed recipe/docs/skill files>
mkdocs build --strict
git diff --check
```
Remove generated documentation artifacts after validation. Confirm every
recipe link resolves and every claimed profile has matching code or evidence.
.claude/skills/add-tts-model/SKILL.md
+8
View File
@@ -499,6 +499,13 @@ Common pre-commit failures, recovery commands for missing sign-off, and the
full `pre-commit run` invocation for a TTS model:
[references/precommit-dco.md](references/precommit-dco.md).
## Documentation and recipe
Follow the [`add-recipe` skill](../add-recipe/SKILL.md) to add or update the
in-repository model-family recipe, the `recipes/README.md` index, the supported
models table, TTS example documentation, and the speech API contract using
validated evidence.
## Integration Checklist
Use this checklist when integrating a new TTS model:
@@ -543,6 +550,7 @@ Use this checklist when integrating a new TTS model:
- [ ] Buildkite entries match level: `test-ready.yml` / `test-merge.yml` or nightly TTS job / `test-nightly.yml`
- [ ] Gradio demo working
- [ ] Documentation added (offline + online docs, nav, supported models)
- [ ] In-repository recipe and `recipes/README.md` row added via the `add-recipe` skill
### Phase 4: Async Chunk
- [ ] Pipeline declares the async handoff processor
.claude/skills/readme.md
+3
View File
@@ -26,6 +26,9 @@ self-check.
- [`add-diffusion-model`](add-diffusion-model/SKILL.md): guides integration of
a new diffusion model into `vllm-omni`
- [`add-recipe`](add-recipe/SKILL.md): creates or updates in-repository model
recipes with verified task/input/output contracts, concrete hardware,
shared feature-guide links, exact commands, and qualification evidence
- [`add-tts-model`](add-tts-model/SKILL.md): covers integration of new TTS
models and related serving workflows
- [`diffusion-perf-opt`](diffusion-perf-opt/SKILL.md): guides diffusion model
docs/contributing/model/adding_omni_model.md
+26 -17
View File
@@ -29,7 +29,7 @@ Each stage is implemented as a separate model class that can be configured indep
When adding a new model, you'll need to create the following structure:
```
```text
vllm_omni/model_executor/models/
└── your_model_name/ # Model directory (e.g., qwen3_omni)
├── __init__.py # Exports main model class
@@ -422,22 +422,25 @@ The stage transition process follows these steps:
Understanding the data structures is crucial for implementing stage transitions:
**Input to your function:**
- `stage_list[source_stage_id].engine_outputs`: List of `EngineCoreOutput` objects
- - Each contains `outputs`: List of `RequestOutput` objects
- Each `RequestOutput` has:
- - - `token_ids`: Generated token IDs
- `multimodal_output`: Dict with keys like `"code_predictor_codes"`, etc.These are the hidden states or intermediate outputs from the model's forward pass
- `prompt_token_ids`: Original prompt token IDs
- Each contains `outputs`: List of `RequestOutput` objects
- Each `RequestOutput` has:
- `token_ids`: Generated token IDs
- `multimodal_output`: Dict with keys like `"code_predictor_codes"`, etc. These are the hidden states or intermediate outputs from the model's forward pass
- `prompt_token_ids`: Original prompt token IDs
**Output from your function:**
- Must return `list[OmniTokensPrompt]` where each `OmniTokensPrompt` contains:
- - `prompt_token_ids`: List[int] - Token IDs for the next stage
- `prompt_token_ids`: List[int] - Token IDs for the next stage
- `additional_information`: Dict[str, Any] - Optional metadata (e.g., embeddings, hidden states)
- `multi_modal_data`: Optional multimodal data if needed
### How Model Outputs Are Stored
The model's `forward()` method returns an `OmniOutput` object that contains:
- `text_hidden_states`: Final hidden states for text generation
- `multimodal_outputs`: Dict containing intermediate outputs
@@ -464,6 +467,7 @@ def forward(self, ...):
```
These keys are then accessible in your stage transition function:
```python
# In stage_input_processors/qwen3_omni.py
thinker_prefill_embeddings = output.multimodal_output["0"] # Access by key
@@ -493,7 +497,7 @@ thinker_hidden_states = output.multimodal_output["24"]
Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`. Each inter-stage edge should provide a **coherent processor set** rather than a single monolithic function:
| Suffix | Role | Runs when |
|--------|------|-----------|
| -------- | ------ | ----------- |
| `*_full_payload` | Worker-side payload producer | `async_chunk=false`; accumulates tensors and ships via connector |
| `*_async_chunk` | Scheduler-side streaming producer | `async_chunk=true`; emits per-chunk payloads |
| `*_token_only` | Orchestrator placeholder builder | `async_chunk=false`; allocates downstream prompt slots only |
@@ -554,7 +558,16 @@ For comprehensive testing guidelines, please refer to the [Test Writing Guide](.
## Adding a Model Recipe
After implementing and testing your model, please add a model recipe to the [vllm-project/recipes](https://github.com/vllm-project/recipes) repository. This helps other users understand how to use your model with vLLM-Omni.
After implementing and testing your model, add a recipe under this repository's
`recipes/` directory. Follow the
[`add-recipe` skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/add-recipe/SKILL.md)
and [`recipes/TEMPLATE.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/TEMPLATE.md),
then update the model row in
[`recipes/README.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/README.md).
The external
[vllm-project/recipes](https://github.com/vllm-project/recipes) repository can
be used as a structural reference, but it is not the destination for a
vLLM-Omni model recipe.
### What to Include
@@ -568,15 +581,10 @@ Your recipe should include:
3. **Usage Examples**: Command-line examples demonstrating how to run the model
4. **Configuration Details**: Important configuration parameters and their meanings
### Example
For reference, see the [LongCat recipe example](https://github.com/vllm-project/recipes/pull/179) which demonstrates the expected format and structure.
### Recipe Location
Create your recipe file in the appropriate directory structure:
- For organization-specific models: `OrganizationName/ModelName.md`
- For general models: `ModelName.md`
Create `recipes/<vendor>/<model-family>.md` and keep commands, feature status,
hardware evidence, and limitations specific to the model family.
The recipe should be a Markdown file that provides clear, reproducible instructions for users to get started with your model.
@@ -590,7 +598,7 @@ Adding a new model to vLLM-Omni involves:
4. **Define pipeline topology and deploy defaults** in `pipeline.py` and `vllm_omni/deploy/`
5. **Implement stage input processors** for stage transitions
6. **Write tests** to verify functionality
7. **Add model recipe** to the [vllm-project/recipes](https://github.com/vllm-project/recipes) repository (see [Adding a Model Recipe](#adding-a-model-recipe) section)
7. **Add an in-repository model recipe** under `recipes/` and update its index (see [Adding a Model Recipe](#adding-a-model-recipe))
### Qwen3-Omni Reference Files
@@ -607,6 +615,7 @@ For a complete reference implementation, see:
- **Testing**: `vllm_omni/tests/e2e/offline_inference/test_qwen3_omni.py`
For more information, see:
- [Architecture Overview](../../design/architecture_overview.md)
- [Supported Models](../../models/supported_models.md)
- [Pipeline and Deploy Configuration Guide](../../configuration/stage_configs.md)
docs/contributing/model/adding_tts_model.md
+8 -5
View File
@@ -991,10 +991,13 @@ git push origin your-branch --force-with-lease
## Adding a Model Recipe
After implementing and testing your model, add a model recipe to the
[vllm-project/recipes](https://github.com/vllm-project/recipes) repository so users can
get started quickly. See [Adding an Omni-Modality Model](./adding_omni_model.md#adding-a-model-recipe)
for the expected format.
After implementing and testing your model, add a recipe under this repository's
`recipes/` directory. Follow the
[`add-recipe` skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/add-recipe/SKILL.md)
and [`recipes/TEMPLATE.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/TEMPLATE.md),
update the model row in
[`recipes/README.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/README.md),
and synchronize the relevant TTS example and speech API documentation.
## Summary
@@ -1011,7 +1014,7 @@ Adding a TTS model to vLLM-Omni involves:
8. **Online serving** - add all 5 integration points to `serving_speech.py` in one commit
9. **Tests** - cover single request, batching, and streaming
10. **Pre-commit + DCO** - run `pre-commit` before pushing; sign every commit with `git commit -s`
11. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
11. **Model recipe** - add under the in-repository `recipes/` directory and update its index
12. **Invariants** - re-check I1–I5 (streaming contract, consumer hygiene, hot-loop discipline, validation pyramid, per-request state) at the end of every phase
### Qwen3-TTS Reference Files
docs/design/figures/magi2/dlo-pareto-1-step.svg
+59
View File
@@ -0,0 +1,59 @@
<svg xmlns="http://www.w3.org/2000/svg" width="760" height="460" viewBox="0 0 760 460" role="img" aria-labelledby="title desc">
<title id="title">MAGI-2 DLO one-step Pareto frontier</title>
<desc id="desc">Warm user latency in seconds versus throughput per device. DP1SP4 is 5.735 seconds and 0.0436 outputs per second per device, DP2SP2 is 6.785 and 0.0737, and DP4 is 9.997 and 0.1000.</desc>
<rect width="760" height="460" fill="#ffffff"/>
<style>
text { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #111827; }
.title { font-size: 20px; font-weight: 600; }
.axis-label { font-size: 14px; font-weight: 500; }
.tick { font-size: 12px; fill: #4b5563; }
.point-label { font-size: 13px; font-weight: 600; }
.value-label { font-size: 11px; fill: #4b5563; }
.grid { stroke: #e5e7eb; stroke-width: 1; }
.axis { stroke: #6b7280; stroke-width: 1.5; }
.frontier { fill: none; stroke: #2563eb; stroke-width: 3; }
.point { fill: #2563eb; stroke: #ffffff; stroke-width: 2; }
</style>
<text x="380" y="28" text-anchor="middle" class="title">One diffusion step</text>
<g aria-hidden="true">
<line x1="85" y1="309" x2="700" y2="309" class="grid"/>
<line x1="85" y1="226" x2="700" y2="226" class="grid"/>
<line x1="85" y1="143" x2="700" y2="143" class="grid"/>
<line x1="85" y1="61" x2="700" y2="61" class="grid"/>
<line x1="85" y1="40" x2="85" y2="350" class="axis"/>
<line x1="85" y1="350" x2="700" y2="350" class="axis"/>
<text x="73" y="313" text-anchor="end" class="tick">0.04</text>
<text x="73" y="230" text-anchor="end" class="tick">0.06</text>
<text x="73" y="147" text-anchor="end" class="tick">0.08</text>
<text x="73" y="65" text-anchor="end" class="tick">0.10</text>
<text x="85" y="370" text-anchor="middle" class="tick">5</text>
<text x="197" y="370" text-anchor="middle" class="tick">6</text>
<text x="309" y="370" text-anchor="middle" class="tick">7</text>
<text x="421" y="370" text-anchor="middle" class="tick">8</text>
<text x="532" y="370" text-anchor="middle" class="tick">9</text>
<text x="644" y="370" text-anchor="middle" class="tick">10</text>
</g>
<text x="392" y="414" text-anchor="middle" class="axis-label">Warm user latency (seconds) — lower is better</text>
<text x="20" y="195" transform="rotate(-90 20 195)" text-anchor="middle" class="axis-label">Throughput/device (outputs/s) — higher is better</text>
<polyline points="167,294 285,169 644,61" class="frontier"/>
<circle cx="167" cy="294" r="7" class="point"/>
<text x="179" y="288" class="point-label">DP1SP4</text>
<text x="179" y="305" class="value-label">5.735 s, 0.0436</text>
<circle cx="285" cy="169" r="7" class="point"/>
<text x="297" y="163" class="point-label">DP2SP2</text>
<text x="297" y="180" class="value-label">6.785 s, 0.0737</text>
<circle cx="644" cy="61" r="7" class="point"/>
<text x="632" y="84" text-anchor="end" class="point-label">DP4</text>
<text x="632" y="101" text-anchor="end" class="value-label">9.997 s, 0.1000</text>
<text x="700" y="444" text-anchor="end" class="tick">4x NVIDIA L20X · warm qualification run</text>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 3.1 KiB

docs/design/figures/magi2/dlo-pareto-4-step.svg
+65
View File
@@ -0,0 +1,65 @@
<svg xmlns="http://www.w3.org/2000/svg" width="760" height="460" viewBox="0 0 760 460" role="img" aria-labelledby="title desc">
<title id="title">MAGI-2 DLO four-step Pareto frontier</title>
<desc id="desc">Warm user latency in seconds versus throughput per device. DP1SP4 is 8.647 seconds and 0.0289 outputs per second per device, DP2SP2 is 11.682 and 0.0428, and DP4 is 20.957 and 0.0477.</desc>
<rect width="760" height="460" fill="#ffffff"/>
<style>
text { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #111827; }
.title { font-size: 20px; font-weight: 600; }
.axis-label { font-size: 14px; font-weight: 500; }
.tick { font-size: 12px; fill: #4b5563; }
.point-label { font-size: 13px; font-weight: 600; }
.value-label { font-size: 11px; fill: #4b5563; }
.grid { stroke: #e5e7eb; stroke-width: 1; }
.axis { stroke: #6b7280; stroke-width: 1.5; }
.frontier { fill: none; stroke: #2563eb; stroke-width: 3; }
.point { fill: #2563eb; stroke: #ffffff; stroke-width: 2; }
</style>
<text x="380" y="28" text-anchor="middle" class="title">Four diffusion steps</text>
<g aria-hidden="true">
<line x1="85" y1="350" x2="700" y2="350" class="grid"/>
<line x1="85" y1="288" x2="700" y2="288" class="grid"/>
<line x1="85" y1="226" x2="700" y2="226" class="grid"/>
<line x1="85" y1="164" x2="700" y2="164" class="grid"/>
<line x1="85" y1="102" x2="700" y2="102" class="grid"/>
<line x1="85" y1="40" x2="700" y2="40" class="grid"/>
<line x1="85" y1="40" x2="85" y2="350" class="axis"/>
<line x1="85" y1="350" x2="700" y2="350" class="axis"/>
<text x="73" y="354" text-anchor="end" class="tick">0.025</text>
<text x="73" y="292" text-anchor="end" class="tick">0.030</text>
<text x="73" y="230" text-anchor="end" class="tick">0.035</text>
<text x="73" y="168" text-anchor="end" class="tick">0.040</text>
<text x="73" y="106" text-anchor="end" class="tick">0.045</text>
<text x="73" y="44" text-anchor="end" class="tick">0.050</text>
<text x="85" y="370" text-anchor="middle" class="tick">8</text>
<text x="173" y="370" text-anchor="middle" class="tick">10</text>
<text x="261" y="370" text-anchor="middle" class="tick">12</text>
<text x="349" y="370" text-anchor="middle" class="tick">14</text>
<text x="436" y="370" text-anchor="middle" class="tick">16</text>
<text x="524" y="370" text-anchor="middle" class="tick">18</text>
<text x="612" y="370" text-anchor="middle" class="tick">20</text>
<text x="700" y="370" text-anchor="middle" class="tick">22</text>
</g>
<text x="392" y="414" text-anchor="middle" class="axis-label">Warm user latency (seconds) — lower is better</text>
<text x="20" y="195" transform="rotate(-90 20 195)" text-anchor="middle" class="axis-label">Throughput/device (outputs/s) — higher is better</text>
<polyline points="113,302 247,129 654,69" class="frontier"/>
<circle cx="113" cy="302" r="7" class="point"/>
<text x="125" y="296" class="point-label">DP1SP4</text>
<text x="125" y="313" class="value-label">8.647 s, 0.0289</text>
<circle cx="247" cy="129" r="7" class="point"/>
<text x="259" y="123" class="point-label">DP2SP2</text>
<text x="259" y="140" class="value-label">11.682 s, 0.0428</text>
<circle cx="654" cy="69" r="7" class="point"/>
<text x="642" y="92" text-anchor="end" class="point-label">DP4</text>
<text x="642" y="109" text-anchor="end" class="value-label">20.957 s, 0.0477</text>
<text x="700" y="444" text-anchor="end" class="tick">4x NVIDIA L20X · warm qualification run</text>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 3.6 KiB

docs/models/supported_models.md
+5 -4
View File
@@ -18,7 +18,7 @@ th {
</style>
| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU | Recipe |
|--------------|--------|-------------------|------------|---------|-----|-----------|--------|
| -------------- | -------- | ------------------- | ------------ | --------- | ----- | ----------- | -------- |
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) |
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — |
@@ -42,7 +42,7 @@ th {
| `SanaWmPipeline` | SANA-WM | `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-WM.md) |
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | ✅︎ | — |
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — |
@@ -55,6 +55,7 @@ th {
| `SanaVideoPipeline` | SANA-Video-2B T2V (native 480p/720p; Diffusers adapter validated at both resolutions) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) |
| `SanaImageToVideoPipeline` | SANA-Video-2B I2V (native and Diffusers adapter validated at 480p/720p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) |
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
| `Magi2Pipeline` | MAGI-2 native Preview-stage T2VA and I2VA (272p/540p) | `sand-ai/MAGI-2-preview` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SandAI/MAGI-2-preview-L20X.md) |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
@@ -101,11 +102,11 @@ th {
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | ✅︎ | — |
| `DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | ✅︎ | — |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — |
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
| `HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
| `HiDreamO1ImagePipeline` | HiDream-O1-Image | `HiDream-ai/HiDream-O1-Image` | ✅︎ | | | | — |
The `Recipe` column links the corresponding deployment recipe. `Published` is
docs/user_guide/diffusion_features.md
+8 -9
View File
@@ -1,5 +1,7 @@
# Diffusion Advanced Features
<!-- markdownlint-disable MD028 MD060 -->
## Table of Contents
- [Overview](#overview)
@@ -30,7 +32,6 @@ Cache methods trade minimal quality for significant speedup. Quality loss is typ
| **[TeaCache](diffusion/cache_acceleration/teacache.md)** | Adaptive caching using modulated inputs | Quick setup, balanced quality/speed on single GPU |
| **[Cache-DiT](diffusion/cache_acceleration/cache_dit.md)** | Multiple caching techniques: DBCache, TaylorSeer, SCM | Fine-grained control, tunable quality-speed tradeoff |
#### Lossless Acceleration
Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU).
@@ -72,7 +73,6 @@ Extension methods add specialized capabilities to diffusion models beyond standa
| **[LoRA Inference](diffusion/lora.md)** | Enables inference with Low-Rank Adaptation (LoRA) adapters weights | Reinforcement learning extensions |
| **[Frame Interpolation](diffusion/frame_interpolation.md)** | Inserts intermediate video frames after generation for smoother motion | Video generation pipelines that need higher temporal smoothness |
### Execution Modes
Execution modes control how the diffusion pipeline processes requests and
@@ -147,6 +147,7 @@ The following tables show which models support each feature:
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
> Notes:
>
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary.
> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload.
@@ -170,18 +171,18 @@ The following tables show which models support each feature:
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **LongCat-Video-Avatar-1.5** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **MiniMax-H3** | ✅ (FL2VA) | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ |
| **MAGI-2 Preview** | ❌ | ✅ | ✅ (Ulysses) | ✅ (2-way) | ✅ | ❌ | ✅ | ✅ (1-GPU/LW; DLO DP-AG/SP no-AG) | ✅ (tile) | ❌ | ❌ |
| **SANA-WM** | ❌ | ❌ | ❌<sup>5</sup> | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
> Notes:
> 5. SANA-WM cannot support sequence parallelism: its bidirectional gated delta
> recurrence carries state across frames, so a rank cannot denoise a slice of
> the token sequence in isolation. Doing so would need a distributed scan or
> an all-gather before every GDN block. The remaining ❌ columns are simply
> unvalidated on this model, not known-broken.
> recurrence carries state across frames, so a rank cannot denoise a slice of
> the token sequence in isolation. Doing so would need a distributed scan or
> an all-gather before every GDN block. The remaining ❌ columns are simply
> unvalidated on this model, not known-broken.
> **Step execution note:** Helios supports single-request step execution only;
> use `max_num_seqs=1`.
**Frame Interpolation Support**
- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines; SANA-Video-2B native text-to-video and image-to-video pipelines
@@ -193,7 +194,6 @@ The following tables show which models support each feature:
|-----------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ |
## Feature Compatibility
**Legend:**
@@ -232,7 +232,6 @@ The following tables show which models support each feature:
5. The compatibility matrix uses FP8 as the representative quantization method.
6. Step Execution is not compatible with any diffusion cache backend. LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
## Multi-Thread Weight Loading
The loading guide now lives at [Diffusion Startup and
examples/offline_inference/image_to_video/README.md
+9 -3
View File
@@ -1,7 +1,7 @@
# Image-To-Video
This shared example generates videos from images with VACE, Wan2.2, LTX-2,
HunyuanVideo-1.5, SANA-Video, Cosmos3, and other compatible pipelines.
HunyuanVideo-1.5, SANA-Video, Cosmos3, MAGI-2, and other compatible pipelines.
- `image_to_video.py`: command-line script for single video generation with advanced options.
@@ -27,9 +27,14 @@ This folder provides a unified CLI script for image-to-video generation using vL
| `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card |
| `Lightricks/LTX-2` | 512 x 768 | 121 | 40 | video 3.0 / audio 7.0 | Memory use depends on frame count and tensor parallelism |
| `sand-ai/MAGI-2-preview` | 512 x 896 | 125 | 100 | Model-fixed | Native four-GPU TP/SP; resident SP4 default; DLO available |
| `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480 x 832 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; Wan VAE |
| `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704 x 1280 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; LTX-2 Video VAE |
MAGI-2 native Preview setup, four-GPU topology and DLO choices, I2VA commands,
and eight-GPU validation status are documented in the
[`MAGI-2 Preview L20X recipe`](../../../recipes/SandAI/MAGI-2-preview-L20X.md).
!!! info
Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading.
@@ -147,6 +152,9 @@ python image_to_video.py \
| `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
| `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
| `--enable-layerwise-offload` | flag | off | Enable layerwise offloading on DiT modules |
| `--enable-distributed-layerwise-offload` | flag | off | Enable distributed layerwise offload |
| `--dlo-use-allgather` / `--dlo-no-use-allgather` | flag | on / off | Select sharded AllGather reconstruction (default) or rank-local streaming |
| `--dlo-resident-layers` | int | `0` | Leading main-DiT blocks kept device-resident during distributed layerwise offload |
| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
| `--tensor-parallel-size` | int | `1` | Tensor parallel size (effective for models that support TP, e.g. LTX2) |
| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree |
@@ -339,8 +347,6 @@ Key arguments:
- `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.
- `--hsdp-replicate-size`: Number of replica groups for HSDP. Each replica holds a full sharded copy. Default 1 means pure sharding (no replication).
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
## Wan2.1 VACE Conditional Tasks
examples/offline_inference/image_to_video/image_to_video.py
+61 -13
View File
@@ -3,7 +3,7 @@
"""
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2/LTX-2.3,
HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE.
HunyuanVideo-1.5, SANA-Video, Cosmos3, MAGI-2, or Wan2.1 VACE.
Supports:
- Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
@@ -64,6 +64,7 @@ import numpy as np
import PIL.Image
import torch
from vllm_omni.diffusion.data import resolve_model_class_name
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -75,6 +76,7 @@ from vllm_omni.model_extras import (
from vllm_omni.model_extras import (
get_extra_body_params,
get_model_class_name,
get_video_generation_defaults,
should_preserve_reference_image_size,
)
from vllm_omni.outputs import OmniRequestOutput
@@ -122,15 +124,15 @@ def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description=(
"Generate a video from one or more images "
"(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE)."
"(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, MAGI-2, or Wan2.1 VACE)."
)
)
parser.add_argument(
"--model",
default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
help=(
"Diffusers I2V model ID or local path "
"(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE)."
"I2V model ID or local path "
"(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, MAGI-2, or Wan2.1 VACE)."
),
)
parser.add_argument(
@@ -209,7 +211,7 @@ def parse_args() -> argparse.Namespace:
help="Diffusion KV-cache quantization skip-layer selector, e.g. '0,1,4-8'.",
)
parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.")
parser.add_argument("--fps", type=float, default=None, help="Frames per second for the output video.")
parser.add_argument(
"--vae-use-slicing",
action="store_true",
@@ -230,6 +232,30 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable layerwise (blockwise) offloading on DiT modules.",
)
parser.add_argument(
"--enable-distributed-layerwise-offload",
action="store_true",
help="Enable distributed layerwise offloading with overlapped host-to-device weight streaming.",
)
parser.add_argument(
"--dlo-use-allgather",
dest="dlo_use_allgather",
action="store_true",
default=True,
help="Use shard + AllGather weight reconstruction for distributed layerwise offload (default: enabled).",
)
parser.add_argument(
"--dlo-no-use-allgather",
dest="dlo_use_allgather",
action="store_false",
help="Stream standard-loader rank-local weights without DLO sharding or AllGather.",
)
parser.add_argument(
"--dlo-resident-layers",
type=int,
default=0,
help="Number of leading main-DiT blocks to keep device-resident during distributed layerwise offload.",
)
parser.add_argument(
"--enforce-eager",
action="store_true",
@@ -391,12 +417,14 @@ def main():
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
model_name = str(args.model).lower() if args.model is not None else ""
model_class_name = args.model_class_name
model_class_name_lower = (model_class_name or "").lower()
resolved_model_class_name = model_class_name or resolve_model_class_name(args.model)
model_class_name_lower = (resolved_model_class_name or "").lower()
video_defaults = get_video_generation_defaults(resolved_model_class_name, args.extra_body)
is_ltx2_distilled = "distilled" in model_class_name_lower or "distilled" in model_name
is_ltx23 = "ltx23" in model_class_name_lower or "ltx-2.3" in model_name
is_ltx2 = is_ltx2_distilled or is_ltx23 or "ltx2" in model_class_name_lower or "ltx-2" in model_name
is_sana = model_class_name == "SanaImageToVideoPipeline" or "sana-video" in model_name
is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
is_sana = resolved_model_class_name == "SanaImageToVideoPipeline" or "sana-video" in model_name
is_cosmos = "cosmos" in model_name or "cosmos" in model_class_name_lower
is_cosmos_edge = is_cosmos and ("edge" in model_name or "edge" in model_class_name_lower)
image = PIL.Image.open(args.image).convert("RGB") if args.image else None
@@ -411,7 +439,17 @@ def main():
# Per-model generation defaults, applied only when the matching flag is omitted.
# Cosmos3 would otherwise silently inherit the Wan2.2 defaults (wrong size/steps/shift).
if is_cosmos_edge:
if video_defaults is not None:
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
video_defaults.fps,
video_defaults.guidance_scale,
video_defaults.num_frames,
video_defaults.num_inference_steps,
video_defaults.flow_shift,
video_defaults.max_area,
video_defaults.dimension_multiple,
)
elif is_cosmos_edge:
# Cosmos3-Edge native defaults: 480x832, gs 5.0, flow_shift 3.0 — NOT the Nano/Super
# 720p / gs 6.0 / flow_shift 10.0 below, which produce degenerate output on Edge.
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
@@ -477,7 +515,10 @@ def main():
# Calculate dimensions if not provided (model-aware max area).
height = args.height
width = args.width
if height is None or width is None:
if video_defaults is not None:
height = height or video_defaults.height
width = width or video_defaults.width
elif height is None or width is None:
calc_height, calc_width = calculate_dimensions(dimension_image, max_area=d_max_area, mod_value=d_mod)
height = height or calc_height
width = width or calc_width
@@ -557,6 +598,9 @@ def main():
cache_config=cache_config,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
profiler_config=args.profiler_config,
enable_distributed_layerwise_offload=args.enable_distributed_layerwise_offload,
dlo_use_allgather=args.dlo_use_allgather,
dlo_resident_layers=args.dlo_resident_layers,
)
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
@@ -613,9 +657,12 @@ def main():
)
negative_prompt = args.negative_prompt
if negative_prompt is None and not is_ltx2:
# Preserve the historical empty-prompt behavior for non-LTX examples.
negative_prompt = ""
if negative_prompt is None:
if video_defaults is not None:
negative_prompt = video_defaults.default_negative_prompt
elif not is_ltx2:
# Preserve the historical empty-prompt behavior for non-LTX examples.
negative_prompt = ""
prompt_dict = build_image_to_video_prompt(
prompt=args.prompt,
negative_prompt=negative_prompt,
@@ -642,6 +689,7 @@ def main():
boundary_ratio=args.boundary_ratio,
num_inference_steps=num_inference_steps,
num_frames=num_frames,
fps=fps,
frame_rate=frame_rate,
extra_args=sampling_extra_args,
)
examples/offline_inference/text_to_video/text_to_video.md
+8
View File
@@ -14,9 +14,14 @@ A unified script for text-to-video generation. Supports multiple models with mod
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
| `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) |
| `BestWishYsh/Helios-Base` / `Helios-Mid` / `Helios-Distilled` | 384x640 | 99 | 50 | 5.0 / 5.0 / 1.0 | — |
| `sand-ai/MAGI-2-preview` | 512x896 | 125 | 100 | Model-fixed | Native four-GPU TP/SP; resident SP4 default; DLO available |
| `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480x832 | 81 | 50 | 6.0 | BF16 DiT + FP32 Wan VAE wrapper |
| `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704x1280 | 81 | 50 | 6.0 | BF16 DiT + LTX-2 Video VAE wrapper |
MAGI-2 native Preview setup, four-GPU topology and DLO choices, request
constraints, and eight-GPU validation status are documented in the
[`MAGI-2 Preview L20X recipe`](../../../recipes/SandAI/MAGI-2-preview-L20X.md).
## Local CLI Usage
### Wan2.2 (default)
@@ -264,6 +269,9 @@ python text_to_video.py \
- `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
- `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
- `--enable-distributed-layerwise-offload`: enable distributed layerwise offload.
- `--dlo-use-allgather` / `--dlo-no-use-allgather`: use sharded AllGather reconstruction (the default) or stream rank-local weights without AllGather.
- `--dlo-resident-layers`: keep this many leading main-DiT blocks device-resident during distributed layerwise offload (default: `0`).
- `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
- `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
- `--quantization`: quantization method (such as `fp8` for FP8).
examples/offline_inference/text_to_video/text_to_video.py
+61 -9
View File
@@ -10,12 +10,18 @@ from typing import Any
import numpy as np
import torch
from vllm_omni.diffusion.data import resolve_model_class_name
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.lora.request import LoRARequest
from vllm_omni.lora.utils import stable_lora_int_id
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name, get_output_tensor_range
from vllm_omni.model_extras import (
get_extra_body_params,
get_model_class_name,
get_output_tensor_range,
get_video_generation_defaults,
)
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
@@ -135,9 +141,17 @@ _MODEL_PRESETS = {
}
def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
def _detect_preset(
model: str,
model_class_name: str | None = None,
extra_body: dict[str, Any] | None = None,
) -> dict:
model_lower = model.lower()
class_lower = (model_class_name or "").lower()
resolved_model_class_name = model_class_name or resolve_model_class_name(model)
class_lower = (resolved_model_class_name or "").lower()
video_defaults = get_video_generation_defaults(resolved_model_class_name, extra_body)
if video_defaults is not None:
return video_defaults.cli_defaults()
if "sana-video" in model_lower or "sana_video" in model_lower or "sanavideo" in class_lower:
return _MODEL_PRESETS["sana_720p" if "720p" in model_lower else "sana_480p"]
if "lingbot" in model_lower or "lingbotvideo" in class_lower:
@@ -221,7 +235,8 @@ def parse_extra_body(value: str) -> dict[str, Any]:
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Generate a video from a text prompt. "
"Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, and other text-to-video models."
"Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, MAGI-2, "
"and other text-to-video models."
)
parser.add_argument(
"--model",
@@ -288,7 +303,7 @@ def parse_args() -> argparse.Namespace:
type=str,
default=None,
choices=["cache_dit"],
help="Cache backend for acceleration (Wan2.2). Default: None.",
help="Cache backend for supported diffusion pipelines. Default: None.",
)
parser.add_argument(
"--enable-cache-dit-summary",
@@ -296,7 +311,7 @@ def parse_args() -> argparse.Namespace:
help="Enable cache-dit summary logging after diffusion forward passes.",
)
parser.add_argument("--output", type=str, default=None, help="Output path (mp4). Default: model-specific.")
parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.")
parser.add_argument("--fps", type=float, default=None, help="Frames per second for the output video.")
parser.add_argument(
"--vae-use-slicing",
action="store_true",
@@ -322,6 +337,30 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable layerwise (blockwise) offloading on DiT modules.",
)
parser.add_argument(
"--enable-distributed-layerwise-offload",
action="store_true",
help="Enable distributed layerwise offloading with overlapped host-to-device weight streaming.",
)
parser.add_argument(
"--dlo-use-allgather",
dest="dlo_use_allgather",
action="store_true",
default=True,
help="Use shard + AllGather weight reconstruction for distributed layerwise offload (default: enabled).",
)
parser.add_argument(
"--dlo-no-use-allgather",
dest="dlo_use_allgather",
action="store_false",
help="Stream standard-loader rank-local weights without DLO sharding or AllGather.",
)
parser.add_argument(
"--dlo-resident-layers",
type=int,
default=0,
help="Number of leading main-DiT blocks to keep device-resident during distributed layerwise offload.",
)
parser.add_argument(
"--audio-sample-rate",
type=int,
@@ -464,15 +503,17 @@ def _extract_peak_memory_mb(result: Any) -> float:
def main():
args = parse_args()
model_class_name = args.model_class_name
resolved_model_class_name = model_class_name or resolve_model_class_name(args.model)
video_defaults = get_video_generation_defaults(resolved_model_class_name, args.extra_body)
preset = _detect_preset(args.model, model_class_name)
preset = _detect_preset(args.model, resolved_model_class_name, args.extra_body)
for key, default_val in preset.items():
if getattr(args, key.replace("-", "_"), None) is None:
setattr(args, key.replace("-", "_"), default_val)
model_class_name = args.model_class_name
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
# Cache-dit config (Wan2.2 only)
# Shared Cache-DiT defaults for pipelines that declare an adapter.
cache_config = None
if args.cache_backend == "cache_dit":
cache_config = {
@@ -501,6 +542,9 @@ def main():
cfg_parallel_size=args.cfg_parallel_size,
tensor_parallel_size=args.tensor_parallel_size,
vae_patch_parallel_size=args.vae_patch_parallel_size,
use_hsdp=args.use_hsdp,
hsdp_shard_size=args.hsdp_shard_size,
hsdp_replicate_size=args.hsdp_replicate_size,
pipeline_parallel_size=args.pipeline_parallel_size,
enable_expert_parallel=args.enable_expert_parallel,
enforce_eager=args.enforce_eager,
@@ -509,6 +553,9 @@ def main():
cache_config=cache_config,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
profiler_config=args.profiler_config,
enable_distributed_layerwise_offload=args.enable_distributed_layerwise_offload,
dlo_use_allgather=args.dlo_use_allgather,
dlo_resident_layers=args.dlo_resident_layers,
)
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
@@ -572,6 +619,8 @@ def main():
)
negative_prompt = args.negative_prompt
if negative_prompt is None and video_defaults is not None:
negative_prompt = video_defaults.default_negative_prompt
if negative_prompt is None and all(preset is not _MODEL_PRESETS[name] for name in ("lingbot", "ltx2", "ltx23")):
# Preserve the historical empty-prompt behavior for non-LTX examples.
negative_prompt = ""
@@ -591,7 +640,10 @@ def main():
num_frames=args.num_frames,
extra_args=extra_args,
)
if args.frame_rate is not None:
if video_defaults is not None:
sampling_kwargs["fps"] = args.fps
sampling_kwargs["frame_rate"] = args.frame_rate or args.fps
elif args.frame_rate is not None:
sampling_kwargs["frame_rate"] = args.frame_rate
if args.guidance_scale_high is not None:
sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
pyproject.toml
+2
View File
@@ -351,6 +351,8 @@ extend-ignore-re = [
]
[tool.typos.default.extend-words]
ue = "ue"
# Einstein-summation index tuple for head/expert/dimension axes.
hed = "hed"
semantics = "semantics"
fullset = "fullset"
Vai = "Vai"
recipes/README.md
+9 -8
View File
@@ -10,7 +10,8 @@ keep naming and layout consistent, organize recipes by model vendor in a way
that is aligned with
[`vllm-project/recipes`](https://github.com/vllm-project/recipes), but treat
that external repository as a reference for structure rather than the place to
add files for this repo. Use one Markdown file per model family by default.
add files for this repo. Use one Markdown file per model and hardware profile,
named `<model>-<hardware>.md`.
> **Note:** The vLLM-Omni TTS recipes (12 models across 8 providers) are also
> published on the rendered recipes site at
@@ -29,16 +30,16 @@ Example layout:
```text
recipes/
Qwen/
Qwen3-Omni.md
Qwen3-TTS.md
Qwen3-Omni-H100.md
Qwen3-TTS-A100.md
Tencent/
Covo-Audio-Chat.md
Covo-Audio-Chat-L40S.md
```
## Available Recipes
| Recipe | Task | Hardware |
|--------|------|----------|
| -------- | ------ | ---------- |
| [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
| [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs |
| [`Boogu/Boogu-Image.md`](./Boogu/Boogu-Image.md) | Text-to-image online serving (Boogu-Image-0.1-Base) | 1x A100/H100 40GB+ |
@@ -62,6 +63,7 @@ recipes/
| [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB |
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) |
| [`SandAI/MAGI-2-preview-L20X.md`](./SandAI/MAGI-2-preview-L20X.md) | Native Preview-stage text/image-to-video-and-audio generation | 1x/4x NVIDIA L20X 140 GiB qualified |
| [`Robbyant/LingBot-World-2.0.md`](./Robbyant/LingBot-World-2.0.md) | Offline and experimental realtime interactive world generation | 1x H200/B200 or 2x B200 for TP=2 |
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
| [`cosmos3/Cosmos3-Edge.md`](./cosmos3/Cosmos3-Edge.md) | T2I / T2V / I2V + action (Physical-AI) generation on the Nemotron-based Edge transformer | 1x GPU (~8 GiB) |
@@ -92,8 +94,7 @@ recipes/
| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
| [`JD/JoyAI-VL-Interaction.md`](./JD/JoyAI-VL-Interaction.md) | Real-time streaming video-language interaction (proactive speak/silence/delegate) | 1x GPU 24GB+ |
Within a single recipe file, include different hardware support sections such
as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like
`1x A100 80GB` or `2x L40S` inside those sections when applicable.
Keep one accelerator model or platform per recipe file. Add a separately
suffixed recipe and table row for another GPU, ROCm accelerator, or NPU.
See [TEMPLATE.md](./TEMPLATE.md) for the recommended format.
recipes/SandAI/MAGI-2-preview-L20X.md
+288
View File
@@ -0,0 +1,288 @@
# MAGI-2 Preview — NVIDIA L20X
> Native text/image-to-video-and-audio generation with the released Preview stage
## Summary
- Vendor: SandAI
- Model: `sand-ai/MAGI-2-preview`
- Runtime: Native vLLM-Omni pipeline; no SandAI runtime import
- Modes: Offline shared examples and OpenAI-compatible video serving
- Hardware: NVIDIA L20X 140 GiB
- Recommended deployment: Four-GPU resident SP4 (`TP=1`, `SP=4`)
- Maintainer: Community
## Supported model contract
### Tasks and inputs
| Task | Text prompt | Media condition | Shared entrypoint |
| --- | --- | --- | --- |
| T2VA | Required | None | [`text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py) |
| I2VA | Required | Exactly one still image | [`image_to_video.py`](../../examples/offline_inference/image_to_video/image_to_video.py) |
The I2VA path preserves the source aspect ratio and applies MAGI-2's native
resize/pad transform. Video or audio conditions, multiple images, and fused
per-rank request batches are not supported.
### Output
| Category | Supported Preview specification |
| --- | --- |
| Duration | Exactly 10 seconds |
| Video | 125 frames at 12.5 fps |
| Native tiers | `272p` (448x256) and `540p` (896x512) |
| Optional final resize | `output_width` and `output_height` must be supplied together |
| Audio | 44.1 kHz stereo, generated with the video |
| Refiner | The separate 1080p refiner is not included in this integration |
### Profiles provided by this recipe
| Profile | Devices | Purpose | Qualification |
| --- | ---: | --- | --- |
| Resident SP4 | 4 | Recommended fidelity/default deployment | Released smoke and reference parity |
| Rank-local DLO SP4 | 4 | Full-quality lower-HBM transformer streaming | Released checkpoint, 100-step T2VA and I2VA |
| DLO DP4 / DP2SP2 | 4 | Concurrent request throughput | Bounded one- and four-step coverage |
| Ordinary layerwise | 1 | Memory-constrained single request | Released checkpoint, one and four steps |
| Online resident SP4 | 4 | OpenAI-compatible video API | Serving path covered |
## References
- Official model card: <https://huggingface.co/sand-ai/MAGI-2-preview>
- Architecture blog: <https://sand.ai/blog/magi-2-preview>
- Supported models: [`docs/models/supported_models.md`](../../docs/models/supported_models.md)
- Diffusion feature matrix: [`docs/user_guide/diffusion_features.md`](../../docs/user_guide/diffusion_features.md)
## Prepare the checkpoint
The native Preview path uses approximately 274 GiB on disk. The separate
refiner is not required.
```bash
export MAGI2_CKPT_ROOT=/path/to/MAGI-2-preview
hf download sand-ai/MAGI-2-preview \
preview/ text_encoder/ vae/ turbo_vae/ stable-audio-open-1.0/ \
--revision 2dea51b64db47ee5b4402d36fd90829a0c58913b \
--local-dir "$MAGI2_CKPT_ROOT"
```
Install vLLM-Omni and ensure `ffmpeg` is on `PATH`. No second model runtime is
required.
## Hardware support
### Locally qualified NVIDIA configuration
| Item | Value |
| --- | --- |
| Accelerator | NVIDIA L20X |
| Per-device memory | 143,771 MiB (about 140.4 GiB) |
| Device count | 4; the one-device profile uses one of the same GPUs |
| Device interconnect | All-to-all NV18 links (`nvidia-smi topo -m`) |
| Host memory | 2.0 TiB |
The official SandAI runtime requires eight NVIDIA Hopper GPUs. That official
eight-GPU system was unavailable for this PR; compatible eight-worker
topologies pass configuration validation but are not locally runtime-qualified.
No H100, H200, B200, or B300 runtime claim is made here.
### Software environment
| Item | Value |
| --- | --- |
| OS | Ubuntu 22.04.5 LTS, Linux 5.10.134 |
| Python | 3.12.13 |
| NVIDIA driver / CUDA runtime | 570.133.20 / 13.0 |
| PyTorch | 2.11.0+cu130 |
| vLLM | 0.26.0 |
| vLLM-Omni | This PR branch; evidence commits `a5af2a8c` and `42bbfb57` |
| Precision | BF16 |
Set `MAGI2_DETERMINISTIC=1` before worker startup when deterministic kernels
are required. The setting is fixed for the worker lifetime.
Each worker process supports one MAGI-2 pipeline instance. Pipeline startup
sets process-wide deterministic state, so multi-pipeline tests and deployments
must use separate worker processes rather than constructing multiple instances
inside one process.
MAGI-2 uses vLLM's bundled FlashAttention 3 kernels by default on Hopper
(compute capability 9.x), FlashAttention 4 on supported Blackwell devices
with FA3 as the fallback, and FlashAttention 2 on earlier architectures. Set
`MAGI2_FLASH_ATTN_VERSION=2`, `3`, or `4` before worker startup to force an
available kernel for diagnostics or output comparisons.
## Offline generation
### Recommended four-GPU T2VA
```bash
export CUDA_VISIBLE_DEVICES=0,1,2,3
python examples/offline_inference/text_to_video/text_to_video.py \
--model "$MAGI2_CKPT_ROOT" \
--model-class-name Magi2Pipeline \
--prompt "A red fox walks through fresh snow while wind moves the pine branches." \
--height 512 --width 896 --num-frames 125 \
--num-inference-steps 100 --fps 12.5 \
--tensor-parallel-size 1 --ulysses-degree 4 \
--extra-body '{"seconds":10,"resolution":"540p"}' \
--output magi2_540p_t2va.mp4
```
For a load-and-kernel smoke, use one inference step. A one-step output is not a
quality evaluation.
### Recommended four-GPU I2VA
Use the same geometry and parallelism with the shared I2V entrypoint:
```bash
python examples/offline_inference/image_to_video/image_to_video.py \
--model "$MAGI2_CKPT_ROOT" \
--model-class-name Magi2Pipeline \
--image /path/to/first_frame.png \
--prompt "The fox looks up, then walks forward as snow falls around it." \
--height 512 --width 896 --num-frames 125 \
--num-inference-steps 100 --fps 12.5 \
--tensor-parallel-size 1 --ulysses-degree 4 \
--extra-body '{"seconds":10,"resolution":"540p"}' \
--output magi2_540p_i2va.mp4
```
### One-device layerwise offload
MAGI-2 already stages Qwen, the image VAE, TurboVAE, and Oobleck from pinned
CPU memory. Ordinary layerwise offload streams the Preview DiT blocks so the
complete transformer does not need to fit in HBM at once.
```bash
export CUDA_VISIBLE_DEVICES=0
python examples/offline_inference/text_to_video/text_to_video.py \
--model "$MAGI2_CKPT_ROOT" \
--model-class-name Magi2Pipeline \
--prompt "A red fox walks through fresh snow while wind moves the pine branches." \
--height 256 --width 448 --num-frames 125 \
--num-inference-steps 4 --fps 12.5 \
--enable-cpu-offload --enable-layerwise-offload \
--extra-body '{"seconds":10,"resolution":"272p"}' \
--output magi2_272p_cpu_layerwise.mp4
```
Standalone `--enable-cpu-offload` is rejected because staging whole modules
does not make the complete Preview DiT fit on one qualified GPU. When both
flags are supplied, the shared offloader selects layerwise mode and MAGI's
native auxiliary staging remains active.
## Online serving
The OpenAI-compatible serving defaults are the full `540p` preset and `100`
denoising steps. The model always generates a 10-second, 125-frame clip at
12.5 fps; pass `resolution=272p` and `num_inference_steps=1` for a low-cost
load smoke. Requests with another duration are rejected before generation.
```bash
export CUDA_VISIBLE_DEVICES=0,1,2,3
vllm serve "$MAGI2_CKPT_ROOT" --omni \
--model-class-name Magi2Pipeline \
--num-gpus 4 \
--tensor-parallel-size 1 \
--ulysses-degree 4 \
--port 8091
```
```bash
curl -X POST http://localhost:8091/v1/videos/sync \
-F 'prompt=A red fox walks through fresh snow.' \
-F 'seconds=10' \
-F 'size=896x512' \
-F 'num_frames=125' \
-F 'fps=12.5' \
-F 'num_inference_steps=100' \
-F 'seed=42' \
-F 'extra_params={"resolution":"540p"}' \
-o magi2_online.mp4
```
For I2VA, add `-F 'input_reference=@first_frame.png;type=image/png'`.
## Supported features
Use the shared guides for launch syntax and feature semantics. The flags below
are the MAGI-2-specific qualified combinations.
| Feature | MAGI-2 status / topology | Guide |
| --- | --- | --- |
| Sequence parallel | Resident SP4 default; compatible SP8 configuration | [Sequence parallel](../../docs/user_guide/diffusion/parallelism/sequence_parallel.md) |
| Tensor parallel | TP4 or TP2SP2 on four workers | [Tensor parallel](../../docs/user_guide/diffusion/parallelism/tensor_parallel.md) |
| DLO | DP4/DP2SP2 AllGather; SP4 rank-local requires `--dlo-no-use-allgather` | [Distributed layerwise offload](../../docs/design/feature/distributed_layerwise_offload.md) |
| HSDP | HSDP4+SP4; alternative to TP and DLO | [HSDP](../../docs/user_guide/diffusion/parallelism/hsdp.md) |
| CFG parallel | CFG2xSP2 on four workers; CFG2xSP4 requires eight | [CFG parallel](../../docs/user_guide/diffusion/parallelism/cfg_parallel.md) |
| VAE patch parallel | TurboVAE tile decode across the complete worker group | [VAE parallelism](../../docs/user_guide/diffusion/parallelism/vae_parallelism.md) |
| Cache-DiT | Repeated Preview transformer layers; approximate | [Cache-DiT](../../docs/user_guide/diffusion/cache_acceleration/cache_dit.md) |
| One-device offload | Ordinary layerwise DiT plus native auxiliary CPU staging | [CPU offload](../../docs/user_guide/diffusion/cpu_offload.md) |
| Quantization | Not supported | [Feature matrix](../../docs/user_guide/diffusion_features.md) |
| Ring / pipeline parallel | Not supported | [Feature matrix](../../docs/user_guide/diffusion_features.md) |
DP request waves must contain exactly `data_parallel_size` requests with the
same explicit step count. Data-parallel replicas require TP=1. SP-only DLO
cannot AllGather because SP ranks own different MoE-head shards.
## Request fields
Common geometry and sampling values use the shared CLI flags. The native
Preview pipeline accepts these model-specific `--extra-body` fields:
| Field | Meaning |
| --- | --- |
| `seconds` | Must be `10`. |
| `resolution` | `272p` or `540p`; default `540p`. |
| `output_width`, `output_height` | Optional final resize; supply both. |
| `deterministic` | Must match `MAGI2_DETERMINISTIC` at worker startup. |
Use `--image` in the shared I2V entrypoint rather than the lower-level
`image_path` field.
## Verification
```bash
ffprobe -v error \
-show_entries stream=codec_type,width,height,avg_frame_rate,nb_frames,sample_rate,channels \
-of json magi2_540p_t2va.mp4
```
A 540p result contains 125 896x512 frames at 12.5 fps and stereo 44.1 kHz
audio. A 272p result contains 125 448x256 frames with the same duration, rate,
and audio contract.
## Local qualification evidence
### Released-checkpoint E2E
| Profile | Workload | Steps | E2E | Peak HBM | Peak host PSS | Output |
| --- | --- | ---: | ---: | ---: | ---: | --- |
| DLO SP4 | 540p T2VA | 100 | 563.855 s | See PR evidence | See PR evidence | `3437d078e...f16fdea` |
| DLO SP4 | 540p I2VA | 100 | 556.533 s | See PR evidence | See PR evidence | `b8c2a4ac...ff2de8` |
| Layerwise 1 GPU | 272p T2VA | 1 | 8.86 s | 49.29 GiB | Not sampled | Valid video + audio |
| Layerwise 1 GPU | 272p T2VA | 4 | 42.55 s | 49.29 GiB | 371.97 GiB | Valid video + audio |
The one-device four-step host-PSS sampler scans the full process tree and adds
observable CPU overhead, so its stage timing is qualification evidence rather
than a latency benchmark.
### Distributed correctness
- Resident TP4, TP2SP2, and SP4 match a single-rank native oracle within the
documented BF16 tolerance; resident SP4 is the reference-aligned default.
- Deterministic rank-local DLO SP4 output matched resident SP4 exactly.
- HSDP4+SP4, HSDP4+CFG2xSP2, Cache-DiT forced hits, and TurboVAE PP4 have
focused four-rank coverage. TurboVAE PP4 matched resident decode exactly.
- Full 100-step 540p quality generation is qualified only for four-device DLO
SP4. Eight-worker layouts remain configuration-only on this host.
These are bounded qualification runs, not committed benchmarks. Generated
artifacts stay outside the repository, and this integration adds no
model-specific example or benchmark script.
recipes/TEMPLATE.md
+40 -36
View File
@@ -1,6 +1,10 @@
# Recipe Title
# Recipe Title — Hardware
> Example: Qwen3-Omni for speech chat
> Example: Qwen3-Omni for speech chat — H100
Name the file `recipes/<vendor>/<model>-<hardware>.md`. Keep exactly one
hardware profile in each recipe; add a separate suffixed recipe when another
accelerator is qualified.
## Summary
@@ -8,28 +12,44 @@
- Model:
- Task:
- Mode:
- Hardware:
- Maintainer:
## When to use this recipe
Briefly describe the concrete scenario this recipe covers.
## Supported model contract
Put user-visible constraints before commands. Use compact tables for:
- supported tasks and shared entrypoints;
- input modalities, counts, duration/geometry limits, and required combinations;
- output duration, geometry, frame/audio rates, and variants;
- the deployment profiles provided on this recipe's named hardware.
Source model specifications from the canonical model card or repository. Do
not copy values from a sibling model or infer support from configuration alone.
## References
- Upstream or canonical docs:
- Related example under `examples/`:
- Related issue or discussion:
## Hardware Support
## Hardware
Add one section per platform, such as `GPU`, `ROCm`, or `NPU`. Under each
platform section, document one or more tested hardware configurations.
- Accelerator model and per-device memory:
- Number of devices:
- Device interconnect (NVLink, PCIe, or other):
- Host memory, when CPU offload or staging is relevant:
- Qualification scope:
## GPU
Do not add a second accelerator to this file. Create another
`<model>-<hardware>.md` recipe instead. Mention unqualified upstream hardware
only as a limitation, not as a provided profile.
### 1x A100 80GB
#### Environment
## Software environment
- OS:
- Python:
@@ -37,7 +57,10 @@ platform section, document one or more tested hardware configurations.
- vLLM version:
- vLLM-Omni version or commit:
#### Command
Keep worker counts, precision, and TP/SP/DP/PP sizes with the command or
deployment profile; they are not software-environment properties.
## Command
Serve commands use `vllm serve <model> --omni`. The `--omni` flag is what
selects the Omni pipeline, so it is required and cannot be dropped.
@@ -46,40 +69,21 @@ selects the Omni pipeline, so it is required and cannot be dropped.
# Add the exact command(s) here
```
#### Verification
## Verification
```bash
# Add a quick validation command or expected output here
```
#### Notes
## Notes
- Memory usage:
- Key flags:
- Known limitations:
### 2x L40S
## Supported features
Repeat the same structure for other hardware setups as needed.
## ROCm
### Example hardware configuration
Repeat the same nested structure for ROCm setups as needed:
- `#### Environment`
- `#### Command`
- `#### Verification`
- `#### Notes`
## NPU
### Example hardware configuration
Repeat the same nested structure for NPU setups as needed:
- `#### Environment`
- `#### Command`
- `#### Verification`
- `#### Notes`
Use a compact table containing the model-specific status/topology and a link to
the corresponding shared feature guide. Do not repeat the general feature
documentation in the recipe. Record unsupported combinations in the same table
instead of a disconnected restrictions section.
tests/config/test_environment_variables.py
+4 -4
View File
@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
EnvironmentVariableCategory.PUBLIC_OMNI: 24,
EnvironmentVariableCategory.INHERITED_VLLM: 20,
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
EnvironmentVariableCategory.MODEL_SPECIFIC: 57,
EnvironmentVariableCategory.MODEL_SPECIFIC: 62,
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
EnvironmentVariableCategory.INTERNAL: 2,
}
@@ -178,10 +178,10 @@ def test_inventory_matches_reviewed_snapshot_counts():
if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
)
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
ModelEnvironmentVariableDisposition.PROMOTE: 35,
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
ModelEnvironmentVariableDisposition.PROMOTE: 37,
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 7,
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
ModelEnvironmentVariableDisposition.INTERNALIZE: 13,
ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5,
}
tests/diffusion/attention/test_attention_config.py
+1
View File
@@ -907,6 +907,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
forward=lambda *args: pytest.fail("unexpected SDPA fallback"),
),
_assert_metadata_compatible=lambda metadata: None,
_has_custom_attention=False,
_scheduler_paged_kv=False,
paged_kv_cache_role=None,
)
tests/diffusion/attention/test_flash_attn.py
+25 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Test script for FlashAttention backend with padding handling.
@@ -29,6 +29,30 @@ HAS_FLASH_ATTN = fa.HAS_FLASH_ATTN
flash_attn_func = fa.flash_attn_func # noqa: N813
@pytest.mark.parametrize(
("device_major", "requested", "supported", "expected"),
[
(8, None, frozenset((2, 3, 4)), 2),
(9, None, frozenset((2, 3, 4)), 3),
(10, None, frozenset((2, 3, 4)), 4),
(12, None, frozenset((2, 3, 4)), 4),
(9, None, frozenset((2,)), 2),
(10, None, frozenset((2, 3)), 3),
(9, "2", frozenset((2, 3, 4)), 2),
(9, "4", frozenset((2, 3, 4)), 4),
],
)
def test_versioned_flash_attention_selection(device_major, requested, supported, expected):
assert fa._choose_vllm_flash_attn_version(device_major, requested, supported) == expected
def test_versioned_flash_attention_selection_rejects_invalid_or_unavailable_version():
with pytest.raises(ValueError, match="must be 2, 3, or 4"):
fa._choose_vllm_flash_attn_version(9, "5", frozenset((2, 3, 4)))
with pytest.raises(RuntimeError, match="FlashAttention 4 is unavailable"):
fa._choose_vllm_flash_attn_version(9, "4", frozenset((2, 3)))
def create_attention_mask(batch_size: int, seq_len: int, valid_len: int, device: torch.device) -> torch.Tensor:
"""
Create attention mask where first valid_len tokens are valid (1) and rest are padding (0).
tests/diffusion/cache/test_cache_backends.py
Vendored
+5 -2
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Unit tests for cache backends (cache-dit and teacache).
@@ -164,7 +164,10 @@ class TestCacheDiTBackend:
num_inference_steps=20,
verbose=True,
)
mock_cache_dit.summary.assert_called_once_with(transformer, details=True)
mock_cache_dit.summary.assert_called_once_with(
mock_block_adapter.return_value,
details=True,
)
@patch("vllm_omni.diffusion.cache.cachedit.backend.logger")
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
tests/diffusion/cache/test_cache_dit.py
Vendored
+41 -1
View File
@@ -59,6 +59,7 @@ def test_custom_cache_dit_enablers_are_registered_explicitly():
"Cosmos3OmniDiffusersPipeline": cd_model_specific.enable_cache_for_cosmos3,
"Cosmos3OmniPipeline": cd_model_specific.enable_cache_for_cosmos3,
"Krea2Pipeline": cd_model_specific.enable_cache_for_krea2,
"Magi2Pipeline": cd_model_specific.enable_cache_for_magi2,
}
with patch.dict(cd_backend.CUSTOM_DIT_ENABLERS, {}, clear=True):
@@ -90,12 +91,50 @@ def test_cosmos3_aliases_use_cosmos3_custom_cache_dit_enabler(pipeline_name: str
assert cd_backend.CUSTOM_DIT_ENABLERS[pipeline_name] is cd_model_specific.enable_cache_for_cosmos3
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.enable_cache_for_dit")
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
def test_magi2_cache_dit_targets_only_nested_repeated_layers(mock_block_adapter, mock_enable_cache):
pipeline = Mock()
transformer_block = pipeline.transformer.block
layers = torch.nn.ModuleList([torch.nn.Identity()])
transformer_block.layers = layers
adapter = mock_block_adapter.return_value
refresh = Mock()
mock_enable_cache.return_value = refresh
result = cd_model_specific.enable_cache_for_magi2(pipeline, SAMPLE_CACHE_CONFIG)
mock_block_adapter.assert_called_once()
adapter_kwargs = mock_block_adapter.call_args.kwargs
assert adapter_kwargs["transformer"] is transformer_block
assert adapter_kwargs["blocks"] == [layers]
assert adapter_kwargs["has_separate_cfg"] is False
assert adapter_kwargs["check_forward_pattern"] is True
assert result.refresh is refresh
assert result.targets == (adapter,)
get_transformer = mock_enable_cache.call_args.kwargs["get_pipeline_transformer"]
assert get_transformer(pipeline) is transformer_block
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit.summary")
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter.is_cached", return_value=True)
def test_cache_summary_uses_custom_nested_targets(mock_is_cached, mock_summary):
target = object()
pipeline = Mock(_cache_dit_targets=(target,))
cd_backend.cache_summary(pipeline, details=True)
mock_is_cached.assert_called_once_with(target)
mock_summary.assert_called_once_with(target, details=True)
def test_cachedit_public_api_is_explicit():
assert set(cd_backend.__all__) == {
"BagelCachedAdapter",
"CUSTOM_DIT_ENABLERS",
"CacheDiTAdapterConfig",
"CacheDiTBackend",
"CacheDiTEnableResult",
"CacheDiTConfig",
"CacheDiTRequestSpec",
"RequestScopedCacheDiTRuntime",
@@ -213,7 +252,7 @@ def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter):
current_omni_platform.is_rocm(),
reason="vLLM ROCm custom ops lack CPU fallback",
)
def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group):
def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group, request: pytest.FixtureRequest):
"""CacheDiT Pattern_0 treats the second positional arg as encoder_hidden_states,
which is a collision for one of the kwargs in LTX2 since we treat the audio
hidden states as encoder_hidden_states.
@@ -253,6 +292,7 @@ def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group):
pipeline.transformer = model
backend = CacheDiTBackend(DiffusionCacheConfig())
backend.enable(pipeline)
request.addfinalizer(lambda: backend.disable(pipeline))
backend.refresh(pipeline, num_inference_steps=5)
# Wrap call_Fn_blocks in CacheDiT so that we can verify the
tests/diffusion/distributed/test_hsdp.py
+47
View File
@@ -12,6 +12,7 @@ from torch.distributed.tensor import DeviceMesh, DTensor
from tests.helpers.runtime import get_distributed_init_method
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.distributed import hsdp as hsdp_module
from vllm_omni.diffusion.distributed.hsdp import (
HSDPInferenceConfig,
_unshardable_parameters,
@@ -73,6 +74,52 @@ def test_hsdp_keeps_packed_and_scalar_parameters_local(cpu_process_group):
assert not isinstance(model.block.input_global_scale, DTensor)
def test_hsdp_resolves_nested_ignored_module_and_preserves_mixed_dtypes(
cpu_process_group,
monkeypatch: pytest.MonkeyPatch,
):
class Group:
world_size = 1
rank_in_group = 0
device_group = dist.group.WORLD
class MixedDtypeModel(nn.Module):
def __init__(self):
super().__init__()
self.layers = nn.Module()
self.layers.sharded = nn.Linear(2, 2, dtype=torch.bfloat16)
self.layers.ignored = nn.Linear(2, 2, dtype=torch.float32)
self._hsdp_shard_conditions = [lambda name, _module: name == "layers.sharded"]
self._hsdp_ignored_modules = ["layers.ignored"]
self._hsdp_preserve_parameter_dtypes = True
group = Group()
monkeypatch.setattr(hsdp_module, "get_world_group", lambda: group)
monkeypatch.setattr(
hsdp_module,
"_create_hsdp_mesh",
lambda **_kwargs: DeviceMesh("cpu", [0]),
)
model = MixedDtypeModel()
ignored_weight = model.layers.ignored.weight
hsdp_module.apply_hsdp_to_model(
model,
HSDPInferenceConfig(
enabled=True,
hsdp_shard_size=1,
param_dtype=torch.float16,
),
target_device=torch.device("cpu"),
)
assert model.layers.ignored.weight is ignored_weight
assert not isinstance(model.layers.ignored.weight, DTensor)
assert model.layers.ignored.weight.dtype == torch.float32
assert isinstance(model.layers.sharded.weight, DTensor)
assert model.layers.sharded.weight.dtype == torch.bfloat16
class TestHSDPInferenceConfig:
"""Tests for HSDPInferenceConfig dataclass."""
tests/diffusion/models/ltx2/test_ltx2_transformer.py
+1
View File
@@ -199,6 +199,7 @@ def _fake_ltx_parallel_attention():
layer.attention = SimpleNamespace(forward=run("native"))
layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa"))
layer.backend_pref = "FLASH_ATTN"
layer._has_custom_attention = False
return layer, calls
tests/diffusion/models/magi2/test_audio_decoder.py
+56
View File
@@ -0,0 +1,56 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import pytest
import torch
from vllm_omni.diffusion.models.magi2.audio_decoder import (
convert_stable_audio_decoder_key,
convert_stable_audio_decoder_state_dict,
)
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@pytest.mark.parametrize(
("source", "target"),
[
("pretransform.model.decoder.layers.0.weight_v", "conv1.weight_v"),
("pretransform.model.decoder.layers.1.layers.0.alpha", "block.0.snake1.alpha"),
("pretransform.model.decoder.layers.3.layers.1.weight_g", "block.2.conv_t1.weight_g"),
(
"pretransform.model.decoder.layers.5.layers.4.layers.2.beta",
"block.4.res_unit3.snake2.beta",
),
(
"pretransform.model.decoder.layers.2.layers.2.layers.3.weight_v",
"block.1.res_unit1.conv2.weight_v",
),
("pretransform.model.decoder.layers.6.alpha", "snake1.alpha"),
("pretransform.model.decoder.layers.7.weight_g", "conv2.weight_g"),
],
)
def test_stable_audio_decoder_key_conversion(source: str, target: str):
assert convert_stable_audio_decoder_key(source) == target
def test_stable_audio_decoder_conversion_filters_and_reshapes_snake():
alpha = torch.arange(4, dtype=torch.float32)
weight = torch.ones(2, 4, 7)
state = {
"pretransform.model.decoder.layers.6.alpha": alpha,
"pretransform.model.decoder.layers.7.weight_v": weight,
"model.model.transformer.layers.0.weight": torch.empty(1),
}
converted = convert_stable_audio_decoder_state_dict(state)
assert set(converted) == {"snake1.alpha", "conv2.weight_v"}
assert converted["snake1.alpha"].shape == (1, 4, 1)
torch.testing.assert_close(converted["snake1.alpha"].flatten(), alpha)
assert converted["conv2.weight_v"] is weight
def test_unknown_decoder_key_fails_closed():
with pytest.raises(ValueError, match="Unsupported Stable Audio decoder key"):
convert_stable_audio_decoder_key("pretransform.model.decoder.future.weight")
tests/diffusion/models/magi2/test_native_distributed_parity.py
+229
View File
@@ -0,0 +1,229 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import os
import tempfile
from contextlib import ExitStack, contextmanager
from unittest.mock import patch
import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import vllm_omni.diffusion.models.magi2.attention as attention_module
import vllm_omni.diffusion.models.magi2.layers as layers_module
import vllm_omni.diffusion.models.magi2.mh_moe as mh_moe_module
import vllm_omni.diffusion.models.magi2.parallel as parallel_module
from vllm_omni.diffusion.models.magi2.attention import VarlenHandler
from vllm_omni.diffusion.models.magi2.configuration_magi2 import (
Magi2MHCConfig,
Magi2MoEConfig,
Magi2PreviewConfig,
)
from vllm_omni.diffusion.models.magi2.modeling_magi2 import (
Magi2PreviewTransformer,
Modality,
)
from vllm_omni.diffusion.models.magi2.parallel import Magi2ParallelGroup
_WORLD_SIZE = 4
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
def _tiny_config() -> Magi2PreviewConfig:
# Four attention, KV, and MoE heads make every head axis divisible by each
# requested TP/SP layout. The MLP and shared-expert widths are likewise
# divisible by TP=4, so this exercises all native row/column collectives.
return Magi2PreviewConfig(
num_layers=1,
hidden_size=32,
head_dim=8,
num_query_groups=4,
video_in_channels=4,
audio_in_channels=4,
text_in_channels=4,
intermediate_factor=1.5,
multimodal_layers=(0,),
params_dtype=torch.float32,
mhc=Magi2MHCConfig(num_streams=2),
moe=Magi2MoEConfig(
num_heads=4,
num_experts=2,
top_k=2,
expert_intermediate_size=8,
shared_expert_intermediate_size=8,
modality_shared_expert_intermediate_size=8,
layers=(0,),
),
)
def _initialize_model(model: Magi2PreviewTransformer) -> None:
generator = torch.Generator(device="cpu").manual_seed(71)
with torch.no_grad():
for name, parameter in model.named_parameters():
if name == "pre_adapter.rope.bands":
continue
parameter.copy_(
torch.randn(
parameter.shape,
dtype=parameter.dtype,
generator=generator,
)
* 0.025
)
moe = model.block.layers[0].mlp.moe_mlp
moe.router.expert_bias.zero_()
moe.router.expert_bias_ema.zero_()
def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, VarlenHandler]:
generator = torch.Generator(device="cpu").manual_seed(113)
packed = torch.randn(8, 4, generator=generator) * 0.2
coordinates = torch.tensor(
[
[0, 0, 0, 2, 2, 2, 2, 2, 2],
[0, 1, 1, 2, 2, 2, 2, 2, 2],
[0, 0, 0, 1, 2, 2, 1, 2, 2],
[0, 1, 1, 1, 2, 2, 1, 2, 2],
[1, 0, 0, 2, 2, 2, 2, 2, 2],
[1, 1, 1, 2, 2, 2, 2, 2, 2],
[0, 0, 0, 1, 1, 1, 1, 1, 1],
[0, 0, 0, 1, 1, 1, 1, 1, 1],
],
dtype=torch.float32,
)
modalities = torch.tensor(
[
Modality.VIDEO,
Modality.AUDIO,
Modality.TEXT,
Modality.TIME,
Modality.VIDEO,
Modality.AUDIO,
Modality.TEXT,
Modality.TIME,
]
)
cumulative = torch.tensor([0, packed.shape[0]], dtype=torch.int32)
varlen = VarlenHandler(cumulative, cumulative, packed.shape[0], packed.shape[0])
return packed, coordinates, modalities, varlen
@contextmanager
def _patched_groups(
tp_group: Magi2ParallelGroup,
sp_group: Magi2ParallelGroup,
):
# TP owns native weight and MoE-head shards when enabled. SP owns Ulysses
# tokens and is also the MoE-head group in the released SP-only layout.
ep_group = tp_group if tp_group.world_size > 1 else sp_group
with ExitStack() as stack:
stack.enter_context(patch.object(layers_module, "get_magi2_tp_group", return_value=tp_group))
stack.enter_context(patch.object(mh_moe_module, "get_magi2_ep_group", return_value=ep_group))
stack.enter_context(patch.object(parallel_module, "get_magi2_ulysses_group", return_value=sp_group))
stack.enter_context(patch.object(attention_module, "get_magi2_ulysses_group", return_value=sp_group))
yield
def _current_group(
rank: int,
rank_groups: tuple[tuple[tuple[int, ...], dist.ProcessGroup], ...],
) -> Magi2ParallelGroup:
for ranks, process_group in rank_groups:
if rank in ranks:
return Magi2ParallelGroup(
process_group,
world_size=len(ranks),
rank=ranks.index(rank),
)
return Magi2ParallelGroup(None, world_size=1, rank=0)
def _new_groups(
rank_sets: tuple[tuple[int, ...], ...],
) -> tuple[tuple[tuple[int, ...], dist.ProcessGroup], ...]:
groups = []
for ranks in rank_sets:
process_group = dist.new_group(ranks=list(ranks), backend="gloo")
groups.append((ranks, process_group))
return tuple(groups)
def _distributed_worker(rank: int, rendezvous: str) -> None:
torch.set_num_threads(1)
dist.init_process_group(
"gloo",
init_method=rendezvous,
rank=rank,
world_size=_WORLD_SIZE,
)
try:
singleton = Magi2ParallelGroup(None, world_size=1, rank=0)
with _patched_groups(singleton, singleton):
oracle = Magi2PreviewTransformer(_tiny_config())
_initialize_model(oracle)
checkpoint = [(name, value.detach().clone()) for name, value in oracle.state_dict().items()]
inputs = _inputs()
with torch.no_grad():
expected = oracle(*inputs)
tp2_groups = _new_groups(((0, 1), (2, 3)))
sp2_groups = _new_groups(((0, 2), (1, 3)))
world_group = Magi2ParallelGroup(
dist.group.WORLD,
world_size=_WORLD_SIZE,
rank=rank,
)
layouts = (
("tp4", world_group, singleton),
("tp2sp2", _current_group(rank, tp2_groups), _current_group(rank, sp2_groups)),
("sp4", singleton, world_group),
)
for layout_name, tp_group, sp_group in layouts:
tp_group = Magi2ParallelGroup(
tp_group.group,
tp_group.world_size,
tp_group.rank,
replicated_sequence=True,
)
with _patched_groups(tp_group, sp_group):
model = Magi2PreviewTransformer(_tiny_config())
assert model.load_weights(checkpoint) == set(model.state_dict())
with torch.no_grad():
actual = model(*inputs)
max_abs_error = (actual - expected).abs().max()
dist.all_reduce(max_abs_error, op=dist.ReduceOp.MAX)
close = torch.tensor(
int(torch.allclose(actual, expected, atol=2e-5, rtol=2e-4)),
dtype=torch.int32,
)
dist.all_reduce(close, op=dist.ReduceOp.MIN)
if not close.item():
raise AssertionError(
f"{layout_name} differs from the single-rank oracle; "
f"global max absolute error={max_abs_error.item():.8g}"
)
dist.barrier()
finally:
dist.destroy_process_group()
@pytest.mark.skipif(
not dist.is_available() or not dist.is_gloo_available(),
reason="requires torch.distributed gloo",
)
def test_four_rank_tp_sp_layouts_match_single_rank_oracle() -> None:
with tempfile.TemporaryDirectory() as temp_dir:
rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}"
mp.spawn(
_distributed_worker,
args=(rendezvous,),
nprocs=_WORLD_SIZE,
join=True,
)
tests/diffusion/models/magi2/test_native_dlo_allgather.py
+294
View File
@@ -0,0 +1,294 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import json
import os
import tempfile
from contextlib import ExitStack, contextmanager
from dataclasses import dataclass
from unittest.mock import patch
import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
from safetensors.torch import load_file, save_file
import vllm_omni.diffusion.models.magi2.layers as layers_module
import vllm_omni.diffusion.models.magi2.mh_moe as mh_moe_module
from vllm_omni.diffusion.model_loader.host_weight_plan import (
build_checkpoint_mmap_plan,
)
from vllm_omni.diffusion.models.magi2.configuration_magi2 import (
Magi2MHCConfig,
Magi2MoEConfig,
Magi2PreviewConfig,
)
from vllm_omni.diffusion.models.magi2.modeling_magi2 import (
Magi2PreviewTransformer,
)
from vllm_omni.diffusion.models.magi2.parallel import Magi2ParallelGroup
from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline
from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
DistributedLayerwiseOffloadBackend,
DistributedLayerwiseOffloadHook,
)
@dataclass(frozen=True)
class _WeightSourceStub:
model_or_path: str
subfolder: str
revision: str | None
prefix: str
@dataclass(frozen=True)
class _PipelineModulesStub:
dits: list[nn.Module]
dit_names: list[str]
_WORLD_SIZE = 4
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
def _tiny_config() -> Magi2PreviewConfig:
return Magi2PreviewConfig(
# The released checkpoint has 40 layers. Keep this tiny fixture at two
# so it exercises the shared backend's streamable-block contract.
num_layers=2,
hidden_size=16,
head_dim=8,
num_query_groups=2,
video_in_channels=4,
audio_in_channels=4,
text_in_channels=4,
intermediate_factor=2,
multimodal_layers=(0, 1),
params_dtype=torch.float32,
mhc=Magi2MHCConfig(num_streams=2),
moe=Magi2MoEConfig(
num_heads=2,
num_experts=4,
top_k=2,
expert_intermediate_size=8,
shared_expert_intermediate_size=8,
modality_shared_expert_intermediate_size=8,
layers=(0, 1),
),
)
@contextmanager
def _patched_weight_groups(ep_group: Magi2ParallelGroup):
singleton = Magi2ParallelGroup(None, world_size=1, rank=0)
with ExitStack() as stack:
stack.enter_context(patch.object(layers_module, "get_magi2_tp_group", return_value=singleton))
stack.enter_context(patch.object(mh_moe_module, "get_magi2_ep_group", return_value=ep_group))
yield
def _initialize_checkpoint(model: Magi2PreviewTransformer) -> None:
generator = torch.Generator(device="cpu").manual_seed(211)
with torch.no_grad():
for parameter in model.parameters():
parameter.copy_(
torch.randn(
parameter.shape,
dtype=parameter.dtype,
generator=generator,
)
* 0.025
)
# The normal loader intentionally mirrors EMA into the runtime router
# bias. Keep both released checkpoint entries equal so this test can
# compare every reconstructed tensor, not only inference-visible ones.
moe = model.block.layers[0].mlp.moe_mlp
moe.router.expert_bias.copy_(moe.router.expert_bias_ema)
def _write_checkpoint(checkpoint_root: str) -> None:
singleton = Magi2ParallelGroup(None, world_size=1, rank=0)
with _patched_weight_groups(singleton):
source = Magi2PreviewTransformer(_tiny_config())
_initialize_checkpoint(source)
weights = {name: tensor.detach().cpu().contiguous() for name, tensor in source.state_dict().items()}
preview_dir = os.path.join(checkpoint_root, "preview")
os.makedirs(preview_dir)
weight_file = os.path.join(preview_dir, "model.safetensors")
save_file(weights, weight_file)
with open(os.path.join(preview_dir, "model.safetensors.index.json"), "w") as index:
json.dump(
{"weight_map": {name: os.path.basename(weight_file) for name in weights}},
index,
)
def _new_groups(
rank_sets: tuple[tuple[int, ...], ...],
) -> tuple[tuple[tuple[int, ...], dist.ProcessGroup], ...]:
return tuple((ranks, dist.new_group(ranks=list(ranks), backend="gloo")) for ranks in rank_sets)
def _current_group(
rank: int,
rank_groups: tuple[tuple[tuple[int, ...], dist.ProcessGroup], ...],
) -> tuple[tuple[int, ...], dist.ProcessGroup]:
return next((ranks, group) for ranks, group in rank_groups if rank in ranks)
def _load_mmap_transform_and_reconstruct(
checkpoint_root: str,
dp_group: dist.ProcessGroup,
dp_size: int,
dp_rank: int,
ep_group: Magi2ParallelGroup,
) -> None:
checkpoint_file = os.path.join(checkpoint_root, "preview", "model.safetensors")
checkpoint = load_file(checkpoint_file, device="cpu")
with _patched_weight_groups(ep_group):
oracle = Magi2PreviewTransformer(_tiny_config())
assert oracle.load_weights(checkpoint.items()) == set(oracle.state_dict())
with torch.device("meta"):
target = Magi2PreviewTransformer(_tiny_config())
target.requires_grad_(False)
pipeline = Magi2Pipeline.__new__(Magi2Pipeline)
nn.Module.__init__(pipeline)
pipeline.checkpoint_root = checkpoint_root
pipeline.transformer = target
plan_result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", target),),
sources=(
_WeightSourceStub(
model_or_path=checkpoint_root,
subfolder="preview",
revision=None,
prefix="transformer.",
),
),
model_path=checkpoint_root,
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert plan_result.fallback_reason is None
assert plan_result.plan is not None
backend = object.__new__(DistributedLayerwiseOffloadBackend)
backend.device = torch.device("cpu")
backend._using_rank_local_mmap = False
backend._mmap_transforms_by_tensor_id = {}
backend._load_weights_via_mmap(
pipeline,
_PipelineModulesStub(dits=[target], dit_names=["transformer"]),
plan_result.plan,
)
target_block = target.block.layers[0]
oracle_block = oracle.block.layers[0]
target_parameters = dict(target_block.named_parameters())
oracle_parameters = dict(oracle_block.named_parameters())
# The mmap loader binds the released full-head tensor. The DLO shard step
# must apply the model-declared EP transform before it computes the DP
# shard, so DP peers gather identical SP-local layouts in DP2SP2.
moe_name = "mlp.moe_mlp.gate"
mmap_moe = target_parameters[moe_name]
mmap_transform = backend._mmap_transforms_by_tensor_id[id(mmap_moe)]
assert callable(mmap_transform)
assert tuple(mmap_moe.shape) == tuple(checkpoint[f"block.layers.0.{moe_name}"].shape)
assert tuple(mmap_transform(mmap_moe).shape) == tuple(oracle_parameters[moe_name].shape)
# Keep this gloo regression accelerator-isolated. Pinned allocation is the
# final, layout-preserving operation in _shard_and_pin and has separate
# backend coverage; enabling it here would initialize CUDA in every worker.
cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin(
target_parameters,
dict(target_block.named_buffers()),
dp_size=dp_size,
rank=dp_rank,
pin_memory=False,
tensor_transforms=backend._mmap_transforms_by_tensor_id,
)
assert all(not shard.requires_grad for shard in cpu_shards.values())
assert all(parameter.numel() == 0 for parameter in target_block.parameters())
for dtype, local_shard in cpu_shards.items():
gathered = torch.empty(
local_shard.numel() * dp_size,
dtype=dtype,
device="cpu",
)
dist.all_gather_into_tensor(gathered, local_shard, group=dp_group)
for tensor_metadata in metadata[dtype]:
name = tensor_metadata["name"]
offset = tensor_metadata["offset"]
numel = tensor_metadata["numel"]
reconstructed = gathered[offset : offset + numel].view(tensor_metadata["shape"])
torch.testing.assert_close(reconstructed, oracle_parameters[name])
def _distributed_worker(rank: int, rendezvous: str, checkpoint_root: str) -> None:
torch.set_num_threads(1)
dist.init_process_group(
"gloo",
init_method=rendezvous,
rank=rank,
world_size=_WORLD_SIZE,
)
try:
dp2_groups = _new_groups(((0, 2), (1, 3)))
sp2_groups = _new_groups(((0, 1), (2, 3)))
_load_mmap_transform_and_reconstruct(
checkpoint_root,
dist.group.WORLD,
dp_size=4,
dp_rank=rank,
ep_group=Magi2ParallelGroup(None, world_size=1, rank=0),
)
dist.barrier()
dp_ranks, dp_group = _current_group(rank, dp2_groups)
sp_ranks, sp_group = _current_group(rank, sp2_groups)
_load_mmap_transform_and_reconstruct(
checkpoint_root,
dp_group,
dp_size=2,
dp_rank=dp_ranks.index(rank),
ep_group=Magi2ParallelGroup(
sp_group,
world_size=2,
rank=sp_ranks.index(rank),
),
)
dist.barrier()
finally:
dist.destroy_process_group()
@pytest.mark.skipif(
not dist.is_available() or not dist.is_gloo_available(),
reason="requires torch.distributed gloo",
)
def test_dlo_mmap_transform_dp_shard_and_allgather_for_dp4_and_dp2sp2() -> None:
with tempfile.TemporaryDirectory() as temp_dir:
_write_checkpoint(temp_dir)
rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}"
with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}):
mp.spawn(
_distributed_worker,
args=(rendezvous, temp_dir),
nprocs=_WORLD_SIZE,
join=True,
)
tests/diffusion/models/magi2/test_native_preview.py
+275
View File
@@ -0,0 +1,275 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import cache_dit
import pytest
import torch
from vllm_omni.diffusion.cache.cachedit.model_specific import enable_cache_for_magi2
from vllm_omni.diffusion.data import DiffusionCacheConfig
from vllm_omni.diffusion.models.magi2.attention import VarlenHandler
from vllm_omni.diffusion.models.magi2.configuration_magi2 import (
Magi2MHCConfig,
Magi2MoEConfig,
Magi2PreviewConfig,
)
from vllm_omni.diffusion.models.magi2.layers import MultiModalityRMSNorm
from vllm_omni.diffusion.models.magi2.mh_moe import Magi2MultiHeadMoE
from vllm_omni.diffusion.models.magi2.modeling_magi2 import (
Magi2PreviewTransformer,
Modality,
)
from vllm_omni.diffusion.models.magi2.preview_data_proxy import (
Magi2DataProxy,
Magi2PreviewDataProxyConfig,
ModelInput,
)
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
def _tiny_config(
params_dtype: torch.dtype = torch.float32,
*,
num_layers: int = 1,
) -> Magi2PreviewConfig:
layer_indices = tuple(range(num_layers))
return Magi2PreviewConfig(
num_layers=num_layers,
hidden_size=16,
head_dim=8,
num_query_groups=2,
video_in_channels=4,
audio_in_channels=4,
text_in_channels=4,
intermediate_factor=2,
multimodal_layers=layer_indices,
params_dtype=params_dtype,
mhc=Magi2MHCConfig(num_streams=2),
moe=Magi2MoEConfig(
num_heads=2,
num_experts=4,
top_k=2,
expert_intermediate_size=8,
shared_expert_intermediate_size=8,
modality_shared_expert_intermediate_size=8,
layers=layer_indices,
),
)
def _initialize_tiny_model(model: Magi2PreviewTransformer, seed: int) -> None:
generator = torch.Generator(device="cpu").manual_seed(seed)
with torch.no_grad():
for parameter in model.parameters():
parameter.copy_(
torch.randn(
parameter.shape,
generator=generator,
dtype=parameter.dtype,
)
* 0.02
)
for module in model.modules():
if isinstance(module, MultiModalityRMSNorm):
module.weight.zero_()
elif isinstance(module, Magi2MultiHeadMoE):
module.router.expert_bias.zero_()
module.router.expert_bias_ema.zero_()
def test_hsdp_policy_shards_individual_preview_layers():
model = Magi2PreviewTransformer(_tiny_config(num_layers=2))
condition = model._hsdp_shard_conditions[0]
assert condition("block.layers.0", model.block.layers[0])
assert condition("block.layers.1", model.block.layers[1])
assert not condition("block", model.block)
assert not condition("pre_adapter", model.pre_adapter)
def test_data_proxy_keeps_output_layout_request_scoped() -> None:
proxy = Magi2DataProxy(Magi2PreviewDataProxyConfig(time_channel_dim=0))
def make_input(width: int) -> ModelInput:
return ModelInput(
x_t=torch.arange(2 * width, dtype=torch.float32).reshape(1, 2, 1, 1, width),
audio_x_t=torch.zeros(1, 1, 2),
audio_feat_len=torch.tensor([1]),
txt_feat=torch.zeros(1, 1, 2),
txt_feat_len=torch.tensor([1]),
t=torch.tensor([0.5]),
)
first = proxy.process_input(make_input(2))
second = proxy.process_input(make_input(3))
first_video, first_audio = proxy.process_output(first.token_sequence, first.output_layout)
second_video, second_audio = proxy.process_output(second.token_sequence, second.output_layout)
assert first_video.shape == (1, 2, 1, 1, 2)
assert second_video.shape == (1, 2, 1, 1, 3)
assert first_audio.shape == second_audio.shape == (1, 1, 2)
def test_tiny_native_preview_runs_through_nested_cachedit_adapter() -> None:
model = Magi2PreviewTransformer(_tiny_config(num_layers=3))
_initialize_tiny_model(model, seed=11)
pipeline = type("Magi2TestPipeline", (), {"transformer": model})()
result = enable_cache_for_magi2(
pipeline,
DiffusionCacheConfig(
Fn_compute_blocks=1,
Bn_compute_blocks=0,
max_warmup_steps=1,
residual_diff_threshold=1.0,
),
)
result.refresh(pipeline, 4, verbose=False)
packed = torch.randn(6, 4)
coordinates = torch.ones(6, 9)
modalities = torch.tensor(
[
Modality.VIDEO,
Modality.VIDEO,
Modality.AUDIO,
Modality.AUDIO,
Modality.TEXT,
Modality.TEXT,
]
)
cumulative = torch.tensor([0, 6], dtype=torch.int32)
varlen = VarlenHandler(cumulative, cumulative, 6, 6)
execution_counts = [0, 0, 0]
hooks = [
layer.register_forward_hook(
lambda _module, _args, _output, layer_index=index: execution_counts.__setitem__(
layer_index, execution_counts[layer_index] + 1
)
)
for index, layer in enumerate(model.block.layers)
]
try:
with torch.no_grad():
first = model(packed, coordinates, modalities, varlen)
second = model(packed, coordinates, modalities, varlen)
assert torch.isfinite(first).all()
assert torch.isfinite(second).all()
# The first layer is the configured Fn block and always executes. The
# zero-residual repeated input hits the cache on the second call, so
# the two middle layers are skipped.
assert execution_counts == [2, 1, 1]
finally:
for hook in hooks:
hook.remove()
cache_dit.disable_cache(result.targets[0])
assert not getattr(model.block, "_is_cached", False)
def test_tiny_native_preview_matches_pinned_reference_golden() -> None:
"""Full-model golden from SandAI reference f68a0f9bbccb.
The reference's unavailable compiler/FA3/MoE CUDA boundaries were replaced
by their eager PyTorch equations when generating this tensor. Keeping the
resulting golden local makes this regression independent of that runtime.
"""
model = Magi2PreviewTransformer(_tiny_config(torch.bfloat16))
with torch.no_grad():
trainable_parameters = (parameter for parameter in model.parameters() if parameter.requires_grad)
for index, parameter in enumerate(trainable_parameters):
values = (torch.arange(parameter.numel(), dtype=torch.float32) % 17 - 8) * 0.002 + (index % 5 - 2) * 0.0001
parameter.copy_(values.reshape(parameter.shape).to(parameter.dtype))
moe = model.block.layers[0].mlp.moe_mlp
moe.router.expert_bias.zero_()
moe.router.expert_bias_ema.zero_()
packed = torch.tensor(
[
[0.1, -0.2, 0.3, -0.4],
[0.5, 0.6, -0.7, 0.8],
[-0.9, 1.0, -0.1, 0.2],
[0.3, -0.4, 0.5, -0.6],
[0.7, 0.8, 0.9, 1.0],
[-0.2, 0.4, -0.6, 0.8],
[0.9, -0.7, 0.5, -0.3],
]
)
coordinates = torch.tensor(
[
[0, 0, 0, 2, 2, 2, 2, 2, 2],
[1, 1, 1, 2, 2, 2, 2, 2, 2],
[0, 0, 0, 1, 2, 2, 1, 2, 2],
[0, 1, 1, 1, 2, 2, 1, 2, 2],
[0, 0, 0, 1, 1, 1, 1, 1, 1],
[0, 0, 0, 1, 1, 1, 1, 1, 1],
[0, 0, 0, 1, 1, 1, 1, 1, 1],
],
dtype=torch.float32,
)
modalities = torch.tensor(
[
Modality.VIDEO,
Modality.VIDEO,
Modality.AUDIO,
Modality.AUDIO,
Modality.TEXT,
Modality.TEXT,
Modality.TIME,
]
)
cumulative = torch.tensor([0, 4, 7], dtype=torch.int32)
time_tokens = torch.tensor(
[
[0.1, 0.2, 0.3],
[0.4, 0.5, 0.6],
[0.7, 0.8, 0.9],
[1.0, 0.9, 0.8],
[0.7, 0.6, 0.5],
[0.4, 0.3, 0.2],
[0.1, 0.0, -0.1],
]
)
expected = torch.tensor(
[
[-0.093289732933, 0.133922040462, 0.149612516165, 0.069642566144],
[-0.118302434683, 0.129068359733, 0.156465008855, 0.084789708257],
[-0.114637844265, 0.137497439981, 0.159646511078, 0.088356778026],
[-0.104430131614, 0.146924808621, 0.152607530355, 0.092736266553],
[0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0],
]
)
varlen = VarlenHandler(cumulative, cumulative, 4, 4)
with torch.no_grad():
actual = model(packed, coordinates, modalities, varlen, time_tokens)
torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5)
def test_strict_loader_covers_all_keys_and_uses_ema_router_bias() -> None:
source = Magi2PreviewTransformer(_tiny_config())
_initialize_tiny_model(source, seed=11)
source_moe = source.block.layers[0].mlp.moe_mlp
source_moe.router.expert_bias.fill_(2.0)
source_moe.router.expert_bias_ema.fill_(3.0)
checkpoint = [(name, value.detach().clone()) for name, value in source.state_dict().items()]
target = Magi2PreviewTransformer(_tiny_config())
loaded = target.load_weights(checkpoint)
assert loaded == set(target.state_dict())
target_moe = target.block.layers[0].mlp.moe_mlp
torch.testing.assert_close(
target_moe.router.expert_bias,
torch.full_like(target_moe.router.expert_bias, 3.0),
)
torch.testing.assert_close(
target_moe.router.expert_bias_ema, torch.full_like(target_moe.router.expert_bias_ema, 3.0)
)
tests/diffusion/models/magi2/test_pipeline_magi2.py
+631
View File
@@ -0,0 +1,631 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
from dataclasses import dataclass, field
from types import MethodType
from typing import Any
from unittest.mock import Mock
import numpy as np
import pytest
import torch
import torch.nn as nn
from PIL import Image
from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
from vllm_omni.diffusion.models.magi2.pipeline_magi2 import (
MAGI2_AUDIO_SAMPLE_RATE,
MAGI2_MODEL_REVISION,
Magi2Pipeline,
_Magi2StagedComponent,
_resolve_checkpoint_root,
_validate_native_topology,
)
from vllm_omni.diffusion.models.magi2.preview_data_proxy import Magi2DataProxy
from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig, Magi2PreviewSampler
from vllm_omni.diffusion.registry import DiffusionModelRegistry
from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
from vllm_omni.errors import OmniClientError
from vllm_omni.model_extras.registry import (
get_extra_body_params,
should_preserve_reference_image_size,
)
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@dataclass
class _SamplingStub:
width: int | None = None
height: int | None = None
fps: float | None = None
frame_rate: float | None = None
resolved_frame_rate: float | None = None
num_frames: int = 1
num_inference_steps: int | None = None
num_outputs_per_prompt: int = 1
seed: int | None = 42
generator: torch.Generator | list[torch.Generator] | None = None
extra_args: dict[str, Any] = field(default_factory=dict)
@dataclass
class _RequestStub:
prompts: list[Any]
sampling_params: _SamplingStub
num_reqs: int = 1
@dataclass
class _ParallelStub:
pipeline_parallel_size: int = 1
data_parallel_size: int = 1
tensor_parallel_size: int = 1
sequence_parallel_size: int = 1
ulysses_degree: int = 1
ring_degree: int = 1
allgather_degree: int = 1
cfg_parallel_size: int = 1
vae_patch_parallel_size: int = 1
text_encoder_tp_size: int = 1
enable_expert_parallel: bool = False
use_hsdp: bool = False
@dataclass
class _TopologyStub:
parallel_config: _ParallelStub = field(default_factory=_ParallelStub)
enable_cpu_offload: bool = False
enable_layerwise_offload: bool = False
enable_distributed_layerwise_offload: bool = False
dlo_use_allgather: bool = True
quantization_config: object = None
cache_backend: str = "none"
custom_pipeline_args: dict[str, object] = field(default_factory=dict)
additional_config: dict[str, object] = field(default_factory=lambda: {"magi2_allow_unsupported_topology": True})
class _FakeNativeRuntime:
def __init__(self) -> None:
self.calls: list[dict] = []
def evaluate(self, **kwargs):
self.calls.append(kwargs)
video = np.zeros((2, 8, 12, 3), dtype=np.uint8)
audio = np.zeros((32, 2), dtype=np.float32)
return video, audio
def _pipeline() -> tuple[Magi2Pipeline, _FakeNativeRuntime]:
pipe = Magi2Pipeline.__new__(Magi2Pipeline)
nn.Module.__init__(pipe)
pipe.deterministic = False
pipe.enable_diffusion_pipeline_profiler = False
runtime = _FakeNativeRuntime()
def evaluate(_self, **kwargs):
return runtime.evaluate(**kwargs)
pipe._evaluate_preview = MethodType(evaluate, pipe)
return pipe, runtime
def _request(prompt, **sampling_overrides):
sampling = _SamplingStub()
for key, value in sampling_overrides.items():
setattr(sampling, key, value)
return _RequestStub(prompts=[prompt], sampling_params=sampling)
def _checkpoint_tree(root) -> None:
files = (
"preview/model.safetensors.index.json",
"text_encoder/config.json",
"text_encoder/model.safetensors.index.json",
"vae/Wan2.2_VAE.pth",
"turbo_vae/TurboV3-Wan22-TinyShallow_7_7.json",
"turbo_vae/checkpoint.ckpt",
"stable-audio-open-1.0/model_config.json",
"stable-audio-open-1.0/model.safetensors",
)
for relative in files:
path = root / relative
path.parent.mkdir(parents=True, exist_ok=True)
path.touch()
def test_official_id_detection_and_metadata():
assert is_diffusion_model("sand-ai/MAGI-2-preview")
assert is_diffusion_model("https://huggingface.co/sand-ai/MAGI-2-preview")
assert resolve_model_class_name("sand-ai/MAGI-2-preview") == "Magi2Pipeline"
assert DiffusionModelRegistry._try_load_model_cls("Magi2Pipeline") is Magi2Pipeline
metadata = get_diffusion_model_metadata("Magi2Pipeline")
assert metadata.supports_multimodal_inputs
assert metadata.max_multimodal_image_inputs == 1
assert should_preserve_reference_image_size(
"Magi2Pipeline",
model="sand-ai/MAGI-2-preview",
)
assert {"seconds", "resolution"} <= get_extra_body_params("Magi2Pipeline")
def test_local_preview_signature_detection_without_refiner(tmp_path):
_checkpoint_tree(tmp_path)
assert is_diffusion_model(str(tmp_path))
assert resolve_model_class_name(str(tmp_path)) == "Magi2Pipeline"
config = OmniDiffusionConfig(model=str(tmp_path))
config.enrich_config()
assert config.model_class_name == "Magi2Pipeline"
assert config.supports_multimodal_inputs
def test_huggingface_url_resolves_pinned_snapshot(tmp_path, monkeypatch):
_checkpoint_tree(tmp_path)
api = Mock()
api.snapshot_download.side_effect = lambda *, repo_id, revision: (
str(tmp_path)
if (repo_id, revision) == ("sand-ai/MAGI-2-preview", MAGI2_MODEL_REVISION)
else pytest.fail("unexpected snapshot request")
)
monkeypatch.setattr(
"vllm.transformers_utils.repo_utils.hf_api",
lambda: api,
)
assert _resolve_checkpoint_root(
"https://huggingface.co/sand-ai/MAGI-2-preview",
None,
) == str(tmp_path.resolve())
def test_forward_uses_native_540p_preview_defaults(monkeypatch):
pipe, runtime = _pipeline()
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
lambda: False,
)
result = pipe(_request("A fox walks through snow"))
call = runtime.calls[0]
assert call["prompt"] == "A fox walks through snow"
assert call["image"] is None
assert (call["width"], call["height"]) == (896, 512)
assert call["num_inference_steps"] == 100
assert result.output["metadata"]["video"]["fps"] == 12.5
assert result.output["metadata"]["audio"]["sample_rate"] == MAGI2_AUDIO_SAMPLE_RATE
assert result.stage_durations["magi2_preview_e2e"] >= 0
def test_forward_maps_272p_i2v_and_output_resize(monkeypatch):
pipe, runtime = _pipeline()
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
lambda: False,
)
image = Image.new("RGB", (16, 9), "white")
prompt = {
"prompt": "The first frame begins moving",
"multi_modal_data": {"image": image},
}
result = pipe(
_request(
prompt,
width=448,
height=256,
num_inference_steps=1,
extra_args={
"output_width": 6,
"output_height": 4,
},
)
)
call = runtime.calls[0]
assert call["image"] is image
assert pipe._ensure_figure_token(call["prompt"]) == (
"The first frame begins moving\nreference_layer:The first frame refers to <Figure 1>"
)
assert (call["width"], call["height"]) == (448, 256)
assert call["num_inference_steps"] == 1
assert result.output["payload"]["video"].shape == (2, 4, 6, 3)
def test_forward_rejects_conflicting_resolution_and_sampling_geometry(monkeypatch):
pipe, runtime = _pipeline()
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
lambda: False,
)
with pytest.raises(OmniClientError, match="requires 896x512"):
pipe(
_request(
"A fox walks through snow",
width=448,
height=256,
extra_args={"resolution": "540p"},
)
)
assert not runtime.calls
@pytest.mark.parametrize(
("prompt", "extra_args", "message"),
[
("", {}, "non-empty prompt"),
("prompt", {"seconds": 5}, "10-second clips"),
("prompt", {"duration": 5}, "10-second clips"),
("prompt", {"resolution": "720p"}, "Unsupported native"),
("prompt", {"resolution": "1080p"}, "Unsupported native"),
("prompt", {"resolution": "540p", "use_refiner": True}, "refiner"),
("prompt", {"output_width": 0, "output_height": 4}, "positive integers"),
],
)
def test_forward_rejects_invalid_preview_requests(
monkeypatch,
prompt,
extra_args,
message,
):
pipe, runtime = _pipeline()
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
lambda: False,
)
with pytest.raises(OmniClientError, match=message):
pipe(_request(prompt, extra_args=extra_args))
assert not runtime.calls
def test_forward_rejects_zero_inference_steps_before_generation(monkeypatch):
pipe, runtime = _pipeline()
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
lambda: False,
)
with pytest.raises(OmniClientError, match="inference steps must be positive"):
pipe(_request("prompt", num_inference_steps=0))
assert not runtime.calls
def test_forward_stops_peak_monitor_when_initial_synchronize_fails(monkeypatch):
pipe, runtime = _pipeline()
pipe.enable_diffusion_pipeline_profiler = True
monitor = Mock(peak_bytes=0)
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_cuda",
lambda: True,
)
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
lambda: True,
)
monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0)
monkeypatch.setattr(torch.accelerator, "synchronize", Mock(side_effect=RuntimeError("sync failed")))
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2._PeakReservedMonitor",
Mock(return_value=monitor),
)
with pytest.raises(RuntimeError, match="sync failed"):
pipe(_request("A fox walks through snow"))
monitor.start.assert_called_once_with()
monitor.stop.assert_called_once_with()
assert not runtime.calls
def test_forward_does_not_start_peak_monitor_without_profiler(monkeypatch):
pipe, runtime = _pipeline()
monitor_factory = Mock(side_effect=AssertionError("unexpected memory monitor"))
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_cuda",
lambda: True,
)
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
lambda: True,
)
monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0)
monkeypatch.setattr(torch.accelerator, "synchronize", Mock())
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2._PeakReservedMonitor",
monitor_factory,
)
result = pipe(_request("A fox walks through snow"))
monitor_factory.assert_not_called()
assert result.peak_memory_mb == 0.0
assert len(runtime.calls) == 1
def test_forward_rejects_multiple_images(monkeypatch):
pipe, _ = _pipeline()
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
lambda: False,
)
prompt = {
"prompt": "animate",
"multi_modal_data": {"image": [Image.new("RGB", (2, 2)), Image.new("RGB", (2, 2))]},
}
with pytest.raises(OmniClientError, match="at most one input image"):
pipe(_request(prompt))
def test_pipeline_rejects_unknown_initialization_arguments():
with pytest.raises(TypeError, match=r"Unexpected MAGI-2.*unknown_option"):
Magi2Pipeline(None, unknown_option=True)
def _topology_config(**overrides):
parallel = _ParallelStub()
config = _TopologyStub(parallel_config=parallel)
for key, value in overrides.items():
if hasattr(parallel, key):
setattr(parallel, key, value)
else:
setattr(config, key, value)
return config
def test_native_topology_accepts_qualified_tensor_parallelism():
_validate_native_topology(_topology_config(tensor_parallel_size=4))
def test_native_topology_rejects_nondivisible_tensor_parallelism():
config = _topology_config(tensor_parallel_size=3)
with pytest.raises(ValueError, match="does not divide"):
_validate_native_topology(config)
def test_native_topology_accepts_single_device_layerwise_with_cpu_staging():
config = _topology_config(
enable_layerwise_offload=True,
additional_config={},
)
_validate_native_topology(config)
config.enable_cpu_offload = True
_validate_native_topology(config)
config.enable_layerwise_offload = False
with pytest.raises(ValueError, match="Combine --enable-cpu-offload"):
_validate_native_topology(config)
def test_native_topology_requires_dlo_rank_local_mode():
config = _topology_config(enable_distributed_layerwise_offload=True)
with pytest.raises(ValueError, match="dlo-no-use-allgather"):
_validate_native_topology(config)
config.dlo_use_allgather = False
_validate_native_topology(config)
def test_native_topology_accepts_four_device_hsdp_cfg_and_vae_layouts():
_validate_native_topology(
_topology_config(
sequence_parallel_size=4,
ulysses_degree=4,
use_hsdp=True,
)
)
_validate_native_topology(
_topology_config(
sequence_parallel_size=2,
ulysses_degree=2,
cfg_parallel_size=2,
)
)
_validate_native_topology(
_topology_config(
sequence_parallel_size=4,
ulysses_degree=4,
vae_patch_parallel_size=4,
)
)
def test_native_topology_rejects_cfg_data_parallel_combination():
with pytest.raises(ValueError, match="CFG parallelism.*DLO data parallelism"):
_validate_native_topology(_topology_config(cfg_parallel_size=2, data_parallel_size=2))
def test_cfg_parallel_branch_adapter_preserves_packed_cfg_math():
sampler = Magi2PreviewSampler(nn.Identity(), Magi2DataProxy(), device="cpu", dtype=torch.float32)
latent = torch.arange(8, dtype=torch.float32).reshape(1, 2, 2, 1, 2)
audio_latent = torch.arange(6, dtype=torch.float32).reshape(1, 3, 2)
positive_text = torch.ones(1, 3, 4)
negative_text = torch.full((1, 2, 4), -1.0)
packed = sampler.prepare_model_input(
latent=latent,
audio_latent=audio_latent,
txt_feat=positive_text,
null_txt_feat=negative_text,
t=torch.tensor(500.0),
cfg_config=CFGConfig(),
)
positive, negative = sampler._split_cfg_model_input(packed)
assert positive.x_t.shape[0] == negative.x_t.shape[0] == 1
assert packed.ref_video_feat.shape == (2, 2, 0, 1, 2)
assert packed.ref_video_feat.numel() == 0
torch.testing.assert_close(positive.txt_feat[:, :3], positive_text)
torch.testing.assert_close(negative.txt_feat[:, :2], negative_text)
video_pos = torch.full_like(latent, 3.0)
video_neg = torch.full_like(latent, 1.0)
audio_pos = torch.full_like(audio_latent, 4.0)
audio_neg = torch.full_like(audio_latent, 2.0)
expected = sampler.cfg_velocity(
(torch.cat((video_pos, video_neg)), torch.cat((audio_pos, audio_neg))),
5.0,
7.0,
CFGConfig(),
latent,
audio_latent,
)
actual = sampler.combine_cfg_noise(
(video_pos, audio_pos),
(video_neg, audio_neg),
1.0,
kwargs={
"video_txt_guidance_scale": 5.0,
"audio_txt_guidance_scale": 7.0,
"cfg_config": CFGConfig(),
"latent": latent,
"audio_latent": audio_latent,
},
)
torch.testing.assert_close(actual[0], expected[0])
torch.testing.assert_close(actual[1], expected[1])
def test_decode_audio_preserves_batch_for_released_preview_shape(monkeypatch):
class ShapeCheckingDecoder(nn.Module):
def __init__(self) -> None:
super().__init__()
self.seen_shape: tuple[int, ...] | None = None
def decode(self, latent: torch.Tensor) -> torch.Tensor:
self.seen_shape = tuple(latent.shape)
if self.seen_shape != (1, 64, 250):
raise AssertionError(f"unexpected decoder input shape {self.seen_shape}")
return torch.zeros(1, 2, 512_000)
decoder = ShapeCheckingDecoder()
stager = Mock()
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2.PinnedModuleStager",
Mock(return_value=stager),
)
component = _Magi2StagedComponent(
decoder,
torch.device("cpu"),
pin_memory=False,
)
pipeline = object.__new__(Magi2Pipeline)
nn.Module.__init__(pipeline)
pipeline._is_output_rank = True
pipeline._offload_aux_after_use = True
pipeline.device_str = "cpu"
pipeline.audio_decoder = component
def fake_resample(audio: np.ndarray, target_length: int) -> np.ndarray:
assert audio.shape == (512_000, 2)
assert target_length == 441_000
return np.zeros((target_length, 2), dtype=audio.dtype)
monkeypatch.setattr("scipy.signal.resample", fake_resample)
audio = pipeline._decode_audio(torch.zeros(1, 250, 64))
assert decoder.seen_shape == (1, 64, 250)
assert audio is not None
assert audio.shape == (441_000, 2)
stager.load.assert_called_once_with()
stager.offload.assert_called_once_with()
class _TinyMagiTransformer(nn.Module):
_layerwise_offload_blocks_attrs = ["block"]
def __init__(self) -> None:
super().__init__()
self.block = nn.ModuleList([nn.Linear(2, 2) for _ in range(3)])
class _NoEagerMoveComponent(_Magi2StagedComponent):
def to(self, *args, **kwargs):
del args, kwargs
raise AssertionError("staged MAGI-2 component was moved eagerly")
def _offload_pipeline(monkeypatch):
stagers = [Mock() for _ in range(4)]
stager_factory = Mock(side_effect=stagers)
monkeypatch.setattr(
"vllm_omni.platforms.current_omni_platform.Stream",
Mock(return_value=Mock()),
)
monkeypatch.setattr(
"vllm_omni.diffusion.models.magi2.pipeline_magi2.PinnedModuleStager",
stager_factory,
)
pipeline = object.__new__(Magi2Pipeline)
nn.Module.__init__(pipeline)
pipeline.transformer = _TinyMagiTransformer()
components = [
_NoEagerMoveComponent(
nn.Linear(2, 2),
torch.device("cuda:0"),
pin_memory=False,
)
for _ in range(4)
]
(
pipeline.text_encoder,
pipeline.image_vae,
pipeline.video_decoder,
pipeline.audio_decoder,
) = components
return pipeline, stagers
class _FakeLayerwiseHook:
def __init__(self) -> None:
self._prev_hook = None
self.current_slot = 0
self.prefetch_layer = Mock()
self.get_weights = Mock()
def test_dlo_no_allgather_enable_preserves_staged_aux_and_streams_blocks(
monkeypatch,
):
from vllm_omni.diffusion.offloader.base import (
OffloadConfig,
OffloadStrategy,
)
from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
DistributedLayerwiseOffloadBackend,
)
pipeline, stagers = _offload_pipeline(monkeypatch)
hooks: list[_FakeLayerwiseHook] = []
def fake_apply(*args, **kwargs):
del args, kwargs
hook = _FakeLayerwiseHook()
hooks.append(hook)
return hook
monkeypatch.setattr(
"vllm_omni.diffusion.offloader.distributed_layerwise_backend.apply_distributed_block_hook",
fake_apply,
)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
dlo_use_allgather=False,
),
torch.device("cpu"),
)
backend._allocate_shared_buffers = Mock(return_value=[{}, {}])
backend._cleanup_after_loading = Mock()
backend._release_mmap_handles = Mock()
backend.enable(pipeline)
assert backend.enabled
assert backend.dp_size == 1
assert len(hooks) == len(pipeline.transformer.block)
assert all(stager.offload.call_count == 1 for stager in stagers)
assert all(stager.load.call_count == 0 for stager in stagers)
hooks[-1].prefetch_layer.assert_called_once_with(
slot=hooks[0].current_slot,
non_blocking=False,
)
tests/diffusion/models/magi2/test_turbo_vae.py
+66
View File
@@ -0,0 +1,66 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import pytest
import torch
import torch.nn as nn
from vllm_omni.diffusion.models.magi2.turbo_vae import (
Magi2TurboVAEDecoder,
extract_turbo_decoder_state_dict,
turbo_unpatchify,
)
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
def test_turbo_unpatchify_matches_wan_channel_order():
# Four patch channels for one output channel. MAGI's order interleaves
# width before height when reconstructing each 2x2 patch.
patched = torch.arange(4, dtype=torch.float32).view(1, 4, 1, 1, 1)
output = turbo_unpatchify(patched, patch_size=2)
expected = torch.tensor([[[[[0.0, 2.0], [1.0, 3.0]]]]])
torch.testing.assert_close(output, expected)
def test_extract_turbo_decoder_prefers_ema_and_drops_training_heads():
decoder_weight = torch.ones(1)
checkpoint = {
"state_dict": {"module.decoder.old.weight": torch.zeros(1)},
"ema_state_dict": {
"module.decoder.conv_in.conv.weight": decoder_weight,
"module.aligned_feature_projection_heads.0.weight": torch.empty(1),
},
}
assert extract_turbo_decoder_state_dict(checkpoint) == {"decoder.conv_in.conv.weight": decoder_weight}
def test_turbo_temporal_tiles_preserve_chunk_overlap_contract():
class RepeatDecoder(nn.Module):
def forward(self, value: torch.Tensor, *, is_first_chunk: bool) -> torch.Tensor:
del is_first_chunk
return value.repeat_interleave(4, dim=2)
decoder = object.__new__(Magi2TurboVAEDecoder)
nn.Module.__init__(decoder)
decoder.z_dim = 1
decoder.first_chunk_size = 3
decoder.step_size = 2
decoder.temporal_compression_ratio = 4
decoder.decoder = RepeatDecoder()
decoder.use_tiling = False
decoder.register_buffer("latent_mean", torch.zeros(1), persistent=False)
decoder.register_buffer("latent_std", torch.ones(1), persistent=False)
latent = torch.arange(7, dtype=torch.float32).view(1, 1, 7, 1, 1)
actual = decoder.decode(latent)
expected = latent.repeat_interleave(4, dim=2)
torch.testing.assert_close(actual, expected)
decoder.latent_mean = torch.zeros(1, device="meta")
decoder.latent_std = torch.ones(1, device="meta")
prepared, _ = decoder._prepare_latent(latent)
assert prepared.device.type == "meta"
tests/e2e/online_serving/test_magi2.py
+77
View File
@@ -0,0 +1,77 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Released-checkpoint MAGI-2 Preview serving smoke on four H100 GPUs.
This is a one-step 272p load/dispatch test, not a quality or performance test.
It exercises the recommended resident SP4 topology through ``/v1/videos``.
From ``tests/``::
pytest -s -v e2e/online_serving/test_magi2.py -m "advanced_model and diffusion" --run-level=advanced_model
"""
import json
import os
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
MODEL = "sand-ai/MAGI-2-preview"
PROMPT = "A red fox walks through fresh snow while wind moves the pine branches."
FOUR_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=4)
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parametrize(
"omni_server",
[
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
"--model-class-name",
"Magi2Pipeline",
"--num-gpus",
"4",
"--tensor-parallel-size",
"1",
"--ulysses-degree",
"4",
],
init_timeout=1800,
),
id="resident_sp4_272p_one_step",
marks=FOUR_CARD_MARKS,
)
],
indirect=True,
)
def test_magi2_preview_serving_smoke(
omni_server: OmniServer,
online_client: OnlineOmniClient,
) -> None:
request_config = {
"model": omni_server.model,
"form_data": {
"prompt": PROMPT,
"height": 256,
"width": 448,
"num_frames": 125,
"fps": 12.5,
"num_inference_steps": 1,
"seed": 42,
"extra_params": json.dumps({"resolution": "272p"}),
},
"expected_audio": {"sample_rate": 44100, "channels": 2},
"fps_tolerance": 0.01,
}
online_client.send_video_diffusion_request(request_config)
tests/entrypoints/openai_api/test_video_server.py
+83 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Unit tests for OpenAI-compatible video generation endpoints.
"""
@@ -635,6 +635,77 @@ def test_video_generation_bridges_request_fields(generation_request, expected_nu
assert sampling.extra_args["duration"] == expected_duration
def test_magi2_i2v_preserves_reference_geometry_for_model_preprocessing(test_client, mocker: MockerFixture):
image_bytes = _make_test_image_bytes((48, 32))
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
return_value=b"fake-video",
)
engine = test_client.app.state.openai_serving_video._engine_client
engine.model_class_name = "Magi2Pipeline"
response = test_client.post(
"/v1/videos",
data={
"prompt": "A bear playing with yarn.",
"width": "96",
"height": "64",
},
files={"input_reference": ("input.png", image_bytes, "image/png")},
)
assert response.status_code == 200
video_id = response.json()["id"]
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
input_image = engine.captured_prompt["multi_modal_data"]["image"]
assert isinstance(input_image, Image.Image)
assert input_image.size == (48, 32)
def test_magi2_serving_applies_native_defaults_and_rejects_explicit_frame_mismatch():
engine = FakeAsyncOmni()
engine.model_class_name = "Magi2Pipeline"
handler = OmniOpenAIServingVideo.for_diffusion(
diffusion_engine=engine,
model_name="sand-ai/MAGI-2-preview",
)
asyncio.run(handler._run_and_extract(VideoGenerationRequest(prompt="A fox walks through snow"), "defaults"))
sampling = engine.captured_sampling_params_list[0]
assert (sampling.width, sampling.height) == (896, 512)
assert sampling.num_frames == 125
assert sampling.fps == sampling.frame_rate == 12.5
assert sampling.num_inference_steps == 100
assert "duration" not in sampling.extra_args
with pytest.raises(HTTPException, match="10-second clips only"):
asyncio.run(
handler._run_and_extract(
VideoGenerationRequest(prompt="A fox walks through snow", seconds="5"),
"bad-duration",
)
)
with pytest.raises(HTTPException, match="10-second clips only"):
asyncio.run(
handler._run_and_extract(
VideoGenerationRequest(
prompt="A fox walks through snow",
extra_params={"duration": 5},
),
"bad-duration-extra",
)
)
with pytest.raises(HTTPException, match="requires 125 frames"):
asyncio.run(
handler._run_and_extract(
VideoGenerationRequest(prompt="A fox walks through snow", num_frames=1),
"bad-frames",
)
)
def test_i2v_video_generation_with_image_reference_form(test_client, mocker: MockerFixture):
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
@@ -1703,6 +1774,11 @@ def test_video_request_validation():
assert req.quality is None
assert req.generate_sound is False
assert req.sound_duration is None
assert VideoGenerationRequest(prompt="test", fps=12.5).resolve_video_params().fps == 12.5
with pytest.raises(ValueError):
VideoGenerationRequest(prompt="test", fps=float("inf"))
with pytest.raises(ValueError):
VideoGenerationRequest(prompt="test", video_params={"fps": float("nan")})
assert VideoGenerationRequest(prompt="test", generate_sound=True, sound_duration=1.5).generate_sound is True
with pytest.raises(ValueError):
VideoGenerationRequest(prompt="test", size="invalid")
@@ -2230,6 +2306,8 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
"/v1/videos/sync",
data={
"prompt": "param pass",
"seconds": "10",
"fps": "12.5",
"num_inference_steps": "30",
"guidance_scale": "6.5",
"seed": "42",
@@ -2243,6 +2321,10 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
assert captured.guidance_scale == 6.5
assert captured.seed == 42
assert captured.quality == "high"
assert captured.num_frames == 125
assert captured.fps == 12.5
assert captured.frame_rate == 12.5
assert captured.extra_args["duration"] == 10.0
def test_sync_sana_wm_extra_params_payload_passes_to_engine_prompt(test_client, mocker: MockerFixture):
tests/examples/offline_inference/test_magi2_shared_defaults.py
+55
View File
@@ -0,0 +1,55 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import importlib.util
from pathlib import Path
from types import ModuleType
import pytest
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
REPO_ROOT = Path(__file__).parents[3]
def _load_example(name: str, relative_path: str) -> ModuleType:
path = REPO_ROOT / relative_path
spec = importlib.util.spec_from_file_location(name, path)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
@pytest.fixture(scope="module")
def text_to_video() -> ModuleType:
return _load_example(
"magi2_text_to_video_example",
"examples/offline_inference/text_to_video/text_to_video.py",
)
@pytest.mark.parametrize(
("extra_body", "expected_size"),
[
(None, (896, 512)),
({"resolution": "540p"}, (896, 512)),
({"resolution": "272p"}, (448, 256)),
],
)
def test_shared_text_to_video_uses_native_preview_defaults(
text_to_video: ModuleType,
extra_body: dict[str, object] | None,
expected_size: tuple[int, int],
) -> None:
preset = text_to_video._detect_preset(
"sand-ai/MAGI-2-preview",
"Magi2Pipeline",
extra_body,
)
assert (preset["width"], preset["height"]) == expected_size
assert preset["num_frames"] == 125
assert preset["num_inference_steps"] == 100
assert preset["fps"] == 12.5
tests/helpers/assertions.py
+56 -4
View File
@@ -269,7 +269,9 @@ def assert_video_diffusion_response(
"height": ...,
"fps": ...,
...
}
},
"expected_audio": {"sample_rate": 44100, "channels": 2},
"fps_tolerance": 0.01,
}
"""
form_data = request_config.get("form_data", {})
@@ -280,7 +282,8 @@ def assert_video_diffusion_response(
expected_frames = _maybe_int(form_data.get("num_frames"))
expected_width = _maybe_int(form_data.get("width"))
expected_height = _maybe_int(form_data.get("height"))
expected_fps = _maybe_int(form_data.get("fps"))
expected_fps = _maybe_float(form_data.get("fps"))
expected_audio = request_config.get("expected_audio")
# Skip num_frames assertion for Helios models because they round up frames
model = request_config.get("model", "")
@@ -294,7 +297,19 @@ def assert_video_diffusion_response(
width=expected_width,
height=expected_height,
fps=expected_fps,
fps_tolerance=request_config.get("fps_tolerance", 1.0),
)
if expected_audio is not None:
assert isinstance(expected_audio, dict), "expected_audio must be an object"
sample_rate = _maybe_int(expected_audio.get("sample_rate"))
channels = _maybe_int(expected_audio.get("channels"))
assert sample_rate is not None, "expected_audio.sample_rate is required"
assert channels is not None, "expected_audio.channels is required"
assert_video_audio_valid(
vid_bytes,
sample_rate=sample_rate,
channels=channels,
)
def assert_video_first_frame_matches(
@@ -346,6 +361,12 @@ def _maybe_int(value: Any) -> int | None:
return int(value)
def _maybe_float(value: Any) -> float | None:
if value is None:
return None
return float(value)
def assert_image_valid(image: Path | Image.Image, *, width: int | None = None, height: int | None = None):
"""Assert the file is a loadable image with optional exact dimensions."""
if isinstance(image, Path):
@@ -367,12 +388,14 @@ def assert_video_valid(
width: int | None = None,
height: int | None = None,
fps: float | None = None,
fps_tolerance: float = 1.0,
) -> dict[str, int | float]:
"""Assert the MP4 has the expected resolution and frame count.
For several diffusion backends, encoded MP4 frame count follows a codec-aligned
convention (e.g. request `num_frames=8` can produce 9 encoded frames). Keep
this compatibility behavior to avoid false negatives in online-serving tests.
Fixed-rate models can request a stricter ``fps_tolerance`` in frames/second.
"""
temp_path = None
cap = None
@@ -408,8 +431,9 @@ def assert_video_valid(
assert actual_width == width, f"Expected width={width}, got {actual_width}"
if height is not None:
assert actual_height == height, f"Expected height={height}, got {actual_height}"
if fps is not None and actual_fps:
assert abs(actual_fps - float(fps)) < 1.0, f"Expected fps~={fps}, got {actual_fps}"
if fps is not None:
assert math.isfinite(actual_fps) and actual_fps > 0, f"Invalid video fps: {actual_fps}"
assert abs(actual_fps - float(fps)) < fps_tolerance, f"Expected fps~={fps}, got {actual_fps}"
if num_frames is not None:
expected_frames = (int(num_frames) // 4) * 4 + 1
assert actual_frames == expected_frames, f"Expected frames={expected_frames}, got {actual_frames}"
@@ -433,6 +457,33 @@ def assert_video_valid(
pass
def assert_video_audio_valid(
video: Path | bytes | BytesIO,
*,
sample_rate: int,
channels: int,
) -> None:
"""Assert that an encoded video carries an audio stream with the requested format."""
if isinstance(video, Path):
source: str | BytesIO = str(video)
elif isinstance(video, bytes):
source = BytesIO(video)
elif isinstance(video, BytesIO):
video.seek(0)
source = video
else:
raise TypeError(f"Unsupported video type: {type(video)}")
with av.open(source) as container:
audio_streams = list(container.streams.audio)
assert audio_streams, "Video response does not contain an audio stream"
audio_stream = audio_streams[0]
actual_sample_rate = int(audio_stream.rate or 0)
actual_channels = int(audio_stream.channels or 0)
assert actual_sample_rate == sample_rate, f"Expected audio sample rate={sample_rate}, got {actual_sample_rate}"
assert actual_channels == channels, f"Expected audio channels={channels}, got {actual_channels}"
def assert_audio_valid(
audio_or_path: Path | np.ndarray,
*,
@@ -1047,5 +1098,6 @@ __all__ = [
"assert_omni_response",
"assert_video_diffusion_response",
"assert_video_valid",
"assert_video_audio_valid",
"assert_audio_valid",
]
tests/helpers/tests/test_video_response_contract.py
+99
View File
@@ -0,0 +1,99 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from fractions import Fraction
from io import BytesIO
from types import SimpleNamespace
import av
import numpy as np
import pytest
from tests.helpers.assertions import assert_video_diffusion_response, assert_video_valid
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _encode_clip(fps=12.5, *, audio_rate=44100, audio_layout="stereo"):
output = BytesIO()
with av.open(output, "w", format="mp4") as container:
video = container.add_stream("mpeg4", rate=Fraction(str(fps)))
video.width = video.height = 32
video.pix_fmt = "yuv420p"
audio = None
if audio_rate is not None:
audio = container.add_stream("aac", rate=audio_rate)
audio.layout = audio_layout
for _ in range(5):
frame = av.VideoFrame.from_ndarray(np.zeros((32, 32, 3), dtype=np.uint8), format="rgb24")
for packet in video.encode(frame):
container.mux(packet)
for packet in video.encode():
container.mux(packet)
if audio is not None:
channels = len(audio.layout.channels)
frame = av.AudioFrame.from_ndarray(
np.zeros((channels, 1024), dtype=np.float32), format="fltp", layout=audio_layout
)
frame.sample_rate = audio_rate
for packet in audio.encode(frame):
container.mux(packet)
for packet in audio.encode():
container.mux(packet)
return output.getvalue()
@pytest.mark.parametrize(
("fps", "audio_rate", "audio_layout", "error"),
[
(12.5, 44100, "stereo", None),
(12.0, 44100, "stereo", "Expected fps"),
(13.0, 44100, "stereo", "Expected fps"),
(12.5, None, "stereo", "does not contain an audio stream"),
(12.5, 48000, "stereo", "Expected audio sample rate"),
(12.5, 44100, "mono", "Expected audio channels"),
],
)
def test_fixed_rate_video_audio_response_contract(fps, audio_rate, audio_layout, error):
response = SimpleNamespace(videos=[_encode_clip(fps, audio_rate=audio_rate, audio_layout=audio_layout)])
request_config = {
"form_data": {"fps": "12.5", "num_frames": 5, "width": 32, "height": 32},
"fps_tolerance": 0.01,
"expected_audio": {"sample_rate": 44100, "channels": 2},
}
if error is None:
assert_video_diffusion_response(response, request_config)
else:
with pytest.raises(AssertionError, match=error):
assert_video_diffusion_response(response, request_config)
def test_video_response_retains_default_fps_tolerance():
response = SimpleNamespace(videos=[_encode_clip(12.0, audio_rate=None)])
assert_video_diffusion_response(response, {"form_data": {"fps": 12.5}})
@pytest.mark.parametrize("reported_fps", [0.0, float("nan"), float("inf")])
def test_video_response_rejects_invalid_reported_fps(monkeypatch, tmp_path, reported_fps):
import cv2
path = tmp_path / "clip.mp4"
path.write_bytes(_encode_clip(audio_rate=None))
real_capture = cv2.VideoCapture
class InvalidFpsCapture:
def __init__(self, path):
self.capture = real_capture(path)
def isOpened(self):
return self.capture.isOpened()
def get(self, prop):
return reported_fps if prop == cv2.CAP_PROP_FPS else self.capture.get(prop)
def release(self):
self.capture.release()
monkeypatch.setattr(cv2, "VideoCapture", InvalidFpsCapture)
with pytest.raises(AssertionError, match="Invalid video fps"):
assert_video_valid(path, fps=12.5, fps_tolerance=0.01)
tests/model_extras/test_model_extras.py
+18 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -194,6 +194,23 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na
assert should_init_extra_args_for_non_diffusion_stages(pipeline_name) is False
@pytest.mark.core_model
@pytest.mark.cpu
def test_magi2_extra_registry_is_preview_only() -> None:
assert get_extra_body_params("Magi2Pipeline") == frozenset(
{
"seconds",
"resolution",
"image_path",
"output_width",
"output_height",
"deterministic",
}
)
assert get_extra_output_params("Magi2Pipeline") == frozenset()
assert should_init_extra_args_for_non_diffusion_stages("Magi2Pipeline") is False
@pytest.mark.core_model
@pytest.mark.cpu
def test_ltx_extra_registry_declares_official_guidance_params() -> None:
tests/model_tests/diffusion/test_alignment.py
+6
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Tests for alignment between _DIFFUSION_MODELS and DIFFUSION_TEST_SETTINGS; if
tests in this file are failing, you are probably adding a new model, and need
@@ -64,6 +67,9 @@ EXCLUDED_MODELS = [
# Requires camera/action assets and persistent AR state; covered by its
# dedicated contract tests and real-checkpoint E2E smoke instead.
"LingBotWorldCausalDMDPipeline",
# The 270+ GiB multimodal checkpoint needs 4/8-way native sequence/head
# parallel execution; covered by focused native contract and GPU tests.
"Magi2Pipeline",
"OmniVoicePipeline",
"OmniVoice",
"Cosmos3OmniDiffusersPipeline",
vllm_omni/config/environment_variable_inventory.py
+5
View File
@@ -153,6 +153,8 @@ _MODEL_PROMOTE = (
"INTERNVLA_A1_COSMOS_DIR",
"INTERNVLA_A1_COSMOS_ENCODER_PATH",
"INTERNVLA_A1_PROCESSOR_DIR",
"MAGI2_DETERMINISTIC",
"MAGI2_FLASH_ATTN_VERSION",
"MIMO_AUDIO_TOKENIZER_CUDA_GRAPH",
"MIMO_AUDIO_TOKENIZER_DEVICE",
"MIMO_AUDIO_TOKENIZER_PATH",
@@ -174,6 +176,7 @@ _MODEL_PROMOTE = (
_MODEL_REQUEST_SCOPE = (
"DOTS_TTS_DIT_NUM_STEPS",
"GR00T_NOISE_SEED",
"MAGI2_NEGATIVE_PROMPT",
"STEP_AUDIO2_DEFAULT_PROMPT_WAV",
"VLLM_GEPARD_END_FADE_MS",
"VLLM_GEPARD_END_SILENCE_MS",
@@ -189,6 +192,8 @@ _MODEL_INTERNALIZE = (
"DOTS_TTS_BETA_TRACE",
"DYNIN_S2U_VENDOR_ROOT",
"DZ_PHASE_TIMING",
"MAGI2_ALLOW_UNSUPPORTED_TOPOLOGY",
"MAGI2_ROUTER_BIAS_SOURCE",
"MING_TTS_STAGE1_FINAL_LOG",
"MOSS_TTS_DEBUG_STOP",
"NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE",
vllm_omni/diffusion/attention/backends/utils/fa.py
+84 -1
View File
@@ -16,7 +16,7 @@
# limitations under the License.
# Adapted from https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_flash_attention_utils.py
from collections.abc import Callable
from functools import lru_cache
from functools import cache, lru_cache
from typing import Any
import torch
@@ -166,6 +166,89 @@ else:
HAS_FLASH_ATTN = flash_attn_func is not None or flash_attn_varlen_func is not None
def _choose_vllm_flash_attn_version(
device_major: int,
requested: str | int | None,
supported_versions: frozenset[int],
) -> int:
"""Choose a vLLM-bundled FA kernel without model-specific policy."""
if requested is not None:
try:
version = int(requested)
except (TypeError, ValueError) as error:
raise ValueError("FlashAttention version must be 2, 3, or 4") from error
if version not in (2, 3, 4):
raise ValueError("FlashAttention version must be 2, 3, or 4")
if version not in supported_versions:
raise RuntimeError(f"FlashAttention {version} is unavailable on this device")
return version
if device_major >= 10 and 4 in supported_versions:
return 4
if device_major >= 9 and 3 in supported_versions:
return 3
if 2 in supported_versions:
return 2
raise RuntimeError("No vLLM-bundled FlashAttention 2/3/4 kernel is available")
@cache
def resolve_vllm_flash_attn_version(requested: str | int | None = None) -> int:
"""Resolve an available vLLM-bundled FA2/FA3/FA4 implementation."""
if not current_omni_platform.is_cuda():
raise RuntimeError("vLLM-bundled versioned FlashAttention requires CUDA")
from vllm.vllm_flash_attn.flash_attn_interface import is_fa_version_supported
capability = current_omni_platform.get_device_capability()
if capability is None:
raise RuntimeError("Cannot resolve FlashAttention without a CUDA device capability")
supported_versions = frozenset(version for version in (2, 3, 4) if is_fa_version_supported(version))
return _choose_vllm_flash_attn_version(capability.major, requested, supported_versions)
def vllm_flash_attn_varlen_with_lse(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
*,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
max_seqlen_q: int,
max_seqlen_k: int,
softmax_scale: float | None = None,
softcap: float = 0.0,
causal: bool = False,
deterministic: bool = False,
fa_version: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Run packed vLLM FlashAttention and retain LSE for post-processing."""
from vllm.vllm_flash_attn import flash_attn_varlen_func as vllm_flash_attn_varlen_func
version = resolve_vllm_flash_attn_version(fa_version)
out, lse = vllm_flash_attn_varlen_func(
q,
k,
v,
int(max_seqlen_q),
cu_seqlens_q,
int(max_seqlen_k),
cu_seqlens_k=cu_seqlens_k,
dropout_p=0.0,
softmax_scale=softmax_scale,
causal=causal,
window_size=None,
softcap=max(float(softcap), 0.0),
deterministic=deterministic,
return_softmax_lse=True,
fa_version=version,
)
return out, lse
@lru_cache(maxsize=1)
def is_flash_attn_installed() -> bool:
"""Return whether a Flash Attention backend package is importable.
vllm_omni/diffusion/attention/layer.py
+101 -76
View File
@@ -46,6 +46,7 @@ def _try_extract_layer_index(prefix: str) -> int | None:
class Attention(nn.Module):
_scheduler_paged_kv = False
_has_custom_attention = False
def __init__(
self,
@@ -74,6 +75,11 @@ class Attention(nn.Module):
# attention remains dense and contributes no native KVCacheSpec.
paged_kv_cache_role: str | None = None,
paged_kv_cache_dtype: torch.dtype | None = None,
# Model-owned kernel for architectures whose attention contract cannot
# be represented by the generic backend interface (for example packed
# varlen attention with learned sink logits). The shared Attention
# layer still owns parallel dispatch and compile boundaries.
custom_attention: nn.Module | None = None,
):
super().__init__()
@@ -93,6 +99,8 @@ class Attention(nn.Module):
self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
self.head_size = head_size
self._has_custom_attention = custom_attention is not None
# Resolve backend via role-aware config.
# The global diffusion config is set during model init via
# set_current_diffusion_config(); no env-var re-parsing needed here.
@@ -108,91 +116,103 @@ class Attention(nn.Module):
model_class_name = getattr(config, "model_class_name", None) if config is not None else None
allow_trtllm_default = get_diffusion_model_metadata(model_class_name).attention_mask_free
attn_backend_cls, spec = get_attn_backend_for_role(
role=role,
head_size=head_size,
attention_config=attention_config,
role_category=role_category,
allow_trtllm_default=allow_trtllm_default,
)
scheduler_paged_kv = (
config is not None
and getattr(config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
is DiffusionKVCacheMode.PAGED_SCHEDULER
)
self._scheduler_paged_kv = scheduler_paged_kv
if (
scheduler_paged_kv
and paged_kv_cache_role is not None
and spec is None
and not attn_backend_cls.supports_paged_kv
):
# FLASH_ATTN is an Omni selector, not a device-specific kernel.
# vLLM resolves it to CUDA FlashAttention on GPU and Ascend native
# attention on NPU. A platform may legitimately default dense
# attention to CUDNN/SDPA (for example when optional dense FA
# extras are absent), but a paged layer must advertise the
# capability so the Worker can register its native cache view.
# Load that selector from the registry: this is not a user-explicit
# dense FLASH_ATTN request, so skip platform explicit-backend
# validation (Blackwell would otherwise require CuTe FA4). Formal
# paged execution still delegates to vLLM's native paged backend.
from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
if custom_attention is None:
attn_backend_cls, spec = get_attn_backend_for_role(
role=role,
head_size=head_size,
attention_config=attention_config,
role_category=role_category,
allow_trtllm_default=allow_trtllm_default,
)
if (
scheduler_paged_kv
and paged_kv_cache_role is not None
and spec is None
and not attn_backend_cls.supports_paged_kv
):
# FLASH_ATTN is an Omni selector, not a device-specific kernel.
# vLLM resolves it to CUDA FlashAttention on GPU and Ascend native
# attention on NPU. A platform may legitimately default dense
# attention to CUDNN/SDPA (for example when optional dense FA
# extras are absent), but a paged layer must advertise the
# capability so the Worker can register its native cache view.
# Load that selector from the registry: this is not a user-explicit
# dense FLASH_ATTN request, so skip platform explicit-backend
# validation (Blackwell would otherwise require CuTe FA4). Formal
# paged execution still delegates to vLLM's native paged backend.
from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
dense_backend_name = attn_backend_cls.get_name()
attn_backend_cls = DiffusionAttentionBackendEnum.FLASH_ATTN.get_class()
logger.info(
"Resolved marked paged diffusion attention role=%r to %r because the platform default %r "
"does not support Scheduler-owned KV",
role,
attn_backend_cls.get_name(),
dense_backend_name,
dense_backend_name = attn_backend_cls.get_name()
attn_backend_cls = DiffusionAttentionBackendEnum.FLASH_ATTN.get_class()
logger.info(
"Resolved marked paged diffusion attention role=%r to %r because the platform default %r "
"does not support Scheduler-owned KV",
role,
attn_backend_cls.get_name(),
dense_backend_name,
)
parallel_config = getattr(config, "parallel_config", None)
allgather_degree = getattr(parallel_config, "allgather_degree", 1)
# TODO: Move AllGather-KV compatibility into an AttentionBackend capability
# so validation does not depend on backend names.
if not skip_sequence_parallel and allgather_degree > 1 and attn_backend_cls.get_name() == "TRTLLM_ATTN":
raise ValueError(
"TRTLLM_ATTN does not support AllGather-KV sequence parallelism. "
"Set --allgather-degree 1 or select another diffusion attention backend."
)
self.attn_spec = spec
if spec is not None:
backend_kwargs = spec.backend_kwargs()
self.backend_pref = spec.backend
self.backend_explicit = True
logger.debug("Attention(role=%s) → backend=%s", role, spec.backend)
else:
# Propagate the resolved platform default so Ring Attention can
# make a compatible automatic selection on the current GPU.
self.backend_pref = attn_backend_cls.get_name()
logger.debug("Attention(role=%s) → platform default (%s)", role, self.backend_pref)
self.attn_backend = attn_backend_cls
self.attn_impl_cls = self.attn_backend.get_impl_cls()
self.attention = self.attn_impl_cls(
num_heads=num_heads,
head_size=head_size,
softmax_scale=softmax_scale,
causal=causal,
num_kv_heads=num_kv_heads,
qkv_layout=qkv_layout,
prefix=prefix,
backend_kwargs=backend_kwargs,
role=role,
backend_explicit=self.backend_explicit,
)
parallel_config = getattr(config, "parallel_config", None)
allgather_degree = getattr(parallel_config, "allgather_degree", 1)
# TODO: Move AllGather-KV compatibility into an AttentionBackend capability
# so validation does not depend on backend names.
if not skip_sequence_parallel and allgather_degree > 1 and attn_backend_cls.get_name() == "TRTLLM_ATTN":
raise ValueError(
"TRTLLM_ATTN does not support AllGather-KV sequence parallelism. "
"Set --allgather-degree 1 or select another diffusion attention backend."
# Some model-specific paths call this helper directly. The main backend
# dispatch below never switches to it implicitly.
self.sdpa_fallback = SDPABackend.get_impl_cls()(
num_heads=num_heads,
head_size=head_size,
softmax_scale=softmax_scale,
causal=causal,
num_kv_heads=num_kv_heads,
qkv_layout=qkv_layout,
)
self.attn_spec = spec
if spec is not None:
backend_kwargs = spec.backend_kwargs()
self.backend_pref = spec.backend
self.backend_explicit = True
logger.debug("Attention(role=%s) → backend=%s", role, spec.backend)
else:
# Propagate the resolved platform default so Ring Attention can
# make a compatible automatic selection on the current GPU.
self.backend_pref = attn_backend_cls.get_name()
logger.debug("Attention(role=%s) → platform default (%s)", role, self.backend_pref)
self.attn_backend = attn_backend_cls
self.attn_impl_cls = self.attn_backend.get_impl_cls()
self.attention = self.attn_impl_cls(
num_heads=num_heads,
head_size=head_size,
softmax_scale=softmax_scale,
causal=causal,
num_kv_heads=num_kv_heads,
qkv_layout=qkv_layout,
prefix=prefix,
backend_kwargs=backend_kwargs,
role=role,
backend_explicit=self.backend_explicit,
)
# Some model-specific paths call this helper directly. The main backend
# dispatch below never switches to it implicitly.
self.sdpa_fallback = SDPABackend.get_impl_cls()(
num_heads=num_heads,
head_size=head_size,
softmax_scale=softmax_scale,
causal=causal,
num_kv_heads=num_kv_heads,
qkv_layout=qkv_layout,
)
if paged_kv_cache_role is not None:
raise ValueError("custom_attention does not support Scheduler-managed paged KV")
if not skip_sequence_parallel:
raise ValueError("custom_attention must own its communication and requires skip_sequence_parallel=True")
self.attn_spec = None
self.attn_backend = None
self.attn_impl_cls = type(custom_attention)
self.attention = custom_attention
self.sdpa_fallback = None
logger.debug("Attention(role=%s) → custom kernel=%s", role, type(custom_attention).__name__)
self.softmax_scale = softmax_scale
self.scatter_idx = scatter_idx
@@ -269,7 +289,7 @@ class Attention(nn.Module):
return self.parallel_strategy
def _init_kv_cache_quantization(self, config) -> None:
if config is None:
if config is None or self._has_custom_attention:
return
dtype = getattr(config, "diffusion_kv_cache_dtype", None)
if dtype == "auto":
@@ -467,6 +487,9 @@ class Attention(nn.Module):
return self.paged_kv_cache_role is not None and self._active_paged_kv_adapter() is not None
def _run_local_attention(self, query, key, value, attn_metadata):
if self._has_custom_attention:
return self.attention(query, key, value, attn_metadata)
self._assert_metadata_compatible(attn_metadata)
in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile
@@ -494,6 +517,8 @@ class Attention(nn.Module):
def _assert_metadata_compatible(self, attn_metadata: AttentionMetadata | None) -> None:
if attn_metadata is None:
return
if self.attn_backend is None:
return
backend_name = self.attn_backend.get_name()
if attn_metadata.attn_mask is not None and not self.attn_backend.supports_attention_mask(
getattr(self, "attn_spec", None)
vllm_omni/diffusion/cache/cachedit/__init__.py
Vendored
+3 -1
View File
@@ -1,11 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Public API for the Cache-DiT diffusion cache backend."""
from vllm_omni.diffusion.cache.cachedit.backend import (
CUSTOM_DIT_ENABLERS,
CacheDiTBackend,
CacheDiTEnableResult,
cache_summary,
enable_cache_for_dit,
)
@@ -32,6 +33,7 @@ __all__ = [
"CUSTOM_DIT_ENABLERS",
"CacheDiTAdapterConfig",
"CacheDiTBackend",
"CacheDiTEnableResult",
"CacheDiTConfig",
"CacheDiTRequestSpec",
"RequestScopedCacheDiTRuntime",
vllm_omni/diffusion/cache/cachedit/backend.py
Vendored
+35 -9
View File
@@ -1,11 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Generic Cache-DiT backend lifecycle and integration."""
from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass
from operator import attrgetter
from typing import TYPE_CHECKING, Any, TypeAlias, cast
@@ -28,7 +29,20 @@ if TYPE_CHECKING:
logger = init_logger(__name__)
RefreshCacheContextFunc: TypeAlias = Callable[[Any, int, bool], None]
CacheDiTEnabler: TypeAlias = Callable[[Any, DiffusionCacheConfig], RefreshCacheContextFunc]
@dataclass(frozen=True)
class CacheDiTEnableResult:
"""Custom-enabler result with the exact targets needed for teardown."""
refresh: RefreshCacheContextFunc
targets: tuple[Any, ...]
CacheDiTEnabler: TypeAlias = Callable[
[Any, DiffusionCacheConfig],
RefreshCacheContextFunc | CacheDiTEnableResult,
]
# Model-specific implementations register themselves when the package loads.
CUSTOM_DIT_ENABLERS: dict[str, CacheDiTEnabler] = {}
@@ -56,13 +70,15 @@ def _dit_module_names(pipeline: SupportsComponentDiscovery) -> tuple[str, ...]:
def cache_summary(pipeline: SupportsComponentDiscovery, details: bool = True) -> None:
"""Log Cache-DiT statistics for every transformer on the pipeline."""
transformers = [attrgetter(name)(pipeline) for name in _dit_module_names(pipeline)]
for transformer in transformers:
if not BlockAdapter.is_cached(transformer):
targets = getattr(pipeline, "_cache_dit_targets", None)
if targets is None:
targets = tuple(attrgetter(name)(pipeline) for name in _dit_module_names(pipeline))
for target in targets:
if not BlockAdapter.is_cached(target):
continue
cache_dit.summary(transformer, details=details)
cache_dit.summary(target, details=details)
if not transformers:
if not targets:
logger.warning("CacheDiT summary failed; this pipeline has no defined transformer attribute")
@@ -261,8 +277,13 @@ class CacheDiTBackend(CacheBackend):
self._cache_targets = []
if custom_enabler is not None:
logger.info("Using custom cache-dit enabler for model: %s", pipeline_name)
self._refresh_funcs = [custom_enabler(pipeline, self.config)]
self._cache_targets = [_default_get_pipeline_transformer(pipeline)]
result = custom_enabler(pipeline, self.config)
if isinstance(result, CacheDiTEnableResult):
self._refresh_funcs = [result.refresh]
self._cache_targets = list(result.targets)
else:
self._refresh_funcs = [result]
self._cache_targets = [_default_get_pipeline_transformer(pipeline)]
else:
for name in _dit_module_names(pipeline):
get_transformer = _make_pipeline_transformer_getter(name)
@@ -282,6 +303,8 @@ class CacheDiTBackend(CacheBackend):
if not self._refresh_funcs:
raise ValueError(f"Pipeline {pipeline_name} has no declared DiT modules for Cache-DiT")
pipeline._cache_dit_targets = tuple(self._cache_targets)
self.enabled = True
logger.info("Cache-dit enabled successfully on %s", pipeline_name)
@@ -302,6 +325,8 @@ class CacheDiTBackend(CacheBackend):
finally:
self._refresh_funcs = []
self._cache_targets = []
if hasattr(pipeline, "_cache_dit_targets"):
del pipeline._cache_dit_targets
self.enabled = False
def refresh(self, pipeline: SupportsComponentDiscovery, num_inference_steps: int, verbose: bool = True) -> None:
@@ -320,6 +345,7 @@ class CacheDiTBackend(CacheBackend):
__all__ = [
"CUSTOM_DIT_ENABLERS",
"CacheDiTEnableResult",
"CacheDiTBackend",
"cache_summary",
"enable_cache_for_dit",
vllm_omni/diffusion/cache/cachedit/model_specific.py
Vendored
+33 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Model-specific Cache-DiT adapters and enablers."""
@@ -25,6 +25,7 @@ from vllm.logger import init_logger
from vllm_omni.diffusion.cache.cachedit.backend import (
CUSTOM_DIT_ENABLERS,
CacheDiTEnableResult,
RefreshCacheContextFunc,
_build_cache_context_refresh,
_default_get_pipeline_transformer,
@@ -836,6 +837,36 @@ def enable_cache_for_krea2(pipeline: Any, cache_config: Any) -> RefreshCacheCont
return enable_cache_for_dit(pipeline, cache_config, block_adapter)
def _get_magi2_transformer_block(pipeline: Any) -> torch.nn.Module:
return pipeline.transformer.block
def enable_cache_for_magi2(pipeline: Any, cache_config: Any) -> CacheDiTEnableResult:
"""Cache only MAGI-2's repeated native transformer-layer stack.
The pre/post adapters still execute on every denoising step. MAGI-2 either
packs both CFG branches into one transformer call or assigns one branch to
each CFG-parallel rank. In both layouts each rank invokes this stack once
per denoising step, so it remains a non-separate-CFG Pattern-3 stack.
"""
transformer_block = _get_magi2_transformer_block(pipeline)
block_adapter = BlockAdapter(
transformer=transformer_block,
blocks=[transformer_block.layers],
forward_pattern=[ForwardPattern.Pattern_3],
has_separate_cfg=False,
check_forward_pattern=True,
)
refresh = enable_cache_for_dit(
pipeline,
cache_config,
block_adapter,
get_pipeline_transformer=_get_magi2_transformer_block,
)
return CacheDiTEnableResult(refresh=refresh, targets=(block_adapter,))
def register_custom_dit_enablers() -> None:
"""Register model-specific Cache-DiT enablers.
@@ -853,6 +884,7 @@ def register_custom_dit_enablers() -> None:
"Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3,
"Cosmos3OmniPipeline": enable_cache_for_cosmos3,
"Krea2Pipeline": enable_cache_for_krea2,
"Magi2Pipeline": enable_cache_for_magi2,
}
)
vllm_omni/diffusion/data.py
+15 -2
View File
@@ -631,10 +631,14 @@ def resolve_model_class_name(
"""
from vllm.transformers_utils.config import get_hf_file_to_dict
from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index
from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index, resolve_native_diffusion_model_class
if not model:
return None
native_model_class = resolve_native_diffusion_model_class(model)
if native_model_class is not None:
return native_model_class
is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"}
# Diffusers models: read _class_name from the pipeline index. Missing
@@ -1332,7 +1336,10 @@ class OmniDiffusionConfig:
"""
from vllm.transformers_utils.config import get_hf_file_to_dict
from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index
from vllm_omni.diffusion.utils.hf_utils import (
get_diffusion_model_index,
resolve_native_diffusion_model_class,
)
assert self.model is not None
try:
@@ -1398,6 +1405,12 @@ class OmniDiffusionConfig:
else:
cfg = get_hf_file_to_dict("config.json", self.model, revision=self.revision)
if cfg is None:
native_model_class = resolve_native_diffusion_model_class(self.model)
if native_model_class is not None:
self.model_class_name = native_model_class
self.set_tf_model_config(TransformerConfig())
self.update_multimodal_support()
return
# Lance ships its top-level config.json one directory above
# the per-checkpoint subfolders (``Lance_3B/`` or
# ``Lance_3B_Video/``). Try to recover that case before
vllm_omni/diffusion/distributed/hsdp.py
+6 -4
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from collections.abc import Callable
from dataclasses import dataclass
@@ -139,8 +139,9 @@ def apply_hsdp_to_model(
# hsdp_config.param_dtype (typically bfloat16). FP8 GEMM kernels expect
# FP8 inputs; an implicit FP8 -> bf16 cast would silently break them.
has_fp8_params = any(p.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) for p in model.parameters())
preserve_parameter_dtypes = bool(getattr(model, "_hsdp_preserve_parameter_dtypes", False))
mp_policy = MixedPrecisionPolicy(
param_dtype=None if has_fp8_params else hsdp_config.param_dtype,
param_dtype=None if has_fp8_params or preserve_parameter_dtypes else hsdp_config.param_dtype,
reduce_dtype=hsdp_config.reduce_dtype,
output_dtype=hsdp_config.output_dtype,
cast_forward_inputs=False,
@@ -182,8 +183,9 @@ def apply_hsdp_to_model(
"the ignored modules can be placed on the worker's execution device."
)
for mod_name in ignored_module_names:
sub_mod = getattr(model, mod_name, None)
if sub_mod is None:
try:
sub_mod = model.get_submodule(mod_name)
except AttributeError:
logger.warning("_hsdp_ignored_modules entry %r not found on model", mod_name)
continue
sub_mod.to(target_device)
vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py
+23 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Loader-side adapters for proven direct-checkpoint mmap layouts."""
from __future__ import annotations
@@ -83,8 +83,30 @@ class _Cosmos3DirectMmapAdapter:
return DirectMmapTensorPolicy(allow_custom_loader=True)
class _Magi2DirectMmapAdapter:
"""Apply MAGI-2's rank-local MoE-head slices during bounded staging."""
def policy_for(
self,
runtime_name: str,
target: torch.Tensor,
) -> DirectMmapTensorPolicy:
del runtime_name
transform = getattr(target, "mmap_weight_transform", None)
return DirectMmapTensorPolicy(
transform=transform if callable(transform) else None,
)
def get_direct_mmap_adapter(pipeline: nn.Module) -> DirectMmapAdapter | None:
"""Return a loader-owned adapter without requiring a pipeline flag."""
from vllm_omni.diffusion.models.magi2.modeling_magi2 import (
Magi2PreviewTransformer,
)
if any(isinstance(module, Magi2PreviewTransformer) for module in pipeline.modules()):
return _Magi2DirectMmapAdapter()
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
MiniMaxH3DiTModel,
)
vllm_omni/diffusion/model_loader/host_weight_plan.py
+38 -5
View File
@@ -270,15 +270,48 @@ def _validate_source_metadata(
source = handle.get_slice(binding.checkpoint_key)
source_shape = tuple(source.get_shape())
source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype())
if source_shape != tuple(target.shape):
raise _PlanIncompatibleError(
f"shape mismatch for {runtime_name!r}: checkpoint={source_shape}, runtime={tuple(target.shape)}"
)
if source_dtype is None or source_dtype != target.dtype:
if source_dtype is None:
raise _PlanIncompatibleError(
f"dtype mismatch for {runtime_name!r}: checkpoint={source.get_dtype()}, runtime={target.dtype}"
)
runtime_shape = source_shape
runtime_dtype = source_dtype
if binding.transform is not None:
# Adapter transforms can map full checkpoint tensors to a
# rank-local runtime shape. Validate that contract without
# materializing the source tensor during preflight.
try:
transformed = binding.transform(
torch.empty(
source_shape,
dtype=source_dtype,
device="meta",
)
)
except (NotImplementedError, RuntimeError, TypeError, ValueError) as exc:
raise _PlanIncompatibleError(
f"checkpoint transform for {runtime_name!r} cannot be validated on tensor metadata: {exc}"
) from exc
if not isinstance(transformed, torch.Tensor):
raise _PlanIncompatibleError(
f"checkpoint transform for {runtime_name!r} returned {type(transformed).__name__}, "
"expected torch.Tensor"
)
runtime_shape = tuple(transformed.shape)
runtime_dtype = transformed.dtype
if runtime_shape != tuple(target.shape):
raise _PlanIncompatibleError(
f"shape mismatch for {runtime_name!r}: checkpoint={source_shape}, "
f"transformed={runtime_shape}, runtime={tuple(target.shape)}"
)
if runtime_dtype != target.dtype:
raise _PlanIncompatibleError(
f"dtype mismatch for {runtime_name!r}: checkpoint={source.get_dtype()}, "
f"transformed={runtime_dtype}, runtime={target.dtype}"
)
def build_checkpoint_mmap_plan(
pipeline: nn.Module,
vllm_omni/diffusion/model_metadata.py
+5
View File
@@ -55,6 +55,11 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
supports_mixed_reference_inputs=True,
final_output_type="video",
),
"Magi2Pipeline": DiffusionModelMetadata(
supports_multimodal_inputs=True,
max_multimodal_image_inputs=1,
final_output_type="video",
),
"WanPipeline": DiffusionModelMetadata(
attention_mask_free=True,
final_output_type="video",
vllm_omni/diffusion/models/magi2/__init__.py
+21
View File
@@ -0,0 +1,21 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""SandAI MAGI-2 Preview diffusion support."""
from __future__ import annotations
from typing import Any
__all__ = ["Magi2Pipeline", "get_magi2_post_process_func"]
def __getattr__(name: str) -> Any:
if name in __all__:
from .pipeline_magi2 import Magi2Pipeline, get_magi2_post_process_func
return {
"Magi2Pipeline": Magi2Pipeline,
"get_magi2_post_process_func": get_magi2_post_process_func,
}[name]
raise AttributeError(name)
vllm_omni/diffusion/models/magi2/attention.py
+293
View File
@@ -0,0 +1,293 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Copyright (c) 2026 SandAI. All Rights Reserved.
"""Native MAGI-2 packed attention with sinks and Ulysses exchange.
The sink and context-parallel math is adapted from SandAI's Apache-2.0
MAGI-2 preview implementation. This version uses vLLM's bundled
FlashAttention extension and vLLM-Omni's existing Ulysses process group; the
PyTorch path is an exact, portable oracle for small tests.
"""
from __future__ import annotations
import logging
import os
from dataclasses import dataclass
from functools import cache
import torch
import torch.nn as nn
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.backends.utils.fa import (
resolve_vllm_flash_attn_version,
vllm_flash_attn_varlen_with_lse,
)
from .parallel import (
Magi2ParallelGroup,
get_magi2_ulysses_group,
scatter_heads_gather_seqlen,
scatter_seqlen_gather_heads,
)
logger = logging.getLogger(__name__)
@cache
def _resolve_flash_attn_version() -> int:
"""Apply MAGI-2's operator override to the shared FA2/FA3/FA4 resolver."""
requested = os.environ.get("MAGI2_FLASH_ATTN_VERSION")
version = resolve_vllm_flash_attn_version(requested)
logger.info("MAGI-2 selected FlashAttention %d", version)
return version
@dataclass(frozen=True)
class VarlenHandler:
"""Packed-sequence metadata consumed by MAGI-2 attention."""
cu_seqlens_q: torch.Tensor | None
cu_seqlens_k: torch.Tensor | None
max_seqlen_q: int | None = None
max_seqlen_k: int | None = None
def resolved(self, q_tokens: int, k_tokens: int) -> tuple[torch.Tensor, torch.Tensor, int, int]:
if self.cu_seqlens_q is None:
cu_q = torch.tensor([0, q_tokens], device="cpu", dtype=torch.int32)
else:
cu_q = self.cu_seqlens_q
if self.cu_seqlens_k is None:
cu_k = torch.tensor([0, k_tokens], device="cpu", dtype=torch.int32)
else:
cu_k = self.cu_seqlens_k
max_q = self.max_seqlen_q
max_k = self.max_seqlen_k
if max_q is None:
max_q = int(torch.diff(cu_q).max().item()) if cu_q.numel() > 1 else 0
if max_k is None:
max_k = int(torch.diff(cu_k).max().item()) if cu_k.numel() > 1 else 0
return cu_q, cu_k, max_q, max_k
def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor:
if not interleaved:
x1, x2 = x.chunk(2, dim=-1)
return torch.cat((-x2, x1), dim=-1)
x1, x2 = x[..., ::2], x[..., 1::2]
return torch.stack((-x2, x1), dim=-1).flatten(-2)
def apply_rotary_emb(
x: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
*,
interleaved: bool = False,
) -> torch.Tensor:
"""Apply the released element-wise RoPE layout to ``[...,H,D]``."""
rotary_dim = cos.shape[-1] * 2
if rotary_dim > x.shape[-1]:
raise ValueError(f"RoPE dimension {rotary_dim} exceeds head dimension {x.shape[-1]}")
if interleaved:
cos = cos.unsqueeze(-2).repeat_interleave(2, dim=-1)
sin = sin.unsqueeze(-2).repeat_interleave(2, dim=-1)
else:
cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2)
sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2)
rotated = x[..., :rotary_dim] * cos + rotate_half(x[..., :rotary_dim], interleaved) * sin
return torch.cat((rotated, x[..., rotary_dim:]), dim=-1)
def correct_out_lse_with_sink(
out: torch.Tensor,
lse: torch.Tensor,
sink: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Add zero-valued attention sinks to an already-computed softmax.
FlashAttention returns ``out[T,H,D]`` and conventionally ``lse[H,T]``.
A MAGI sink contains additional logits ``[num_sink,H]`` whose values are
zero vectors, so only the denominator and LSE change.
"""
if sink is None or sink.numel() == 0:
return out, lse
if out.ndim != 3 or lse.ndim != 2 or sink.ndim != 2:
raise ValueError(
"expected out[T,H,D], lse[H,T], sink[num_sink,H], got "
f"{tuple(out.shape)}, {tuple(lse.shape)}, {tuple(sink.shape)}"
)
old_lse = lse.float().transpose(0, 1)
sink_lse = torch.logsumexp(sink.float(), dim=0).unsqueeze(0)
if old_lse.shape[-1] != sink_lse.shape[-1]:
raise ValueError("attention sink and FlashAttention head counts differ")
new_lse = torch.logaddexp(old_lse, sink_lse)
delta = old_lse - new_lse
delta = torch.where(torch.isfinite(delta), delta, torch.full_like(delta, -torch.inf))
corrected = out * torch.exp(delta).unsqueeze(-1).to(out.dtype)
return corrected, new_lse.transpose(0, 1).contiguous()
def _repeat_kv_heads(tensor: torch.Tensor, query_heads: int) -> torch.Tensor:
kv_heads = tensor.shape[1]
if kv_heads == query_heads:
return tensor
if query_heads % kv_heads:
raise ValueError(f"query heads {query_heads} must be divisible by KV heads {kv_heads}")
return tensor.repeat_interleave(query_heads // kv_heads, dim=1)
def torch_varlen_attention_with_sink(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
*,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
softcap: float = -1.0,
sink: torch.Tensor | None = None,
) -> torch.Tensor:
"""Reference packed attention, including GQA and sink logits."""
if q.ndim != 3 or k.ndim != 3 or v.ndim != 3:
raise ValueError("packed attention expects q/k/v shaped [tokens,heads,dim]")
if cu_seqlens_q.numel() != cu_seqlens_k.numel():
raise ValueError("query and key cumulative-length arrays must contain the same batch count")
output = torch.empty_like(q)
scale = q.shape[-1] ** -0.5
for batch_idx in range(cu_seqlens_q.numel() - 1):
q_start, q_end = (int(v) for v in cu_seqlens_q[batch_idx : batch_idx + 2].tolist())
k_start, k_end = (int(v) for v in cu_seqlens_k[batch_idx : batch_idx + 2].tolist())
q_part = q[q_start:q_end].float()
k_part = _repeat_kv_heads(k[k_start:k_end], q.shape[1]).float()
v_part = _repeat_kv_heads(v[k_start:k_end], q.shape[1]).float()
scores = torch.einsum("qhd,khd->hqk", q_part, k_part) * scale
if softcap > 0:
scores = softcap * torch.tanh(scores / softcap)
if sink is not None and sink.numel() > 0:
sink_scores = sink.float().transpose(0, 1).unsqueeze(1).expand(-1, q_part.shape[0], -1)
probabilities = torch.softmax(torch.cat((scores, sink_scores), dim=-1), dim=-1)[..., : k_part.shape[0]]
else:
probabilities = torch.softmax(scores, dim=-1)
output[q_start:q_end] = torch.einsum("hqk,khd->qhd", probabilities, v_part).to(output.dtype)
return output
def packed_attention_with_sink(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
varlen: VarlenHandler,
*,
softcap: float = -1.0,
sink: torch.Tensor | None = None,
) -> torch.Tensor:
"""Run packed attention on one rank after Ulysses head exchange."""
cu_q, cu_k, max_q, max_k = varlen.resolved(q.shape[0], k.shape[0])
cu_q = cu_q.to(device=q.device, dtype=torch.int32).contiguous()
cu_k = cu_k.to(device=q.device, dtype=torch.int32).contiguous()
if q.is_cuda:
out, lse = vllm_flash_attn_varlen_with_lse(
q,
k,
v,
cu_seqlens_q=cu_q,
cu_seqlens_k=cu_k,
max_seqlen_q=max_q,
max_seqlen_k=max_k,
softcap=softcap,
deterministic=os.environ.get("MAGI2_DETERMINISTIC", "0") == "1",
fa_version=_resolve_flash_attn_version(),
)
return correct_out_lse_with_sink(out, lse, sink)[0]
return torch_varlen_attention_with_sink(
q,
k,
v,
cu_seqlens_q=cu_q,
cu_seqlens_k=cu_k,
softcap=softcap,
sink=sink,
)
def ulysses_packed_attention_with_sink(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
varlen: VarlenHandler,
split_sizes: list[int] | torch.Tensor,
*,
softcap: float = -1.0,
sink: torch.Tensor | None = None,
group: Magi2ParallelGroup | None = None,
) -> torch.Tensor:
"""MAGI-2 attention with overlapping Ulysses CP/head exchange."""
group = group or get_magi2_ulysses_group()
if isinstance(split_sizes, torch.Tensor):
split_sizes = [int(v) for v in split_sizes.detach().cpu().tolist()]
if group.world_size > 1:
q, k, v = scatter_heads_gather_seqlen((q, k, v), split_sizes, group)
if sink is not None:
if sink.shape[-1] % group.world_size:
raise ValueError("attention sink heads must divide across Ulysses ranks")
sink = torch.chunk(sink, group.world_size, dim=-1)[group.rank].contiguous()
output = packed_attention_with_sink(q, k, v, varlen, softcap=softcap, sink=sink)
if group.world_size > 1:
output = scatter_seqlen_gather_heads(output.contiguous(), split_sizes, group)
assert isinstance(output, torch.Tensor)
return output
class Magi2PackedAttentionKernel(nn.Module):
"""Model kernel plugged into the shared diffusion Attention contract."""
def __init__(self, softcap: float) -> None:
super().__init__()
self.softcap = softcap
def forward(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: AttentionMetadata | None = None,
) -> torch.Tensor:
if attn_metadata is None:
raise ValueError("MAGI-2 packed attention requires attention metadata")
varlen = attn_metadata.extra.get("magi2_varlen")
split_sizes = attn_metadata.extra.get("magi2_split_sizes")
sink = attn_metadata.extra.get("magi2_sink")
if not isinstance(varlen, VarlenHandler):
raise TypeError("magi2_varlen must be a VarlenHandler")
if not isinstance(split_sizes, (list, torch.Tensor)):
raise TypeError("magi2_split_sizes must be a list or tensor")
if sink is not None and not isinstance(sink, torch.Tensor):
raise TypeError("magi2_sink must be a tensor or None")
return ulysses_packed_attention_with_sink(
query,
key,
value,
varlen,
split_sizes,
softcap=self.softcap,
sink=sink,
)
__all__ = [
"VarlenHandler",
"Magi2PackedAttentionKernel",
"apply_rotary_emb",
"correct_out_lse_with_sink",
"packed_attention_with_sink",
"torch_varlen_attention_with_sink",
"ulysses_packed_attention_with_sink",
]
vllm_omni/diffusion/models/magi2/audio_decoder.py
+195
View File
@@ -0,0 +1,195 @@
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2026 SandAI. All Rights Reserved.
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Native Stable Audio Open decoder used by MAGI-2.
The MAGI checkpoint embeds Stable Audio's Oobleck VAE under the
``pretransform.model`` prefix and uses the sequential module names from
``stable-audio-tools``. This module converts those names to Diffusers'
``AutoencoderOobleck`` layout and loads only the decoder tensors.
"""
from __future__ import annotations
import json
from collections.abc import Mapping
from pathlib import Path
from typing import Any
import torch
import torch.nn as nn
from diffusers import AutoencoderOobleck
from safetensors import safe_open
_CHECKPOINT_PREFIX = "pretransform.model.decoder."
_BLOCK_COMPONENTS = {
0: "snake1",
1: "conv_t1",
}
_RESIDUAL_COMPONENTS = {
0: "snake1",
1: "conv1",
2: "snake2",
3: "conv2",
}
def convert_stable_audio_decoder_key(key: str) -> str | None:
"""Map one Stable Audio decoder key to Diffusers' Oobleck decoder.
Unrelated checkpoint tensors return ``None``. A key inside the decoder
namespace that does not match the released Stable Audio layout raises so a
changed checkpoint cannot be loaded partially by accident.
"""
if not key.startswith(_CHECKPOINT_PREFIX):
return None
suffix = key[len(_CHECKPOINT_PREFIX) :]
parts = suffix.split(".")
if len(parts) < 3 or parts[0] != "layers":
raise ValueError(f"Unsupported Stable Audio decoder key: {key}")
try:
layer = int(parts[1])
except ValueError as exc:
raise ValueError(f"Unsupported Stable Audio decoder key: {key}") from exc
tail = parts[2:]
if layer == 0 and len(tail) == 1:
return f"conv1.{tail[0]}"
if 1 <= layer <= 5:
block = layer - 1
if len(tail) == 3 and tail[0] == "layers":
try:
component = _BLOCK_COMPONENTS[int(tail[1])]
except (KeyError, ValueError) as exc:
raise ValueError(f"Unsupported Stable Audio decoder key: {key}") from exc
return f"block.{block}.{component}.{tail[2]}"
if len(tail) == 5 and tail[0] == "layers" and tail[2] == "layers":
try:
residual = int(tail[1]) - 1
component = _RESIDUAL_COMPONENTS[int(tail[3])]
except (KeyError, ValueError) as exc:
raise ValueError(f"Unsupported Stable Audio decoder key: {key}") from exc
if residual not in (1, 2, 3):
raise ValueError(f"Unsupported Stable Audio decoder key: {key}")
return f"block.{block}.res_unit{residual}.{component}.{tail[4]}"
if layer == 6 and len(tail) == 1:
return f"snake1.{tail[0]}"
if layer == 7 and len(tail) == 1:
return f"conv2.{tail[0]}"
raise ValueError(f"Unsupported Stable Audio decoder key: {key}")
def convert_stable_audio_decoder_state_dict(
state_dict: Mapping[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
"""Convert and filter a Stable Audio checkpoint state dictionary."""
converted: dict[str, torch.Tensor] = {}
for key, value in state_dict.items():
target_key = convert_stable_audio_decoder_key(key)
if target_key is None:
continue
if target_key in converted:
raise ValueError(f"Duplicate converted Stable Audio key: {target_key}")
if target_key.endswith((".alpha", ".beta")):
value = value.reshape(1, -1, 1)
converted[target_key] = value
return converted
def _load_oobleck_config(model_path: Path) -> tuple[dict[str, Any], int]:
config_path = model_path / "model_config.json"
if not config_path.is_file():
raise FileNotFoundError(f"Stable Audio config does not exist: {config_path}")
with config_path.open(encoding="utf-8") as file:
full_config = json.load(file)
try:
pretransform = full_config["model"]["pretransform"]
config = pretransform["config"]
decoder = config["decoder"]
decoder_config = decoder["config"]
sample_rate = int(full_config["sample_rate"])
except (KeyError, TypeError, ValueError) as exc:
raise ValueError(f"Invalid Stable Audio model config: {config_path}") from exc
if pretransform.get("type") != "autoencoder" or decoder.get("type") != "oobleck":
raise ValueError("MAGI-2 requires an Oobleck Stable Audio pretransform")
if decoder_config.get("use_nearest_upsample", False):
raise ValueError("Diffusers Oobleck does not support nearest-neighbor decoder upsampling")
if decoder_config.get("final_tanh", True):
raise ValueError("MAGI-2 requires the Stable Audio decoder with final_tanh=false")
strides = list(decoder_config["strides"])
downsampling_ratio = int(config["downsampling_ratio"])
if int(torch.tensor(strides).prod().item()) != downsampling_ratio:
raise ValueError("Stable Audio strides do not match downsampling_ratio")
kwargs = {
"encoder_hidden_size": int(config["encoder"]["config"]["channels"]),
"downsampling_ratios": strides,
"channel_multiples": list(decoder_config["c_mults"]),
"decoder_channels": int(decoder_config["channels"]),
"decoder_input_channels": int(decoder_config["latent_dim"]),
"audio_channels": int(decoder_config["out_channels"]),
"sampling_rate": sample_rate,
}
return kwargs, downsampling_ratio
class Magi2AudioDecoder(nn.Module):
"""Decode MAGI-2 audio latents with the bundled Stable Audio VAE weights."""
def __init__(
self,
model_path: str | Path,
*,
device: str | torch.device = "cpu",
dtype: torch.dtype = torch.float32,
) -> None:
super().__init__()
model_path = Path(model_path)
if not model_path.is_dir():
raise ValueError(f"Stable Audio model path must be a local directory: {model_path}")
oobleck_kwargs, self.downsampling_ratio = _load_oobleck_config(model_path)
self.sample_rate = int(oobleck_kwargs["sampling_rate"])
# Construct the Diffusers implementation for architectural compatibility,
# then retain only the decoder because MAGI-2 has no audio input encoder.
autoencoder = AutoencoderOobleck(**oobleck_kwargs)
self.decoder = autoencoder.decoder.to(device=device, dtype=dtype)
del autoencoder
weights_path = model_path / "model.safetensors"
if not weights_path.is_file():
raise FileNotFoundError(f"Stable Audio weights do not exist: {weights_path}")
selected: dict[str, torch.Tensor] = {}
with safe_open(weights_path, framework="pt", device=str(device)) as checkpoint:
for key in checkpoint.keys():
if key.startswith(_CHECKPOINT_PREFIX):
selected[key] = checkpoint.get_tensor(key)
converted = convert_stable_audio_decoder_state_dict(selected)
converted = {key: value.to(dtype=dtype) for key, value in converted.items()}
self.decoder.load_state_dict(converted, strict=True)
self.eval().requires_grad_(False)
@property
def latent_fps(self) -> float:
return self.sample_rate / self.downsampling_ratio
@torch.inference_mode()
def decode(self, latents: torch.Tensor) -> torch.Tensor:
parameter = next(self.decoder.parameters())
return self.decoder(latents.to(device=parameter.device, dtype=parameter.dtype))
forward = decode
__all__ = [
"Magi2AudioDecoder",
"convert_stable_audio_decoder_key",
"convert_stable_audio_decoder_state_dict",
]
vllm_omni/diffusion/models/magi2/conditioning.py
+321
View File
@@ -0,0 +1,321 @@
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2026 SandAI. All Rights Reserved.
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Native MAGI-2 prompt conditioning.
Adapted from SandAI's MAGI-2 Preview inference implementation. The model
loading remains in-tree and uses the Transformers Qwen3.5 implementation; no
MAGI-2 reference checkout is imported at runtime.
"""
from __future__ import annotations
import json
from collections.abc import Iterable
from typing import Any
import torch
import torch.nn as nn
from tokenizers import Regex, pre_tokenizers
from transformers import AutoTokenizer
def _strip_empty(obj: Any) -> Any:
if isinstance(obj, dict):
return {
key: value
for key, value in ((key, _strip_empty(value)) for key, value in obj.items())
if value is not None and value != [] and value != {}
}
if isinstance(obj, list):
return [_strip_empty(value) for value in obj if value is not None]
return obj
def _one_line(value: Any) -> str:
if value is None:
return ""
if not isinstance(value, str):
value = str(value)
return " ".join(value.split())
def _as_dict(value: Any) -> dict[str, Any]:
return value if isinstance(value, dict) else {}
def _as_list(value: Any) -> list[Any]:
if value is None:
return []
return value if isinstance(value, list) else [value]
def _pick(values: dict[str, Any], keys: Iterable[str]) -> list[tuple[str, Any]]:
return [(key, values[key]) for key in keys if key in values and values[key] not in (None, [], {})]
def json_to_compact_markdown(raw_json: str | dict[str, Any]) -> str:
"""Convert MAGI's structured prompt JSON into its training-time text form."""
if isinstance(raw_json, str):
obj = json.loads(raw_json.strip())
else:
obj = raw_json
obj = _strip_empty(obj)
if not isinstance(obj, dict) or "global_layer" not in obj:
return raw_json if isinstance(raw_json, str) else json.dumps(raw_json, ensure_ascii=False)
global_layer = _as_dict(obj.get("global_layer"))
dynamic_layer = _as_dict(obj.get("dynamic_layer"))
reference_layer = obj.get("reference_layer", [])
lines: list[str] = []
context = _one_line(global_layer.get("context"))
description = _one_line(global_layer.get("description"))
if context or description:
lines.append(f"context: {context}" if context else "context")
if description:
lines.append(description)
aesthetics = _as_dict(global_layer.get("aesthetics"))
aesthetic_parts = []
for label, key in (("style", "style"), ("mood", "mood_atmosphere"), ("color", "color_scheme")):
value = _one_line(aesthetics.get(key))
if value:
aesthetic_parts.append(f"{label}={value}")
if aesthetic_parts:
lines.append("aesthetics: " + "; ".join(aesthetic_parts))
audio = _as_dict(global_layer.get("audio_baseline"))
dialogue = _as_dict(audio.get("dialogue"))
audio_parts = []
language = _one_line(dialogue.get("language"))
speakers = _as_list(dialogue.get("speaker_tags"))
ambience = _one_line(audio.get("ambience"))
if language:
audio_parts.append(f"language={language}")
if speakers:
audio_parts.append("speakers=" + ",".join(map(_one_line, speakers)))
if ambience:
audio_parts.append(f"ambience={ambience}")
if audio_parts:
lines.append("audio: " + "; ".join(audio_parts))
subjects = global_layer.get("alive_subjects_static")
if isinstance(subjects, list) and subjects:
lines.append("subjects:")
for subject in subjects:
if not isinstance(subject, dict):
continue
sid = _one_line(subject.get("subject_id"))
desc = _one_line(subject.get("description"))
position = _one_line(subject.get("position"))
orientation = _one_line(subject.get("orientation"))
attrs = []
for key, value in _pick(
_as_dict(subject.get("visual_attributes")),
["gender", "age_group", "ethnicity", "clothing", "appearance_details"],
):
value = _one_line(value)
if value:
attrs.append(f"{key}={value}")
row = " - " + " | ".join(part for part in (sid, desc) if part)
extra = "; ".join(part for part in (position, orientation) if part)
if extra:
row += f" ({extra})"
if attrs:
row += " :: " + "; ".join(attrs)
lines.append(row)
objects = global_layer.get("objects_static")
if isinstance(objects, list) and objects:
lines.append("objects:")
for obj_item in objects:
if not isinstance(obj_item, dict):
continue
oid = _one_line(obj_item.get("object_id"))
desc = _one_line(obj_item.get("description"))
shape = _one_line(obj_item.get("shape_and_color"))
position = _one_line(obj_item.get("position"))
row = " - " + " | ".join(part for part in (oid, desc) if part)
details = "; ".join(part for part in (shape, position) if part)
if details:
row += " :: " + details
lines.append(row)
segments = dynamic_layer.get("timeline_segments")
if isinstance(segments, list) and segments:
lines.append("timeline:")
for segment in segments:
if not isinstance(segment, dict):
continue
basic = _as_dict(segment.get("segment_basic_info"))
timestamp = _one_line(basic.get("timestamp_range"))
desc = _one_line(basic.get("segment_description"))
head = f" - {timestamp}" if timestamp else " -"
if desc:
head += f" {desc}"
lines.append(head.rstrip())
segment_audio = _as_dict(segment.get("audio"))
for dialogue_line in segment_audio.get("dialogue_lines", []) or []:
if not isinstance(dialogue_line, dict):
continue
speaker = _one_line(dialogue_line.get("speaker"))
text = _one_line(dialogue_line.get("text"))
timestamp = _one_line(dialogue_line.get("timestamp"))
if text:
line = (f" - dialogue {speaker}: " if speaker else " - dialogue: ") + text
if timestamp:
line += f" ({timestamp})"
lines.append(line)
for alive in segment.get("alive_subjects", []) or []:
if not isinstance(alive, dict):
continue
sid = _one_line(alive.get("subject_id"))
action = _as_dict(alive.get("action"))
parts = [
_one_line(action.get("primary_action")),
_one_line(action.get("interaction")),
_one_line(action.get("facial_expression")),
]
parts = [part for part in parts if part]
if parts:
lines.append(f" - action {sid}: " + "; ".join(parts))
for obj_item in segment.get("objects", []) or []:
if not isinstance(obj_item, dict):
continue
oid = _one_line(obj_item.get("object_id"))
state = _as_dict(obj_item.get("dynamic_state"))
parts = [_one_line(state.get("state_change")), _one_line(state.get("motion_detail"))]
parts = [part for part in parts if part]
if parts:
lines.append(f" - objects {oid}: " + "; ".join(parts))
if isinstance(reference_layer, list):
lines.extend(str(line) for line in reference_layer if str(line).strip())
return "\n".join(line for line in lines if line.strip())
def normalize_prompt(prompt: str) -> str:
"""Normalize structured prompts while leaving ordinary text byte-for-byte intact."""
try:
parsed = json.loads(prompt)
except (json.JSONDecodeError, TypeError):
return prompt
if isinstance(parsed, dict | list):
return json_to_compact_markdown(prompt)
return prompt
class Magi2Qwen35TextEncoder(nn.Module):
"""Qwen3.5-27B feature extractor used by MAGI-2 Preview."""
def __init__(
self,
model_path: str,
*,
dtype: torch.dtype = torch.bfloat16,
max_length: int = 7000,
skip_layer: int = 2,
local_files_only: bool = True,
) -> None:
super().__init__()
self.max_length = max_length
self.skip_layer = skip_layer
self.tokenizer = AutoTokenizer.from_pretrained(
model_path,
padding_side="right",
local_files_only=local_files_only,
)
cjk_split = pre_tokenizers.Split(
pattern=Regex(
r"([\u1100-\u11ff\u2e80-\ua4cf\ua840-\uD7AF"
r"\uF900-\uFAFF\uFE30-\uFE4F\uFF65-\uFFDC\U00020000-\U0002FFFF])"
),
behavior="isolated",
)
original = self.tokenizer.backend_tokenizer.pre_tokenizer
self.tokenizer.backend_tokenizer.pre_tokenizer = pre_tokenizers.Sequence([cjk_split, original])
from transformers import Qwen3_5TextModel
self.text_model = Qwen3_5TextModel.from_pretrained(
model_path,
torch_dtype=dtype,
local_files_only=local_files_only,
)
self.text_model.eval().requires_grad_(False)
def _tokenize(self, prompt: str, *, offsets: bool = False):
return self.tokenizer(
[normalize_prompt(prompt)],
return_tensors="pt",
padding="longest",
return_offsets_mapping=offsets,
max_length=self.max_length,
truncation=True,
)
def get_target_token_indices(self, prompt: str, target: str | None) -> list[int] | None:
if not target:
return None
normalized = normalize_prompt(prompt)
offsets = self._tokenize(prompt, offsets=True)["offset_mapping"][0]
start = normalized.find(target)
if start < 0:
return None
end = start + len(target)
indices = []
for index, (token_start, token_end) in enumerate(offsets.tolist()):
if token_start == 0 and token_end == 0 and index != 0:
continue
if max(token_start, start) < min(token_end, end):
indices.append(index)
return indices
def pool_figure_tokens(
self,
prompt: str,
targets: list[str],
text_features: torch.Tensor,
) -> torch.Tensor:
"""Mean-pool the token features spelling each ``<Figure N>`` marker."""
embeddings = []
for target in targets:
indices = self.get_target_token_indices(prompt, target)
if indices:
embeddings.append(text_features[0, indices].mean(dim=0).clone())
else:
embeddings.append(text_features.new_zeros(text_features.shape[-1]))
return torch.stack(embeddings)
@torch.inference_mode()
def encode(self, prompt: str) -> torch.Tensor:
inputs = self._tokenize(prompt)
device = next(self.text_model.parameters()).device
inputs = inputs.to(device)
outputs = self.text_model(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
output_hidden_states=True,
return_dict=True,
)
if self.skip_layer == 0:
return outputs.last_hidden_state
return outputs.hidden_states[-(self.skip_layer + 1)]
forward = encode
__all__ = [
"Magi2Qwen35TextEncoder",
"json_to_compact_markdown",
"normalize_prompt",
]
vllm_omni/diffusion/models/magi2/configuration_magi2.py
+120
View File
@@ -0,0 +1,120 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Copyright (c) 2026 SandAI. All Rights Reserved.
"""Native configuration for the released MAGI-2 Preview checkpoint.
The defaults mirror ``sand-ai/MAGI-2-preview``. Keeping the architecture in
vLLM-Omni makes model construction independent of SandAI's Python package and
also gives tests a small-config entry point.
"""
from __future__ import annotations
from dataclasses import dataclass, field
import torch
@dataclass(frozen=True)
class Magi2MHCConfig:
enabled: bool = True
num_streams: int = 4
alpha_init: float = 0.01
sinkhorn_iterations: int = 20
sinkhorn_epsilon: float = 1e-12
@dataclass(frozen=True)
class Magi2MoEConfig:
num_heads: int = 12
num_experts: int = 256
top_k: int = 6
expert_intermediate_size: int = 1280
shared_expert_intermediate_size: int = 1280
modality_shared_expert_intermediate_size: int = 1280
score_function: str = "sigmoid"
normalize_routing_weights: bool = True
routing_scale: float = 4.9
layers: tuple[int, ...] = tuple(range(2, 38))
def __post_init__(self) -> None:
if self.num_heads <= 0 or self.num_experts <= 0:
raise ValueError("MAGI-2 MoE head and expert counts must be positive")
if not 0 < self.top_k <= self.num_experts:
raise ValueError("MAGI-2 MoE top_k must be in [1, num_experts]")
if self.score_function != "sigmoid":
raise ValueError("The released MAGI-2 checkpoint uses sigmoid routing")
@dataclass(frozen=True)
class Magi2PreviewConfig:
num_layers: int = 40
hidden_size: int = 3072
head_dim: int = 128
num_query_groups: int = 24
video_in_channels: int = 48
audio_in_channels: int = 64
text_in_channels: int = 5120
intermediate_factor: float = 4.0
multimodal_layers: tuple[int, ...] = (0, 1, 38, 39)
params_dtype: torch.dtype = torch.bfloat16
attention_softcap: float = -1.0
attention_sink_tokens: int = 1
attention_gating: bool = True
mhc: Magi2MHCConfig = field(default_factory=Magi2MHCConfig)
moe: Magi2MoEConfig = field(default_factory=Magi2MoEConfig)
@property
def num_attention_heads(self) -> int:
if self.hidden_size % self.head_dim:
raise ValueError("hidden_size must be divisible by head_dim")
return self.hidden_size // self.head_dim
@property
def num_heads_q(self) -> int:
return self.num_attention_heads
@property
def num_heads_kv(self) -> int:
return self.num_query_groups
@property
def virtual_width(self) -> int:
return self.hidden_size * self.mhc.num_streams
def validate(self) -> None:
if not self.mhc.enabled:
raise ValueError("The released MAGI-2 Preview checkpoint requires mHC")
if not self.attention_gating:
raise ValueError("The released MAGI-2 Preview checkpoint requires attention gating")
if self.attention_sink_tokens < 1:
raise ValueError("The released MAGI-2 Preview checkpoint requires an attention sink")
if self.hidden_size % self.moe.num_heads:
raise ValueError("hidden_size must be divisible by the number of MoE heads")
layer_ids = set(range(self.num_layers))
if not set(self.multimodal_layers).issubset(layer_ids):
raise ValueError("multimodal layer index is outside the transformer")
if not set(self.moe.layers).issubset(layer_ids):
raise ValueError("MoE layer index is outside the transformer")
@dataclass(frozen=True)
class Magi2GenerationConfig:
duration_seconds: float = 10.0
fps: float = 12.5
output_frames: int = 125
preview_steps: int = 100
shift: float = 7.0
video_guidance_scale: float = 5.0
audio_guidance_scale: float = 7.0
audio_latent_fps: float = 25.0
video_vae_stride: tuple[int, int, int] = (8, 16, 16)
patch_size: tuple[int, int, int] = (1, 1, 1)
video_latent_channels: int = 48
audio_latent_channels: int = 64
audio_sample_rate: int = 44_100
MAGI2_PREVIEW_CONFIG = Magi2PreviewConfig()
MAGI2_GENERATION_CONFIG = Magi2GenerationConfig()
vllm_omni/diffusion/models/magi2/layers.py
+437
View File
@@ -0,0 +1,437 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Copyright (c) 2026 SandAI. All Rights Reserved.
"""Checkpoint-compatible native layers for MAGI-2 Preview.
The grouped-linear, normalization, Fourier, and mHC formulas are adapted from
SandAI's Apache-2.0 preview implementation. They have been modified to remove
MagiCompiler and external Triton runtime dependencies.
"""
from __future__ import annotations
import math
from collections.abc import Callable
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
from .parallel import Magi2ParallelGroup, get_magi2_tp_group
def swiglu7(
x: torch.Tensor,
alpha: float = 1.702,
limit: float = 7.0,
out_dtype: torch.dtype | None = None,
) -> torch.Tensor:
"""Released GPT-OSS-style clamped SwiGLU activation."""
out_dtype = x.dtype if out_dtype is None else out_dtype
x = x.to(torch.float32)
gate, linear = x[..., ::2], x[..., 1::2]
gate = gate.clamp(max=limit)
linear = linear.clamp(min=-limit, max=limit)
return (gate * torch.sigmoid(alpha * gate) * (linear + 1.0)).to(out_dtype)
class ModalityDispatcher:
"""Precompute stable modality grouping and inverse permutation metadata."""
def __init__(self, modality_mapping: torch.Tensor, num_modalities: int) -> None:
if modality_mapping.ndim != 1:
raise ValueError("modality mapping must be a one-dimensional tensor")
self.modality_mapping = modality_mapping
self.num_modalities = num_modalities
self.permute_mapping = torch.argsort(modality_mapping, stable=True)
self.inv_permute_mapping = torch.argsort(self.permute_mapping)
self.permuted_modality_mapping = modality_mapping.index_select(0, self.permute_mapping)
self.group_size = torch.bincount(self.permuted_modality_mapping.long(), minlength=num_modalities).to(
torch.int32
)
self.group_size_cpu = [int(value) for value in self.group_size.cpu().tolist()]
self.cu_group_sizes = F.pad(torch.cumsum(self.group_size, dim=0), (1, 0))
def dispatch(self, tensor: torch.Tensor) -> list[torch.Tensor]:
return list(torch.split(tensor, self.group_size_cpu, dim=0))
@staticmethod
def undispatch(*groups: torch.Tensor) -> torch.Tensor:
return torch.cat(groups, dim=0)
def permute(self, tensor: torch.Tensor) -> torch.Tensor:
return tensor.index_select(0, self.permute_mapping)
def inverse_permute(self, tensor: torch.Tensor) -> torch.Tensor:
return tensor.index_select(0, self.inv_permute_mapping)
class Magi2GroupedLinear(nn.Module):
"""Modality-grouped linear with the released flattened weight layout."""
def __init__(
self,
in_features: int,
out_features: int,
*,
num_experts: int = 1,
bias: bool = False,
dtype: torch.dtype | None = None,
device: torch.device | str | None = None,
parallel_mode: str | None = None,
qkv_splits: tuple[int, int, int] | None = None,
tp_group: Magi2ParallelGroup | None = None,
) -> None:
super().__init__()
if parallel_mode not in (None, "column", "row"):
raise ValueError(f"unknown MAGI-2 linear parallel mode {parallel_mode!r}")
if qkv_splits is not None and parallel_mode != "column":
raise ValueError("segmented QKV slicing requires column parallelism")
if qkv_splits is not None and sum(qkv_splits) != out_features:
raise ValueError("QKV splits must sum to out_features")
self.in_features = in_features
self.out_features = out_features
self.num_experts = num_experts
self.parallel_mode = parallel_mode
self.qkv_splits = qkv_splits
self.tp_group = tp_group or get_magi2_tp_group()
self.local_in_features = in_features
self.local_out_features = out_features
if parallel_mode == "column" and self.tp_group.world_size > 1:
split_dims = qkv_splits or (out_features,)
if any(size % self.tp_group.world_size for size in split_dims):
raise ValueError(
f"column-parallel output dimensions {split_dims} must divide TP={self.tp_group.world_size}"
)
self.local_out_features = sum(size // self.tp_group.world_size for size in split_dims)
elif parallel_mode == "row" and self.tp_group.world_size > 1:
if in_features % self.tp_group.world_size:
raise ValueError(f"row-parallel input {in_features} must divide TP={self.tp_group.world_size}")
self.local_in_features = in_features // self.tp_group.world_size
self.weight = nn.Parameter(
torch.empty(
num_experts * self.local_out_features,
self.local_in_features,
dtype=dtype,
device=device,
)
)
if bias:
bias_features = self.local_out_features if parallel_mode == "column" else out_features
self.bias = nn.Parameter(torch.empty(num_experts * bias_features, dtype=dtype, device=device))
else:
self.register_parameter("bias", None)
if self.tp_group.world_size > 1 and parallel_mode is not None:
self.weight.checkpoint_weight_transform = self.shard_checkpoint_weight
if self.bias is not None and parallel_mode == "column":
self.bias.checkpoint_weight_transform = self.shard_checkpoint_bias
def _column_slices(self) -> tuple[tuple[int, int], ...]:
splits = self.qkv_splits or (self.out_features,)
offsets: list[tuple[int, int]] = []
base = 0
for size in splits:
local = size // self.tp_group.world_size
start = base + self.tp_group.rank * local
offsets.append((start, start + local))
base += size
return tuple(offsets)
def shard_checkpoint_weight(self, checkpoint_tensor: torch.Tensor) -> torch.Tensor:
"""Convert the released grouped weight into this TP rank's shard."""
if tuple(checkpoint_tensor.shape) == tuple(self.weight.shape):
return checkpoint_tensor
expected = (self.num_experts * self.out_features, self.in_features)
if tuple(checkpoint_tensor.shape) != expected:
raise ValueError(
f"grouped linear checkpoint has shape {tuple(checkpoint_tensor.shape)}, expected {expected}"
)
grouped = checkpoint_tensor.view(self.num_experts, self.out_features, self.in_features)
if self.parallel_mode == "column":
shards = [grouped[:, start:end] for start, end in self._column_slices()]
return torch.cat(shards, dim=1).reshape(self.num_experts * self.local_out_features, self.in_features)
if self.parallel_mode == "row":
start = self.tp_group.rank * self.local_in_features
return grouped[:, :, start : start + self.local_in_features].reshape(
self.num_experts * self.out_features,
self.local_in_features,
)
return checkpoint_tensor
def shard_checkpoint_bias(self, checkpoint_tensor: torch.Tensor) -> torch.Tensor:
if self.bias is None:
raise ValueError("cannot load bias into a bias-free grouped linear")
if tuple(checkpoint_tensor.shape) == tuple(self.bias.shape):
return checkpoint_tensor
expected = (self.num_experts * self.out_features,)
if tuple(checkpoint_tensor.shape) != expected:
raise ValueError(f"grouped linear bias has shape {tuple(checkpoint_tensor.shape)}, expected {expected}")
grouped = checkpoint_tensor.view(self.num_experts, self.out_features)
shards = [grouped[:, start:end] for start, end in self._column_slices()]
return torch.cat(shards, dim=1).reshape(-1)
def forward(
self,
tensor: torch.Tensor,
modality_dispatcher: ModalityDispatcher | None = None,
) -> torch.Tensor:
weight = self.weight.view(self.num_experts, self.local_out_features, self.local_in_features)
bias_features = self.local_out_features if self.parallel_mode == "column" else self.out_features
bias = self.bias.view(self.num_experts, bias_features) if self.bias is not None else None
linear_bias = None if self.parallel_mode == "row" else bias
if self.num_experts == 1:
output = F.linear(tensor, weight[0], None if linear_bias is None else linear_bias[0])
else:
if modality_dispatcher is None:
raise ValueError("modality_dispatcher is required for grouped linear")
if modality_dispatcher.num_modalities != self.num_experts:
raise ValueError("grouped linear expert count does not match modality dispatcher")
inputs = modality_dispatcher.dispatch(tensor)
outputs = [
F.linear(part, weight[index], None if linear_bias is None else linear_bias[index])
for index, part in enumerate(inputs)
]
output = torch.cat(outputs, dim=0)
if self.parallel_mode == "row" and self.tp_group.world_size > 1:
dist.all_reduce(output, group=self.tp_group.group)
if self.parallel_mode == "row" and bias is not None:
if self.num_experts == 1:
output = output + bias[0]
else:
assert modality_dispatcher is not None
outputs = [part + bias[index] for index, part in enumerate(modality_dispatcher.dispatch(output))]
output = torch.cat(outputs, dim=0)
return output
def make_grouped_linear(
in_features: int,
out_features: int,
*,
num_experts: int = 1,
bias: bool = False,
dtype: torch.dtype | None = None,
parallel_mode: str | None = None,
qkv_splits: tuple[int, int, int] | None = None,
) -> Magi2GroupedLinear:
return Magi2GroupedLinear(
in_features,
out_features,
num_experts=num_experts,
bias=bias,
dtype=dtype,
parallel_mode=parallel_mode,
qkv_splits=qkv_splits,
)
class MultiModalityRMSNorm(nn.Module):
"""RMSNorm with independent modality and mHC-stream scales."""
def __init__(
self,
dim: int,
*,
eps: float = 1e-6,
num_modality: int = 1,
num_patterns: int = 1,
out_dtype: torch.dtype | None = None,
device: torch.device | str | None = None,
) -> None:
super().__init__()
self.dim = dim
self.eps = eps
self.num_modality = num_modality
self.num_patterns = num_patterns
self.out_dtype = out_dtype
self.weight = nn.Parameter(torch.zeros(num_patterns * dim * num_modality, dtype=torch.float32, device=device))
def forward(
self,
tensor: torch.Tensor,
modality_dispatcher: ModalityDispatcher | None = None,
) -> torch.Tensor:
original_dtype = tensor.dtype
normalized = tensor.float()
normalized = normalized * torch.rsqrt(normalized.square().mean(dim=-1, keepdim=True) + self.eps)
if self.num_modality == 1:
weight = self.weight.view(self.num_patterns, self.dim) + 1.0
result = normalized * weight
else:
if modality_dispatcher is None:
raise ValueError("modality_dispatcher is required for multimodal RMSNorm")
inputs = modality_dispatcher.dispatch(normalized)
weights = self.weight.view(self.num_modality, self.num_patterns, self.dim)
result = modality_dispatcher.undispatch(
*(part * (weights[index] + 1.0) for index, part in enumerate(inputs))
)
return result.to(self.out_dtype or original_dtype)
def _frequency_bands(
num_bands: int,
*,
temperature: float,
device: torch.device | str | None,
) -> torch.Tensor:
exponent = torch.arange(num_bands, dtype=torch.float32, device=device) / num_bands
return 1.0 / temperature**exponent
class ElementWiseFourierEmbed(nn.Module):
"""Nine-coordinate Fourier embedding used as MAGI's element-wise RoPE."""
def __init__(
self,
dim: int,
*,
temperature: float = 10000.0,
learnable: bool = False,
device: torch.device | str | None = None,
dtype: torch.dtype = torch.float32,
) -> None:
super().__init__()
self.dim = dim
self.temperature = temperature
bands = _frequency_bands(dim // 8, temperature=temperature, device=device).to(dtype)
if learnable:
self.bands = nn.Parameter(bands)
else:
# ``bands`` is part of the released checkpoint. Keep it in the
# parameter namespace so the generic DLO mmap loader can bind it
# without first materializing a full CPU transformer. It remains
# immutable during inference.
self.bands = nn.Parameter(bands, requires_grad=False)
def forward(self, coords: torch.Tensor) -> torch.Tensor:
if coords.ndim != 2 or coords.shape[-1] != 9:
raise ValueError("MAGI coordinates must have shape [tokens,9]")
xyz = coords[:, :3]
sizes = coords[:, 3:6]
references = coords[:, 6:9]
scales = (references - 1) / (sizes - 1)
scales = torch.where((references == 1) & (sizes == 1), torch.ones_like(scales), scales)
if not torch.isfinite(scales).all():
raise ValueError("invalid MAGI coordinate scale")
centers = (sizes - 1) / 2
centers = centers.clone()
centers[:, 0] = 0
projection = (xyz - centers).unsqueeze(-1) * scales.unsqueeze(-1) * self.bands
return torch.cat((projection.sin(), projection.cos()), dim=1).flatten(1)
MHCTensorTuple = tuple[torch.Tensor, torch.Tensor, torch.Tensor]
def sinkhorn_knopp(matrix_logits: torch.Tensor, iterations: int, epsilon: float) -> torch.Tensor:
matrix = torch.exp(matrix_logits - matrix_logits.amax(dim=(-2, -1), keepdim=True))
for _ in range(iterations):
matrix = matrix / (matrix.sum(dim=-2, keepdim=True) + epsilon)
matrix = matrix / (matrix.sum(dim=-1, keepdim=True) + epsilon)
return matrix
class MHCHandler:
"""Exact four-stream manifold-constrained hyper-connection math."""
def __init__(
self,
num_streams: int,
hidden_size: int,
*,
sinkhorn_iterations: int = 20,
sinkhorn_epsilon: float = 1e-12,
dtype: torch.dtype = torch.float32,
) -> None:
self.num_streams = num_streams
self.hidden_size = hidden_size
self.sinkhorn_iterations = sinkhorn_iterations
self.sinkhorn_epsilon = sinkhorn_epsilon
self.dtype = dtype
self.matmul_scale = 1.0 / math.sqrt(float(num_streams * hidden_size))
def flatten(self, tensor: torch.Tensor) -> torch.Tensor:
self._check_multi(tensor)
return tensor.view(tensor.shape[0], -1)
def compute_logits(
self,
flattened: torch.Tensor,
norm: Callable[[torch.Tensor], torch.Tensor],
phi_fused: torch.Tensor,
) -> MHCTensorTuple:
if flattened.ndim != 2 or flattened.shape[-1] != self.num_streams * self.hidden_size:
raise ValueError("invalid flattened mHC shape")
fused = norm(flattened).to(self.dtype) @ phi_fused
pre, post, residual = torch.split(
fused,
(self.num_streams, self.num_streams, self.num_streams**2),
dim=-1,
)
return pre, post, residual.view(-1, self.num_streams, self.num_streams)
def apply_pre(
self,
streams: torch.Tensor,
alpha_bias_logits: MHCTensorTuple,
*,
out_dtype: torch.dtype | None = None,
) -> torch.Tensor:
self._check_multi(streams)
alpha, bias, logits = alpha_bias_logits
coefficients = torch.sigmoid(alpha * self.matmul_scale * logits + bias.unsqueeze(0))
return torch.einsum("tn,tnc->tc", coefficients.to(out_dtype or streams.dtype), streams)
def compute_post_residual(
self,
post: MHCTensorTuple,
residual: MHCTensorTuple,
*,
out_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
alpha_post, bias_post, post_logits = post
alpha_residual, bias_residual, residual_logits = residual
post_coefficients = 2.0 * torch.sigmoid(alpha_post * self.matmul_scale * post_logits + bias_post.unsqueeze(0))
residual_matrix = sinkhorn_knopp(
alpha_residual * self.matmul_scale * residual_logits.float() + bias_residual.unsqueeze(0).float(),
self.sinkhorn_iterations,
self.sinkhorn_epsilon,
)
return post_coefficients.to(out_dtype), residual_matrix.to(out_dtype)
def hyper_connect(
self,
residual_streams: torch.Tensor,
branch_output: torch.Tensor,
post_coefficients: torch.Tensor,
residual_matrix: torch.Tensor,
) -> torch.Tensor:
self._check_multi(residual_streams)
if branch_output.ndim != 2 or branch_output.shape[-1] != self.hidden_size:
raise ValueError("invalid mHC branch-output shape")
branch = torch.einsum("tn,tc->tnc", post_coefficients, branch_output)
mixed = torch.einsum("tij,tjc->tic", residual_matrix, residual_streams)
return mixed + branch
def _check_multi(self, tensor: torch.Tensor) -> None:
expected = (self.num_streams, self.hidden_size)
if tensor.ndim != 3 or tensor.shape[1:] != expected:
raise ValueError(f"expected mHC tensor [tokens,{expected[0]},{expected[1]}], got {tuple(tensor.shape)}")
__all__ = [
"ElementWiseFourierEmbed",
"MHCHandler",
"Magi2GroupedLinear",
"ModalityDispatcher",
"MultiModalityRMSNorm",
"make_grouped_linear",
"sinkhorn_knopp",
"swiglu7",
]
vllm_omni/diffusion/models/magi2/mh_moe.py
+547
View File
@@ -0,0 +1,547 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Copyright (c) 2025-2026 SandAI. All Rights Reserved.
"""Native multi-head MoE used by MAGI-2 Preview.
Adapted from SandAI's Apache-2.0 ``flash_mh_moe`` implementation and modified
to use vLLM's existing expert-parallel group. MAGI's routing is unusual: each
of twelve 256-wide hidden-state heads independently selects experts from its
own 256-expert bank. It is therefore not representable by vLLM's conventional
whole-token :class:`FusedMoE` primitive.
"""
from __future__ import annotations
import math
import os
from dataclasses import dataclass
from typing import Literal
import torch
import torch.nn as nn
import torch.nn.functional as F
from vllm.triton_utils import tl, triton
from vllm_omni.platforms import current_omni_platform
from .parallel import Magi2ParallelGroup, ep_dispatch, ep_undispatch, get_magi2_ep_group
RoutingScore = Literal["softmax", "sigmoid"]
def swiglu7_pair(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
"""Released clamped SwiGLU7 expert activation, evaluated in fp32."""
dtype = gate.dtype
gate = gate.float().clamp(max=7.0)
up = up.float().clamp(min=-7.0, max=7.0)
return (gate * torch.sigmoid(1.702 * gate) * (up + 1.0)).to(dtype)
def compute_topk_probs_and_indices(
router_logits: torch.Tensor,
top_k: int,
*,
score_func: RoutingScore = "sigmoid",
expert_bias: torch.Tensor | None = None,
route_norm: bool = True,
norm_eps: float = 1e-12,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Route independently for every ``[head, token]`` pair.
The auxiliary-free bias affects expert selection but deliberately does not
affect the returned routing probability, matching the training recipe.
"""
if router_logits.ndim != 3:
raise ValueError("router_logits must be [heads,tokens,experts]")
if not 0 < top_k <= router_logits.shape[-1]:
raise ValueError("top_k must be in [1, num_experts]")
if score_func == "sigmoid":
router_scores = torch.sigmoid(router_logits)
elif score_func == "softmax":
router_scores = torch.softmax(router_logits, dim=-1)
else:
raise ValueError(f"unsupported routing score function {score_func!r}")
selection_scores = router_scores
if expert_bias is not None:
selection_scores = selection_scores + expert_bias.view(router_logits.shape[0], 1, -1)
# Keep the reference's default sorted=True behavior. Besides defining the
# route order for ties, this also fixes the reduction order used by the
# following L1 normalization.
topk_indices = torch.topk(selection_scores, top_k, dim=-1).indices
topk_probs = router_scores.gather(-1, topk_indices)
if route_norm:
topk_probs = F.normalize(topk_probs, p=1, dim=-1, eps=norm_eps)
return topk_probs, topk_indices
def global_sort_routes(
topk_probs: torch.Tensor,
topk_indices: torch.Tensor,
num_experts: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Convert per-head routes into a stable flattened-expert CSR layout."""
if topk_probs.shape != topk_indices.shape or topk_indices.ndim != 3:
raise ValueError("top-k probabilities and indices must have the same [H,S,K] shape")
heads, sequence, top_k = topk_indices.shape
device = topk_indices.device
head_offset = torch.arange(heads, device=device).view(heads, 1, 1) * num_experts
flattened_experts = (topk_indices + head_offset).reshape(-1)
flat_probs = topk_probs.reshape(-1)
flat_tokens = torch.arange(sequence, device=device).view(1, sequence, 1).expand(heads, sequence, top_k).reshape(-1)
order = flattened_experts.argsort(stable=True)
gather_ids = flat_tokens[order].to(torch.int32)
sorted_probs = flat_probs[order].float()
counts = torch.bincount(flattened_experts, minlength=heads * num_experts)
offsets = torch.zeros(heads * num_experts + 1, device=device, dtype=torch.long)
offsets[1:] = counts.cumsum(0)
return gather_ids, sorted_probs, offsets
def torch_mh_moe_forward(
x: torch.Tensor,
gather_ids: torch.Tensor,
probs: torch.Tensor,
expert_offsets: torch.Tensor,
w_gate: torch.Tensor,
w_up: torch.Tensor,
w_down: torch.Tensor,
) -> torch.Tensor:
"""Small-shape correctness oracle for the fused expert kernel."""
if x.ndim != 3:
raise ValueError("multi-head MoE input must be [tokens,heads,head_dim]")
output = torch.zeros_like(x)
experts_per_head = (expert_offsets.numel() - 1) // x.shape[1]
for flat_expert in range(expert_offsets.numel() - 1):
begin = int(expert_offsets[flat_expert].item())
end = int(expert_offsets[flat_expert + 1].item())
if begin == end:
continue
head = flat_expert // experts_per_head
token_ids = gather_ids[begin:end].long()
expert_input = x.index_select(0, token_ids)[:, head]
gate = expert_input @ w_gate[flat_expert]
up = expert_input @ w_up[flat_expert]
hidden = swiglu7_pair(gate, up)
expert_output = hidden @ w_down[flat_expert]
expert_output = expert_output * probs[begin:end, None].to(expert_output.dtype)
output[:, head].index_add_(0, token_ids, expert_output)
return output
_SWIGLU7_ALPHA = tl.constexpr(1.702)
_SWIGLU7_LIMIT = tl.constexpr(7.0)
_SWIGLU7_BIAS = tl.constexpr(1.0)
@triton.jit
def _swiglu7_kernel(gate, up, out_dtype: tl.constexpr):
gate_clamped = tl.minimum(gate, _SWIGLU7_LIMIT)
up_clamped = tl.maximum(tl.minimum(up, _SWIGLU7_LIMIT), -_SWIGLU7_LIMIT)
sigmoid = tl.sigmoid(_SWIGLU7_ALPHA * gate_clamped)
swish = gate_clamped * sigmoid
return (swish * (up_clamped + _SWIGLU7_BIAS)).to(out_dtype)
@triton.jit
def _binary_search_expert(
cumulative_tiles,
tile_id,
num_experts: tl.constexpr,
log2_num_experts: tl.constexpr,
):
lo = 0
hi = num_experts
for _ in tl.static_range(0, log2_num_experts + 1):
mid = (lo + hi + 1) // 2
below = tl.load(cumulative_tiles + mid) <= tile_id
lo = tl.where(below, mid, lo)
hi = tl.where(below, hi, mid - 1)
return lo
@triton.jit
def _mh_moe_kernel(
x_ptr,
wg_ptr,
wu_ptr,
wd_ptr,
y_ptr,
gather_ids_ptr,
probs_ptr,
expert_offsets_ptr,
cumulative_tiles_ptr,
stride_x_s,
stride_x_h,
stride_x_dh,
stride_wg_e,
stride_wg_dh,
stride_wg_de,
stride_wu_e,
stride_wu_dh,
stride_wu_de,
stride_wd_e,
stride_wd_de,
stride_wd_dh,
stride_y_s,
stride_y_h,
stride_y_dh,
d_head: tl.constexpr,
d_expert: tl.constexpr,
num_heads: tl.constexpr,
num_flat_experts: tl.constexpr,
log2_num_experts: tl.constexpr,
block_t: tl.constexpr,
block_dh: tl.constexpr,
block_de: tl.constexpr,
acc_dtype: tl.constexpr = tl.float32,
deterministic: tl.constexpr = False,
):
tile_id = tl.program_id(0)
total_tiles = tl.load(cumulative_tiles_ptr + num_flat_experts)
if tile_id >= total_tiles:
return
expert = _binary_search_expert(cumulative_tiles_ptr, tile_id, num_flat_experts, log2_num_experts)
expert_i64 = expert.to(tl.int64)
head = expert // (num_flat_experts // num_heads)
tile_in_expert = tile_id - tl.load(cumulative_tiles_ptr + expert)
token_start = tl.load(expert_offsets_ptr + expert) + tile_in_expert * block_t
expert_end = tl.load(expert_offsets_ptr + expert + 1)
count = tl.minimum(token_start + block_t, expert_end) - token_start
dh_block_offsets = tl.arange(0, block_dh)
de_block_offsets = tl.arange(0, block_de)
token_offsets = tl.arange(0, block_t)
dh_offsets = tl.arange(0, d_head)
token_positions = token_start + token_offsets
token_mask = token_offsets < count
# Token indices fit in int32, but multiplying a large packed-batch index by
# the hidden-width stride does not. Promote before computing element offsets.
gather_ids = tl.load(gather_ids_ptr + token_positions, mask=token_mask, other=0).to(tl.int64)
probabilities = tl.load(probs_ptr + token_positions, mask=token_mask, other=0.0)
x_base = gather_ids * stride_x_s + head * stride_x_h
output_acc = tl.zeros([block_t, d_head], dtype=acc_dtype)
for de_start in tl.range(0, d_expert, block_de):
de_offsets = de_start + de_block_offsets
gate_acc = tl.zeros([block_t, block_de], dtype=acc_dtype)
up_acc = tl.zeros([block_t, block_de], dtype=acc_dtype)
for dh_start in tl.static_range(0, d_head, block_dh):
local_dh = dh_start + dh_block_offsets
x_block = tl.load(
x_ptr + x_base[:, None] + local_dh[None, :] * stride_x_dh,
mask=token_mask[:, None],
other=0.0,
)
wg = tl.load(
wg_ptr
+ expert_i64 * stride_wg_e
+ local_dh[:, None] * stride_wg_dh
+ de_offsets[None, :] * stride_wg_de
)
wu = tl.load(
wu_ptr
+ expert_i64 * stride_wu_e
+ local_dh[:, None] * stride_wu_dh
+ de_offsets[None, :] * stride_wu_de
)
gate_acc += tl.dot(x_block, wg)
up_acc += tl.dot(x_block, wu)
hidden = _swiglu7_kernel(gate_acc, up_acc, wd_ptr.dtype.element_ty)
down = tl.load(
wd_ptr + expert_i64 * stride_wd_e + de_offsets[:, None] * stride_wd_de + dh_offsets[None, :] * stride_wd_dh
)
output_acc += tl.dot(hidden, down)
output_acc = output_acc * probabilities[:, None]
if deterministic:
output_ptrs = y_ptr + token_positions[:, None] * stride_y_s + dh_offsets[None, :] * stride_y_dh
tl.store(output_ptrs, output_acc.to(y_ptr.dtype.element_ty), mask=token_mask[:, None])
else:
output_base = gather_ids * stride_y_s + head * stride_y_h
output_ptrs = y_ptr + output_base[:, None] + dh_offsets[None, :] * stride_y_dh
tl.atomic_add(output_ptrs, output_acc.to(y_ptr.dtype.element_ty), mask=token_mask[:, None])
def _deterministic_scatter(
sorted_output: torch.Tensor,
reference: torch.Tensor,
gather_ids: torch.Tensor,
expert_offsets: torch.Tensor,
) -> torch.Tensor:
num_flat_experts = expert_offsets.numel() - 1
experts_per_head = num_flat_experts // reference.shape[1]
expert_lengths = torch.diff(expert_offsets)
head_values = torch.arange(num_flat_experts, device=gather_ids.device) // experts_per_head
head_ids = torch.repeat_interleave(head_values, expert_lengths)
scatter_ids = gather_ids.long() * reference.shape[1] + head_ids.long()
output = torch.zeros_like(reference).view(-1, reference.shape[-1])
output.scatter_add_(0, scatter_ids[:, None].expand_as(sorted_output), sorted_output.to(output.dtype))
return output.view_as(reference)
def _select_block_config() -> tuple[int, int, int, int, int]:
"""Return the reference kernel config, capped for pre-Blackwell GPUs."""
capability = current_omni_platform.get_device_capability()
if capability is not None and capability.major >= 10: # Blackwell
return (128, 64, 32, 2, 8)
# BLOCK_T=128 needs 122,880 bytes of shared memory and is not safe on the
# qualified L20X path. This is the reference kernel's portable config.
return (64, 64, 32, 2, 4)
def triton_mh_moe_forward(
x: torch.Tensor,
gather_ids: torch.Tensor,
probs: torch.Tensor,
expert_offsets: torch.Tensor,
w_gate: torch.Tensor,
w_up: torch.Tensor,
w_down: torch.Tensor,
*,
deterministic: bool = False,
) -> torch.Tensor:
"""Fused gather/expert/scatter kernel for released MAGI dimensions."""
routed_tokens = gather_ids.numel()
if routed_tokens == 0:
return torch.zeros_like(x)
d_head, d_expert = x.shape[-1], w_down.shape[1]
block_t, block_dh, block_de, num_stages, num_warps = _select_block_config()
if d_head % block_dh or d_expert % block_de:
return torch_mh_moe_forward(x, gather_ids, probs, expert_offsets, w_gate, w_up, w_down)
if deterministic:
output = torch.empty((routed_tokens, 1, d_head), device=x.device, dtype=x.dtype)
else:
output = torch.zeros_like(x)
num_flat_experts = expert_offsets.numel() - 1
expert_tiles = (torch.diff(expert_offsets) + block_t - 1) // block_t
cumulative_tiles = torch.cat(
(torch.zeros(1, dtype=torch.int32, device=x.device), expert_tiles.cumsum(0, dtype=torch.int32))
)
# Match the reference launch bound. Empty/excess programs return after
# comparing against ``cumulative_tiles[-1]`` inside the kernel.
grid = ((routed_tokens + block_t - 1) // block_t + num_flat_experts,)
log2_experts = max(1, math.ceil(math.log2(max(num_flat_experts, 1) + 1)))
_mh_moe_kernel[grid](
x,
w_gate,
w_up,
w_down,
output,
gather_ids,
probs,
expert_offsets,
cumulative_tiles,
x.stride(0),
x.stride(1),
x.stride(2),
w_gate.stride(0),
w_gate.stride(1),
w_gate.stride(2),
w_up.stride(0),
w_up.stride(1),
w_up.stride(2),
w_down.stride(0),
w_down.stride(1),
w_down.stride(2),
output.stride(0),
output.stride(1),
output.stride(2),
d_head,
d_expert,
x.shape[1],
num_flat_experts,
log2_experts,
block_t,
block_dh,
block_de,
tl.float32,
deterministic,
num_stages=num_stages,
num_warps=num_warps,
)
if deterministic:
return _deterministic_scatter(output.view(routed_tokens, d_head), x, gather_ids, expert_offsets)
return output
@dataclass(frozen=True)
class Magi2MultiHeadMoEConfig:
hidden_size: int
num_heads: int
num_experts: int
top_k: int
expert_intermediate_size: int
params_dtype: torch.dtype
score_func: RoutingScore = "sigmoid"
route_norm: bool = True
route_scale: float = 1.0
class Magi2MultiHeadMoE(nn.Module):
"""Checkpoint-compatible MAGI-2 head-routed expert layer."""
_EP_SHARDED_PARAMETER_NAMES = frozenset(
{"gate", "W_gate", "W_up", "W_down", "router.expert_bias", "router.expert_bias_ema"}
)
def __init__(
self,
config: Magi2MultiHeadMoEConfig,
*,
ep_group: Magi2ParallelGroup | None = None,
) -> None:
super().__init__()
if config.hidden_size % config.num_heads:
raise ValueError("hidden_size must be divisible by the number of MoE heads")
self.config = config
self.num_heads = config.num_heads
self.num_experts = config.num_experts
self.top_k = config.top_k
self.d_head = config.hidden_size // config.num_heads
self.d_expert = config.expert_intermediate_size
self.ep_group = ep_group or get_magi2_ep_group()
self.padded_num_heads = math.ceil(self.num_heads / self.ep_group.world_size) * self.ep_group.world_size
self.local_num_heads = self.padded_num_heads // self.ep_group.world_size
self.local_flatten_num_experts = self.local_num_heads * self.num_experts
self.ep_pad_heads = self.padded_num_heads - self.num_heads
self.local_head_start = self.ep_group.rank * self.local_num_heads
self.has_real_moe_heads = self.local_head_start < self.num_heads
self.gate = nn.Parameter(torch.empty(self.local_flatten_num_experts, self.d_head, dtype=torch.float32))
self.W_gate = nn.Parameter(
torch.empty(self.local_flatten_num_experts, self.d_head, self.d_expert, dtype=config.params_dtype)
)
self.W_up = nn.Parameter(
torch.empty(self.local_flatten_num_experts, self.d_head, self.d_expert, dtype=config.params_dtype)
)
self.W_down = nn.Parameter(
torch.empty(self.local_flatten_num_experts, self.d_expert, self.d_head, dtype=config.params_dtype)
)
self.router = nn.Module()
# Both tensors are released checkpoint entries. Non-trainable
# Parameters let the DLO mmap path bind them on a meta-constructed
# model; persistent buffers are intentionally not mmap-loaded by the
# generic backend.
self.router.expert_bias = nn.Parameter(
torch.zeros(self.local_flatten_num_experts, dtype=torch.float32),
requires_grad=False,
)
self.router.expert_bias_ema = nn.Parameter(
torch.zeros(self.local_flatten_num_experts, dtype=torch.float32),
requires_grad=False,
)
for name in self._EP_SHARDED_PARAMETER_NAMES:
target: nn.Module | Magi2MultiHeadMoE = self
parts = name.split(".")
for part in parts[:-1]:
target = getattr(target, part)
parameter = getattr(target, parts[-1])
parameter.mmap_weight_transform = self.ep_slice
def ep_slice(self, checkpoint_tensor: torch.Tensor) -> torch.Tensor:
"""Slice flattened ``(head,expert)`` checkpoint rows for this rank."""
if checkpoint_tensor.shape[0] == self.local_flatten_num_experts:
return checkpoint_tensor
start = self.local_head_start * self.num_experts
end = min(start + self.local_flatten_num_experts, checkpoint_tensor.shape[0])
if start >= checkpoint_tensor.shape[0]:
return torch.zeros(
(self.local_flatten_num_experts, *checkpoint_tensor.shape[1:]),
dtype=checkpoint_tensor.dtype,
device=checkpoint_tensor.device,
)
local = checkpoint_tensor[start:end]
if local.shape[0] < self.local_flatten_num_experts:
# Uneven EP/head partitions require materialized zero padding;
# divisible production layouts keep the mmap-backed slice above.
padding = torch.zeros(
(self.local_flatten_num_experts - local.shape[0], *local.shape[1:]),
dtype=local.dtype,
device=local.device,
)
local = torch.cat((local, padding), dim=0)
return local
def _route(self, x_heads: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
gate = self.gate.view(self.local_num_heads, self.num_experts, self.d_head).float()
logits = torch.einsum("shd,hed->hse", x_heads.float(), gate)
bias_source = (os.environ.get("MAGI2_ROUTER_BIAS_SOURCE") or "ema").strip().lower()
bias_tensor = self.router.expert_bias if bias_source == "main" else self.router.expert_bias_ema
bias = bias_tensor.view(self.local_num_heads, self.num_experts)
probs, indices = compute_topk_probs_and_indices(
logits,
self.top_k,
score_func=self.config.score_func,
expert_bias=bias,
route_norm=self.config.route_norm,
)
return probs * self.config.route_scale, indices
def _local_forward(self, x_heads: torch.Tensor) -> torch.Tensor:
probabilities, indices = self._route(x_heads)
gather_ids, sorted_probs, offsets = global_sort_routes(probabilities, indices, self.num_experts)
if x_heads.is_cuda:
return triton_mh_moe_forward(
x_heads,
gather_ids,
sorted_probs,
offsets,
self.W_gate,
self.W_up,
self.W_down,
deterministic=os.environ.get("MAGI2_DETERMINISTIC", "0") == "1",
)
return torch_mh_moe_forward(x_heads, gather_ids, sorted_probs, offsets, self.W_gate, self.W_up, self.W_down)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.ep_group.world_size > 1 and self.ep_group.replicated_sequence:
# TP column-parallel ``split_linear`` already emits exactly this
# rank's contiguous MoE-head slice. Compute it once and leave it
# sharded for the row-parallel ``merge_linear``; no token dispatch
# or head all-gather belongs on the true TP path.
local_hidden_size = self.local_num_heads * self.d_head
if x.shape[-1] != local_hidden_size:
raise ValueError(f"TP-local MAGI MoE input has width {x.shape[-1]}, expected {local_hidden_size}")
local = x.view(-1, self.local_num_heads, self.d_head)
output = self._local_forward(local) if self.has_real_moe_heads else torch.zeros_like(local)
return output.reshape(-1, local_hidden_size)
x_heads = x.view(-1, self.num_heads, self.d_head)
if self.ep_pad_heads:
padding = x_heads.new_zeros((x_heads.shape[0], self.ep_pad_heads, self.d_head))
x_heads = torch.cat((x_heads, padding), dim=1)
sequence_split_sizes: list[int] | None = None
if self.ep_group.world_size > 1:
local_size = torch.tensor([x_heads.shape[0]], dtype=torch.int64, device=x_heads.device)
gathered_sizes = [torch.empty_like(local_size) for _ in range(self.ep_group.world_size)]
torch.distributed.all_gather(gathered_sizes, local_size, group=self.ep_group.group)
sequence_split_sizes = [int(size.item()) for size in gathered_sizes]
x_heads = ep_dispatch(x_heads, self.ep_group, sequence_split_sizes)
output = self._local_forward(x_heads) if self.has_real_moe_heads else torch.zeros_like(x_heads)
if self.ep_group.world_size > 1:
output = ep_undispatch(output, self.ep_group, sequence_split_sizes)
if self.ep_pad_heads:
output = output[:, : self.num_heads]
return output.reshape(-1, self.num_heads * self.d_head)
__all__ = [
"Magi2MultiHeadMoE",
"Magi2MultiHeadMoEConfig",
"compute_topk_probs_and_indices",
"global_sort_routes",
"torch_mh_moe_forward",
"triton_mh_moe_forward",
]
vllm_omni/diffusion/models/magi2/modeling_magi2.py
+739
View File
@@ -0,0 +1,739 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Copyright (c) 2026 SandAI. All Rights Reserved.
"""Native MAGI-2 Preview diffusion transformer.
The architecture and mHC/MoE equations are adapted from SandAI's Apache-2.0
MAGI-2 Preview implementation. This modified vLLM-Omni version removes the
external ``inference`` package, MagiCompiler, and SandAI process manager while
preserving the released checkpoint module names and tensor layouts.
"""
from __future__ import annotations
import os
from collections.abc import Iterable
from enum import IntEnum
from functools import partial
import torch
import torch.nn as nn
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.layer import Attention
from .attention import Magi2PackedAttentionKernel, VarlenHandler, apply_rotary_emb
from .configuration_magi2 import Magi2PreviewConfig
from .layers import (
ElementWiseFourierEmbed,
MHCHandler,
ModalityDispatcher,
MultiModalityRMSNorm,
make_grouped_linear,
swiglu7,
)
from .mh_moe import Magi2MultiHeadMoE, Magi2MultiHeadMoEConfig
from .parallel import Magi2SequenceDispatcher
class Modality(IntEnum):
VIDEO = 0
AUDIO = 1
TEXT = 2
TIME = 3
class Magi2Attention(nn.Module):
def __init__(
self,
config: Magi2PreviewConfig,
*,
num_modality: int,
) -> None:
super().__init__()
self.config = config
self.num_modality = num_modality
self.head_dim = config.head_dim
global_q_size = config.num_heads_q * self.head_dim
global_kv_size = config.num_heads_kv * self.head_dim
self.pre_norm = MultiModalityRMSNorm(
config.hidden_size,
num_modality=num_modality,
)
self.linear_g = make_grouped_linear(
config.hidden_size,
config.num_heads_q,
num_experts=num_modality,
bias=False,
dtype=config.params_dtype,
parallel_mode="column",
)
self.linear_qkv = make_grouped_linear(
config.hidden_size,
global_q_size + 2 * global_kv_size,
num_experts=num_modality,
bias=False,
dtype=config.params_dtype,
parallel_mode="column",
qkv_splits=(global_q_size, global_kv_size, global_kv_size),
)
self.linear_proj = make_grouped_linear(
global_q_size,
config.hidden_size,
num_experts=num_modality,
bias=False,
dtype=config.params_dtype,
parallel_mode="row",
)
self.tp_group = self.linear_qkv.tp_group
self.num_heads_q = config.num_heads_q // self.tp_group.world_size
self.num_heads_kv = config.num_heads_kv // self.tp_group.world_size
self.q_size = self.num_heads_q * self.head_dim
self.kv_size = self.num_heads_kv * self.head_dim
self.sinks = nn.Parameter(torch.empty(config.attention_sink_tokens, self.num_heads_q, dtype=torch.float32))
if self.tp_group.world_size > 1:
self.sinks.checkpoint_weight_transform = self._shard_sinks
self.q_norm = MultiModalityRMSNorm(
self.head_dim,
num_modality=num_modality,
out_dtype=torch.float32,
)
self.k_norm = MultiModalityRMSNorm(
self.head_dim,
num_modality=num_modality,
out_dtype=torch.float32,
)
# MAGI-2's packed CFG sequences, learned sink logits, softcap, and
# uneven Ulysses splits require a model kernel. Route it through the
# framework Attention layer so compile/dispatch ownership remains
# shared while the model kernel owns its specialized communication.
self.packed_attention = Attention(
num_heads=self.num_heads_q,
num_kv_heads=self.num_heads_kv,
head_size=self.head_dim,
causal=False,
softmax_scale=self.head_dim**-0.5,
qkv_layout="THD",
skip_sequence_parallel=True,
disable_kv_quant=True,
custom_attention=Magi2PackedAttentionKernel(config.attention_softcap),
)
def _shard_sinks(self, checkpoint_tensor: torch.Tensor) -> torch.Tensor:
if tuple(checkpoint_tensor.shape) == tuple(self.sinks.shape):
return checkpoint_tensor
expected = (self.config.attention_sink_tokens, self.config.num_heads_q)
if tuple(checkpoint_tensor.shape) != expected:
raise ValueError(f"attention sinks have shape {tuple(checkpoint_tensor.shape)}, expected {expected}")
start = self.tp_group.rank * self.num_heads_q
return checkpoint_tensor[:, start : start + self.num_heads_q]
def forward(
self,
hidden_states: torch.Tensor,
rope: torch.Tensor,
varlen_handler: VarlenHandler,
modality_dispatcher: ModalityDispatcher,
cp_split_sizes: list[int] | torch.Tensor,
) -> torch.Tensor:
normalized = self.pre_norm(hidden_states, modality_dispatcher)
gates = self.linear_g(normalized, modality_dispatcher)
qkv = self.linear_qkv(normalized, modality_dispatcher)
q, k, v = torch.split(qkv, (self.q_size, self.kv_size, self.kv_size), dim=-1)
q = q.view(-1, self.num_heads_q, self.head_dim)
k = k.view(-1, self.num_heads_kv, self.head_dim)
v = v.view(-1, self.num_heads_kv, self.head_dim)
gates = gates.view(-1, self.num_heads_q, 1)
q = self.q_norm(q, modality_dispatcher)
k = self.k_norm(k, modality_dispatcher)
q = modality_dispatcher.inverse_permute(q).unsqueeze(0)
k = modality_dispatcher.inverse_permute(k).unsqueeze(0)
v = modality_dispatcher.inverse_permute(v).unsqueeze(0)
sin, cos = rope.tensor_split(2, dim=-1)
# The released implementation normalizes Q/K in fp32, applies RoPE,
# then converts Q/K/V to the checkpoint dtype at the attention
# boundary. This conversion is also required by FlashAttention.
q = apply_rotary_emb(q, cos, sin).squeeze(0).to(self.config.params_dtype)
k = apply_rotary_emb(k, cos, sin).squeeze(0).to(self.config.params_dtype)
v = v.squeeze(0).to(self.config.params_dtype)
output = self.packed_attention(
q,
k,
v,
AttentionMetadata(
extra={
"magi2_varlen": varlen_handler,
"magi2_split_sizes": cp_split_sizes,
"magi2_sink": self.sinks,
}
),
)
output = modality_dispatcher.permute(output)
output = output * torch.sigmoid(gates)
output = output.reshape(-1, self.q_size).to(self.config.params_dtype)
return self.linear_proj(output, modality_dispatcher)
class Magi2MLP(nn.Module):
def __init__(
self,
config: Magi2PreviewConfig,
*,
num_modality: int,
) -> None:
super().__init__()
intermediate_size = int(config.hidden_size * config.intermediate_factor * 2 / 3) // 128 * 128
# Tiny unit-test configurations can be narrower than the release's
# 128-channel rounding unit while retaining the same formula family.
if intermediate_size == 0:
intermediate_size = max(1, int(config.hidden_size * config.intermediate_factor * 2 / 3))
self.intermediate_size = intermediate_size
self.pre_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality)
self.up_gate_proj = make_grouped_linear(
config.hidden_size,
2 * intermediate_size,
num_experts=num_modality,
bias=False,
dtype=config.params_dtype,
parallel_mode="column",
)
self.down_proj = make_grouped_linear(
intermediate_size,
config.hidden_size,
num_experts=num_modality,
bias=False,
dtype=config.params_dtype,
parallel_mode="row",
)
def forward(self, hidden_states: torch.Tensor, dispatcher: ModalityDispatcher) -> torch.Tensor:
hidden_states = self.pre_norm(hidden_states, dispatcher)
hidden_states = self.up_gate_proj(hidden_states, dispatcher)
hidden_states = swiglu7(hidden_states)
return self.down_proj(hidden_states, dispatcher)
class Magi2MultiHeadMoELayer(nn.Module):
def __init__(self, config: Magi2PreviewConfig) -> None:
super().__init__()
moe = config.moe
self.config = config
self.pre_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=3)
self.split_linear = make_grouped_linear(
config.hidden_size,
config.hidden_size,
bias=False,
dtype=config.params_dtype,
parallel_mode="column",
)
self.moe_mlp = Magi2MultiHeadMoE(
Magi2MultiHeadMoEConfig(
hidden_size=config.hidden_size,
num_heads=moe.num_heads,
num_experts=moe.num_experts,
top_k=moe.top_k,
expert_intermediate_size=moe.expert_intermediate_size,
params_dtype=config.params_dtype,
score_func=moe.score_function, # type: ignore[arg-type]
route_norm=moe.normalize_routing_weights,
route_scale=moe.routing_scale,
)
)
self.merge_linear = make_grouped_linear(
config.hidden_size,
config.hidden_size,
bias=False,
dtype=config.params_dtype,
parallel_mode="row",
)
self.shared_expert_fc1 = make_grouped_linear(
config.hidden_size,
2 * moe.shared_expert_intermediate_size,
bias=False,
dtype=config.params_dtype,
parallel_mode="column",
)
self.shared_expert_fc2 = make_grouped_linear(
moe.shared_expert_intermediate_size,
config.hidden_size,
bias=False,
dtype=config.params_dtype,
parallel_mode="row",
)
self.modality_specific_shared_expert_fc1 = make_grouped_linear(
config.hidden_size,
2 * moe.modality_shared_expert_intermediate_size,
num_experts=3,
bias=False,
dtype=config.params_dtype,
parallel_mode="column",
)
self.modality_specific_shared_expert_fc2 = make_grouped_linear(
moe.modality_shared_expert_intermediate_size,
config.hidden_size,
num_experts=3,
bias=False,
dtype=config.params_dtype,
parallel_mode="row",
)
tp_size = self.split_linear.tp_group.world_size
self.local_shared_expert_intermediate_size = moe.shared_expert_intermediate_size // tp_size
self.local_modality_shared_expert_intermediate_size = moe.modality_shared_expert_intermediate_size // tp_size
def _shared_experts(
self,
normalized: torch.Tensor,
dispatcher: ModalityDispatcher,
) -> torch.Tensor:
shared = self.shared_expert_fc1(normalized)
modality = self.modality_specific_shared_expert_fc1(normalized, dispatcher)
activated = swiglu7(torch.cat((shared, modality), dim=-1))
shared, modality = activated.split(
(
self.local_shared_expert_intermediate_size,
self.local_modality_shared_expert_intermediate_size,
),
dim=-1,
)
return self.shared_expert_fc2(shared) + self.modality_specific_shared_expert_fc2(
modality.contiguous(), dispatcher
)
def forward(self, hidden_states: torch.Tensor, dispatcher: ModalityDispatcher) -> torch.Tensor:
normalized = self.pre_norm(hidden_states, dispatcher)
routed = self.split_linear(normalized)
routed = self.moe_mlp(routed)
routed = self.merge_linear(routed)
return routed + self._shared_experts(normalized, dispatcher)
class Magi2PreAdapter(nn.Module):
def __init__(self, config: Magi2PreviewConfig) -> None:
super().__init__()
self.config = config
self.adapter_dim = config.virtual_width
self.video_embedder = nn.Linear(
config.video_in_channels,
self.adapter_dim,
bias=True,
dtype=torch.float32,
)
self.text_embedder = nn.Linear(
config.text_in_channels,
self.adapter_dim,
bias=True,
dtype=torch.float32,
)
self.audio_embedder = nn.Linear(
config.audio_in_channels,
self.adapter_dim,
bias=True,
dtype=torch.float32,
)
self.rope = ElementWiseFourierEmbed(config.head_dim, learnable=False)
def forward(
self,
packed: torch.Tensor,
coords_mapping: torch.Tensor,
video_indices: torch.Tensor,
audio_indices: torch.Tensor,
text_indices: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
rope = self.rope(coords_mapping)
output = torch.zeros(
packed.shape[0],
self.adapter_dim,
dtype=torch.float32,
device=packed.device,
)
if text_indices.numel():
output.index_copy_(
0,
text_indices,
self.text_embedder(packed.index_select(0, text_indices)[:, : self.config.text_in_channels].float()),
)
if audio_indices.numel():
output.index_copy_(
0,
audio_indices,
self.audio_embedder(packed.index_select(0, audio_indices)[:, : self.config.audio_in_channels].float()),
)
if video_indices.numel():
output.index_copy_(
0,
video_indices,
self.video_embedder(packed.index_select(0, video_indices)[:, : self.config.video_in_channels].float()),
)
return output, rope
class Magi2PostAdapter(nn.Module):
def __init__(self, config: Magi2PreviewConfig) -> None:
super().__init__()
self.config = config
self.adapter_dim = config.virtual_width
self.final_norm_video = MultiModalityRMSNorm(self.adapter_dim)
self.final_norm_audio = MultiModalityRMSNorm(self.adapter_dim)
self.final_linear_video = nn.Linear(
self.adapter_dim,
config.video_in_channels,
bias=False,
dtype=torch.float32,
)
self.final_linear_audio = nn.Linear(
self.adapter_dim,
config.audio_in_channels,
bias=False,
dtype=torch.float32,
)
self.final_out_dim = max(config.video_in_channels, config.audio_in_channels)
def forward(
self,
hidden_states: torch.Tensor,
video_indices: torch.Tensor,
audio_indices: torch.Tensor,
) -> torch.Tensor:
output = torch.zeros(
hidden_states.shape[0],
self.final_out_dim,
dtype=torch.float32,
device=hidden_states.device,
)
if video_indices.numel():
video = self.final_norm_video(hidden_states.index_select(0, video_indices).float())
video = self.final_linear_video(video).float()
output[:, : self.config.video_in_channels].index_copy_(0, video_indices, video)
if audio_indices.numel():
audio = self.final_norm_audio(hidden_states.index_select(0, audio_indices).float())
audio = self.final_linear_audio(audio).float()
output[:, : self.config.audio_in_channels].index_copy_(0, audio_indices, audio)
return output
class Magi2TransformerLayer(nn.Module):
def __init__(self, config: Magi2PreviewConfig, layer_index: int) -> None:
super().__init__()
num_modality = 3 if layer_index in config.multimodal_layers else 1
self.config = config
self.attention = Magi2Attention(config, num_modality=num_modality)
self.mlp: nn.Module
if layer_index in config.moe.layers:
self.mlp = Magi2MultiHeadMoELayer(config)
else:
self.mlp = Magi2MLP(config, num_modality=num_modality)
self._init_mhc(num_modality)
def _init_mhc(self, num_modality: int) -> None:
streams = self.config.mhc.num_streams
hidden = self.config.hidden_size
alpha = self.config.mhc.alpha_init
for branch in ("attn", "mlp"):
for coefficient in ("pre", "post", "res"):
self.register_parameter(
f"mhc_alpha_{coefficient}_{branch}",
nn.Parameter(torch.full((1,), alpha, dtype=torch.float32)),
)
self.register_parameter(
f"mhc_bias_pre_{branch}",
nn.Parameter(torch.empty(streams, dtype=torch.float32)),
)
self.register_parameter(
f"mhc_bias_post_{branch}",
nn.Parameter(torch.empty(streams, dtype=torch.float32)),
)
self.register_parameter(
f"mhc_bias_res_{branch}",
nn.Parameter(torch.empty(streams, streams, dtype=torch.float32)),
)
self.register_parameter(
f"mhc_phi_fused_{branch}",
nn.Parameter(
torch.empty(
streams * hidden,
streams + streams + streams * streams,
dtype=torch.float32,
)
),
)
self.mhc_norm = MultiModalityRMSNorm(
streams * hidden,
num_modality=num_modality,
out_dtype=torch.float32,
)
self.mhc_handler = MHCHandler(
streams,
hidden,
sinkhorn_iterations=self.config.mhc.sinkhorn_iterations,
sinkhorn_epsilon=self.config.mhc.sinkhorn_epsilon,
)
def _branch_logits(
self,
streams: torch.Tensor,
branch: str,
dispatcher: ModalityDispatcher,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
return self.mhc_handler.compute_logits(
self.mhc_handler.flatten(streams),
partial(self.mhc_norm, modality_dispatcher=dispatcher),
getattr(self, f"mhc_phi_fused_{branch}"),
)
def _branch_input(
self,
streams: torch.Tensor,
branch: str,
logits: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
) -> torch.Tensor:
pre_logits, _, _ = logits
return self.mhc_handler.apply_pre(
streams,
(
getattr(self, f"mhc_alpha_pre_{branch}"),
getattr(self, f"mhc_bias_pre_{branch}"),
pre_logits,
),
out_dtype=streams.dtype,
)
def _connect(
self,
streams: torch.Tensor,
output: torch.Tensor,
branch: str,
logits: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
) -> torch.Tensor:
_, post_logits, residual_logits = logits
post, residual = self.mhc_handler.compute_post_residual(
(
getattr(self, f"mhc_alpha_post_{branch}"),
getattr(self, f"mhc_bias_post_{branch}"),
post_logits,
),
(
getattr(self, f"mhc_alpha_res_{branch}"),
getattr(self, f"mhc_bias_res_{branch}"),
residual_logits,
),
out_dtype=streams.dtype,
)
return self.mhc_handler.hyper_connect(streams, output, post, residual)
def forward(
self,
hidden_states: torch.Tensor,
rope: torch.Tensor,
varlen_handler: VarlenHandler,
modality_dispatcher: ModalityDispatcher,
cp_split_sizes: list[int] | torch.Tensor,
) -> torch.Tensor:
streams = hidden_states.reshape(hidden_states.shape[0], self.config.mhc.num_streams, self.config.hidden_size)
attention_logits = self._branch_logits(streams, "attn", modality_dispatcher)
attention_input = self._branch_input(streams, "attn", attention_logits)
attention_output = self.attention(
attention_input,
rope,
varlen_handler,
modality_dispatcher,
cp_split_sizes,
)
streams = self._connect(streams, attention_output, "attn", attention_logits)
mlp_logits = self._branch_logits(streams, "mlp", modality_dispatcher)
mlp_input = self._branch_input(streams, "mlp", mlp_logits)
mlp_output = self.mlp(mlp_input, modality_dispatcher)
streams = self._connect(streams, mlp_output, "mlp", mlp_logits)
return streams.reshape(streams.shape[0], -1)
class Magi2TransformerBlock(nn.Module):
def __init__(self, config: Magi2PreviewConfig) -> None:
super().__init__()
self.config = config
self.layers = nn.ModuleList(Magi2TransformerLayer(config, index) for index in range(config.num_layers))
def __iter__(self):
"""Expose checkpoint-nested layers to the layerwise offload backends."""
return iter(self.layers)
def forward(
self,
hidden_states: torch.Tensor,
rope: torch.Tensor,
varlen_handler: VarlenHandler,
modality_dispatcher: ModalityDispatcher,
cp_split_sizes: list[int] | torch.Tensor,
) -> torch.Tensor:
hidden_states = hidden_states.to(self.config.params_dtype)
hidden_states = modality_dispatcher.permute(hidden_states)
for layer in self.layers:
hidden_states = layer(
hidden_states,
rope,
varlen_handler,
modality_dispatcher,
cp_split_sizes,
)
return modality_dispatcher.inverse_permute(hidden_states)
def _is_magi2_transformer_layer(_name: str, module: nn.Module) -> bool:
"""Shard one Preview layer at a time under the shared FSDP2/HSDP path."""
return isinstance(module, Magi2TransformerLayer)
class Magi2PreviewTransformer(nn.Module):
"""Native preview DiT with the released checkpoint hierarchy."""
# ``block`` must remain the registered child name because it is part of the
# released checkpoint hierarchy. The ``layers`` property below exposes its
# ModuleList through the standard layerwise-offload contract without
# changing the registered module names.
_layerwise_offload_blocks_attrs = ["layers"]
_hsdp_shard_conditions = [_is_magi2_transformer_layer]
_hsdp_preserve_parameter_dtypes = True
_EP_SHARDED_SUFFIXES = (
".moe_mlp.gate",
".moe_mlp.W_gate",
".moe_mlp.W_up",
".moe_mlp.W_down",
".moe_mlp.router.expert_bias",
".moe_mlp.router.expert_bias_ema",
)
def __init__(self, config: Magi2PreviewConfig | None = None) -> None:
super().__init__()
self.config = config or Magi2PreviewConfig()
self.config.validate()
self.pre_adapter = Magi2PreAdapter(self.config)
self.post_adapter = Magi2PostAdapter(self.config)
self.block = Magi2TransformerBlock(self.config)
# SP doubles as MAGI's MoE-head parallel axis. These modules therefore
# contain different checkpoint slices on each SP rank and must remain
# rank-local while HSDP shards the replicated parameters around them.
self._hsdp_ignored_modules = [
f"block.layers.{index}.mlp.moe_mlp"
for index, layer in enumerate(self.block.layers)
if isinstance(layer.mlp, Magi2MultiHeadMoELayer)
]
@property
def layers(self) -> nn.ModuleList:
"""Expose Preview layers to the shared offload block discovery API."""
return self.block.layers
def forward(
self,
x: torch.Tensor,
coords_mapping: torch.Tensor,
modality_mapping: torch.Tensor,
varlen_handler: VarlenHandler,
time_token_sequence: torch.Tensor | None = None,
) -> torch.Tensor:
dispatcher = Magi2SequenceDispatcher()
x = dispatcher.dispatch(x)
coords_mapping = dispatcher.dispatch(coords_mapping)
modality_mapping = dispatcher.dispatch(modality_mapping)
if time_token_sequence is not None:
time_token_sequence = dispatcher.dispatch(time_token_sequence)
assert dispatcher.split_sizes is not None
cp_split_sizes = dispatcher.split_sizes
time_mask = modality_mapping == int(Modality.TIME)
modality_mapping = torch.where(time_mask, int(Modality.TEXT), modality_mapping)
modality_dispatcher = ModalityDispatcher(modality_mapping, 3)
video_indices = torch.nonzero(modality_mapping == int(Modality.VIDEO)).flatten()
audio_indices = torch.nonzero(modality_mapping == int(Modality.AUDIO)).flatten()
text_indices = torch.nonzero(modality_mapping == int(Modality.TEXT)).flatten()
hidden_states, rope = self.pre_adapter(
x,
coords_mapping,
video_indices,
audio_indices,
text_indices,
)
if time_token_sequence is not None and time_token_sequence.shape[-1] > 0:
hidden_states[:, : time_token_sequence.shape[-1]] = time_token_sequence.to(hidden_states.dtype)
hidden_states = self.block(
hidden_states,
rope,
varlen_handler,
modality_dispatcher,
cp_split_sizes,
)
output = self.post_adapter(hidden_states, video_indices, audio_indices)
return dispatcher.undispatch(output)
def _moe_for_weight(self, name: str) -> Magi2MultiHeadMoE | None:
if ".moe_mlp." not in name:
return None
module_name = name.split(".moe_mlp.", 1)[0] + ".moe_mlp"
module = dict(self.named_modules()).get(module_name)
return module if isinstance(module, Magi2MultiHeadMoE) else None
def validate_loaded_weights(self, loaded_names: set[str]) -> None:
"""Fail closed when the mmap loader misses a Preview tensor."""
loaded_names = {name.removeprefix("transformer.") for name in loaded_names}
loaded_names = {f"block.{name}" if name.startswith("layers.") else name for name in loaded_names}
expected = {name for name, _ in self.named_parameters()}
missing = expected - loaded_names
if missing:
raise ValueError(f"MAGI-2 Preview mmap loading is missing {len(missing)} weights: {sorted(missing)[:8]}")
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
"""Strictly load and TP/MoE-slice the released Preview checkpoint.
This is the ordinary loader used by resident and rank-local DLO
deployments. DLO+AllGather uses the pipeline's mmap mapping and the
same per-parameter transforms before orthogonal DP sharding.
"""
targets = self.state_dict(keep_vars=True)
loaded: set[str] = set()
for raw_name, checkpoint_tensor in weights:
name = raw_name.removeprefix("transformer.")
if name not in targets:
raise ValueError(f"unexpected MAGI-2 preview checkpoint weight {raw_name!r}")
target = targets[name]
checkpoint_transform = getattr(target, "checkpoint_weight_transform", None)
moe = self._moe_for_weight(name)
if callable(checkpoint_transform):
checkpoint_tensor = checkpoint_transform(checkpoint_tensor)
elif moe is not None and name.endswith(self._EP_SHARDED_SUFFIXES):
checkpoint_tensor = moe.ep_slice(checkpoint_tensor)
if tuple(target.shape) != tuple(checkpoint_tensor.shape):
raise ValueError(
f"MAGI-2 weight {name!r} has shape {tuple(checkpoint_tensor.shape)}, expected {tuple(target.shape)}"
)
with torch.no_grad():
target.copy_(checkpoint_tensor.to(device=target.device, dtype=target.dtype))
loaded.add(name)
if (os.environ.get("MAGI2_ROUTER_BIAS_SOURCE") or "ema").strip().lower() != "main":
for module_name, module in self.named_modules():
if not isinstance(module, Magi2MultiHeadMoE):
continue
module.router.expert_bias.copy_(module.router.expert_bias_ema)
loaded.add(f"{module_name}.router.expert_bias")
missing = set(targets) - loaded
if missing:
raise ValueError(f"MAGI-2 preview checkpoint is missing {len(missing)} weights: {sorted(missing)[:8]}")
return loaded
__all__ = [
"Magi2PreviewTransformer",
"Modality",
"VarlenHandler",
]
vllm_omni/diffusion/models/magi2/parallel.py
+392
View File
@@ -0,0 +1,392 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Copyright (c) 2025-2026 SandAI. All Rights Reserved.
"""Distributed primitives for the native MAGI-2 preview transformer.
This file is adapted from SandAI's Apache-2.0 MAGI-2 preview context- and
expert-parallel primitives. It has been modified to reuse vLLM-Omni's
already-initialized sequence-parallel and expert-parallel process groups.
MAGI-2 supports two equivalent layouts for its two parallel views:
* Ulysses context parallelism shards tokens and gathers attention heads.
* Multi-head expert parallelism shards the *MoE head* axis, not experts.
With SP-only deployment both views use the SP group. With TP enabled,
Ulysses keeps using SP while native column/row tensor parallelism and the MoE
heads use TP. This gives TP4 and TP2SP2 independent weight/head and token
axes while preserving the released equations and checkpoint hierarchy.
These helpers intentionally do not create or own process groups.
"""
from __future__ import annotations
from collections.abc import Iterable
from dataclasses import dataclass
import torch
import torch.distributed as dist
@dataclass(frozen=True)
class Magi2ParallelGroup:
"""The small process-group surface needed by MAGI-2 kernels."""
group: dist.ProcessGroup | None
world_size: int
rank: int
replicated_sequence: bool = False
def _dist_group_info(group: dist.ProcessGroup | None) -> Magi2ParallelGroup:
if group is None or not dist.is_available() or not dist.is_initialized():
return Magi2ParallelGroup(None, 1, 0)
return Magi2ParallelGroup(
group=group,
world_size=dist.get_world_size(group),
rank=dist.get_rank(group),
)
def get_magi2_ulysses_group() -> Magi2ParallelGroup:
"""Return vLLM-Omni's Ulysses subgroup, or a rank-local fallback.
MAGI-2 does not support Ring or AllGather-KV attention. The caller is
expected to validate that ``ulysses_degree == sequence_parallel_size``.
"""
if not dist.is_available() or not dist.is_initialized():
return Magi2ParallelGroup(None, 1, 0)
try:
from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
coordinator = get_sp_group()
except (AssertionError, RuntimeError):
return Magi2ParallelGroup(None, 1, 0)
return Magi2ParallelGroup(
group=coordinator.ulysses_group,
world_size=coordinator.ulysses_world_size,
rank=coordinator.ulysses_rank,
)
def get_magi2_ep_group() -> Magi2ParallelGroup:
"""Return the process group used for MAGI's MoE-head parallelism.
TP is the explicit MoE-head axis when it is larger than one. Otherwise
the released SP-only layout overlaps head parallelism with Ulysses. Both
groups are initialized and owned by vLLM-Omni; MAGI creates no ad-hoc
process groups.
"""
if not dist.is_available() or not dist.is_initialized():
return Magi2ParallelGroup(None, 1, 0)
try:
from vllm.distributed.parallel_state import get_tp_group
coordinator = get_tp_group()
except (AssertionError, RuntimeError):
return get_magi2_ulysses_group()
if coordinator.world_size <= 1:
return get_magi2_ulysses_group()
return Magi2ParallelGroup(
group=coordinator.device_group,
world_size=coordinator.world_size,
rank=coordinator.rank_in_group,
replicated_sequence=True,
)
def get_magi2_tp_group() -> Magi2ParallelGroup:
"""Return vLLM's tensor-parallel group without an SP fallback."""
if not dist.is_available() or not dist.is_initialized():
return Magi2ParallelGroup(None, 1, 0, replicated_sequence=True)
try:
from vllm.distributed.parallel_state import get_tp_group
coordinator = get_tp_group()
except (AssertionError, RuntimeError):
return Magi2ParallelGroup(None, 1, 0, replicated_sequence=True)
return Magi2ParallelGroup(
group=coordinator.device_group,
world_size=coordinator.world_size,
rank=coordinator.rank_in_group,
replicated_sequence=True,
)
def get_magi2_replica_group(data_parallel_size: int) -> Magi2ParallelGroup:
"""Return the complete TP x SP group for one data-parallel replica.
The pipeline uses this group for conditioning broadcasts and output-rank
ownership. With DP=1 the diffusion world is exactly one TP x SP replica.
With DP>1, MAGI currently requires TP=1, so the existing SP group is the
complete per-replica group. CFG x DP is deliberately rejected by
``_validate_native_topology``; if that combination is enabled in the
future, this helper must select a group fixed to one DP replica and include
its CFG ranks rather than returning the SP group alone.
"""
if not dist.is_available() or not dist.is_initialized():
return Magi2ParallelGroup(None, 1, 0)
if data_parallel_size > 1:
try:
from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
coordinator = get_sp_group()
except (AssertionError, RuntimeError):
return Magi2ParallelGroup(None, 1, 0)
return Magi2ParallelGroup(
group=coordinator.device_group,
world_size=coordinator.world_size,
rank=coordinator.rank_in_group,
)
return _dist_group_info(dist.group.WORLD)
def balanced_split_sizes(length: int, world_size: int) -> list[int]:
"""Split ``length`` contiguously, putting one extra token on early ranks."""
if length < 0:
raise ValueError(f"length must be non-negative, got {length}")
if world_size < 1:
raise ValueError(f"world_size must be positive, got {world_size}")
base, remainder = divmod(length, world_size)
return [base + int(rank < remainder) for rank in range(world_size)]
def shard_sequence(
tensor: torch.Tensor,
split_sizes: list[int] | None = None,
group: Magi2ParallelGroup | None = None,
) -> tuple[torch.Tensor, list[int]]:
"""Return this Ulysses rank's contiguous sequence slice.
Inputs are replicated at the pipeline boundary, matching the diffusion
runner contract, so no collective is needed for dispatch.
"""
group = group or get_magi2_ulysses_group()
split_sizes = split_sizes or balanced_split_sizes(tensor.shape[0], group.world_size)
if len(split_sizes) != group.world_size or sum(split_sizes) != tensor.shape[0]:
raise ValueError(
f"split_sizes={split_sizes} do not partition sequence length "
f"{tensor.shape[0]} across {group.world_size} ranks"
)
start = sum(split_sizes[: group.rank])
return tensor.narrow(0, start, split_sizes[group.rank]).contiguous(), split_sizes
def gather_sequence(
tensor: torch.Tensor,
split_sizes: list[int],
group: Magi2ParallelGroup | None = None,
) -> torch.Tensor:
"""Gather uneven contiguous sequence shards in rank order."""
group = group or get_magi2_ulysses_group()
if group.world_size == 1:
return tensor
if len(split_sizes) != group.world_size:
raise ValueError("split_sizes length must equal the Ulysses world size")
if tensor.shape[0] != split_sizes[group.rank]:
raise ValueError(f"rank {group.rank} owns {tensor.shape[0]} tokens, expected {split_sizes[group.rank]}")
output_shape = (sum(split_sizes), *tensor.shape[1:])
output = torch.empty(output_shape, dtype=tensor.dtype, device=tensor.device)
chunks = list(torch.split(output, split_sizes, dim=0))
dist.all_gather(chunks, tensor.contiguous(), group=group.group)
return output
def scatter_seqlen_gather_heads(
tensor: torch.Tensor,
split_sizes: list[int],
group: Magi2ParallelGroup | None = None,
) -> torch.Tensor:
"""Ulysses ``[sum(S_r), H, D] -> [S_rank, world*H, D]`` exchange."""
group = group or get_magi2_ulysses_group()
if group.world_size == 1:
return tensor
if tensor.ndim != 3 or not tensor.is_contiguous():
raise ValueError("Ulysses attention input must be contiguous [S,H,D]")
if len(split_sizes) != group.world_size or sum(split_sizes) != tensor.shape[0]:
raise ValueError("split_sizes must partition the input sequence")
local_tokens = split_sizes[group.rank]
output = torch.empty(
(group.world_size * local_tokens, tensor.shape[1], tensor.shape[2]),
dtype=tensor.dtype,
device=tensor.device,
)
dist.all_to_all_single(
output,
tensor,
output_split_sizes=[local_tokens] * group.world_size,
input_split_sizes=split_sizes,
group=group.group,
)
output = (
output.view(group.world_size, local_tokens, tensor.shape[1], tensor.shape[2])
.permute(1, 0, 2, 3)
.reshape(local_tokens, group.world_size * tensor.shape[1], tensor.shape[2])
)
return output
def scatter_heads_gather_seqlen(
tensors: Iterable[torch.Tensor],
split_sizes: list[int],
group: Magi2ParallelGroup | None = None,
) -> list[torch.Tensor]:
"""Inverse batched Ulysses exchange for Q/K/V.
Each input is ``[S_rank, world*H_i, D]`` and each output is
``[sum(S_r), H_i, D]``. Fusing Q/K/V into one all-to-all preserves the
reference communication ordering and avoids three independent collectives.
"""
group = group or get_magi2_ulysses_group()
tensors = list(tensors)
if group.world_size == 1:
return tensors
if not tensors:
return []
local_tokens = split_sizes[group.rank]
if len(split_sizes) != group.world_size:
raise ValueError("split_sizes length must equal the Ulysses world size")
if any(t.ndim != 3 or t.shape[0] != local_tokens for t in tensors):
raise ValueError("all Ulysses inputs must be [local_tokens, heads, dim]")
reshaped: list[torch.Tensor] = []
local_head_counts: list[int] = []
head_dim = tensors[0].shape[-1]
for tensor in tensors:
if tensor.shape[-1] != head_dim or tensor.shape[1] % group.world_size:
raise ValueError("attention heads must divide evenly across Ulysses ranks")
local_heads = tensor.shape[1] // group.world_size
local_head_counts.append(local_heads)
reshaped.append(
tensor.view(local_tokens, group.world_size, local_heads, head_dim)
.permute(1, 0, 2, 3)
.reshape(group.world_size * local_tokens, local_heads, head_dim)
)
fused = torch.cat(reshaped, dim=1).contiguous()
output = torch.empty(
(sum(split_sizes), fused.shape[1], head_dim),
dtype=fused.dtype,
device=fused.device,
)
dist.all_to_all_single(
output,
fused,
output_split_sizes=split_sizes,
input_split_sizes=[local_tokens] * group.world_size,
group=group.group,
)
return list(torch.split(output, local_head_counts, dim=1))
def ep_dispatch(
tensor: torch.Tensor,
group: Magi2ParallelGroup | None = None,
sequence_split_sizes: list[int] | None = None,
) -> torch.Tensor:
"""Dispatch ``[S,H,D]`` so each rank evaluates a contiguous head shard."""
group = group or get_magi2_ep_group()
if group.world_size == 1:
return tensor
if tensor.ndim != 3 or tensor.shape[1] % group.world_size:
raise ValueError(
f"MoE head count {tensor.shape[1] if tensor.ndim >= 2 else '?'} must divide by EP size {group.world_size}"
)
sequence, heads, dim = tensor.shape
local_heads = heads // group.world_size
if sequence_split_sizes is None:
local_size = torch.tensor([sequence], dtype=torch.int64, device=tensor.device)
gathered_sizes = [torch.empty_like(local_size) for _ in range(group.world_size)]
dist.all_gather(gathered_sizes, local_size, group=group.group)
sequence_split_sizes = [int(size.item()) for size in gathered_sizes]
if len(sequence_split_sizes) != group.world_size or sequence_split_sizes[group.rank] != sequence:
raise ValueError("EP sequence split sizes do not describe the local tensor")
send = tensor.contiguous().view(sequence, group.world_size, local_heads, dim).permute(1, 0, 2, 3).contiguous()
output = torch.empty(
(sum(sequence_split_sizes), local_heads, dim),
dtype=tensor.dtype,
device=tensor.device,
)
row_width = local_heads * dim
dist.all_to_all_single(
output.view(-1),
send.view(-1),
output_split_sizes=[size * row_width for size in sequence_split_sizes],
input_split_sizes=[sequence * row_width] * group.world_size,
group=group.group,
)
return output
def ep_undispatch(
tensor: torch.Tensor,
group: Magi2ParallelGroup | None = None,
sequence_split_sizes: list[int] | None = None,
) -> torch.Tensor:
"""Undo :func:`ep_dispatch` and restore the complete MoE head axis."""
group = group or get_magi2_ep_group()
if group.world_size == 1:
return tensor
if tensor.ndim != 3:
raise ValueError("EP-dispatched tensor must be [global_tokens,local_heads,dim]")
if sequence_split_sizes is None:
if tensor.shape[0] % group.world_size:
raise ValueError("uneven EP sequence requires explicit split sizes")
sequence_split_sizes = [tensor.shape[0] // group.world_size] * group.world_size
if len(sequence_split_sizes) != group.world_size or sum(sequence_split_sizes) != tensor.shape[0]:
raise ValueError("EP sequence split sizes do not partition the global tensor")
sequence = sequence_split_sizes[group.rank]
local_heads, dim = tensor.shape[1:]
send = tensor.contiguous()
output = torch.empty(
(group.world_size, sequence, local_heads, dim),
dtype=tensor.dtype,
device=tensor.device,
)
row_width = local_heads * dim
dist.all_to_all_single(
output.view(-1),
send.view(-1),
output_split_sizes=[sequence * row_width] * group.world_size,
input_split_sizes=[size * row_width for size in sequence_split_sizes],
group=group.group,
)
return output.permute(1, 0, 2, 3).contiguous().view(sequence, group.world_size * local_heads, dim)
class Magi2SequenceDispatcher:
"""Request-scoped dispatcher that enforces one consistent token split."""
def __init__(self, group: Magi2ParallelGroup | None = None) -> None:
self.group = group or get_magi2_ulysses_group()
self.split_sizes: list[int] | None = None
def dispatch(self, tensor: torch.Tensor) -> torch.Tensor:
split_sizes = balanced_split_sizes(tensor.shape[0], self.group.world_size)
if self.split_sizes is not None and split_sizes != self.split_sizes:
raise ValueError(
f"all packed inputs must share a sequence length; got {split_sizes}, expected {self.split_sizes}"
)
self.split_sizes = split_sizes
return shard_sequence(tensor, split_sizes, self.group)[0]
def undispatch(self, tensor: torch.Tensor) -> torch.Tensor:
if self.split_sizes is None:
raise RuntimeError("dispatch must be called before undispatch")
output = gather_sequence(tensor, self.split_sizes, self.group)
self.split_sizes = None
return output
vllm_omni/diffusion/models/magi2/pipeline_magi2.py
+1124
View File
File diff suppressed because it is too large. Load diff
vllm_omni/diffusion/models/magi2/preview_data_proxy.py
+691
View File
@@ -0,0 +1,691 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Copyright (c) 2026 SandAI. All Rights Reserved.
"""Native packing helpers for the MAGI-2 Preview transformer.
The released model consumes one varlen sequence per sample. Tokens are laid
out in VIDEO -> AUDIO -> TEXT order, followed by zero or more reference-image
special-token/image-token pairs. This module keeps that layout and its 9-D
coordinate metadata independent of any particular distributed topology.
The packing math is adapted from SandAI's Apache-2.0 MAGI-2 Preview inference
implementation. It intentionally has no dependency on that implementation at
runtime.
"""
from __future__ import annotations
import math
from collections.abc import Sequence
from dataclasses import dataclass
from enum import IntEnum
from itertools import chain
from typing import Literal
import torch
from einops import rearrange
from torch.nn import functional as F
from .attention import VarlenHandler
class Modality(IntEnum):
"""Checkpoint modality IDs used by the pre/post adapters."""
VIDEO = 0
AUDIO = 1
TEXT = 2
# TIME is remapped to TEXT by the transformer after selecting time tokens.
TIME = 3
@dataclass(frozen=True)
class Magi2PreviewDataProxyConfig:
"""Packing configuration for the released preview checkpoint."""
t_patch_size: int = 1
patch_size: int = 1
spatial_rope_interpolation: Literal["inter", "extra"] = "extra"
add_time_token: bool = False
time_channel_dim: int = 64
time_aligned_rope: bool = False
audio_latent_fps: float = 25.0
time_pos_fps: float = 3.125
vae_first_latent_is_image: bool = True
video_fps: float = 25.0
def __post_init__(self) -> None:
if self.t_patch_size <= 0 or self.patch_size <= 0:
raise ValueError("MAGI-2 patch sizes must be positive")
if self.time_channel_dim < 0:
raise ValueError("time_channel_dim must be non-negative")
@dataclass
class ModelInput:
"""Unpacked tensors for one native preview-transformer invocation."""
x_t: torch.Tensor
audio_x_t: torch.Tensor
audio_feat_len: torch.Tensor | Sequence[int]
txt_feat: torch.Tensor
txt_feat_len: torch.Tensor | Sequence[int]
t: torch.Tensor
ref_audio_feat: torch.Tensor | None = None
ref_audio_feat_len: torch.Tensor | Sequence[int] | None = None
ref_video_feat: torch.Tensor | None = None
ref_video_feat_len: torch.Tensor | Sequence[int] | None = None
per_token_video_t: torch.Tensor | None = None
per_token_audio_t: torch.Tensor | None = None
# Images are [B, M, C, 1, H, W], lengths are [B, M, 2] holding
# the latent patch-grid (H, W), and special tokens are [B, M, D].
ref_image_feat: torch.Tensor | None = None
ref_image_feat_len: torch.Tensor | None = None
ref_image_special_token_embedding: torch.Tensor | None = None
def _to_int(value: int | torch.Tensor) -> int:
if isinstance(value, torch.Tensor):
return int(value.detach().reshape(-1)[0].item())
return int(value)
def _length_at(value: torch.Tensor | Sequence[int], index: int) -> int:
if isinstance(value, torch.Tensor) and value.ndim == 0:
if index:
raise IndexError("a scalar length only describes one sample")
return _to_int(value)
return _to_int(value[index])
def _pad_cat(tensors: list[torch.Tensor], device: torch.device, dtype: torch.dtype) -> torch.Tensor:
if not tensors:
return torch.empty(0, 0, device=device, dtype=dtype)
max_channel = max(tensor.shape[-1] for tensor in tensors)
return torch.cat(
[F.pad(tensor.to(device=device, dtype=dtype), (0, max_channel - tensor.shape[-1])) for tensor in tensors],
dim=0,
)
def _segment_paint(
values: list[int],
seqlens: list[int],
*,
dtype: torch.dtype,
device: torch.device,
output_size: int,
) -> torch.Tensor:
mapping = torch.empty(output_size, dtype=dtype, device=device)
offset = 0
for value, seqlen in zip(values, seqlens, strict=True):
mapping[offset : offset + seqlen] = int(value)
offset += seqlen
if offset != output_size:
raise ValueError(f"segment lengths total {offset}, expected {output_size}")
return mapping
def _seqlens_to_cu_seqlens(seqlens: list[int], device: torch.device | None = None) -> torch.Tensor:
lengths = torch.tensor(seqlens, dtype=torch.int32, device=device)
return F.pad(torch.cumsum(lengths, dim=0), (1, 0))
def _len_to_list(value: torch.Tensor) -> list[int]:
return [int(item) for item in value.detach().to(torch.long).reshape(-1).tolist()]
def get_coords(
shape: tuple[int, int, int],
ref_feat_shape: tuple[int, int, int],
offset_thw: tuple[int, int, int] = (0, 0, 0),
*,
device: torch.device = torch.device("cpu"),
dtype: torch.dtype = torch.float32,
time_positions: torch.Tensor | None = None,
) -> torch.Tensor:
"""Build MAGI-2's ``(t,h,w,T,H,W,ref_T,ref_H,ref_W)`` rows."""
ori_t, ori_h, ori_w = shape
ref_t, ref_h, ref_w = ref_feat_shape
offset_t, offset_h, offset_w = offset_thw
if time_positions is None:
time_range = torch.arange(ori_t, device=device, dtype=dtype) + offset_t
else:
if time_positions.numel() != ori_t:
raise ValueError(f"received {time_positions.numel()} time positions for {ori_t} time steps")
time_range = time_positions.to(device=device, dtype=dtype) + offset_t
height_range = torch.arange(ori_h, device=device, dtype=dtype) + offset_h
width_range = torch.arange(ori_w, device=device, dtype=dtype) + offset_w
time_grid, height_grid, width_grid = torch.meshgrid(
time_range,
height_range,
width_range,
indexing="ij",
)
coords = torch.stack((time_grid, height_grid, width_grid), dim=-1).reshape(-1, 3)
metadata = torch.tensor(
(ori_t, ori_h, ori_w, ref_t, ref_h, ref_w),
device=device,
dtype=dtype,
)
return torch.cat((coords, metadata.expand(coords.shape[0], -1)), dim=-1)
def sinusoidal_embedding_1d(dim: int, position: torch.Tensor) -> torch.Tensor:
"""The preview checkpoint's FP32 diffusion-time channel embedding."""
position = position.to(torch.float32) * 1000.0
half = dim // 2
frequencies = torch.exp(
-math.log(10000) * torch.arange(0, half, dtype=torch.float32, device=position.device) / half
)
arguments = position[:, None].float() * frequencies[None]
embedding = torch.cat((torch.cos(arguments), torch.sin(arguments)), dim=-1)
if dim % 2:
embedding = torch.cat((embedding, torch.zeros_like(embedding[:, :1])), dim=-1)
return embedding
def _patchify_3d(x: torch.Tensor, t_patch_size: int, patch_size: int) -> torch.Tensor:
"""Non-overlapping 3-D unfold with the released channel-major order."""
kernel = (t_patch_size, patch_size, patch_size)
if x.ndim != 5:
raise ValueError(f"expected [B,C,T,H,W], got {tuple(x.shape)}")
if not all(size >= kernel_size for size, kernel_size in zip(x.shape[2:], kernel, strict=True)):
return torch.empty(
x.shape[0],
0,
x.shape[1] * math.prod(kernel),
device=x.device,
dtype=x.dtype,
)
patches = x.unfold(2, t_patch_size, t_patch_size)
patches = patches.unfold(3, patch_size, patch_size)
patches = patches.unfold(4, patch_size, patch_size)
# [B,C,T,H,W,pT,pH,pW] -> [B,T*H*W,C*pT*pH*pW]
patches = patches.permute(0, 2, 3, 4, 1, 5, 6, 7).contiguous()
return patches.reshape(x.shape[0], -1, x.shape[1] * math.prod(kernel))
@dataclass
class SingleData:
"""Packed metadata for one sample."""
video_x_t: torch.Tensor
audio_x_t: torch.Tensor
audio_feat_len: int
txt_feat: torch.Tensor
txt_feat_len: int
t: int
h: int
w: int
patch_size: int
t_patch_size: int
spatial_rope_interpolation: Literal["inter", "extra"]
diffusion_t: torch.Tensor | None = None
per_token_video_t: torch.Tensor | None = None
per_token_audio_t: torch.Tensor | None = None
time_channel_dim: int = 0
vae_first_latent_is_image: bool = True
video_fps: float = 25.0
time_pos_fps: float = 3.125
ref_image_feats: list[torch.Tensor] | None = None
ref_image_feat_lens: list[list[int]] | None = None
ref_image_special_tokens: list[torch.Tensor] | None = None
def __post_init__(self) -> None:
self.video_token_num = self.video_x_t.shape[0]
self.origin_audio_feat_len = self.audio_x_t.shape[0]
self.audio_x_t = self.audio_x_t[: self.audio_feat_len]
self.txt_feat = self.txt_feat[: self.txt_feat_len]
if self.per_token_audio_t is not None:
self.per_token_audio_t = self.per_token_audio_t[: self.audio_feat_len]
self.ref_image_feats = self.ref_image_feats or []
self.ref_image_feat_lens = self.ref_image_feat_lens or []
self.ref_image_special_tokens = self.ref_image_special_tokens or []
if not (len(self.ref_image_feats) == len(self.ref_image_feat_lens) == len(self.ref_image_special_tokens)):
raise ValueError("reference-image features, lengths, and special tokens must have equal counts")
self.ref_image_token_nums = [math.prod(feat_len) for feat_len in self.ref_image_feat_lens]
self.ref_image_feats = [
feature[:token_num]
for feature, token_num in zip(self.ref_image_feats, self.ref_image_token_nums, strict=True)
]
self.num_ref_images = len(self.ref_image_feats)
self.total_ref_image_feat_len = sum(self.ref_image_token_nums)
self.video_channel = self.video_x_t.shape[-1]
self.audio_channel = self.audio_x_t.shape[-1]
@property
def device(self) -> torch.device:
return self.video_x_t.device
@property
def default_dtype(self) -> torch.dtype:
return self.video_x_t.dtype
@property
def add_time_token(self) -> bool:
return self.diffusion_t is not None
@property
def total_token_num(self) -> int:
total = self.video_token_num + self.audio_feat_len + self.txt_feat_len
total += self.total_ref_image_feat_len + self.num_ref_images
return total + int(self.add_time_token)
@property
def feat_to_cat(self) -> list[torch.Tensor]:
tensors = [self.video_x_t, self.audio_x_t, self.txt_feat]
for index in range(self.num_ref_images):
tensors.append(self.ref_image_special_tokens[index])
tensors.append(self.ref_image_feats[index])
if self.add_time_token:
tensors.append(self.diffusion_t.to(device=self.device, dtype=self.default_dtype).reshape(1, 1))
return tensors
@property
def token_sequence(self) -> torch.Tensor:
return _pad_cat(self.feat_to_cat, self.device, self.default_dtype)
@property
def modality_map_seqlens(self) -> tuple[list[int], list[int]]:
seqlens = [self.video_token_num, self.audio_feat_len, self.txt_feat_len]
modalities = [Modality.VIDEO, Modality.AUDIO, Modality.TEXT]
for index in range(self.num_ref_images):
seqlens.extend((1, self.ref_image_token_nums[index]))
modalities.extend((Modality.TEXT, Modality.VIDEO))
if self.add_time_token:
seqlens.append(1)
modalities.append(Modality.TIME)
return seqlens, modalities
@property
def modality_mapping(self) -> torch.Tensor:
seqlens, modalities = self.modality_map_seqlens
return _segment_paint(
modalities,
seqlens,
dtype=torch.int32,
device=self.device,
output_size=self.total_token_num,
)
def _default_coords(
self,
shape: tuple[int, int, int],
ref_feat_shape: tuple[int, int, int],
offset_thw: tuple[int, int, int] = (0, 0, 0),
time_positions: torch.Tensor | None = None,
) -> torch.Tensor:
return get_coords(
shape,
ref_feat_shape,
offset_thw,
device=self.device,
dtype=self.default_dtype,
time_positions=time_positions,
)
@property
def coords_to_cat(self) -> list[torch.Tensor]:
t_steps = self.t // self.t_patch_size
h_steps = self.h // self.patch_size
w_steps = self.w // self.patch_size
if self.spatial_rope_interpolation == "inter":
video_h_ref, video_w_ref = 32, 32
else:
video_h_ref, video_w_ref = h_steps, w_steps
video_coords = self._default_coords(
(t_steps, h_steps, w_steps),
(t_steps, video_h_ref, video_w_ref),
)
magic_audio_ref_t = (self.audio_feat_len - 1) // 8 + 1
audio_coords = self._default_coords(
(self.audio_feat_len, 1, 1),
(magic_audio_ref_t // self.t_patch_size, 1, 1),
)
coords = [
video_coords,
audio_coords,
self._default_coords(
(self.txt_feat_len, 1, 1),
(1, 1, 1),
offset_thw=(-self.txt_feat_len, 0, 0),
),
]
for index in range(self.num_ref_images):
token_len = self.ref_image_token_nums[index]
if len(self.ref_image_feat_lens[index]) >= 2:
image_h, image_w = self.ref_image_feat_lens[index][:2]
else:
image_h = image_w = math.ceil(math.sqrt(token_len))
# Keep image context after the generated clip with a one-step gap.
time_offset = t_steps + 2 + index
coords.append(
torch.tensor(
[[time_offset, -1, -1, 1, image_h, image_w, 1, image_h, image_w]],
device=self.device,
dtype=self.default_dtype,
)
)
coords.append(
self._default_coords(
(1, image_h, image_w),
(1, image_h, image_w),
offset_thw=(time_offset, 0, 0),
)[:token_len]
)
if self.add_time_token:
coords.append(self._default_coords((1, 1, 1), (1, 1, 1))[:1])
return coords
@property
def coords_mapping(self) -> torch.Tensor:
return torch.cat(self.coords_to_cat, dim=0)
@property
def time_token_sequence(self) -> torch.Tensor:
if self.time_channel_dim == 0:
return torch.empty(self.total_token_num, 0, device=self.device)
if self.per_token_video_t is None or self.per_token_audio_t is None:
raise ValueError("per-token video and audio timesteps are required when time_channel_dim is non-zero")
parts = [
self.per_token_video_t.squeeze(-1),
self.per_token_audio_t.squeeze(-1),
torch.zeros(self.txt_feat_len, device=self.device),
]
for index in range(self.num_ref_images):
parts.append(torch.zeros(1, device=self.device))
parts.append(torch.zeros(self.ref_image_token_nums[index], device=self.device))
if self.add_time_token:
parts.append(self.diffusion_t.reshape(1).to(self.device))
raw_t = torch.cat(parts, dim=0)
if self.time_channel_dim == 1:
return raw_t.unsqueeze(-1)
return sinusoidal_embedding_1d(self.time_channel_dim, raw_t)
@dataclass
class SimplePackedData:
"""A concatenation of independently addressable packed samples."""
items: list[SingleData]
def __post_init__(self) -> None:
if not self.items:
raise ValueError("SimplePackedData must contain at least one item")
self._total_token_num_list = [item.total_token_num for item in self.items]
self._total_token_num_sum = sum(self._total_token_num_list)
self._total_token_num_max = max(self._total_token_num_list)
self._total_token_cu_seqlens = _seqlens_to_cu_seqlens(self._total_token_num_list)
@property
def device(self) -> torch.device:
return self.items[0].device
@property
def default_dtype(self) -> torch.dtype:
return self.items[0].default_dtype
@property
def token_sequence(self) -> torch.Tensor:
tensors = list(chain.from_iterable(item.feat_to_cat for item in self.items))
return _pad_cat(tensors, self.device, self.default_dtype)
@property
def modality_mapping(self) -> torch.Tensor:
seqlens: list[int] = []
modalities: list[int] = []
for item in self.items:
item_seqlens, item_modalities = item.modality_map_seqlens
seqlens.extend(item_seqlens)
modalities.extend(item_modalities)
return _segment_paint(
modalities,
seqlens,
dtype=torch.int32,
device=self.device,
output_size=self.total_token_num,
)
@property
def coords_mapping(self) -> torch.Tensor:
coords = list(chain.from_iterable(item.coords_to_cat for item in self.items))
return torch.cat(coords, dim=0)
@property
def time_token_sequence(self) -> torch.Tensor:
return torch.cat([item.time_token_sequence for item in self.items], dim=0)
@property
def total_token_num(self) -> int:
return self._total_token_num_sum
@property
def cu_seqlen(self) -> torch.Tensor:
return self._total_token_cu_seqlens.clone()
@property
def max_seqlen(self) -> int:
return self._total_token_num_max
def __getitem__(self, index: int) -> SingleData:
return self.items[index]
def depack_token_sequence(self, token_sequence: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
if token_sequence.shape[0] != self.total_token_num:
raise ValueError(f"model returned {token_sequence.shape[0]} tokens, expected {self.total_token_num}")
videos: list[torch.Tensor] = []
audios: list[torch.Tensor] = []
token_slices = torch.split(token_sequence, self._total_token_num_list, dim=0)
for item, token_slice in zip(self.items, token_slices, strict=True):
video_flat = token_slice[: item.video_token_num, : item.video_channel]
output_channels = item.video_channel // (item.t_patch_size * item.patch_size * item.patch_size)
video = rearrange(
video_flat,
"(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)",
T=item.t // item.t_patch_size,
H=item.h // item.patch_size,
W=item.w // item.patch_size,
pT=item.t_patch_size,
pH=item.patch_size,
pW=item.patch_size,
C=output_channels,
).contiguous()
audio = torch.zeros(
item.origin_audio_feat_len,
item.audio_channel,
device=token_sequence.device,
dtype=token_sequence.dtype,
)
audio[: item.audio_feat_len] = token_slice[
item.video_token_num : item.video_token_num + item.audio_feat_len,
: item.audio_channel,
]
videos.append(video)
audios.append(audio)
return torch.stack(videos, dim=0), torch.stack(audios, dim=0)
@dataclass(frozen=True)
class PackedModelInput:
"""Request-owned packed transformer arguments and output layout."""
token_sequence: torch.Tensor
coords_mapping: torch.Tensor
modality_mapping: torch.Tensor
varlen_handler: VarlenHandler
time_token_sequence: torch.Tensor
output_layout: SimplePackedData
@property
def model_args(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, VarlenHandler, torch.Tensor]:
return (
self.token_sequence,
self.coords_mapping,
self.modality_mapping,
self.varlen_handler,
self.time_token_sequence,
)
class Magi2DataProxy:
"""Convert dense modality tensors to/from the transformer's packed ABI.
The proxy is request-stateless: all packing state needed by
``process_output`` is returned as ``PackedModelInput.output_layout`` rather
than stored on this shared instance, so concurrent requests cannot
overwrite one another.
"""
def __init__(self, config: Magi2PreviewDataProxyConfig | None = None) -> None:
self.config = config or Magi2PreviewDataProxyConfig()
self.patch_size = self.config.patch_size
self.t_patch_size = self.config.t_patch_size
def img2tokens(self, x_t: torch.Tensor) -> torch.Tensor:
return _patchify_3d(x_t, self.t_patch_size, self.patch_size)
@staticmethod
def _validate_batch(data: ModelInput) -> None:
if data.x_t.ndim != 5:
raise ValueError(f"x_t must be [B,C,T,H,W], got {tuple(data.x_t.shape)}")
if data.audio_x_t.ndim != 3 or data.txt_feat.ndim != 3:
raise ValueError("audio_x_t and txt_feat must be [B,L,C]")
batch_size = data.x_t.shape[0]
if data.audio_x_t.shape[0] != batch_size or data.txt_feat.shape[0] != batch_size:
raise ValueError("video, audio, and text batch sizes must match")
ref_values = (
data.ref_image_feat,
data.ref_image_feat_len,
data.ref_image_special_token_embedding,
)
if any(value is not None for value in ref_values) and not all(value is not None for value in ref_values):
raise ValueError("reference-image features, lengths, and special tokens must be provided together")
if data.ref_image_feat is not None:
if data.ref_image_feat.ndim != 6:
raise ValueError("ref_image_feat must be [B,M,C,T,H,W]")
if data.ref_image_feat_len.ndim != 3 or data.ref_image_feat_len.shape[-1] < 2:
raise ValueError("ref_image_feat_len must be [B,M,2]")
if data.ref_image_special_token_embedding.ndim != 3:
raise ValueError("ref_image_special_token_embedding must be [B,M,D]")
expected = data.ref_image_feat.shape[:2]
if data.ref_image_feat_len.shape[:2] != expected:
raise ValueError("reference-image feature and length batch/image counts differ")
if data.ref_image_special_token_embedding.shape[:2] != expected:
raise ValueError("reference-image feature and special-token batch/image counts differ")
def process_input(
self,
data: ModelInput,
) -> PackedModelInput:
self._validate_batch(data)
batch_size, _, video_t, video_h, video_w = data.x_t.shape
video_tokens = self.img2tokens(data.x_t)
audio_tokens = data.audio_x_t.contiguous()
text_tokens = data.txt_feat.contiguous()
per_token_video_tokens = None
per_token_audio = None
if self.config.time_channel_dim > 0 and data.per_token_video_t is not None:
per_token_video_tokens = self.img2tokens(data.per_token_video_t)[:, :, :1]
per_token_audio = data.per_token_audio_t
has_ref_image = data.ref_image_feat is not None
num_ref_images = data.ref_image_feat.shape[1] if has_ref_image else 0
ref_device, ref_dtype = text_tokens.device, text_tokens.dtype
items: list[SingleData] = []
for batch_index in range(batch_size):
ref_image_feats: list[torch.Tensor] = []
ref_image_feat_lens: list[list[int]] = []
ref_image_special_tokens: list[torch.Tensor] = []
for image_index in range(num_ref_images):
feature = self.img2tokens(data.ref_image_feat[batch_index, image_index].unsqueeze(0)).squeeze(0)
ref_image_feats.append(feature)
ref_image_feat_lens.append(_len_to_list(data.ref_image_feat_len[batch_index, image_index]))
ref_image_special_tokens.append(
data.ref_image_special_token_embedding[batch_index, image_index]
.to(device=ref_device, dtype=ref_dtype)
.unsqueeze(0)
)
if isinstance(data.t, torch.Tensor) and data.t.ndim == 0:
diffusion_t = data.t
else:
diffusion_t = data.t[batch_index]
items.append(
SingleData(
video_x_t=video_tokens[batch_index],
audio_x_t=audio_tokens[batch_index],
audio_feat_len=_length_at(data.audio_feat_len, batch_index),
txt_feat=text_tokens[batch_index],
txt_feat_len=_length_at(data.txt_feat_len, batch_index),
t=video_t,
h=video_h,
w=video_w,
patch_size=self.patch_size,
t_patch_size=self.t_patch_size,
spatial_rope_interpolation=self.config.spatial_rope_interpolation,
diffusion_t=diffusion_t if self.config.add_time_token else None,
per_token_video_t=(
per_token_video_tokens[batch_index] if per_token_video_tokens is not None else None
),
per_token_audio_t=(per_token_audio[batch_index] if per_token_audio is not None else None),
time_channel_dim=self.config.time_channel_dim,
time_pos_fps=self.config.time_pos_fps,
vae_first_latent_is_image=self.config.vae_first_latent_is_image,
video_fps=self.config.video_fps,
ref_image_feats=ref_image_feats,
ref_image_feat_lens=ref_image_feat_lens,
ref_image_special_tokens=ref_image_special_tokens,
)
)
packed = SimplePackedData(items)
cu_seqlens = packed.cu_seqlen.to(device=data.x_t.device, dtype=torch.int32)
varlen_handler = VarlenHandler(
cu_seqlens_q=cu_seqlens,
cu_seqlens_k=cu_seqlens,
max_seqlen_q=packed.max_seqlen,
max_seqlen_k=packed.max_seqlen,
)
return PackedModelInput(
token_sequence=packed.token_sequence,
coords_mapping=packed.coords_mapping,
modality_mapping=packed.modality_mapping,
varlen_handler=varlen_handler,
time_token_sequence=packed.time_token_sequence,
output_layout=packed,
)
@staticmethod
def process_output(
x: torch.Tensor,
output_layout: SimplePackedData,
) -> tuple[torch.Tensor, torch.Tensor]:
return output_layout.depack_token_sequence(x)
__all__ = [
"Magi2DataProxy",
"Magi2PreviewDataProxyConfig",
"Modality",
"ModelInput",
"PackedModelInput",
"SimplePackedData",
"SingleData",
"VarlenHandler",
"get_coords",
"sinusoidal_embedding_1d",
]
vllm_omni/diffusion/models/magi2/sampler_magi2.py
+682
View File
@@ -0,0 +1,682 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Copyright (c) 2026 SandAI. All Rights Reserved.
"""Native MAGI-2 Preview classifier-free-guidance sampling.
The CFG and denoising math is adapted from SandAI's Apache-2.0 MAGI-2
Preview inference implementation. Model placement is deliberately not
managed here: vLLM-Omni's loader/offloader (including DLO) owns placement,
while this sampler only prepares tensors and invokes the already placed model.
"""
from __future__ import annotations
import math
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Any
import torch
import torch.nn.functional as F
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
from .preview_data_proxy import Magi2DataProxy, ModelInput
def _pad_or_trim(
tensor: torch.Tensor,
target_size: int,
dim: int,
pad_value: float = 0.0,
) -> tuple[torch.Tensor, int]:
current_size = tensor.size(dim)
if current_size < target_size:
padding_amount = target_size - current_size
padding = [0] * (2 * tensor.dim())
padding_dim_index = tensor.dim() - 1 - dim
padding[2 * padding_dim_index + 1] = padding_amount
return F.pad(tensor, tuple(padding), "constant", pad_value), current_size
slices = [slice(None)] * tensor.dim()
slices[dim] = slice(0, target_size)
return tensor[tuple(slices)], current_size
@dataclass(frozen=True)
class CFGConfig:
"""Preview CFG controls, defaulted to the released MAGI-2 recipe."""
use_cfg_trick: bool = False
cfg_trick_start_frame: int = 13
cfg_trick_value: float = 2.0
use_dynamic_cfg: bool = False
dynamic_cfg_start_t: int = 500
dynamic_cfg_cutoff_value: float = 2.0
video_txt_guidance_scale: float = 5.0
audio_txt_guidance_scale: float = 7.0
# The released I2V config keeps image context in both CFG branches, so CFG
# changes text conditioning without pushing away from the input image.
use_ref_for_uncond: bool = True
use_skimmed_cfg_linear: bool = False
skimmed_cfg_scale: float = 3.0
cfg_rescale: float = 0.0
def __post_init__(self) -> None:
finite_values = {
"cfg_trick_value": self.cfg_trick_value,
"dynamic_cfg_cutoff_value": self.dynamic_cfg_cutoff_value,
"video_txt_guidance_scale": self.video_txt_guidance_scale,
"audio_txt_guidance_scale": self.audio_txt_guidance_scale,
"skimmed_cfg_scale": self.skimmed_cfg_scale,
"cfg_rescale": self.cfg_rescale,
}
for name, value in finite_values.items():
if not math.isfinite(float(value)):
raise ValueError(f"{name} must be finite, got {value}")
if self.cfg_trick_start_frame < 0:
raise ValueError("cfg_trick_start_frame must be non-negative")
if self.skimmed_cfg_scale < 0 or self.cfg_rescale < 0:
raise ValueError("skimmed_cfg_scale and cfg_rescale must be non-negative")
@dataclass
class SamplerInput:
video_t_list: Sequence[torch.Tensor]
audio_t_list: Sequence[torch.Tensor]
latent: torch.Tensor # [B,C,T,H,W]
audio_latent: torch.Tensor # [B,L,C]
txt_feat: torch.Tensor # [B,L,D]
null_txt_feat: torch.Tensor # [B,L,D]
ref_audio_feat: torch.Tensor | None
ref_video_feat: torch.Tensor | None
video_scheduler: SchedulerMixin
audio_scheduler: SchedulerMixin
cfg_config: CFGConfig
ref_image_feat: torch.Tensor | None = None # [B,M,C,1,H,W]
ref_image_feat_len: torch.Tensor | None = None # [B,M,2]
ref_image_special_token_embedding: torch.Tensor | None = None # [B,M,D]
def build_magi2_preview_schedulers(
num_inference_steps: int,
*,
device: torch.device | str,
shift: float = 7.0,
) -> tuple[FlowUniPCMultistepScheduler, FlowUniPCMultistepScheduler]:
"""Create independent, identically configured video/audio schedulers."""
if num_inference_steps <= 0:
raise ValueError("num_inference_steps must be positive")
if not math.isfinite(shift) or shift <= 0:
raise ValueError(f"shift must be positive and finite, got {shift}")
video_scheduler = FlowUniPCMultistepScheduler()
audio_scheduler = FlowUniPCMultistepScheduler()
video_scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
audio_scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
return video_scheduler, audio_scheduler
class Magi2PreviewSampler(CFGParallelMixin):
"""Run MAGI-2 Preview's joint video/audio denoising loop."""
def __init__(
self,
model: torch.nn.Module,
data_proxy: Magi2DataProxy | None = None,
device: torch.device | str | None = None,
dtype: torch.dtype | None = None,
) -> None:
self.model = model
self.data_proxy = data_proxy or Magi2DataProxy()
self.device = torch.device(device) if device is not None else None
# Retained for the reference-compatible constructor. Model placement
# and parameter dtype are owned by the generic loader/offloader.
self.dtype = dtype
@torch.inference_mode()
def forward(self, model_input: ModelInput) -> tuple[torch.Tensor, torch.Tensor]:
packed_input = self.data_proxy.process_input(model_input)
model_output = self.model(*packed_input.model_args)
if not isinstance(model_output, torch.Tensor):
raise TypeError(f"MAGI-2 preview transformer must return a tensor, got {type(model_output)!r}")
return self.data_proxy.process_output(model_output, packed_input.output_layout)
def predict_noise(self, model_input: ModelInput) -> tuple[torch.Tensor, torch.Tensor]:
"""Run one conditional or unconditional branch for shared CFG dispatch."""
return self.forward(model_input)
@torch.inference_mode()
def sample(self, sampler_input: SamplerInput) -> tuple[torch.Tensor, torch.Tensor]:
video_timesteps = list(sampler_input.video_t_list)
audio_timesteps = list(sampler_input.audio_t_list)
if not video_timesteps:
raise ValueError("MAGI-2 sampling needs at least one timestep")
if len(video_timesteps) != len(audio_timesteps):
raise ValueError("video and audio timestep schedules must have equal lengths")
video_cfgs, audio_cfgs = self.precalculate_cfg(
video_timesteps,
sampler_input.latent.shape[2],
sampler_input.cfg_config,
device=sampler_input.latent.device,
)
latent = sampler_input.latent.clone()
audio_latent = sampler_input.audio_latent.clone()
for timestep, video_cfg, audio_cfg in zip(
video_timesteps,
video_cfgs,
audio_cfgs,
strict=True,
):
model_input = self.prepare_model_input(
latent=latent,
audio_latent=audio_latent,
txt_feat=sampler_input.txt_feat,
null_txt_feat=sampler_input.null_txt_feat,
ref_audio_feat=sampler_input.ref_audio_feat,
ref_video_feat=sampler_input.ref_video_feat,
ref_image_feat=sampler_input.ref_image_feat,
ref_image_feat_len=sampler_input.ref_image_feat_len,
ref_image_special_token_embedding=(sampler_input.ref_image_special_token_embedding),
t=timestep,
cfg_config=sampler_input.cfg_config,
)
if get_classifier_free_guidance_world_size() > 1:
positive_input, negative_input = self._split_cfg_model_input(model_input)
guided = self.predict_noise_maybe_with_cfg(
do_true_cfg=True,
true_cfg_scale=1.0,
positive_kwargs={"model_input": positive_input},
negative_kwargs={"model_input": negative_input},
cfg_normalize=False,
kwargs={
"video_txt_guidance_scale": video_cfg,
"audio_txt_guidance_scale": audio_cfg,
"cfg_config": sampler_input.cfg_config,
"latent": latent,
"audio_latent": audio_latent,
},
)
if not isinstance(guided, tuple) or len(guided) != 2:
raise RuntimeError("MAGI-2 CFG parallel combine must return video and audio predictions")
latent, audio_latent = self._step_guided(
guided,
latent,
audio_latent,
sampler_input.video_scheduler,
sampler_input.audio_scheduler,
timestep,
)
else:
model_pred = self.forward(model_input)
latent, audio_latent, _, _ = self.step(
model_pred,
latent,
audio_latent,
video_cfg,
audio_cfg,
sampler_input.video_scheduler,
sampler_input.audio_scheduler,
timestep,
cfg_config=sampler_input.cfg_config,
)
return latent, audio_latent
@staticmethod
def _validate_prepare_inputs(
latent: torch.Tensor,
audio_latent: torch.Tensor,
txt_feat: torch.Tensor,
null_txt_feat: torch.Tensor,
) -> None:
if latent.ndim != 5:
raise ValueError("latent must be [B,C,T,H,W]")
if audio_latent.ndim != 3 or txt_feat.ndim != 3 or null_txt_feat.ndim != 3:
raise ValueError("audio and text features must be [B,L,C]")
batch_size = latent.shape[0]
if not (audio_latent.shape[0] == txt_feat.shape[0] == null_txt_feat.shape[0] == batch_size):
raise ValueError("video, audio, positive-text, and negative-text batch sizes must match")
def prepare_model_input(
self,
latent: torch.Tensor,
audio_latent: torch.Tensor,
txt_feat: torch.Tensor,
null_txt_feat: torch.Tensor,
ref_audio_feat: torch.Tensor | None = None,
ref_video_feat: torch.Tensor | None = None,
ref_image_feat: torch.Tensor | None = None,
ref_image_feat_len: torch.Tensor | None = None,
ref_image_special_token_embedding: torch.Tensor | None = None,
t: torch.Tensor | float | None = None,
cfg_config: CFGConfig | None = None,
) -> ModelInput:
self._validate_prepare_inputs(latent, audio_latent, txt_feat, null_txt_feat)
batch_size = latent.shape[0]
audio_latent_len = audio_latent.shape[1]
txt_feat_len = txt_feat.shape[1]
null_txt_feat_len = null_txt_feat.shape[1]
target_length = max(txt_feat_len, null_txt_feat_len)
txt_feat, _ = _pad_or_trim(txt_feat, target_length, 1)
null_txt_feat, _ = _pad_or_trim(null_txt_feat, target_length, 1)
if ref_audio_feat is None:
ref_audio_feat = audio_latent.new_empty(batch_size, 0, audio_latent.shape[-1])
if ref_audio_feat.ndim != 3 or ref_audio_feat.shape[0] != batch_size:
raise ValueError("ref_audio_feat must be [B,L,C] with the same batch size")
ref_audio_feat_len = ref_audio_feat.shape[1]
if ref_video_feat is None:
# Preserve the 5-D ABI without allocating or reading a full-size
# placeholder that the zero reference length excludes downstream.
ref_video_feat = latent[:, :, :0]
ref_video_feat_len = 0
else:
if ref_video_feat.ndim != 5 or ref_video_feat.shape[0] != batch_size:
raise ValueError("ref_video_feat must be [B,C,T,H,W] with the same batch size")
ref_video_feat_len = (ref_video_feat.shape[3] * ref_video_feat.shape[4]) // 4
ref_image_feat_cfg, ref_image_feat_len_cfg, ref_image_special_tokens_cfg = self._prepare_ref_image_cfg(
ref_image_feat,
ref_image_feat_len,
ref_image_special_token_embedding,
cfg_config,
)
if t is None:
t_value = torch.tensor(0.0, device=latent.device, dtype=latent.dtype)
elif isinstance(t, torch.Tensor):
t_value = t.reshape(-1)[0].to(device=latent.device, dtype=latent.dtype)
else:
t_value = torch.tensor(float(t), device=latent.device, dtype=latent.dtype)
t_normalized = t_value / 1000.0
cfg_batch = batch_size * 2
t_batch = t_normalized.expand(cfg_batch)
_, _, video_t, video_h, video_w = latent.shape
audio_t = audio_latent.shape[1]
# Keep zero-stride views here. Patchification/materialization happens
# once in the data proxy, so cloning the full token grids every denoise
# step only increases peak memory and allocator pressure.
per_token_video_t = t_batch.view(-1, 1, 1, 1, 1).expand(cfg_batch, 1, video_t, video_h, video_w)
per_token_audio_t = t_batch.view(-1, 1, 1).expand(cfg_batch, audio_t, 1)
audio_lengths = torch.full(
(cfg_batch,),
audio_latent_len,
device=latent.device,
dtype=torch.long,
)
text_lengths = torch.tensor(
[txt_feat_len] * batch_size + [null_txt_feat_len] * batch_size,
device=latent.device,
dtype=torch.long,
)
ref_audio_lengths = torch.full(
(cfg_batch,),
ref_audio_feat_len,
device=latent.device,
dtype=torch.long,
)
ref_video_lengths = torch.full(
(cfg_batch,),
ref_video_feat_len,
device=latent.device,
dtype=torch.long,
)
return ModelInput(
x_t=torch.cat((latent, latent), dim=0),
audio_x_t=torch.cat((audio_latent, audio_latent), dim=0),
audio_feat_len=audio_lengths,
txt_feat=torch.cat((txt_feat, null_txt_feat), dim=0),
txt_feat_len=text_lengths,
t=t_batch,
per_token_video_t=per_token_video_t,
per_token_audio_t=per_token_audio_t,
ref_audio_feat=torch.cat((ref_audio_feat, torch.zeros_like(ref_audio_feat)), dim=0),
ref_audio_feat_len=ref_audio_lengths,
ref_video_feat=torch.cat((ref_video_feat, torch.zeros_like(ref_video_feat)), dim=0),
ref_video_feat_len=ref_video_lengths,
ref_image_feat=ref_image_feat_cfg,
ref_image_feat_len=ref_image_feat_len_cfg,
ref_image_special_token_embedding=ref_image_special_tokens_cfg,
)
@staticmethod
def _split_cfg_model_input(model_input: ModelInput) -> tuple[ModelInput, ModelInput]:
"""Split the packed positive/negative batch for CFG rank dispatch."""
batch = model_input.x_t.shape[0]
if batch % 2:
raise ValueError("MAGI-2 CFG model input must contain equal positive and negative halves")
half = batch // 2
def split(value: torch.Tensor | Sequence[int] | None) -> tuple[torch.Tensor | Sequence[int] | None, ...]:
if value is None:
return None, None
if isinstance(value, torch.Tensor):
if value.ndim == 0 or value.shape[0] != batch:
raise ValueError("MAGI-2 CFG tensor fields must use the packed CFG batch dimension")
return value[:half], value[half:]
if len(value) != batch:
raise ValueError("MAGI-2 CFG sequence fields must use the packed CFG batch dimension")
return value[:half], value[half:]
field_pairs = {
name: split(getattr(model_input, name))
for name in (
"x_t",
"audio_x_t",
"audio_feat_len",
"txt_feat",
"txt_feat_len",
"t",
"ref_audio_feat",
"ref_audio_feat_len",
"ref_video_feat",
"ref_video_feat_len",
"per_token_video_t",
"per_token_audio_t",
"ref_image_feat",
"ref_image_feat_len",
"ref_image_special_token_embedding",
)
}
branches = []
for branch in range(2):
branches.append(ModelInput(**{name: pair[branch] for name, pair in field_pairs.items()}))
return branches[0], branches[1]
@staticmethod
def _prepare_ref_image_cfg(
ref_image_feat: torch.Tensor | None,
ref_image_feat_len: torch.Tensor | None,
ref_image_special_token_embedding: torch.Tensor | None,
cfg_config: CFGConfig | None,
) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
if ref_image_feat is None:
if ref_image_feat_len is not None or ref_image_special_token_embedding is not None:
raise ValueError("reference-image features, lengths, and special tokens must be provided together")
return None, None, None
if ref_image_feat_len is None or ref_image_special_token_embedding is None:
raise ValueError("reference-image features, lengths, and special tokens must be provided together")
keep_ref = bool(cfg_config is not None and cfg_config.use_ref_for_uncond)
uncond_feature = ref_image_feat if keep_ref else torch.zeros_like(ref_image_feat)
return (
torch.cat((ref_image_feat, uncond_feature), dim=0),
torch.cat((ref_image_feat_len, ref_image_feat_len), dim=0),
torch.cat(
(ref_image_special_token_embedding, ref_image_special_token_embedding),
dim=0,
),
)
def precalculate_cfg(
self,
t_list: Sequence[torch.Tensor],
latent_length: int,
cfg_config: CFGConfig,
*,
device: torch.device | str | None = None,
) -> tuple[list[torch.Tensor], list[float]]:
if latent_length < 0:
raise ValueError("latent_length must be non-negative")
if device is None:
if self.device is not None:
device = self.device
elif t_list:
device = t_list[0].device
else:
device = torch.device("cpu")
if cfg_config.use_cfg_trick:
video_cfg_base = (
torch.tensor(cfg_config.video_txt_guidance_scale, device=device)
.expand(1, 1, latent_length, 1, 1)
.clone()
)
video_cfg_base[:, :, : cfg_config.cfg_trick_start_frame] = min(
cfg_config.cfg_trick_value,
cfg_config.video_txt_guidance_scale,
)
else:
video_cfg_base = torch.tensor([cfg_config.video_txt_guidance_scale], device=device)
all_video_cfgs: list[torch.Tensor] = []
all_audio_cfgs: list[float] = []
for timestep in t_list:
video_cfg = video_cfg_base.clone()
if cfg_config.use_dynamic_cfg:
timestep_value = float(timestep.detach().reshape(-1)[0].item())
if timestep_value < cfg_config.dynamic_cfg_start_t:
video_cfg[video_cfg > cfg_config.dynamic_cfg_cutoff_value] = cfg_config.dynamic_cfg_cutoff_value
all_video_cfgs.append(video_cfg)
all_audio_cfgs.append(cfg_config.audio_txt_guidance_scale)
return all_video_cfgs, all_audio_cfgs
@staticmethod
def _cfg_scale_tensor(
guidance_scale: torch.Tensor | float,
like: torch.Tensor,
) -> torch.Tensor:
if isinstance(guidance_scale, torch.Tensor):
return guidance_scale.to(device=like.device, dtype=like.dtype)
return torch.tensor(float(guidance_scale), device=like.device, dtype=like.dtype)
@classmethod
def _get_skimming_mask(
cls,
x_orig: torch.Tensor,
cond: torch.Tensor,
uncond: torch.Tensor,
guidance_scale: torch.Tensor | float,
) -> torch.Tensor:
guidance_scale = cls._cfg_scale_tensor(guidance_scale, cond)
x_orig = x_orig.to(device=cond.device, dtype=cond.dtype)
denoised = x_orig - ((x_orig - uncond) + guidance_scale * ((x_orig - cond) - (x_orig - uncond)))
matching_pred_signs = (cond - uncond).sign() == cond.sign()
matching_diff_after = cond.sign() == (cond * guidance_scale - uncond * (guidance_scale - 1)).sign()
deviation_influence = denoised.sign() == (denoised - x_orig).sign()
return matching_pred_signs & matching_diff_after & deviation_influence
@classmethod
def _apply_skimmed_cfg_linear(
cls,
x_orig: torch.Tensor | None,
cond: torch.Tensor,
uncond: torch.Tensor,
guidance_scale: torch.Tensor | float,
cfg_config: CFGConfig | None,
) -> torch.Tensor:
if (
cfg_config is None
or not cfg_config.use_skimmed_cfg_linear
or x_orig is None
or not torch.any(uncond).item()
):
return uncond
guidance_scale = cls._cfg_scale_tensor(guidance_scale, cond)
valid_scale = guidance_scale > 1
if not torch.any(valid_scale).item():
return uncond
scale_delta = torch.where(
valid_scale,
guidance_scale - 1,
torch.ones_like(guidance_scale),
)
fallback_weight = (cfg_config.skimmed_cfg_scale - 1) / scale_delta
target_uncond = cond * (1 - fallback_weight) + uncond * fallback_weight
skim_mask = cls._get_skimming_mask(x_orig, cond, uncond, guidance_scale)
uncond = torch.where(skim_mask & valid_scale, target_uncond, uncond)
target_uncond = cond * (1 - fallback_weight) + uncond * fallback_weight
skim_mask = cls._get_skimming_mask(x_orig, uncond, cond, guidance_scale)
return torch.where(skim_mask & valid_scale, target_uncond, uncond)
@staticmethod
def _cfg_rescale_dims(tensor: torch.Tensor) -> tuple[int, ...]:
if tensor.ndim <= 1:
return ()
if tensor.ndim in (2, 3):
return (1,)
return tuple(range(2, tensor.ndim))
@classmethod
def _apply_cfg_rescale(
cls,
pos: torch.Tensor,
cfg: torch.Tensor,
cfg_config: CFGConfig | None,
) -> torch.Tensor:
if cfg_config is None or cfg_config.cfg_rescale <= 0 or cfg.numel() == 0:
return cfg
spatial_dims = cls._cfg_rescale_dims(cfg)
if not spatial_dims:
return cfg
pos_std = pos.float().std(dim=spatial_dims, keepdim=True, unbiased=False)
cfg_std = cfg.float().std(dim=spatial_dims, keepdim=True, unbiased=False).clamp_min(1e-6)
factor = cfg_config.cfg_rescale * (pos_std / cfg_std) + (1 - cfg_config.cfg_rescale)
return cfg * factor.to(device=cfg.device, dtype=cfg.dtype)
def cfg_velocity(
self,
model_output: tuple[torch.Tensor, torch.Tensor],
video_txt_guidance_scale: torch.Tensor | float,
audio_txt_guidance_scale: torch.Tensor | float,
cfg_config: CFGConfig | None = None,
latent: torch.Tensor | None = None,
audio_latent: torch.Tensor | None = None,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
cond_uncond_video, cond_uncond_audio = model_output
if cond_uncond_video.shape[0] % 2 or cond_uncond_audio.shape[0] % 2:
raise ValueError("MAGI-2 CFG outputs must contain equal conditional/unconditional halves")
video_batch = cond_uncond_video.shape[0] // 2
audio_batch = cond_uncond_audio.shape[0] // 2
if video_batch != audio_batch:
raise ValueError("video and audio CFG output batch sizes differ")
if cond_uncond_video.numel() > 0:
cond_video = cond_uncond_video[:video_batch]
uncond_video = cond_uncond_video[video_batch:]
uncond_video = self._apply_skimmed_cfg_linear(
latent,
cond_video,
uncond_video,
video_txt_guidance_scale,
cfg_config,
)
cfg_video = uncond_video + video_txt_guidance_scale * (cond_video - uncond_video)
cfg_video = self._apply_cfg_rescale(cond_video, cfg_video, cfg_config)
else:
cfg_video = None
if cond_uncond_audio.numel() > 0:
cond_audio = cond_uncond_audio[:audio_batch]
uncond_audio = cond_uncond_audio[audio_batch:]
uncond_audio = self._apply_skimmed_cfg_linear(
audio_latent,
cond_audio,
uncond_audio,
audio_txt_guidance_scale,
cfg_config,
)
cfg_audio = uncond_audio + audio_txt_guidance_scale * (cond_audio - uncond_audio)
cfg_audio = self._apply_cfg_rescale(cond_audio, cfg_audio, cfg_config)
else:
cfg_audio = None
return cfg_video, cfg_audio
def combine_cfg_noise(
self,
positive_noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
true_cfg_scale: float,
cfg_normalize: bool = False,
kwargs: dict[str, Any] | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Preserve MAGI-2's dual-modality guidance under shared CFG dispatch."""
del true_cfg_scale, cfg_normalize
if not isinstance(positive_noise_pred, tuple) or not isinstance(negative_noise_pred, tuple):
raise TypeError("MAGI-2 CFG predictions must contain video and audio tensors")
if len(positive_noise_pred) != 2 or len(negative_noise_pred) != 2:
raise ValueError("MAGI-2 CFG predictions must contain exactly two tensors")
if kwargs is None:
raise ValueError("MAGI-2 CFG combine requires guidance and latent context")
cfg_video, cfg_audio = self.cfg_velocity(
(
torch.cat((positive_noise_pred[0], negative_noise_pred[0]), dim=0),
torch.cat((positive_noise_pred[1], negative_noise_pred[1]), dim=0),
),
kwargs["video_txt_guidance_scale"],
kwargs["audio_txt_guidance_scale"],
kwargs.get("cfg_config"),
kwargs.get("latent"),
kwargs.get("audio_latent"),
)
if cfg_video is None or cfg_audio is None:
raise RuntimeError("MAGI-2 Preview must produce both video and audio CFG predictions")
return cfg_video, cfg_audio
@staticmethod
def _step_guided(
guided_output: tuple[torch.Tensor, torch.Tensor],
latent: torch.Tensor,
audio_latent: torch.Tensor,
video_scheduler: SchedulerMixin,
audio_scheduler: SchedulerMixin,
t: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
cfg_video, cfg_audio = guided_output
latent = video_scheduler.step(cfg_video, t, latent, return_dict=False)[0]
audio_latent = audio_scheduler.step(cfg_audio, t, audio_latent, return_dict=False)[0]
return latent, audio_latent
def step(
self,
model_output: tuple[torch.Tensor, torch.Tensor],
latent: torch.Tensor,
audio_latent: torch.Tensor,
video_txt_guidance_scale: torch.Tensor | float,
audio_txt_guidance_scale: torch.Tensor | float,
video_scheduler: SchedulerMixin,
audio_scheduler: SchedulerMixin,
t: torch.Tensor,
cfg_config: CFGConfig | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
cfg_video, cfg_audio = self.cfg_velocity(
model_output,
video_txt_guidance_scale,
audio_txt_guidance_scale,
cfg_config,
latent,
audio_latent,
)
if cfg_video is not None:
latent = video_scheduler.step(cfg_video, t, latent, return_dict=False)[0]
if cfg_audio is not None:
audio_latent = audio_scheduler.step(
cfg_audio,
t,
audio_latent,
return_dict=False,
)[0]
return latent, audio_latent, cfg_video, cfg_audio
__all__ = [
"CFGConfig",
"Magi2PreviewSampler",
"SamplerInput",
"build_magi2_preview_schedulers",
]
vllm_omni/diffusion/models/magi2/turbo_vae.py
+582
View File
@@ -0,0 +1,582 @@
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2025-2026 SandAI. All Rights Reserved.
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Native decoder for the TurboVAED checkpoint shipped with MAGI-2.
Adapted from SandAI's MAGI-2 Preview inference implementation. Training-only
modules and MagiCompiler integration are intentionally omitted.
"""
from __future__ import annotations
import json
from collections.abc import Mapping
from pathlib import Path
from typing import Any
import torch
import torch.distributed as dist
import torch.nn as nn
from einops import rearrange
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
DistributedOperator,
DistributedVaeMixin,
GridSpec,
TileTask,
)
_WAN22_LATENT_MEAN = (
-0.2289,
-0.0052,
-0.1323,
-0.2339,
-0.2799,
0.0174,
0.1838,
0.1557,
-0.1382,
0.0542,
0.2813,
0.0891,
0.1570,
-0.0098,
0.0375,
-0.1825,
-0.2246,
-0.1207,
-0.0698,
0.5109,
0.2665,
-0.2108,
-0.2158,
0.2502,
-0.2055,
-0.0322,
0.1109,
0.1567,
-0.0729,
0.0899,
-0.2799,
-0.1230,
-0.0313,
-0.1649,
0.0117,
0.0723,
-0.2839,
-0.2083,
-0.0520,
0.3748,
0.0152,
0.1957,
0.1433,
-0.2944,
0.3573,
-0.0548,
-0.1681,
-0.0667,
)
_WAN22_LATENT_STD = (
0.4765,
1.0364,
0.4514,
1.1677,
0.5313,
0.4990,
0.4818,
0.5013,
0.8158,
1.0344,
0.5894,
1.0901,
0.6885,
0.6165,
0.8454,
0.4978,
0.5759,
0.3523,
0.7135,
0.6804,
0.5833,
1.4146,
0.8986,
0.5659,
0.7069,
0.5338,
0.4889,
0.4917,
0.4069,
0.4999,
0.6866,
0.4093,
0.5709,
0.6065,
0.6415,
0.4944,
0.5726,
1.2042,
0.5458,
1.6887,
0.3971,
1.0600,
0.3943,
0.5537,
0.5444,
0.4089,
0.7468,
0.7744,
)
def turbo_unpatchify(x: torch.Tensor, patch_size: int) -> torch.Tensor:
if patch_size == 1:
return x
if x.ndim != 5:
raise ValueError(f"Expected a 5-D TurboVAED output, got {tuple(x.shape)}")
batch, patched_channels, frames, height, width = x.shape
patch_area = patch_size * patch_size
if patched_channels % patch_area:
raise ValueError("TurboVAED output channels must be divisible by patch_size squared")
channels = patched_channels // patch_area
x = x.view(batch, channels, patch_size, patch_size, frames, height, width)
x = x.permute(0, 1, 4, 5, 3, 6, 2).contiguous()
return x.view(batch, channels, frames, height * patch_size, width * patch_size)
class _RMSNorm(nn.Module):
def __init__(self, eps: float = 1e-8) -> None:
super().__init__()
self.eps = eps
def forward(self, x: torch.Tensor) -> torch.Tensor:
dtype = x.dtype
variance = x.float().pow(2).mean(1, keepdim=True)
return (x * torch.rsqrt(variance + self.eps)).to(dtype)
class _CausalConv3d(nn.Module):
"""Checkpoint-compatible non-causal 3-D convolution with edge time padding."""
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int | tuple[int, int, int] = 3,
stride: int | tuple[int, int, int] = 1,
*,
is_causal: bool = False,
) -> None:
super().__init__()
if is_causal:
raise ValueError("The released MAGI-2 TurboVAED checkpoint is non-causal")
self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size,) * 3
stride = stride if isinstance(stride, tuple) else (stride,) * 3
self.conv = nn.Conv3d(
in_channels,
out_channels,
self.kernel_size,
stride=stride,
padding=(0, self.kernel_size[1] // 2, self.kernel_size[2] // 2),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
temporal_pad = (self.kernel_size[0] - 1) // 2
if temporal_pad:
x = torch.cat(
[
x[:, :, :1].repeat(1, 1, temporal_pad, 1, 1),
x,
x[:, :, -1:].repeat(1, 1, temporal_pad, 1, 1),
],
dim=2,
)
return self.conv(x)
class _ResNetBlock3d(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int | None = None,
*,
is_upsampler_modified: bool = False,
is_causal: bool = False,
) -> None:
super().__init__()
out_channels = out_channels or in_channels
self.nonlinearity = nn.SiLU()
self.replace_nonlinearity = nn.ReLU()
self.is_upsampler_modified = is_upsampler_modified
self.norm1 = _RMSNorm()
self.conv1 = _CausalConv3d(in_channels, out_channels, is_causal=is_causal)
self.norm2 = _RMSNorm()
self.conv2 = _CausalConv3d(out_channels, out_channels, is_causal=is_causal)
self.norm3 = _RMSNorm(eps=1e-6) if in_channels != out_channels else None
self.conv_shortcut = (
_CausalConv3d(in_channels, out_channels, kernel_size=1, is_causal=is_causal)
if in_channels != out_channels
else None
)
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
x = self.norm1(inputs)
x = self.replace_nonlinearity(x) if self.is_upsampler_modified else self.nonlinearity(x)
x = self.conv1(x)
x = self.conv2(self.nonlinearity(self.norm2(x)))
residual = self.norm3(inputs) if self.norm3 is not None else inputs
if self.conv_shortcut is not None:
residual = self.conv_shortcut(residual)
return x + residual
class _WanUpsample(nn.Upsample):
def forward(self, x: torch.Tensor) -> torch.Tensor:
return super().forward(x.float()).type_as(x)
class _WanResample(nn.Module):
def __init__(self, channels: int, mode: str) -> None:
super().__init__()
self.mode = mode
if mode in ("upsample2d", "upsample3d"):
self.resample = nn.Sequential(
_WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
nn.Conv2d(channels, channels, 3, padding=1),
)
else:
self.resample = nn.Identity()
if mode == "upsample3d":
self.time_conv = _CausalConv3d(channels, channels * 2, (3, 1, 1))
def forward(self, x: torch.Tensor, *, is_first_chunk: bool) -> torch.Tensor:
batch, channels, frames, _height, _width = x.shape
if self.mode == "upsample3d":
x = self.time_conv(x)
x = rearrange(x, "b (split c) t h w -> b c (t split) h w", split=2)
if x.shape[1:3] != (channels, frames * 2):
raise RuntimeError("Unexpected temporal upsample shape")
if is_first_chunk:
x = x[:, :, 1:]
x = rearrange(x, "b c t h w -> (b t) c h w")
x = self.resample(x)
return rearrange(x, "(b t) c h w -> b c t h w", b=batch)
class _MidBlock3d(nn.Module):
def __init__(self, channels: int, num_layers: int, *, is_causal: bool) -> None:
super().__init__()
self.resnets = nn.ModuleList([_ResNetBlock3d(channels, is_causal=is_causal) for _ in range(num_layers)])
def forward(self, x: torch.Tensor) -> torch.Tensor:
for resnet in self.resnets:
x = resnet(x)
return x
class _UpBlock3d(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
num_layers: int,
*,
spatio_temporal_scale: bool,
spatio_only: bool,
is_causal: bool,
) -> None:
super().__init__()
self.conv_in = (
_ResNetBlock3d(in_channels, out_channels, is_causal=is_causal) if in_channels != out_channels else None
)
self.upsamplers = (
nn.ModuleList([_WanResample(out_channels, "upsample2d" if spatio_only else "upsample3d")])
if spatio_temporal_scale
else None
)
self.resnets = nn.ModuleList(
[
_ResNetBlock3d(
out_channels,
is_upsampler_modified=spatio_temporal_scale,
is_causal=is_causal,
)
for _ in range(num_layers)
]
)
def forward(self, x: torch.Tensor, *, is_first_chunk: bool) -> torch.Tensor:
if self.conv_in is not None:
x = self.conv_in(x)
if self.upsamplers is not None:
for upsampler in self.upsamplers:
x = upsampler(x, is_first_chunk=is_first_chunk)
for resnet in self.resnets:
x = resnet(x)
return x
class _TurboDecoder3d(nn.Module):
def __init__(
self,
*,
latent_channels: int,
out_channels: int,
block_out_channels: tuple[int, ...],
layers_per_block: tuple[int, ...],
spatio_temporal_scaling: tuple[bool, ...],
spatio_only: tuple[bool, ...],
patch_size: int,
decoder_causal: bool,
use_unpatchify: bool,
) -> None:
super().__init__()
if patch_size != 2 or not use_unpatchify:
raise ValueError("The released MAGI-2 TurboVAED requires patch_size=2 and use_unpatchify=true")
self.patch_size = patch_size
blocks = tuple(reversed(block_out_channels))
scaling = tuple(reversed(spatio_temporal_scaling))
layers = tuple(reversed(layers_per_block))
spatial_only = tuple(reversed(spatio_only))
current = blocks[0]
self.conv_in = _CausalConv3d(latent_channels, current, is_causal=decoder_causal)
self.mid_block = _MidBlock3d(current, layers[0], is_causal=decoder_causal)
self.up_blocks = nn.ModuleList()
for index, output in enumerate(blocks):
self.up_blocks.append(
_UpBlock3d(
current,
output,
layers[index + 1],
spatio_temporal_scale=scaling[index],
spatio_only=spatial_only[index],
is_causal=decoder_causal,
)
)
current = output
self.conv_act = nn.SiLU()
self.conv_out = _CausalConv3d(
current,
out_channels * patch_size * patch_size,
is_causal=decoder_causal,
)
def forward(self, x: torch.Tensor, *, is_first_chunk: bool) -> torch.Tensor:
x = self.mid_block(self.conv_in(x))
for block in self.up_blocks:
x = block(x, is_first_chunk=is_first_chunk)
dtype = x.dtype
x = (x * torch.rsqrt(x.float().pow(2).mean(1, keepdim=True) + 1e-8)).to(dtype)
return turbo_unpatchify(self.conv_out(self.conv_act(x)), self.patch_size)
def _load_turbo_config(config_path: Path) -> dict[str, Any]:
if not config_path.is_file():
raise FileNotFoundError(f"TurboVAED config does not exist: {config_path}")
with config_path.open(encoding="utf-8") as file:
config = json.load(file)
unsupported = {
"decoder_causal": config.get("decoder_causal", False),
"decoder_is_dw_conv": any(config.get("decoder_is_dw_conv", [])),
"use_unpatchify": not config.get("use_unpatchify", False),
}
enabled = [name for name, value in unsupported.items() if value]
if enabled:
raise ValueError(f"Unsupported TurboVAED configuration: {', '.join(enabled)}")
return config
def extract_turbo_decoder_state_dict(checkpoint: Mapping[str, Any]) -> dict[str, torch.Tensor]:
"""Select decoder EMA tensors and strip the training module prefix."""
state: Any = checkpoint
if "ema_state_dict" in state:
state = state["ema_state_dict"]
elif "state_dict" in state:
state = state["state_dict"]
if not isinstance(state, dict):
raise ValueError("TurboVAED checkpoint does not contain a state dictionary")
result = {}
for key, value in state.items():
target = key[7:] if key.startswith("module.") else key
if target.startswith("decoder."):
result[target] = value
if not result:
raise ValueError("TurboVAED checkpoint contains no decoder tensors")
return result
class Magi2TurboVAEDecoder(nn.Module, DistributedVaeMixin):
"""Load and run MAGI-2's default distilled video decoder."""
def __init__(
self,
config_path: str | Path,
checkpoint_path: str | Path,
*,
device: str | torch.device = "cpu",
dtype: torch.dtype = torch.float32,
) -> None:
super().__init__()
config = _load_turbo_config(Path(config_path))
self.z_dim = int(config["latent_channels"])
self.first_chunk_size = int(config["first_chunk_size"])
self.step_size = int(config["step_size"])
self.temporal_compression_ratio = int(config["temporal_compression_ratio"])
self.spatial_compression_ratio = int(config["spatial_compression_ratio"])
self.decoder = _TurboDecoder3d(
latent_channels=self.z_dim,
out_channels=int(config["out_channels"]),
block_out_channels=tuple(config["decoder_block_out_channels"]),
layers_per_block=tuple(config["decoder_layers_per_block"]),
spatio_temporal_scaling=tuple(config["decoder_spatio_temporal_scaling"]),
spatio_only=tuple(config["decoder_spatio_only"]),
patch_size=int(config["patch_size"]),
decoder_causal=bool(config["decoder_causal"]),
use_unpatchify=bool(config["use_unpatchify"]),
).to(device=device, dtype=dtype)
self.register_buffer(
"latent_mean",
torch.tensor(_WAN22_LATENT_MEAN, dtype=torch.float32, device=device),
persistent=False,
)
self.register_buffer(
"latent_std",
torch.tensor(_WAN22_LATENT_STD, dtype=torch.float32, device=device),
persistent=False,
)
checkpoint_path = Path(checkpoint_path)
if not checkpoint_path.is_file():
raise FileNotFoundError(f"TurboVAED checkpoint does not exist: {checkpoint_path}")
checkpoint = torch.load(
checkpoint_path,
map_location="cpu",
weights_only=True,
mmap=True,
)
state = extract_turbo_decoder_state_dict(checkpoint)
self.load_state_dict(state, strict=True)
self.eval().requires_grad_(False)
self.use_tiling = False
self.use_slicing = False
if dist.is_available() and dist.is_initialized():
self.init_distributed()
def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None:
if mode != "tile":
raise ValueError(f"MAGI-2 TurboVAE supports temporal tile parallelism only, got {mode!r}")
if not hasattr(self, "distributed_executor"):
if parallel_size > 1:
raise RuntimeError("TurboVAE patch parallelism requires initialized torch.distributed")
return
super().set_parallel_size(parallel_size, mode=mode)
def is_distributed_enabled(self) -> bool:
return hasattr(self, "distributed_executor") and super().is_distributed_enabled()
def _prepare_latent(self, z: torch.Tensor) -> tuple[torch.Tensor, int]:
z = z.to(device=self.latent_std.device)
dtype = z.dtype
z = z * self.latent_std.view(1, self.z_dim, 1, 1, 1) + self.latent_mean.view(1, self.z_dim, 1, 1, 1)
z = z.to(dtype)
first = self.first_chunk_size
step = self.step_size
frames = z.shape[2]
if frames < first:
padding = first - frames
elif (frames - first) % step:
padding = step - (frames - first) % step
else:
padding = 0
if padding:
z = torch.cat([z, z[:, :, -1:].repeat(1, 1, padding, 1, 1)], dim=2)
return z, padding
def _chunk_tasks(self, z: torch.Tensor, padding: int) -> tuple[list[TileTask], GridSpec]:
first = self.first_chunk_size
step = self.step_size
overlap = self.temporal_compression_ratio
frames = z.shape[2]
descriptors: list[tuple[int, int, bool, int, int]] = []
if frames == first:
descriptors.append((0, frames, True, 0, 0))
else:
descriptors.append((0, first + 1, True, 0, overlap))
for index in range(first, frames, step):
last = index + step == frames
descriptors.append(
(
index - 1,
index + step if last else index + step + 1,
False,
overlap,
0 if last else overlap,
)
)
tasks = [
TileTask(
tile_id=chunk_index,
grid_coord=(chunk_index, int(is_first), crop_left, crop_right),
tensor=z[:, :, left:right],
workload=(right - left) * z.shape[3] * z.shape[4],
)
for chunk_index, (left, right, is_first, crop_left, crop_right) in enumerate(descriptors)
]
return tasks, GridSpec(
split_dims=(2,),
grid_shape=(len(tasks),),
tile_spec={"padding": padding},
output_dtype=z.dtype,
)
def _decode_chunk(self, task: TileTask) -> torch.Tensor:
_index, is_first, crop_left, crop_right = task.grid_coord
assert isinstance(task.tensor, torch.Tensor)
output = self.decoder(task.tensor, is_first_chunk=bool(is_first))
right = output.shape[2] - crop_right if crop_right else output.shape[2]
return output[:, :, crop_left:right]
def _merge_chunks(
self,
chunks: dict[tuple[int, ...], torch.Tensor],
grid_spec: GridSpec,
) -> torch.Tensor:
output = torch.cat([tensor for _coord, tensor in sorted(chunks.items())], dim=2)
padding = int(grid_spec.tile_spec["padding"])
if padding:
output = output[:, :, : -padding * self.temporal_compression_ratio]
return output
@torch.inference_mode()
def decode(self, z: torch.Tensor, *, output_offload: bool = False) -> torch.Tensor:
z, padding = self._prepare_latent(z)
tasks, grid_spec = self._chunk_tasks(z, padding)
if self.is_distributed_enabled():
output = self.distributed_executor.execute(
z,
DistributedOperator(
split=lambda _z: (tasks, grid_spec),
exec=self._decode_chunk,
merge=self._merge_chunks,
),
broadcast_result=False,
)
return output.cpu() if output_offload and output.numel() else output
chunks = []
for task in tasks:
tensor = self._decode_chunk(task)
chunks.append((task.grid_coord, tensor.cpu() if output_offload else tensor))
return self._merge_chunks(dict(chunks), grid_spec)
forward = decode
__all__ = [
"Magi2TurboVAEDecoder",
"extract_turbo_decoder_state_dict",
"turbo_unpatchify",
]
vllm_omni/diffusion/registry.py
+6
View File
@@ -296,6 +296,11 @@ _DIFFUSION_MODELS = {
"pipeline_sana_video_i2v",
"SanaImageToVideoPipeline",
),
"Magi2Pipeline": (
"magi2",
"pipeline_magi2",
"Magi2Pipeline",
),
"OmniVoicePipeline": (
"omnivoice",
"pipeline_omnivoice",
@@ -588,6 +593,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
"LingBotVideoPipeline": "get_lingbot_video_post_process_func",
"SanaVideoPipeline": "get_sana_video_post_process_func",
"SanaImageToVideoPipeline": "get_sana_video_i2v_post_process_func",
"Magi2Pipeline": "get_magi2_post_process_func",
"OmniVoicePipeline": "get_omnivoice_post_process_func",
"SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func",
"Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func",
vllm_omni/diffusion/utils/hf_utils.py
+36
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import os
from collections.abc import Mapping
from functools import lru_cache
@@ -12,6 +15,14 @@ DIFFUSION_MODEL_INDEX_FILES = (
"modular_model_index.json",
)
_MAGI2_MODEL_ID = "sand-ai/MAGI-2-preview"
_MAGI2_REQUIRED_FILES = (
"preview/model.safetensors.index.json",
"text_encoder/config.json",
"vae/Wan2.2_VAE.pth",
"turbo_vae/checkpoint.ckpt",
)
def get_diffusion_model_index(
model_name: str,
@@ -46,6 +57,26 @@ def _looks_like_bagel(model_name: str) -> bool:
return False
def _looks_like_magi2(model_name: str) -> bool:
"""Detect the official Hub ID or a complete local Preview checkpoint."""
normalized = str(model_name).strip().rstrip("/")
hub_prefix = "https://huggingface.co/"
if normalized.lower().startswith(hub_prefix):
normalized = normalized[len(hub_prefix) :]
normalized = normalized.split("/tree/", 1)[0]
if normalized.lower() == _MAGI2_MODEL_ID.lower():
return True
return os.path.isdir(normalized) and all(
os.path.isfile(os.path.join(normalized, *relative.split("/"))) for relative in _MAGI2_REQUIRED_FILES
)
def resolve_native_diffusion_model_class(model_name: str) -> str | None:
"""Resolve native checkpoints that have no root HF or Diffusers config."""
return "Magi2Pipeline" if _looks_like_magi2(model_name) else None
def _looks_like_dreamzero(model_name: str) -> bool:
"""Best-effort detection for DreamZero-style VLA diffusion checkpoints."""
try:
@@ -104,6 +135,11 @@ def is_diffusion_model(model_name: str) -> bool:
2. Check using vllm's get_hf_file_to_dict utility
3. Try the standard diffusers approach (may fail due to import issues)
"""
# Some native diffusion checkpoints have neither a standard Diffusers
# index nor a root HF config. Recognize their exact Hub ID/local signature.
if resolve_native_diffusion_model_class(model_name) is not None:
return True
# Strategy 1: Check local file system first (fastest, avoids import issues)
if os.path.isdir(model_name):
for filename in DIFFUSION_MODEL_INDEX_FILES:
vllm_omni/entrypoints/openai/api_server.py
+1 -1
View File
@@ -3348,7 +3348,7 @@ async def _parse_video_form(
width: int | None = Form(default=None),
height: int | None = Form(default=None),
num_frames: int | None = Form(default=None),
fps: int | None = Form(default=None),
fps: float | None = Form(default=None, ge=1, allow_inf_nan=False),
aspect_ratio: str | None = Form(default=None),
short_edge: int | None = Form(default=None, ge=1),
num_outputs_per_prompt: int = Form(default=1, ge=1, le=10),
vllm_omni/entrypoints/openai/protocol/videos.py
+31 -6
View File
@@ -55,7 +55,13 @@ class VideoParams(BaseModel):
width: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video width in pixels")
height: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video height in pixels")
num_frames: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Number of frames")
fps: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Frames per second for output video")
fps: float | None = Field(
default=None,
ge=1,
le=_INT64_MAX,
allow_inf_nan=False,
description="Frames per second for output video",
)
@property
def size(self) -> str | None:
@@ -146,7 +152,13 @@ class VideoGenerationRequest(BaseModel):
# Video-specific fields (top-level for OpenAI-style compatibility)
width: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video width in pixels")
height: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video height in pixels")
fps: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Frames per second for output video")
fps: float | None = Field(
default=None,
ge=1,
le=_INT64_MAX,
allow_inf_nan=False,
description="Frames per second for output video",
)
num_frames: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Number of frames to generate")
aspect_ratio: str | None = Field(
default=None,
@@ -279,7 +291,17 @@ class VideoGenerationRequest(BaseModel):
description=("Optional model-specific parameters passed directly to the model's extra_args. "),
)
def resolve_video_params(self) -> VideoParams:
def resolve_video_params(
self,
*,
default_fps: float | None = DEFAULT_FPS,
default_num_frames: int | None = None,
) -> VideoParams:
"""Resolve explicit fields with optional model-owned defaults.
Callers that know the active diffusion pipeline can supply its native
frame contract. Other callers retain the generic 24 fps behavior.
"""
vp = VideoParams(width=self.width, height=self.height, fps=self.fps, num_frames=self.num_frames)
if self.video_params is not None:
@@ -292,10 +314,13 @@ class VideoGenerationRequest(BaseModel):
vp.width, vp.height = parse_size(self.size)
if vp.fps is None:
vp.fps = DEFAULT_FPS
vp.fps = default_fps
if vp.num_frames is None and self.seconds is not None:
vp.num_frames = int(self.seconds) * int(vp.fps)
if vp.num_frames is None:
if default_num_frames is not None:
vp.num_frames = default_num_frames
elif self.seconds is not None and vp.fps is not None:
vp.num_frames = int(float(self.seconds) * float(vp.fps))
return vp
vllm_omni/entrypoints/openai/serving_video.py
+84 -14
View File
@@ -1,9 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import copy
import math
import time
from collections.abc import Mapping
from dataclasses import dataclass
@@ -36,7 +37,8 @@ from vllm_omni.entrypoints.openai.video_api_utils import (
encode_video_base64,
)
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
from vllm_omni.model_extras import should_preserve_reference_image_size
from vllm_omni.model_extras import get_video_generation_defaults, should_preserve_reference_image_size
from vllm_omni.model_extras.video_generation import VideoGenerationDefaults
from vllm_omni.outputs.output_metadata import (
DiffusionMetadataMapping,
DiffusionMultimodalOutput,
@@ -82,7 +84,7 @@ class VideoGenerationArtifacts:
audios: list[DiffusionPayloadValue | None]
actions: list[VideoAction | None]
audio_sample_rate: int
output_fps: int
output_fps: float
stage_durations: dict[str, float]
peak_memory_mb: float
@@ -111,6 +113,21 @@ class OmniOpenAIServingVideo:
return get_od_config()
return getattr(self._engine_client, "od_config", None)
def _resolve_video_generation_defaults(
self,
request: VideoGenerationRequest,
) -> VideoGenerationDefaults | None:
"""Resolve defaults owned by the active diffusion pipeline."""
od_config = self._resolve_diffusion_od_config()
model_class_name = None if od_config is None else getattr(od_config, "model_class_name", None)
try:
return get_video_generation_defaults(model_class_name, request.extra_params)
except (TypeError, ValueError) as exc:
raise HTTPException(
status_code=HTTPStatus.BAD_REQUEST.value,
detail=str(exc),
) from exc
@property
def model_name(self) -> str | None:
return self._model_name
@@ -206,7 +223,49 @@ class OmniOpenAIServingVideo:
)
provided_fields = request.model_fields_set
fps_provided = self._request_fps_provided(request)
vp = request.resolve_video_params()
num_frames_provided = self._request_num_frames_provided(request)
video_defaults = self._resolve_video_generation_defaults(request)
if video_defaults is None:
vp = request.resolve_video_params()
else:
vp = request.resolve_video_params(
default_fps=video_defaults.fps,
default_num_frames=video_defaults.num_frames,
)
vp.width = vp.width or video_defaults.width
vp.height = vp.height or video_defaults.height
if num_frames_provided and vp.num_frames != video_defaults.num_frames:
raise HTTPException(
status_code=HTTPStatus.BAD_REQUEST.value,
detail=(f"This diffusion model requires {video_defaults.num_frames} frames; got {vp.num_frames}."),
)
if "num_inference_steps" not in provided_fields and gen_params.num_inference_steps is None:
gen_params.num_inference_steps = video_defaults.num_inference_steps
# Some native pipelines have a fixed duration. Validate both the
# OpenAI top-level field and model-specific aliases before dispatching
# the request; otherwise ``seconds`` would reach MAGI-2 as a duration
# override and fail only after the worker has started generation.
if video_defaults is not None and video_defaults.duration_seconds is not None:
requested_durations: list[float] = []
try:
if request.seconds is not None:
requested_durations.append(float(request.seconds))
if request.extra_params is not None:
for key in ("seconds", "duration"):
if key in request.extra_params:
requested_durations.append(float(request.extra_params[key]))
except (TypeError, ValueError) as exc:
raise HTTPException(
status_code=HTTPStatus.BAD_REQUEST.value,
detail="The requested video duration must be a finite number of seconds.",
) from exc
expected_duration = video_defaults.duration_seconds
if any(not math.isfinite(duration) or duration != expected_duration for duration in requested_durations):
raise HTTPException(
status_code=HTTPStatus.BAD_REQUEST.value,
detail=f"This diffusion model supports {expected_duration:g}-second clips only.",
)
if (
input_image is not None
and vp.width is not None
@@ -236,15 +295,17 @@ class OmniOpenAIServingVideo:
gen_params.num_frames = vp.num_frames
gen_params.num_outputs_per_prompt = request.num_outputs_per_prompt
if request.seconds is not None:
gen_params.extra_args.setdefault("duration", float(request.seconds))
if video_defaults is None or video_defaults.duration_seconds is None:
gen_params.extra_args.setdefault("duration", float(request.seconds))
if request.aspect_ratio is not None:
gen_params.extra_args["aspect_ratio"] = request.aspect_ratio
if request.short_edge is not None:
gen_params.extra_args["short_edge"] = request.short_edge
if request.start_time_seconds is not None:
gen_params.extra_args["start_time_seconds"] = request.start_time_seconds
# Leave fps/frame_rate as None when the user did not provide fps.
if fps_provided and vp.fps is not None:
# Model-owned defaults are part of the serving contract. Other models
# preserve their engine defaults when the user did not provide fps.
if (fps_provided or video_defaults is not None) and vp.fps is not None:
gen_params.fps = vp.fps
gen_params.frame_rate = float(vp.fps)
if "enable_frame_interpolation" in provided_fields:
@@ -464,6 +525,15 @@ class OmniOpenAIServingVideo:
return False
return "fps" in video_params.model_fields_set and video_params.fps is not None
@staticmethod
def _request_num_frames_provided(request: VideoGenerationRequest) -> bool:
if "num_frames" in request.model_fields_set and request.num_frames is not None:
return True
video_params = request.video_params
if video_params is None or "video_params" not in request.model_fields_set:
return False
return "num_frames" in video_params.model_fields_set and video_params.num_frames is not None
def _resolve_default_sampling_params(self) -> OmniDiffusionSamplingParams:
default_sampling_params_list = getattr(self._engine_client, "default_sampling_params_list", None)
if default_sampling_params_list:
@@ -747,15 +817,15 @@ class OmniOpenAIServingVideo:
return None
@staticmethod
def _resolve_fps(result: object) -> int | None:
def _resolve_fps(result: object) -> float | None:
"""Extract fps from multimodal_output if the model reported it."""
multimodal_output = getattr(result, "multimodal_output", None)
if isinstance(multimodal_output, Mapping):
fps = multimodal_output.get("fps")
if fps is not None:
try:
fps_val = fps.item() if hasattr(fps, "item") else int(fps)
if fps_val > 0:
fps_val = float(fps.item() if hasattr(fps, "item") else fps)
if math.isfinite(fps_val) and fps_val > 0:
return fps_val
except (TypeError, ValueError):
pass
@@ -767,8 +837,8 @@ class OmniOpenAIServingVideo:
fps = mm.get("fps")
if fps is not None:
try:
fps_val = fps.item() if hasattr(fps, "item") else int(fps)
if fps_val > 0:
fps_val = float(fps.item() if hasattr(fps, "item") else fps)
if math.isfinite(fps_val) and fps_val > 0:
return fps_val
except (TypeError, ValueError):
pass
@@ -778,8 +848,8 @@ class OmniOpenAIServingVideo:
fps = mm.get("fps")
if fps is not None:
try:
fps_val = fps.item() if hasattr(fps, "item") else int(fps)
if fps_val > 0:
fps_val = float(fps.item() if hasattr(fps, "item") else fps)
if math.isfinite(fps_val) and fps_val > 0:
return fps_val
except (TypeError, ValueError):
pass
vllm_omni/entrypoints/openai/video_api_utils.py
+2 -2
View File
@@ -786,7 +786,7 @@ def _encode_video_bytes_legacy(
def _encode_video_bytes(
video: Any,
fps: int,
fps: int | float,
audio: Any | None = None,
audio_sample_rate: int | None = None,
video_codec_options: dict[str, str] | None = None,
@@ -908,7 +908,7 @@ def create_streaming_video_encoder(
def encode_video_base64(
video: Any,
fps: int,
fps: int | float,
audio: Any | None = None,
audio_sample_rate: int | None = None,
video_codec_options: dict[str, str] | None = None,
vllm_omni/entrypoints/utils.py
+5 -1
View File
@@ -26,6 +26,7 @@ from vllm_omni.config.yaml_util import create_config, load_yaml_config
from vllm_omni.diffusion.utils.hf_utils import (
_looks_like_dreamzero,
get_diffusion_model_index,
resolve_native_diffusion_model_class,
)
from vllm_omni.entrypoints.stage_utils import _to_dict
from vllm_omni.inputs.data import OmniSamplingParams
@@ -294,7 +295,10 @@ def resolve_model_config_path(model: str) -> str | None:
model_type = hf_config.model_type
except (ValueError, Exception):
# If standard transformers format fails, try diffusers format
if get_diffusion_model_index(config_source) is not None:
native_model_class = resolve_native_diffusion_model_class(config_source)
if native_model_class is not None:
model_type = native_model_class
elif get_diffusion_model_index(config_source) is not None:
model_type = _try_get_class_name_from_diffusers_config(config_source)
if model_type is None:
raise ValueError(
vllm_omni/inputs/data.py
+4 -1
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import copy
import pprint
from dataclasses import asdict, dataclass, field
@@ -251,7 +254,7 @@ class OmniDiffusionSamplingParams:
# Original dimensions (before VAE scaling)
height: int | None = None
width: int | None = None
fps: int | None = None # Integer FPS kept for request compatibility and output encoding.
fps: float | None = None # Request/output-encoding FPS; fractional rates are supported.
frame_rate: float | None = None # Floating-point rate used by the diffusion model when it differs from `fps`.
height_not_provided: bool = False
width_not_provided: bool = False
vllm_omni/model_extras/__init__.py
+3 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from vllm_omni.model_extras.registry import (
build_image_to_image_prompt,
@@ -11,6 +11,7 @@ from vllm_omni.model_extras.registry import (
get_model_class_name,
get_output_tensor_range,
get_transformer_config_subfolder,
get_video_generation_defaults,
get_x_to_text_model_family,
should_init_extra_args_for_non_diffusion_stages,
should_preserve_reference_image_size,
@@ -26,6 +27,7 @@ __all__ = [
"get_model_class_name",
"get_output_tensor_range",
"get_transformer_config_subfolder",
"get_video_generation_defaults",
"get_x_to_text_model_family",
"should_init_extra_args_for_non_diffusion_stages",
"should_preserve_reference_image_size",
vllm_omni/model_extras/magi2.py
+55
View File
@@ -0,0 +1,55 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Request fields supported by the MAGI-2 Preview pipeline."""
from __future__ import annotations
from collections.abc import Mapping
from typing import Any
from vllm_omni.model_extras.video_generation import VideoGenerationDefaults
MAGI2_EXTRA_BODY_PARAMS = frozenset(
{
"seconds",
"resolution",
"image_path",
"output_width",
"output_height",
"deterministic",
}
)
MAGI2_EXTRA_OUTPUT_PARAMS: frozenset[str] = frozenset()
def magi2_preserves_reference_image_size(
*,
model: str | None,
revision: str | None = None,
) -> bool:
"""MAGI-2 owns checkpoint-aligned reference resize and padding."""
del model, revision
return True
def get_magi2_video_generation_defaults(
extra_body: Mapping[str, Any] | None = None,
) -> VideoGenerationDefaults:
resolution = str((extra_body or {}).get("resolution", "540p")).lower()
if resolution == "272p":
width, height = 448, 256
elif resolution == "540p":
width, height = 896, 512
else:
raise ValueError("MAGI-2 Preview resolution must be '272p' or '540p'.")
return VideoGenerationDefaults(
width=width,
height=height,
num_frames=125,
num_inference_steps=100,
fps=12.5,
duration_seconds=10.0,
output="magi2_output.mp4",
default_negative_prompt=None,
)
vllm_omni/model_extras/registry.py
+26 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -36,6 +36,12 @@ from vllm_omni.model_extras.ltx2 import (
ltx_preserves_reference_image_size,
ltx_transformer_config_subfolder,
)
from vllm_omni.model_extras.magi2 import (
MAGI2_EXTRA_BODY_PARAMS,
MAGI2_EXTRA_OUTPUT_PARAMS,
get_magi2_video_generation_defaults,
magi2_preserves_reference_image_size,
)
from vllm_omni.model_extras.mammothmodal2_preview import (
MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS,
MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS,
@@ -70,6 +76,7 @@ from vllm_omni.model_extras.vace import (
from vllm_omni.model_extras.vace import (
build_image_to_video_prompt as build_vace_image_to_video_prompt,
)
from vllm_omni.model_extras.video_generation import VideoGenerationDefaults
TextToImagePromptBuilder = Callable[
[str, str | None, int | None, int | None],
@@ -192,6 +199,12 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
"extra_body_params": COSMOS3_EXTRA_BODY_PARAMS,
"extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS,
},
"Magi2Pipeline": {
"extra_body_params": MAGI2_EXTRA_BODY_PARAMS,
"extra_output_params": MAGI2_EXTRA_OUTPUT_PARAMS,
"video_generation_defaults_builder": get_magi2_video_generation_defaults,
"reference_image_size_resolver": magi2_preserves_reference_image_size,
},
"HeliosPipeline": {
"extra_body_params": HELIOS_EXTRA_BODY_PARAMS,
"extra_output_params": HELIOS_EXTRA_OUTPUT_PARAMS,
@@ -286,6 +299,18 @@ def get_extra_output_params(model_class_name: str | None) -> frozenset[str]:
return spec.get("extra_output_params", frozenset()) if spec is not None else frozenset()
def get_video_generation_defaults(
model_class_name: str | None,
extra_body: Mapping[str, Any] | None = None,
) -> VideoGenerationDefaults | None:
"""Return model-owned defaults for the shared video examples, if declared."""
spec = _get_spec(model_class_name)
if spec is None:
return None
builder = spec.get("video_generation_defaults_builder")
return builder(extra_body) if builder is not None else None
def get_output_tensor_range(model_class_name: str | None) -> OutputTensorRange:
"""Return the declared range for floating-point tensor outputs.
vllm_omni/model_extras/video_generation.py
+44
View File
@@ -0,0 +1,44 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
@dataclass(frozen=True)
class VideoGenerationDefaults:
"""Model-owned defaults consumed by shared video examples."""
width: int
height: int
num_frames: int
num_inference_steps: int
fps: float
output: str | None = None
guidance_scale: float | None = None
flow_shift: float | None = None
dimension_multiple: int = 16
default_negative_prompt: str | None = ""
duration_seconds: float | None = None
@property
def max_area(self) -> int:
return self.width * self.height
def cli_defaults(self) -> dict[str, Any]:
defaults: dict[str, Any] = {
"height": self.height,
"width": self.width,
"num_frames": self.num_frames,
"num_inference_steps": self.num_inference_steps,
"fps": self.fps,
}
if self.output is not None:
defaults["output"] = self.output
if self.guidance_scale is not None:
defaults["guidance_scale"] = self.guidance_scale
if self.flow_shift is not None:
defaults["flow_shift"] = self.flow_shift
return defaults