From d4096e01febb59e82b9af237ba377befb4be8dff Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Sat, 5 Sep 2026 15:32:30 +0800 Subject: [PATCH] [Model] Add native MAGI-2 Preview diffusion support (#5918) Signed-off-by: Hongsheng Liu Signed-off-by: hsliuustc0106 --- .buildkite/cuda/test-nightly.yml | 6 + .claude/skills/add-diffusion-model/SKILL.md | 297 ++--- .../references/cache-dit-patterns.md | 10 +- .../references/custom-model-patterns.md | 22 +- .../native-model-integration-checklist.md | 131 ++ .claude/skills/add-recipe/SKILL.md | 126 ++ .claude/skills/add-tts-model/SKILL.md | 8 + .claude/skills/readme.md | 3 + docs/contributing/model/adding_omni_model.md | 43 +- docs/contributing/model/adding_tts_model.md | 13 +- .../figures/magi2/dlo-pareto-1-step.svg | 59 + .../figures/magi2/dlo-pareto-4-step.svg | 65 + docs/models/supported_models.md | 9 +- docs/user_guide/diffusion_features.md | 17 +- .../image_to_video/README.md | 12 +- .../image_to_video/image_to_video.py | 74 +- .../text_to_video/text_to_video.md | 8 + .../text_to_video/text_to_video.py | 70 +- pyproject.toml | 2 + recipes/README.md | 17 +- recipes/SandAI/MAGI-2-preview-L20X.md | 288 +++++ recipes/TEMPLATE.md | 76 +- tests/config/test_environment_variables.py | 8 +- .../attention/test_attention_config.py | 1 + tests/diffusion/attention/test_flash_attn.py | 26 +- tests/diffusion/cache/test_cache_backends.py | 7 +- tests/diffusion/cache/test_cache_dit.py | 42 +- tests/diffusion/distributed/test_hsdp.py | 47 + .../models/ltx2/test_ltx2_transformer.py | 1 + .../models/magi2/test_audio_decoder.py | 56 + .../magi2/test_native_distributed_parity.py | 229 ++++ .../models/magi2/test_native_dlo_allgather.py | 294 +++++ .../models/magi2/test_native_preview.py | 275 ++++ .../models/magi2/test_pipeline_magi2.py | 631 +++++++++ .../diffusion/models/magi2/test_turbo_vae.py | 66 + tests/e2e/online_serving/test_magi2.py | 77 ++ .../openai_api/test_video_server.py | 84 +- .../test_magi2_shared_defaults.py | 55 + tests/helpers/assertions.py | 60 +- .../tests/test_video_response_contract.py | 99 ++ tests/model_extras/test_model_extras.py | 19 +- tests/model_tests/diffusion/test_alignment.py | 6 + .../config/environment_variable_inventory.py | 5 + .../diffusion/attention/backends/utils/fa.py | 85 +- vllm_omni/diffusion/attention/layer.py | 177 +-- .../diffusion/cache/cachedit/__init__.py | 4 +- vllm_omni/diffusion/cache/cachedit/backend.py | 44 +- .../cache/cachedit/model_specific.py | 34 +- vllm_omni/diffusion/data.py | 17 +- vllm_omni/diffusion/distributed/hsdp.py | 10 +- .../checkpoint_adapters/direct_mmap.py | 24 +- .../model_loader/host_weight_plan.py | 43 +- vllm_omni/diffusion/model_metadata.py | 5 + vllm_omni/diffusion/models/magi2/__init__.py | 21 + vllm_omni/diffusion/models/magi2/attention.py | 293 +++++ .../diffusion/models/magi2/audio_decoder.py | 195 +++ .../diffusion/models/magi2/conditioning.py | 321 +++++ .../models/magi2/configuration_magi2.py | 120 ++ vllm_omni/diffusion/models/magi2/layers.py | 437 +++++++ vllm_omni/diffusion/models/magi2/mh_moe.py | 547 ++++++++ .../diffusion/models/magi2/modeling_magi2.py | 739 +++++++++++ vllm_omni/diffusion/models/magi2/parallel.py | 392 ++++++ .../diffusion/models/magi2/pipeline_magi2.py | 1124 +++++++++++++++++ .../models/magi2/preview_data_proxy.py | 691 ++++++++++ .../diffusion/models/magi2/sampler_magi2.py | 682 ++++++++++ vllm_omni/diffusion/models/magi2/turbo_vae.py | 582 +++++++++ vllm_omni/diffusion/registry.py | 6 + vllm_omni/diffusion/utils/hf_utils.py | 36 + vllm_omni/entrypoints/openai/api_server.py | 2 +- .../entrypoints/openai/protocol/videos.py | 37 +- vllm_omni/entrypoints/openai/serving_video.py | 98 +- .../entrypoints/openai/video_api_utils.py | 4 +- vllm_omni/entrypoints/utils.py | 6 +- vllm_omni/inputs/data.py | 5 +- vllm_omni/model_extras/__init__.py | 4 +- vllm_omni/model_extras/magi2.py | 55 + vllm_omni/model_extras/registry.py | 27 +- vllm_omni/model_extras/video_generation.py | 44 + 78 files changed, 9896 insertions(+), 459 deletions(-) create mode 100644 .claude/skills/add-diffusion-model/references/native-model-integration-checklist.md create mode 100644 .claude/skills/add-recipe/SKILL.md create mode 100644 docs/design/figures/magi2/dlo-pareto-1-step.svg create mode 100644 docs/design/figures/magi2/dlo-pareto-4-step.svg create mode 100644 recipes/SandAI/MAGI-2-preview-L20X.md create mode 100644 tests/diffusion/models/magi2/test_audio_decoder.py create mode 100644 tests/diffusion/models/magi2/test_native_distributed_parity.py create mode 100644 tests/diffusion/models/magi2/test_native_dlo_allgather.py create mode 100644 tests/diffusion/models/magi2/test_native_preview.py create mode 100644 tests/diffusion/models/magi2/test_pipeline_magi2.py create mode 100644 tests/diffusion/models/magi2/test_turbo_vae.py create mode 100644 tests/e2e/online_serving/test_magi2.py create mode 100644 tests/examples/offline_inference/test_magi2_shared_defaults.py create mode 100644 tests/helpers/tests/test_video_response_contract.py create mode 100644 vllm_omni/diffusion/models/magi2/__init__.py create mode 100644 vllm_omni/diffusion/models/magi2/attention.py create mode 100644 vllm_omni/diffusion/models/magi2/audio_decoder.py create mode 100644 vllm_omni/diffusion/models/magi2/conditioning.py create mode 100644 vllm_omni/diffusion/models/magi2/configuration_magi2.py create mode 100644 vllm_omni/diffusion/models/magi2/layers.py create mode 100644 vllm_omni/diffusion/models/magi2/mh_moe.py create mode 100644 vllm_omni/diffusion/models/magi2/modeling_magi2.py create mode 100644 vllm_omni/diffusion/models/magi2/parallel.py create mode 100644 vllm_omni/diffusion/models/magi2/pipeline_magi2.py create mode 100644 vllm_omni/diffusion/models/magi2/preview_data_proxy.py create mode 100644 vllm_omni/diffusion/models/magi2/sampler_magi2.py create mode 100644 vllm_omni/diffusion/models/magi2/turbo_vae.py create mode 100644 vllm_omni/model_extras/magi2.py create mode 100644 vllm_omni/model_extras/video_generation.py diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 31bc23cda..c85e9abc1 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -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: diff --git a/.claude/skills/add-diffusion-model/SKILL.md b/.claude/skills/add-diffusion-model/SKILL.md index 16b68b06e..bda30f32e 100644 --- a/.claude/skills/add-diffusion-model/SKILL.md +++ b/.claude/skills/add-diffusion-model/SKILL.md @@ -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//pipeline_.py` | Always required | | **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models//_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` | | **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models//` as separate files | `autoencoder.py`, `fusion.py` | -| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | 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//download_.py` | `download_.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//` 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//` 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 diff --git a/.claude/skills/add-diffusion-model/references/cache-dit-patterns.md b/.claude/skills/add-diffusion-model/references/cache-dit-patterns.md index 2385a70ee..2edfc26de 100644 --- a/.claude/skills/add-diffusion-model/references/cache-dit-patterns.md +++ b/.claude/skills/add-diffusion-model/references/cache-dit-patterns.md @@ -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 diff --git a/.claude/skills/add-diffusion-model/references/custom-model-patterns.md b/.claude/skills/add-diffusion-model/references/custom-model-patterns.md index 4e25f9762..62d8edc83 100644 --- a/.claude/skills/add-diffusion-model/references/custom-model-patterns.md +++ b/.claude/skills/add-diffusion-model/references/custom-model-patterns.md @@ -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// @@ -28,7 +28,10 @@ examples/offline_inference// └── download_.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//`. 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: diff --git a/.claude/skills/add-diffusion-model/references/native-model-integration-checklist.md b/.claude/skills/add-diffusion-model/references/native-model-integration-checklist.md new file mode 100644 index 000000000..894a0c035 --- /dev/null +++ b/.claude/skills/add-diffusion-model/references/native-model-integration-checklist.md @@ -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. diff --git a/.claude/skills/add-recipe/SKILL.md b/.claude/skills/add-recipe/SKILL.md new file mode 100644 index 000000000..25ed06a51 --- /dev/null +++ b/.claude/skills/add-recipe/SKILL.md @@ -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//-.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 +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. diff --git a/.claude/skills/add-tts-model/SKILL.md b/.claude/skills/add-tts-model/SKILL.md index 8c11d04ba..9193b9bf6 100644 --- a/.claude/skills/add-tts-model/SKILL.md +++ b/.claude/skills/add-tts-model/SKILL.md @@ -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 diff --git a/.claude/skills/readme.md b/.claude/skills/readme.md index 54008b82c..d73dea3e8 100644 --- a/.claude/skills/readme.md +++ b/.claude/skills/readme.md @@ -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 diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index 42b9a53d7..1f1de15be 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/docs/contributing/model/adding_omni_model.md @@ -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//.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) diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index 264287306..7a07bd5ed 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -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 diff --git a/docs/design/figures/magi2/dlo-pareto-1-step.svg b/docs/design/figures/magi2/dlo-pareto-1-step.svg new file mode 100644 index 000000000..6d3e8eca3 --- /dev/null +++ b/docs/design/figures/magi2/dlo-pareto-1-step.svg @@ -0,0 +1,59 @@ + + MAGI-2 DLO one-step Pareto frontier + 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. + + + + One diffusion step + + + + Warm user latency (seconds) — lower is better + Throughput/device (outputs/s) — higher is better + + + + + DP1SP4 + 5.735 s, 0.0436 + + + DP2SP2 + 6.785 s, 0.0737 + + + DP4 + 9.997 s, 0.1000 + + 4x NVIDIA L20X · warm qualification run + diff --git a/docs/design/figures/magi2/dlo-pareto-4-step.svg b/docs/design/figures/magi2/dlo-pareto-4-step.svg new file mode 100644 index 000000000..7fc67c465 --- /dev/null +++ b/docs/design/figures/magi2/dlo-pareto-4-step.svg @@ -0,0 +1,65 @@ + + MAGI-2 DLO four-step Pareto frontier + 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. + + + + Four diffusion steps + + + + Warm user latency (seconds) — lower is better + Throughput/device (outputs/s) — higher is better + + + + + DP1SP4 + 8.647 s, 0.0289 + + + DP2SP2 + 11.682 s, 0.0428 + + + DP4 + 20.957 s, 0.0477 + + 4x NVIDIA L20X · warm qualification run + diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 5d0ca6dbc..9e3a5a9f7 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -18,7 +18,7 @@ th { | 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 + imagegen1) | `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 diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 7d443c942..00c245489 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -1,5 +1,7 @@ # Diffusion Advanced Features + + ## 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** | ❌ | ❌ | ❌5 | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | > 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 diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md index 02f233642..24bc52164 100644 --- a/examples/offline_inference/image_to_video/README.md +++ b/examples/offline_inference/image_to_video/README.md @@ -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 diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 65e7389f0..3f8776971 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -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, ) diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index 0cb75a8d4..99a7a9d9f 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -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). diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 02383c6e4..559bbbb50 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 92dcb1a05..73d21cbe8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/recipes/README.md b/recipes/README.md index c5f4166ce..88788e3d3 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -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 `-.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. diff --git a/recipes/SandAI/MAGI-2-preview-L20X.md b/recipes/SandAI/MAGI-2-preview-L20X.md new file mode 100644 index 000000000..c240360ea --- /dev/null +++ b/recipes/SandAI/MAGI-2-preview-L20X.md @@ -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: +- Architecture blog: +- 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. diff --git a/recipes/TEMPLATE.md b/recipes/TEMPLATE.md index f2a145660..cf6fcb184 100644 --- a/recipes/TEMPLATE.md +++ b/recipes/TEMPLATE.md @@ -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//-.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 +`-.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 --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. diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 86328b46f..52ff1e274 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -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, } diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index b695d59fe..656161589 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/tests/diffusion/attention/test_attention_config.py @@ -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, ) diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index 3b5c877cb..be78ad690 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/tests/diffusion/attention/test_flash_attn.py @@ -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). diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py index 85bee131b..29cda91a3 100644 --- a/tests/diffusion/cache/test_cache_backends.py +++ b/tests/diffusion/cache/test_cache_backends.py @@ -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") diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index a24595c9b..c1d061ef2 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/tests/diffusion/cache/test_cache_dit.py @@ -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 diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py index 25762c356..433dec02f 100644 --- a/tests/diffusion/distributed/test_hsdp.py +++ b/tests/diffusion/distributed/test_hsdp.py @@ -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.""" diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index f71b84d3e..937d487e6 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py @@ -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 diff --git a/tests/diffusion/models/magi2/test_audio_decoder.py b/tests/diffusion/models/magi2/test_audio_decoder.py new file mode 100644 index 000000000..71129cb4b --- /dev/null +++ b/tests/diffusion/models/magi2/test_audio_decoder.py @@ -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") diff --git a/tests/diffusion/models/magi2/test_native_distributed_parity.py b/tests/diffusion/models/magi2/test_native_distributed_parity.py new file mode 100644 index 000000000..5c3de50f0 --- /dev/null +++ b/tests/diffusion/models/magi2/test_native_distributed_parity.py @@ -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, + ) diff --git a/tests/diffusion/models/magi2/test_native_dlo_allgather.py b/tests/diffusion/models/magi2/test_native_dlo_allgather.py new file mode 100644 index 000000000..0c097cfc0 --- /dev/null +++ b/tests/diffusion/models/magi2/test_native_dlo_allgather.py @@ -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, + ) diff --git a/tests/diffusion/models/magi2/test_native_preview.py b/tests/diffusion/models/magi2/test_native_preview.py new file mode 100644 index 000000000..9167d72ee --- /dev/null +++ b/tests/diffusion/models/magi2/test_native_preview.py @@ -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) + ) diff --git a/tests/diffusion/models/magi2/test_pipeline_magi2.py b/tests/diffusion/models/magi2/test_pipeline_magi2.py new file mode 100644 index 000000000..7acb6b209 --- /dev/null +++ b/tests/diffusion/models/magi2/test_pipeline_magi2.py @@ -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
" + ) + 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, + ) diff --git a/tests/diffusion/models/magi2/test_turbo_vae.py b/tests/diffusion/models/magi2/test_turbo_vae.py new file mode 100644 index 000000000..487b0fb0a --- /dev/null +++ b/tests/diffusion/models/magi2/test_turbo_vae.py @@ -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" diff --git a/tests/e2e/online_serving/test_magi2.py b/tests/e2e/online_serving/test_magi2.py new file mode 100644 index 000000000..22a30b495 --- /dev/null +++ b/tests/e2e/online_serving/test_magi2.py @@ -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) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 6d8977c55..969691ebc 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -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): diff --git a/tests/examples/offline_inference/test_magi2_shared_defaults.py b/tests/examples/offline_inference/test_magi2_shared_defaults.py new file mode 100644 index 000000000..b8c2c233b --- /dev/null +++ b/tests/examples/offline_inference/test_magi2_shared_defaults.py @@ -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 diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index 4a7d1af24..65d28109f 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -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", ] diff --git a/tests/helpers/tests/test_video_response_contract.py b/tests/helpers/tests/test_video_response_contract.py new file mode 100644 index 000000000..73d1f94cb --- /dev/null +++ b/tests/helpers/tests/test_video_response_contract.py @@ -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) diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 463bedfc3..6bf823339 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -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: diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index acc1067a6..509f113b7 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -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", diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 78212777a..1a1d40141 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -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", diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index 54a27e217..e2f96abab 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/vllm_omni/diffusion/attention/backends/utils/fa.py @@ -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. diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 3b2d76ba1..ac541808b 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -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) diff --git a/vllm_omni/diffusion/cache/cachedit/__init__.py b/vllm_omni/diffusion/cache/cachedit/__init__.py index 7166f7098..348ffe1b8 100644 --- a/vllm_omni/diffusion/cache/cachedit/__init__.py +++ b/vllm_omni/diffusion/cache/cachedit/__init__.py @@ -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", diff --git a/vllm_omni/diffusion/cache/cachedit/backend.py b/vllm_omni/diffusion/cache/cachedit/backend.py index 9442ac46e..a0789dee5 100644 --- a/vllm_omni/diffusion/cache/cachedit/backend.py +++ b/vllm_omni/diffusion/cache/cachedit/backend.py @@ -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", diff --git a/vllm_omni/diffusion/cache/cachedit/model_specific.py b/vllm_omni/diffusion/cache/cachedit/model_specific.py index f7af987ce..0a25abd3c 100644 --- a/vllm_omni/diffusion/cache/cachedit/model_specific.py +++ b/vllm_omni/diffusion/cache/cachedit/model_specific.py @@ -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, } ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index d069d95d1..10a2f9a19 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -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 diff --git a/vllm_omni/diffusion/distributed/hsdp.py b/vllm_omni/diffusion/distributed/hsdp.py index c496804bc..82c1d1986 100644 --- a/vllm_omni/diffusion/distributed/hsdp.py +++ b/vllm_omni/diffusion/distributed/hsdp.py @@ -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) diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py index 86b23429b..b3e74e339 100644 --- a/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py +++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py @@ -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, ) diff --git a/vllm_omni/diffusion/model_loader/host_weight_plan.py b/vllm_omni/diffusion/model_loader/host_weight_plan.py index ca755b082..e6301b0ec 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_plan.py +++ b/vllm_omni/diffusion/model_loader/host_weight_plan.py @@ -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, diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index fcb4075ec..6b51e5753 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -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", diff --git a/vllm_omni/diffusion/models/magi2/__init__.py b/vllm_omni/diffusion/models/magi2/__init__.py new file mode 100644 index 000000000..4c09fc3f3 --- /dev/null +++ b/vllm_omni/diffusion/models/magi2/__init__.py @@ -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) diff --git a/vllm_omni/diffusion/models/magi2/attention.py b/vllm_omni/diffusion/models/magi2/attention.py new file mode 100644 index 000000000..ca052c786 --- /dev/null +++ b/vllm_omni/diffusion/models/magi2/attention.py @@ -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", +] diff --git a/vllm_omni/diffusion/models/magi2/audio_decoder.py b/vllm_omni/diffusion/models/magi2/audio_decoder.py new file mode 100644 index 000000000..cc80f8b36 --- /dev/null +++ b/vllm_omni/diffusion/models/magi2/audio_decoder.py @@ -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", +] diff --git a/vllm_omni/diffusion/models/magi2/conditioning.py b/vllm_omni/diffusion/models/magi2/conditioning.py new file mode 100644 index 000000000..00849cf75 --- /dev/null +++ b/vllm_omni/diffusion/models/magi2/conditioning.py @@ -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 ``
`` 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", +] diff --git a/vllm_omni/diffusion/models/magi2/configuration_magi2.py b/vllm_omni/diffusion/models/magi2/configuration_magi2.py new file mode 100644 index 000000000..84b934064 --- /dev/null +++ b/vllm_omni/diffusion/models/magi2/configuration_magi2.py @@ -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() diff --git a/vllm_omni/diffusion/models/magi2/layers.py b/vllm_omni/diffusion/models/magi2/layers.py new file mode 100644 index 000000000..3fe77ab69 --- /dev/null +++ b/vllm_omni/diffusion/models/magi2/layers.py @@ -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", +] diff --git a/vllm_omni/diffusion/models/magi2/mh_moe.py b/vllm_omni/diffusion/models/magi2/mh_moe.py new file mode 100644 index 000000000..6419cfae8 --- /dev/null +++ b/vllm_omni/diffusion/models/magi2/mh_moe.py @@ -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", +] diff --git a/vllm_omni/diffusion/models/magi2/modeling_magi2.py b/vllm_omni/diffusion/models/magi2/modeling_magi2.py new file mode 100644 index 000000000..87d3402d3 --- /dev/null +++ b/vllm_omni/diffusion/models/magi2/modeling_magi2.py @@ -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", +] diff --git a/vllm_omni/diffusion/models/magi2/parallel.py b/vllm_omni/diffusion/models/magi2/parallel.py new file mode 100644 index 000000000..caada6a2d --- /dev/null +++ b/vllm_omni/diffusion/models/magi2/parallel.py @@ -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 diff --git a/vllm_omni/diffusion/models/magi2/pipeline_magi2.py b/vllm_omni/diffusion/models/magi2/pipeline_magi2.py new file mode 100644 index 000000000..c7c480876 --- /dev/null +++ b/vllm_omni/diffusion/models/magi2/pipeline_magi2.py @@ -0,0 +1,1124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# Copyright (c) 2026 SandAI. All Rights Reserved. + +"""Native vLLM-Omni pipeline for ``sand-ai/MAGI-2-preview``. + +The model, packing, scheduler, checkpoint mapping, and decoder integration are +owned by vLLM-Omni. SandAI's Apache-2.0 implementation was used as the +accuracy reference, but is neither imported nor required at runtime. + +This first native integration intentionally supports the released Preview +stage (272p and 540p). The separately checkpointed 1080p refiner is not +silently routed through an external implementation. +""" + +from __future__ import annotations + +import json +import logging +import math +import os +import random +import threading +import time +from collections.abc import Iterable, Mapping +from contextlib import contextmanager +from pathlib import Path +from typing import Any, ClassVar + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn +from diffusers.utils import load_image +from diffusers.video_processor import VideoProcessor +from PIL import Image + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.interface import ( + SupportAudioOutput, + SupportImageInput, + SupportsComponentDiscovery, +) +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.offloader import OffloadPlan, PinnedModuleStager +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( + DiffusionPipelineProfilerMixin, +) +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.errors import OmniClientError +from vllm_omni.platforms import current_omni_platform + +from .audio_decoder import Magi2AudioDecoder +from .conditioning import Magi2Qwen35TextEncoder +from .configuration_magi2 import ( + MAGI2_GENERATION_CONFIG, + MAGI2_PREVIEW_CONFIG, +) +from .parallel import get_magi2_replica_group +from .preview_data_proxy import Magi2DataProxy +from .sampler_magi2 import ( + CFGConfig, + Magi2PreviewSampler, + SamplerInput, + build_magi2_preview_schedulers, +) +from .turbo_vae import Magi2TurboVAEDecoder + +logger = logging.getLogger(__name__) + +MAGI2_MODEL_ID = "sand-ai/MAGI-2-preview" +MAGI2_MODEL_REVISION = "2dea51b64db47ee5b4402d36fd90829a0c58913b" +MAGI2_AUDIO_SAMPLE_RATE = MAGI2_GENERATION_CONFIG.audio_sample_rate + +_RESOLUTION_PRESETS: dict[str, tuple[int, int]] = { + "272p": (448, 256), + "540p": (896, 512), +} +_PRESET_BY_SIZE = {size: name for name, size in _RESOLUTION_PRESETS.items()} + +_REQUIRED_CHECKPOINT_PATHS = ( + "preview/model.safetensors.index.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", +) + +DEFAULT_NEGATIVE_PROMPT = ( + "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, " + "overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly " + "drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy " + "background, three legs, many people in the background, walking backwards" +) +DEFAULT_NEGATIVE_PROMPT += ( + ", low quality, worst quality, poor quality, noise, background noise, hiss, hum, buzz, crackle, static, " + "compression artifacts, MP3 artifacts, digital clipping, distortion, muffled, muddy, unclear, echo, reverb, " + "room echo, over-reverberated, hollow sound, distant, washed out, harsh, shrill, piercing, grating, tinny, " + "thin sound, boomy, bass-heavy, flat EQ, over-compressed, abrupt cut, jarring transition, sudden silence, " + "looping artifact, music, instrumental, sirens, alarms, crowd noise, unrelated sound effects, chaotic, " + "disorganized, messy, cheap sound" +) +DEFAULT_NEGATIVE_PROMPT += ( + ", emotionless, flat delivery, deadpan, lifeless, apathetic, robotic, mechanical, monotone, flat intonation, " + "undynamic, boring, reading from a script, AI voice, synthetic, text-to-speech, TTS, insincere, fake emotion, " + "exaggerated, overly dramatic, melodramatic, cheesy, cringey, hesitant, unconfident, tired, weak voice, " + "stuttering, stammering, mumbling, slurred speech, mispronounced, bad articulation, lisp, vocal fry, creaky " + "voice, mouth clicks, lip smacks, wet mouth sounds, heavy breathing, audible inhales, plosives, p-pops, " + "coughing, clearing throat, sneezing, speaking too fast, rushed, speaking too slow, dragged out, unnatural " + "pauses, awkward silence, choppy, disconnected, multiple speakers, two voices, background talking, out of tune, " + "off-key, autotune artifacts" +) + + +class _PeakReservedMonitor: + def __init__(self, device: int) -> None: + self.device = device + self.peak_bytes = torch.accelerator.memory_reserved(device) + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + + def _run(self) -> None: + while not self._stop.is_set(): + self.peak_bytes = max( + self.peak_bytes, + torch.accelerator.memory_reserved(self.device), + ) + self._stop.wait(0.05) + + def start(self) -> None: + self._thread.start() + + def stop(self) -> None: + self._stop.set() + self._thread.join() + self.peak_bytes = max( + self.peak_bytes, + torch.accelerator.memory_reserved(self.device), + ) + + +def _config_value( + od_config: OmniDiffusionConfig, + key: str, + env_key: str, + default: Any = None, +) -> Any: + custom = od_config.custom_pipeline_args or {} + additional = od_config.additional_config or {} + if key in custom: + return custom[key] + if key in additional: + return additional[key] + return os.environ.get(env_key, default) + + +def _env_flag(value: object, *, default: bool = False) -> bool: + if value is None: + return default + if isinstance(value, bool): + return value + return str(value).strip().lower() in {"1", "true", "yes", "on"} + + +def _resolve_native_resolution(sampling: object, extra: Mapping[str, object]) -> tuple[str, int, int]: + """Resolve MAGI-2 generation geometry from the shared sampling fields.""" + resolution_value = extra.get("resolution") + resolution = None if resolution_value is None else str(resolution_value).lower() + if resolution is not None and resolution not in _RESOLUTION_PRESETS: + raise OmniClientError( + f"Unsupported native MAGI-2 resolution {resolution!r}; choose " + f"from {sorted(_RESOLUTION_PRESETS)}. The 1080p refiner is not " + "part of the Preview-only native PR." + ) + + requested_width = getattr(sampling, "width", None) + requested_height = getattr(sampling, "height", None) + if (requested_width is None) != (requested_height is None): + raise OmniClientError("MAGI-2 native width and height must be supplied together") + + requested_size: tuple[int, int] | None = None + if requested_width is not None: + try: + requested_size = (int(requested_width), int(requested_height)) + except (TypeError, ValueError) as exc: + raise OmniClientError("MAGI-2 native dimensions must be positive integers") from exc + if requested_size[0] <= 0 or requested_size[1] <= 0: + raise OmniClientError("MAGI-2 native dimensions must be positive integers") + + if resolution is None: + if requested_size is None: + resolution = "540p" + else: + resolution = _PRESET_BY_SIZE.get(requested_size) + if resolution is None: + raise OmniClientError( + f"Unsupported native MAGI-2 size {requested_size[0]}x{requested_size[1]}; " + f"choose from {sorted(_PRESET_BY_SIZE)} or use output_width/output_height for final resizing." + ) + elif requested_size is not None and requested_size != _RESOLUTION_PRESETS[resolution]: + expected_width, expected_height = _RESOLUTION_PRESETS[resolution] + raise OmniClientError( + f"MAGI-2 resolution {resolution!r} requires {expected_width}x{expected_height}; " + f"got {requested_size[0]}x{requested_size[1]}." + ) + + width, height = _RESOLUTION_PRESETS[resolution] + return resolution, width, height + + +def _resolve_checkpoint_root(model: str, revision: str | None) -> str: + path = Path(model).expanduser() + if path.is_dir(): + root = path.resolve() + else: + normalized = model.strip().rstrip("/") + prefix = "https://huggingface.co/" + if normalized.lower().startswith(prefix): + normalized = normalized[len(prefix) :].split("/tree/", 1)[0] + if normalized.lower() != MAGI2_MODEL_ID.lower(): + raise ValueError( + "MAGI-2 expects a local checkpoint directory or the official " + f"model ID {MAGI2_MODEL_ID!r}; got {model!r}." + ) + from vllm.transformers_utils.repo_utils import hf_api + + pinned_revision = revision or MAGI2_MODEL_REVISION + logger.warning( + "Resolving the pinned MAGI-2 snapshot from Hugging Face. Pre-download it for predictable startup." + ) + root = Path(hf_api().snapshot_download(repo_id=normalized, revision=pinned_revision)).resolve() + + missing = [relative for relative in _REQUIRED_CHECKPOINT_PATHS if not (root / relative).is_file()] + if missing: + raise ValueError(f"Incomplete MAGI-2 checkpoint at {root}: missing " + ", ".join(missing)) + return str(root) + + +def _validate_native_topology(od_config: OmniDiffusionConfig) -> None: + parallel = od_config.parallel_config + unsupported: list[str] = [] + for name in ( + "pipeline_parallel_size", + "ring_degree", + "allgather_degree", + "text_encoder_tp_size", + ): + value = getattr(parallel, name, 1) + if value not in (None, 1): + unsupported.append(f"{name}={value}") + if getattr(parallel, "enable_expert_parallel", False): + unsupported.append("enable_expert_parallel") + if od_config.quantization_config is not None: + unsupported.append("quantization") + if unsupported: + raise ValueError( + "MAGI-2 Preview uses Ulysses sequence parallelism and MoE-head " + "parallelism. Unsupported options: " + ", ".join(unsupported) + ) + + dp_size = int(getattr(parallel, "data_parallel_size", 1) or 1) + tp_size = int(getattr(parallel, "tensor_parallel_size", 1) or 1) + sp_size = int(getattr(parallel, "sequence_parallel_size", 1) or 1) + cfg_size = int(getattr(parallel, "cfg_parallel_size", 1) or 1) + vae_pp_size = int(getattr(parallel, "vae_patch_parallel_size", 1) or 1) + ulysses = int(getattr(parallel, "ulysses_degree", 1) or 1) + if cfg_size not in (1, 2): + raise ValueError(f"MAGI-2 has exactly two CFG branches; cfg_parallel_size must be 1 or 2, got {cfg_size}.") + if sp_size != ulysses: + raise ValueError(f"MAGI-2 requires sequence_parallel_size == ulysses_degree; got {sp_size} and {ulysses}.") + + attention_shards = tp_size * sp_size + if MAGI2_PREVIEW_CONFIG.num_heads_q % attention_shards: + raise ValueError( + "MAGI-2 query heads must divide TP x SP; got " + f"{MAGI2_PREVIEW_CONFIG.num_heads_q} heads, TP={tp_size}, SP={sp_size}." + ) + if MAGI2_PREVIEW_CONFIG.num_heads_kv % attention_shards: + raise ValueError( + "MAGI-2 KV heads must divide TP x SP; got " + f"{MAGI2_PREVIEW_CONFIG.num_heads_kv} heads, TP={tp_size}, SP={sp_size}." + ) + dense_intermediate = ( + int(MAGI2_PREVIEW_CONFIG.hidden_size * MAGI2_PREVIEW_CONFIG.intermediate_factor * 2 / 3) // 128 * 128 + ) + tp_dimensions = { + "hidden_size": MAGI2_PREVIEW_CONFIG.hidden_size, + "dense_intermediate_size": dense_intermediate, + "MoE heads": MAGI2_PREVIEW_CONFIG.moe.num_heads, + "shared_expert_intermediate_size": MAGI2_PREVIEW_CONFIG.moe.shared_expert_intermediate_size, + "modality_shared_expert_intermediate_size": (MAGI2_PREVIEW_CONFIG.moe.modality_shared_expert_intermediate_size), + } + invalid_tp_dimensions = [name for name, size in tp_dimensions.items() if size % tp_size] + if invalid_tp_dimensions: + raise ValueError(f"MAGI-2 tensor_parallel_size={tp_size} does not divide: " + ", ".join(invalid_tp_dimensions)) + + configured_world_size = dp_size * cfg_size * tp_size * sp_size + cpu_offload = bool(od_config.enable_cpu_offload) + layerwise_offload = bool(od_config.enable_layerwise_offload) + distributed_offload = bool(od_config.enable_distributed_layerwise_offload) + if cpu_offload and not layerwise_offload: + raise ValueError( + "MAGI-2 already stages its auxiliary components from CPU, while " + "the complete Preview transformer cannot fit on one qualified GPU. " + "Combine --enable-cpu-offload with --enable-layerwise-offload, or " + "use --enable-layerwise-offload alone." + ) + if layerwise_offload and distributed_offload: + raise ValueError("MAGI-2 ordinary and distributed layerwise offload are mutually exclusive") + if layerwise_offload and configured_world_size != 1: + raise ValueError( + "MAGI-2 ordinary layerwise offload is a single-worker path; " + f"got world_size={configured_world_size}. Use distributed layerwise " + "offload for multi-worker layouts." + ) + if layerwise_offload and getattr(parallel, "use_hsdp", False): + raise ValueError("MAGI-2 ordinary layerwise offload cannot be combined with HSDP") + if dist.is_available() and dist.is_initialized() and dist.get_world_size() != configured_world_size: + raise ValueError( + "MAGI-2 parallel dimensions do not cover the worker world; got " + f"world_size={dist.get_world_size()}, DP={dp_size}, CFG={cfg_size}, TP={tp_size}, SP={sp_size}." + ) + + if vae_pp_size not in (1, configured_world_size): + raise ValueError( + "MAGI-2 VAE patch parallelism currently uses the complete DiT group; " + f"expected vae_patch_parallel_size=1 or {configured_world_size}, got {vae_pp_size}." + ) + + dlo_allgather = bool(getattr(od_config, "dlo_use_allgather", True)) + if cfg_size > 1 and dp_size > 1: + raise ValueError("MAGI-2 CFG parallelism is not yet combined with DLO data parallelism") + if distributed_offload and getattr(parallel, "use_hsdp", False): + raise ValueError("MAGI-2 HSDP and distributed layerwise offload are alternative transformer memory modes") + if dp_size > 1 and not distributed_offload: + raise ValueError("MAGI-2 data parallelism currently requires distributed layerwise offload") + if dp_size > 1 and tp_size > 1: + raise ValueError("MAGI-2 DLO data-parallel replicas currently require tensor_parallel_size=1") + if distributed_offload and dlo_allgather: + if dp_size <= 1: + raise ValueError( + "MAGI-2 DLO AllGather requires data_parallel_size > 1. SP ranks " + "own different MoE-head shards; use --dlo-no-use-allgather for SP-only DLO." + ) + if tp_size > 1: + raise ValueError("MAGI-2 DLO AllGather does not support tensor_parallel_size > 1") + + allow_unsupported = _env_flag( + _config_value( + od_config, + "magi2_allow_unsupported_topology", + "MAGI2_ALLOW_UNSUPPORTED_TOPOLOGY", + ) + ) + supported_single_worker = layerwise_offload and configured_world_size == 1 + if configured_world_size not in {4, 8} and not supported_single_worker and not allow_unsupported: + raise ValueError( + "MAGI-2 Preview is qualified for one worker with ordinary layerwise " + "offload, or four/eight workers; got " + f"DP={dp_size}, CFG={cfg_size}, TP={tp_size}, SP={sp_size} " + f"(world_size={configured_world_size}). Set " + "MAGI2_ALLOW_UNSUPPORTED_TOPOLOGY=1 only for controlled bring-up." + ) + + +def _seed_request(seed: int) -> None: + random.seed(seed) + np.random.seed(seed % (2**32)) + torch.manual_seed(seed) + + +def _resolve_request_seed(sampling: object) -> int: + seed = getattr(sampling, "seed", None) + if seed is not None: + return int(seed) + generator = getattr(sampling, "generator", None) + if isinstance(generator, list): + generator = generator[0] if generator else None + if generator is not None and hasattr(generator, "initial_seed"): + return int(generator.initial_seed()) + return 42 + + +def _single_image(value: object) -> str | Image.Image | None: + if value is None: + return None + if isinstance(value, list | tuple): + if len(value) != 1: + raise OmniClientError(f"MAGI-2 accepts at most one input image, got {len(value)}") + value = value[0] + if isinstance(value, os.PathLike): + return os.fspath(value) + if isinstance(value, str | Image.Image): + return value + raise OmniClientError( + f"MAGI-2 image input must be a PIL image, path, or one-element image list; got {type(value).__name__}." + ) + + +def _resizepad(image: Image.Image, height: int, width: int) -> Image.Image: + if image.width <= 0 or image.height <= 0: + raise ValueError(f"Invalid input image size {image.size}") + scale = min(width / image.width, height / image.height) + target_width = max(1, int(round(image.width * scale))) + target_height = max(1, int(round(image.height * scale))) + resized = image.convert("RGB").resize( + (target_width, target_height), + resample=Image.Resampling.LANCZOS, + ) + canvas = Image.new("RGB", (width, height), (255, 255, 255)) + canvas.paste( + resized, + ((width - target_width) // 2, (height - target_height) // 2), + ) + return canvas + + +def _resize_video(video: np.ndarray, width: int, height: int) -> np.ndarray: + frames = torch.from_numpy(video) + if frames.ndim != 4: + raise ValueError(f"Expected video shape (T,H,W,C), got {tuple(frames.shape)}") + was_uint8 = frames.dtype == torch.uint8 + frames = frames.float().div_(255.0) if was_uint8 else frames.float() + frames = torch.nn.functional.interpolate( + frames.permute(3, 0, 1, 2), + size=(height, width), + mode="bilinear", + align_corners=False, + ).permute(1, 2, 3, 0) + if was_uint8: + return frames.clamp_(0, 1).mul_(255).round_().byte().numpy() + return frames.numpy() + + +def _magi2_post_process(output: object) -> object: + return dict(output) if isinstance(output, Mapping) else output + + +def get_magi2_post_process_func(_od_config: OmniDiffusionConfig): + return _magi2_post_process + + +class _Magi2StagedComponent(nn.Module): + """Keep one immutable CPU copy and materialize a component per stage. + + MAGI-2's text encoder and three codec components are mutually exclusive at + runtime. Giving each one an explicit lifecycle lets both layerwise + offload backends leave them on the host at startup instead of eagerly + co-residing roughly 60 GiB of auxiliary weights on the output rank. + """ + + def __init__( + self, + module: nn.Module, + device: str | torch.device, + *, + pin_memory: bool, + ) -> None: + super().__init__() + self.module = module + self._stager = PinnedModuleStager( + module, + torch.device(device), + pin_memory=pin_memory, + ) + + def load_to_device(self) -> None: + self._stager.load() + + def offload_to_cpu(self) -> None: + self._stager.offload() + + def to(self, *_args: object, **_kwargs: object) -> _Magi2StagedComponent: + """Keep lifecycle-managed auxiliaries on the host between stages. + + The generic HSDP loader places discovered non-DiT components on the + execution device after sharding. MAGI-2 owns those components through + ``PinnedModuleStager`` instead, so a direct placement request must not + make Qwen and all three codecs resident together. + """ + + return self + + def forward(self, *args: object, **kwargs: object) -> object: + return self.module(*args, **kwargs) + + +class Magi2Pipeline( + nn.Module, + ProgressBarMixin, + SupportsComponentDiscovery, + DiffusionPipelineProfilerMixin, + SupportImageInput, + SupportAudioOutput, +): + """Native MAGI-2 Preview text/image-to-video-and-audio pipeline. + + One pipeline instance is supported per worker process. Initialization sets + process-wide deterministic state, so tests or deployments that need more + than one pipeline must isolate them in separate worker processes. + """ + + support_image_input: ClassVar[bool] = True + support_audio_output: ClassVar[bool] = True + audio_sample_rate: ClassVar[int] = MAGI2_AUDIO_SAMPLE_RATE + dummy_run_num_frames: ClassVar[int] = 0 + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = [ + "image_vae", + "video_decoder", + "audio_decoder", + ] + _offload_plan: ClassVar[OffloadPlan] = OffloadPlan( + on_demand_component_paths=frozenset( + { + "text_encoder", + "image_vae", + "video_decoder", + "audio_decoder", + } + ) + ) + + @staticmethod + def _remap_ckpt_key(checkpoint_key: str) -> str | None: + """Map released Preview keys to the native pipeline namespace. + + The DLO host-weight planner looks this hook up by name when it builds + a direct mmap plan; the ordinary loader uses ``load_weights`` below. + """ + + checkpoint_key = checkpoint_key.removeprefix("transformer.") + if checkpoint_key.startswith(("block.", "pre_adapter.", "post_adapter.")): + return f"transformer.{checkpoint_key}" + return None + + def __init__(self, od_config: OmniDiffusionConfig, **kwargs: object) -> None: + if kwargs: + names = ", ".join(sorted(kwargs)) + raise TypeError(f"Unexpected MAGI-2 pipeline initialization argument(s): {names}") + super().__init__() + if not od_config.model: + raise ValueError("MAGI-2 requires od_config.model") + _validate_native_topology(od_config) + if not current_omni_platform.is_cuda() or not current_omni_platform.is_available(): + raise RuntimeError("MAGI-2 Preview requires CUDA GPUs") + + self.od_config = od_config + self.dtype = od_config.dtype or torch.bfloat16 + self.device_str = f"cuda:{torch.accelerator.current_device_index()}" + self.checkpoint_root = _resolve_checkpoint_root( + str(od_config.model), + od_config.revision, + ) + self.deterministic = _env_flag( + _config_value( + od_config, + "magi2_deterministic", + "MAGI2_DETERMINISTIC", + ), + default=bool(getattr(od_config, "fa_deterministic", False)), + ) + os.environ["MAGI2_DETERMINISTIC"] = str(int(self.deterministic)) + # Each diffusion worker owns one startup-fixed pipeline. Configure the + # process-wide PyTorch mode once instead of mutating it per request; + # warn-only avoids rejecting otherwise valid third-party CUDA kernels. + torch.use_deterministic_algorithms(self.deterministic, warn_only=True) + self._parallel_group = get_magi2_replica_group( + int(getattr(od_config.parallel_config, "data_parallel_size", 1) or 1) + ) + self._is_output_rank = self._parallel_group.rank == 0 + self._offload_aux_after_use = True + self._transformer_is_layerwise_offloaded = bool( + od_config.enable_layerwise_offload or od_config.enable_distributed_layerwise_offload + ) + self._transformer_is_hsdp = bool(getattr(od_config.parallel_config, "use_hsdp", False)) + self._distributed_video_decode = int(od_config.parallel_config.vae_patch_parallel_size) > 1 + + # Importing here keeps config-only model detection light. The class is + # an in-tree implementation; there is no dynamic remote-code import. + from .modeling_magi2 import Magi2PreviewTransformer + + MAGI2_PREVIEW_CONFIG.validate() + mmap_dlo = bool( + od_config.enable_distributed_layerwise_offload and getattr(od_config, "dlo_use_allgather", True) + ) + if mmap_dlo: + # AllGather DLO binds checkpoint tensors as mmap views and copies + # only this rank's orthogonal DP shard. Constructing the 212-GiB + # Preview transformer on CPU first would defeat that memory model. + with torch.device("meta"): + self.transformer = Magi2PreviewTransformer(MAGI2_PREVIEW_CONFIG) + # The mmap path is inference-only. Removing autograd ownership here + # lets the shared DLO backend shard views without a MAGI-specific + # detach branch. + self.transformer.requires_grad_(False) + else: + self.transformer = Magi2PreviewTransformer(MAGI2_PREVIEW_CONFIG) + self.data_proxy = Magi2DataProxy() + self.sampler = Magi2PreviewSampler( + self.transformer, + self.data_proxy, + device=self.device_str, + dtype=self.dtype, + ) + + root = Path(self.checkpoint_root) + # Auxiliaries start on CPU and are staged one at a time. Distributed + # TurboVAE decode is the only case that materializes a codec on every + # rank; Qwen, image VAE, and Oobleck remain output-rank-only. + with torch.device("cpu"): + if self._is_output_rank: + text_encoder: nn.Module = Magi2Qwen35TextEncoder( + str(root / "text_encoder"), + dtype=self.dtype, + ) + from vllm_omni.diffusion.models.lance.wan_vae import LanceWanVAE + + image_vae: nn.Module = LanceWanVAE( + str(root / "vae" / "Wan2.2_VAE.pth"), + dtype=torch.float32, + device="cpu", + ) + audio_decoder: nn.Module = Magi2AudioDecoder( + root / "stable-audio-open-1.0", + device="cpu", + dtype=torch.float32, + ) + else: + text_encoder = nn.Identity() + image_vae = nn.Identity() + audio_decoder = nn.Identity() + + if self._is_output_rank or self._distributed_video_decode: + video_decoder: nn.Module = Magi2TurboVAEDecoder( + root / "turbo_vae" / "TurboV3-Wan22-TinyShallow_7_7.json", + root / "turbo_vae" / "checkpoint.ckpt", + device="cpu", + dtype=self.dtype, + ) + else: + video_decoder = nn.Identity() + + pin_memory = bool(getattr(od_config, "pin_cpu_memory", True)) + self.text_encoder = _Magi2StagedComponent( + text_encoder, + self.device_str, + pin_memory=pin_memory, + ) + self.image_vae = _Magi2StagedComponent( + image_vae, + self.device_str, + pin_memory=pin_memory, + ) + self.video_decoder = _Magi2StagedComponent( + video_decoder, + self.device_str, + pin_memory=pin_memory, + ) + self.audio_decoder = _Magi2StagedComponent( + audio_decoder, + self.device_str, + pin_memory=pin_memory, + ) + + self.video_processor = VideoProcessor(vae_scale_factor=2 * MAGI2_GENERATION_CONFIG.video_vae_stride[1]) + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=self.checkpoint_root, + subfolder="preview", + revision=None, + prefix="transformer.", + fall_back_to_pt=False, + ) + ] + self.setup_diffusion_pipeline_profiler( + profiler_targets=[ + "_encode_prompts", + "_encode_reference_image", + "_pool_figure_token", + "sampler.sample", + "_decode_video", + "_decode_audio", + ], + enable_diffusion_pipeline_profiler=bool(getattr(od_config, "enable_diffusion_pipeline_profiler", False)), + ) + + @property + def vae(self) -> nn.Module: + """Expose TurboVAE through the shared distributed-VAE contract.""" + + return self.video_decoder.module + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> set[str]: + prefix = "transformer." + + def stripped_weights(): + for name, tensor in weights: + if not name.startswith(prefix): + raise ValueError(f"Unexpected MAGI-2 checkpoint key {name!r}") + yield name[len(prefix) :], tensor + + if not hasattr(self.transformer, "load_weights"): + raise RuntimeError("Native MAGI-2 transformer has no strict weight loader") + loaded = self.transformer.load_weights(stripped_weights()) + return {prefix + name for name in loaded} + + @contextmanager + def _component_on_device(self, component: _Magi2StagedComponent): + component.load_to_device() + try: + yield component.module + finally: + if self._offload_aux_after_use: + component.offload_to_cpu() + + @contextmanager + def _conditioning_memory_window(self): + """Make room for the 27B text encoder on the output rank. + + The text encoder and a resident Preview shard may not fit together on + the qualified devices. Layerwise offloading already keeps the DiT on + the host; otherwise stage the output rank's shard to CPU for the two + prompt encodes and restore it before any model collective begins. + """ + + stage_transformer = ( + self._is_output_rank + and not self._transformer_is_layerwise_offloaded + and not getattr(self, "_transformer_is_hsdp", False) + ) + if stage_transformer: + self.transformer.to("cpu") + torch.accelerator.empty_cache() + try: + yield + finally: + if stage_transformer: + self.transformer.to(self.device_str) + + def _broadcast_tensor(self, tensor: torch.Tensor | None) -> torch.Tensor: + group = self._parallel_group + if group.world_size == 1: + if tensor is None: + raise RuntimeError("source rank did not provide a tensor") + return tensor.to(self.device_str) + + source_global_rank = dist.get_global_rank(group.group, 0) + metadata: list[tuple[tuple[int, ...], torch.dtype] | None] = [ + (tuple(tensor.shape), tensor.dtype) if group.rank == 0 else None + ] + dist.broadcast_object_list( + metadata, + src=source_global_rank, + group=group.group, + device=torch.device(self.device_str), + ) + resolved_metadata = metadata[0] + if resolved_metadata is None: + raise RuntimeError("source rank did not broadcast tensor metadata") + shape, dtype = resolved_metadata + if group.rank == 0: + assert tensor is not None + output = tensor.to(device=self.device_str, dtype=dtype).contiguous() + else: + output = torch.empty(shape, device=self.device_str, dtype=dtype) + dist.broadcast(output, src=source_global_rank, group=group.group) + return output + + def _encode_prompts(self, prompts: tuple[str, ...]) -> tuple[torch.Tensor, ...]: + encoded: list[torch.Tensor | None] = [None] * len(prompts) + if self._is_output_rank: + with self._component_on_device(self.text_encoder) as text_encoder: + # Call ``encode`` directly so framework on-demand hooks do not + # evict the encoder between the positive and negative prompt. + encoded = [text_encoder.encode(prompt).to(self.dtype) for prompt in prompts] + return tuple(self._broadcast_tensor(value) for value in encoded) + + def encode_prompt(self, prompt: str) -> torch.Tensor: + return self._encode_prompts((prompt,))[0] + + @staticmethod + def _ensure_figure_token(prompt: str) -> str: + try: + value = json.loads(prompt) + if not isinstance(value, dict): + raise ValueError + value["reference_layer"] = ["The first frame refers to
"] + return json.dumps(value, ensure_ascii=False) + except (json.JSONDecodeError, TypeError, ValueError): + return prompt + "\nreference_layer:The first frame refers to
" + + def _encode_reference_image( + self, + image: str | Image.Image, + height: int, + width: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + latent: torch.Tensor | None = None + if self._is_output_rank: + pil_image = load_image(image) + max_length = max(height, width) + if pil_image.width > pil_image.height: + target_width = max_length + target_height = int(pil_image.height * max_length / pil_image.width) + else: + target_height = max_length + target_width = int(pil_image.width * max_length / pil_image.height) + padded = _resizepad(pil_image, target_height, target_width) + pixels = self.video_processor.preprocess( + padded, + height=target_height, + width=target_width, + ) + pixels = pixels[:, :3].to(torch.float32).unsqueeze(2) + with self._component_on_device(self.image_vae) as image_vae: + latent = image_vae.encode_video( + pixels.to(self.device_str), + use_sample=False, + ) + latent = self._broadcast_tensor(latent) + feature = latent.unsqueeze(1) + lengths = torch.tensor( + [[[latent.shape[-2], latent.shape[-1]]]], + dtype=torch.long, + device=latent.device, + ) + return feature, lengths + + def _pool_figure_token( + self, + prompt: str, + context: torch.Tensor, + ) -> torch.Tensor: + special: torch.Tensor | None = None + if self._is_output_rank: + # This helper only tokenizes the prompt and pools rows that are + # already present in ``context``; it does not read encoder weights. + # Keep it outside the staged encoder window so CPU-only pooling + # cannot trigger an unnecessary device transfer. + special = self.text_encoder.module.pool_figure_tokens( + prompt, + ["
"], + context, + ).unsqueeze(0) + return self._broadcast_tensor(special) + + def _decode_video(self, latent: torch.Tensor) -> np.ndarray | None: + distributed_video_decode = getattr(self, "_distributed_video_decode", False) + if not self._is_output_rank and not distributed_video_decode: + return None + with self._component_on_device(self.video_decoder) as video_decoder: + decoded = video_decoder.decode( + latent.to(self.device_str, dtype=self.dtype), + output_offload=True, + ) + if not self._is_output_rank: + return None + if decoded.ndim != 5: + raise RuntimeError(f"TurboVAE returned unexpected shape {tuple(decoded.shape)}") + video = decoded[0].float().mul_(0.5).add_(0.5).clamp_(0, 1) + return video.permute(1, 2, 3, 0).mul_(255).byte().cpu().numpy() + + def _decode_audio(self, latent: torch.Tensor) -> np.ndarray | None: + if not self._is_output_rank: + return None + with self._component_on_device(self.audio_decoder) as audio_decoder: + waveform = audio_decoder.decode(latent.transpose(1, 2).to(self.device_str)) + audio = waveform.squeeze(0).transpose(0, 1).float().cpu().numpy() + # MAGI emits 25 latent frames/s while Stable Audio's native latent + # rate is 44,100/2,048. Preserve the reference Fourier resampling. + from scipy.signal import resample + + target_length = int(audio.shape[0] * 441 / 512) + return resample(audio, target_length) + + @torch.inference_mode() + def _evaluate_preview( + self, + *, + prompt: str, + image: str | Image.Image | None, + width: int, + height: int, + num_inference_steps: int, + ) -> tuple[np.ndarray | None, np.ndarray | None]: + ref_image_feat = None + ref_image_feat_len = None + if image is not None: + ref_image_feat, ref_image_feat_len = self._encode_reference_image( + image, + height, + width, + ) + prompt = self._ensure_figure_token(prompt) + + negative_prompt = os.environ.get( + "MAGI2_NEGATIVE_PROMPT", + DEFAULT_NEGATIVE_PROMPT, + ) + with self._conditioning_memory_window(): + context, context_null = self._encode_prompts((prompt, negative_prompt)) + special = self._pool_figure_token(prompt, context) if ref_image_feat is not None else None + + latent_height = height // MAGI2_GENERATION_CONFIG.video_vae_stride[1] + latent_width = width // MAGI2_GENERATION_CONFIG.video_vae_stride[2] + public_frame_length = int(round(MAGI2_GENERATION_CONFIG.duration_seconds * MAGI2_GENERATION_CONFIG.fps)) + if public_frame_length != MAGI2_GENERATION_CONFIG.output_frames: + raise RuntimeError( + "MAGI-2 generation config is inconsistent: duration_seconds * fps must equal output_frames" + ) + # The Preview DiT/decoder operates on the internal 25-FPS timeline, + # while the public output contract is 12.5 FPS. The internal timeline + # is therefore twice as long; the decoded result is muxed at 12.5 FPS. + video_frame_length = public_frame_length * 2 + latent_length = (video_frame_length - 1) // MAGI2_GENERATION_CONFIG.video_vae_stride[0] + 1 + audio_length = int(round(MAGI2_GENERATION_CONFIG.duration_seconds * MAGI2_GENERATION_CONFIG.audio_latent_fps)) + # Draw order is parity-critical: video noise precedes audio noise. + latent = torch.randn( + 1, + MAGI2_GENERATION_CONFIG.video_latent_channels, + latent_length, + latent_height, + latent_width, + device=self.device_str, + dtype=torch.float32, + ) + audio_latent = torch.randn( + 1, + audio_length, + MAGI2_GENERATION_CONFIG.audio_latent_channels, + device=self.device_str, + dtype=torch.float32, + ) + ref_audio = torch.zeros( + 1, + 0, + MAGI2_GENERATION_CONFIG.audio_latent_channels, + device=self.device_str, + dtype=torch.float32, + ) + video_scheduler, audio_scheduler = build_magi2_preview_schedulers( + num_inference_steps, + device=self.device_str, + shift=MAGI2_GENERATION_CONFIG.shift, + ) + sampler_input = SamplerInput( + video_t_list=video_scheduler.timesteps, + audio_t_list=audio_scheduler.timesteps, + latent=latent, + audio_latent=audio_latent, + txt_feat=context, + null_txt_feat=context_null, + ref_audio_feat=ref_audio, + ref_video_feat=None, + video_scheduler=video_scheduler, + audio_scheduler=audio_scheduler, + cfg_config=CFGConfig( + video_txt_guidance_scale=(MAGI2_GENERATION_CONFIG.video_guidance_scale), + audio_txt_guidance_scale=(MAGI2_GENERATION_CONFIG.audio_guidance_scale), + ), + ref_image_feat=ref_image_feat, + ref_image_feat_len=ref_image_feat_len, + ref_image_special_token_embedding=special, + ) + latent, audio_latent = self.sampler.sample(sampler_input) + video = self._decode_video(latent) + audio = self._decode_audio(audio_latent) + return video, audio + + @torch.inference_mode() + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + if req.num_reqs != 1: + raise OmniClientError("MAGI-2 currently supports one request at a time") + + raw_prompt = req.prompts[0] + image_value: object = None + if isinstance(raw_prompt, str): + prompt = raw_prompt + multimodal: Mapping[str, object] = {} + elif isinstance(raw_prompt, Mapping): + prompt = str(raw_prompt.get("prompt") or "") + multimodal = raw_prompt.get("multi_modal_data") or {} + image_value = raw_prompt.get("image_path") + else: + raise OmniClientError(f"Unsupported MAGI-2 prompt type: {type(raw_prompt).__name__}") + if not prompt: + raise OmniClientError("MAGI-2 requires a non-empty prompt") + if not isinstance(multimodal, Mapping): + raise OmniClientError("MAGI-2 multi_modal_data must be a mapping") + unexpected = set(multimodal) - {"image"} + if unexpected: + raise OmniClientError( + "MAGI-2 accepts text and at most one still image; unsupported " + "modalities: " + ", ".join(sorted(unexpected)) + ) + if "image" in multimodal: + image_value = multimodal["image"] + + sampling = req.sampling_params + num_outputs = int(getattr(sampling, "num_outputs_per_prompt", 1) or 1) + if num_outputs != 1: + raise OmniClientError(f"MAGI-2 supports exactly one output per prompt; got {num_outputs}") + extra = sampling.extra_args or {} + if image_value is None: + image_value = extra.get("image_path") + image = _single_image(image_value) + + _, width, height = _resolve_native_resolution(sampling, extra) + if _env_flag(extra.get("use_refiner")): + raise OmniClientError( + "The MAGI-2 1080p refiner is a separate model and is not yet enabled by the native Preview pipeline." + ) + seconds = float(extra.get("seconds", extra.get("duration", 10.0))) + if not math.isfinite(seconds) or seconds != 10.0: + raise OmniClientError("MAGI-2 Preview supports 10-second clips only") + + requested_fps = getattr(sampling, "resolved_frame_rate", None) + if requested_fps is None: + requested_fps = getattr(sampling, "frame_rate", None) + if requested_fps is None: + requested_fps = getattr(sampling, "fps", None) + if requested_fps is not None and float(requested_fps) != 12.5: + raise OmniClientError(f"MAGI-2 Preview output is fixed at 12.5 fps; got {requested_fps}.") + requested_frames = getattr(sampling, "num_frames", None) + # ``1`` is the shared image-model engine default and therefore means + # unset here. Serving validates an explicitly supplied frame count + # before applying MAGI-2's model-owned default. + if requested_frames not in {None, 1, 125}: + raise OmniClientError(f"MAGI-2 Preview output is fixed at 125 frames; got {requested_frames}.") + requested_steps = sampling.num_inference_steps + steps = MAGI2_GENERATION_CONFIG.preview_steps if requested_steps is None else int(requested_steps) + if steps <= 0: + raise OmniClientError("MAGI-2 inference steps must be positive") + + output_width = extra.get("output_width") + output_height = extra.get("output_height") + if (output_width is None) != (output_height is None): + raise OmniClientError("MAGI-2 output resize requires output_width and output_height") + if output_width is not None: + try: + output_width = int(output_width) + output_height = int(output_height) + except (TypeError, ValueError) as exc: + raise OmniClientError("MAGI-2 output dimensions must be positive integers") from exc + if output_width <= 0 or output_height <= 0: + raise OmniClientError("MAGI-2 output dimensions must be positive integers") + + requested_deterministic = extra.get("deterministic") + if requested_deterministic is not None and _env_flag(requested_deterministic) != self.deterministic: + raise OmniClientError( + "MAGI-2 deterministic mode is fixed at worker startup; restart " + "with MAGI2_DETERMINISTIC set to the requested value." + ) + seed = _resolve_request_seed(sampling) + _seed_request(seed) + + has_cuda = current_omni_platform.is_cuda() and current_omni_platform.is_available() + device_index = torch.accelerator.current_device_index() if has_cuda else None + # Sampling reserved memory is qualification instrumentation, not part + # of ordinary serving. It starts only when the pipeline profiler is + # explicitly enabled, so every CUDA request avoids a 20 Hz thread. + monitor = ( + _PeakReservedMonitor(device_index) + if device_index is not None and getattr(self, "enable_diffusion_pipeline_profiler", False) + else None + ) + monitor_started = False + started = time.perf_counter() + try: + if monitor is not None: + monitor.start() + monitor_started = True + torch.accelerator.synchronize() + started = time.perf_counter() + video, audio = self._evaluate_preview( + prompt=prompt, + image=image, + width=width, + height=height, + num_inference_steps=steps, + ) + if has_cuda: + torch.accelerator.synchronize() + finally: + if monitor_started: + monitor.stop() + elapsed = time.perf_counter() - started + + if video is not None and output_width is not None: + video = _resize_video(video, output_width, output_height) + + peak_memory_mb = monitor.peak_bytes / 1024**2 if monitor is not None else 0.0 + if has_cuda and dist.is_available() and dist.is_initialized() and self._parallel_group.world_size > 1: + peak = torch.tensor(peak_memory_mb, device=self.device_str) + dist.all_reduce( + peak, + op=dist.ReduceOp.MAX, + group=self._parallel_group.group, + ) + peak_memory_mb = float(peak.item()) + + stage_durations = {"magi2_preview_e2e": elapsed} + if getattr(self, "enable_diffusion_pipeline_profiler", False): + stage_durations.update(self.stage_durations) + + return DiffusionOutput( + output={ + "payload": {"video": video, "audio": audio}, + "metadata": { + "video": {"fps": 12.5}, + "audio": {"sample_rate": MAGI2_AUDIO_SAMPLE_RATE}, + }, + }, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ) diff --git a/vllm_omni/diffusion/models/magi2/preview_data_proxy.py b/vllm_omni/diffusion/models/magi2/preview_data_proxy.py new file mode 100644 index 000000000..13ed1766e --- /dev/null +++ b/vllm_omni/diffusion/models/magi2/preview_data_proxy.py @@ -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", +] diff --git a/vllm_omni/diffusion/models/magi2/sampler_magi2.py b/vllm_omni/diffusion/models/magi2/sampler_magi2.py new file mode 100644 index 000000000..994fdd3e4 --- /dev/null +++ b/vllm_omni/diffusion/models/magi2/sampler_magi2.py @@ -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", +] diff --git a/vllm_omni/diffusion/models/magi2/turbo_vae.py b/vllm_omni/diffusion/models/magi2/turbo_vae.py new file mode 100644 index 000000000..e600811e1 --- /dev/null +++ b/vllm_omni/diffusion/models/magi2/turbo_vae.py @@ -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", +] diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 297177052..46ddaad1d 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -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", diff --git a/vllm_omni/diffusion/utils/hf_utils.py b/vllm_omni/diffusion/utils/hf_utils.py index 5c41685f0..9f5e6f06f 100644 --- a/vllm_omni/diffusion/utils/hf_utils.py +++ b/vllm_omni/diffusion/utils/hf_utils.py @@ -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: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 5f190d8d7..80e438ee4 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -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), diff --git a/vllm_omni/entrypoints/openai/protocol/videos.py b/vllm_omni/entrypoints/openai/protocol/videos.py index f71e9feed..fb4c395d7 100644 --- a/vllm_omni/entrypoints/openai/protocol/videos.py +++ b/vllm_omni/entrypoints/openai/protocol/videos.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 398c7b81f..f3f4b7570 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 34bc23384..09a3d284d 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -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, diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index adc6cbabe..4b91f752a 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -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( diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 6ccc1099a..3ef934bf5 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -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 diff --git a/vllm_omni/model_extras/__init__.py b/vllm_omni/model_extras/__init__.py index f10ca9806..143a9b224 100644 --- a/vllm_omni/model_extras/__init__.py +++ b/vllm_omni/model_extras/__init__.py @@ -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", diff --git a/vllm_omni/model_extras/magi2.py b/vllm_omni/model_extras/magi2.py new file mode 100644 index 000000000..360a0228d --- /dev/null +++ b/vllm_omni/model_extras/magi2.py @@ -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, + ) diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index a23dd016c..4d2ad010a 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -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. diff --git a/vllm_omni/model_extras/video_generation.py b/vllm_omni/model_extras/video_generation.py new file mode 100644 index 000000000..86d116266 --- /dev/null +++ b/vllm_omni/model_extras/video_generation.py @@ -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