[skip ci][Doc] Remove obsolete DiffusionParallelConfig in many example scripts (#6347)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/models/supported_models.md

@@ -39,6 +39,7 @@ th {
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
 | `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `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` | ✅︎ |   |   |   | — |

docs/user_guide/diffusion/parallelism/hsdp.md

@@ -152,7 +152,8 @@ In `DiffusionParallelConfig`:
 **Constraints:**
 
 - `hsdp_replicate_size × hsdp_shard_size == world_size`
-- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
+- HSDP cannot be used with Tensor Parallelism, ordinary Data Parallelism,
+  Pipeline Parallelism, or Expert Parallelism.
 
 ---
 

docs/user_guide/diffusion_features.md

@@ -168,6 +168,14 @@ The following tables show which models support each feature:
 | **Cosmos3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **LongCat-Video-Avatar-1.5** |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **MiniMax-H3**               | ✅ (FL2VA) |     ✅      |           ✅           |       ❌        |       ✅ (DiT/TE)  |         ❌         |   ✅    |             ✅             |       ✅ (tile)       |      ✅ (DiT)      |        ❌         |
+| **SANA-WM**                  |     ❌     |     ❌      |          ❌<sup>5</sup> |       ✅        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+
+> Notes:
+> 5. SANA-WM cannot support sequence parallelism: its bidirectional gated delta
+>    recurrence carries state across frames, so a rank cannot denoise a slice of
+>    the token sequence in isolation. Doing so would need a distributed scan or
+>    an all-gather before every GDN block. The remaining ❌ columns are simply
+>    unvalidated on this model, not known-broken.
 
 > **Step execution note:** Helios supports single-request step execution only;
 > use `max_num_seqs=1`.
@@ -175,7 +183,7 @@ The following tables show which models support each feature:
 **Frame Interpolation Support**
 
 - **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
-- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni
+- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni, SANA-WM
 
 ### AudioGen
 

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -53,7 +53,6 @@ from typing import Any
 import torch
 from PIL import Image
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -139,12 +138,6 @@ async def main():
 
     # ---- Torch setup ----
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
-    parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree,
-        ring_degree=args.ring_degree,
-        cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
-    )
 
     # ---- Cache Config ----
     if args.cache_backend == "cache_dit":
@@ -171,7 +164,10 @@ async def main():
         vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
-        parallel_config=parallel_config,
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
         diffusion_load_format="dummy",

examples/offline_inference/diffusion/lingbot_world_v2.py

@@ -154,8 +154,6 @@ def _positive_finite(value: float, flag: str) -> float:
 def build_omni_kwargs(
     args: argparse.Namespace,
     paths: LingBotPaths,
-    *,
-    parallel_config: Any,
 ) -> dict[str, Any]:
     """Build Omni configuration without overriding checkpoint class discovery."""
 
@@ -167,7 +165,7 @@ def build_omni_kwargs(
     return {
         "model": model,
         "flow_shift": flow_shift,
-        "parallel_config": parallel_config,
+        "tensor_parallel_size": args.tensor_parallel_size,
         "enforce_eager": args.enforce_eager,
         "model_config": {"lingbot_action_root": str(paths.action_root)},
     }
@@ -262,12 +260,10 @@ def main(argv: Sequence[str] | None = None) -> Path:
     # helper tests work on machines without the native vLLM runtime.
     from diffusers.utils import export_to_video
 
-    from vllm_omni.diffusion.data import DiffusionParallelConfig
     from vllm_omni.entrypoints.omni import Omni
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-    parallel_config = DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size)
+    omni_kwargs = build_omni_kwargs(args, paths)
-    omni_kwargs = build_omni_kwargs(args, paths, parallel_config=parallel_config)
     omni = Omni(**omni_kwargs)
     try:
         outputs = omni.generate(

examples/offline_inference/diffusion/lingbot_world_v2_realtime.py

@@ -125,7 +125,6 @@ async def run(argv: Sequence[str] | None = None) -> Path:
     # require a CUDA-enabled vLLM installation.
     import torch
 
-    from vllm_omni.diffusion.data import DiffusionParallelConfig
     from vllm_omni.diffusion.models.lingbot_world.actions import LingBotCameraControlReducer
     from vllm_omni.entrypoints.async_omni import AsyncOmni
     from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
@@ -141,7 +140,7 @@ async def run(argv: Sequence[str] | None = None) -> Path:
         model=args.model,
         engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine",
         enforce_eager=args.enforce_eager,
-        parallel_config=DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size),
+        tensor_parallel_size=args.tensor_parallel_size,
         max_num_seqs=1,
         model_config={
             "ar_diffusion_height": args.height,

examples/offline_inference/image_to_image/image_edit.py

@@ -96,7 +96,6 @@ from typing import Any
 import torch
 from PIL import Image
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
@@ -452,14 +451,6 @@ def main():
 
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
-    parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree,
-        ring_degree=args.ring_degree,
-        cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
-        enable_expert_parallel=args.enable_expert_parallel,
-    )
-
     # Configure cache based on backend type
     cache_config = None
     if args.cache_backend == "cache_dit":
@@ -490,7 +481,11 @@ def main():
         vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
-        parallel_config=parallel_config,
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+        enable_expert_parallel=args.enable_expert_parallel,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         profiler_config=args.profiler_config,

examples/offline_inference/image_to_video/image_to_video.py

@@ -56,7 +56,6 @@ import numpy as np
 import PIL.Image
 import torch
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 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
@@ -503,17 +502,6 @@ def main():
         }
 
     profiler_enabled = args.profiler_config is not None
-    parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree,
-        ring_degree=args.ring_degree,
-        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,
-    )
     omni_kwargs = dict(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
@@ -524,7 +512,15 @@ def main():
         diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
         diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers,
         enable_cpu_offload=args.enable_cpu_offload,
-        parallel_config=parallel_config,
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        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,
         enforce_eager=args.enforce_eager,
         model_class_name=model_class_name,
         cache_backend=args.cache_backend,

examples/offline_inference/longcat_video/end2end.py

@@ -15,7 +15,6 @@ import imageio
 import numpy as np
 import torch
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.models.longcat_video.pipeline_longcat_video_avatar import (
     prepare_longcat_video_avatar_model_for_omni,
 )
@@ -236,14 +235,6 @@ def main():
         dtype="bfloat16",
         enforce_eager=args.enforce_eager,
         additional_config=additional_config,
-        parallel_config=DiffusionParallelConfig(
-            ulysses_degree=1,
-            ring_degree=1,
-            cfg_parallel_size=1,
-            tensor_parallel_size=1,
-            vae_patch_parallel_size=1,
-            pipeline_parallel_size=1,
-        ),
     )
 
     audio_input: str | dict[str, str] | None = args.audio

examples/offline_inference/speech_to_video/speech_to_video.py

@@ -23,7 +23,6 @@ from typing import Any
 
 import torch
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -250,17 +249,6 @@ def main():
             "scm_steps_policy": "dynamic",
         }
 
-    parallel_config = DiffusionParallelConfig(
-        tensor_parallel_size=args.tensor_parallel_size,
-        cfg_parallel_size=args.cfg_parallel_size,
-        vae_patch_parallel_size=args.vae_patch_parallel_size,
-        ring_degree=args.ring_degree,
-        ulysses_degree=args.ulysses_degree,
-        use_hsdp=args.use_hsdp,
-        hsdp_shard_size=args.hsdp_shard_size,
-        hsdp_replicate_size=args.hsdp_replicate_size,
-    )
-
     profiler_enabled = args.profiler_config is not None
 
     omni_kwargs = dict(
@@ -271,7 +259,14 @@ def main():
         vae_use_tiling=args.vae_use_tiling,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_layerwise_offload=args.enable_layerwise_offload,
-        parallel_config=parallel_config,
+        tensor_parallel_size=args.tensor_parallel_size,
+        cfg_parallel_size=args.cfg_parallel_size,
+        vae_patch_parallel_size=args.vae_patch_parallel_size,
+        ring_degree=args.ring_degree,
+        ulysses_degree=args.ulysses_degree,
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
         enforce_eager=args.enforce_eager,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         cache_backend=args.cache_backend,

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -38,7 +38,6 @@ from typing import Any
 import numpy as np
 import torch
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 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
@@ -306,17 +305,13 @@ def main():
     print(f"  Seed: {args.seed}")
     print(f"{'=' * 60}\n")
 
-    parallel_config = DiffusionParallelConfig(
-        use_hsdp=args.use_hsdp,
-        hsdp_shard_size=args.hsdp_shard_size,
-        hsdp_replicate_size=args.hsdp_replicate_size,
-    )
-
     # Initialize Omni. AudioX requires an explicit pipeline class; Stable Audio
     # is resolved from its config.
     omni_kwargs: dict[str, Any] = dict(
         model=args.model,
-        parallel_config=parallel_config,
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,

examples/offline_inference/text_to_speech/ming_tts/end2end.py

@@ -13,6 +13,7 @@ import yaml
 from transformers import AutoTokenizer
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor
 from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
     KEY_CFG,
     KEY_MAX_DECODE_STEPS,
@@ -21,7 +22,6 @@ from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
     SAMPLE_RATE,
 )
 from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt
-from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor
 
 try:
     from .runner import (
@@ -142,7 +142,8 @@ def _resolve_speaker_embedding(args, case: dict, paths: list[str]):
         return _load_speaker_embedding(args.speaker_embedding)
     if not (case.get("auto_extract_speaker_embeddings") or args.extract_speaker_embeddings) or not paths:
         return None
-    embs = MingSpeakerEmbeddingExtractor(args.model).extract_many(paths)
+    # Runs before the engine is built, so nothing has pre-fetched campplus.onnx.
+    embs = SpeakerEmbeddingExtractor(args.model, allow_download=True).extract_many(paths)
     if not embs:
         raise RuntimeError("Speaker extraction produced no embeddings")
     return embs[0] if len(embs) == 1 else torch.stack(embs, dim=0)

examples/offline_inference/text_to_video/text_to_video.py

@@ -10,7 +10,6 @@ from typing import Any
 import numpy as np
 import torch
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 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
@@ -462,17 +461,6 @@ def main():
             "scm_steps_policy": "dynamic",
         }
 
-    # Configure parallel settings
-    parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree,
-        ring_degree=args.ring_degree,
-        cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
-        vae_patch_parallel_size=args.vae_patch_parallel_size,
-        pipeline_parallel_size=args.pipeline_parallel_size,
-        enable_expert_parallel=args.enable_expert_parallel,
-    )
-
     profiler_enabled = args.profiler_config is not None
 
     omni_kwargs = dict(
@@ -481,7 +469,13 @@ def main():
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         enable_cpu_offload=args.enable_cpu_offload,
-        parallel_config=parallel_config,
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+        vae_patch_parallel_size=args.vae_patch_parallel_size,
+        pipeline_parallel_size=args.pipeline_parallel_size,
+        enable_expert_parallel=args.enable_expert_parallel,
         enforce_eager=args.enforce_eager,
         model_class_name=model_class_name,
         cache_backend=args.cache_backend,

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -11,7 +11,6 @@ import numpy as np
 from PIL import Image
 from vllm.multimodal.media.audio import load_audio
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -197,14 +196,6 @@ def main() -> None:
             },
         )
 
-    parallel_config = DiffusionParallelConfig(
-        cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
-        use_hsdp=args.use_hsdp,
-        hsdp_shard_size=args.hsdp_shard_size,
-        hsdp_replicate_size=args.hsdp_replicate_size,
-    )
-
     cache_config = None
     if args.cache_backend == "cache_dit":
         cache_config = {
@@ -222,7 +213,11 @@ def main() -> None:
 
     omni_kwargs = dict(
         model=args.model,
-        parallel_config=parallel_config,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
         model_type=args.model_type,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_layerwise_offload=args.enable_layerwise_offload,

recipes/NVIDIA/SANA-WM.md

@@ -0,0 +1,193 @@
+# SANA-WM
+
+> Camera-controllable first-frame image-to-video world model.
+
+## Summary
+
+- Vendor: Efficient-Large-Model / NVlabs SANA
+- Model: `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` (standard diffusers layout, converted offline from the NVlabs release; Stage-1 transformer + VAE only)
+- Task: First-frame image-to-video generation with camera control
+- Mode: Online serving with the OpenAI-compatible video API
+- Model weights: about 13 GB for the Stage-1 transformer (10 GB) and VAE (2.3 GB)
+- Local disk: reserve about 40 GB for the Hugging Face cache and runtime artifacts
+- Recommended GPU: 24 GB or larger CUDA GPU
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want to serve SANA-WM through `/v1/videos` or
+`/v1/videos/sync`. The model takes a text prompt, a first-frame image, and
+either an action DSL string or explicit camera poses. vLLM-Omni serves the
+SANA-WM Stage-1 DiT, decoded through the SANA VAE. The optional LTX-2 refiner
+stage is not supported by this integration; it is a planned follow-up.
+
+## References
+
+- Upstream model card: <https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional>
+- Video API: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
+
+## Hardware Support
+
+## GPU
+
+### 1x NVIDIA RTX PRO 6000 Blackwell 96GB
+
+#### Capacity
+
+- Model storage: the Stage-1 transformer is about 10 GB and the VAE about
+  2.3 GB. The Gemma text encoder is a separate 4.9 GB repo. The engine prefetches
+  the whole model repo at startup (`allow_patterns=["*"]`), which is why the
+  Stage-1 weights live in their own repo — the two-stage one carries an
+  additional 84 GB `refiner/` that this path never loads.
+- Text encoder: the pipeline tries `google/gemma-2-2b-it` first, then falls back
+  to the ungated mirror `Efficient-Large-Model/gemma-2-2b-it`. The first repo is
+  gated, so without an accepted licence and `hf auth login` you will see one
+  failed load before the fallback succeeds — that is expected, not an error to
+  chase. Set `VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER` to pin a specific repo or a
+  local path and skip both.
+- Disk sizing: provision about 40 GB of local disk or Hugging Face cache volume
+  so the model, temporary downloads, and generated artifacts fit without cache
+  eviction.
+- GPU sizing: the default 1280x704, 161-frame, 60-step serving profile peaks at
+  22.6 GB of device memory and takes about 133 s to generate on one RTX PRO
+  6000 Blackwell. The peak lands in the VAE decode, which is why the pipeline
+  forces VAE tiling on regardless of `vae_use_tiling` — without it the same
+  request costs about 9 GB more, and a 321-frame one OOMs outright. On smaller
+  GPUs, lower `width`, `height`, or `num_frames` before serving production
+  requests.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Driver / runtime: NVIDIA driver with CUDA runtime supported by your PyTorch
+  build
+- Recommended operator library: Triton, installed through the vLLM/vLLM-Omni
+  Python environment
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+The repo ships the standard Diffusers layout (`model_index.json` +
+`transformer/`, `vae/`), and its `model_index.json` names `SanaWmPipeline`, so
+the pipeline class resolves on its own.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+vllm serve BBBBruce/SANA-WM_bidirectional-stage1-diffusers \
+  --omni \
+  --host 0.0.0.0 \
+  --port 8091
+```
+
+No deploy config: single-stage diffusion models are deliberately absent from
+`OMNI_PIPELINES` (`vllm_omni/config/pipeline_registry.py`), so stage resolution
+falls back to the default stage config and a YAML's stage settings — including
+`default_sampling_params` — would not be applied. The production generation
+settings therefore live in the model (`num_inference_steps=60`,
+`guidance_scale=5.0`); a request that omits a field gets them. The examples
+below still pass every field explicitly so the numbers are visible.
+
+If you point this at the older two-stage repo
+(`BBBBruce/SANA-WM_bidirectional-diffusers`), startup fails with `Model class
+SanaWmTwoStagesPipeline not found in diffusion model registry`, because that
+repo's `model_index.json` names a class this build does not register. Add
+`--model-class-name SanaWmPipeline` to override it.
+
+#### Verification
+
+Use a short smoke request first:
+
+```bash
+curl -sS -X POST http://localhost:8091/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "prompt=A slow forward camera move through a quiet city street." \
+  -F "negative_prompt=blurry, low quality, distorted, watermark" \
+  -F "input_reference=@/path/to/first_frame.png;type=image/png" \
+  -F "width=1280" \
+  -F "height=704" \
+  -F "num_frames=9" \
+  -F "fps=16" \
+  -F "num_inference_steps=2" \
+  -F "guidance_scale=5.0" \
+  -F "seed=42" \
+  --form-string 'extra_params={"sana_wm":{"action":"w-8","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}' \
+  -o sana_wm_smoke.mp4
+```
+
+For a production-length request, note that the action durations must sum to
+`num_frames - 1` — the rollout includes the identity start pose — and a
+mismatch is rejected rather than padded or truncated:
+
+```bash
+curl -sS -X POST http://localhost:8091/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "prompt=A slow forward camera move through a quiet city street." \
+  -F "negative_prompt=blurry, low quality, distorted, watermark" \
+  -F "input_reference=@/path/to/first_frame.png;type=image/png" \
+  -F "width=1280" \
+  -F "height=704" \
+  -F "num_frames=161" \
+  -F "fps=16" \
+  -F "num_inference_steps=60" \
+  -F "guidance_scale=5.0" \
+  -F "seed=42" \
+  --form-string 'extra_params={"sana_wm":{"action":"w-160","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}' \
+  -o sana_wm_output.mp4
+```
+
+Use `POST /v1/videos` instead when you want job storage and polling rather than
+inline MP4 bytes. It accepts the same form fields as `/v1/videos/sync`.
+
+```bash
+create_response=$(curl -sS -X POST http://localhost:8091/v1/videos \
+  -H "Accept: application/json" \
+  -F "prompt=A slow forward camera move through a quiet city street." \
+  -F "negative_prompt=blurry, low quality, distorted, watermark" \
+  -F "input_reference=@/path/to/first_frame.png;type=image/png" \
+  -F "width=1280" \
+  -F "height=704" \
+  -F "num_frames=161" \
+  -F "fps=16" \
+  -F "num_inference_steps=60" \
+  -F "guidance_scale=5.0" \
+  -F "seed=42" \
+  --form-string 'extra_params={"sana_wm":{"action":"w-160","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}')
+
+video_id=$(echo "$create_response" | jq -r '.id')
+curl -sS "http://localhost:8091/v1/videos/${video_id}" | jq .
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o sana_wm_output.mp4
+```
+
+#### Notes
+
+- Sequence parallelism is not supported. The bidirectional gated delta
+  recurrence carries state across frames, so a rank cannot denoise a slice
+  of the token sequence in isolation; supporting it needs a distributed scan
+  or an all-gather before the GDN blocks.
+
+- `input_reference` is required for the first frame. Use `image_reference` only
+  when you need a JSON-safe image URL or data URL instead of a multipart upload.
+- `sana_wm` must provide exactly one of `action` or `camera`.
+- Action strings use comma-separated `<keys>-<duration>` segments. Supported
+  keys are `w`, `a`, `s`, `d` for translation and `i`, `j`, `k`, `l` for
+  pitch/yaw rotation. The durations must sum to `num_frames - 1`.
+- Explicit camera control (alternative to `action`): pass
+  `"camera": {"poses": [...]}` where `poses` is a list of `num_frames`
+  camera-to-world 4x4 matrices (row-major, OpenCV `+X right, +Y down, +Z forward`
+  convention), e.g.
+  `extra_params={"sana_wm":{"camera":{"poses":[[[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]], ...]},"intrinsics":{...}}}`.
+  Most callers should prefer `action`; explicit poses exist for callers that
+  already have a per-frame trajectory.
+- Explicit `intrinsics` are recommended and take the mapping form
+  `{"fx":640,"fy":640,"cx":640,"cy":352}` (for 1280x704). This `{fx,fy,cx,cy}`
+  mapping is the only accepted intrinsics form; omit `intrinsics` to derive them
+  from the output resolution. All four values must be finite, and `fx`/`fy`
+  must be positive — the ray map divides by them.
+- The video API returns decoded MP4 bytes and has no `output_type` field, so
+  raw Stage-1 latents are reachable only from the offline API
+  (`OmniDiffusionSamplingParams(output_type="latent")`); see
+  [`tests/e2e/offline_inference/test_sana_wm.py`](../../tests/e2e/offline_inference/test_sana_wm.py).
+  Putting `output_type` in `extra_params` does not work: it lands in
+  `sampling_params.extra_args`, while the pipeline reads the top-level field.

recipes/README.md

@@ -70,6 +70,7 @@ recipes/
 | [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
 | [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
 | [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
+| [`NVIDIA/SANA-WM.md`](./NVIDIA/SANA-WM.md) | First-frame image-to-video serving with camera control | 1x 24GB+ CUDA GPU; ~40GB disk |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |

recipes/inclusionAI/Ming-omni-tts.md

@@ -5,20 +5,21 @@
 ## Summary
 
 - Vendor: inclusionAI
-- Model: `inclusionAI/Ming-omni-tts-0.5B`
+- Model: `inclusionAI/Ming-omni-tts-16.8B-A3B`, `inclusionAI/Ming-omni-tts-0.5B`
 - Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls
 - Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API
 - Maintainer: Community
 
 ## References
 
-- [Huggingface Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
+- [Huggingface: inclusionAI/Ming-omni-tts-16.8B-A3B](https://huggingface.co/inclusionAI/Ming-omni-tts-16.8B-A3B)
+- [Huggingface: inclusionAI/Ming-omni-tts-0.5B](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
 - Upstream repository [inclusionAI/Ming-omni-tts](https://github.com/inclusionAI/Ming-omni-tts)
 
 
 ## Hardware Support
 
-This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B two-stage TTS pipeline deployment.
+This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B and the MoE 16.8B two-stage TTS pipeline deployment.
 Other hardware is welcome as community validation lands.
 
 ## CUDA
@@ -39,11 +40,7 @@ Other hardware is welcome as community validation lands.
 Launch the two-stage talker:
 
 ```bash
-vllm serve inclusionAI/Ming-omni-tts-0.5B \
+vllm serve inclusionAI/Ming-omni-tts-16.8B-A3B --omni --port 8091
-    --deploy-config vllm_omni/deploy/ming_tts.yaml \
-    --omni \
-    --port 8091 \
-    --enforce-eager
 ```
 
 #### Verification
@@ -54,7 +51,7 @@ Basic synthesis (save the WAV bytes):
 curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
-      "model": "inclusionAI/Ming-omni-tts-0.5B",
+      "model": "inclusionAI/Ming-omni-tts-16.8B-A3B",
       "input": "你好，这是 Ming 在线语音合成测试。",
       "response_format": "wav"
     }' --output ming_tts_basic.wav
@@ -74,7 +71,7 @@ Reference-audio zero-shot cloning:
 REF_AUDIO="$BASE/10002287-00000094.wav"
 jq -n --arg ref_audio "$REF_AUDIO" \
     '{
-       model: "inclusionAI/Ming-omni-tts-0.5B",
+       model: "inclusionAI/Ming-omni-tts-16.8B-A3B",
        input: "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。",
        ref_audio: $ref_audio,
        ref_text: "在此奉劝大家别乱打美白针。",
@@ -107,7 +104,7 @@ jq -n \
     --arg a "$REF_A" \
     --arg b "$REF_B" \
     '{
-       model: "inclusionAI/Ming-omni-tts-0.5B",
+       model: "inclusionAI/Ming-omni-tts-16.8B-A3B",
        input: $input,
        ref_audio: [$a, $b],
        ref_text: $ref_text,

tests/config/test_omni_config.py

@@ -733,8 +733,7 @@ def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae
 
 def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
     cfg = OmniStageDiffusionParallelConfig(
-        pipeline_parallel_size=2,
+        ulysses_degree=4,
-        ulysses_degree=2,
         use_hsdp=True,
         hsdp_shard_size=-1,
         hsdp_replicate_size=2,
@@ -745,10 +744,10 @@ def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
 
 
 def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp():
-    with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+    with pytest.raises(ValueError, match="not compatible with TP"):
         OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
 
-    with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+    with pytest.raises(ValueError, match="not compatible with DP"):
         OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
 
 

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -74,14 +74,6 @@ def mock_dist(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(dist, "broadcast", lambda *args, **kwargs: None)
 
 
-@pytest.fixture(autouse=True)
-def mock_dit_group(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group",
-        lambda: None,
-    )
-
-
 @pytest.fixture(autouse=True)
 def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(DistributedVaeExecutor, "gather_tensors", lambda self, x: [x])

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -92,9 +92,8 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
     monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
     monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group)
-    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
 
-    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
         monkeypatch.setattr(parallel_state, name, None)
     for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
         monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
@@ -229,9 +228,8 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
     monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
     monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
     monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
-    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
 
-    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
         monkeypatch.setattr(parallel_state, name, None)
     for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
         monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
@@ -259,6 +257,33 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
     ]
 
 
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_destroy_model_parallel_clears_vllm_pipeline_group(monkeypatch):
+    """A destroyed diffusion PP group must not block the next initialization."""
+
+    class _DestroyableGroup:
+        def __init__(self) -> None:
+            self.destroy_calls = 0
+
+        def destroy(self) -> None:
+            self.destroy_calls += 1
+
+    pipeline_group = _DestroyableGroup()
+    for name in ("_DP", "_CFG", "_SP"):
+        monkeypatch.setattr(parallel_state, name, None)
+    monkeypatch.setattr(parallel_state, "_PP", pipeline_group)
+    monkeypatch.setattr(parallel_state, "_EXPERT_PARALLEL_GROUP_RANKS", None)
+    for name in ("_DP", "_PCP", "_TP", "_EP"):
+        monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+    monkeypatch.setattr(parallel_state.vllm_parallel_state, "_PP", pipeline_group, raising=False)
+
+    parallel_state.destroy_model_parallel()
+
+    assert pipeline_group.destroy_calls == 1
+    assert parallel_state.vllm_parallel_state._PP is None
+
+
 @pytest.mark.cpu
 @pytest.mark.core_model
 def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
@@ -295,9 +320,8 @@ def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
     monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
     monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
     monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
-    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
 
-    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
         monkeypatch.setattr(parallel_state, name, None)
     for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
         monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)

tests/diffusion/distributed/test_hsdp.py

@@ -234,7 +234,7 @@ class TestDiffusionParallelConfigHSDP:
 
     def test_hsdp_cannot_use_with_tp(self):
         """Test that HSDP and Tensor Parallelism cannot be used together."""
-        with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+        with pytest.raises(ValueError, match="not compatible with TP"):
             DiffusionParallelConfig(
                 tensor_parallel_size=2,
                 use_hsdp=True,
@@ -243,7 +243,7 @@ class TestDiffusionParallelConfigHSDP:
 
     def test_hsdp_cannot_use_with_dp(self):
         """Test that HSDP and Data Parallelism cannot be used together."""
-        with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+        with pytest.raises(ValueError, match="not compatible with DP"):
             DiffusionParallelConfig(
                 data_parallel_size=2,
                 use_hsdp=True,
@@ -264,158 +264,6 @@ class TestDiffusionParallelConfigHSDP:
         assert config.hsdp_shard_size == 2  # auto: 4 // 2
 
 
-class TestStandaloneHSDPDetection:
-    """Tests for standalone HSDP detection and dit_parallel_size calculation.
-
-    These tests verify the logic used in initialize_model_parallel() to detect
-    standalone HSDP mode and calculate effective parallel sizes.
-
-    Standalone HSDP is when all non-HSDP parallelism dimensions are 1.
-    """
-
-    @staticmethod
-    def compute_standalone_hsdp_params(
-        data_parallel_size: int = 1,
-        cfg_parallel_size: int = 1,
-        sequence_parallel_size: int = 1,
-        pipeline_parallel_size: int = 1,
-        tensor_parallel_size: int = 1,
-        fully_shard_degree: int = 1,
-        hsdp_replicate_size: int = 1,
-    ) -> dict:
-        """Compute standalone HSDP detection parameters.
-
-        This mirrors the logic in initialize_model_parallel().
-        """
-        dit_parallel_size = (
-            data_parallel_size
-            * cfg_parallel_size
-            * sequence_parallel_size
-            * pipeline_parallel_size
-            * tensor_parallel_size
-        )
-
-        # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
-        is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1
-
-        # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size)
-        if is_standalone_hsdp:
-            effective_dit_parallel_size = fully_shard_degree * hsdp_replicate_size
-        else:
-            effective_dit_parallel_size = dit_parallel_size
-
-        effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size
-
-        return {
-            "original_dit_parallel_size": dit_parallel_size,
-            "is_standalone_hsdp": is_standalone_hsdp,
-            "effective_dit_parallel_size": effective_dit_parallel_size,
-            "effective_dp_size": effective_dp_size,
-        }
-
-    def test_standalone_hsdp_basic(self):
-        """Test basic standalone HSDP detection (shard_size=4, replicate=1)."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 1
-        assert result["is_standalone_hsdp"] is True
-        assert result["effective_dit_parallel_size"] == 4
-        assert result["effective_dp_size"] == 4
-
-    def test_standalone_hsdp_with_replicate(self):
-        """Test standalone HSDP with replication (shard_size=4, replicate=2)."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=4,
-            hsdp_replicate_size=2,
-        )
-        assert result["original_dit_parallel_size"] == 1
-        assert result["is_standalone_hsdp"] is True
-        assert result["effective_dit_parallel_size"] == 8  # 4 * 2
-        assert result["effective_dp_size"] == 8
-
-    def test_combined_hsdp_sp_not_standalone(self):
-        """Test HSDP combined with SP is NOT detected as standalone.
-
-        This is a regression test for the bug where the condition
-        `dit_parallel_size == fully_shard_degree` incorrectly matched
-        combined modes like SP=4 + HSDP=4.
-        """
-        result = self.compute_standalone_hsdp_params(
-            sequence_parallel_size=4,
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 4
-        assert result["is_standalone_hsdp"] is False
-        # Should NOT override dp_size for combined mode
-        assert result["effective_dp_size"] == 1  # original data_parallel_size
-
-    def test_combined_hsdp_cfg_not_standalone(self):
-        """Test HSDP combined with CFG is NOT detected as standalone."""
-        result = self.compute_standalone_hsdp_params(
-            cfg_parallel_size=2,
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 2
-        assert result["is_standalone_hsdp"] is False
-        assert result["effective_dp_size"] == 1
-
-    def test_combined_hsdp_pp_not_standalone(self):
-        """Test HSDP combined with PP is NOT detected as standalone."""
-        result = self.compute_standalone_hsdp_params(
-            pipeline_parallel_size=2,
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 2
-        assert result["is_standalone_hsdp"] is False
-        assert result["effective_dp_size"] == 1
-
-    def test_no_hsdp_not_standalone(self):
-        """Test that no HSDP (fully_shard_degree=1) is NOT standalone."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=1,
-        )
-        assert result["original_dit_parallel_size"] == 1
-        assert result["is_standalone_hsdp"] is False
-        assert result["effective_dp_size"] == 1
-
-    def test_combined_multiple_parallelism_not_standalone(self):
-        """Test HSDP combined with multiple parallelism is NOT standalone."""
-        result = self.compute_standalone_hsdp_params(
-            sequence_parallel_size=2,
-            cfg_parallel_size=2,
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 4  # 2 * 2
-        assert result["is_standalone_hsdp"] is False
-        assert result["effective_dp_size"] == 1
-
-    def test_standalone_hsdp_large_shard(self):
-        """Test standalone HSDP with large shard size."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=8,
-            hsdp_replicate_size=1,
-        )
-        assert result["is_standalone_hsdp"] is True
-        assert result["effective_dit_parallel_size"] == 8
-        assert result["effective_dp_size"] == 8
-
-    def test_standalone_hsdp_large_replicate(self):
-        """Test standalone HSDP with large replicate size."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=4,
-            hsdp_replicate_size=4,
-        )
-        assert result["is_standalone_hsdp"] is True
-        assert result["effective_dit_parallel_size"] == 16  # 4 * 4
-        assert result["effective_dp_size"] == 16
-
-
 class TestHSDPShardConditions:
     """Tests for _hsdp_shard_conditions matching logic."""
 

tests/diffusion/models/sana_wm/test_cam_reduce_scatter.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for SANA-WM's camera-branch reduce-scatter.
+
+``out_proj_cam`` is row-parallel, so every rank produces a partial sum over the
+full hidden width, while the GDN stream it is added to is TP-local. Instead of
+all-reducing and then discarding all but one slice, the branch reduce-scatters
+straight to the local shard. These tests pin the shard each rank ends up with,
+including the bias slice, on both code paths:
+
+* the process-group path, which drives ``torch.distributed.reduce_scatter_tensor``
+  directly because vLLM-Omni's diffusion ``GroupCoordinator`` does not implement
+  ``reduce_scatter``;
+* the all-reduce-and-slice fallback taken when the coordinator exposes no
+  ``device_group``.
+
+CPU-only: the collectives and the rank/world-size lookups are mocked, so a
+single process can simulate an arbitrary TP degree.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import SanaWmSelfAttention
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer"
+
+HIDDEN = 24
+
+
+def _host(bias: torch.Tensor | None) -> SimpleNamespace:
+    """A stand-in ``self`` exposing only what the method reads."""
+    return SimpleNamespace(out_proj_cam=SimpleNamespace(bias=bias))
+
+
+def _fake_reduce_scatter_tensor(rank: int, tp_size: int):
+    """Model ``reduce_scatter_tensor``: split dim 0 and keep chunk ``rank``.
+
+    The test drives a single process, so the other ranks' partials are already
+    folded into the input; scattering is all this has to reproduce.
+    """
+
+    def _impl(output: torch.Tensor, input_: torch.Tensor, group=None) -> None:
+        output.copy_(input_.chunk(tp_size, dim=0)[rank])
+
+    return _impl
+
+
+def _tp_group(with_device_group: bool):
+    return SimpleNamespace(device_group=object() if with_device_group else None)
+
+
+@pytest.mark.parametrize("tp_size", [2, 4])
+@pytest.mark.parametrize("rank", [0, 1])
+def test_process_group_path_matches_all_reduce_then_slice(tp_size: int, rank: int) -> None:
+    """The scattered shard equals the slice the all-reduce path would take."""
+    local_width = HIDDEN // tp_size
+    torch.manual_seed(0)
+    contrib = torch.randn(2, 3, HIDDEN)
+    bias = torch.randn(HIDDEN)
+    reference = torch.zeros(2, 3, local_width)
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size),
+        patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank),
+        patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(True)),
+        patch("torch.distributed.reduce_scatter_tensor", side_effect=_fake_reduce_scatter_tensor(rank, tp_size)),
+    ):
+        out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference)
+
+    start = rank * local_width
+    expected = (contrib + bias)[..., start : start + local_width]
+    assert out.shape == (2, 3, local_width)
+    torch.testing.assert_close(out, expected)
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_fallback_path_matches_the_process_group_path(rank: int) -> None:
+    """Without a device_group the method all-reduces and slices, same result."""
+    tp_size = 2
+    local_width = HIDDEN // tp_size
+    torch.manual_seed(1)
+    contrib = torch.randn(2, 3, HIDDEN)
+    bias = torch.randn(HIDDEN)
+    reference = torch.zeros(2, 3, local_width)
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size),
+        patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank),
+        patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(False)),
+        patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=lambda t: t),
+    ):
+        out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference)
+
+    start = rank * local_width
+    torch.testing.assert_close(out, (contrib + bias)[..., start : start + local_width])
+
+
+def test_tp1_adds_bias_without_a_collective() -> None:
+    contrib = torch.randn(2, 3, HIDDEN)
+    bias = torch.randn(HIDDEN)
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1),
+        patch("torch.distributed.reduce_scatter_tensor") as collective,
+    ):
+        out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, contrib)
+
+    assert collective.call_count == 0
+    torch.testing.assert_close(out, contrib + bias)
+
+
+def test_rejects_a_width_that_does_not_shard() -> None:
+    """A local width that is not 1/tp of the contribution is a wiring bug."""
+    contrib = torch.randn(1, 2, HIDDEN)
+    reference = torch.zeros(1, 2, HIDDEN // 3)  # tp=2 would need HIDDEN // 2
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
+        patch("torch.distributed.reduce_scatter_tensor"),
+        pytest.raises(ValueError, match="camera contribution width mismatch"),
+    ):
+        SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(None), contrib, reference)

tests/diffusion/models/sana_wm/test_fused_qk_rms_norm.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for SANA-WM's TensorParallelRMSNorm q/k fusion.
+
+`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor
+and issues a single tensor-parallel all-reduce instead of two. These tests lock
+in that it is numerically identical to applying `norm_q`/`norm_k` separately,
+and that it really issues exactly one (fused) all-reduce.
+
+CPU-only: the collective is mocked, so a single process simulates one rank
+(all-reduce is identity). This exercises the cat/split logic a real multi-rank
+run depends on.
+"""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import (
+    TensorParallelRMSNorm,
+    fused_qk_rms_norm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer"
+
+
+def _make_norms(dim: int, tp_size: int) -> tuple[TensorParallelRMSNorm, TensorParallelRMSNorm]:
+    """Two norms with distinct random weights so a swapped q/k slice is caught."""
+    torch.manual_seed(0)
+    norm_q = TensorParallelRMSNorm(dim, tp_size=tp_size)
+    norm_k = TensorParallelRMSNorm(dim, tp_size=tp_size)
+    norm_q.weight.data.normal_()
+    norm_k.weight.data.normal_()
+    return norm_q, norm_k
+
+
+def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor:
+    """Single-rank all-reduce: the reduction over one rank is the identity."""
+    return tensor
+
+
+def test_fused_matches_separate() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim, tp_size=2)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce):
+        ref_q, ref_k = norm_q(q), norm_k(k)
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
+
+
+def test_fused_issues_one_all_reduce() -> None:
+    dim = 32
+    norm_q, norm_k = _make_norms(dim, tp_size=4)
+    q = torch.randn(1, 3, dim)
+    k = torch.randn(1, 3, dim)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+        fused_qk_rms_norm(norm_q, norm_k, q, k)
+        assert collective.call_count == 1
+        # Both sums travel in one tensor: last dim is 2, not 1.
+        assert collective.call_args.args[0].shape[-1] == 2
+
+        collective.reset_mock()
+        norm_q(q), norm_k(k)
+        assert collective.call_count == 2
+
+
+@pytest.mark.parametrize(
+    "norms",
+    [
+        (RMSNorm(16), RMSNorm(16)),
+        (nn.Identity(), nn.Identity()),
+    ],
+    ids=["plain_rms_norm", "identity"],
+)
+def test_falls_back_for_non_tp_norms(norms: tuple[nn.Module, nn.Module]) -> None:
+    """qk_norm=False and TP=1 both hand in norms the fused path cannot use."""
+    norm_q, norm_k = norms
+    q = torch.randn(2, 4, 16)
+    k = torch.randn(2, 4, 16)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+        assert collective.call_count == 0
+
+    torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0)
+
+
+def test_tp1_skips_the_collective() -> None:
+    norm_q, norm_k = _make_norms(16, tp_size=1)
+    q = torch.randn(2, 4, 16)
+    k = torch.randn(2, 4, 16)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+        assert collective.call_count == 0
+
+    torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0)

tests/diffusion/models/sana_wm/test_request.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for SANA-WM request normalization."""
+
+import pytest
+
+from vllm_omni.diffusion.models.sana_wm.request import (
+    SANA_WM_DEFAULT_HEIGHT,
+    SANA_WM_DEFAULT_NUM_FRAMES,
+    SANA_WM_DEFAULT_WIDTH,
+    normalize_sana_wm_payload,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _prompt(**sana_wm):
+    """A valid request, with the action rollout kept in step with num_frames."""
+    block = {"num_frames": SANA_WM_DEFAULT_NUM_FRAMES, **sana_wm}
+    block.setdefault("action", f"w-{block['num_frames'] - 1}")
+    return {
+        "prompt": "a quiet city street",
+        "multi_modal_data": {"image": object()},
+        "sana_wm": block,
+    }
+
+
+def _identity_poses(count):
+    return [[[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]] for _ in range(count)]
+
+
+def _camera_prompt(**camera):
+    prompt = _prompt(num_frames=9)
+    prompt["sana_wm"].pop("action")
+    prompt["sana_wm"]["camera"] = {"poses": _identity_poses(9), **camera}
+    return prompt
+
+
+def test_defaults_are_latent_aligned():
+    payload = normalize_sana_wm_payload(_prompt())["additional_information"]["sana_wm"]
+    assert payload["num_frames"] == SANA_WM_DEFAULT_NUM_FRAMES
+    assert payload["height"] == SANA_WM_DEFAULT_HEIGHT
+    assert payload["width"] == SANA_WM_DEFAULT_WIDTH
+
+
+@pytest.mark.parametrize("num_frames", [9, 33, 161])
+def test_accepts_aligned_geometry(num_frames):
+    payload = normalize_sana_wm_payload(_prompt(num_frames=num_frames))["additional_information"]["sana_wm"]
+    assert payload["num_frames"] == num_frames
+
+
+@pytest.mark.parametrize("num_frames", [8, 10, 160])
+def test_rejects_unaligned_num_frames(num_frames):
+    # The VAE compresses 8x temporally; an unaligned count would be floored.
+    with pytest.raises(ValueError, match="num_frames"):
+        normalize_sana_wm_payload(_prompt(num_frames=num_frames))
+
+
+@pytest.mark.parametrize(("height", "width"), [(700, 1280), (704, 1270), (703, 1279)])
+def test_rejects_unaligned_resolution(height, width):
+    with pytest.raises(ValueError, match="divisible by 32"):
+        normalize_sana_wm_payload(_prompt(height=height, width=width))
+
+
+def test_normalization_is_idempotent():
+    once = normalize_sana_wm_payload(_prompt(num_frames=33))
+    twice = normalize_sana_wm_payload(once)
+    assert once["additional_information"]["sana_wm"] == twice["additional_information"]["sana_wm"]
+
+
+def test_rejects_action_shorter_than_num_frames():
+    # Used to shrink the camera tensor silently, then fail deep in the
+    # transformer on a camera/latent frame mismatch.
+    with pytest.raises(ValueError, match="rolls out 33 frames but num_frames is 161"):
+        normalize_sana_wm_payload(_prompt(action="w-32"))
+
+
+def test_rejects_action_longer_than_num_frames():
+    with pytest.raises(ValueError, match="rolls out 161 frames but num_frames is 9"):
+        normalize_sana_wm_payload(_prompt(num_frames=9, action="w-160"))
+
+
+def test_rejects_oversized_action_without_expanding_it():
+    # A 12-byte request body: expanding this rollout before checking the frame
+    # count allocated ~16 GB. Rejection has to be arithmetic, so this test also
+    # fails (by hanging/OOM) if the duration is ever expanded again.
+    with pytest.raises(ValueError, match="rolls out 1000000001 frames"):
+        normalize_sana_wm_payload(_prompt(action="w-1000000000"))
+
+
+@pytest.mark.parametrize("action", ["", "   ", "w-0", "w-abc", "w", "x-4"])
+def test_rejects_malformed_action(action):
+    with pytest.raises(ValueError):
+        normalize_sana_wm_payload(_prompt(action=action))
+
+
+@pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf")])
+def test_rejects_unusable_focal_length(value):
+    # compute_raymap divides by fx/fy; a zero or non-finite focal length turns
+    # the whole Plucker map into NaN and the response into a NaN video.
+    intrinsics = {"fx": value, "fy": 640, "cx": 640, "cy": 352}
+    with pytest.raises(ValueError, match="fx"):
+        normalize_sana_wm_payload(_prompt(intrinsics=intrinsics))
+
+
+@pytest.mark.parametrize("value", [float("nan"), float("inf")])
+def test_rejects_non_finite_principal_point(value):
+    intrinsics = {"fx": 640, "fy": 640, "cx": value, "cy": 352}
+    with pytest.raises(ValueError, match="must be finite"):
+        normalize_sana_wm_payload(_prompt(intrinsics=intrinsics))
+
+
+@pytest.mark.parametrize("value", [float("nan"), float("inf")])
+def test_rejects_non_finite_speeds(value):
+    # `parsed <= 0` is False for NaN, so the positivity check alone let it pass.
+    with pytest.raises(ValueError, match="must be finite"):
+        normalize_sana_wm_payload(_prompt(translation_speed=value))
+
+
+def test_rejects_non_finite_pose_entry():
+    prompt = _camera_prompt()
+    prompt["sana_wm"]["camera"]["poses"][3][2][1] = float("nan")
+    with pytest.raises(ValueError, match="must be finite"):
+        normalize_sana_wm_payload(prompt)
+
+
+def test_accepts_default_camera_metadata():
+    payload = normalize_sana_wm_payload(_camera_prompt())["additional_information"]["sana_wm"]
+    assert payload["camera"]["poses"]
+    # The unread fields are not carried into the canonical payload.
+    assert set(payload["camera"]) == {"poses"}
+
+
+@pytest.mark.parametrize(
+    ("field", "value"),
+    [("format", "w2c_4x4"), ("coordinate_system", "opengl")],
+)
+def test_rejects_unsupported_camera_convention(field, value):
+    # Nothing downstream reads these, so accepting one would mean interpreting
+    # the poses under the wrong convention.
+    with pytest.raises(ValueError, match=field.split("_")[0]):
+        normalize_sana_wm_payload(_camera_prompt(**{field: value}))

tests/diffusion/test_diffusion_config_propagation.py

@@ -92,6 +92,30 @@ class TestParallelConfigPropagation:
         od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
         assert od.num_gpus == 2
 
+    def test_dp_is_inferred_from_num_gpus(self):
+        pc = DiffusionParallelConfig(tensor_parallel_size=2, ulysses_degree=2)
+        od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=8)
+        assert od.parallel_config.data_parallel_size == 2
+        assert od.parallel_config.world_size == 8
+
+    def test_explicit_dp_is_validated_against_num_gpus(self):
+        pc = DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2)
+        od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4)
+        assert od.parallel_config.data_parallel_size == 2
+
+        with pytest.raises(ValueError, match="does not match WORLD-derived value"):
+            OmniDiffusionConfig.from_kwargs(
+                model="x",
+                parallel_config=DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2),
+                num_gpus=8,
+            )
+
+    def test_hsdp_does_not_infer_ordinary_dp(self):
+        pc = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=4)
+        od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4)
+        assert od.parallel_config.data_parallel_size == 1
+        assert od.parallel_config.world_size == 4
+
 
 class TestCreateDefaultDiffusion:
     """Verify engine_args structure from create_default_diffusion."""

tests/e2e/offline_inference/test_sana_wm.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import numpy as np
+import pytest
+
+pytestmark = [
+    pytest.mark.advanced_model,
+    pytest.mark.diffusion,
+    pytest.mark.gpu,
+]
+
+
+def _make_sana_wm_e2e_image() -> Any:
+    from PIL import Image
+
+    from vllm_omni.diffusion.models.sana_wm import SANA_WM_OUTPUT_HEIGHT, SANA_WM_OUTPUT_WIDTH
+
+    return Image.new("RGB", (SANA_WM_OUTPUT_WIDTH, SANA_WM_OUTPUT_HEIGHT), (96, 128, 160))
+
+
+def _sana_wm_e2e_tensor_parallel_size() -> int:
+    return int(
+        os.environ.get(
+            "SANA_WM_E2E_TENSOR_PARALLEL_SIZE",
+            os.environ.get("SANA_WM_E2E_TP", "1"),
+        )
+    )
+
+
+def _coerce_video_array(video: Any) -> np.ndarray:
+    import torch
+
+    if isinstance(video, list):
+        assert video, "SANA-WM e2e produced an empty video list."
+        video = video[0]
+    if isinstance(video, torch.Tensor):
+        video = video.detach().cpu().float().numpy()
+    video_array = np.asarray(video)
+    if video_array.ndim == 5:
+        assert video_array.shape[0] == 1
+        video_array = video_array[0]
+    assert video_array.ndim == 4
+    return video_array
+
+
+def _assert_sana_wm_e2e_shape(
+    video: np.ndarray,
+    *,
+    output_type: str,
+    num_frames: int,
+) -> None:
+    if output_type == "latent":
+        # Latent output is channel-first: (C, T, H, W) after removing the
+        # batch dimension. Stage-1 emits LTX-2 latents with 128 channels.
+        assert video.shape[0] == 128
+        assert 0 < video.shape[1] <= num_frames
+        return
+    assert 0 < video.shape[0] <= num_frames
+
+
+def _run_sana_wm_e2e(*, output_type: str) -> np.ndarray:
+    import torch
+
+    from vllm_omni.diffusion.models.sana_wm import (
+        SANA_WM_MODEL_ID,
+        SANA_WM_OUTPUT_HEIGHT,
+        SANA_WM_OUTPUT_WIDTH,
+    )
+    from vllm_omni.entrypoints.omni import Omni
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.outputs import OmniRequestOutput
+
+    if not torch.cuda.is_available():
+        pytest.skip("Sana-WM e2e requires CUDA.")
+
+    model = os.environ.get("SANA_WM_E2E_MODEL", SANA_WM_MODEL_ID)
+    model_class_name = os.environ.get("SANA_WM_E2E_MODEL_CLASS", "SanaWmPipeline")
+    num_frames = int(os.environ.get("SANA_WM_E2E_NUM_FRAMES", "9"))
+    stage1_steps = int(
+        os.environ.get(
+            "SANA_WM_E2E_STAGE1_STEPS",
+            os.environ.get("SANA_WM_E2E_NUM_INFERENCE_STEPS", "1"),
+        )
+    )
+    action = os.environ.get("SANA_WM_E2E_ACTION", f"w-{max(num_frames - 1, 1)}")
+    tensor_parallel_size = _sana_wm_e2e_tensor_parallel_size()
+    image = _make_sana_wm_e2e_image()
+    omni = Omni(
+        model=model,
+        model_class_name=model_class_name,
+        enforce_eager=True,
+        tensor_parallel_size=tensor_parallel_size,
+    )
+    extra_args: dict[str, Any] = {}
+    # The native Stage-1 path is latent-token capped; raise it for full-size
+    # e2e runs (704x1280 exceeds the default 4096-token cap).
+    native_max_tokens = os.environ.get("SANA_WM_E2E_NATIVE_MAX_TOKENS", "").strip()
+    if native_max_tokens:
+        extra_args["sana_wm_native_max_tokens"] = int(native_max_tokens)
+    # Return the requested output type so the shape assertion matches the
+    # pipeline output instead of the raw Stage-1 latent default.
+
+    output = omni.generate(
+        {
+            "prompt": "A slow forward camera move through a quiet city street.",
+            "multi_modal_data": {"image": image},
+            "sana_wm": {
+                "action": action,
+                "num_frames": num_frames,
+                "translation_speed": 0.055,
+                "rotation_speed_deg": 1.2,
+                # Avoid the optional Pi3X dependency by passing deterministic
+                # camera intrinsics directly.
+                "intrinsics": {
+                    "fx": SANA_WM_OUTPUT_WIDTH / 2,
+                    "fy": SANA_WM_OUTPUT_WIDTH / 2,
+                    "cx": SANA_WM_OUTPUT_WIDTH / 2,
+                    "cy": SANA_WM_OUTPUT_HEIGHT / 2,
+                },
+            },
+        },
+        OmniDiffusionSamplingParams(
+            height=SANA_WM_OUTPUT_HEIGHT,
+            width=SANA_WM_OUTPUT_WIDTH,
+            num_frames=num_frames,
+            seed=0,
+            fps=16,
+            num_inference_steps=stage1_steps,
+            guidance_scale=1.0,
+            guidance_scale_provided=True,
+            output_type=output_type,
+            extra_args=extra_args,
+        ),
+    )
+
+    request_output = output[0] if isinstance(output, list) else output
+    assert isinstance(request_output, OmniRequestOutput)
+    assert request_output.error is None
+    assert request_output.images
+
+    frames = request_output.images[0]
+    video = _coerce_video_array(frames)
+    _assert_sana_wm_e2e_shape(video, output_type=output_type, num_frames=num_frames)
+    return video
+
+
+def test_sana_wm_native_generates_video() -> None:
+    output_type = os.environ.get("SANA_WM_E2E_OUTPUT_TYPE", "np")
+    video = _run_sana_wm_e2e(output_type=output_type)
+    if output_type == "latent":
+        assert video.ndim == 4
+    else:
+        assert video.shape[-1] in (3, 4)

tests/e2e/online_serving/test_sana_wm.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Online serving smoke for SANA-WM image-to-video (first-frame I2V via ``/v1/videos``).
+
+Boots the server straight from the model id — the Stage-1 repo's
+``model_index.json`` names ``SanaWmPipeline``, so the class resolves on its own
+— submits one async ``/v1/videos`` job with a first-frame reference image, and
+asserts the job completes and returns video bytes.
+
+From ``tests/``::
+
+    pytest -s -v e2e/online_serving/test_sana_wm.py -m "advanced_model and diffusion" --run-level=advanced_model
+"""
+
+import base64
+import json
+import os
+from io import BytesIO
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from vllm_omni.diffusion.models.sana_wm import (
+    SANA_WM_MODEL_ID,
+    SANA_WM_OUTPUT_HEIGHT,
+    SANA_WM_OUTPUT_WIDTH,
+)
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = SANA_WM_MODEL_ID
+PROMPT = "A slow forward camera move through a quiet city street."
+NEGATIVE_PROMPT = "blurry, low quality, distorted, watermark"
+
+# Smoke knobs: a handful of frames / steps so the job finishes quickly.
+SMOKE_NUM_FRAMES = 9
+SMOKE_NUM_INFERENCE_STEPS = 2
+# First-frame I2V needs a camera trajectory; "w-<n>" is the forward-move action DSL.
+# Passed via the ``sana_wm`` form field (a JSON string parsed server-side); explicit
+# intrinsics avoid the optional Pi3X camera-calibration dependency.
+SANA_WM_PARAMS = {
+    "action": f"w-{SMOKE_NUM_FRAMES - 1}",
+    "translation_speed": 0.055,
+    "rotation_speed_deg": 1.2,
+    "intrinsics": {
+        "fx": SANA_WM_OUTPUT_WIDTH / 2,
+        "fy": SANA_WM_OUTPUT_WIDTH / 2,
+        "cx": SANA_WM_OUTPUT_WIDTH / 2,
+        "cy": SANA_WM_OUTPUT_HEIGHT / 2,
+    },
+}
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _first_frame_data_url() -> str:
+    """A solid-color first frame encoded as a PNG ``data:`` URL for I2V."""
+    from PIL import Image
+
+    image = Image.new("RGB", (SANA_WM_OUTPUT_WIDTH, SANA_WM_OUTPUT_HEIGHT), (96, 128, 160))
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
+    return f"data:image/png;base64,{encoded}"
+
+
+def _get_diffusion_feature_cases(model: str):
+    """Single default server row."""
+    return [
+        pytest.param(
+            OmniServerParams(model=model),
+            id="default",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_image_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    """Default SANA-WM I2V smoke: async ``/v1/videos`` job completes and returns video bytes."""
+    request_config = {
+        "model": omni_server.model,
+        "image_reference": _first_frame_data_url(),
+        "form_data": {
+            "prompt": PROMPT,
+            "negative_prompt": NEGATIVE_PROMPT,
+            # The camera payload rides inside extra_params; there is no
+            # top-level sana_wm form field on the video endpoints, so sending
+            # one is silently dropped and the request fails in the
+            # preprocessor for missing camera control.
+            "extra_params": json.dumps({"sana_wm": SANA_WM_PARAMS}),
+            "height": SANA_WM_OUTPUT_HEIGHT,
+            "width": SANA_WM_OUTPUT_WIDTH,
+            "num_frames": SMOKE_NUM_FRAMES,
+            "fps": 16,
+            "num_inference_steps": SMOKE_NUM_INFERENCE_STEPS,
+            "guidance_scale": 5.0,
+            "seed": 42,
+        },
+    }
+    openai_client.send_video_diffusion_request(request_config)

tests/entrypoints/openai_api/test_api_server_guards.py

@@ -0,0 +1,810 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""API-server surface guards for the OpenAI Omni entrypoint.
+
+Purpose
+-------
+Freeze the *observable* server contract so helper moves (P0.2), route
+extraction (P0.3), and later endpoint-family PRs cannot silently change:
+
+* which routes exist,
+* how upstream routes are overridden at assembly time,
+* unavailable / missing-handler envelopes,
+* shared engine-error JSON fields,
+* app-state keys wired by ``omni_init_app_state``.
+
+CPU-only (fakes / TestClient). No checked-in OpenAPI golden JSON.
+
+Common wrong changes these tests are meant to catch
+----------------------------------------------------
+* Route accidentally removed, renamed, method-changed, or double-registered
+  after a move (census / assembled-app checks).
+* OpenAPI drops an Omni HTTP path while the router still has it (or the reverse).
+* Assembly forgets to strip upstream ``/v1/chat/completions`` (batch) /
+  ``/v1/models``, or also deletes unrelated upstream routes.
+* Profiler / Omni router not mounted; storage not started; Omni engine
+  exception handlers not registered.
+* Timestamp middleware stops stamping HTTP, or starts mutating WebSocket scopes.
+* Missing-handler WS/HTTP responses change status, JSON shape, or wording
+  (clients treat these as capability probes).
+* Engine error payloads lose ``request_id`` / ``error_stage_id``.
+* Pure-diffusion vs multi-stage init drops an ``app.state`` key, or leaves a
+  handler as ``None`` while the route is still registered.
+
+Not covered here (add elsewhere if needed)
+------------------------------------------
+* Deep per-endpoint success behavior (family UTs).
+* Byte-identical helper bodies after a move (P0.2 AST integrity).
+* Live GPU / speech byte-identity checks.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+
+import pytest
+from fastapi import FastAPI, HTTPException
+from fastapi.routing import APIRoute, APIWebSocketRoute
+from fastapi.testclient import TestClient
+from starlette.datastructures import State
+from starlette.requests import Request
+from starlette.websockets import WebSocketDisconnect
+from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
+
+from vllm_omni.entrypoints.openai import api_server
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# FastAPI/Starlette auto-adds HEAD on every GET route. Including HEAD here
+# would invent ("HEAD", path) entries for every GET and fail the current census.
+# TODO: extend this set and the expected census if @router.head / @router.options appear.
+_HTTP_METHODS = {"GET", "POST", "DELETE", "PUT", "PATCH"}
+
+# Exact Omni router census on main. Update intentionally when routes change.
+_EXPECTED_ROUTER_ROUTES = {
+    ("POST", "/v1/chat/completions"),
+    ("POST", "/v1/chat/completions/batch"),
+    ("POST", "/v1/audio/speech"),
+    ("POST", "/v1/audio/speech/batch"),
+    ("POST", "/v1/audio/generate"),
+    ("GET", "/v1/audio/voices"),
+    ("POST", "/v1/audio/voices"),
+    ("DELETE", "/v1/audio/voices/{name}"),
+    ("WEBSOCKET", "/v1/audio/speech/stream"),
+    ("WEBSOCKET", "/v1/video/chat/stream"),
+    ("WEBSOCKET", "/v1/realtime/video"),
+    ("WEBSOCKET", "/v1/realtime"),
+    ("WEBSOCKET", "/v1/realtime/robot/openpi"),
+    ("WEBSOCKET", "/v1/duplex"),
+    ("GET", "/health"),
+    ("GET", "/v1/models"),
+    ("POST", "/v1/images/generations"),
+    ("POST", "/v1/images/edits"),
+    ("POST", "/v1/videos"),
+    ("POST", "/v1/videos/sync"),
+    ("GET", "/v1/videos"),
+    ("GET", "/v1/videos/{video_id}"),
+    ("DELETE", "/v1/videos/{video_id}"),
+    ("GET", "/v1/videos/{video_id}/content"),
+    ("POST", "/v1/omni/sleep"),
+    ("POST", "/v1/omni/wakeup"),
+}
+_EXPECTED_PROFILER_ROUTES = {
+    ("POST", "/start_profile"),
+    ("POST", "/stop_profile"),
+}
+
+# App-state keys that must remain available across refactor phases.
+# Presence is not enough: "not wired" is written as ``state.X = None``.
+_DIFFUSION_APP_STATE_KEYS = {
+    "engine_client",
+    "log_stats",
+    "args",
+    "sleeping_stages",
+    "stage_configs",
+    "vllm_config",
+    "diffusion_engine",
+    "openai_serving_models",
+    "serving_tokenization",
+    "openai_serving_chat",
+    "openai_serving_chat_batch",
+    "openai_serving_speech",
+    "openai_serving_audio_generate",
+    "openai_serving_video",
+    "openai_streaming_video_output",
+    "openai_serving_duplex",
+    "openai_streaming_speech",
+    "openai_streaming_video",
+    "openai_serving_realtime_robot",
+    "enable_server_load_tracking",
+    "server_load_metrics",
+}
+_DIFFUSION_MUST_BE_NONE = {
+    "vllm_config",
+    "serving_tokenization",
+    "openai_serving_duplex",
+    "openai_streaming_speech",
+    "openai_streaming_video",
+}
+_DIFFUSION_MUST_BE_WIRED = _DIFFUSION_APP_STATE_KEYS - _DIFFUSION_MUST_BE_NONE
+_MULTISTAGE_APP_STATE_KEYS = {
+    "engine_client",
+    "log_stats",
+    "args",
+    "sleeping_stages",
+    "stage_configs",
+    "vllm_config",
+    "openai_serving_models",
+    "serving_tokenization",
+    "online_renderer",
+    "openai_serving_chat",
+    "openai_serving_chat_batch",
+    "openai_serving_speech",
+    "openai_serving_audio_generate",
+    "openai_streaming_speech",
+    "openai_streaming_video",
+    "openai_serving_duplex",
+    "openai_serving_realtime",
+    "openai_serving_video",
+    "openai_serving_realtime_robot",
+    "enable_server_load_tracking",
+    "server_load_metrics",
+}
+_MULTISTAGE_MUST_BE_NONE = {
+    "openai_serving_duplex",
+    "openai_serving_realtime_robot",
+}
+_MULTISTAGE_MUST_BE_WIRED = _MULTISTAGE_APP_STATE_KEYS - _MULTISTAGE_MUST_BE_NONE
+
+
+def _route_entries(routes) -> list[tuple[str, str]]:
+    entries: list[tuple[str, str]] = []
+    for route in routes:
+        if isinstance(route, APIRoute):
+            for method in sorted((route.methods or set()) & _HTTP_METHODS):
+                entries.append((method, route.path))
+        elif isinstance(route, APIWebSocketRoute):
+            entries.append(("WEBSOCKET", route.path))
+    return entries
+
+
+def _route_manifest(routes) -> set[tuple[str, str]]:
+    return set(_route_entries(routes))
+
+
+def _assert_unique_route_entries(routes) -> None:
+    entries = _route_entries(routes)
+    duplicates = sorted({item for item in entries if entries.count(item) > 1})
+    assert not duplicates, f"duplicate route registrations: {duplicates}"
+
+
+def _matching_http_routes(routes, method: str, path: str) -> list[APIRoute]:
+    return [
+        route
+        for route in routes
+        if isinstance(route, APIRoute) and route.path == path and method in (route.methods or set())
+    ]
+
+
+def _single_http_route(routes, method: str, path: str) -> APIRoute:
+    matches = _matching_http_routes(routes, method, path)
+    assert len(matches) == 1, f"expected exactly one {method} {path}, got {len(matches)}"
+    return matches[0]
+
+
+def _assert_app_state_snapshot(
+    state,
+    *,
+    expected_keys: set[str],
+    must_be_wired: set[str],
+    must_be_none: set[str],
+) -> None:
+    assert must_be_wired.isdisjoint(must_be_none)
+    assert must_be_wired | must_be_none <= expected_keys
+    present = {key for key in expected_keys if hasattr(state, key)}
+    assert present == expected_keys
+    not_wired = sorted(key for key in must_be_wired if getattr(state, key) is None)
+    assert not not_wired, f"app.state keys registered but not wired: {not_wired}"
+    unexpectedly_set = sorted(key for key in must_be_none if getattr(state, key) is not None)
+    assert not unexpectedly_set, f"app.state keys expected to stay None: {unexpectedly_set}"
+
+
+def _minimal_args(**overrides) -> SimpleNamespace:
+    args = SimpleNamespace(
+        tool_parser_plugin=None,
+        reasoning_parser_plugin=None,
+        reasoning_parser=None,
+        structured_outputs_config=SimpleNamespace(reasoning_parser=None),
+        enable_ssl_refresh=False,
+        host="127.0.0.1",
+        port=8000,
+        uvicorn_log_level="info",
+        disable_uvicorn_access_log=True,
+        ssl_keyfile=None,
+        ssl_certfile=None,
+        ssl_ca_certs=None,
+        ssl_cert_reqs=0,
+        ssl_ciphers=None,
+        h11_max_incomplete_event_size=None,
+        h11_max_header_count=None,
+        log_config_file=None,
+        disable_access_log_for_endpoints=None,
+        disable_log_stats=True,
+        enable_log_requests=False,
+        max_log_len=None,
+        served_model_name=None,
+        model="demo-model",
+        enable_server_load_tracking=False,
+        log_error_stack=False,
+        chat_template=None,
+        chat_template_content_format="auto",
+        trust_request_chat_template=False,
+        tool_server=None,
+        lora_modules=None,
+        response_role="assistant",
+        default_chat_template_kwargs=None,
+        return_tokens_as_token_ids=False,
+        enable_auto_tool_choice=False,
+        exclude_tools_when_tool_choice_none=False,
+        tool_call_parser=None,
+        enable_prompt_tokens_details=False,
+        enable_force_include_usage=False,
+        enable_log_outputs=False,
+        enable_log_deltas=False,
+        tokens_only=False,
+        shutdown_timeout=0,
+        stage_configs_path=None,
+        deploy_config=None,
+    )
+    for key, value in overrides.items():
+        setattr(args, key, value)
+    return args
+
+
+def _request_for(app: FastAPI, *, method: str = "GET", path: str = "/") -> Request:
+    return Request(
+        {
+            "type": "http",
+            "method": method,
+            "path": path,
+            "headers": [],
+            "app": app,
+            "state": {},
+        }
+    )
+
+
+class _FakeSocket:
+    def __init__(self) -> None:
+        self.closed = False
+
+    def close(self) -> None:
+        self.closed = True
+
+
+class _FakeEngineClient:
+    def __init__(self, *, stage_configs=None, endpoint_restrictions=None, vllm_config=None) -> None:
+        self.stage_configs = stage_configs if stage_configs is not None else []
+        self.endpoint_restrictions = endpoint_restrictions if endpoint_restrictions is not None else {}
+        self.model_config = SimpleNamespace()
+        self.input_processor = object()
+        self.renderer = object()
+        self.errored = False
+        self.vllm_config = vllm_config
+
+    async def get_vllm_config(self):
+        return self.vllm_config
+
+    async def get_supported_tasks(self) -> tuple[str, ...]:
+        return ("generate",)
+
+    async def get_tokenizer(self):
+        return SimpleNamespace(chat_template="dummy")
+
+    async def check_health(self) -> None:
+        return None
+
+
+def _marker(name: str):
+    return SimpleNamespace(name=name)
+
+
+def test_router_route_manifest_is_explicit_and_complete() -> None:
+    """Lock the Omni + profiler router census.
+
+    Fails if a route is added/removed/renamed, a method changes, or a
+    WebSocket path disappears during helper/router moves. Also fails if the
+    same public route is registered more than once, which a set manifest alone
+    would collapse.
+    """
+    # TODO(P0.3): retarget this object-level census when routes move onto family
+    # routers. The assembled-app census is the client-facing lock that must survive.
+    _assert_unique_route_entries(api_server.router.routes)
+    _assert_unique_route_entries(api_server.profiler_router.routes)
+    assert _route_manifest(api_server.router.routes) == _EXPECTED_ROUTER_ROUTES
+    assert _route_manifest(api_server.profiler_router.routes) == _EXPECTED_PROFILER_ROUTES
+
+
+def test_router_openapi_paths_cover_http_manifest() -> None:
+    """Lock that OpenAPI advertises every HTTP census route/method.
+
+    Fails if a decorator still registers a route but OpenAPI loses the path
+    (or method), e.g. after moving route bodies / response models.
+    """
+    app = FastAPI()
+    app.include_router(api_server.router)
+    app.include_router(api_server.profiler_router)
+    paths = app.openapi().get("paths", {})
+
+    http_routes = {
+        (method, path) for method, path in _EXPECTED_ROUTER_ROUTES | _EXPECTED_PROFILER_ROUTES if method != "WEBSOCKET"
+    }
+    for method, path in sorted(http_routes):
+        assert path in paths, f"missing OpenAPI path {path}"
+        assert method.lower() in paths[path], f"missing OpenAPI method {method} {path}"
+
+
+@pytest.mark.asyncio
+async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_router(monkeypatch) -> None:
+    """Lock worker assembly: override, mount, storage, handlers, full census.
+
+    Fails if assembly stops replacing upstream chat/batch/models, deletes
+    unrelated upstream routes, forgets Omni/profiler mounts, skips storage
+    start, or drops Omni ``EngineDeadError`` / ``EngineGenerateError`` handlers.
+    """
+    captured: dict[str, object] = {}
+
+    def fake_build_openai_app(args, supported_tasks):
+        app = FastAPI()
+
+        @app.post("/v1/chat/completions")
+        async def upstream_chat():
+            return {"owner": "upstream"}
+
+        @app.post("/v1/chat/completions/batch")
+        async def upstream_batch():
+            return {"owner": "upstream"}
+
+        @app.get("/v1/models")
+        async def upstream_models():
+            return {"owner": "upstream"}
+
+        @app.get("/v1/upstream-only")
+        async def upstream_only():
+            return {"owner": "upstream"}
+
+        captured["supported_tasks"] = supported_tasks
+        return app
+
+    @asynccontextmanager
+    async def fake_build_async_omni(*_args, **_kwargs):
+        yield _FakeEngineClient(
+            stage_configs=[{"engine_args": {"profiler_config": {"profiler": "torch"}}}],
+            vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(_api_process_rank=0)),
+        )
+
+    async def fake_omni_init_app_state(engine_client, state, args):
+        state.engine_client = engine_client
+        state.initialized_by_omni = True
+
+    async def fake_storage_start():
+        captured["storage_started"] = True
+
+    async def fake_serve_http(app, **_kwargs):
+        captured["served_app"] = app
+        task = asyncio.create_task(asyncio.sleep(0))
+        await asyncio.sleep(0)
+        return task
+
+    monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app)
+    monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
+    monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state)
+    monkeypatch.setattr(
+        api_server,
+        "shutdown_unsupported_routes",
+        lambda app, restrictions: captured.setdefault("restrictions", restrictions),
+    )
+    monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=fake_storage_start))
+    monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
+
+    sock = _FakeSocket()
+    await api_server.omni_run_server_worker("127.0.0.1:8000", sock, _minimal_args())
+
+    served_wrapper = captured["served_app"]
+    # Timestamp middleware wraps the FastAPI app.
+    assert hasattr(served_wrapper, "_inner")
+    served_app = served_wrapper._inner
+    routes = served_app.routes
+    _assert_unique_route_entries(routes)
+    manifest = _route_manifest(routes)
+
+    assert captured["supported_tasks"] == ("generate",)
+    assert captured["storage_started"] is True
+    assert captured["restrictions"] == {}
+    assert sock.closed is True
+    assert served_app.state.initialized_by_omni is True
+
+    # Upstream override: unrelated route kept; Omni owns overridden pairs once.
+    # Endpoint identity ensures the surviving route is the Omni handler, not
+    # the fake upstream handler with the same method/path.
+    assert ("GET", "/v1/upstream-only") in manifest
+    assert _single_http_route(routes, "POST", "/v1/chat/completions").endpoint is api_server.create_chat_completion
+    assert (
+        _single_http_route(routes, "POST", "/v1/chat/completions/batch").endpoint
+        is api_server.create_batch_chat_completion
+    )
+    assert _single_http_route(routes, "GET", "/v1/models").endpoint is api_server.show_available_models
+
+    # Assembled app includes Omni router HTTP/WS surface + profiler.
+    for item in _EXPECTED_ROUTER_ROUTES:
+        assert item in manifest, f"assembled app missing {item}"
+    assert ("POST", "/start_profile") in manifest
+    assert ("POST", "/stop_profile") in manifest
+
+    assert served_app.exception_handlers[EngineGenerateError]
+    assert served_app.exception_handlers[EngineDeadError]
+
+    # Light OpenAPI check on the assembled FastAPI app (not the ASGI wrapper).
+    openapi_paths = served_app.openapi().get("paths", {})
+    assert "/v1/chat/completions" in openapi_paths
+    assert "/v1/models" in openapi_paths
+    assert "/health" in openapi_paths
+
+
+@pytest.mark.asyncio
+async def test_timestamp_middleware_stamps_http_and_passes_websocket(monkeypatch) -> None:
+    """Lock outermost timestamp middleware behavior.
+
+    Fails if HTTP requests lose ``request_timestamp``, or WebSocket scopes
+    start getting stamped / mutated incorrectly.
+    """
+    captured: dict[str, object] = {}
+
+    def fake_build_openai_app(args, supported_tasks):
+        return FastAPI()
+
+    @asynccontextmanager
+    async def fake_build_async_omni(*_args, **_kwargs):
+        yield _FakeEngineClient()
+
+    async def fake_omni_init_app_state(engine_client, state, args):
+        state.engine_client = engine_client
+
+    async def fake_serve_http(app, **_kwargs):
+        captured["served_app"] = app
+        task = asyncio.create_task(asyncio.sleep(0))
+        await asyncio.sleep(0)
+        return task
+
+    monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app)
+    monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
+    monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state)
+    monkeypatch.setattr(api_server, "shutdown_unsupported_routes", lambda *_a, **_k: None)
+    monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=lambda: asyncio.sleep(0)))
+    monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
+
+    await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), _minimal_args())
+    middleware = captured["served_app"]
+
+    http_scope = {"type": "http", "state": {}}
+    ws_scope = {"type": "websocket"}
+
+    async def _receive():
+        return {"type": "http.disconnect"}
+
+    async def _send(_message):
+        return None
+
+    called = {"http": False, "websocket": False}
+
+    async def inner(scope, receive, send):
+        called[scope["type"]] = True
+
+    middleware._inner = inner
+    await middleware(http_scope, _receive, _send)
+    await middleware(ws_scope, _receive, _send)
+
+    assert called["http"] is True
+    assert called["websocket"] is True
+    assert "request_timestamp" in http_scope["state"]
+    assert isinstance(http_scope["state"]["request_timestamp"], float)
+    assert "request_timestamp" not in ws_scope.get("state", {})
+
+
+@pytest.mark.parametrize(
+    ("path", "payload"),
+    [
+        ("/v1/audio/speech/stream", {"type": "error", "message": "Streaming speech is not available"}),
+        ("/v1/video/chat/stream", {"type": "error", "message": "Streaming video chat is not available"}),
+        ("/v1/realtime/video", {"type": "error", "message": "Streaming video generation is not available"}),
+        ("/v1/realtime", {"type": "error", "error": "Realtime API is not available", "code": "unsupported"}),
+        (
+            "/v1/realtime/robot/openpi",
+            {"type": "error", "error": "Robot policy not available", "code": "unsupported"},
+        ),
+        ("/v1/duplex", {"type": "error", "error": "Duplex API is not available", "code": "unsupported"}),
+    ],
+)
+def test_websocket_routes_emit_stable_unavailable_frames_and_close(path: str, payload: dict[str, str]) -> None:
+    """Lock missing-handler WebSocket capability probes.
+
+    Fails if unavailable JSON changes (keys/wording/code) or the socket stays
+    open after the error frame — common when WS wrappers move between modules.
+    """
+    app = FastAPI()
+    app.include_router(api_server.router)
+
+    with TestClient(app) as client:
+        with client.websocket_connect(path) as websocket:
+            assert websocket.receive_json() == payload
+            with pytest.raises(WebSocketDisconnect):
+                websocket.receive_text()
+
+
+def test_health_without_engine_returns_stable_unhealthy_response() -> None:
+    """Lock ``/health`` with no engine initialized.
+
+    Fails if status drifts from 503 or the unhealthy payload shape/reason
+    string changes during health-route extraction.
+    """
+    app = FastAPI()
+    response = asyncio.run(api_server.health(_request_for(app, path="/health")))
+
+    assert response.status_code == 503
+    assert json.loads(response.body) == {
+        "status": "unhealthy",
+        "reason": "No engine initialized",
+    }
+
+
+def test_models_without_handler_returns_empty_openai_list() -> None:
+    """Lock ``/v1/models`` when ``openai_serving_models`` is unset.
+
+    Fails if the empty listing becomes an error, or the OpenAI list envelope
+    changes (``object`` / ``data``).
+    """
+    app = FastAPI()
+    response = asyncio.run(api_server.show_available_models(_request_for(app, path="/v1/models")))
+
+    assert response.status_code == 200
+    assert json.loads(response.body) == {"object": "list", "data": []}
+
+
+def test_images_generation_without_engine_preserves_service_unavailable_error() -> None:
+    """Lock images generation when the multi-stage engine is missing.
+
+    Fails if the 503 detail string changes — clients/tests use it as a
+    capability signal when bootstrap failed.
+    """
+    app = FastAPI()
+    raw_request = _request_for(app, method="POST", path="/v1/images/generations")
+    request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model")
+
+    with pytest.raises(HTTPException) as exc_info:
+        asyncio.run(api_server.generate_images(request, raw_request))
+
+    assert exc_info.value.status_code == 503
+    assert exc_info.value.detail == "Multi-stage engine not initialized. Start server with a multi-stage omni model."
+
+
+def test_images_generation_without_multistage_chat_handler_preserves_unavailable_error(monkeypatch) -> None:
+    """Lock images generation when chat serving was not wired.
+
+    Fails if multi-stage images stop requiring ``openai_serving_chat``, or the
+    503 detail changes after images/chat ownership splits.
+    """
+    app = FastAPI()
+    app.state.engine_client = SimpleNamespace(stage_configs=[object(), object()])
+    app.state.stage_configs = app.state.engine_client.stage_configs
+    app.state.openai_serving_models = SimpleNamespace(base_model_paths=[SimpleNamespace(name="demo-model")])
+    app.state.openai_serving_chat = None
+    app.state.args = SimpleNamespace(max_generated_image_size=None)
+
+    raw_request = _request_for(app, method="POST", path="/v1/images/generations")
+    request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model")
+    monkeypatch.setattr(api_server, "get_stage_type", lambda _stage_cfg: "diffusion")
+
+    with pytest.raises(HTTPException) as exc_info:
+        asyncio.run(api_server.generate_images(request, raw_request))
+
+    assert exc_info.value.status_code == 503
+    assert exc_info.value.detail == "openai_serving_chat is not initialized for multi-stage image generation."
+
+
+def test_speech_without_handler_preserves_not_found_http_error() -> None:
+    """Lock speech HTTP when ``openai_serving_speech`` is unset.
+
+    Fails if missing speech becomes a different status/detail after TTS route
+    wrappers move out of ``api_server``.
+    """
+    app = FastAPI()
+    app.state.openai_serving_speech = None
+    raw_request = _request_for(app, method="POST", path="/v1/audio/speech")
+    request = api_server.OpenAICreateSpeechRequest(input="hello", model="demo-model", voice="alloy")
+
+    with pytest.raises(HTTPException) as exc_info:
+        asyncio.run(api_server.create_speech(request, raw_request))
+
+    assert exc_info.value.status_code == 404
+    assert exc_info.value.detail == "The model does not support Speech API"
+
+
+def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatch) -> None:
+    """Lock shared ``EngineGenerateError`` JSON fields used by Omni handlers.
+
+    Fails if ``request_id`` / ``error_stage_id`` are dropped when error helpers
+    move (P0.2) or exception-handler wiring changes.
+
+    Exercises the envelope through the registered handler, not the private
+    ``_create_engine_error_json_response`` helper, so relocating that helper
+    in P0.2 is not a false red.
+    """
+    app = FastAPI()
+    # TODO(P0.2): retarget this call if ``_register_omni_exception_handlers``
+    # leaves api_server (red here is a rename, not a surface regression).
+    api_server._register_omni_exception_handlers(app)
+    app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False))
+    app.state.server = object()
+    app.state.args = SimpleNamespace(log_error_stack=False)
+
+    req = _request_for(app, method="POST", path="/v1/chat/completions")
+    req.state.request_metadata = SimpleNamespace(request_id="req-123")
+
+    monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None)
+
+    exc = EngineGenerateError("boom")
+    exc.error_stage_id = 2  # type: ignore[attr-defined]
+    handler = app.exception_handlers[EngineGenerateError]
+    response = asyncio.run(handler(req, exc))
+    payload = json.loads(response.body)
+
+    assert "error" in payload
+    assert payload["error"]["request_id"] == "req-123"
+    assert payload["error"]["error_stage_id"] == 2
+    assert response.status_code >= 400
+
+
+def test_engine_dead_error_handler_registered_returns_json(monkeypatch) -> None:
+    """Lock Omni ``EngineDeadError`` handler registration + JSON envelope.
+
+    Fails if registration is skipped or the handler stops returning OpenAI-style
+    JSON with ``request_id``.
+    """
+    app = FastAPI()
+    # TODO(P0.2): retarget this call if ``_register_omni_exception_handlers``
+    # leaves api_server (red here is a rename, not a surface regression).
+    api_server._register_omni_exception_handlers(app)
+    app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False))
+    app.state.server = object()
+    app.state.args = SimpleNamespace(log_error_stack=False)
+
+    monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None)
+
+    handler = app.exception_handlers[EngineDeadError]
+    req = _request_for(app)
+    req.state.request_metadata = SimpleNamespace(request_id="dead-1")
+    response = asyncio.run(handler(req, EngineDeadError()))
+    payload = json.loads(response.body)
+    assert payload["error"]["request_id"] == "dead-1"
+    assert response.status_code >= 400
+
+
+@pytest.mark.asyncio
+async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None:
+    """Lock pure-diffusion ``app.state`` keys after init, including live vs None.
+
+    Fails if bootstrap drops keys route owners still read (e.g. video/speech/
+    models), wires a handler that must stay None, or leaves a required
+    handler as None.
+    """
+    stage = SimpleNamespace(engine_args={})
+    engine = _FakeEngineClient(stage_configs=[stage])
+
+    def _for_diffusion_factory(label: str):
+        @classmethod
+        def _factory(cls, *args, **kwargs):
+            return _marker(label)
+
+        return _factory
+
+    monkeypatch.setattr(api_server, "get_stage_type", lambda _cfg: "diffusion")
+    monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat"))
+    monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch"))
+    monkeypatch.setattr(
+        api_server.OmniOpenAIServingAudioGenerate,
+        "for_diffusion",
+        _for_diffusion_factory("audio_generate"),
+    )
+    monkeypatch.setattr(api_server.OmniOpenAIServingVideo, "for_diffusion", _for_diffusion_factory("video"))
+    monkeypatch.setattr(api_server.OmniStreamingVideoOutputHandler, "__init__", lambda self, *a, **k: None)
+    monkeypatch.setattr(api_server.OmniOpenAIServingSpeech, "for_diffusion", _for_diffusion_factory("speech"))
+    monkeypatch.setattr(
+        api_server.ServingRealtimeRobotOpenPI,
+        "create_policy_server",
+        classmethod(lambda cls, *a, **k: _marker("openpi")),
+    )
+
+    state = State()
+    await api_server.omni_init_app_state(engine, state, _minimal_args())
+
+    _assert_app_state_snapshot(
+        state,
+        expected_keys=_DIFFUSION_APP_STATE_KEYS,
+        must_be_wired=_DIFFUSION_MUST_BE_WIRED,
+        must_be_none=_DIFFUSION_MUST_BE_NONE,
+    )
+    assert state.diffusion_engine is engine
+
+
+@pytest.mark.asyncio
+async def test_multistage_app_state_key_snapshot(monkeypatch) -> None:
+    """Lock multi-stage ``app.state`` keys after init, including live vs None.
+
+    Fails if chat/speech/video/realtime/tokenization keys disappear or are
+    left as None after init refactors — classic “route still registered,
+    handler never wired” failure mode.
+    """
+    engine = _FakeEngineClient(
+        stage_configs=[object(), object()],
+        vllm_config=SimpleNamespace(
+            lora_config=None,
+            model_config=SimpleNamespace(),
+            parallel_config=SimpleNamespace(_api_process_rank=0),
+        ),
+    )
+
+    class _FakeModels:
+        def __init__(self, *args, **kwargs):
+            self.base_model_paths = kwargs.get("base_model_paths") or []
+
+        async def init_static_loras(self):
+            return None
+
+    class _FakeCtor:
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def warmup(self):
+            return None
+
+    class _FakeSpeech(_FakeCtor):
+        async def warmup(self):
+            return None
+
+    monkeypatch.setattr(api_server, "load_chat_template", lambda *_a, **_k: None)
+    monkeypatch.setattr(api_server, "process_lora_modules", lambda modules, _defaults: modules or [])
+    monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels)
+    monkeypatch.setattr(api_server, "OnlineRenderer", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingResponses", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingChat", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingChatBatch", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingCompletion", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingPooling", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingEmbedding", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingClassification", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingScores", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingTokenization", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingTranscription", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingTranslation", _FakeCtor)
+    monkeypatch.setattr(api_server, "AnthropicServingMessages", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingTokens", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingSpeech", _FakeSpeech)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingAudioGenerate", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniStreamingSpeechHandler", _FakeCtor)
+    monkeypatch.setattr(api_server, "create_streaming_video_handler", lambda **_k: _marker("streaming_video"))
+    monkeypatch.setattr(api_server, "OpenAIServingRealtime", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingVideo", _FakeCtor)
+    monkeypatch.setattr(api_server, "should_enable_duplex_endpoint", lambda *_a, **_k: False)
+
+    state = State()
+    await api_server.omni_init_app_state(engine, state, _minimal_args())
+
+    _assert_app_state_snapshot(
+        state,
+        expected_keys=_MULTISTAGE_APP_STATE_KEYS,
+        must_be_wired=_MULTISTAGE_MUST_BE_WIRED,
+        must_be_none=_MULTISTAGE_MUST_BE_NONE,
+    )

tests/entrypoints/openai_api/test_tts_adapter.py

@@ -22,6 +22,10 @@ from vllm_omni.entrypoints.openai.tts_adapters.indextts2 import (
     IndexTTS25Adapter,
     indextts2_conditioning_cache_salt,
 )
+from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import (
+    MossTTSAdapter,
+    MossTTSNanoAdapter,
+)
 from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter
 from vllm_omni.model_executor.models.indextts2 import prompt_utils
 from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
@@ -94,6 +98,20 @@ def test_all_adapters_are_ar_or_diffusion():
         assert cls.backend in ("ar", "diffusion")
 
 
+@pytest.mark.parametrize("adapter_cls", [MossTTSAdapter, MossTTSNanoAdapter])
+def test_moss_tts_applies_request_max_new_tokens(adapter_cls):
+    adapter = adapter_cls(SimpleNamespace(server=object()))
+    stage_defaults = [SimpleNamespace(max_tokens=4096)]
+
+    overridden = adapter.apply_sampling_overrides(
+        stage_defaults,
+        SimpleNamespace(max_new_tokens=512),
+    )
+
+    assert overridden[0].max_tokens == 512
+    assert stage_defaults[0].max_tokens == 4096
+
+
 def test_qwen3_tts_metadata():
     assert Qwen3TTSAdapter.backend == "ar"
     assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)

tests/entrypoints/openai_api/test_video_server.py

@@ -2168,6 +2168,31 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
     assert captured.quality == "high"
 
 
+def test_sync_sana_wm_extra_params_payload_passes_to_engine_prompt(test_client, mocker: MockerFixture):
+    _mock_encode_video_bytes(mocker)
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={
+            "prompt": "drive forward",
+            "extra_params": json.dumps(
+                {
+                    "sana_wm": {"action": "d-4", "rotation_speed_deg": 1.5},
+                    "sana_wm_native_max_tokens": 30000,
+                }
+            ),
+        },
+    )
+
+    assert response.status_code == 200
+    engine = test_client.app.state.openai_serving_video._engine_client
+    captured_params = engine.captured_sampling_params_list[0]
+    # The camera block reaches the model through extra_args; the Sana preprocess
+    # hook is what lifts it onto the prompt.
+    assert captured_params.extra_args["sana_wm"]["action"] == "d-4"
+    assert captured_params.extra_args["sana_wm"]["rotation_speed_deg"] == 1.5
+    assert captured_params.extra_args["sana_wm_native_max_tokens"] == 30000
+
+
 def test_sync_frame_interpolation_params_pass_to_sampling_params(test_client, mocker: MockerFixture):
     """Frame interpolation parameters should be forwarded on the sync path."""
     encode_mock = _mock_encode_video_bytes(mocker)

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -58,6 +58,23 @@ def test_default_stage_devices_from_sequence_parallel():
     assert stage_cfg["runtime"]["devices"] == "0,1,2,3"
 
 
+def test_default_stage_devices_and_dp_from_num_gpus():
+    """Resolve omitted DP before deriving devices from an explicit GPU count."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "num_gpus": 8,
+            "tensor_parallel_size": 2,
+            "ulysses_degree": 2,
+        }
+    )[0]
+
+    parallel_config = stage_cfg["engine_args"]["parallel_config"]
+    assert parallel_config.data_parallel_size == 2
+    assert parallel_config.world_size == 8
+    assert stage_cfg["engine_args"]["num_gpus"] == 8
+    assert stage_cfg["runtime"]["devices"] == "0,1,2,3,4,5,6,7"
+
+
 def test_default_stage_config_uses_parallel_size_kwargs():
     """Ensure default diffusion parallel_config uses CLI/API parallel sizes."""
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
@@ -76,8 +93,8 @@ def test_default_stage_config_uses_parallel_size_kwargs():
     assert parallel_config.enable_expert_parallel is True
 
 
-def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
+def test_default_stage_config_preserves_omitted_dp_for_runtime_inference():
-    """Ensure nullified diffusion parallel-size kwargs fall back to defaults."""
+    """Keep omitted DP unresolved until the runtime WORLD size is known."""
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
         {
             "pipeline_parallel_size": None,
@@ -92,7 +109,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
 
     parallel_config = stage_cfg["engine_args"]["parallel_config"]
     assert parallel_config.pipeline_parallel_size == 1
-    assert parallel_config.data_parallel_size == 1
+    assert parallel_config.data_parallel_size is None
     assert parallel_config.tensor_parallel_size == 1
     assert parallel_config.enable_expert_parallel is False
     assert stage_cfg["engine_args"]["enforce_eager"] is False

tests/model_executor/model_loader/test_resolve_model_path.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for resolve_model_to_local_path."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import huggingface_hub
+import pytest
+from huggingface_hub.errors import LocalEntryNotFoundError
+
+from vllm_omni.model_executor.model_loader import weight_utils
+from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Both resolution paths are patched in every test, so nothing is ever fetched
+REPO_ID = "test-org/test-model"
+RESOLVED = "/cache/snapshots/deadbeef"
+
+
+@pytest.fixture
+def paths(mocker):
+    """The resolver's two ways out, both patched."""
+    return SimpleNamespace(
+        cache_only=mocker.patch.object(huggingface_hub, "snapshot_download", return_value=RESOLVED),
+        download=mocker.patch.object(weight_utils, "download_weights_from_hf_specific", return_value=RESOLVED),
+    )
+
+
+@pytest.mark.parametrize("allow_download", [False, True])
+def test_local_directory_short_circuits(tmp_path, paths, allow_download):
+    """A local dir resolves to itself without either path running."""
+    assert resolve_model_to_local_path(str(tmp_path), allow_download=allow_download) == str(tmp_path)
+
+    paths.cache_only.assert_not_called()
+    paths.download.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "allow_patterns", "cache_dir"),
+    [
+        ({}, None, None),
+        ({"allow_patterns": ["subdir/*"], "cache_dir": "/cache"}, ["subdir/*"], "/cache"),
+    ],
+)
+def test_repo_id_is_cache_only_by_default(paths, kwargs, allow_patterns, cache_dir):
+    """Without `allow_download` nothing is fetched, and narrowing is forwarded."""
+    assert resolve_model_to_local_path(REPO_ID, **kwargs) == RESOLVED
+
+    paths.cache_only.assert_called_once_with(
+        REPO_ID,
+        cache_dir=cache_dir,
+        allow_patterns=allow_patterns,
+        local_files_only=True,
+    )
+    paths.download.assert_not_called()
+
+
+def test_cache_miss_propagates(paths):
+    """A failed resolution raises instead of degrading into the raw repo id."""
+    paths.cache_only.side_effect = LocalEntryNotFoundError("cold cache")
+
+    with pytest.raises(LocalEntryNotFoundError):
+        resolve_model_to_local_path(REPO_ID)
+
+    paths.download.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "allow_patterns"),
+    [
+        ({}, ["*"]),
+        ({"allow_patterns": ["subdir/model*.safetensors"]}, ["subdir/model*.safetensors"]),
+    ],
+)
+def test_download_routes_through_the_locked_helper(paths, kwargs, allow_patterns):
+    """Downloads go through the vLLM helper holding the node-wide lock.
+
+    That helper is also what honors `VLLM_USE_MODELSCOPE`;
+    calling `snapshot_download` directly would drop both.
+    """
+    assert resolve_model_to_local_path(REPO_ID, allow_download=True, cache_dir="/cache", **kwargs) == RESOLVED
+
+    paths.download.assert_called_once_with(
+        model_name_or_path=REPO_ID,
+        cache_dir="/cache",
+        allow_patterns=allow_patterns,
+        require_all=True,
+    )
+    paths.cache_only.assert_not_called()

tests/model_executor/models/ming_tts/test_ming_shared_modules.py

@@ -7,8 +7,14 @@ import pytest
 import torch
 
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
+from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
 from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig
-from vllm_omni.model_executor.models.common.ming.fm import Solver
+from vllm_omni.model_executor.models.common.ming.cfm_cudagraph import CFMGraphExecutor, CFMSampler
+from vllm_omni.model_executor.models.common.ming.cfm_dit import DiT
+from vllm_omni.model_executor.models.common.ming.cfm_head import FlowLoss
+from vllm_omni.model_executor.models.common.ming.cfm_solver import Solver, get_epss_timesteps
+from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor
+from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor
 from vllm_omni.model_executor.models.ming_tts.constants import (
     AGGREGATOR_HIDDEN_SIZE,
     HISTORY_PATCH_SIZE,
@@ -24,7 +30,6 @@ from vllm_omni.model_executor.models.ming_tts.constants import (
     SPEAKER_EMBEDDING_DIM,
     VAE_PATCH_SIZE,
 )
-from vllm_omni.model_executor.models.ming_tts.flowloss_head import FlowLoss
 from vllm_omni.model_executor.models.ming_tts.ming_tts_llm import MingLLMModel
 from vllm_omni.model_executor.models.ming_tts.speaker_extractor import _resolve_speaker_embeddings
 from vllm_omni.model_executor.models.ming_tts.validation import validate_ming_tts_config
@@ -44,7 +49,87 @@ def test_ming_tts_audio_vae_uses_common_config():
 
 def test_ming_tts_cfm_solver_uses_common_implementation():
     """Ming dense imports the shared solver implementation directly."""
-    assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.fm"
+    assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.cfm_solver"
+
+
+@pytest.mark.parametrize(
+    ("symbol", "expected_module"),
+    [
+        (Aggregator, "vllm_omni.model_executor.models.common.ming.aggregator"),
+        (DiT, "vllm_omni.model_executor.models.common.ming.cfm_dit"),
+        (FlowLoss, "vllm_omni.model_executor.models.common.ming.cfm_head"),
+        (CFMGraphExecutor, "vllm_omni.model_executor.models.common.ming.cfm_cudagraph"),
+        (CFMSampler, "vllm_omni.model_executor.models.common.ming.cfm_cudagraph"),
+        (SpeakerEmbeddingExtractor, "vllm_omni.model_executor.models.common.ming.speaker_extractor"),
+        (SpkembExtractor, "vllm_omni.model_executor.models.common.ming.spk_embedding"),
+    ],
+)
+def test_ming_tts_reusable_modules_live_in_common(symbol, expected_module):
+    """The pieces a second Ming-family TTS model reuses live under common/ming.
+
+    Guards the promotion: if an implementation drifts back into `ming_tts`
+    (or into another model package), a second model would have to reach across
+    packages to get it.
+    """
+    assert symbol.__module__ == expected_module
+
+
+def test_ming_flash_talker_uses_the_shared_aggregator():
+    """Ming flash no longer carries its own copy of the aggregator.
+
+    Both Ming packages ran near-identical aggregators one directory apart,
+    which is the duplication this module layout exists to remove.
+    """
+    from vllm_omni.model_executor.models.ming_flash_omni import talker_module
+
+    assert talker_module.Aggregator is Aggregator
+    assert not any(
+        isinstance(obj, type) and obj.__module__ == talker_module.__name__ and name == "Aggregator"
+        for name, obj in vars(talker_module).items()
+    )
+
+
+def test_shared_aggregator_matches_the_ming_flash_construction():
+    """The shared aggregator reproduces Ming flash's former default geometry.
+
+    `talker_module.Aggregator` defaulted to `in_channels=64` while the
+    shared class defaults to 4, so the call site now passes `latent_dim`
+    explicitly. This pins the resulting shapes rather than the default.
+    """
+    latent_dim, llm_hidden = 64, 32
+    aggregator = Aggregator(
+        in_channels=latent_dim,
+        llm_input_dim=llm_hidden,
+        hidden_size=16,
+        depth=1,
+        num_heads=2,
+    )
+
+    assert aggregator.x_embedder.in_features == latent_dim
+    assert aggregator.final_layer.linear.out_features == llm_hidden
+
+    out = aggregator(torch.randn(1, 3, latent_dim))
+    # Keeps only the CLS projection: [B, T, D] -> [B, 1, llm_hidden].
+    assert out.shape == (1, 1, llm_hidden)
+
+
+def test_shared_aggregator_rejects_mismatched_latent_width():
+    """Wrong-width input fails loudly instead of silently projecting incorrect."""
+    aggregator = Aggregator(in_channels=64, llm_input_dim=32, hidden_size=16, depth=1, num_heads=2)
+
+    with pytest.raises(ValueError, match="feature dim mismatch"):
+        aggregator(torch.randn(1, 3, 4))
+
+
+def test_cfm_solver_owns_the_timestep_schedule():
+    """get_epss_timesteps sits with the solver, not with the DiT blocks."""
+    from vllm_omni.model_executor.models.common.ming import dit_blocks
+
+    assert not hasattr(dit_blocks, "get_epss_timesteps")
+
+    t = get_epss_timesteps(10, device="cpu", dtype=torch.float32)
+    assert t.shape == (11,)
+    assert float(t[0]) == 0.0 and float(t[-1]) == 1.0
 
 
 def test_ming_tts_flowloss_preserves_checkpoint_prefix():

tests/model_executor/models/minicpmo_4_5/test_streaming_audio_cache.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MiniCPM-o 4.5 streaming audio cache compatibility."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from transformers.models.whisper.modeling_whisper import WhisperConfig
+
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
+    WHISPER_ATTENTION_CLASSES,
+    MiniCPMWhisperEncoder,
+    MiniCPMWhisperEncoderLayer,
+    _get_audio_cache_length,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _LegacyWhisperAttention(nn.Module):
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.seen_cache = None
+
+    def forward(self, hidden_states, past_key_value=None, **kwargs):
+        self.seen_cache = past_key_value
+        return hidden_states, None, past_key_value
+
+
+class _CurrentWhisperAttention(nn.Module):
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.seen_cache = None
+
+    def forward(self, hidden_states, past_key_values=None, **kwargs):
+        self.seen_cache = past_key_values
+        return hidden_states, None
+
+
+@pytest.mark.parametrize(
+    ("implementation", "attention_cls"),
+    [
+        ("test_legacy_cache", _LegacyWhisperAttention),
+        ("test_current_cache", _CurrentWhisperAttention),
+    ],
+)
+def test_whisper_attention_preserves_streaming_cache(
+    monkeypatch: pytest.MonkeyPatch,
+    implementation: str,
+    attention_cls: type[nn.Module],
+) -> None:
+    monkeypatch.setitem(WHISPER_ATTENTION_CLASSES, implementation, attention_cls)
+    config = WhisperConfig(
+        d_model=4,
+        encoder_attention_heads=1,
+        encoder_ffn_dim=8,
+        encoder_layers=1,
+    )
+    config._attn_implementation = implementation
+    layer = MiniCPMWhisperEncoderLayer(config, layer_idx=0)
+    cache = object()
+
+    outputs = layer(
+        torch.randn(1, 2, config.d_model),
+        attention_mask=None,
+        layer_head_mask=None,
+        past_key_values=cache,
+        use_cache=True,
+    )
+
+    assert layer.self_attn.seen_cache is cache
+    assert outputs[-1] is cache
+
+
+def test_audio_cache_length_uses_current_cache_api() -> None:
+    dynamic_cache = SimpleNamespace(get_seq_length=lambda: 7)
+    encoder_decoder_cache = SimpleNamespace(self_attention_cache=dynamic_cache)
+
+    assert _get_audio_cache_length(encoder_decoder_cache) == 7
+
+
+def test_audio_cache_length_supports_legacy_cache() -> None:
+    key = torch.zeros(1, 1, 5, 2)
+    value = torch.zeros_like(key)
+
+    assert _get_audio_cache_length(((key, value),)) == 5
+
+
+def _tiny_encoder_config(attn_implementation: str) -> WhisperConfig:
+    config = WhisperConfig(
+        num_mel_bins=16,
+        d_model=32,
+        encoder_layers=2,
+        encoder_attention_heads=2,
+        encoder_ffn_dim=64,
+        max_source_positions=100,
+        dropout=0.0,
+        attention_dropout=0.0,
+    )
+    config._attn_implementation = attn_implementation
+    return config
+
+
+@pytest.mark.parametrize("attn_implementation", ["eager", "sdpa"])
+def test_encoder_reuses_cache_across_streaming_chunks(attn_implementation: str) -> None:
+    """Encode two chunks against the installed Whisper attention and cache classes.
+
+    Unlike the stub-based tests above, this drives the real transformers cache
+    API that ``get_audio_embedding_streaming`` relies on, so it fails whenever a
+    removed cache method or renamed attention kwarg breaks streaming reuse.
+    """
+    torch.manual_seed(0)
+    config = _tiny_encoder_config(attn_implementation)
+    encoder = MiniCPMWhisperEncoder(config).eval()
+
+    first_frames, second_frames = 20, 14
+    first_len = (first_frames - 1) // 2 + 1
+    second_len = (second_frames - 1) // 2 + 1
+    first_mel = torch.randn(1, config.num_mel_bins, first_frames)
+    second_mel = torch.randn(1, config.num_mel_bins, second_frames)
+
+    with torch.no_grad():
+        first = encoder(
+            first_mel,
+            attention_mask=torch.zeros(1, 1, first_len, first_len),
+            past_key_values=None,
+            use_cache=True,
+            output_hidden_states=True,
+        )
+        cache = first.past_key_values
+        assert cache is not None
+        assert _get_audio_cache_length(cache) == first_len
+
+        second = encoder(
+            second_mel,
+            attention_mask=torch.zeros(1, 1, second_len, first_len + second_len),
+            past_key_values=cache,
+            use_cache=True,
+            output_hidden_states=True,
+        )
+        assert _get_audio_cache_length(second.past_key_values) == first_len + second_len
+
+        stateless = encoder(
+            second_mel,
+            attention_mask=torch.zeros(1, 1, second_len, second_len),
+            past_key_values=None,
+            use_cache=True,
+            output_hidden_states=True,
+        )
+
+    assert torch.isfinite(second.last_hidden_state).all()
+    assert not torch.allclose(second.last_hidden_state, stateless.last_hidden_state)

tests/model_executor/models/moss_tts/test_compact_topk.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU regression tests for MOSS-TTS compact top-k sampling."""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts import modeling_moss_tts_local
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
+
+
+def test_compact_topk_only_samples_retained_candidates() -> None:
+    logits = torch.randn(32, 1024, generator=torch.Generator().manual_seed(7))
+
+    sampled = modeling_moss_tts_local._sample_token(
+        logits,
+        temperature=1.7,
+        top_k=25,
+        top_p=0.8,
+        do_sample=True,
+        generator=torch.Generator().manual_seed(42),
+    )
+
+    retained = logits.topk(25, dim=-1).indices
+    assert torch.all((retained == sampled.unsqueeze(-1)).any(dim=-1))
+
+
+def test_compact_topk_tiny_nucleus_selects_argmax() -> None:
+    logits = torch.randn(32, 1024, generator=torch.Generator().manual_seed(17))
+
+    sampled = modeling_moss_tts_local._sample_token(
+        logits,
+        temperature=1.7,
+        top_k=25,
+        top_p=1e-9,
+        do_sample=True,
+        generator=torch.Generator().manual_seed(42),
+    )
+
+    torch.testing.assert_close(sampled, logits.argmax(dim=-1))

tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU parity tests for MOSS-TTS Local Depth RoPE caching."""
+
+import pytest
+import torch
+from transformers import GPT2Config
+
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+    MossTTSLocalDepthTransformer,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
+
+
+def _test_config() -> GPT2Config:
+    return GPT2Config(
+        n_embd=32,
+        n_head=4,
+        n_inner=64,
+        layer_norm_epsilon=1e-6,
+        rope_base=10000.0,
+    )
+
+
+def test_rope_cache_matches_reference_and_reuses_storage() -> None:
+    model = MossTTSLocalDepthTransformer(_test_config()).eval()
+    attn = model.h[0].attn
+    device = torch.device("cpu")
+
+    attn.prepare_rope_cache(12, device, torch.float32)
+    cache_ptr = attn._rope_cos_cache.data_ptr()
+    cos, sin = attn._rope_cos_sin(4, device, torch.float32)
+
+    positions = torch.arange(4, dtype=torch.float32)
+    freqs = torch.einsum("s,d->sd", positions, attn.inv_freq.float())
+    expected_cos = freqs.cos().repeat_interleave(2, dim=-1).view(1, 4, 1, attn.head_dim)
+    expected_sin = freqs.sin().repeat_interleave(2, dim=-1).view(1, 4, 1, attn.head_dim)
+    torch.testing.assert_close(cos, expected_cos)
+    torch.testing.assert_close(sin, expected_sin)
+
+    # Smaller slices retain the fixed allocation used by CUDA Graph.
+    attn.prepare_rope_cache(6, device, torch.float32)
+    assert attn._rope_cos_cache.data_ptr() == cache_ptr

tests/model_executor/models/moss_tts/test_talker_logits_processor.py

@@ -23,6 +23,11 @@ class _RecordingLogitsProcessor:
         return torch.zeros((1, 4))
 
 
+class _EmbeddingBackbone(nn.Module):
+    def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return input_ids.to(dtype=torch.float32).unsqueeze(-1).repeat(1, 4)
+
+
 def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> None:
     from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
         MossTTSDelayTalkerForGeneration,
@@ -41,3 +46,74 @@ def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> N
     assert logits.shape == (1, 4)
     assert model.logits_processor.args == (model.text_lm_head, hidden_states)
     assert model.logits_processor.kwargs == {}
+
+
+def test_moss_tts_local_keeps_fixed_shape_rows_and_routes_stop_logits() -> None:
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSLocalTalkerForGeneration,
+    )
+
+    model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+    nn.Module.__init__(model)
+    model.audio_pad_token_id = 1024
+    model.audio_assistant_slot_token_id = 2
+    model.im_end_token_id = 3
+    model.text_vocab_size = 8
+    model.n_vq = 4
+    model._batch_state = None
+    model._batch_state_spans = None
+    model._batch_should_continue = None
+
+    valid = torch.tensor([[1, 2, 3, 4]], dtype=torch.long)
+    stopped = torch.full((1, 4), model.audio_pad_token_id, dtype=torch.long)
+    infos = [
+        {"audio_state": {}, "audio_codes": {"current": valid}},
+        {"audio_state": {}, "audio_codes": {"current": stopped}},
+    ]
+    spans = [(0, 3), (3, 4)]
+
+    output = model.make_omni_output(
+        torch.randn(4, 8),
+        model_intermediate_buffer=infos,
+        request_token_spans=spans,
+    )
+
+    audio = output.multimodal_outputs["codes"]["audio"]
+    assert len(audio) == 2
+    torch.testing.assert_close(audio[0], valid)
+    # Keeping this row avoids CUDA boolean-index shape discovery; the
+    # talker2codec raw async-chunk processor drops all-pad rows on CPU.
+    torch.testing.assert_close(audio[1], stopped)
+    assert model._batch_should_continue.tolist() == [True, False]
+
+    logits = model.compute_logits(torch.randn(4, 8))
+
+    assert logits is not None
+    assert logits.argmax(dim=-1).tolist() == [2, 2, 2, 3]
+
+
+def test_moss_tts_local_single_token_prefill_does_not_advance_audio() -> None:
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSLocalTalkerForGeneration,
+    )
+
+    model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+    nn.Module.__init__(model)
+    model.model = _EmbeddingBackbone()
+    model.n_vq = 4
+
+    input_ids = torch.tensor([7], dtype=torch.long)
+    _, embeds, update = model.preprocess(
+        input_ids,
+        None,
+        _omni_is_prefill=True,
+        audio_state={"is_stopping": True},
+        ref_offset=5,
+    )
+
+    torch.testing.assert_close(embeds, torch.full((1, 4), 7.0))
+    assert update == {
+        "audio_state": {"is_stopping": False},
+        "ref_offset": 6,
+    }
+    assert "mtp_inputs" not in update

tests/model_tests/diffusion/test_alignment.py

@@ -80,6 +80,7 @@ EXCLUDED_MODELS = [
     "Pi0Pipeline",
     "SoulXSingerPipeline",
     "SoulXSingerSVCPipeline",
+    "SanaWmPipeline",
 ]
 
 

vllm_omni/config/omni_config.py

@@ -600,12 +600,17 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
             * self.cfg_parallel_size
         )
         if self.use_hsdp:
-            if self.tensor_parallel_size > 1 or self.data_parallel_size > 1:
+            incompatible = []
-                raise ValueError(
+            if self.tensor_parallel_size > 1:
-                    "HSDP (use_hsdp=True) cannot be used with TP or DP "
+                incompatible.append("TP")
-                    f"(tensor_parallel_size={self.tensor_parallel_size}, "
+            if self.data_parallel_size > 1:
-                    f"data_parallel_size={self.data_parallel_size})"
+                incompatible.append("DP")
-                )
+            if self.pipeline_parallel_size > 1:
+                incompatible.append("PP")
+            if self.enable_expert_parallel:
+                incompatible.append("EP")
+            if incompatible:
+                raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible))
             if self.hsdp_shard_size == -1:
                 if other_parallel_world_size == 1:
                     raise ValueError("Cannot auto-calculate hsdp_shard_size when other parallelism is all 1")

vllm_omni/diffusion/data.py

@@ -154,8 +154,8 @@ class DiffusionParallelConfig:
     pipeline_parallel_size: int = 1
     """Number of pipeline parallel stages."""
 
-    data_parallel_size: int = 1
+    data_parallel_size: int | None = None
-    """Number of data parallel groups."""
+    """Number of data parallel groups. Inferred from WORLD when omitted."""
 
     tensor_parallel_size: int = 1
     """Number of tensor parallel groups."""
@@ -243,7 +243,7 @@ class DiffusionParallelConfig:
     def _validate_parallel_config(self) -> Self:
         """Validates the config relationships among the parallel strategies."""
         assert self.pipeline_parallel_size > 0, "Pipeline parallel size must be > 0"
-        assert self.data_parallel_size > 0, "Data parallel size must be > 0"
+        assert self.data_parallel_size is None or self.data_parallel_size > 0, "Data parallel size must be > 0"
         assert self.tensor_parallel_size > 0, "Tensor parallel size must be > 0"
         assert self.sequence_parallel_size > 0, "Sequence parallel size must be > 0"
         assert self.ulysses_degree > 0, "Ulysses degree must be > 0"
@@ -283,10 +283,12 @@ class DiffusionParallelConfig:
                 self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
             )
 
-        # Calculate world_size from other parallelism dimensions
+        # Until the runtime WORLD size is known, an omitted DP dimension means
+        # one replica. OmniDiffusionConfig resolves it against num_gpus below.
+        data_parallel_size = self.data_parallel_size or 1
         other_parallel_world_size = (
             self.pipeline_parallel_size
-            * self.data_parallel_size
+            * data_parallel_size
             * self.tensor_parallel_size
             * self.sequence_parallel_size
             * self.cfg_parallel_size
@@ -297,13 +299,17 @@ class DiffusionParallelConfig:
         # 1. Standalone: when other parallelism is all 1, HSDP determines world_size
         # 2. Combined: HSDP overlays on top of other parallelism
         if self.use_hsdp:
-            if self.tensor_parallel_size > 1 or self.data_parallel_size > 1:
+            incompatible = []
-                raise ValueError(
+            if self.tensor_parallel_size > 1:
-                    "HSDP (use_hsdp=True) cannot be used with TP or DP "
+                incompatible.append("TP")
-                    f"(tensor_parallel_size={self.tensor_parallel_size}, "
+            if data_parallel_size > 1:
-                    f"data_parallel_size={self.data_parallel_size}). "
+                incompatible.append("DP")
-                    "Set tensor_parallel_size=1 and data_parallel_size=1 when using HSDP."
+            if self.pipeline_parallel_size > 1:
-                )
+                incompatible.append("PP")
+            if self.enable_expert_parallel:
+                incompatible.append("EP")
+            if incompatible:
+                raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible))
             if self.hsdp_shard_size == -1:
                 # Auto-calculate: use other_parallel_world_size as shard_size
                 if self.hsdp_replicate_size <= 0:
@@ -338,6 +344,42 @@ class DiffusionParallelConfig:
         else:
             self.world_size = other_parallel_world_size
 
+    def resolve_data_parallel_size(self, world_size: int) -> int:
+        """Resolve or validate diffusion DP against the actual WORLD size."""
+        if world_size <= 0:
+            raise ValueError(f"WORLD size must be > 0, but got {world_size}")
+
+        if self.use_hsdp:
+            if self.data_parallel_size not in (None, 1):
+                raise ValueError("HSDP (FSDP2) requires data_parallel_size to be 1")
+            expected_world_size = self.hsdp_replicate_size * self.hsdp_shard_size
+            if world_size != expected_world_size:
+                raise ValueError(
+                    f"WORLD size ({world_size}) must equal HSDP size "
+                    f"({self.hsdp_replicate_size} x {self.hsdp_shard_size} = {expected_world_size})"
+                )
+            self.data_parallel_size = 1
+            self.world_size = world_size
+            return 1
+
+        non_dp_size = (
+            self.pipeline_parallel_size
+            * self.tensor_parallel_size
+            * self.sequence_parallel_size
+            * self.cfg_parallel_size
+        )
+        if world_size % non_dp_size != 0:
+            raise ValueError(f"WORLD size ({world_size}) must be divisible by non-DP parallel size ({non_dp_size})")
+        inferred_data_parallel_size = world_size // non_dp_size
+        if self.data_parallel_size is not None and self.data_parallel_size != inferred_data_parallel_size:
+            raise ValueError(
+                f"data_parallel_size ({self.data_parallel_size}) does not match WORLD-derived value "
+                f"({inferred_data_parallel_size})"
+            )
+        self.data_parallel_size = inferred_data_parallel_size
+        self.world_size = world_size
+        return inferred_data_parallel_size
+
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> "DiffusionParallelConfig":
         """
@@ -999,10 +1041,7 @@ class OmniDiffusionConfig:
             else:
                 self.num_gpus = 1
 
-        if self.num_gpus < self.parallel_config.world_size:
+        self.parallel_config.resolve_data_parallel_size(self.num_gpus)
-            raise ValueError(
-                f"num_gpus ({self.num_gpus}) < parallel_config.world_size ({self.parallel_config.world_size})"
-            )
 
         if self.diffusion_compile_granularity == "full":
             incompatible_features = []

vllm_omni/diffusion/distributed/__init__.py

@@ -6,11 +6,6 @@ from vllm_omni.diffusion.distributed.hsdp import (
     HSDPInferenceConfig,
     apply_hsdp_to_model,
 )
-from vllm_omni.diffusion.distributed.parallel_state import (
-    get_fs_group,
-    get_fully_shard_rank,
-    get_fully_shard_world_size,
-)
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelConfig,
     SequenceParallelInput,
@@ -46,8 +41,4 @@ __all__ = [
     "get_sharding_validator",
     "HSDPInferenceConfig",
     "apply_hsdp_to_model",
-    # FS (Fully Shard) utilities
-    "get_fs_group",
-    "get_fully_shard_rank",
-    "get_fully_shard_world_size",
 ]

vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py

@@ -9,8 +9,6 @@ import torch
 import torch.distributed as dist
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
-
 logger = init_logger(__name__)
 
 
@@ -45,9 +43,10 @@ class DistributedVaeExecutor:
     """
 
     def __init__(self):
-        self.group = get_dit_group()
+        # VAE tiles are scheduled across the complete worker WORLD.
-        self.world_size = dist.get_world_size(self.group)
+        self.group = None
-        self.rank = dist.get_rank(self.group)
+        self.world_size = dist.get_world_size()
+        self.rank = dist.get_rank()
         self.parallel_size = 1
         self.parallel_mode = "tile"
 
@@ -189,7 +188,7 @@ class DistributedVaeMixin:
         if self.distributed_executor.parallel_size > pp_size:
             logger.warning(
                 f"vae_patch_parallel_size={self.distributed_executor.parallel_size} "
-                f"is greater than dit_group={world_size};"
+                f"is greater than WORLD={world_size};"
-                f" using dit_group size={world_size}"
+                f" using WORLD size={world_size}"
             )
         return True

vllm_omni/diffusion/distributed/group_coordinator.py

@@ -991,6 +991,18 @@ class PipelineGroupCoordinator(GroupCoordinator):
     def _pipeline_isend_skip(self, tensor: torch.tensor):
         return torch.distributed.isend(tensor, dst=self.skip_rank, group=self.skip_device_group)
 
+    def destroy(self):
+        groups = [*self.device_groups, *self.cpu_groups, self.skip_device_group]
+        seen = {id(self.device_group), id(self.cpu_group)}
+        for group in groups:
+            if group is not None and id(group) not in seen:
+                torch.distributed.destroy_process_group(group)
+                seen.add(id(group))
+        self.device_groups = []
+        self.cpu_groups = []
+        self.skip_device_group = None
+        super().destroy()
+
 
 class SequenceParallelGroupCoordinator(GroupCoordinator):
     def __init__(
@@ -1031,3 +1043,15 @@ class SequenceParallelGroupCoordinator(GroupCoordinator):
         self.ring_rank = torch.distributed.get_rank(self.ring_group)
         self.allgather_world_size = torch.distributed.get_world_size(self.allgather_group)
         self.allgather_rank = torch.distributed.get_rank(self.allgather_group)
+
+    def destroy(self):
+        groups = [self.ulysses_group, self.ring_group, self.allgather_group]
+        seen = {id(self.device_group), id(self.cpu_group)}
+        for group in groups:
+            if group is not None and id(group) not in seen:
+                torch.distributed.destroy_process_group(group)
+                seen.add(id(group))
+        self.ulysses_group = None
+        self.ring_group = None
+        self.allgather_group = None
+        super().destroy()

vllm_omni/diffusion/distributed/hsdp.py

@@ -15,12 +15,7 @@ from torch.distributed.fsdp import (
 )
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.distributed.parallel_state import (
+from vllm_omni.diffusion.distributed.parallel_state import get_world_group
-    get_fs_group,
-    get_fully_shard_rank,
-    get_fully_shard_world_size,
-    get_world_group,
-)
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
@@ -54,36 +49,36 @@ class HSDPInferenceConfig:
 def _create_hsdp_mesh(
     device_type: str,
     replicate_size: int,
-    shard_pg: torch.distributed.ProcessGroup,
+    shard_size: int,
 ) -> DeviceMesh:
-    """Create a 2D DeviceMesh for HSDP using an existing ProcessGroup for the shard dimension.
+    """Create the FSDP2 DeviceMesh; it owns the HSDP process groups.
 
     Args:
         device_type: The device type (e.g., "cuda", "npu")
         replicate_size: Number of replica groups
-        shard_pg: The ProcessGroup for the shard dimension (from FS GroupCoordinator)
+        shard_size: Number of ranks in each FSDP shard group
 
     Returns:
         A 2D DeviceMesh with dimensions ("replicate", "shard")
     """
-    shard_size = torch.distributed.get_world_size(shard_pg)
     world_size = replicate_size * shard_size
+    actual_world_size = torch.distributed.get_world_size()
+    if world_size != actual_world_size:
+        raise ValueError(
+            f"HSDP mesh dimensions ({replicate_size} x {shard_size} = {world_size}) "
+            f"must equal WORLD size ({actual_world_size})"
+        )
 
     # Build 2D mesh tensor: shape (replicate_size, shard_size)
     # Ranks are arranged so that each row is a shard group
     mesh_tensor = torch.arange(world_size).reshape(replicate_size, shard_size)
 
-    # Create DeviceMesh with the shard ProcessGroup
-    # For the shard dimension, we reuse the existing FS ProcessGroup
     device_mesh = init_device_mesh(
         device_type,
         mesh_shape=(replicate_size, shard_size),
         mesh_dim_names=("replicate", "shard"),
     )
 
-    # Note: init_device_mesh creates new ProcessGroups internally.
-    # For consistency, we verify the mesh structure matches our FS group.
-    # In a future optimization, we could pass the existing ProcessGroups directly.
     logger.debug(
         "Created HSDP mesh: replicate_size=%d, shard_size=%d, mesh=%s",
         replicate_size,
@@ -119,34 +114,24 @@ def apply_hsdp_to_model(
     if not hsdp_config.enabled:
         raise ValueError("HSDP is not enabled in config")
 
-    # Use GroupCoordinator for distributed info
     world_group = get_world_group()
-    fs_group = get_fs_group()
-
     world_size = world_group.world_size
     rank = world_group.rank_in_group
-    fs_world_size = get_fully_shard_world_size()
-    fs_rank = get_fully_shard_rank()
 
     hsdp_replicate_size = hsdp_config.hsdp_replicate_size
     hsdp_shard_size = hsdp_config.hsdp_shard_size
 
-    # Validate that the FS group matches the HSDP shard size
+    if hsdp_replicate_size * hsdp_shard_size != world_size:
-    if fs_world_size != hsdp_shard_size:
         raise ValueError(
-            f"FS group world_size ({fs_world_size}) does not match "
+            f"HSDP mesh dimensions ({hsdp_replicate_size} x {hsdp_shard_size}) must equal WORLD size ({world_size})"
-            f"HSDP shard_size ({hsdp_shard_size}). "
-            "Ensure fully_shard_degree is set correctly in initialize_model_parallel."
         )
 
     logger.info(
-        "HSDP Inference: replicate_size=%d, shard_size=%d, world_size=%d, rank=%d, fs_world_size=%d, fs_rank=%d",
+        "HSDP Inference (FSDP2): replicate_size=%d, shard_size=%d, world_size=%d, rank=%d",
         hsdp_replicate_size,
         hsdp_shard_size,
         world_size,
         rank,
-        fs_world_size,
-        fs_rank,
     )
 
     # When the model contains FP8 parameters (online quantization), let FSDP
@@ -163,14 +148,14 @@ def apply_hsdp_to_model(
 
     device_type = current_omni_platform.device_type
 
-    # Create 2D DeviceMesh for HSDP using the FS group's ProcessGroup for shard dimension
+    # DeviceMesh is the single source of truth for FSDP2 shard/replicate groups
     # The mesh shape is (replicate, shard) where:
     # - replicate: groups of ranks that hold the same shard (for gradient all-reduce in training)
     # - shard: groups of ranks that each hold different shards (for parameter all-gather)
     device_mesh = _create_hsdp_mesh(
         device_type=device_type,
         replicate_size=hsdp_replicate_size,
-        shard_pg=fs_group.device_group,
+        shard_size=hsdp_shard_size,
     )
 
     hsdp_shard_conditions = getattr(model, "_hsdp_shard_conditions", None)

vllm_omni/diffusion/distributed/parallel_state.py

@@ -29,11 +29,11 @@ If you only need to use the distributed environment without model parallelism,
 """
 
 import inspect
+from math import prod
 
 import torch
 import torch.distributed
 import vllm.distributed.parallel_state as vllm_parallel_state
-from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion import envs
@@ -59,8 +59,6 @@ _SP: SequenceParallelGroupCoordinator | None = None
 _PP: PipelineGroupCoordinator | None = None
 _CFG: GroupCoordinator | None = None
 _DP: GroupCoordinator | None = None
-_FS: GroupCoordinator | None = None  # Fully Sharded (HSDP shard dimension)
-_DIT: GroupCoordinator | None = None
 
 # Rank-layout metadata for expert parallelism. This is not a process group;
 # it is reused by platform-specific runtimes that must build companion groups
@@ -184,7 +182,6 @@ class RankGenerator:
         pp: int,
         cfg: int,
         dp: int,
-        fs: int = 1,
         order: str = "tp-sp-pp-cfg-dp",
         rank_offset: int = 0,
     ) -> None:
@@ -193,9 +190,7 @@ class RankGenerator:
         self.pp = pp
         self.cfg = cfg
         self.dp = dp
-        self.fs = fs
         self.rank_offset = rank_offset
-        self.world_size = tp * sp * pp * cfg * dp
 
         self.name_to_size = {
             "tp": self.tp,
@@ -203,15 +198,10 @@ class RankGenerator:
             "pp": self.pp,
             "cfg": self.cfg,
             "dp": self.dp,
-            "fs": self.fs,
         }
         order = order.lower()
 
         for name in self.name_to_size.keys():
-            # Skip 'fs' validation - it's handled separately with independent_ranks=True
-            # and doesn't participate in the main orthogonal rank generation
-            if name == "fs":
-                continue
             if name not in order and self.name_to_size[name] != 1:
                 raise RuntimeError(
                     f"The size of ({name}) is ({self.name_to_size[name]}), "
@@ -234,7 +224,7 @@ class RankGenerator:
             mask[ordered_token.index(t)] = True
         return mask
 
-    def get_ranks(self, token, independent_ranks: bool = False):
+    def get_ranks(self, token):
         """Get rank group by input token.
 
         Arguments:
@@ -243,27 +233,11 @@ class RankGenerator:
                 to obtain multiple parallel types, we can use a hyphen
                 '-' to separate them. For example, if we want to obtain
                 the TP_DP group, the token should be 'tp-dp'.
-            independent_ranks (bool):
-                If True, generate independent rank groups that divide the world
-                into groups of the specified size. Used for FS (fully shard) groups
-                which operate independently from the main parallelism hierarchy.
         """
-        if independent_ranks and token == "fs":
-            # FS groups divide world into groups of size fs
-            # e.g., world_size=8, fs=4 -> [[0,1,2,3], [4,5,6,7]]
-            ranks = []
-            num_groups = self.world_size // self.fs
-            for i in range(num_groups):
-                group = list(range(i * self.fs + self.rank_offset, (i + 1) * self.fs + self.rank_offset))
-                ranks.append(group)
-            return ranks
-
         mask = self.get_mask(self.order, token)
-        ranks = generate_masked_orthogonal_rank_groups(self.world_size, self.ordered_size, mask)
+        ranks = generate_masked_orthogonal_rank_groups(prod(self.ordered_size), self.ordered_size, mask)
-        if self.rank_offset > 0:
+        if self.rank_offset:
-            for rank_group in ranks:
+            return [[rank + self.rank_offset for rank in rank_group] for rank_group in ranks]
-                for i in range(len(rank_group)):
-                    rank_group[i] += self.rank_offset
         return ranks
 
 
@@ -376,22 +350,6 @@ def get_data_parallel_rank():
     return get_dp_group().rank_in_group
 
 
-# FS (Fully Shard / HSDP shard dimension)
-def get_fs_group() -> GroupCoordinator:
-    assert _FS is not None, "fully shard group is not initialized"
-    return _FS
-
-
-def get_fully_shard_world_size():
-    """Return world size for the fully shard group."""
-    return get_fs_group().world_size
-
-
-def get_fully_shard_rank():
-    """Return my rank for the fully shard group."""
-    return get_fs_group().rank_in_group
-
-
 def is_dp_last_group():
     """Return True if in the last data parallel group, False otherwise."""
     return (
@@ -401,17 +359,6 @@ def is_dp_last_group():
     )
 
 
-def get_dit_world_size():
-    """Return world size for the DiT model."""
-    return (
-        get_data_parallel_world_size()
-        * get_classifier_free_guidance_world_size()
-        * get_sequence_parallel_world_size()
-        * get_pipeline_parallel_world_size()
-        * get_tensor_model_parallel_world_size()
-    )
-
-
 # * SET
 
 
@@ -498,7 +445,6 @@ def init_model_parallel_group(
         "expert",
         "sequence",
         "classifier_free_guidance",
-        "fully_shard",
     ], f"parallel_mode {parallel_mode} is not supported"
     if parallel_mode == "pipeline":
         return PipelineGroupCoordinator(
@@ -546,19 +492,6 @@ def init_vllm_model_parallel_group(
     )
 
 
-def init_dit_group(
-    dit_parallel_size: int,
-    backend: str,
-):
-    global _DIT
-    _DIT = torch.distributed.new_group(ranks=list(range(dit_parallel_size)), backend=backend)
-
-
-def get_dit_group():
-    assert _DIT is not None, "DIT group is not initialized"
-    return _DIT
-
-
 # adapted from https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py
 def set_seq_parallel_pg(
     sp_ulysses_degree: int,
@@ -737,8 +670,8 @@ def set_seq_parallel_pg(
     return ulysses_pg, ring_pg, allgather_pg
 
 
-def initialize_model_parallel(
+def _initialize_model_parallel(
-    data_parallel_size: int = 1,
+    data_parallel_size: int | None = None,
     cfg_parallel_size: int = 1,
     sequence_parallel_size: int | None = None,
     ulysses_degree: int = 1,
@@ -746,9 +679,8 @@ def initialize_model_parallel(
     allgather_degree: int = 1,
     tensor_parallel_size: int = 1,
     pipeline_parallel_size: int = 1,
-    fully_shard_degree: int = 1,
-    hsdp_replicate_size: int = 1,
     enable_expert_parallel: bool = False,
+    use_hsdp: bool = False,
     backend: str | None = None,
 ) -> None:
     if backend is None:
@@ -768,7 +700,6 @@ def initialize_model_parallel(
             (causal=False only). Mutually exclusive with ulysses/ring in v1.
         tensor_parallel_size: number of GPUs used for tensor parallelism.
         pipeline_parallel_size: number of GPUs used for pipeline parallelism.
-        fully_shard_degree: number of GPUs used for fully sharded data parallelism (HSDP shard dimension).
         backend: distributed backend of pytorch collective comm.
 
     Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
@@ -798,16 +729,20 @@ def initialize_model_parallel(
     ranks 8 to 15 belong to the second box.
     """
     # Get world size and rank. Ensure some consistencies.
-    assert torch.distributed.is_initialized()
+    if not torch.distributed.is_initialized():
+        raise RuntimeError("torch.distributed must be initialized before model parallel initialization")
+    if model_parallel_is_initialized():
+        raise RuntimeError("model parallel groups are already initialized")
     world_size: int = torch.distributed.get_world_size()
     backend = backend or torch.distributed.get_backend(get_world_group().device_group)
 
     if allgather_degree > 1:
-        assert ulysses_degree == 1 and ring_degree == 1, (
+        if ulysses_degree != 1 or ring_degree != 1:
-            "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. "
+            raise ValueError(
-            f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, "
+                "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. "
-            f"allgather_degree={allgather_degree}."
+                f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, "
-        )
+                f"allgather_degree={allgather_degree}."
+            )
 
     expected_sequence_parallel_size = allgather_degree if allgather_degree > 1 else ring_degree * ulysses_degree
     if sequence_parallel_size is None:
@@ -820,42 +755,38 @@ def initialize_model_parallel(
             f"but got {sequence_parallel_size}"
         )
 
-    dit_parallel_size = (
+    non_dp_size = cfg_parallel_size * sequence_parallel_size * pipeline_parallel_size * tensor_parallel_size
-        data_parallel_size * cfg_parallel_size * sequence_parallel_size * pipeline_parallel_size * tensor_parallel_size
+    if world_size % non_dp_size != 0:
-    )
+        raise ValueError(f"WORLD size ({world_size}) must be divisible by non-DP parallel size ({non_dp_size})")
-
+    if use_hsdp:
-    # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
+        if data_parallel_size not in (None, 1):
-    is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1
+            raise ValueError("HSDP (FSDP2) requires data_parallel_size to be 1")
-
+        if non_dp_size not in (1, world_size):
-    # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size) as dit_parallel_size
+            raise ValueError(f"HSDP non-DP parallel size must be 1 or WORLD size ({world_size}), but got {non_dp_size}")
-    # This ensures orthogonal rank generation works correctly for all HSDP workers
+        data_parallel_size = 1
-    if is_standalone_hsdp:
+    else:
-        dit_parallel_size = fully_shard_degree * hsdp_replicate_size
+        inferred_data_parallel_size = world_size // non_dp_size
-
+        if data_parallel_size is not None and data_parallel_size != inferred_data_parallel_size:
-    if world_size < dit_parallel_size:
+            raise ValueError(
-        raise RuntimeError(
+                f"data_parallel_size ({data_parallel_size}) does not match WORLD-derived value "
-            f"world_size ({world_size}) is less than "
+                f"({inferred_data_parallel_size})"
-            f"tensor_parallel_size ({tensor_parallel_size}) x "
+            )
-            f"pipeline_parallel_size ({pipeline_parallel_size}) x"
+        data_parallel_size = inferred_data_parallel_size
-            f"sequence_parallel_size ({sequence_parallel_size}) x"
-            f"cfg_parallel_size "
-            f"({cfg_parallel_size}) x"
-            f"data_parallel_size ({data_parallel_size})"
-        )
-
-    # For standalone HSDP, use (fully_shard_degree * hsdp_replicate_size) as data_parallel_size
-    # so that RankGenerator.world_size matches the actual number of workers
-    effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size
 
     rank_generator: RankGenerator = RankGenerator(
         tensor_parallel_size,
         sequence_parallel_size,
         pipeline_parallel_size,
         cfg_parallel_size,
-        effective_dp_size,
+        data_parallel_size,
-        fs=fully_shard_degree,
         order="tp-sp-pp-cfg-dp",
     )
+
+    def get_rank_groups(token: str) -> list[list[int]]:
+        if use_hsdp and non_dp_size == 1:
+            return [[rank] for rank in range(world_size)]
+        return rank_generator.get_ranks(token)
+
     use_moe_parallel_mapping = False
     if enable_expert_parallel:
         od_config = get_forward_context().omni_diffusion_config
@@ -863,11 +794,11 @@ def initialize_model_parallel(
         if not use_moe_parallel_mapping:
             raise RuntimeError("Expert parallelism enabled for a non-MoE model")
 
-    sp_group_ranks = rank_generator.get_ranks("sp")
+    sp_group_ranks = get_rank_groups("sp")
     global _DP
     assert _DP is None, "data parallel group is already initialized"
     _DP = init_model_parallel_group(
-        group_ranks=rank_generator.get_ranks("dp"),
+        group_ranks=get_rank_groups("dp"),
         local_rank=get_world_group().local_rank,
         backend=backend,
         parallel_mode="data",
@@ -877,7 +808,7 @@ def initialize_model_parallel(
     global _CFG
     assert _CFG is None, "classifier_free_guidance group is already initialized"
     _CFG = init_model_parallel_group(
-        group_ranks=rank_generator.get_ranks("cfg"),
+        group_ranks=get_rank_groups("cfg"),
         local_rank=get_world_group().local_rank,
         backend=backend,
         parallel_mode="classifier_free_guidance",
@@ -885,7 +816,7 @@ def initialize_model_parallel(
     global _PP
     assert _PP is None, "pipeline model parallel group is already initialized"
     _PP = init_model_parallel_group(
-        group_ranks=rank_generator.get_ranks("pp"),
+        group_ranks=get_rank_groups("pp"),
         local_rank=get_world_group().local_rank,
         backend=backend,
         parallel_mode="pipeline",
@@ -899,7 +830,7 @@ def initialize_model_parallel(
         sp_ring_degree=ring_degree,
         sp_allgather_degree=allgather_degree,
         rank=get_world_group().rank_in_group,
-        world_size=dit_parallel_size,
+        world_size=world_size,
         sp_group_ranks=sp_group_ranks,
     )
     _SP = init_model_parallel_group(
@@ -925,7 +856,7 @@ def initialize_model_parallel(
         )
 
     assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized"
-    tp_group_ranks = rank_generator.get_ranks("tp")
+    tp_group_ranks = get_rank_groups("tp")
     if use_moe_parallel_mapping:
         vllm_parallel_state._TP = init_vllm_model_parallel_group(
             group_ranks=tp_group_ranks,
@@ -944,22 +875,14 @@ def initialize_model_parallel(
         # CFG is a diffusion-specific replica dimension. Fold it into vLLM DP
         # only when constructing the vLLM EP layout for expert-parallel paths.
         vllm_parallel_state._DP = init_vllm_model_parallel_group(
-            group_ranks=rank_generator.get_ranks("cfg-dp"),
+            group_ranks=get_rank_groups("cfg-dp"),
             local_rank=get_world_group().local_rank,
             backend=backend,
             group_name="dp",
         )
 
-    global _FS, _EXPERT_PARALLEL_GROUP_RANKS
+    global _EXPERT_PARALLEL_GROUP_RANKS
-    assert _FS is None, "fully shard group is already initialized"
+    _EXPERT_PARALLEL_GROUP_RANKS = get_rank_groups("tp-sp-cfg-dp")
-    _FS = init_model_parallel_group(
-        group_ranks=rank_generator.get_ranks("fs", independent_ranks=True),
-        local_rank=get_world_group().local_rank,
-        backend=backend,
-        parallel_mode="fully_shard",
-    )
-
-    _EXPERT_PARALLEL_GROUP_RANKS = rank_generator.get_ranks("tp-sp-cfg-dp")
     if use_moe_parallel_mapping:
         vllm_parallel_state._EP = init_vllm_model_parallel_group(
             group_ranks=_EXPERT_PARALLEL_GROUP_RANKS,
@@ -969,12 +892,63 @@ def initialize_model_parallel(
             use_all2all=True,
         )
 
-    init_dit_group(dit_parallel_size, backend)
+
+def initialize_model_parallel(
+    data_parallel_size: int | None = None,
+    cfg_parallel_size: int = 1,
+    sequence_parallel_size: int | None = None,
+    ulysses_degree: int = 1,
+    ring_degree: int = 1,
+    allgather_degree: int = 1,
+    tensor_parallel_size: int = 1,
+    pipeline_parallel_size: int = 1,
+    enable_expert_parallel: bool = False,
+    use_hsdp: bool = False,
+    backend: str | None = None,
+) -> None:
+    """Atomically initialize diffusion parallel groups.
+
+    Configuration is validated before group creation. Any failure tears down
+    groups created by this attempt, including vLLM-bound coordinators.
+    """
+    existing_groups = {
+        "dp": _DP,
+        "cfg": _CFG,
+        "sp": _SP,
+        "pp": _PP,
+        "vllm_tp": vllm_parallel_state._TP,
+        "vllm_dp": vllm_parallel_state._DP,
+        "vllm_pp": vllm_parallel_state._PP,
+        "vllm_ep": vllm_parallel_state._EP,
+        "vllm_pcp": vllm_parallel_state._PCP,
+    }
+    initialized_names = [name for name, group in existing_groups.items() if group is not None]
+    if initialized_names:
+        raise RuntimeError(
+            "model parallel state must be empty before initialization; found " + ", ".join(initialized_names)
+        )
+    try:
+        _initialize_model_parallel(
+            data_parallel_size=data_parallel_size,
+            cfg_parallel_size=cfg_parallel_size,
+            sequence_parallel_size=sequence_parallel_size,
+            ulysses_degree=ulysses_degree,
+            ring_degree=ring_degree,
+            allgather_degree=allgather_degree,
+            tensor_parallel_size=tensor_parallel_size,
+            pipeline_parallel_size=pipeline_parallel_size,
+            enable_expert_parallel=enable_expert_parallel,
+            use_hsdp=use_hsdp,
+            backend=backend,
+        )
+    except BaseException:
+        destroy_model_parallel()
+        raise
 
 
 def destroy_model_parallel():
     """Set the groups to none and destroy them."""
-    global _DP, _CFG, _SP, _PP, _FS, _EXPERT_PARALLEL_GROUP_RANKS
+    global _DP, _CFG, _SP, _PP, _EXPERT_PARALLEL_GROUP_RANKS
 
     if vllm_parallel_state._DP and vllm_parallel_state._DP is not _DP:
         vllm_parallel_state._DP.destroy()
@@ -1005,14 +979,14 @@ def destroy_model_parallel():
     vllm_parallel_state._EP = None
     _EXPERT_PARALLEL_GROUP_RANKS = None
 
+    if vllm_parallel_state._PP and vllm_parallel_state._PP is not _PP:
+        vllm_parallel_state._PP.destroy()
+    vllm_parallel_state._PP = None
+
     if _PP:
         _PP.destroy()
     _PP = None
 
-    if _FS:
-        _FS.destroy()
-    _FS = None
-
 
 def destroy_distributed_environment():
     global _WORLD

vllm_omni/diffusion/model_metadata.py

@@ -56,6 +56,10 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
     "WanImageToVideoPipeline": DiffusionModelMetadata(attention_mask_free=True),
     "WanVACEPipeline": DiffusionModelMetadata(attention_mask_free=True),
     "WanS2VPipeline": DiffusionModelMetadata(attention_mask_free=True),
+    "SanaWmPipeline": DiffusionModelMetadata(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=1,
+    ),
 }
 
 

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -29,7 +29,7 @@ from vllm_omni.diffusion.cache.cachedit import (
 )
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
-    get_dit_group,
+    get_world_group,
     init_world_group,
 )
 from vllm_omni.diffusion.distributed.utils import get_local_device
@@ -423,7 +423,7 @@ def _validate_reference_image(image: Image.Image) -> None:
 def _dit_rank_world() -> tuple[Any, int, int]:
     if not dist.is_initialized():
         return None, 0, 1
-    group = get_dit_group()
+    group = get_world_group().device_group
     return group, dist.get_rank(group), dist.get_world_size(group)
 
 
@@ -1715,7 +1715,7 @@ class MiniMaxH3Pipeline(
                     dist.broadcast(
                         has_audio_tensor,
                         src=0,
-                        group=get_dit_group(),
+                        group=get_world_group().device_group,
                     )
                 has_audio = [bool(value) for value in has_audio_tensor.tolist()]
 

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -18,7 +18,7 @@ from transformers.dynamic_module_utils import get_class_from_dynamic_module
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
     DistributedVaeMixin,
 )
-from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
+from vllm_omni.diffusion.distributed.parallel_state import get_world_group
 from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
 
 from .packed_tokens import minimax_h3_patchify_video_latent
@@ -157,7 +157,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
     ) -> None:
         if mode != "tile":
             raise ValueError(f"MiniMax H3 VAE supports its native tile parallel mode only, got {mode!r}")
-        group = get_dit_group()
+        group = get_world_group().device_group
         world_size = dist.get_world_size(group)
         rank = dist.get_rank(group)
         parallel_size = int(parallel_size)

vllm_omni/diffusion/models/sana_wm/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""SANA-WM diffusion model integration."""
+
+from vllm_omni.diffusion.models.sana_wm.config import SanaWmConfig
+from vllm_omni.diffusion.models.sana_wm.pipeline_sana_wm import (
+    SANA_WM_MODEL_ID,
+    SANA_WM_OUTPUT_HEIGHT,
+    SANA_WM_OUTPUT_WIDTH,
+    SanaWmPipeline,
+    get_sana_wm_pre_process_func,
+)
+from vllm_omni.diffusion.models.sana_wm.request import normalize_sana_wm_payload
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import SanaWmTransformer3DModel
+
+__all__ = [
+    "SANA_WM_MODEL_ID",
+    "SANA_WM_OUTPUT_HEIGHT",
+    "SANA_WM_OUTPUT_WIDTH",
+    "SanaWmConfig",
+    "SanaWmPipeline",
+    "SanaWmTransformer3DModel",
+    "get_sana_wm_pre_process_func",
+    "normalize_sana_wm_payload",
+]

vllm_omni/diffusion/models/sana_wm/camera_control.py

@@ -0,0 +1,356 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Camera-control helpers for SANA-WM.
+
+This module mirrors the public NVlabs/Sana SANA-WM camera preparation path:
+WASD/IJKL action rollout or explicit camera-to-world poses are converted into
+relative poses, per-latent-frame ray metadata, and per-VAE-chunk Plucker maps.
+Model-local projection layers consume these tensors in the Stage-1 DiT.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+
+SANA_WM_DEFAULT_TRANSLATION_SPEED = 0.05
+SANA_WM_DEFAULT_ROTATION_SPEED_DEG = 1.2
+SANA_WM_DEFAULT_PITCH_LIMIT_DEG = 85.0
+SANA_WM_ALLOWED_ACTION_KEYS = frozenset("wasdijkl")
+SANA_WM_DEFAULT_VAE_STRIDE = (8, 32, 32)
+
+
+@dataclass(frozen=True)
+class SanaWmCameraCondition:
+    poses: Any | None = None
+    intrinsics: Any | None = None
+    action: str | None = None
+    num_frames: int | None = None
+    height: int = 704
+    width: int = 1280
+    translation_speed: float = SANA_WM_DEFAULT_TRANSLATION_SPEED
+    rotation_speed_deg: float = SANA_WM_DEFAULT_ROTATION_SPEED_DEG
+    pitch_limit_deg: float = SANA_WM_DEFAULT_PITCH_LIMIT_DEG
+
+
+def _rot_x(angle_rad: float) -> np.ndarray:
+    c, s = np.cos(angle_rad), np.sin(angle_rad)
+    return np.array([[1.0, 0.0, 0.0], [0.0, c, -s], [0.0, s, c]], dtype=np.float64)
+
+
+def _rot_y(angle_rad: float) -> np.ndarray:
+    c, s = np.cos(angle_rad), np.sin(angle_rad)
+    return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float64)
+
+
+def _parse_action_segments(action: str) -> list[tuple[list[str], int]]:
+    """Parse an action string into ``(keys, duration)`` pairs without expanding it.
+
+    A single segment names an unbounded rollout (``"w-1000000000"`` is valid
+    syntax), so expanding here would allocate one list entry per frame before
+    anything checks the frame count. Callers total the durations and validate
+    against ``num_frames`` first; only :func:`_parse_action_string` expands.
+    """
+    cleaned = "".join(action.replace("，", ",").split())
+    if not cleaned:
+        raise ValueError("Sana-WM action string is empty.")
+
+    segments: list[tuple[list[str], int]] = []
+    for segment in cleaned.split(","):
+        if not segment or "-" not in segment:
+            raise ValueError(f"Invalid Sana-WM action segment {segment!r}: expected '<keys>-<duration>'.")
+        keys_part, duration = segment.rsplit("-", 1)
+        if not duration.isdigit() or int(duration) <= 0:
+            raise ValueError(f"Sana-WM action segment {segment!r} has a non-positive duration {duration!r}.")
+        keys_lower = keys_part.lower()
+        if keys_lower == "none":
+            keys: list[str] = []
+        else:
+            bad = sorted({key for key in keys_lower if key not in SANA_WM_ALLOWED_ACTION_KEYS})
+            if bad:
+                allowed = "".join(sorted(SANA_WM_ALLOWED_ACTION_KEYS))
+                raise ValueError(f"Sana-WM action segment {segment!r} contains unknown keys {bad}; allowed: {allowed}.")
+            keys = sorted(set(keys_lower))
+        segments.append((keys, int(duration)))
+    return segments
+
+
+def action_rollout_num_frames(action: str) -> int:
+    """Frames ``action`` rolls out: the identity start pose plus one per step.
+
+    Totals the durations arithmetically, so an oversized request is rejected
+    before anything allocates per-frame state.
+    """
+    return sum(duration for _, duration in _parse_action_segments(action)) + 1
+
+
+def _parse_action_string(action: str) -> list[list[str]]:
+    per_frame: list[list[str]] = []
+    for keys, duration in _parse_action_segments(action):
+        per_frame.extend([keys] * duration)
+    return per_frame
+
+
+def action_string_to_c2w(
+    action: str,
+    *,
+    translation_speed: float = SANA_WM_DEFAULT_TRANSLATION_SPEED,
+    rotation_speed_deg: float = SANA_WM_DEFAULT_ROTATION_SPEED_DEG,
+    pitch_limit_deg: float = SANA_WM_DEFAULT_PITCH_LIMIT_DEG,
+) -> np.ndarray:
+    """Roll out an OpenCV-convention ``(N + 1, 4, 4)`` c2w trajectory."""
+
+    per_frame = _parse_action_string(action)
+    rotate_rad = math.radians(rotation_speed_deg)
+    pitch_limit_rad = math.radians(pitch_limit_deg)
+    current = np.eye(4, dtype=np.float64)
+    current_pitch = 0.0
+    poses = [current.copy()]
+
+    for keys in per_frame:
+        held = set(keys)
+        rotation = current[:3, :3]
+        translation = current[:3, 3]
+
+        pitch_delta = (rotate_rad if "i" in held else 0.0) - (rotate_rad if "k" in held else 0.0)
+        new_pitch = current_pitch + pitch_delta
+        if -pitch_limit_rad <= new_pitch <= pitch_limit_rad:
+            current_pitch = new_pitch
+        else:
+            pitch_delta = 0.0
+
+        yaw_delta = (rotate_rad if "l" in held else 0.0) - (rotate_rad if "j" in held else 0.0)
+        new_rotation = _rot_y(yaw_delta) @ rotation @ _rot_x(pitch_delta)
+
+        forward = new_rotation[:, 2].copy()
+        forward[1] = 0.0
+        right = new_rotation[:, 0].copy()
+        right[1] = 0.0
+        forward_norm = float(np.linalg.norm(forward))
+        right_norm = float(np.linalg.norm(right))
+        if forward_norm > 0:
+            forward /= forward_norm + 1e-6
+        if right_norm > 0:
+            right /= right_norm + 1e-6
+
+        move = np.zeros(3, dtype=np.float64)
+        if "w" in held:
+            move += forward * translation_speed
+        if "s" in held:
+            move -= forward * translation_speed
+        if "d" in held:
+            move += right * translation_speed
+        if "a" in held:
+            move -= right * translation_speed
+
+        current = np.eye(4, dtype=np.float64)
+        current[:3, :3] = new_rotation
+        current[:3, 3] = translation + move
+        poses.append(current.copy())
+
+    return np.stack(poses, axis=0).astype(np.float32)
+
+
+def get_pose_inverse(transform: torch.Tensor) -> torch.Tensor:
+    """Invert homogeneous rigid transforms with shape ``(..., 4, 4)``."""
+
+    rotation = transform[..., :3, :3]
+    translation = transform[..., :3, 3]
+    rotation_inv = rotation.transpose(-1, -2)
+    translation_inv = -torch.matmul(rotation_inv, translation.unsqueeze(-1)).squeeze(-1)
+    output = torch.eye(4, dtype=transform.dtype, device=transform.device).repeat(transform.shape[:-2] + (1, 1))
+    output[..., :3, :3] = rotation_inv
+    output[..., :3, 3] = translation_inv
+    return output
+
+
+def compute_raymap(
+    intrinsics: torch.Tensor,
+    poses: torch.Tensor,
+    height: int,
+    width: int,
+    *,
+    use_plucker: bool = True,
+) -> torch.Tensor:
+    """Compute SANA-WM geometry ray maps.
+
+    Args:
+        intrinsics: ``(T, 4)`` tensor in ``[fx, fy, cx, cy]`` order.
+        poses: ``(T, 4, 4)`` camera-to-world matrices.
+        height: raymap height in latent pixels.
+        width: raymap width in latent pixels.
+        use_plucker: return ``[direction, moment]`` if true, otherwise
+            ``[origin, direction]``.
+    """
+
+    if intrinsics.ndim != 2 or intrinsics.shape[-1] != 4:
+        raise ValueError(f"Sana-WM intrinsics must have shape (T, 4), got {tuple(intrinsics.shape)}.")
+    if poses.ndim != 3 or poses.shape[-2:] != (4, 4):
+        raise ValueError(f"Sana-WM poses must have shape (T, 4, 4), got {tuple(poses.shape)}.")
+    if intrinsics.shape[0] != poses.shape[0]:
+        raise ValueError("Sana-WM intrinsics and poses must have the same temporal length.")
+
+    frames = intrinsics.shape[0]
+    device = intrinsics.device
+    dtype = intrinsics.dtype
+    y_grid, x_grid = torch.meshgrid(
+        torch.arange(height, device=device, dtype=dtype),
+        torch.arange(width, device=device, dtype=dtype),
+        indexing="ij",
+    )
+    x_grid = x_grid[None].expand(frames, -1, -1)
+    y_grid = y_grid[None].expand(frames, -1, -1)
+
+    fx = intrinsics[:, 0].view(frames, 1, 1)
+    fy = intrinsics[:, 1].view(frames, 1, 1)
+    cx = intrinsics[:, 2].view(frames, 1, 1)
+    cy = intrinsics[:, 3].view(frames, 1, 1)
+    dirs_cam = torch.stack([(x_grid - cx) / fx, (y_grid - cy) / fy, torch.ones_like(x_grid)], dim=-1)
+
+    rotation = poses[:, :3, :3]
+    translation = poses[:, :3, 3]
+    dirs_world = torch.einsum("tij,thwj->thwi", rotation, dirs_cam)
+    dirs_world = dirs_world / torch.norm(dirs_world, dim=-1, keepdim=True).clamp_min(1e-12)
+    origins = translation.view(frames, 1, 1, 3).expand_as(dirs_world)
+
+    if use_plucker:
+        moments = torch.cross(origins, dirs_world, dim=-1)
+        return torch.cat([dirs_world, moments], dim=-1)
+    return torch.cat([origins, dirs_world], dim=-1)
+
+
+def intrinsics_to_vec4_array(intrinsics: Any, *, num_frames: int, height: int, width: int) -> np.ndarray:
+    """Normalize the intrinsics payload to an ``(F, 4)`` ``[fx, fy, cx, cy]`` array.
+
+    Accepts either ``None`` (derive from the output resolution) or the
+    ``{fx, fy, cx, cy}`` mapping — the only form the request contract exposes.
+    """
+
+    if intrinsics is None:
+        focal = float(max(height, width))
+        vec = np.array([focal, focal, width / 2.0, height / 2.0], dtype=np.float32)
+        return np.broadcast_to(vec, (num_frames, 4)).copy()
+
+    if isinstance(intrinsics, dict):
+        vec = np.array(
+            [intrinsics["fx"], intrinsics["fy"], intrinsics["cx"], intrinsics["cy"]],
+            dtype=np.float32,
+        )
+        return np.broadcast_to(vec, (num_frames, 4)).copy()
+
+    raise ValueError("Sana-WM intrinsics must be a {fx, fy, cx, cy} mapping.")
+
+
+def _pack_camera_conditions(
+    poses: torch.Tensor,
+    intrinsics_latent: torch.Tensor,
+    *,
+    num_frames: int,
+    latent_height: int,
+    latent_width: int,
+    vae_time_stride: int,
+) -> dict[str, torch.Tensor]:
+    # This yields exactly the latent frame count the VAE produces:
+    # ``len(arange(0, n, s)) == ceil(n / s) == (n - 1) // s + 1``.
+    time_indices = torch.arange(0, num_frames, vae_time_stride)
+
+    raymap = torch.cat([poses[time_indices].reshape(len(time_indices), -1), intrinsics_latent[time_indices]], dim=-1)
+
+    chunks: list[torch.Tensor] = []
+    spatial_raymap_frames: list[torch.Tensor] = []
+    chunk_starts = time_indices - (vae_time_stride - 1)
+    for start, t_idx in zip(chunk_starts, time_indices):
+        start_idx = max(0, int(start))
+        end_idx = start_idx + vae_time_stride
+        chunk_poses = poses[start_idx:end_idx]
+        chunk_intrinsics = intrinsics_latent[start_idx:end_idx]
+        if chunk_poses.shape[0] < vae_time_stride:
+            pad = vae_time_stride - chunk_poses.shape[0]
+            chunk_poses = torch.cat([chunk_poses, chunk_poses[-1:].repeat(pad, 1, 1)], dim=0)
+            chunk_intrinsics = torch.cat([chunk_intrinsics, chunk_intrinsics[-1:].repeat(pad, 1)], dim=0)
+        plucker = compute_raymap(chunk_intrinsics, chunk_poses, latent_height, latent_width, use_plucker=True)
+        chunks.append(plucker.permute(0, 3, 1, 2).reshape(-1, latent_height, latent_width))
+
+        # Per-frame spatial ray-direction map for raymap_embedder: take the
+        # representative frame in each chunk (the anchor at t_idx) and use only
+        # the direction channels (last 3 of 6 in the origin+direction layout).
+        t = int(t_idx)
+        od = compute_raymap(
+            intrinsics_latent[t : t + 1],
+            poses[t : t + 1],
+            latent_height,
+            latent_width,
+            use_plucker=False,
+        )  # [1, H, W, 6]
+        spatial_raymap_frames.append(od[0, :, :, 3:])  # [H, W, 3] — direction only
+
+    chunk_plucker = torch.stack(chunks).permute(1, 0, 2, 3)
+    # spatial_raymap: [3, F_latent, H, W] — consumed by raymap_embedder
+    spatial_raymap = torch.stack(spatial_raymap_frames).permute(3, 0, 1, 2)
+    return {"raymap": raymap, "chunk_plucker": chunk_plucker, "spatial_raymap": spatial_raymap}
+
+
+def build_plucker_condition(
+    condition: SanaWmCameraCondition,
+    *,
+    vae_stride: tuple[int, int, int] | list[int] = SANA_WM_DEFAULT_VAE_STRIDE,
+) -> dict[str, torch.Tensor]:
+    """Build native Stage-1 camera tensors from normalized request metadata."""
+
+    if bool(condition.action) == (condition.poses is not None):
+        raise ValueError("Sana-WM camera condition requires exactly one of action or poses.")
+
+    if condition.action:
+        poses_c2w = action_string_to_c2w(
+            condition.action,
+            translation_speed=condition.translation_speed,
+            rotation_speed_deg=condition.rotation_speed_deg,
+            pitch_limit_deg=condition.pitch_limit_deg,
+        )
+    else:
+        poses_c2w = np.asarray(condition.poses, dtype=np.float32)
+
+    if poses_c2w.ndim != 3 or poses_c2w.shape[-2:] != (4, 4):
+        raise ValueError(f"Sana-WM c2w poses must have shape (F, 4, 4), got {poses_c2w.shape}.")
+
+    num_frames = condition.num_frames or int(poses_c2w.shape[0])
+    if poses_c2w.shape[0] < num_frames:
+        if condition.action:
+            num_frames = int(poses_c2w.shape[0])
+        else:
+            raise ValueError(f"Sana-WM c2w poses length {poses_c2w.shape[0]} is shorter than num_frames {num_frames}.")
+    poses_c2w = poses_c2w[:num_frames]
+
+    intrinsics = intrinsics_to_vec4_array(
+        condition.intrinsics,
+        num_frames=num_frames,
+        height=condition.height,
+        width=condition.width,
+    )
+
+    vae_time_stride = int(vae_stride[0])
+    vae_spatial_stride = int(vae_stride[-1])
+    latent_height = condition.height // vae_spatial_stride
+    latent_width = condition.width // vae_spatial_stride
+
+    poses = torch.from_numpy(poses_c2w).float()
+    first_inv = get_pose_inverse(poses[0:1]).squeeze(0)
+    poses_rel = torch.matmul(first_inv, poses[1:])
+    poses = torch.cat([torch.eye(4, dtype=poses.dtype).unsqueeze(0), poses_rel], dim=0)
+
+    intrinsics_latent = torch.from_numpy(intrinsics).float()
+    intrinsics_latent[:, [0, 2]] *= latent_width / float(condition.width)
+    intrinsics_latent[:, [1, 3]] *= latent_height / float(condition.height)
+
+    return _pack_camera_conditions(
+        poses,
+        intrinsics_latent,
+        num_frames=num_frames,
+        latent_height=latent_height,
+        latent_width=latent_width,
+        vae_time_stride=vae_time_stride,
+    )

vllm_omni/diffusion/models/sana_wm/config.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Config for the SANA-WM Stage-1 transformer.
+
+Loaded from the standard diffusers ``transformer/config.json`` (a flat dict
+keyed by the dataclass field names). The bespoke ``config.yaml`` release format
+is handled once, offline, at checkpoint-conversion time — the runtime only ever
+sees the converted, diffusers-standard config.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+from dataclasses import dataclass, field, fields
+from pathlib import Path
+from typing import Any
+
+
+def _as_tuple3(value: Any, default: tuple[int, int, int]) -> tuple[int, int, int]:
+    if value is None:
+        return default
+    if isinstance(value, int):
+        return (value, value, value)
+    if isinstance(value, (list, tuple)):
+        if len(value) == 2:
+            return (1, int(value[0]), int(value[1]))
+        if len(value) == 3:
+            return (int(value[0]), int(value[1]), int(value[2]))
+    return default
+
+
+# LTX-2 VAE compression ratios (SANA-WM ships ``AutoencoderKLLTX2Video``).
+# They live here rather than in the pipeline so the request-normalisation module
+# can validate latent geometry without importing the pipeline.
+SANA_WM_VAE_SPATIAL_COMPRESSION = 32
+SANA_WM_VAE_TEMPORAL_COMPRESSION = 8
+
+
+@dataclass(frozen=True)
+class SanaWmConfig:
+    """SANA-WM Stage-1 architecture and runtime config.
+
+    Defaults mirror the first public HF release. ``from_json`` (reading
+    ``transformer/config.json``) is the normal construction path.
+    """
+
+    architecture_name: str | None = "SanaMSVideoCamCtrl_1600M_P1_D20"
+    num_blocks: int = 20
+    hidden_size: int = 2240
+    mlp_ratio: float = 3.0
+    attn_type: str = "BidirectionalGDNTriton"
+    softmax_every_n: int = 4
+    linear_head_dim: int = 112
+    conv_kernel_size: int = 4
+    t_kernel_size: int = 3
+    k_conv_only: bool = True
+    ffn_type: str = "GLUMBConvTemp"
+    pos_embed_type: str = "wan_rope"
+    patch_size: tuple[int, int, int] = (1, 1, 1)
+    qk_norm: bool = True
+    cross_norm: bool = True
+    mixed_precision: str = "bf16"
+    fp32_attention: bool = True
+    image_size: int = 720
+    cam_attn_compress: int = 1
+    use_chunk_plucker_post_attn: bool = True
+    chunk_plucker_channels: int = 48
+    chunk_plucker_post_attn_blocks: int = 20
+    inference_flow_shift: float = 9.8
+    scheduler_type: str = "flow_dpm-solver"
+    chi_prompt: list[str] = field(default_factory=list)
+    model_max_length: int = 300
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> SanaWmConfig:
+        """Build from a flat mapping of field name -> value.
+
+        Unknown keys (e.g. the diffusers ``_class_name`` marker or the
+        ``latent_channels`` / ``prompt_channels`` constructor kwargs) are
+        ignored. ``patch_size`` is coerced from a JSON list back to a tuple.
+        """
+        field_names = {f.name for f in fields(cls)}
+        known = {k: v for k, v in data.items() if k in field_names}
+        if "patch_size" in known:
+            known["patch_size"] = _as_tuple3(known["patch_size"], cls.patch_size)
+        return cls(**known)
+
+    @classmethod
+    def from_json(cls, path: str | Path) -> SanaWmConfig:
+        with open(path, encoding="utf-8") as f:
+            data = json.load(f)
+        return cls.from_dict(data if isinstance(data, Mapping) else {})

vllm_omni/diffusion/models/sana_wm/pipeline_sana_wm.py

@@ -0,0 +1,942 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Sana-WM pipeline integration.
+
+This module wires the registry-visible surface, release-layout validation, and
+the Stage-1 sampling loop. The DiT, Gated-DeltaNet and camera stack are
+vLLM-Omni-native layers (``sana_wm_transformer.py`` / ``ucpe.py``); nothing here
+calls into NVlabs code at runtime.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, ClassVar
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from vllm.logger import init_logger
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
+from vllm_omni.diffusion.models.ltx2.ltx2_latents import (
+    denormalize_latents,
+    normalize_latents,
+    pack_latents,
+    resolve_video_latent_shape,
+    unpack_latents,
+)
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.models.sana_wm.camera_control import (
+    SanaWmCameraCondition,
+    build_plucker_condition,
+)
+from vllm_omni.diffusion.models.sana_wm.config import (
+    SANA_WM_VAE_SPATIAL_COMPRESSION,
+    SANA_WM_VAE_TEMPORAL_COMPRESSION,
+    SanaWmConfig,
+)
+from vllm_omni.diffusion.models.sana_wm.request import normalize_sana_wm_payload
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import (
+    SANA_WM_STAGE1_PROMPT_CHANNELS,
+    SanaWmTransformer3DModel,
+)
+from vllm_omni.diffusion.models.schedulers import FlowMatchEulerDiscreteScheduler
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
+    DiffusionPipelineProfilerMixin,
+)
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+
+logger = init_logger(__name__)
+
+# The loader expects the standard diffusers layout (transformer/config.json +
+# transformer/diffusion_pytorch_model.safetensors + model_index.json), so this
+# repo is an offline conversion of the NVlabs release into that layout. It holds
+# Stage-1 only — transformer/ + vae/, with model_index.json naming
+# SanaWmPipeline — because startup prefetches the whole model repo
+# (allow_patterns=["*"]); a repo that also carried the refiner would pull tens of
+# GB this path never loads.
+SANA_WM_MODEL_ID = "BBBBruce/SANA-WM_bidirectional-stage1-diffusers"
+
+SANA_WM_STAGE1_DIT_FILE = "transformer/diffusion_pytorch_model.safetensors"
+SANA_WM_CONFIG_FILE = "transformer/config.json"
+SANA_WM_VAE_CONFIG_FILE = "vae/config.json"
+SANA_WM_VAE_WEIGHT_FILE = "vae/diffusion_pytorch_model.safetensors"
+SANA_WM_STAGE1_TEXT_ENCODER_ID = "google/gemma-2-2b-it"
+SANA_WM_STAGE1_TEXT_ENCODER_FALLBACK_ID = "Efficient-Large-Model/gemma-2-2b-it"
+SANA_WM_STAGE1_TEXT_ENCODER_ENV = "VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER"
+SANA_WM_OUTPUT_HEIGHT = 704
+SANA_WM_OUTPUT_WIDTH = 1280
+SANA_WM_NUM_TRAIN_TIMESTEPS = 1000
+# Maximum supported clip length (20s @16fps) — the native envelope, not a
+# request default. Request-side defaults live in
+# ``models/sana_wm/request.py``.
+SANA_WM_NATIVE_NUM_FRAMES = 321
+# Fail-fast latent-token cap for the native Stage-1 path, sized to the native
+# 704x1280 output at the maximum clip length so it only rejects genuinely
+# oversized requests. Overridable per-request via ``sana_wm_native_max_tokens``.
+SANA_WM_NATIVE_MAX_TOKENS = (
+    ((SANA_WM_NATIVE_NUM_FRAMES - 1) // SANA_WM_VAE_TEMPORAL_COMPRESSION + 1)
+    * (SANA_WM_OUTPUT_HEIGHT // SANA_WM_VAE_SPATIAL_COMPRESSION)
+    * (SANA_WM_OUTPUT_WIDTH // SANA_WM_VAE_SPATIAL_COMPRESSION)
+)
+# The NVlabs production inference settings for this release. They belong in the
+# model rather than in a deploy YAML: single-stage diffusion models resolve
+# through ``create_default_diffusion``, which never reads a YAML's
+# ``default_sampling_params``, so a caller that omits either field lands here.
+SANA_WM_DEFAULT_NUM_INFERENCE_STEPS = 60
+SANA_WM_DEFAULT_GUIDANCE_SCALE = 5.0
+
+SANA_WM_STAGE1_PATTERNS = (
+    SANA_WM_CONFIG_FILE,
+    SANA_WM_STAGE1_DIT_FILE,
+    SANA_WM_VAE_CONFIG_FILE,
+    SANA_WM_VAE_WEIGHT_FILE,
+)
+SANA_WM_STAGE1_DIT_SUBFOLDER = "transformer"
+SANA_WM_STAGE1_DIT_BASENAME = Path(SANA_WM_STAGE1_DIT_FILE).name
+
+
+@dataclass(frozen=True)
+class SanaWmLocalPaths:
+    """Resolved local file paths for a SANA-WM snapshot."""
+
+    root: Path
+    config: Path
+    stage1_dit: Path
+    vae_config: Path
+    vae_weights: Path
+
+
+@dataclass(frozen=True)
+class SanaWmNativeParams:
+    """Resolved generation settings for one Stage-1 request.
+
+    Built by ``_native_params`` from the request payload and sampling params,
+    and consumed by ``_run_native_backend`` — this is the production path.
+    """
+
+    height: int
+    width: int
+    num_frames: int
+    num_inference_steps: int
+    seed: int
+    # Defaults to the NVlabs production scale, so a caller that never mentions
+    # guidance still gets the two-branch CFG forward the model was tuned for.
+    # An explicit ``guidance_scale <= 1.0`` still selects the single-branch path.
+    cfg_scale: float = SANA_WM_DEFAULT_GUIDANCE_SCALE
+
+
+def build_sana_wm_download_patterns() -> tuple[str, ...]:
+    """Return the minimal HF allow-patterns needed for SANA-WM."""
+
+    return tuple(SANA_WM_STAGE1_PATTERNS)
+
+
+def resolve_sana_wm_local_paths(snapshot_dir: str | Path) -> SanaWmLocalPaths:
+    root = Path(snapshot_dir)
+    return SanaWmLocalPaths(
+        root=root,
+        config=root / SANA_WM_CONFIG_FILE,
+        stage1_dit=root / SANA_WM_STAGE1_DIT_FILE,
+        vae_config=root / SANA_WM_VAE_CONFIG_FILE,
+        vae_weights=root / SANA_WM_VAE_WEIGHT_FILE,
+    )
+
+
+def validate_sana_wm_local_paths(paths: SanaWmLocalPaths) -> None:
+    required = [
+        paths.config,
+        paths.stage1_dit,
+        paths.vae_config,
+        paths.vae_weights,
+    ]
+    missing = [path for path in required if not path.exists()]
+    if missing:
+        joined = ", ".join(str(path) for path in missing)
+        raise FileNotFoundError(f"Missing SANA-WM checkpoint files: {joined}")
+
+
+def resolve_or_download_sana_wm_checkpoint(
+    model: str = SANA_WM_MODEL_ID,
+    *,
+    revision: str | None = None,
+    cache_dir: str | None = None,
+) -> SanaWmLocalPaths:
+    """Resolve a local SANA-WM tree or download the required HF files."""
+
+    root = Path(model)
+    if root.is_dir():
+        snapshot_dir = root
+    else:
+        snapshot_dir = Path(
+            download_weights_from_hf_specific(
+                model,
+                cache_dir,
+                list(build_sana_wm_download_patterns()),
+                revision=revision,
+                require_all=True,
+            )
+        )
+    paths = resolve_sana_wm_local_paths(snapshot_dir)
+    validate_sana_wm_local_paths(paths)
+    return paths
+
+
+def build_sana_wm_output_envelope(
+    *,
+    output: Any,
+    output_type: str,
+    metadata: dict[str, Any],
+) -> dict[str, Any]:
+    """Wrap a Stage-1 result in the canonical output envelope.
+
+    ``normalize_diffusion_postprocess_output`` splits ``{"payload": ...,
+    "metadata": ...}`` into the API-facing payload and the metadata groups, so
+    the model-specific diagnostics ride along under a ``sana_wm`` group instead
+    of the removed ``DiffusionOutput.custom_output`` field.
+    """
+    payload_key = "latents" if output_type == "latent" else "video"
+    return {"payload": {payload_key: output}, "metadata": {"sana_wm": metadata}}
+
+
+def get_sana_wm_pre_process_func(od_config: OmniDiffusionConfig):
+    del od_config
+
+    def pre_process_func(request: Any) -> Any:
+        from vllm_omni.inputs.data import OmniTextPrompt
+
+        sampling_params = getattr(request, "sampling_params", None)
+        prompt = request.prompt
+        prompt_mapping = OmniTextPrompt(prompt=prompt) if isinstance(prompt, str) else prompt
+        prompt_mapping = dict(prompt_mapping)
+
+        # The camera block travels as extra_params["sana_wm"], which serving
+        # merges into extra_args. Offline callers may also put it straight on
+        # the prompt.
+        if sampling_params is not None:
+            extra_args = getattr(sampling_params, "extra_args", None)
+            if isinstance(extra_args, Mapping):
+                block = extra_args.get("sana_wm")
+                if isinstance(block, Mapping):
+                    merged = dict(block)
+                    merged.update(prompt_mapping.get("sana_wm") or {})
+                    prompt_mapping["sana_wm"] = merged
+            # Geometry always comes from the standard sampling fields.
+            for field in ("num_frames", "height", "width"):
+                value = getattr(sampling_params, field, None)
+                if value is not None and not (field == "num_frames" and value == 1):
+                    prompt_mapping[field] = value
+
+        request.prompt = normalize_sana_wm_payload(prompt_mapping)
+
+        # Preprocess is the only place the request is normalised, so publish the
+        # canonical geometry back onto sampling_params for the rest of the stack.
+        if sampling_params is not None:
+            payload = request.prompt["additional_information"]["sana_wm"]
+            sampling_params.num_frames = payload["num_frames"]
+            sampling_params.height = payload["height"]
+            sampling_params.width = payload["width"]
+        return request
+
+    return pre_process_func
+
+
+class SanaWmPipeline(
+    nn.Module,
+    CFGParallelMixin,
+    SupportImageInput,
+    SupportsComponentDiscovery,
+    ProgressBarMixin,
+    DiffusionPipelineProfilerMixin,
+):
+    """Stage-1 SANA-WM image-to-video pipeline."""
+
+    support_image_input: ClassVar[bool] = True
+    color_format: ClassVar[str] = "RGB"
+    # Skip the startup warmup: the generic dummy request carries a blank image
+    # with no camera poses and no action string, which SANA-WM's request
+    # validation always rejects. The first real request does the warming.
+    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]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = []
+
+    def __init__(self, *, od_config: OmniDiffusionConfig | None = None, prefix: str = "") -> None:
+        super().__init__()
+        self.od_config = od_config
+        self.prefix = prefix
+
+        self.sana_wm_config = SanaWmConfig()
+        self.quant_config = getattr(od_config, "quantization_config", None) if od_config is not None else None
+
+        self.tokenizer: Any | None = None
+        self.text_encoder: nn.Module | None = None
+        self.vae: nn.Module | None = None
+        self.release_paths: SanaWmLocalPaths | None = None
+        self.weights_sources = []
+        if od_config is not None and od_config.model is not None:
+            # Resolve the checkpoint once at construction (the standard
+            # startup-time flow, like every other diffusion pipeline) so the
+            # transformer below is built with the transformer/config.json shape
+            # rather than defaults.
+            self.resolve_checkpoint()
+            self.weights_sources = [
+                DiffusersPipelineLoader.ComponentSource(
+                    model_or_path=od_config.model,
+                    subfolder=SANA_WM_STAGE1_DIT_SUBFOLDER,
+                    revision=od_config.revision,
+                    # The subfolder file carries module-local keys (standard
+                    # diffusers layout); prefix them into the pipeline
+                    # namespace for the loader's strict-coverage check.
+                    prefix="transformer.",
+                    fall_back_to_pt=False,
+                    allow_patterns_overrides=[SANA_WM_STAGE1_DIT_BASENAME],
+                )
+            ]
+        # Built eagerly: the loader constructs the pipeline under the target
+        # device / default-dtype context, and the weight loader then streams
+        # checkpoint tensors directly into these modules.
+        self.transformer = SanaWmTransformer3DModel(
+            config=self.sana_wm_config,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.transformer" if prefix else "transformer",
+        )
+        self.device = get_local_device()
+        if od_config is not None and od_config.model is not None:
+            self._build_aux_components()
+
+    def resolve_checkpoint(self) -> SanaWmLocalPaths:
+        if self.release_paths is not None:
+            return self.release_paths
+        if self.od_config is None or self.od_config.model is None:
+            raise ValueError("Sana-WM checkpoint resolution requires od_config.model.")
+        self.release_paths = resolve_or_download_sana_wm_checkpoint(
+            self.od_config.model,
+            revision=self.od_config.revision,
+        )
+        self.sana_wm_config = SanaWmConfig.from_json(self.release_paths.config)
+        return self.release_paths
+
+    @staticmethod
+    def _extra_args(sampling_params: Any | None) -> dict[str, Any]:
+        extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None
+        return dict(extra_args or {})
+
+    def _native_params(self, payload: dict[str, Any], sampling_params: Any | None) -> SanaWmNativeParams:
+        extra_args = self._extra_args(sampling_params)
+        height = int(getattr(sampling_params, "height", None) or payload["height"])
+        width = int(getattr(sampling_params, "width", None) or payload["width"])
+        num_frames = int(getattr(sampling_params, "num_frames", None) or payload["num_frames"])
+        steps = int(
+            getattr(sampling_params, "num_inference_steps", None)
+            or extra_args.get("num_inference_steps")
+            or SANA_WM_DEFAULT_NUM_INFERENCE_STEPS
+        )
+        seed = int(getattr(sampling_params, "seed", None) or extra_args.get("seed", 0))
+        if min(height, width, num_frames, steps) <= 0:
+            raise ValueError("Sana-WM native height, width, num_frames, and steps must be positive.")
+        cfg_scale = SANA_WM_DEFAULT_GUIDANCE_SCALE
+        if sampling_params is not None and getattr(sampling_params, "guidance_scale_provided", False):
+            cfg_scale = float(getattr(sampling_params, "guidance_scale", 1.0) or 1.0)
+        return SanaWmNativeParams(
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            num_inference_steps=steps,
+            seed=seed,
+            cfg_scale=cfg_scale,
+        )
+
+    def _native_dtype(self, device: torch.device) -> torch.dtype:
+        dtype = getattr(self.od_config, "dtype", None) if self.od_config is not None else None
+        if isinstance(dtype, torch.dtype):
+            return dtype
+        return torch.float32 if device.type == "cpu" else torch.bfloat16
+
+    def _runtime_device_dtype(self) -> tuple[torch.device, torch.dtype]:
+        return self.device, self._native_dtype(self.device)
+
+    def _build_aux_components(self) -> None:
+        """Construct the text encoder and VAE at startup, on CPU.
+
+        ``ModuleDiscovery`` skips attributes that are ``None``, so components
+        created lazily on the first request are invisible to sequential
+        offload, HSDP auxiliary placement, and the shared device/memory
+        lifecycle. Building them here and deferring placement to the common
+        logic mirrors LTX-2's ``initialize_pipeline_components``.
+        """
+        dtype = self._native_dtype(self.device)
+        self._load_stage1_text_encoder(dtype=dtype)
+        self._load_vae(dtype=dtype)
+        self._place_aux_components()
+
+    def _place_aux_components(self) -> None:
+        """Move the discovered auxiliary components onto the local device.
+
+        Skipped when offload or HSDP owns placement — those paths expect the
+        components to still be on CPU when they take over.
+        """
+        parallel_config = getattr(self.od_config, "parallel_config", None)
+        if (
+            getattr(self.od_config, "enable_cpu_offload", False)
+            or getattr(self.od_config, "enable_layerwise_offload", False)
+            or getattr(parallel_config, "use_hsdp", False)
+        ):
+            return
+        modules = ModuleDiscovery.discover(self)
+        for module in (*modules.encoders, *modules.vaes, *modules.resident_modules):
+            module.to(self.device)
+
+    @staticmethod
+    def _component_local_files_only(model_id: Any) -> bool:
+        """Decide download policy per component, from its own model id.
+
+        The Stage-1 release directory is local once ``resolve_checkpoint`` has
+        run, but the Gemma text encoder lives in a separate Hub repo. Deriving
+        locality from the release paths forced ``local_files_only=True`` onto
+        Gemma as well, which forbids the download on a box that has not cached
+        it yet.
+        """
+        return Path(str(model_id)).expanduser().exists()
+
+    def _load_stage1_text_encoder(self, *, dtype: torch.dtype) -> None:
+        from transformers import AutoModel, AutoTokenizer
+
+        model_ids = [os.environ.get(SANA_WM_STAGE1_TEXT_ENCODER_ENV, "").strip()]
+        model_ids.extend([SANA_WM_STAGE1_TEXT_ENCODER_ID, SANA_WM_STAGE1_TEXT_ENCODER_FALLBACK_ID])
+        errors: list[str] = []
+        for model_id in dict.fromkeys(model_id for model_id in model_ids if model_id):
+            try:
+                self.tokenizer = AutoTokenizer.from_pretrained(
+                    model_id,
+                    local_files_only=self._component_local_files_only(model_id),
+                )
+                text_encoder_model_id = model_id
+                break
+            except OSError as exc:
+                errors.append(f"{model_id}: {exc}")
+        else:
+            raise OSError("Could not load SANA-WM Stage-1 text encoder. Tried: " + "; ".join(errors))
+        if getattr(self.tokenizer, "pad_token", None) is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = "right"
+        with torch.device("cpu"):
+            # AutoModel, not AutoModelForCausalLM: only the final hidden state is
+            # used, so the LM head would compute a [seq, 256k] logits tensor
+            # that is thrown away on every prompt encode.
+            self.text_encoder = AutoModel.from_pretrained(
+                text_encoder_model_id,
+                torch_dtype=dtype,
+                local_files_only=self._component_local_files_only(text_encoder_model_id),
+            )
+
+    def _ensure_stage1_text_encoder(self, *, device: torch.device, dtype: torch.dtype) -> None:
+        if self.text_encoder is None:
+            self._load_stage1_text_encoder(dtype=dtype)
+        self.text_encoder.to(device=device, dtype=dtype)
+
+    def _load_vae(self, *, dtype: torch.dtype) -> None:
+        if self.release_paths is None and self.od_config is not None and self.od_config.model is not None:
+            self.resolve_checkpoint()
+        if self.release_paths is None:
+            raise ValueError("Sana-WM VAE loading requires a resolved checkpoint.")
+
+        from diffusers import AutoencoderKLLTX2Video
+
+        with torch.device("cpu"):
+            self.vae = AutoencoderKLLTX2Video.from_pretrained(
+                str(self.release_paths.root),
+                subfolder="vae",
+                torch_dtype=dtype,
+                local_files_only=self._component_local_files_only(self.release_paths.root),
+            )
+        self.vae.use_framewise_encoding = True
+        self.vae.use_framewise_decoding = True
+        # NVlabs reads these from its YAML VAE config with 64/96 defaults; the
+        # diffusers VAE config.json only ships scaling_factor, so the defaults
+        # are load-bearing (the keys are genuinely absent here).
+        self.vae.tile_sample_stride_num_frames = int(getattr(self.vae.config, "tile_sample_stride_num_frames", 64))
+        self.vae.tile_sample_min_num_frames = int(getattr(self.vae.config, "tile_sample_min_num_frames", 96))
+
+    def _ensure_vae(self, *, device: torch.device, dtype: torch.dtype) -> None:
+        if self.vae is None:
+            self._load_vae(dtype=dtype)
+        # Tiling is re-asserted here, not in _load_vae: registry.py assigns
+        # ``model.vae.use_tiling = od_config.vae_use_tiling`` after the pipeline
+        # is constructed, and that flag defaults to False. Matching NVlabs'
+        # inference_video_scripts/inference_sana_wm.py, the SANA VAE has to
+        # tile — a 161-frame decode costs about 9 GiB more without it, and a
+        # 321-frame request (41 latent frames) decodes as one 3D volume and
+        # OOMs. Overriding is deliberate, but say so rather than dropping the
+        # caller's setting silently.
+        if not self.vae.use_tiling:
+            logger.warning("Sana-WM forces VAE tiling on; ignoring vae_use_tiling=False (the VAE OOMs without it).")
+        self.vae.use_tiling = True
+        self.vae.to(device=device, dtype=dtype)
+
+    @staticmethod
+    def _preprocess_first_frame(
+        image: Any,
+        *,
+        height: int,
+        width: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Return the first-frame image as ``[1, 3, 1, H, W]`` in ``[-1, 1]``.
+
+        ``VideoProcessor.preprocess`` covers the PIL / ndarray / tensor cases.
+        ``resample="bicubic"`` keeps the filter PIL's ``Image.resize`` defaulted
+        to; the processor's own default is lanczos, which would move the
+        conditioning latent.
+        """
+        import numpy as np
+        from diffusers.video_processor import VideoProcessor
+
+        # The processor normalises to [-1, 1] assuming float inputs in [0, 1].
+        if isinstance(image, np.ndarray) and image.dtype == np.uint8:
+            image = image.astype(np.float32) / 255.0
+
+        processor = VideoProcessor(vae_scale_factor=SANA_WM_VAE_SPATIAL_COMPRESSION, resample="bicubic")
+        frame = processor.preprocess(image, height=height, width=width)
+        return frame.unsqueeze(2).to(device=device, dtype=dtype)
+
+    def _vae_normalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        """LTX-2 VAE per-channel latent normalisation matching NVlabs.
+
+        ``z_norm = (z_raw - latents_mean) * scaling_factor / latents_std``
+
+        The LTX-2 VAE (``AutoencoderKLLTX2Video`` in diffusers) ships with
+        per-channel ``latents_mean`` and ``latents_std`` tensors that must be
+        applied alongside ``scaling_factor`` for the round-trip to be
+        identity. Applying `* scaling_factor` alone gives decoded videos a
+        systematic G/B colour shift (pred RGB means
+        ~(110, 220, 208) vs ref ~(94, 114, 138)).
+
+        Returns the normalised latent in the same dtype as the input.
+        """
+        if self.vae is None:
+            raise RuntimeError("Sana-WM VAE did not initialize.")
+        return normalize_latents(
+            latent,
+            self.vae.latents_mean,
+            self.vae.latents_std,
+            float(self.vae.config.scaling_factor),
+        )
+
+    def _vae_denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        """Inverse of :meth:`_vae_normalize_latent` — match NVlabs decode."""
+        if self.vae is None:
+            raise RuntimeError("Sana-WM VAE did not initialize.")
+        return denormalize_latents(
+            latent,
+            self.vae.latents_mean,
+            self.vae.latents_std,
+            float(self.vae.config.scaling_factor),
+        )
+
+    def _vae_encode_first_frame(
+        self,
+        image: Any,
+        *,
+        height: int,
+        width: int,
+        latent_height: int,
+        latent_width: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Encode the first-frame image to latent space via the LTX-2 VAE.
+
+        Returns shape ``[1, 128, 1, latent_height, latent_width]``.
+        """
+        self._ensure_vae(device=device, dtype=dtype)
+        if self.vae is None:
+            raise RuntimeError("Sana-WM VAE did not initialize.")
+        frame = self._preprocess_first_frame(image, height=height, width=width, device=device, dtype=dtype)
+        with torch.inference_mode():
+            encoded = self.vae.encode(frame)
+            first_latent = encoded.latent_dist.mean
+            first_latent = self._vae_normalize_latent(first_latent.to(dtype=dtype))
+        # The encoder's own rounding can land a pixel off the latent grid the
+        # noise tensor was built on, and the two are concatenated below.
+        if first_latent.shape[-2:] != (latent_height, latent_width):
+            first_latent = F.interpolate(
+                first_latent.squeeze(2),
+                size=(latent_height, latent_width),
+                mode="bilinear",
+                align_corners=False,
+            ).unsqueeze(2)
+        return first_latent  # [1, 128, 1, lh, lw]
+
+    def _stage1_prompt_text(self, prompt: dict[str, Any]) -> str:
+        user_prompt = str(prompt.get("prompt") or "")
+        chi_prompt = "\n".join(part for part in self.sana_wm_config.chi_prompt if part)
+        if chi_prompt:
+            return chi_prompt + user_prompt
+        return user_prompt
+
+    def _native_prompt_embeds(
+        self,
+        prompt: dict[str, Any],
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        prompt_text = self._stage1_prompt_text(prompt)
+        self._ensure_stage1_text_encoder(device=device, dtype=dtype)
+        tokenizer = self.tokenizer
+        if tokenizer is None or self.text_encoder is None:
+            raise RuntimeError("Sana-WM Stage-1 text encoder did not initialize.")
+        chi_prompt = "\n".join(part for part in self.sana_wm_config.chi_prompt if part)
+        model_max_length = int(self.sana_wm_config.model_max_length)
+        max_length_all = model_max_length
+        if chi_prompt:
+            # NVlabs encodes the chi-prompt-prefixed text with extra room for
+            # the chi tokens, then keeps BOS plus the final model window.
+            max_length_all = len(tokenizer.encode(chi_prompt)) + model_max_length - 2
+        encoded = tokenizer(
+            [prompt_text],
+            padding="max_length",
+            max_length=max_length_all,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        ).to(device)
+        # last_hidden_state is the same tensor the CausalLM wrapper exposed as
+        # hidden_states[-1] (both are post-final-norm), without retaining every
+        # intermediate layer.
+        hidden_states = self.text_encoder(**encoded).last_hidden_state
+        attention_mask = encoded.attention_mask
+        if chi_prompt:
+            select_index = [0] + list(range(-model_max_length + 1, 0))
+            hidden_states = hidden_states[:, select_index]
+            attention_mask = attention_mask[:, select_index]
+        if hidden_states.shape[-1] != SANA_WM_STAGE1_PROMPT_CHANNELS:
+            raise ValueError(
+                "Sana-WM Stage-1 Gemma hidden size mismatch: expected "
+                f"{SANA_WM_STAGE1_PROMPT_CHANNELS}, got {hidden_states.shape[-1]}."
+            )
+        # Return RAW Gemma hidden states: the transformer's ``attention_y_norm``
+        # RMSNorm does the normalisation, so normalising here double-normalises.
+        prompt_attention_mask = attention_mask.to(device=device, dtype=torch.float32)
+        # An all-ones mask carries no information. Keep it ``None`` so
+        # cross-attention takes the mask-free attention path instead of handing
+        # every block, on every step, a mask the backend must reshape to no
+        # effect. Evaluated once per prompt encode, so the device sync is paid
+        # twice per request (cond + uncond), not once per denoising step.
+        if bool(prompt_attention_mask.all()):
+            prompt_attention_mask = None
+        self._last_prompt_attention_mask = prompt_attention_mask
+        return hidden_states.to(device=device, dtype=dtype)
+
+    def predict_noise(self, **kwargs: Any) -> torch.Tensor:
+        """Single transformer forward.
+
+        Overrides the mixin default, which assumes a tuple-returning transformer
+        and takes ``result[0]``; SANA-WM's returns the noise prediction directly.
+        """
+        return self.transformer(**kwargs)
+
+    def _decode_native_latents(
+        self,
+        latents: torch.Tensor,
+        *,
+        output_type: str,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> Any:
+        if output_type == "latent":
+            return latents
+        self._ensure_vae(device=device, dtype=dtype)
+        if self.vae is None:
+            raise RuntimeError("Sana-WM VAE did not initialize.")
+        # NVlabs always passes ``temb=None`` to ``vae.decode`` for this VAE
+        # (``builder.py::vae_decode -> LTX2VAE_diffusers``); match that.
+        denorm = self._vae_denormalize_latent(latents.to(getattr(self.vae, "dtype", dtype)))
+        video = self.vae.decode(denorm, temb=None, return_dict=False)[0]
+        # Post-process to the requested output_type (np/pil/pt). Let failures
+        # surface: silently returning the raw decoder tensor would hand the
+        # caller a wrong-format/range video with no error.
+        from diffusers.video_processor import VideoProcessor
+
+        processor = VideoProcessor(vae_scale_factor=getattr(self.vae, "spatial_compression_ratio", 32))
+        return processor.postprocess_video(video, output_type=output_type)
+
+    def _run_native_backend(
+        self,
+        *,
+        prompt: dict[str, Any],
+        payload: dict[str, Any],
+        sampling_params: Any | None,
+    ) -> DiffusionOutput:
+        params = self._native_params(payload, sampling_params)
+        latent_frames, latent_height, latent_width = resolve_video_latent_shape(
+            params.height,
+            params.width,
+            params.num_frames,
+            vae_spatial_compression_ratio=SANA_WM_VAE_SPATIAL_COMPRESSION,
+            vae_temporal_compression_ratio=SANA_WM_VAE_TEMPORAL_COMPRESSION,
+        )
+        token_count = latent_frames * latent_height * latent_width
+        extra_args = self._extra_args(sampling_params)
+        max_tokens = int(extra_args.get("sana_wm_native_max_tokens", SANA_WM_NATIVE_MAX_TOKENS))
+        if token_count > max_tokens:
+            raise ValueError(
+                "Sana-WM native latent token count exceeds the configured cap. "
+                f"Requested latent tokens={token_count}, max={max_tokens}. Request a smaller "
+                "`height`/`width`/`num_frames`, or raise the cap via "
+                "`sana_wm_native_max_tokens`."
+            )
+
+        device, dtype = self._runtime_device_dtype()
+        generator = torch.Generator(device=device)
+        generator.manual_seed(params.seed)
+        noise = torch.randn(
+            (1, 128, latent_frames, latent_height, latent_width),
+            device=device,
+            dtype=dtype,
+            generator=generator,
+        )
+
+        # Shared flow-matching Euler scheduler. NVlabs ``LTXFlowEuler`` drives
+        # diffusers' ``FlowMatchEulerDiscreteScheduler``; this is vLLM-Omni's
+        # vendored copy of it, including the ``per_token_timesteps`` step branch.
+        # The schedule applies the shift twice (once via the shifted ``sigma_min``
+        # persisted by ``__init__``, once in ``set_timesteps``), which is what
+        # reproduces the NVlabs timestep table (e.g. shift=9.8, N=3 ->
+        # [1000, 909.0, 87.7]).
+        scheduler = FlowMatchEulerDiscreteScheduler(
+            num_train_timesteps=SANA_WM_NUM_TRAIN_TIMESTEPS,
+            shift=self.sana_wm_config.inference_flow_shift,
+        )
+        scheduler.set_timesteps(params.num_inference_steps, device=device)
+        timesteps = scheduler.timesteps
+
+        # ``normalize_sana_wm_payload`` guarantees the key is present but says
+        # nothing about its type. Falling back to pure noise on an unusable
+        # image would silently drop the conditioning and produce a
+        # plausible-looking but wrong video, so this raises instead.
+        import numpy as _np
+
+        first_frame_image = (prompt.get("multi_modal_data") or {}).get("image")
+        if not (hasattr(first_frame_image, "convert") or isinstance(first_frame_image, (_np.ndarray, torch.Tensor))):
+            raise TypeError(
+                "Sana-WM first-frame image must be a PIL Image, numpy ndarray, or "
+                f"torch.Tensor; got {type(first_frame_image).__name__}."
+            )
+        # Per-frame timestep contract matching NVlabs ``LTXFlowEuler.sample``:
+        # frame 0 is the clean VAE-encoded conditioning latent and its timestep
+        # is held at 0 while the sampling loop drives the other frames from
+        # noise.
+        first_latent = self._vae_encode_first_frame(
+            first_frame_image,
+            height=params.height,
+            width=params.width,
+            latent_height=latent_height,
+            latent_width=latent_width,
+            device=device,
+            dtype=dtype,
+        )
+        # Place the CLEAN VAE-encoded first frame at frame 0 and rely on
+        # per-frame timesteps + the `condition_mask` torch.where restore (below)
+        # to keep it invariant across denoising steps.
+        latents = torch.cat([first_latent, noise[:, :, 1:]], dim=2)
+
+        prompt_embeds = self._native_prompt_embeds(
+            prompt,
+            device=device,
+            dtype=dtype,
+        )
+        prompt_attention_mask = self._last_prompt_attention_mask
+
+        # Under CFG, encode the negative prompt up front; each denoise step
+        # then runs a two-branch transformer forward. The negative prompt comes
+        # from the request (``prompt["negative_prompt"]``, set by
+        # ``serving_video``); absent, it defaults to the empty string matching
+        # NVlabs ``GenerationParams.negative_prompt = ""``.
+        do_cfg = params.cfg_scale > 1.0
+        if do_cfg:
+            negative_prompt_obj = {"prompt": str(prompt.get("negative_prompt") or "")}
+            negative_prompt_embeds = self._native_prompt_embeds(
+                negative_prompt_obj,
+                device=device,
+                dtype=dtype,
+            )
+            negative_prompt_attention_mask = self._last_prompt_attention_mask
+            # Restore the cond mask so later readers of
+            # `_last_prompt_attention_mask` see the positive branch.
+            self._last_prompt_attention_mask = prompt_attention_mask
+        else:
+            negative_prompt_embeds = None
+            negative_prompt_attention_mask = None
+
+        camera = payload.get("camera") or {}
+        condition = SanaWmCameraCondition(
+            poses=camera.get("poses") if isinstance(camera, dict) else None,
+            intrinsics=payload.get("intrinsics"),
+            action=payload.get("action"),
+            num_frames=params.num_frames,
+            height=params.height,
+            width=params.width,
+            translation_speed=float(payload.get("translation_speed", 0.05)),
+            rotation_speed_deg=float(payload.get("rotation_speed_deg", 1.2)),
+        )
+        camera_tensors = build_plucker_condition(condition)
+        plucker = camera_tensors["chunk_plucker"].to(device=device, dtype=dtype)
+        raymap = camera_tensors["raymap"].to(device=device, dtype=dtype)
+        spatial_raymap = camera_tensors.get("spatial_raymap")
+        if spatial_raymap is not None:
+            spatial_raymap = spatial_raymap.to(device=device, dtype=dtype)
+
+        # Build a per-frame condition mask matching NVlabs ``LTXFlowEuler``:
+        # frame 0 is the conditioning frame at sigma=0 (preserved after
+        # each ``scheduler.step``), all other frames carry the current
+        # sampling sigma. We omit the legacy `add_noise_to_image_conditioning_latents`
+        # motion-continuity term because the public SANA-WM config sets
+        # ``condition_frame_info={0: 0.0}`` (image_cond_noise_scale=0).
+        cam_batch = latents.shape[0]
+        cam_frames = latents.shape[2]
+        condition_mask = torch.zeros(cam_batch, 1, cam_frames, 1, 1, device=latents.device, dtype=latents.dtype)
+        condition_mask[:, :, 0] = 1.0
+
+        for timestep in timesteps:
+            # (B, 1, F) per-frame timestep, frame 0 forced to 0. Stays fp32:
+            # casting through the bf16 latent dtype would quantise it before
+            # ``SanaWmTimestepEmbedder``'s sinusoidal embed.
+            model_timestep = timestep.float().expand(cam_batch, 1, cam_frames).clone()
+            model_timestep[:, :, 0] = 0.0
+            positive_kwargs = {
+                "hidden_states": latents,
+                "timestep": model_timestep,
+                "encoder_hidden_states": prompt_embeds,
+                "encoder_attention_mask": prompt_attention_mask,
+                "plucker": plucker,
+                "raymap": raymap,
+                "spatial_raymap": spatial_raymap,
+            }
+            negative_kwargs = (
+                {
+                    **positive_kwargs,
+                    "encoder_hidden_states": negative_prompt_embeds,
+                    "encoder_attention_mask": negative_prompt_attention_mask,
+                }
+                if do_cfg
+                else None
+            )
+            # ``cfg_normalize=False`` keeps the plain NVlabs combine,
+            # ``uncond + scale * (cond - uncond)``. Under cfg_parallel_size > 1
+            # the mixin runs one branch per rank and all-gathers instead of
+            # evaluating both branches serially.
+            noise_pred = self.predict_noise_maybe_with_cfg(
+                do_cfg,
+                params.cfg_scale,
+                positive_kwargs,
+                negative_kwargs,
+                cfg_normalize=False,
+            )
+
+            # The per-frame contract pairs with a per-token flow-matching Euler
+            # step. Both step branches end at ``prev = sample + dt *
+            # model_output``, but the scheduler defines dt as
+            # ``sigma - sigma_next`` per-token against ``sigma_next - sigma``
+            # standard, so the per-token call takes the sign-flipped prediction
+            # (matching NVlabs ``LTXFlowEuler.sample``); passing the standard
+            # sign here would integrate towards more noise.
+            #
+            # Broadcast the (B, 1, F) timestep to (B, 1, F, H, W) and flatten
+            # F*H*W; frame-0 conditioning tokens are already 0. The flatten
+            # order (F, H, W) matches ``pack_latents``.
+            pt_t = (
+                model_timestep.unsqueeze(-1)
+                .unsqueeze(-1)
+                .expand(cam_batch, 1, cam_frames, latents.shape[3], latents.shape[4])
+                .reshape(cam_batch, -1)
+            )
+            # The per-token branch consumes/returns the packed (B, N, C) layout
+            # and keeps its result in fp32; cast back so the latent dtype stays
+            # stable across steps.
+            stepped_packed = scheduler.step(
+                -pack_latents(noise_pred),
+                timestep,
+                pack_latents(latents),
+                per_token_timesteps=pt_t,
+                return_dict=False,
+            )[0]
+            stepped = unpack_latents(
+                stepped_packed.to(latents.dtype),
+                cam_frames,
+                latents.shape[3],
+                latents.shape[4],
+            )
+
+            # Match NVlabs LTXFlowEuler exactly. This is stricter than a plain
+            # condition-frame restore: with bf16 latents, the `t=1000`
+            # comparison rounds `1 - 1e-6` to `1`, so the first full-noise step
+            # is discarded for generated tokens as well.
+            tokens_to_denoise_mask = timestep / float(SANA_WM_NUM_TRAIN_TIMESTEPS) - 1e-6 < (1.0 - condition_mask)
+            latents = torch.where(tokens_to_denoise_mask, stepped, latents)
+
+        output_type = getattr(sampling_params, "output_type", None) or "np"
+        output = self._decode_native_latents(latents, output_type=output_type, device=device, dtype=dtype)
+        return DiffusionOutput(
+            output=build_sana_wm_output_envelope(
+                output=output,
+                output_type=output_type,
+                metadata={
+                    "backend": "native_gdn",
+                    "output_space": output_type,
+                    "chi_prompt_applied": bool(self.sana_wm_config.chi_prompt),
+                    "num_frames": params.num_frames,
+                    "height": params.height,
+                    "width": params.width,
+                    "latent_tokens": token_count,
+                    "sampling_steps": params.num_inference_steps,
+                },
+            ),
+        )
+
+    def forward(self, req: DiffusionRequestBatch, *args: Any, **kwargs: Any) -> DiffusionOutput:
+        del args, kwargs
+        if len(req.prompts) != 1:
+            raise ValueError("Sana-WM native backend currently supports exactly one prompt per request.")
+
+        prompt = req.prompts[0]
+        if isinstance(prompt, str):
+            raise ValueError("Sana-WM requires a mapping prompt with first-frame image and camera/action metadata.")
+        # get_sana_wm_pre_process_func is the single canonical entry point: it
+        # normalises the request and republishes the geometry onto
+        # sampling_params. normalize_sana_wm_payload is idempotent, so this only
+        # covers offline callers that bypass the preprocess hook.
+        prompt = normalize_sana_wm_payload(prompt)
+        payload = prompt["additional_information"]["sana_wm"]
+
+        # The checkpoint is resolved once in __init__; this is a cached no-op
+        # kept as a guard for pipelines constructed without od_config.model.
+        if self.od_config is not None and self.od_config.model is not None:
+            self.resolve_checkpoint()
+
+        return self._run_native_backend(
+            prompt=prompt,
+            payload=payload,
+            sampling_params=req.sampling_params,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, Any]]) -> set[str]:
+        # Keys arrive pipeline-prefixed (``transformer.*``, per the component
+        # source). AutoWeightsLoader streams them to the submodule and returns
+        # the prefixed names the loader's strict-coverage check compares.
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/diffusion/models/sana_wm/request.py

@@ -0,0 +1,269 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Request normalization for Sana-WM image-to-video.
+
+The released SANA-WM checkpoint is first-frame image-to-video. This module keeps
+request-side normalization lightweight: it validates the image/camera/action
+contract and stores a canonical payload under
+``additional_information["sana_wm"]``. Model-internal raymap / Plucker
+projection lives with the transformer.
+
+This lives beside the model rather than under ``stage_input_processors``
+because nothing loads it as one: the pipeline's ``pre_process_func`` is
+registered from ``pipeline_sana_wm``, the same as every other diffusion model,
+and no deploy config wires a ``custom_process_input_func`` here. Keeping it in
+the stage package forced the pipeline to import *upwards* — the only such
+import in ``diffusion/models/`` — which in turn made the VAE compression
+constants impossible to share and so duplicated.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from vllm_omni.diffusion.models.sana_wm.camera_control import action_rollout_num_frames
+from vllm_omni.diffusion.models.sana_wm.config import (
+    SANA_WM_VAE_SPATIAL_COMPRESSION,
+    SANA_WM_VAE_TEMPORAL_COMPRESSION,
+)
+
+SANA_WM_CANONICAL_KEY = "sana_wm"
+# Request fallbacks used whenever the field is absent, mirroring the deploy YAML
+# (161 frames) and the model's native resolution. Both are genuinely reachable:
+# a served request that omits num_frames never gets one injected, and callers
+# may ask for a different resolution.
+SANA_WM_DEFAULT_NUM_FRAMES = 161
+SANA_WM_DEFAULT_HEIGHT = 704
+SANA_WM_DEFAULT_WIDTH = 1280
+SANA_WM_DEFAULT_CAMERA_FORMAT = "c2w_4x4"
+SANA_WM_DEFAULT_COORDINATE_SYSTEM = "official"
+
+
+def _unwrap_single(value: Any) -> Any:
+    if isinstance(value, list):
+        return value[0] if value else None
+    return value
+
+
+def _as_dict(value: Any, *, name: str) -> dict[str, Any]:
+    if value is None:
+        return {}
+    if not isinstance(value, Mapping):
+        raise ValueError(f"Sana-WM {name} must be a mapping, got {type(value).__name__}.")
+    return dict(value)
+
+
+def _is_sequence(value: Any) -> bool:
+    return isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray))
+
+
+def _as_finite_float(value: Any, *, name: str) -> float:
+    """Parse a float and reject NaN/inf.
+
+    Non-finite geometry does not fail loudly downstream: it propagates through
+    the raymap into the Plucker maps and comes back as an all-NaN video with a
+    200 response.
+    """
+    try:
+        parsed = float(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Sana-WM {name} must be a number, got {value!r}.") from exc
+    if not math.isfinite(parsed):
+        raise ValueError(f"Sana-WM {name} must be finite, got {parsed!r}.")
+    return parsed
+
+
+def _validate_numeric_matrix4x4(matrix: Any, *, name: str) -> None:
+    if not _is_sequence(matrix) or len(matrix) != 4:
+        raise ValueError(f"Sana-WM {name} must be a 4x4 matrix.")
+    for row_idx, row in enumerate(matrix):
+        if not _is_sequence(row) or len(row) != 4:
+            raise ValueError(f"Sana-WM {name}[{row_idx}] must contain 4 values.")
+        for col_idx, value in enumerate(row):
+            _as_finite_float(value, name=f"{name}[{row_idx}][{col_idx}]")
+
+
+def _validate_camera_poses(poses: Any, *, num_frames: int | None = None) -> None:
+    if not _is_sequence(poses) or len(poses) == 0:
+        raise ValueError("Sana-WM camera poses must be a non-empty sequence of 4x4 matrices.")
+    if num_frames is not None and len(poses) != num_frames:
+        raise ValueError(f"Sana-WM camera poses length {len(poses)} must equal num_frames {num_frames}.")
+    for idx, pose in enumerate(poses):
+        _validate_numeric_matrix4x4(pose, name=f"camera.poses[{idx}]")
+
+
+def _validate_intrinsics(intrinsics: Any) -> None:
+    # Only the ``{fx, fy, cx, cy}`` mapping is accepted (the one form the recipe
+    # and examples use). Camera intrinsics are optional; when omitted the model
+    # derives them from the output resolution.
+    if intrinsics is None:
+        return
+    if not isinstance(intrinsics, Mapping):
+        raise ValueError("Sana-WM intrinsics must be a {fx, fy, cx, cy} mapping.")
+    required = {"fx", "fy", "cx", "cy"}
+    missing = sorted(required - set(intrinsics))
+    if missing:
+        raise ValueError(f"Sana-WM intrinsics mapping is missing keys: {missing}.")
+    for key in required:
+        value = _as_finite_float(intrinsics[key], name=f"intrinsics[{key!r}]")
+        # ``compute_raymap`` divides by fx/fy. cx/cy stay unconstrained beyond
+        # finiteness — a principal point outside the frame is unusual, not wrong.
+        if key in ("fx", "fy") and value <= 0:
+            raise ValueError(f"Sana-WM intrinsics[{key!r}] must be positive, got {value}.")
+
+
+def _as_positive_int(value: Any, *, name: str) -> int:
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Sana-WM {name} must be an integer, got {value!r}.") from exc
+    if parsed <= 0:
+        raise ValueError(f"Sana-WM {name} must be positive, got {parsed}.")
+    return parsed
+
+
+def _as_positive_float(value: Any, *, name: str) -> float:
+    # Finiteness first: every comparison against NaN is False, so a bare
+    # ``parsed <= 0`` would accept it.
+    parsed = _as_finite_float(value, name=name)
+    if parsed <= 0:
+        raise ValueError(f"Sana-WM {name} must be positive, got {parsed}.")
+    return parsed
+
+
+def _extract_image(mm_data: Mapping[str, Any]) -> Any:
+    return _unwrap_single(mm_data.get("image"))
+
+
+def _first_present(*values: Any) -> Any:
+    for value in values:
+        if value is not None:
+            return value
+    return None
+
+
+def _validate_latent_geometry(*, num_frames: int, height: int, width: int) -> None:
+    """Reject geometry the VAE would silently floor.
+
+    The SANA-WM VAE compresses 32x spatially and 8x temporally, so a request
+    that is not aligned loses pixels/frames instead of failing.
+    """
+    if height % SANA_WM_VAE_SPATIAL_COMPRESSION or width % SANA_WM_VAE_SPATIAL_COMPRESSION:
+        raise ValueError(
+            f"Sana-WM height and width must be divisible by {SANA_WM_VAE_SPATIAL_COMPRESSION}, "
+            f"got {height} and {width}."
+        )
+    if (num_frames - 1) % SANA_WM_VAE_TEMPORAL_COMPRESSION:
+        raise ValueError(
+            f"Sana-WM num_frames must satisfy (num_frames - 1) % {SANA_WM_VAE_TEMPORAL_COMPRESSION} == 0, "
+            f"got {num_frames}."
+        )
+
+
+def normalize_sana_wm_payload(prompt: Mapping[str, Any]) -> dict[str, Any]:
+    """Return a prompt copy with canonical Sana-WM request metadata.
+
+    The ``sana_wm`` block is read from the top-level ``sana_wm`` key, falling
+    back to ``additional_information["sana_wm"]`` so the function is idempotent
+    when called again on its own output. The first-frame image is read from
+    ``multi_modal_data["image"]``.
+    """
+
+    if not isinstance(prompt, Mapping):
+        raise ValueError(f"Sana-WM prompt must be a mapping, got {type(prompt).__name__}.")
+
+    result = dict(prompt)
+    mm_data = _as_dict(result.get("multi_modal_data"), name="multi_modal_data")
+    image = _extract_image(mm_data)
+    if image is None:
+        raise ValueError("Sana-WM requires a first-frame image in `multi_modal_data['image']`.")
+
+    additional = _as_dict(result.get("additional_information"), name="additional_information")
+    raw = _first_present(
+        result.get(SANA_WM_CANONICAL_KEY),
+        additional.get(SANA_WM_CANONICAL_KEY),
+        {},
+    )
+    raw = _as_dict(raw, name=SANA_WM_CANONICAL_KEY)
+
+    num_frames = _as_positive_int(
+        _first_present(raw.get("num_frames"), result.get("num_frames"), SANA_WM_DEFAULT_NUM_FRAMES),
+        name="num_frames",
+    )
+    height = _as_positive_int(
+        _first_present(raw.get("height"), result.get("height"), SANA_WM_DEFAULT_HEIGHT),
+        name="height",
+    )
+    width = _as_positive_int(_first_present(raw.get("width"), result.get("width"), SANA_WM_DEFAULT_WIDTH), name="width")
+    _validate_latent_geometry(num_frames=num_frames, height=height, width=width)
+
+    camera = raw.get("camera")
+    camera = _as_dict(camera, name="camera") if camera is not None else None
+    action = raw.get("action")
+    if action is not None:
+        if not isinstance(action, str) or not action.strip():
+            raise ValueError("Sana-WM action must be a non-empty string when provided.")
+        # Rejected here rather than in the model: a mismatch used to shorten the
+        # camera tensor silently and then fail deep in the transformer on a
+        # camera/latent frame mismatch, and an unbounded duration used to be
+        # expanded one entry per frame before any of that.
+        rollout_frames = action_rollout_num_frames(action)
+        if rollout_frames != num_frames:
+            raise ValueError(
+                f"Sana-WM action rolls out {rollout_frames} frames but num_frames is {num_frames}. "
+                "Action segment durations must sum to num_frames - 1."
+            )
+    if camera is not None and action is not None:
+        raise ValueError("Sana-WM accepts exactly one of `camera` or `action`, not both.")
+    if camera is None and action is None:
+        raise ValueError("Sana-WM requires either explicit camera poses or an action DSL string.")
+
+    canonical_camera = None
+    if camera is not None:
+        poses = camera.get("poses")
+        if poses is None:
+            raise ValueError("Sana-WM camera payload requires `poses`.")
+        _validate_camera_poses(poses, num_frames=num_frames)
+        # Only the defaults are accepted: nothing downstream reads these two
+        # fields, so honouring another value would mean interpreting w2c poses
+        # as c2w and silently generating the wrong trajectory.
+        camera_format = camera.get("format", SANA_WM_DEFAULT_CAMERA_FORMAT)
+        if camera_format != SANA_WM_DEFAULT_CAMERA_FORMAT:
+            raise ValueError(
+                f"Sana-WM only accepts camera format {SANA_WM_DEFAULT_CAMERA_FORMAT!r}; "
+                f"convert the poses before sending, got {camera_format!r}."
+            )
+        coordinate_system = camera.get("coordinate_system", SANA_WM_DEFAULT_COORDINATE_SYSTEM)
+        if coordinate_system != SANA_WM_DEFAULT_COORDINATE_SYSTEM:
+            raise ValueError(
+                f"Sana-WM only accepts coordinate system {SANA_WM_DEFAULT_COORDINATE_SYSTEM!r}, "
+                f"got {coordinate_system!r}."
+            )
+        canonical_camera = {"poses": poses}
+
+    intrinsics = raw.get("intrinsics")
+    _validate_intrinsics(intrinsics)
+    translation_speed = raw.get("translation_speed")
+    rotation_speed_deg = raw.get("rotation_speed_deg")
+
+    canonical = {
+        "num_frames": num_frames,
+        "height": height,
+        "width": width,
+        "camera": canonical_camera,
+        "action": action,
+        "intrinsics": intrinsics,
+    }
+    if translation_speed is not None:
+        canonical["translation_speed"] = _as_positive_float(translation_speed, name="translation_speed")
+    if rotation_speed_deg is not None:
+        canonical["rotation_speed_deg"] = _as_positive_float(rotation_speed_deg, name="rotation_speed_deg")
+
+    mm_data["image"] = image
+    result["multi_modal_data"] = mm_data
+    additional[SANA_WM_CANONICAL_KEY] = canonical
+    result["additional_information"] = additional
+    result.pop(SANA_WM_CANONICAL_KEY, None)
+    return result

vllm_omni/diffusion/models/sana_wm/ucpe.py

@@ -0,0 +1,536 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""SANA-WM UCPE (Unified Camera Pose Embedding) per-block attention transforms.
+
+Ported from NVlabs/Sana ``sana_camctrl_blocks.py`` for the inference-only
+SANA-WM bidirectional 1600M release. Scope is intentionally narrow:
+
+* Pinhole camera only (xi=0); the UCM ``xi`` parameter is fixed.
+* Inference path only — no training-time camera-branch dropout.
+* Online computation only — the precomputed ``cam_pos_embeds`` shortcut
+  used by NVlabs' fused kernels is omitted because vLLM-Omni builds the
+  matrices once per transformer forward and shares them across all blocks.
+* ``apply_vo=True`` only — the SANA-WM checkpoint always uses the
+  inverse-output transform.
+
+Public surface:
+    ``prepare_cam_geometry(...)`` -> :class:`SanaWmCamGeometry`, built once per
+    transformer forward and threaded through every block. It carries both the
+    ``(apply_q, apply_kv, apply_output)`` closures used by the softmax camera
+    branch and the projection/RoPE tensors used by ``cam_prep_func(...)``, the
+    NVlabs fused camera-prep contract on the GDN camera branch.
+
+Each closure transforms a tensor of shape ``(B, num_heads, N, D)`` where
+``N = T * H * W`` (matching the GDN token layout) and ``D == head_dim``,
+returning a tensor of the same shape.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+
+# ---------------------------------------------------------------------------
+# Pinhole ray-grid construction
+# ---------------------------------------------------------------------------
+
+
+def _make_pixel_grid(
+    height: int,
+    width: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Return ``(H, W, 3)`` grid ``(x, y, 1)`` of pixel coordinates."""
+    xs = torch.arange(width, device=device, dtype=dtype)
+    ys = torch.arange(height, device=device, dtype=dtype)
+    y_grid, x_grid = torch.meshgrid(ys, xs, indexing="ij")
+    return torch.stack((x_grid, y_grid, torch.ones_like(x_grid)), dim=-1)
+
+
+def _unproject_pinhole(
+    fx: torch.Tensor,
+    fy: torch.Tensor,
+    cx: torch.Tensor,
+    cy: torch.Tensor,
+    height: int,
+    width: int,
+) -> torch.Tensor:
+    """Unproject a pixel grid into camera-frame direction vectors.
+
+    Args:
+        fx, fy, cx, cy: shape ``(B, T)`` intrinsics in token-grid units — the
+            caller divides all four by patch size before calling.
+        height, width: token-grid size.
+
+    Returns:
+        ``(B, T, H, W, 3)`` direction tensor (NOT normalised).
+    """
+    B, T = fx.shape
+    device = fx.device
+    dtype = fx.dtype
+    grid = _make_pixel_grid(height, width, device, dtype)  # (H, W, 3)
+    u = grid[..., 0]  # (H, W)
+    v = grid[..., 1]
+    fx_e = fx.view(B, T, 1, 1)
+    fy_e = fy.view(B, T, 1, 1)
+    cx_e = cx.view(B, T, 1, 1)
+    cy_e = cy.view(B, T, 1, 1)
+    x = (u - cx_e) / fx_e
+    y = (v - cy_e) / fy_e
+    z = torch.ones_like(x)
+    # Match NVlabs UCM normalisation at xi=0: alpha = 1, gamma = 1/(1+r^2),
+    # X = gamma*x, Y = gamma*y, Z = gamma. The resulting vector has unit
+    # norm. Computed compactly via direct stack + normalise.
+    d_cam = torch.stack((x, y, z), dim=-1)  # (B, T, H, W, 3)
+    d_cam = d_cam / d_cam.norm(dim=-1, keepdim=True).clamp_min(1e-12)
+    return d_cam
+
+
+def _world_to_ray_mats(
+    d_cam: torch.Tensor,  # (B, T, H, W, 3) unit rays in camera frame
+    c2w: torch.Tensor,  # (B, T, 4, 4) camera-to-world SE(3)
+) -> torch.Tensor:
+    """Build per-pixel ``ray <- world`` SE(3) transforms ``(B, T, H, W, 4, 4)``.
+
+    The ray frame is anchored at the camera origin with the +Z axis along
+    the per-pixel viewing direction and a +Y axis defined by the camera-Y
+    cross product. This matches NVlabs ``world_to_ray_mats``.
+    """
+    B, T, H, W, _ = d_cam.shape
+    device = d_cam.device
+    dtype = d_cam.dtype
+    rot_cam = c2w[..., :3, :3]
+    trans_cam = c2w[..., :3, 3]
+    d_world = torch.einsum("btij,bthwj->bthwi", rot_cam, d_cam)
+    cam_y = rot_cam[..., :, 1].unsqueeze(2).unsqueeze(3).expand(B, T, H, W, 3)
+    z_ray = F.normalize(d_world, dim=-1, eps=1e-6)
+    x_ray = F.normalize(torch.cross(cam_y, z_ray, dim=-1), dim=-1, eps=1e-6)
+    y_ray = F.normalize(torch.cross(z_ray, x_ray, dim=-1), dim=-1, eps=1e-6)
+    rot_local_to_world = torch.stack((x_ray, y_ray, z_ray), dim=-1)  # (B,T,H,W,3,3)
+    rot_world_to_local = rot_local_to_world.transpose(-1, -2)
+    t_world = trans_cam.unsqueeze(2).unsqueeze(3).expand(B, T, H, W, 3)
+    t_w2l = -torch.einsum("bthwij,bthwj->bthwi", rot_world_to_local, t_world)
+    raymats = torch.zeros(B, T, H, W, 4, 4, device=device, dtype=dtype)
+    raymats[..., :3, :3] = rot_world_to_local
+    raymats[..., :3, 3] = t_w2l
+    raymats[..., 3, 3] = 1.0
+    # Guard NaNs from degenerate pixels (e.g. nan direction at infinity).
+    nan_mask = torch.isnan(d_world).any(-1)
+    if nan_mask.any():
+        eye = torch.eye(4, device=device, dtype=dtype)
+        raymats[nan_mask] = eye
+    return raymats
+
+
+def _process_camera_conditions(
+    camera_conditions: torch.Tensor,  # (B, T, 20)
+    hw: tuple[int, int, int],
+    patch_size: tuple[int, int, int] = (1, 1, 1),
+) -> torch.Tensor:
+    """Convert ``(B, T, 20)`` SANA-WM camera conditions to ``(B, T, H, W, 4, 4)``
+    ``ray <- world`` transforms.
+
+    Layout of the last axis is ``[c2w_flat (16), fx, fy, cx, cy]``.
+    Intrinsics follow the NVlabs UCPE contract: the packed camera payload is
+    expressed on the pre-patch latent image grid, while ``HW`` is the
+    transformer token grid. ``patch_size`` maps those units to the token grid
+    before unprojection. The released SANA-WM P1 config uses ``(1, 1, 1)``.
+
+    Args:
+        camera_conditions: ``(B, T, 20)`` from ``_pack_camera_conditions``.
+        hw: ``(T, H, W)`` token-grid shape. Equal to the latent grid only
+            because the released config patches by ``(1, 1, 1)``.
+    """
+    B, T_cond, last = camera_conditions.shape
+    if last != 20:
+        raise ValueError(f"SANA-WM camera_conditions last dim must be 20, got {last}.")
+    T_latent, H_latent, W_latent = hw
+    if T_cond != T_latent:
+        raise ValueError(f"SANA-WM camera_conditions frames {T_cond} != latent frames {T_latent}.")
+    c2w = camera_conditions[..., :16].reshape(B, T_latent, 4, 4)
+    patch_t, patch_h, patch_w = patch_size
+    del patch_t
+    fx = camera_conditions[..., 16] / float(patch_w)
+    fy = camera_conditions[..., 17] / float(patch_h)
+    cx = camera_conditions[..., 18] / float(patch_w)
+    cy = camera_conditions[..., 19] / float(patch_h)
+    d_cam = _unproject_pinhole(fx, fy, cx, cy, H_latent, W_latent)
+    raymats = _world_to_ray_mats(d_cam, c2w)
+    return raymats
+
+
+def _invert_se3(transforms: torch.Tensor) -> torch.Tensor:
+    """Closed-form inverse of a stack of 4x4 SE(3) matrices."""
+    if transforms.shape[-2:] != (4, 4):
+        raise ValueError(f"SE3 inverse expects (..., 4, 4), got {tuple(transforms.shape)}.")
+    rot_inv = transforms[..., :3, :3].transpose(-1, -2)
+    out = torch.zeros_like(transforms)
+    out[..., :3, :3] = rot_inv
+    out[..., :3, 3] = -torch.einsum("...ij,...j->...i", rot_inv, transforms[..., :3, 3])
+    out[..., 3, 3] = 1.0
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Per-token transform primitives
+# ---------------------------------------------------------------------------
+
+
+def _apply_ray_projmat(
+    feats: torch.Tensor,  # (B, num_heads, N, D)
+    matrix: torch.Tensor,  # (B, N, 4, 4)
+) -> torch.Tensor:
+    """Apply a per-token 4x4 projection to features grouped by 4 channels.
+
+    ``D`` must be a multiple of 4. Out-of-place; preserves ``feats.shape``.
+    """
+    B, num_heads, N, D = feats.shape
+    if D % 4 != 0:
+        raise ValueError(f"UCPE projmat expects D divisible by 4, got D={D}.")
+    feats_grouped = feats.reshape(B, num_heads, N, D // 4, 4)
+    # matrix: (B, N, 4, 4); feats_grouped: (B, num_heads, N, K, 4)
+    out = torch.einsum("bnij,bhnkj->bhnki", matrix, feats_grouped)
+    return out.reshape(B, num_heads, N, D)
+
+
+def _apply_complex_rope(
+    hidden_states: torch.Tensor,
+    freqs: torch.Tensor,
+    *,
+    inverse: bool = False,
+) -> torch.Tensor:
+    """Apply complex RoPE: treat last-dim pairs as ``(real, imag)`` and multiply
+    by ``freqs``. Computation runs in fp64 to match the NVlabs reference.
+
+    Args:
+        hidden_states: ``(..., D)`` with ``D % 2 == 0``.
+        freqs: complex tensor broadcastable to ``(..., D // 2)``.
+        inverse: when ``True`` apply ``freqs.conj()`` (used for ``apply_o``).
+    """
+    x_real = hidden_states.to(torch.float64).contiguous()
+    x_complex = torch.view_as_complex(x_real.unflatten(-1, (-1, 2)))
+    if inverse:
+        freqs = freqs.conj()
+    x_out = torch.view_as_real(x_complex * freqs).flatten(-2, -1)
+    return x_out.to(hidden_states.dtype)
+
+
+def _apply_block_diagonal(
+    feats: torch.Tensor,
+    func_size_pairs: list[tuple[Callable[[torch.Tensor], torch.Tensor], int]],
+) -> torch.Tensor:
+    """Split ``feats`` along the last axis by ``block_sizes`` and apply one
+    callable per block, then concatenate. Preserves ``feats.shape``.
+    """
+    funcs, block_sizes = zip(*func_size_pairs)
+    if feats.shape[-1] != sum(block_sizes):
+        raise ValueError(f"UCPE block-diagonal block sum {sum(block_sizes)} != feats D {feats.shape[-1]}.")
+    blocks = torch.split(feats, list(block_sizes), dim=-1)
+    return torch.cat([func(block) for func, block in zip(funcs, blocks)], dim=-1)
+
+
+def _slice_rope_for_cam(
+    rotary_emb: torch.Tensor | None,
+    rope_dim: int,
+) -> torch.Tensor | None:
+    """Re-slice WAN-style RoPE frequencies for a smaller ``rope_dim`` using the
+    same ``(T, H, W)`` split as the main branch. Mirrors NVlabs slicing logic.
+
+    The T/H/W band layout is read off the input shape
+    (``rotary_emb.shape[-1] * 2``) rather than a separately supplied head_dim,
+    so the split stays correct without the caller passing it. ``cam_head_dim``
+    still has to equal the main ``head_dim`` — ``SanaWmSelfAttention.__init__``
+    rejects anything else.
+    """
+    if rotary_emb is None:
+        return None
+    # Derive WAN RoPE T/H/W band offsets from the actual input tensor size,
+    # not from a head_dim argument. This is correct regardless of whether
+    # cam_head_dim == main_head_dim.
+    in_head_dim = rotary_emb.shape[-1] * 2
+    orig_t = in_head_dim // 2 - 2 * (in_head_dim // 6)
+    orig_h = in_head_dim // 6
+    new_t = rope_dim // 2 - 2 * (rope_dim // 6)
+    new_h = rope_dim // 6
+    new_w = rope_dim // 6
+    t_part = rotary_emb[..., :new_t]
+    h_part = rotary_emb[..., orig_t : orig_t + new_h]
+    w_part = rotary_emb[..., orig_t + orig_h : orig_t + orig_h + new_w]
+    return torch.cat([t_part, h_part, w_part], dim=-1)
+
+
+# ---------------------------------------------------------------------------
+# Apply-fn preparation
+# ---------------------------------------------------------------------------
+
+
+def _identity(x: torch.Tensor) -> torch.Tensor:
+    return x
+
+
+def _prepare_ray_apply_fns(
+    head_dim: int,
+    proj: torch.Tensor,  # (B, N, 4, 4) ray <- world
+    proj_t: torch.Tensor,  # (B, N, 4, 4) ray <- world transpose
+    proj_inv: torch.Tensor,  # (B, N, 4, 4) world <- ray
+    rotary_emb: torch.Tensor | None = None,
+) -> tuple[Callable, Callable, Callable]:
+    """Build ``(apply_q, apply_kv, apply_o)`` block-diagonal callables.
+
+    Each callable transforms a ``(B, num_heads, N, head_dim)`` tensor:
+    the first ``head_dim // 2`` channels go through the per-token 4x4
+    projection, the remaining ``head_dim // 2`` channels go through complex
+    RoPE (or identity when ``rotary_emb`` is None).
+    """
+    if rotary_emb is not None:
+        rope_fwd = partial(_apply_complex_rope, freqs=rotary_emb, inverse=False)
+        rope_inv = partial(_apply_complex_rope, freqs=rotary_emb, inverse=True)
+    else:
+        rope_fwd = _identity
+        rope_inv = _identity
+
+    half = head_dim // 2
+
+    transforms_q = [
+        (partial(_apply_ray_projmat, matrix=proj_t), half),
+        (rope_fwd, half),
+    ]
+    transforms_kv = [
+        (partial(_apply_ray_projmat, matrix=proj_inv), half),
+        (rope_fwd, half),
+    ]
+    transforms_o = [
+        (partial(_apply_ray_projmat, matrix=proj), half),
+        (rope_inv, half),
+    ]
+
+    apply_q = partial(_apply_block_diagonal, func_size_pairs=transforms_q)
+    apply_kv = partial(_apply_block_diagonal, func_size_pairs=transforms_kv)
+    apply_o = partial(_apply_block_diagonal, func_size_pairs=transforms_o)
+    return apply_q, apply_kv, apply_o
+
+
+@dataclass(frozen=True)
+class SanaWmCamGeometry:
+    """UCPE geometry shared by every block of one transformer forward.
+
+    The ray grid, its SE(3) inverse and the camera RoPE tables depend only on
+    the camera payload, the latent grid, the patch size and the head dim — all
+    fixed for the whole forward — yet each block's camera branch used to
+    rebuild them. At 20 blocks x 60 steps x 2 CFG branches that is ~2400
+    rebuilds of an 18480-token ray grid and its 4x4 inverses per request.
+
+    ``apply_q``/``apply_kv``/``apply_output`` serve the softmax camera branch;
+    ``proj_q``/``proj_kv``/``rope_cos``/``rope_sin`` serve the fused
+    ``cam_prep_func`` contract on the GDN branch. Both are derived from the
+    same projection matrices, so they are built together.
+    """
+
+    proj_q: torch.Tensor
+    proj_kv: torch.Tensor
+    rope_cos: torch.Tensor
+    rope_sin: torch.Tensor
+    apply_q: Callable[[torch.Tensor], torch.Tensor]
+    apply_kv: Callable[[torch.Tensor], torch.Tensor]
+    apply_output: Callable[[torch.Tensor], torch.Tensor]
+
+
+def _prepare_ucpe_rope_tables(
+    rotary_emb_cam: torch.Tensor,
+    token_count: int,
+    half_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Convert complex RoPE freqs to NVlabs' interleaved cos/sin tables.
+
+    The output convention is ``out[d] = x[d] * cos[d] + x[d ^ 1] * sin[d]``.
+    For every complex pair, ``sin[2i] = -imag(freq[i])`` and
+    ``sin[2i + 1] = imag(freq[i])``.
+    """
+    freqs = rotary_emb_cam
+    while freqs.ndim > 2 and freqs.shape[0] == 1:
+        freqs = freqs.squeeze(0)
+    if freqs.ndim != 2:
+        raise ValueError(f"Sana-WM cam RoPE freqs must be 2D after squeeze, got {tuple(freqs.shape)}.")
+    if freqs.shape[0] != token_count:
+        raise ValueError(f"Sana-WM cam RoPE token count {freqs.shape[0]} != {token_count}.")
+    if freqs.shape[1] * 2 != half_dim:
+        raise ValueError(f"Sana-WM cam RoPE dim {freqs.shape[1] * 2} != half_dim {half_dim}.")
+    cos_half = freqs.real.float()
+    sin_half = freqs.imag.float()
+    rope_cos = cos_half.repeat_interleave(2, dim=-1).contiguous()
+    rope_sin = torch.stack((-sin_half, sin_half), dim=-1).reshape(token_count, half_dim).contiguous()
+    return rope_cos, rope_sin
+
+
+def prepare_cam_geometry(
+    *,
+    camera_conditions: torch.Tensor,
+    spatial_shape: tuple[int, int, int],
+    patch_size: tuple[int, int, int],
+    head_dim: int,
+    rotary_emb: torch.Tensor | None,
+) -> SanaWmCamGeometry:
+    """Build the per-forward UCPE geometry: ray projections, RoPE tables, apply fns.
+
+    Call once per transformer forward and share the result across every block —
+    see :class:`SanaWmCamGeometry` for why.
+    """
+    if head_dim % 4 != 0:
+        raise ValueError(f"UCPE head_dim must be divisible by 4, got {head_dim}.")
+    if head_dim % 2 != 0 or (head_dim // 2) % 4 != 0:
+        raise ValueError(f"Sana-WM cam prep head_dim={head_dim} must satisfy D % 2 == 0 and (D/2) % 4 == 0.")
+
+    batch = camera_conditions.shape[0]
+    frames, height, width = spatial_shape
+    token_count = frames * height * width
+    half_dim = head_dim // 2
+
+    raymats = _process_camera_conditions(camera_conditions, spatial_shape, patch_size=patch_size)
+    proj = raymats.reshape(batch, token_count, 4, 4).contiguous()
+    proj_q = proj.transpose(-1, -2).contiguous()
+    proj_kv = _invert_se3(proj).contiguous()
+
+    rotary_emb_cam = _slice_rope_for_cam(rotary_emb, half_dim)
+    if rotary_emb_cam is None:
+        rope_cos = torch.ones(token_count, half_dim, device=camera_conditions.device, dtype=torch.float32)
+        rope_sin = torch.zeros(token_count, half_dim, device=camera_conditions.device, dtype=torch.float32)
+    else:
+        rope_cos, rope_sin = _prepare_ucpe_rope_tables(rotary_emb_cam, token_count, half_dim)
+
+    apply_q, apply_kv, apply_output = _prepare_ray_apply_fns(head_dim, proj, proj_q, proj_kv, rotary_emb=rotary_emb_cam)
+    return SanaWmCamGeometry(
+        proj_q=proj_q,
+        proj_kv=proj_kv,
+        rope_cos=rope_cos,
+        rope_sin=rope_sin,
+        apply_q=apply_q,
+        apply_kv=apply_kv,
+        apply_output=apply_output,
+    )
+
+
+def _validate_cam_prep_inputs(
+    q_raw: torch.Tensor,
+    k_raw: torch.Tensor,
+    v_raw: torch.Tensor,
+    proj_q: torch.Tensor,
+    proj_kv: torch.Tensor,
+    rope_cos: torch.Tensor,
+    rope_sin: torch.Tensor,
+) -> int:
+    """Validate the fused cam-prep contract and return ``head_dim // 2``."""
+    if q_raw.shape != k_raw.shape or q_raw.shape != v_raw.shape:
+        raise ValueError(f"Sana-WM cam prep q/k/v shapes must match, got {q_raw.shape}, {k_raw.shape}, {v_raw.shape}.")
+    if q_raw.ndim != 4:
+        raise ValueError(f"Sana-WM cam prep q/k/v must be (B,N,H,D), got {tuple(q_raw.shape)}.")
+    batch, token_count, _num_heads, head_dim = q_raw.shape
+    if head_dim % 2 != 0 or (head_dim // 2) % 4 != 0:
+        raise ValueError(f"Sana-WM cam prep head_dim={head_dim} must satisfy D % 2 == 0 and (D/2) % 4 == 0.")
+    half_dim = head_dim // 2
+    if proj_q.shape != (batch, token_count, 4, 4):
+        raise ValueError(f"Sana-WM cam prep proj_q shape {tuple(proj_q.shape)} != {(batch, token_count, 4, 4)}.")
+    if proj_kv.shape != (batch, token_count, 4, 4):
+        raise ValueError(f"Sana-WM cam prep proj_kv shape {tuple(proj_kv.shape)} != {(batch, token_count, 4, 4)}.")
+    if rope_cos.shape != (token_count, half_dim) or rope_sin.shape != (token_count, half_dim):
+        raise ValueError(
+            "Sana-WM cam prep rope table shapes must be "
+            f"{(token_count, half_dim)}, got {tuple(rope_cos.shape)} and {tuple(rope_sin.shape)}."
+        )
+    return half_dim
+
+
+def _apply_ray_projection(feats: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor:
+    batch, token_count, num_heads, half_dim = feats.shape
+    grouped = feats.reshape(batch, token_count, num_heads, half_dim // 4, 4)
+    out = torch.einsum("bnij,bnhgj->bnhgi", matrix.float(), grouped)
+    return out.reshape(batch, token_count, num_heads, half_dim)
+
+
+def _apply_interleaved_rope(feats: torch.Tensor, rope_cos: torch.Tensor, rope_sin: torch.Tensor) -> torch.Tensor:
+    half_dim = feats.shape[-1]
+    pair_index = torch.arange(half_dim, device=feats.device) ^ 1
+    paired = feats.index_select(-1, pair_index)
+    return feats * rope_cos.reshape(1, rope_cos.shape[0], 1, half_dim) + paired * rope_sin.reshape(
+        1, rope_sin.shape[0], 1, half_dim
+    )
+
+
+def cam_prep_func(
+    q_normed: torch.Tensor,
+    k_normed: torch.Tensor,
+    v_raw: torch.Tensor,
+    *,
+    proj_q: torch.Tensor,
+    proj_kv: torch.Tensor,
+    rope_cos: torch.Tensor,
+    rope_sin: torch.Tensor,
+    k_scale: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Native equivalent of NVlabs ``cam_prep_func``.
+
+    The caller is responsible for the q/k RMSNorm, which must run through the
+    model's norm modules so a tensor-parallel shard still normalises over the
+    global channel width. This function only applies ReLU, the K scale, the
+    UCPE ray projection and RoPE.
+
+    Args:
+        q_normed: ``(B, N, H, D)`` camera Q, already RMSNorm'd.
+        k_normed: ``(B, N, H, D)`` camera K, already RMSNorm'd. K must also
+            include the temporal short convolution.
+        v_raw: ``(B, N, H, D)`` raw camera V.
+        proj_q: ``(B, N, 4, 4)`` Q ray projection matrices.
+        proj_kv: ``(B, N, 4, 4)`` K/V ray projection matrices.
+        rope_cos: ``(N, D//2)`` interleaved-pair RoPE cosine table.
+        rope_sin: ``(N, D//2)`` interleaved-pair RoPE sine table.
+        k_scale: ``D^-0.5 * spatial_tokens^-0.5``.
+
+    Returns:
+        ``q_trans``, ``k_trans``, ``v_trans`` in ``(B, H, D, N)`` layout, and
+        ``inflation_sq`` in ``(B, H, N)``.
+    """
+    half_dim = _validate_cam_prep_inputs(
+        q_normed,
+        k_normed,
+        v_raw,
+        proj_q,
+        proj_kv,
+        rope_cos,
+        rope_sin,
+    )
+
+    q_norm = torch.relu(q_normed.float())
+    k_norm = torch.relu(k_normed.float()) * float(k_scale)
+    value = v_raw.float()
+
+    q_half = _apply_ray_projection(q_norm[..., :half_dim], proj_q)
+    k_half = _apply_ray_projection(k_norm[..., :half_dim], proj_kv)
+    v_half = _apply_ray_projection(value[..., :half_dim], proj_kv)
+
+    q_rope = _apply_interleaved_rope(q_norm[..., half_dim:], rope_cos.float(), rope_sin.float())
+    k_rope = _apply_interleaved_rope(k_norm[..., half_dim:], rope_cos.float(), rope_sin.float())
+    v_rope = _apply_interleaved_rope(value[..., half_dim:], rope_cos.float(), rope_sin.float())
+
+    k_pre_sq = k_norm.square().sum(dim=-1).permute(0, 2, 1).contiguous()
+    k_post = torch.cat((k_half, k_rope), dim=-1)
+    k_post_sq = k_post.square().sum(dim=-1).permute(0, 2, 1).contiguous()
+    inflation_sq = k_post_sq.clamp_min(1e-12) / k_pre_sq.clamp_min(1e-12)
+
+    def to_bhdn(tensor: torch.Tensor) -> torch.Tensor:
+        return tensor.permute(0, 2, 3, 1).contiguous().to(v_raw.dtype)
+
+    q_out = to_bhdn(torch.cat((q_half, q_rope), dim=-1))
+    k_out = to_bhdn(k_post)
+    v_out = to_bhdn(torch.cat((v_half, v_rope), dim=-1))
+    return q_out, k_out, v_out, inflation_sq
+
+
+__all__ = [
+    "SanaWmCamGeometry",
+    "cam_prep_func",
+    "prepare_cam_geometry",
+]

vllm_omni/diffusion/offloader/base.py

@@ -62,7 +62,7 @@ class OffloadConfig:
         dp_size = 1
         if parallel_config is not None:
             dp_size = getattr(parallel_config, "data_parallel_size", 1)
-            # HSDP's fully_shard_degree also contributes to effective DP
+            # HSDP shard and replica sizes determine the effective group size.
             hsdp_shard_size = getattr(parallel_config, "hsdp_shard_size", -1) if use_hsdp else -1
             hsdp_replicate_size = getattr(parallel_config, "hsdp_replicate_size", 1) if use_hsdp else 1
             if use_hsdp and hsdp_shard_size > 0:

vllm_omni/diffusion/registry.py

@@ -171,6 +171,11 @@ _DIFFUSION_MODELS = {
         "pipeline_ming_imagegen",
         "MingImagePipeline",
     ),
+    "SanaWmPipeline": (
+        "sana_wm",
+        "pipeline_sana_wm",
+        "SanaWmPipeline",
+    ),
     "InternVLAA1Pipeline": (
         "internvla_a1",
         "pipeline_internvla_a1",
@@ -633,6 +638,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "LingBotVideoPipeline": "get_lingbot_video_pre_process_func",
     "HunyuanImage3ForCausalMM": "get_hunyuan_image_3_pre_process_func",
     "MagiHumanPipeline": "get_magi_human_pre_process_func",
+    "SanaWmPipeline": "get_sana_wm_pre_process_func",
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func",
     "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func",
     "SoulXSingerPipeline": "get_soulxsinger_pre_process_func",

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -279,9 +279,8 @@ class DiffusionWorker:
                 allgather_degree=parallel_config.allgather_degree,
                 tensor_parallel_size=parallel_config.tensor_parallel_size,
                 pipeline_parallel_size=parallel_config.pipeline_parallel_size,
-                fully_shard_degree=parallel_config.hsdp_shard_size if parallel_config.use_hsdp else 1,
-                hsdp_replicate_size=parallel_config.hsdp_replicate_size if parallel_config.use_hsdp else 1,
                 enable_expert_parallel=parallel_config.enable_expert_parallel,
+                use_hsdp=parallel_config.use_hsdp,
             )
             if (
                 getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)

vllm_omni/engine/async_omni_engine.py

@@ -971,7 +971,7 @@ class AsyncOmniEngine:
             ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict"
             sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size")
             pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
-            data_parallel_size = normalized_kwargs.get("data_parallel_size") or 1
+            data_parallel_size = normalized_kwargs.get("data_parallel_size")
             tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1
             cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1
             pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
@@ -1004,6 +1004,11 @@ class AsyncOmniEngine:
                 hsdp_replicate_size=hsdp_replicate_size,
             )
 
+        num_gpus = normalized_kwargs.get("num_gpus")
+        if num_gpus is not None:
+            num_gpus = int(num_gpus)
+            parallel_config.resolve_data_parallel_size(num_gpus)
+
         num_devices = max(1, int(parallel_config.world_size))
         devices = ",".join(str(i) for i in range(num_devices))
         model_class_name = kwargs.get("model_class_name", None)
@@ -1092,6 +1097,8 @@ class AsyncOmniEngine:
                 else {}
             ),
         }
+        if num_gpus is not None:
+            stage_engine_args["num_gpus"] = num_gpus
         # Only set dtype if it was already explicitly passed and normalized
         if "dtype" in normalized_kwargs:
             stage_engine_args["dtype"] = normalized_kwargs["dtype"]

vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py

@@ -10,7 +10,12 @@ from typing import TYPE_CHECKING
 from vllm.inputs import tokens_input
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, conditioning_cache_salt
+from vllm_omni.entrypoints.openai.tts_adapters.base import (
+    ARTTSAdapter,
+    PreparedRequest,
+    apply_max_new_tokens,
+    conditioning_cache_salt,
+)
 
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -104,6 +109,15 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         prompt["cache_salt"] = conditioning_cache_salt(request, tts_params)
         return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name)
 
+    def apply_sampling_overrides(
+        self,
+        sampling_params_list: list,
+        request: "OpenAICreateSpeechRequest",
+        prompt: dict | None = None,
+        request_id: str | None = None,
+    ) -> list:
+        return apply_max_new_tokens(sampling_params_list, request)
+
 
 @register_tts_adapter
 class MossTTSNanoAdapter(_MossTTSAdapterBase):

vllm_omni/model_executor/model_loader/weight_utils.py

@@ -1,3 +1,4 @@
+import os
 import time
 from pathlib import Path
 
@@ -85,3 +86,46 @@ def download_weights_from_hf_specific(
                 time_taken,
             )
     return hf_folder
+
+
+def resolve_model_to_local_path(
+    model: str,
+    *,
+    allow_download: bool = False,
+    allow_patterns: list[str] | None = None,
+    cache_dir: str | None = None,
+) -> str:
+    """Resolve a model reference to a local directory.
+    Failure is always raised: whether a failure is fatal is the caller's decision.
+
+    Args:
+        model: A local directory or an HF Hub repo ID.
+        allow_download: Fetch missing files instead of requiring a warm cache.
+            Downloads go through :func:`download_weights_from_hf_specific`
+        allow_patterns: Restrict the fetch to these glob patterns.
+            Defaults to the whole repository. Ignored when model is already a directory.
+        cache_dir: Download cache location; None uses the HF default.
+
+    Returns:
+        Path to a local directory containing the model files.
+    """
+    if os.path.isdir(model):
+        return model
+
+    if allow_download:
+        return download_weights_from_hf_specific(
+            model_name_or_path=model,
+            cache_dir=cache_dir,
+            allow_patterns=allow_patterns or ["*"],
+            require_all=True,
+        )
+
+    # Cache-only lookups stay on huggingface_hub even under VLLM_USE_MODELSCOPE:
+    # ModelScope's snapshot_download has no local_files_only equivalent,
+    # so a ModelScope deployment that needs resolution here must pass allow_download.
+    return huggingface_hub.snapshot_download(
+        model,
+        cache_dir=cache_dir,
+        allow_patterns=allow_patterns,
+        local_files_only=True,
+    )

vllm_omni/model_executor/models/common/ming/__init__.py

@@ -1,2 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Components shared by the Ming-family model packages.
+
+Naming: The flow-matching pieces all use the cfm_ prefix,
+because "flow matching", "fm" and "cfm" were three names for one concept here:
+
+* dit_blocks    - DiT building blocks (DiTBlock, FinalLayer, CondEmbedder)
+* cfm_dit       - the DiT used as the CFM velocity model
+* cfm_solver    - timestep schedules and the ODE/SDE integrator
+* cfm_head      - inference-only CFM wrapper preserving cfm.model.* names
+* cfm_cudagraph - batch-1 CUDA-graph executor for the sampling loop
+"""

vllm_omni/model_executor/models/common/ming/aggregator.py

@@ -7,7 +7,7 @@ import torch
 import torch.nn as nn
 from x_transformers.x_transformers import RotaryEmbedding
 
-from vllm_omni.model_executor.models.common.ming.dit import DiTBlock, FinalLayer
+from vllm_omni.model_executor.models.common.ming.dit_blocks import DiTBlock, FinalLayer
 
 
 class Aggregator(nn.Module):

vllm_omni/model_executor/models/common/ming/cfm_cudagraph.py

@@ -17,7 +17,7 @@ from __future__ import annotations
 import torch
 import torch.nn as nn
 
-from vllm_omni.model_executor.models.common.ming.fm import build_timesteps
+from vllm_omni.model_executor.models.common.ming.cfm_solver import build_timesteps
 
 
 class CFMSampler(nn.Module):

vllm_omni/model_executor/models/common/ming/cfm_dit.py

@@ -7,7 +7,7 @@ import torch.nn as nn
 from x_transformers.x_transformers import RotaryEmbedding
 
 from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding
-from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer
+from vllm_omni.model_executor.models.common.ming.dit_blocks import CondEmbedder, DiTBlock, FinalLayer
 
 
 class DiT(nn.Module):

vllm_omni/model_executor/models/common/ming/cfm_head.py

@@ -5,9 +5,8 @@
 import torch
 import torch.nn as nn
 
-from vllm_omni.model_executor.models.common.ming.fm import Solver, build_timesteps
+from vllm_omni.model_executor.models.common.ming.cfm_dit import DiT
-
+from vllm_omni.model_executor.models.common.ming.cfm_solver import Solver, build_timesteps
-from .fm.dit import DiT
 
 
 class FlowLoss(nn.Module):

vllm_omni/model_executor/models/common/ming/cfm_solver.py

@@ -3,7 +3,21 @@
 
 import torch
 
-from .dit import get_epss_timesteps
+
+def get_epss_timesteps(n, device, dtype):
+    dt = 1 / 32
+    predefined_timesteps = {
+        5: [0, 2, 4, 8, 16, 32],
+        6: [0, 2, 4, 6, 8, 16, 32],
+        7: [0, 2, 4, 6, 8, 16, 24, 32],
+        10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],
+        12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],
+        16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],
+    }
+    t = predefined_timesteps.get(n, [])
+    if not t:
+        return torch.linspace(0, 1, n + 1, device=device, dtype=dtype)
+    return dt * torch.tensor(t, device=device, dtype=dtype)
 
 
 def apply_sway_sampling(t, sway_sampling_coef):

vllm_omni/model_executor/models/common/ming/dit_blocks.py

@@ -216,19 +216,3 @@ class CondEmbedder(nn.Module):
 
     def forward(self, llm_cond):
         return self.cond_embedder(llm_cond)
-
-
-def get_epss_timesteps(n, device, dtype):
-    dt = 1 / 32
-    predefined_timesteps = {
-        5: [0, 2, 4, 8, 16, 32],
-        6: [0, 2, 4, 6, 8, 16, 32],
-        7: [0, 2, 4, 6, 8, 16, 24, 32],
-        10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],
-        12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],
-        16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],
-    }
-    t = predefined_timesteps.get(n, [])
-    if not t:
-        return torch.linspace(0, 1, n + 1, device=device, dtype=dtype)
-    return dt * torch.tensor(t, device=device, dtype=dtype)

vllm_omni/model_executor/models/common/ming/speaker_extractor.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/spkemb_extractor.py
+"""CAM++ speaker-embedding extraction, shared by the Ming model packages."""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import torchaudio
+from vllm.multimodal.media.audio import load_audio
+
+from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
+from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor
+
+
+class SpeakerEmbeddingExtractor:
+    """CAM++ ONNX speaker embeddings, resampled to the extractor's rate."""
+
+    def __init__(self, model, target_sr=16000, *, allow_download: bool = False):
+        local_model_path = resolve_model_to_local_path(
+            model,
+            allow_download=allow_download,
+            allow_patterns=["campplus.onnx"],
+        )
+        campplus_path = os.path.join(local_model_path, "campplus.onnx")
+        if not os.path.exists(campplus_path):
+            raise RuntimeError(f"Missing Ming speaker extractor model: {campplus_path}")
+
+        self.target_sr = int(target_sr)
+        self._core = SpkembExtractor(campplus_path, target_sr=self.target_sr)
+
+    def extract_from_waveform(self, waveform, sample_rate):
+        if not isinstance(waveform, torch.Tensor):
+            waveform = torch.as_tensor(waveform)
+
+        tensor = waveform.detach().to(torch.float32)
+        if tensor.ndim == 1:
+            tensor = tensor.unsqueeze(0)
+        if int(sample_rate) != self.target_sr:
+            tensor = torchaudio.transforms.Resample(orig_freq=int(sample_rate), new_freq=self.target_sr)(tensor)
+
+        embedding = self._core._extract_spk_embedding(tensor)
+        return embedding.squeeze(0).to(dtype=torch.float32)
+
+    def extract_from_file(self, audio_path):
+        audio, sample_rate = load_audio(audio_path, sr=None, mono=False)
+        waveform = torch.from_numpy(audio).contiguous()
+        return self.extract_from_waveform(waveform, sample_rate)
+
+    def extract_many(self, audio_paths):
+        return [self.extract_from_file(path) for path in audio_paths]
+
+
+__all__ = ["SpeakerEmbeddingExtractor"]

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py

@@ -25,14 +25,15 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.sequence import IntermediateTensors
 
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
-from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
+from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
 from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE, AudioVAEConfig
 from vllm_omni.model_executor.models.ming_tts.constants import SPEAKER_EMBEDDING_DIM
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.transformers_utils.configs.ming_flash_omni import MingFlashOmniTalkerConfig
 
 from .prompt_utils import DEFAULT_PROMPT as MING_DEFAULT_PROMPT
-from .talker_module import CFM, Aggregator, DiT, MingAudioGenerator, build_tts_input
+from .talker_module import CFM, DiT, MingAudioGenerator, build_tts_input
 from .text_processing import segment_and_normalize
 from .voice_presets import VoicePresetRegistry
 
@@ -119,7 +120,11 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin)
             DiT(llm_input_dim=self.hidden_size, **config.flowmodel),
             steps=config.steps,
         )
-        self.aggregator = Aggregator(llm_input_dim=self.hidden_size, **config.aggregator)
+        # config.aggregator still wins `in_channels` if the checkpoint states it.
+        self.aggregator = Aggregator(
+            llm_input_dim=self.hidden_size,
+            **{"in_channels": self.latent_dim, **config.aggregator},
+        )
         self.stop_head = nn.Linear(self.hidden_size, 2, bias=True)
         self.spk_head = nn.Linear(SPEAKER_EMBEDDING_DIM, self.hidden_size, bias=True)
 
@@ -525,27 +530,22 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin)
 
     def _iter_talker_safetensors(self) -> Iterable[tuple[str, torch.Tensor]]:
         """Yield (name, tensor) pairs from talker/model*.safetensors."""
-        model_path = self._model_path
+        model_root = resolve_model_to_local_path(
-        # Try local path first
+            self._model_path,
-        for candidate in (os.path.join(model_path, "talker"), model_path):
+            allow_download=True,
+            allow_patterns=["talker/model*.safetensors"],
+            cache_dir=self.vllm_config.load_config.download_dir,
+        )
+
+        # Nested talker/ layout first, then a flat checkpoint at the root.
+        for candidate in (os.path.join(model_root, "talker"), model_root):
             sf_files = sorted(glob_module.glob(os.path.join(candidate, "model*.safetensors")))
             if sf_files:
                 for sf_path in sf_files:
                     yield from load_file(sf_path, device="cpu").items()
                 return
 
-        # HF hub fallback: download only the talker checkpoint files
+        raise RuntimeError(f"No talker safetensors found under {model_root}. Expected talker/model*.safetensors.")
-        model_root = download_weights_from_hf_specific(
-            model_path,
-            self.vllm_config.load_config.download_dir,
-            allow_patterns=["talker/model*.safetensors"],
-        )
-        talker_dir = os.path.join(model_root, "talker")
-        sf_files = sorted(glob_module.glob(os.path.join(talker_dir, "model*.safetensors")))
-        if not sf_files:
-            raise RuntimeError(f"No talker safetensors found under {model_root}. Expected talker/model*.safetensors.")
-        for sf_path in sf_files:
-            yield from load_file(sf_path, device="cpu").items()
 
     def _load_vae_weights(self) -> set[str]:
         """Load AudioVAE weights from talker/vae/model.safetensors."""

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py

@@ -64,6 +64,7 @@ from vllm.sequence import IntermediateTensors
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
+from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config
 from vllm_omni.transformers_utils.processors.ming import (
@@ -715,12 +716,8 @@ class MingFlashOmniThinkerForConditionalGeneration(
 
         p = Path(model_path) / "mlp" / "model.safetensors"
         if not p.exists() and not os.path.exists(model_path):
-            from vllm_omni.model_executor.model_loader.weight_utils import (
-                download_weights_from_hf_specific,
-            )
-
             try:
-                local_path = download_weights_from_hf_specific(model_path, None, ["mlp/*"])
+                local_path = resolve_model_to_local_path(model_path, allow_download=True, allow_patterns=["mlp/*"])
                 p = Path(local_path) / "mlp" / "model.safetensors"
             except Exception:
                 logger.exception("[MingFlashOmniThinker] failed to download mlp/ from %s", model_path)

vllm_omni/model_executor/models/ming_flash_omni/talker_module.py

@@ -33,9 +33,14 @@ from vllm.platforms import current_platform
 from x_transformers.x_transformers import RotaryEmbedding
 
 from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding
+from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
 from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE
-from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer, get_epss_timesteps
+from vllm_omni.model_executor.models.common.ming.cfm_solver import (
-from vllm_omni.model_executor.models.common.ming.fm import apply_sway_sampling, integrate_cfm_steps
+    apply_sway_sampling,
+    get_epss_timesteps,
+    integrate_cfm_steps,
+)
+from vllm_omni.model_executor.models.common.ming.dit_blocks import CondEmbedder, DiTBlock, FinalLayer
 
 logger = init_logger(__name__)
 
@@ -458,58 +463,6 @@ def silence_holder(
     return speech_out, sil_cache
 
 
-########################################################################
-# Audio Postprocess
-# Ported from:
-# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/aggregator.py
-########################################################################
-
-
-class Aggregator(nn.Module):
-    """Maps generated audio latent patches back to LLM embedding space."""
-
-    def __init__(
-        self,
-        in_channels: int = 64,
-        hidden_size: int = 1152,
-        depth: int = 28,
-        num_heads: int = 16,
-        mlp_ratio: float = 4.0,
-        llm_input_dim: int = 896,
-        **kwargs,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = in_channels
-        self.num_heads = num_heads
-
-        self.word_embedder = nn.Embedding(1, hidden_size)
-        self.x_embedder = nn.Linear(in_channels, hidden_size)
-        self.hidden_size = hidden_size
-
-        self.rotary_embed = RotaryEmbedding(hidden_size // num_heads)
-
-        self.blocks = nn.ModuleList(
-            [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)]
-        )
-        self.final_layer = FinalLayer(hidden_size, llm_input_dim)
-
-    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
-        x = self.x_embedder(x)
-        cls_embed = self.word_embedder(torch.zeros((x.shape[0], 1), dtype=torch.long, device=x.device))
-        x = torch.cat([cls_embed, x], dim=1)
-
-        rope = self.rotary_embed.forward_from_seq_len(x.shape[1])
-        if mask is not None:
-            mask_pad = mask.clone().detach()[:, :1]
-            mask = torch.cat([mask_pad, mask], dim=-1)
-        for block in self.blocks:
-            x = block(x, mask, rope)
-        x = self.final_layer(x)
-        x = x[:, :1, :]
-        return x
-
-
 ########################################################################
 # Prompt Builder
 # Adapted from:

vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py

@@ -16,16 +16,15 @@ import torch
 from transformers.utils.hub import cached_file
 from vllm.logger import init_logger
 
-from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
+from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor
 
-from .spk_embedding import SpkembExtractor
 from .talker_module import resample
 
 if TYPE_CHECKING:
+    from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
     from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE
 
-    from .talker_module import Aggregator
-
 logger = init_logger(__name__)
 
 
@@ -273,18 +272,19 @@ class VoicePresetRegistry:
             if os.path.isfile(path):
                 return path, candidate
 
-        if not os.path.isdir(self._model_path):
+        try:
-            try:
+            hf_root = resolve_model_to_local_path(
-                hf_root = download_weights_from_hf_specific(
+                self._model_path,
-                    self._model_path,
+                allow_download=True,
-                    self._download_dir,
+                allow_patterns=["talker/data/**"],
-                    allow_patterns=["talker/data/**"],
+                cache_dir=self._download_dir,
-                    require_all=True,
+            )
-                )
+        except Exception as e:
-                candidate = os.path.join(hf_root, "talker", "data", "voice_name.json")
+            logger.info("Could not download voice presets from HF: %s", e)
-                if os.path.isfile(candidate):
+            return None, None
-                    return candidate, os.path.join(hf_root, "talker")
+
-            except Exception as e:  # pragma: no cover
+        candidate = os.path.join(hf_root, "talker", "data", "voice_name.json")
-                logger.info("Could not download voice presets from HF: %s", e)
+        if os.path.isfile(candidate):
+            return candidate, os.path.join(hf_root, "talker")
 
         return None, None

vllm_omni/model_executor/models/ming_tts/fm/__init__.py

@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py

@@ -11,15 +11,21 @@ import torch
 import torch.nn as nn
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
+from vllm.model_executor.models.bailing_moe import BailingMoeModel
 from vllm.model_executor.models.qwen2 import Qwen2Model
 from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
 from vllm.sequence import IntermediateTensors
 from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 
+from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
+from vllm_omni.model_executor.models.common.ming.cfm_cudagraph import (
+    CFMGraphExecutor,
+    CFMSampler,
+)
+from vllm_omni.model_executor.models.common.ming.cfm_head import FlowLoss
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 
-from .aggregator import Aggregator
 from .config_ming_tts import (
     KEY_CFG,
     KEY_DECODE_STEP,
@@ -34,7 +40,6 @@ from .config_ming_tts import (
     MingTTSConfig,
 )
 from .constants import SPEAKER_EMBEDDING_DIM
-from .flowloss_head import FlowLoss
 from .patch_emission import (
     MING_STOP_REASON_CODES,
     MING_STOP_REASON_CONTINUE,
@@ -82,8 +87,6 @@ class MingLLMModel(nn.Module):
         self.prefix = prefix
         self.fm_dtype = _resolve_ming_runtime_dtype(vllm_config)
         if self.ming_config.model_variant == "moe":
-            from vllm.model_executor.models.bailing_moe import BailingMoeModel
-
             # BailingMoeModel reads ``vllm_config.model_config.hf_config`` directly
             # (no get_text_config()), so re-wrap with the nested bailing_moe config.
             llm_config = vllm_config.model_config.hf_config.get_text_config()
@@ -335,8 +338,6 @@ class MingLLMModel(nn.Module):
         if self._cfm_graph is not None:
             return self._cfm_graph
         try:
-            from .fm.cfm_cudagraph import CFMGraphExecutor, CFMSampler
-
             # TODO(perf):
             #   (1) compile per-block (dit_model.blocks) instead of the
             #   whole forward for tighter fusion / fewer graph breaks;

vllm_omni/model_executor/models/ming_tts/prompt_encoder.py

@@ -12,7 +12,7 @@ import torch
 from safetensors import safe_open
 from vllm.logger import init_logger
 
-from vllm_omni.engine.stage_init_utils import _resolve_model_to_local_path
+from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
 from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE
 
 from .audio_prep import (
@@ -78,7 +78,7 @@ def _load_prompt_encoder(wrapper: Any) -> AudioVAE:
         loaded_encoder_params = set()
         with torch.no_grad():
             for shard_path in _iter_model_safetensors(
-                _resolve_model_to_local_path(str(wrapper.vllm_config.model_config.model))
+                resolve_model_to_local_path(str(wrapper.vllm_config.model_config.model))
             ):
                 with safe_open(str(shard_path), framework="pt", device="cpu") as handle:
                     for key in handle.keys():

vllm_omni/model_executor/models/ming_tts/speaker_extractor.py

@@ -3,16 +3,13 @@
 # Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/spkemb_extractor.py
 from __future__ import annotations
 
-import os
 import threading
 from typing import TYPE_CHECKING, Any
 
 import torch
-import torchaudio
 from vllm.logger import init_logger
-from vllm.multimodal.media.audio import load_audio
 
-from vllm_omni.model_executor.models.ming_flash_omni.spk_embedding import SpkembExtractor
+from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor
 
 from .audio_prep import coerce_prompt_waveform, coerce_speaker_embeddings
 from .constants import (
@@ -29,48 +26,9 @@ logger = init_logger(__name__)
 _EXTRACTOR_LOAD_LOCK = threading.Lock()
 
 
-def resolve_model_to_local_path(model):
+# Ming's bundle ships the CAM++ artifact as `campplus.onnx` in the model
-    if os.path.isdir(model):
+# directory, which is the shared extractor's default.
-        return model
+def _load_speaker_extractor(wrapper: MingTTSForConditionalGeneration) -> SpeakerEmbeddingExtractor:
-
-    from huggingface_hub import snapshot_download
-
-    return snapshot_download(model)
-
-
-class MingSpeakerEmbeddingExtractor:
-    def __init__(self, model, target_sr=16000):
-        local_model_path = resolve_model_to_local_path(model)
-        campplus_path = os.path.join(local_model_path, "campplus.onnx")
-        if not os.path.exists(campplus_path):
-            raise RuntimeError(f"Missing Ming speaker extractor model: {campplus_path}")
-
-        self.target_sr = int(target_sr)
-        self._core = SpkembExtractor(campplus_path, target_sr=self.target_sr)
-
-    def extract_from_waveform(self, waveform, sample_rate):
-        if not isinstance(waveform, torch.Tensor):
-            waveform = torch.as_tensor(waveform)
-
-        tensor = waveform.detach().to(torch.float32)
-        if tensor.ndim == 1:
-            tensor = tensor.unsqueeze(0)
-        if int(sample_rate) != self.target_sr:
-            tensor = torchaudio.transforms.Resample(orig_freq=int(sample_rate), new_freq=self.target_sr)(tensor)
-
-        embedding = self._core._extract_spk_embedding(tensor)
-        return embedding.squeeze(0).to(dtype=torch.float32)
-
-    def extract_from_file(self, audio_path):
-        audio, sample_rate = load_audio(audio_path, sr=None, mono=False)
-        waveform = torch.from_numpy(audio).contiguous()
-        return self.extract_from_waveform(waveform, sample_rate)
-
-    def extract_many(self, audio_paths):
-        return [self.extract_from_file(path) for path in audio_paths]
-
-
-def _load_speaker_extractor(wrapper: MingTTSForConditionalGeneration) -> MingSpeakerEmbeddingExtractor:
     extractor = getattr(wrapper, "_speaker_extractor", None)
     if extractor is not None:
         return extractor
@@ -79,7 +37,7 @@ def _load_speaker_extractor(wrapper: MingTTSForConditionalGeneration) -> MingSpe
         extractor = getattr(wrapper, "_speaker_extractor", None)
         if extractor is None:
             model_path = wrapper.vllm_config.model_config.model
-            extractor = MingSpeakerEmbeddingExtractor(model_path, target_sr=16000)
+            extractor = SpeakerEmbeddingExtractor(model_path, target_sr=16000)
             wrapper._speaker_extractor = extractor
     return extractor
 
@@ -174,8 +132,6 @@ def _resolve_speaker_embeddings(
 
 
 __all__ = [
-    "MingSpeakerEmbeddingExtractor",
     "_load_speaker_extractor",
     "_resolve_speaker_embeddings",
-    "resolve_model_to_local_path",
 ]

vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py

@@ -58,6 +58,25 @@ class BatchedToken2WavState:
     hift_cache: dict[str, torch.Tensor]
 
 
+def _undecorate_dynamo(module: nn.Module, method: str) -> None:
+    """Restore ``method`` on ``module`` if TorchDynamo wrapped it.
+
+    ``cosyvoice2`` decorates ``UpsampleConformerEncoderV2.forward_chunk`` with
+    ``torch.compile(backend="eager")``. That backend performs no Inductor
+    optimisation, so the wrapper only adds tracing and guard construction, and
+    duplex pays it again on every unseen chunk shape -- seconds inside a live
+    response. Dropping the wrapper leaves the original implementation, which is
+    what the eager backend was executing anyway.
+    """
+    bound = getattr(module, method, None)
+    original = getattr(bound, "_torchdynamo_orig_callable", None) or getattr(bound, "__wrapped__", None)
+    if original is None:
+        return
+    function = getattr(original, "__func__", original)
+    module.__dict__[method] = function.__get__(module, type(module))
+    logger.info("Bypassed TorchDynamo wrapper on %s.%s", type(module).__name__, method)
+
+
 class BatchedToken2Wav(nn.Module):
     """Drive Token2wav's modules with dynamically-sized, request-owned caches.
 
@@ -82,6 +101,7 @@ class BatchedToken2Wav(nn.Module):
         self._trt_stepper = trt_stepper
         self.flow = token2wav.flow
         self.hift = token2wav.hift
+        _undecorate_dynamo(self.flow.encoder, "forward_chunk")
         hift_parameter = next(self.hift.parameters(), None)
         if hift_parameter is not None and hift_parameter.device.type == "cuda":
             # Prime the CUDA state used by HiFT during backend construction.

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -17,6 +17,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
 import math
 import os
 import warnings
@@ -2505,6 +2506,17 @@ def _in_projection(
     return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)
 
 
+def _get_audio_cache_length(past_key_values: Any) -> int:
+    cache = getattr(past_key_values, "self_attention_cache", past_key_values)
+    get_seq_length = getattr(cache, "get_seq_length", None)
+    if callable(get_seq_length):
+        try:
+            return int(get_seq_length())
+        except TypeError:
+            return int(get_seq_length(0))
+    return int(cache[0][0].shape[2])
+
+
 # Copied from transformers.models.whisper.modeling_whisper.WhisperEncoderLayer and add use_cache for streaming inference
 class MiniCPMWhisperEncoderLayer(nn.Module):
     def __init__(self, config: WhisperConfig, layer_idx: int = None):
@@ -2517,6 +2529,10 @@ class MiniCPMWhisperEncoderLayer(nn.Module):
             config=config,
             layer_idx=layer_idx,
         )
+        attention_parameters = inspect.signature(self.self_attn.forward).parameters
+        self._past_key_values_kwarg = (
+            "past_key_values" if "past_key_values" in attention_parameters else "past_key_value"
+        )
         self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
         self.dropout = config.dropout
         self.activation_fn = ACT2FN[config.activation_function]
@@ -2554,16 +2570,18 @@ class MiniCPMWhisperEncoderLayer(nn.Module):
         """
         residual = hidden_states
         hidden_states = self.self_attn_layer_norm(hidden_states)
-        attn_out = self.self_attn(
+        attention_kwargs = {
-            hidden_states=hidden_states,
+            "hidden_states": hidden_states,
-            attention_mask=attention_mask,
+            "attention_mask": attention_mask,
-            layer_head_mask=layer_head_mask,
+            "layer_head_mask": layer_head_mask,
-            output_attentions=output_attentions,
+            "output_attentions": output_attentions,
-            past_key_value=past_key_values,
+            self._past_key_values_kwarg: past_key_values,
-        )
+        }
+        attn_out = self.self_attn(**attention_kwargs)
         hidden_states = attn_out[0]
         attn_weights = attn_out[1] if len(attn_out) > 1 else None
-        past_key_values = attn_out[2] if len(attn_out) > 2 else None
+        if len(attn_out) > 2:
+            past_key_values = attn_out[2]
         hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
         hidden_states = residual + hidden_states
 
@@ -2769,7 +2787,7 @@ class MiniCPMWhisperEncoder(WhisperEncoder):
                 past_key_values = EncoderDecoderCache(past_key_values, DynamicCache())
             else:
                 pass
-            past_key_values_length = past_key_values.self_attention_cache.get_usable_length(inputs_embeds.shape[1])
+            past_key_values_length = _get_audio_cache_length(past_key_values)
             if inputs_embeds.shape[1] + past_key_values_length > embed_pos.shape[0]:
                 logger.warning("seems the audio is longer than 30s. repeating the last part of the audio")
                 embed_pos_front = embed_pos[past_key_values_length:, :]
@@ -4468,7 +4486,7 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
             return []
 
         if self.audio_past_key_values is not None:
-            cache_length = self.audio_past_key_values[0][0].shape[2]
+            cache_length = _get_audio_cache_length(self.audio_past_key_values)
             apm_max_len = self.apm.embed_positions.weight.shape[0]
             if cache_length + current_seq_len >= apm_max_len:
                 logger.warning(
@@ -4480,7 +4498,7 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
 
         past_len = 0
         if self.audio_past_key_values is not None:
-            past_len = self.audio_past_key_values[0][0].shape[2]
+            past_len = _get_audio_cache_length(self.audio_past_key_values)
         total_seq_len = past_len + current_seq_len
         audio_attention_mask = torch.zeros(
             (batch_size, 1, current_seq_len, total_seq_len),

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py

@@ -128,32 +128,45 @@ def _sample_token(
     do_sample: bool,
     generator: torch.Generator | None = None,
 ) -> torch.Tensor:
-    """Top-k + top-p sampling (matches upstream's ``sample_token`` for the
+    """Top-k + top-p sampling for the upstream inference branch.
-    inference branch).
+
+    Nucleus filtering and multinomial sampling operate only on the retained
+    top-k candidates before the result is mapped back to the original
+    vocabulary. This preserves the categorical distribution when the top-k
+    boundary has no ties, but it is not seed/bit equivalent to multinomial over
+    a full-vocabulary tensor because random-number mapping depends on width.
     """
     if not do_sample or temperature <= 0:
         return logits.argmax(dim=-1)
 
     logits = logits / max(temperature, 1e-6)
+    compact_indices = None
     if top_k and top_k > 0 and top_k < logits.shape[-1]:
-        top_vals, _ = torch.topk(logits, top_k, dim=-1)
+        top_vals, top_indices = torch.topk(logits, top_k, dim=-1, sorted=True)
-        thresh = top_vals[..., -1:].expand_as(logits)
+        # topk is already descending: keep the following nucleus filter,
-        logits = torch.where(logits < thresh, torch.full_like(logits, float("-inf")), logits)
+        # softmax, and multinomial at width k instead of the full vocab.
+        logits = top_vals
+        compact_indices = top_indices
 
     if 0.0 < top_p < 1.0:
-        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
+        sorted_indices = None
-        probs = F.softmax(sorted_logits, dim=-1)
+        if compact_indices is None:
+            logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+        probs = F.softmax(logits, dim=-1)
         cum = probs.cumsum(dim=-1)
         # Drop tail beyond top_p (keep at least one token).
         drop = cum > top_p
         drop[..., 1:] = drop[..., :-1].clone()
         drop[..., 0] = False
-        sorted_logits = sorted_logits.masked_fill(drop, float("-inf"))
+        logits = logits.masked_fill(drop, float("-inf"))
-        logits = torch.full_like(logits, float("-inf")).scatter_(-1, sorted_idx, sorted_logits)
+        if sorted_indices is not None:
+            logits = torch.full_like(logits, float("-inf")).scatter_(-1, sorted_indices, logits)
 
     probs = F.softmax(logits, dim=-1)
     flat = probs.reshape(-1, probs.shape[-1])
     sampled = torch.multinomial(flat, num_samples=1, generator=generator).reshape(probs.shape[:-1])
+    if compact_indices is not None:
+        return compact_indices.gather(-1, sampled.unsqueeze(-1)).squeeze(-1)
     return sampled
 
 

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py

@@ -52,15 +52,48 @@ class _MossTTSLocalAttention(nn.Module):
         self.c_proj = nn.Linear(hidden_size, hidden_size, bias=True)
         inv_freq = 1.0 / (rope_base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32) / self.head_dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("_rope_cos_cache", torch.empty(0), persistent=False)
+        self.register_buffer("_rope_sin_cache", torch.empty(0), persistent=False)
 
-    def _rope_cos_sin(
+    def prepare_rope_cache(
-        self, seq_len: int, device: torch.device, dtype: torch.dtype
+        self,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+        max_seq_len: int,
-        position_ids = torch.arange(seq_len, device=device, dtype=torch.float32)
+        device: torch.device,
-        freqs = torch.einsum("s,d->sd", position_ids, self.inv_freq.to(device=device))
+        dtype: torch.dtype,
+    ) -> None:
+        """Materialize frame-local RoPE values once per device/dtype.
+
+        Local Depth positions always restart at zero and never exceed the
+        number of RVQ codebooks (12 for MOSS-TTS Local v1.5).  Building
+        arange/einsum/cos/sin inside every one of the 12 depth steps creates
+        many tiny kernels and also bloats the enclosing talker CUDA graph.
+        Keep one non-persistent cache and slice it for prefix execution.
+        """
+        if max_seq_len <= 0:
+            return
+        cache = self._rope_cos_cache
+        if cache.numel() > 0 and cache.device == device and cache.dtype == dtype and int(cache.shape[1]) >= max_seq_len:
+            return
+
+        position_ids = torch.arange(max_seq_len, device=device, dtype=torch.float32)
+        inv_freq = self.inv_freq.to(device=device, dtype=torch.float32)
+        freqs = torch.einsum("s,d->sd", position_ids, inv_freq)
         cos = freqs.cos().repeat_interleave(2, dim=-1).to(dtype)
         sin = freqs.sin().repeat_interleave(2, dim=-1).to(dtype)
-        return cos.view(1, seq_len, 1, self.head_dim), sin.view(1, seq_len, 1, self.head_dim)
+        self._rope_cos_cache = cos.view(1, max_seq_len, 1, self.head_dim)
+        self._rope_sin_cache = sin.view(1, max_seq_len, 1, self.head_dim)
+
+    def _rope_cos_sin(
+        self,
+        seq_len: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        self.prepare_rope_cache(seq_len, device, dtype)
+        return (
+            self._rope_cos_cache[:, :seq_len],
+            self._rope_sin_cache[:, :seq_len],
+        )
 
     @staticmethod
     def _rotate_half(x: torch.Tensor) -> torch.Tensor:
@@ -200,9 +233,15 @@ class MossTTSLocalDepthTransformer(nn.Module):
         batch_size = backbone_last_hidden.shape[0]
         dtype = self.ln_f.weight.dtype
 
+        # Populate the complete [0, n_vq) table before torch.compile traces
+        # _forward_prefix or the runner captures talker_mtp.  The compiled /
+        # captured body then contains only fixed-address cache slices rather
+        # than cache construction or buffer replacement.
+        for block in self.h:
+            block.attn.prepare_rope_cache(n_vq, backbone_last_hidden.device, dtype)
+
         embeds = backbone_last_hidden.new_zeros((batch_size, n_vq, self.hidden_size), dtype=dtype)
         embeds[:, 0, :] = backbone_last_hidden.to(dtype)
-
         hidden = self._run_prefix(embeds[:, :1, :])
         local_hidden = hidden[:, 0, :]
 

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -1323,7 +1323,11 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
         self.mtp_hidden_size = hidden_size
         self.talker_mtp_graph_safe = True
         self.talker_mtp_output_key = ("audio_codes", "current")
-        self.use_async_omni_output = False
+        # ``make_omni_output`` keeps code rows fixed-shape and performs all
+        # state updates eagerly, so the runner can safely pack/snapshot them
+        # before CUDA-graph buffers are reused.  This moves the pinned D2H and
+        # output materialization off the Stage-0 execution thread.
+        self.use_async_omni_output = True
         self.eager_omni_postprocess_before_async_output = True
         self.omni_pooler_payload_include_hidden = False
         self.postprocess_uses_multimodal_outputs = False
@@ -1466,7 +1470,7 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
     ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
         """Build per-step input embeddings (text + audio additive fusion).
 
-        Prefill: initialise the per-request ``{is_stopping, step}`` state.
+        Prefill: initialise the per-request ``{is_stopping}`` state.
         Decode: combine the forced text token (from ``compute_logits``) with
         the audio codes the local transformer produced last step.
         """
@@ -1474,8 +1478,9 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
         span_len = int(input_ids.shape[0])
         audio_state = info_dict.get("audio_state")
         is_first_call = not isinstance(audio_state, dict)
+        is_prefill = bool(info_dict.get("_omni_is_prefill", False))
 
-        if span_len > 1 or is_first_call:
+        if span_len > 1 or is_first_call or is_prefill:
             embeds = self.model.embed_tokens(input_ids)
 
             ref_codes = (info_dict.get("codes", {}) or {}).get("ref")
@@ -1668,9 +1673,14 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
                 should_continue = codes.ne(self.audio_pad_token_id).any(dim=-1)
                 should_values.append(should_continue.reshape(-1)[:1])
                 have_mtp_output = True
-                emitted_codes = codes[should_continue]
+                # Do not boolean-index a CUDA tensor here.  Even for the
-                per_req_codes.append(emitted_codes)
+                # common single-row case, ``codes[should_continue]`` must
-                have_codes = have_codes or emitted_codes.numel() > 0
+                # discover a dynamic output size on the host and therefore
+                # synchronizes the CUDA stream once per request.  Keep the
+                # fixed-shape row (including an all-pad stop row); the async
+                # chunk processor already filters all-pad rows on CPU.
+                per_req_codes.append(codes)
+                have_codes = True
             else:
                 should_values.append(torch.ones((1,), device=hidden.device, dtype=torch.bool))
                 per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))