Add normalized online FP8 cache for DLO

Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -105,6 +105,31 @@ flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
 transforms. Model-specific direct-layout knowledge, when required, lives in a
 checkpoint adapter beside the ordinary loader.
 
+### Phase-I normalized online-FP8 cache
+
+The direct-checkpoint plan above cannot quantize a BF16 source online. Phase I
+adds an explicit, opt-in normalized cache for the narrower case of one TP1
+transformer with per-tensor online FP8 and DLO AllGather:
+
+```text
+cache miss:  BF16 source -> ordinary online FP8 loader -> atomic FP8 cache
+cache hit:   FP8 cache   -> direct mmap plan -> DLO host shard -> AllGather
+```
+
+The cache stores the serialized runtime representation, including finalized FP8
+weights, generated scales, and persistent transformer buffers. Its manifest
+records the schema, component, quantization contract, tensor count, source
+files, and a source fingerprint. A missing, stale, incomplete, incompatible, or
+invalid artifact is a cache miss; cache publication is best-effort and falls
+back to the ordinary loader on failure.
+
+Phase I keeps the cache binding deliberately simple: one source per transformer
+component, with existing per-tensor transforms applied before publication. The
+cache is not selected for DLO-disabled, no-AllGather, TP>1, or HSDP deployments,
+and it normalizes the transformer only. Cache sharing and compatibility for
+other components, transformed TP coordinates, and broader DP/SP layouts remain
+in [RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
+
 ### AllGather path
 
 With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -211,10 +236,10 @@ reconstructs those tensors with their recorded layouts. Other online methods
 must use `--dlo-no-use-allgather` or disable online quantization until their
 runtime layouts are validated.
 
-A normalized runtime mmap cache, built through the ordinary loader, is the
-proposed general mechanism for sharing transformed TP or quantized layouts.
-That cache and its publication/lifecycle protocol are intentionally outside
-this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
+The Phase-I normalized online-FP8 cache is the supported TP1 transformer case.
+The general cache-compatibility contract for additional components, transformed
+TP coordinates, and cross-DP/SP reuse remains in
+[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
 
 ## Validation coverage
 

docs/models/supported_models.md

@@ -39,7 +39,6 @@ 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/offloader/distributed_layerwise_offload.md

@@ -42,6 +42,12 @@ vllm serve /path/to/model --omni \
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --usp 4
+
+# Opt-in normalized online-FP8 cache (TP1 + DLO AllGather)
+vllm serve /path/to/model --omni \
+  --quantization fp8 \
+  --enable-distributed-layerwise-offload \
+  --dlo-fp8-cache-dir /path/to/node-local/fp8-cache
 ```
 
 ```python
@@ -63,6 +69,7 @@ omni = Omni(
 | `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
 | `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
 | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
+| `--dlo-fp8-cache-dir PATH` | Opt into the Phase-I normalized online-FP8 cache for the transformer | unset |
 
 ## Host-weight loading
 
@@ -104,6 +111,24 @@ process.
 When the effective DLO group size is one, `dlo_use_allgather=True` does not
 perform a collective and uses the same rank-local transfer behavior.
 
+## Normalized online-FP8 cache
+
+`--dlo-fp8-cache-dir` is an opt-in Phase-I cache for per-tensor online FP8 with
+TP1 and DLO AllGather. On a cache miss, the ordinary loader quantizes the
+transformer once and publishes finalized FP8 weights and scales as an atomic,
+serialized safetensors artifact. On a valid cache hit, the loader selects that
+serialized source before DLO planning, so the transformer can be loaded through
+the direct checkpoint-mmap path without re-materializing the full BF16 source.
+
+The cache is keyed to the source files through the manifest's schema, component,
+quantization contract, and source fingerprint. Incomplete, incompatible, or
+stale artifacts are treated as cache misses. Cache creation is best-effort: a
+publication failure logs a warning and keeps the ordinary online-FP8 path.
+
+Phase I normalizes the transformer only. Other pipeline components keep their
+ordinary loading and quantization behavior. The cache is not selected when DLO
+is disabled, AllGather is disabled, TP is greater than one, or HSDP is enabled.
+
 ## Declarative topology
 
 Models may declare an `OffloadPlan` instead of embedding offload logic:
@@ -147,10 +172,9 @@ must enter each collective.
   `OffloadPlan` that declares eligible `resident_dit_paths`.
 - DP concurrency requires an explicit, identical inference-step count.
 
-Sharing transformed TP or quantized runtime layouts through a normalized mmap
-cache is a follow-up design in
-[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
-the direct-checkpoint path.
+Sharing transformed TP or quantized runtime layouts beyond this TP1 transformer
+cache remains a follow-up design in
+[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
 
 See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
 for an end-to-end example.

docs/user_guide/diffusion_features.md

@@ -168,14 +168,6 @@ 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`.
@@ -183,7 +175,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, SANA-WM
+- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni
 
 ### AudioGen
 

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -53,6 +53,7 @@ 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
@@ -138,6 +139,12 @@ 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":
@@ -164,10 +171,7 @@ async def main():
         vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_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,
+        parallel_config=parallel_config,
         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,6 +154,8 @@ 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."""
 
@@ -165,7 +167,7 @@ def build_omni_kwargs(
     return {
         "model": model,
         "flow_shift": flow_shift,
-        "tensor_parallel_size": args.tensor_parallel_size,
+        "parallel_config": parallel_config,
         "enforce_eager": args.enforce_eager,
         "model_config": {"lingbot_action_root": str(paths.action_root)},
     }
@@ -260,10 +262,12 @@ 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
 
-    omni_kwargs = build_omni_kwargs(args, paths)
+    parallel_config = DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size)
+    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,6 +125,7 @@ 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
@@ -140,7 +141,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,
-        tensor_parallel_size=args.tensor_parallel_size,
+        parallel_config=DiffusionParallelConfig(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,6 +96,7 @@ 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
@@ -451,6 +452,14 @@ 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":
@@ -481,11 +490,7 @@ def main():
         vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_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,
+        parallel_config=parallel_config,
         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,6 +56,7 @@ 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
@@ -502,6 +503,17 @@ 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,
@@ -512,15 +524,7 @@ 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,
-        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,
+        parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
         model_class_name=model_class_name,
         cache_backend=args.cache_backend,

examples/offline_inference/longcat_video/end2end.py

@@ -15,6 +15,7 @@ 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,
 )
@@ -235,6 +236,14 @@ 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,6 +23,7 @@ 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
@@ -249,6 +250,17 @@ 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(
@@ -259,14 +271,7 @@ def main():
         vae_use_tiling=args.vae_use_tiling,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_layerwise_offload=args.enable_layerwise_offload,
-        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,
+        parallel_config=parallel_config,
         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,6 +38,7 @@ 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
@@ -305,13 +306,17 @@ 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,
-        use_hsdp=args.use_hsdp,
-        hsdp_shard_size=args.hsdp_shard_size,
-        hsdp_replicate_size=args.hsdp_replicate_size,
+        parallel_config=parallel_config,
         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,7 +13,6 @@ 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,
@@ -22,6 +21,7 @@ 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,8 +142,7 @@ 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
-    # Runs before the engine is built, so nothing has pre-fetched campplus.onnx.
-    embs = SpeakerEmbeddingExtractor(args.model, allow_download=True).extract_many(paths)
+    embs = MingSpeakerEmbeddingExtractor(args.model).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,6 +10,7 @@ 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
@@ -461,6 +462,17 @@ 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(
@@ -469,13 +481,7 @@ def main():
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         enable_cpu_offload=args.enable_cpu_offload,
-        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,
+        parallel_config=parallel_config,
         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,6 +11,7 @@ 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
@@ -196,6 +197,14 @@ 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 = {
@@ -213,11 +222,7 @@ def main() -> None:
 
     omni_kwargs = dict(
         model=args.model,
-        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,
+        parallel_config=parallel_config,
         model_type=args.model_type,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_layerwise_offload=args.enable_layerwise_offload,

recipes/NVIDIA/SANA-WM.md

@@ -1,193 +0,0 @@
-# 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,7 +70,6 @@ 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,21 +5,20 @@
 ## Summary
 
 - Vendor: inclusionAI
-- Model: `inclusionAI/Ming-omni-tts-16.8B-A3B`, `inclusionAI/Ming-omni-tts-0.5B`
+- Model: `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: 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)
+- [Huggingface Model card](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 and the MoE 16.8B two-stage TTS pipeline deployment.
+This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B two-stage TTS pipeline deployment.
 Other hardware is welcome as community validation lands.
 
 ## CUDA
@@ -40,7 +39,11 @@ Other hardware is welcome as community validation lands.
 Launch the two-stage talker:
 
 ```bash
-vllm serve inclusionAI/Ming-omni-tts-16.8B-A3B --omni --port 8091
+vllm serve inclusionAI/Ming-omni-tts-0.5B \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --omni \
+    --port 8091 \
+    --enforce-eager
 ```
 
 #### Verification
@@ -51,7 +54,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-16.8B-A3B",
+      "model": "inclusionAI/Ming-omni-tts-0.5B",
       "input": "你好，这是 Ming 在线语音合成测试。",
       "response_format": "wav"
     }' --output ming_tts_basic.wav
@@ -71,7 +74,7 @@ Reference-audio zero-shot cloning:
 REF_AUDIO="$BASE/10002287-00000094.wav"
 jq -n --arg ref_audio "$REF_AUDIO" \
     '{
-       model: "inclusionAI/Ming-omni-tts-16.8B-A3B",
+       model: "inclusionAI/Ming-omni-tts-0.5B",
        input: "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。",
        ref_audio: $ref_audio,
        ref_text: "在此奉劝大家别乱打美白针。",
@@ -104,7 +107,7 @@ jq -n \
     --arg a "$REF_A" \
     --arg b "$REF_B" \
     '{
-       model: "inclusionAI/Ming-omni-tts-16.8B-A3B",
+       model: "inclusionAI/Ming-omni-tts-0.5B",
        input: $input,
        ref_audio: [$a, $b],
        ref_text: $ref_text,

tests/dfx/conftest.py

@@ -653,7 +653,6 @@ def run_benchmark(
     random_input_len: Any | None = None,
     random_output_len: Any | None = None,
     resource_label: str | None = None,
-    num_warmups: int = 2,
 ) -> dict[str, Any]:
     """Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
 
@@ -678,7 +677,7 @@ def run_benchmark(
         + args
         + [
             "--num-warmups",
-            str(num_warmups),
+            "2",
             "--save-result",
             "--result-dir",
             os.environ.get("BENCHMARK_DIR", "tests"),

tests/dfx/perf/scripts/run_benchmark.py

@@ -279,6 +279,5 @@ def test_performance_benchmark(omni_server, benchmark_params):
             random_input_len=params.get("random_input_len"),
             random_output_len=params.get("random_output_len"),
             resource_label=resource_label,
-            num_warmups=max(2, int(concurrency)),
         )
         assert_result(result, params, num_prompt)

tests/diffusion/model_loader/test_fp8_cache.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for the Phase-I normalized online-FP8 cache contract."""
+
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import load_file, save_file
+from torch import nn
+
+from vllm_omni.diffusion.model_loader.fp8_cache import (
+    CACHE_COMPLETE_NAME,
+    CACHE_COMPONENT_SUBFOLDER,
+    CACHE_MANIFEST_NAME,
+    _source_fingerprint,
+    _transformer_quant_config,
+    build_normalized_fp8_cache,
+    is_online_fp8_config,
+    load_fp8_cache_manifest,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def _write_valid_cache(tmp_path):
+    source = tmp_path / "source.safetensors"
+    save_file({"weight": torch.ones(2, 2)}, str(source))
+
+    cache = tmp_path / "cache"
+    component = cache / CACHE_COMPONENT_SUBFOLDER
+    component.mkdir(parents=True)
+    shard = component / "part-00000.safetensors"
+    save_file({"weight": torch.ones(2, 2, dtype=torch.float8_e4m3fn)}, str(shard))
+    (component / "model.safetensors.index.json").write_text(
+        json.dumps(
+            {
+                "metadata": {"total_size": 4},
+                "weight_map": {"weight": shard.name},
+            }
+        )
+    )
+
+    manifest = {
+        "schema_version": 1,
+        "cache_kind": "normalized_fp8_runtime",
+        "component": "transformer",
+        "component_subfolder": CACHE_COMPONENT_SUBFOLDER,
+        "source_prefix": "transformer.",
+        "quantization": {"method": "fp8", "serialized": True, "activation_scheme": "dynamic"},
+        "tensor_count": 1,
+        "source_files": [str(source)],
+        "source_fingerprint": _source_fingerprint({str(source)}),
+    }
+    cache.mkdir(exist_ok=True)
+    (cache / CACHE_MANIFEST_NAME).write_text(json.dumps(manifest))
+    (cache / CACHE_COMPLETE_NAME).write_text("phase1\n")
+    return cache, source, manifest
+
+
+def test_online_and_serialized_fp8_configs_are_classified():
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    assert is_online_fp8_config(Fp8Config())
+    assert not is_online_fp8_config(Fp8Config(is_checkpoint_fp8_serialized=True))
+
+    cache_config = _transformer_quant_config()
+    assert not is_online_fp8_config(cache_config)
+    assert cache_config.resolve("transformer").is_checkpoint_fp8_serialized
+    assert cache_config.resolve("text_encoder") is None
+
+
+def test_load_fp8_cache_manifest_validates_artifact_and_sources(tmp_path):
+    cache, source, expected = _write_valid_cache(tmp_path)
+
+    assert load_fp8_cache_manifest(cache) == expected
+
+    source.write_bytes(source.read_bytes() + b"changed")
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_load_fp8_cache_manifest_rejects_wrong_serialization_contract(tmp_path):
+    cache, _, manifest = _write_valid_cache(tmp_path)
+
+    manifest["quantization"]["serialized"] = False
+    (cache / CACHE_MANIFEST_NAME).write_text(json.dumps(manifest))
+
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_load_fp8_cache_manifest_rejects_path_traversal(tmp_path):
+    cache, _, manifest = _write_valid_cache(tmp_path)
+    index_path = cache / CACHE_COMPONENT_SUBFOLDER / "model.safetensors.index.json"
+    index_path.write_text(json.dumps({"weight_map": {"weight": "../outside.safetensors"}}))
+
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_load_fp8_cache_manifest_rejects_non_object_json(tmp_path):
+    cache, _, _ = _write_valid_cache(tmp_path)
+    (cache / CACHE_MANIFEST_NAME).write_text("[]")
+
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_build_normalized_fp8_cache_serializes_weight_and_scale(tmp_path):
+    from vllm.model_executor.layers.quantization.online.fp8 import (
+        Fp8PerTensorOnlineLinearMethod,
+    )
+
+    source_dir = tmp_path / "source"
+    source_dir.mkdir()
+    source_file = source_dir / "model.safetensors"
+    save_file({"block.weight": torch.ones((3, 2), dtype=torch.bfloat16)}, str(source_file))
+
+    model = nn.Module()
+    model._dit_modules = ("transformer",)
+    model.transformer = nn.Module()
+    model.transformer.block = nn.Module()
+    model.transformer.block.weight = nn.Parameter(torch.ones((2, 3), dtype=torch.float8_e4m3fn), requires_grad=False)
+    model.transformer.block.weight_scale = nn.Parameter(torch.tensor([0.5]), requires_grad=False)
+    model.transformer.block.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod)
+
+    source = SimpleNamespace(
+        model_or_path=str(source_dir),
+        subfolder=None,
+        revision=None,
+        prefix="transformer.",
+    )
+    cache_dir = tmp_path / "cache"
+
+    manifest = build_normalized_fp8_cache(
+        model,
+        sources=(source,),
+        cache_dir=cache_dir,
+    )
+
+    assert manifest["cache_kind"] == "normalized_fp8_runtime"
+    tensors = load_file(str(cache_dir / CACHE_COMPONENT_SUBFOLDER / "part-00000.safetensors"))
+    scale_tensors = load_file(str(cache_dir / CACHE_COMPONENT_SUBFOLDER / "part-00001.safetensors"))
+    assert tensors["block.weight"].dtype == torch.float8_e4m3fn
+    assert tuple(tensors["block.weight"].shape) == (3, 2)
+    assert torch.equal(scale_tensors["block.weight_scale"], torch.tensor([0.5]))
+
+
+@pytest.mark.parametrize("missing", [CACHE_COMPLETE_NAME, CACHE_MANIFEST_NAME])
+def test_load_fp8_cache_manifest_requires_completion_marker(tmp_path, missing):
+    cache, _, _ = _write_valid_cache(tmp_path)
+    (cache / missing).unlink()
+
+    assert load_fp8_cache_manifest(cache) is None

tests/diffusion/models/sana_wm/test_cam_reduce_scatter.py

@@ -1,129 +0,0 @@
-# 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

@@ -1,111 +0,0 @@
-# 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

@@ -1,142 +0,0 @@
-# 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/e2e/offline_inference/test_sana_wm.py

@@ -1,159 +0,0 @@
-# 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

@@ -1,106 +0,0 @@
-# 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

@@ -1,810 +0,0 @@
-# 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,10 +22,6 @@ 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 (
@@ -98,20 +94,6 @@ 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,31 +2168,6 @@ 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

@@ -347,6 +347,8 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "--dlo-no-use-allgather",
             "--dlo-resident-layers",
             "20",
+            "--dlo-fp8-cache-dir",
+            "/tmp/dlo-fp8-cache",
         ]
     )
 
@@ -357,9 +359,11 @@ def test_serve_cli_forwards_distributed_offload_residency():
     assert args.enable_distributed_layerwise_offload is True
     assert args.dlo_use_allgather is False
     assert args.dlo_resident_layers == 20
+    assert args.dlo_fp8_cache_dir == "/tmp/dlo-fp8-cache"
     assert engine_args["enable_distributed_layerwise_offload"] is True
     assert engine_args["dlo_use_allgather"] is False
     assert engine_args["dlo_resident_layers"] == 20
+    assert engine_args["dlo_fp8_cache_dir"] == "/tmp/dlo-fp8-cache"
 
 
 def test_serve_cli_accepts_diffusion_compile_controls():

tests/model_executor/model_loader/test_resolve_model_path.py

@@ -1,92 +0,0 @@
-# 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,14 +7,8 @@ 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.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.common.ming.fm import Solver
 from vllm_omni.model_executor.models.ming_tts.constants import (
     AGGREGATOR_HIDDEN_SIZE,
     HISTORY_PATCH_SIZE,
@@ -30,6 +24,7 @@ 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
@@ -49,87 +44,7 @@ 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.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
+    assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.fm"
 
 
 def test_ming_tts_flowloss_preserves_checkpoint_prefix():

tests/model_executor/models/minicpmo_4_5/test_streaming_audio_cache.py

@@ -1,156 +0,0 @@
-# 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

@@ -1,41 +0,0 @@
-# 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

@@ -1,44 +0,0 @@
-# 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,11 +23,6 @@ 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,
@@ -46,74 +41,3 @@ 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,7 +80,6 @@ EXCLUDED_MODELS = [
     "Pi0Pipeline",
     "SoulXSingerPipeline",
     "SoulXSingerSVCPipeline",
-    "SanaWmPipeline",
 ]
 
 

vllm_omni/config/omni_config.py

@@ -693,6 +693,7 @@ class _DiffusionConfigProjection:
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = Field(default=0, ge=0)
+    dlo_fp8_cache_dir: str | None = None
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
     diffusion_compile_dynamic: bool = Field(default=True, strict=True)

vllm_omni/config/stage_config.py

@@ -383,6 +383,7 @@ class StageDeployConfig:
     enable_distributed_layerwise_offload: bool | None = None
     dlo_use_allgather: bool | None = None
     dlo_resident_layers: int | None = None
+    dlo_fp8_cache_dir: str | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None
 

vllm_omni/diffusion/data.py

@@ -752,6 +752,10 @@ class OmniDiffusionConfig:
     dlo_use_allgather: bool = True
     # Leading main-DiT blocks kept resident by distributed layerwise offload.
     dlo_resident_layers: int = 0
+    # Optional Phase-I normalized online-FP8 cache. A cache miss builds the
+    # artifact after the ordinary online loader; a valid cache hit uses the
+    # serialized FP8 checkpoint-mmap path.
+    dlo_fp8_cache_dir: str | None = None
 
     pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading
 

vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py

@@ -41,6 +41,8 @@ class _MiniMaxH3DirectMmapAdapter:
         target: torch.Tensor,
     ) -> DirectMmapTensorPolicy:
         del target
+        if getattr(self.pipeline, "_dlo_fp8_cache_normalized", False):
+            return DirectMmapTensorPolicy(allow_custom_loader=True)
         transform = None
         suffix = ".qkv_proj.weight"
         if runtime_name.endswith(suffix):

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -8,7 +8,7 @@ import re
 import time
 from collections.abc import Generator, Iterable, Sequence
 from pathlib import Path
-from typing import cast
+from typing import Any, cast
 
 import torch
 from torch import nn
@@ -34,6 +34,12 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     get_checkpoint_adapter,
 )
+from vllm_omni.diffusion.model_loader.fp8_cache import (
+    _transformer_quant_config,
+    build_normalized_fp8_cache,
+    is_online_fp8_config,
+    load_fp8_cache_manifest,
+)
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
     build_checkpoint_mmap_plan,
@@ -141,6 +147,10 @@ class DiffusersPipelineLoader:
         self.quant_config = od_config.quantization_config
         self.parallel_config = od_config.parallel_config
         self.host_weight_plan: HostWeightPlan | None = None
+        self._fp8_cache_dir: Path | None = None
+        self._fp8_cache_requested = False
+        self._fp8_cache_hit = False
+        self._fp8_cache_manifest: dict[str, Any] | None = None
 
     def take_host_weight_plan(self) -> HostWeightPlan | None:
         """Transfer the loader-produced plan to the offload backend."""
@@ -367,12 +377,61 @@ class DiffusersPipelineLoader:
         )
 
     def _get_weight_sources(self, model: nn.Module) -> tuple["ComponentSource", ...]:
-        return tuple(
+        sources = tuple(
             cast(
                 Iterable[DiffusersPipelineLoader.ComponentSource],
                 getattr(model, "weights_sources", ()),
             )
         )
+        if not self._fp8_cache_hit or self._fp8_cache_dir is None or self._fp8_cache_manifest is None:
+            return sources
+
+        cache_prefix = str(self._fp8_cache_manifest["source_prefix"])
+        cache_subfolder = str(self._fp8_cache_manifest.get("component_subfolder", "transformer"))
+        return tuple(
+            dataclasses.replace(
+                source,
+                model_or_path=str(self._fp8_cache_dir),
+                subfolder=cache_subfolder,
+                revision=None,
+            )
+            if source.prefix == cache_prefix
+            else source
+            for source in sources
+        )
+
+    def _prepare_fp8_cache(self) -> None:
+        """Select the Phase-I cache path before constructing the pipeline."""
+        cache_dir = getattr(self.od_config, "dlo_fp8_cache_dir", None)
+        if not cache_dir or not is_online_fp8_config(self.quant_config):
+            return
+        if not getattr(self.od_config, "enable_distributed_layerwise_offload", False):
+            logger.warning("Ignoring --dlo-fp8-cache-dir because distributed layerwise offload is disabled")
+            return
+        if not getattr(self.od_config, "dlo_use_allgather", True):
+            logger.warning("Ignoring --dlo-fp8-cache-dir because DLO AllGather is disabled")
+            return
+        if int(getattr(self.parallel_config, "tensor_parallel_size", 1)) != 1:
+            logger.warning("Ignoring --dlo-fp8-cache-dir because Phase-I cache requires TP=1")
+            return
+        if bool(getattr(self.parallel_config, "use_hsdp", False)):
+            logger.warning("Ignoring --dlo-fp8-cache-dir because Phase-I cache does not support HSDP")
+            return
+
+        self._fp8_cache_dir = Path(cache_dir)
+        self._fp8_cache_requested = True
+        self._fp8_cache_manifest = load_fp8_cache_manifest(self._fp8_cache_dir)
+        if self._fp8_cache_manifest is None:
+            logger.info("DLO normalized FP8 cache miss: %s", self._fp8_cache_dir)
+            return
+
+        self._fp8_cache_hit = True
+        self.quant_config = _transformer_quant_config()
+        self.od_config.quantization_config = self.quant_config
+        logger.info(
+            "DLO normalized FP8 cache hit: %s; selecting serialized FP8 checkpoint path",
+            self._fp8_cache_dir,
+        )
 
     def _get_expected_parameter_names(self, model: nn.Module) -> set[str]:
         """Return parameter names that should be covered by strict load checks."""
@@ -403,6 +462,7 @@ class DiffusersPipelineLoader:
         self.host_weight_plan = None
         if load_format is None:
             load_format = "default"
+        self._prepare_fp8_cache()
         # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8),
         # weights are already quantized in the checkpoint — load directly on CPU.
         # For online quantization, load on device so quantization can run on accelerator,
@@ -413,7 +473,9 @@ class DiffusersPipelineLoader:
             is_offline = getattr(quant_cfg, "data_type", None) == "mx_fp" or getattr(
                 quant_cfg, "is_checkpoint_quantized", False
             )
-            if not is_offline:
+            if self._fp8_cache_hit:
+                logger.info("Normalized FP8 cache is already serialized; keeping model initialization on CPU")
+            elif not is_offline:
                 load_device = device.type
                 offload_after_quant = True
                 logger.info(
@@ -431,6 +493,11 @@ class DiffusersPipelineLoader:
                 )
             else:
                 model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False)
+                if self._fp8_cache_hit:
+                    # The direct-mmap adapter must see this before plan
+                    # construction so normalized QKV tensors bypass the raw
+                    # checkpoint grouped-QKV transform.
+                    setattr(model, "_dlo_fp8_cache_normalized", True)
 
                 _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
                 _use_ag = getattr(self.od_config, "dlo_use_allgather", True)
@@ -520,6 +587,20 @@ class DiffusersPipelineLoader:
                 model.to("cpu")
                 logger.info("Quantization complete, offloaded model back to CPU")
 
+        if self._fp8_cache_requested and self._fp8_cache_dir is not None and not self._fp8_cache_hit:
+            try:
+                build_normalized_fp8_cache(
+                    model,
+                    sources=self._get_weight_sources(model),
+                    cache_dir=self._fp8_cache_dir,
+                )
+            except Exception as exc:
+                logger.warning(
+                    "DLO normalized FP8 cache build failed; retaining ordinary-loader path: %s",
+                    exc,
+                    exc_info=True,
+                )
+
         self._apply_skip_softmax_calibration(model)
         return model.eval()
 

vllm_omni/diffusion/model_loader/fp8_cache.py

@@ -0,0 +1,349 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Phase-I normalized FP8 cache support for DLO checkpoint mmap."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import shutil
+import uuid
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.model_loader.host_weight_plan import (
+    _build_source_map,
+    _collect_required_dit_tensors,
+)
+from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+
+logger = init_logger(__name__)
+
+CACHE_SCHEMA_VERSION = 1
+CACHE_KIND = "normalized_fp8_runtime"
+CACHE_MANIFEST_NAME = "manifest.json"
+CACHE_COMPLETE_NAME = "COMPLETE"
+CACHE_COMPONENT_SUBFOLDER = "transformer"
+
+
+def _transformer_quant_config() -> ComponentQuantizationConfig:
+    """Build the serialized FP8 config used by a normalized cache."""
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    return ComponentQuantizationConfig(
+        {
+            "transformer": Fp8Config(
+                is_checkpoint_fp8_serialized=True,
+                activation_scheme="dynamic",
+            )
+        }
+    )
+
+
+def is_online_fp8_config(quant_config: Any) -> bool:
+    """Return whether the transformer component requests online FP8."""
+    if quant_config is None:
+        return False
+    if hasattr(quant_config, "resolve"):
+        quant_config = quant_config.resolve("transformer")
+    get_name = getattr(quant_config, "get_name", None)
+    if not callable(get_name) or str(get_name()).lower() != "fp8":
+        return False
+    return not bool(getattr(quant_config, "is_checkpoint_fp8_serialized", False))
+
+
+def _source_fingerprint(files: set[str]) -> str:
+    entries: list[dict[str, int | str]] = []
+    for filename in sorted(files):
+        stat = os.stat(filename)
+        entries.append(
+            {
+                "path": os.path.abspath(filename),
+                "size": stat.st_size,
+                "mtime_ns": stat.st_mtime_ns,
+            }
+        )
+    payload = json.dumps(entries, sort_keys=True, separators=(",", ":")).encode()
+    return hashlib.sha256(payload).hexdigest()
+
+
+def load_fp8_cache_manifest(cache_dir: str | os.PathLike[str] | None) -> dict[str, Any] | None:
+    """Return a validated Phase-I manifest, or ``None`` for a cache miss."""
+    if cache_dir is None:
+        return None
+    root = Path(cache_dir)
+    manifest_path = root / CACHE_MANIFEST_NAME
+    if not (root / CACHE_COMPLETE_NAME).is_file() or not manifest_path.is_file():
+        return None
+    try:
+        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+    if not isinstance(manifest, dict):
+        return None
+    if manifest.get("schema_version") != CACHE_SCHEMA_VERSION:
+        return None
+    if manifest.get("cache_kind") != CACHE_KIND:
+        return None
+    if manifest.get("component") != CACHE_COMPONENT_SUBFOLDER:
+        return None
+    quantization = manifest.get("quantization")
+    if not isinstance(quantization, dict) or quantization != {
+        "method": "fp8",
+        "serialized": True,
+        "activation_scheme": "dynamic",
+    }:
+        return None
+    if manifest.get("component_subfolder") != CACHE_COMPONENT_SUBFOLDER:
+        return None
+    if manifest.get("source_prefix") != f"{CACHE_COMPONENT_SUBFOLDER}.":
+        return None
+
+    index_path = root / CACHE_COMPONENT_SUBFOLDER / "model.safetensors.index.json"
+    if not index_path.is_file():
+        return None
+    try:
+        index = json.loads(index_path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+    if not isinstance(index, dict):
+        return None
+    weight_map = index.get("weight_map")
+    if not isinstance(weight_map, dict) or not weight_map or not all(isinstance(key, str) for key in weight_map):
+        return None
+    if manifest.get("tensor_count") != len(weight_map):
+        return None
+    for filename in weight_map.values():
+        if not isinstance(filename, str):
+            return None
+        relative_path = Path(filename)
+        if (
+            relative_path.is_absolute()
+            or relative_path.name != filename
+            or relative_path.suffix != ".safetensors"
+            or not (root / CACHE_COMPONENT_SUBFOLDER / relative_path).is_file()
+        ):
+            return None
+
+    source_files = manifest.get("source_files")
+    source_fingerprint = manifest.get("source_fingerprint")
+    if not isinstance(source_files, list) or not source_files or not isinstance(source_fingerprint, str):
+        return None
+    try:
+        if _source_fingerprint({str(filename) for filename in source_files}) != source_fingerprint:
+            return None
+    except (OSError, TypeError):
+        return None
+    return manifest
+
+
+def _cache_key_for_runtime_name(runtime_name: str, prefix: str) -> str:
+    if not runtime_name.startswith(prefix):
+        raise ValueError(f"runtime tensor {runtime_name!r} is outside cache source prefix {prefix!r}")
+    return runtime_name[len(prefix) :]
+
+
+def build_normalized_fp8_cache(
+    model: nn.Module,
+    *,
+    sources: tuple[object, ...],
+    cache_dir: str | os.PathLike[str],
+) -> dict[str, Any]:
+    """Write a normalized FP8 cache from a completed online-FP8 model.
+
+    The builder writes tensors one at a time into individual safetensors shards,
+    so it does not create a second full-model state dict. The first cache build
+    still follows the existing online loader; warm launches consume the cache
+    through the direct DLO mmap path.
+    """
+    cache_root = Path(cache_dir)
+    existing_manifest = load_fp8_cache_manifest(cache_root)
+    if existing_manifest is not None:
+        return existing_manifest
+    if cache_root.exists() and any(cache_root.iterdir()):
+        raise FileExistsError(f"FP8 cache directory is not empty and is not valid: {cache_root}")
+    cache_root.parent.mkdir(parents=True, exist_ok=True)
+
+    dit_names = tuple(getattr(model, "_dit_modules", ()))
+    dit_names = tuple(name for name in dit_names if hasattr(model, name))
+    if len(dit_names) != 1:
+        raise ValueError(f"Phase-I FP8 cache requires exactly one DiT module, got {dit_names}")
+    dit_name = dit_names[0]
+    if dit_name != CACHE_COMPONENT_SUBFOLDER:
+        raise ValueError(
+            f"Phase-I FP8 cache requires a top-level {CACHE_COMPONENT_SUBFOLDER!r} module, got {dit_name!r}"
+        )
+    source_prefix = f"{dit_name}."
+    dit_module = model.get_submodule(dit_name)
+    planned_sources = tuple(source for source in sources if getattr(source, "prefix", "") == source_prefix)
+    if len(planned_sources) != 1:
+        raise ValueError(f"Phase-I FP8 cache requires one source for {source_prefix!r}, got {len(planned_sources)}")
+
+    source_map = _build_source_map(
+        sources=planned_sources,
+        model_path=None,
+        remap_fn=getattr(type(model), "_remap_ckpt_key", None),
+    )
+    required = _collect_required_dit_tensors(((dit_name, dit_module),))
+    quantized_modules: dict[str, nn.Module] = {}
+    try:
+        from vllm.model_executor.layers.quantization.online.fp8 import Fp8PerTensorOnlineLinearMethod
+    except ImportError as exc:  # pragma: no cover - vLLM 0.27 provides this class
+        raise RuntimeError("Phase-I FP8 cache requires vLLM's online FP8 implementation") from exc
+
+    for module_name, module in model.named_modules():
+        if not module_name.startswith(source_prefix):
+            continue
+        if isinstance(getattr(module, "quant_method", None), Fp8PerTensorOnlineLinearMethod):
+            if not hasattr(module, "weight_scale"):
+                raise ValueError(f"online FP8 module {module_name!r} has no generated weight_scale")
+            quantized_modules[module_name] = module
+
+    tmp_root = cache_root.parent / f".{cache_root.name}.tmp-{uuid.uuid4().hex}"
+    component_root = tmp_root / CACHE_COMPONENT_SUBFOLDER
+    component_root.mkdir(parents=True, exist_ok=False)
+    weight_map: dict[str, str] = {}
+    source_files: set[str] = set()
+    total_size = 0
+    part_index = 0
+
+    file_cache: dict[str, Any] = {}
+
+    def write_tensor(key: str, tensor: torch.Tensor) -> None:
+        nonlocal part_index, total_size
+        if key in weight_map:
+            raise ValueError(f"normalized FP8 cache has duplicate tensor key {key!r}")
+        tensor = tensor.detach().to("cpu").contiguous()
+        filename = f"part-{part_index:05d}.safetensors"
+        part_index += 1
+        save_file({key: tensor}, str(component_root / filename))
+        weight_map[key] = filename
+        total_size += tensor.numel() * tensor.element_size()
+
+    try:
+        for runtime_name, target in required.items():
+            binding_info = source_map.get(runtime_name)
+            module_name = runtime_name.removesuffix(".weight")
+            quant_module = quantized_modules.get(module_name) if runtime_name.endswith(".weight") else None
+            if runtime_name.endswith(".weight_scale"):
+                owner_name = runtime_name.removesuffix(".weight_scale")
+                if owner_name in quantized_modules:
+                    # Generated online scales are published in the dedicated
+                    # pass below; they have no raw checkpoint binding.
+                    continue
+            if quant_module is not None:
+                if binding_info is None:
+                    raise ValueError(f"no source binding for quantized weight {runtime_name!r}")
+                qweight = quant_module.weight.detach()
+                if qweight.ndim != 2:
+                    raise ValueError(f"Phase-I FP8 cache only supports 2D weights, got {runtime_name!r}")
+                # Online FP8 stores the final runtime weight as qweight.T. The
+                # serialized FP8 method expects the pre-process qweight shape.
+                qweight = qweight.t().contiguous()
+                # Normalized caches contain the post-checkpoint-adapter QKV
+                # order. The cache-hit adapter bypasses the raw-checkpoint
+                # grouped-QKV transform, then the serialized FP8 finalizer
+                # performs only the standard weight transpose.
+                write_tensor(binding_info[0], qweight)
+                source_files.add(binding_info[1])
+                continue
+
+            if binding_info is not None:
+                checkpoint_key, file_path = binding_info
+                source_files.add(file_path)
+                handle = file_cache.get(file_path)
+                if handle is None:
+                    handle = safe_open(file_path, framework="pt", device="cpu")
+                    file_cache[file_path] = handle
+                write_tensor(checkpoint_key, handle.get_tensor(checkpoint_key))
+            else:
+                # Constructor-derived persistent buffers have no checkpoint
+                # entry; preserve them as canonical cache tensors.
+                write_tensor(_cache_key_for_runtime_name(runtime_name, source_prefix), target)
+
+        for module_name, module in quantized_modules.items():
+            scale = module.weight_scale.detach()
+            logical_widths = getattr(module, "logical_widths", None)
+            if scale.numel() == 1 and logical_widths is not None and len(logical_widths) > 1:
+                # Serialized Fp8LinearMethod represents a fused per-tensor
+                # weight with one scale slot per logical shard. Online FP8
+                # produces one scalar for the fused tensor, so replicate that
+                # scalar to the offline runtime shape without changing its
+                # value.
+                scale = scale.reshape(1).expand(len(logical_widths)).contiguous()
+            write_tensor(
+                _cache_key_for_runtime_name(f"{module_name}.weight_scale", source_prefix),
+                scale,
+            )
+
+        index = {
+            "metadata": {"total_size": total_size},
+            "weight_map": dict(sorted(weight_map.items())),
+        }
+        (component_root / "model.safetensors.index.json").write_text(
+            json.dumps(index, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+        manifest = {
+            "schema_version": CACHE_SCHEMA_VERSION,
+            "cache_kind": CACHE_KIND,
+            "component": "transformer",
+            "component_subfolder": CACHE_COMPONENT_SUBFOLDER,
+            "source_prefix": source_prefix,
+            "quantization": {
+                "method": "fp8",
+                "serialized": True,
+                "activation_scheme": "dynamic",
+            },
+            "tensor_count": len(weight_map),
+            "source_fingerprint": _source_fingerprint(source_files),
+            "source_files": sorted(source_files),
+        }
+        (tmp_root / CACHE_MANIFEST_NAME).write_text(
+            json.dumps(manifest, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+        (tmp_root / CACHE_COMPLETE_NAME).write_text("phase1\n", encoding="utf-8")
+        for handle in file_cache.values():
+            close = getattr(handle, "__exit__", None)
+            if callable(close):
+                close(None, None, None)
+        try:
+            os.replace(tmp_root, cache_root)
+        except FileExistsError:
+            # Multiple DLO workers may build the same cache concurrently. If
+            # another worker published a complete artifact first, treat this
+            # publisher as a successful cache race rather than a build error.
+            published = load_fp8_cache_manifest(cache_root)
+            if published is None:
+                raise
+            shutil.rmtree(tmp_root, ignore_errors=True)
+            logger.info("Normalized FP8 cache was published concurrently at %s", cache_root)
+            return published
+        logger.info("Published normalized FP8 cache at %s (%d tensors)", cache_root, len(weight_map))
+        return manifest
+    except BaseException:
+        for handle in file_cache.values():
+            close = getattr(handle, "__exit__", None)
+            if callable(close):
+                close(None, None, None)
+        shutil.rmtree(tmp_root, ignore_errors=True)
+        raise
+
+
+__all__ = [
+    "CACHE_SCHEMA_VERSION",
+    "CACHE_KIND",
+    "CACHE_COMPONENT_SUBFOLDER",
+    "build_normalized_fp8_cache",
+    "is_online_fp8_config",
+    "load_fp8_cache_manifest",
+    "_transformer_quant_config",
+]

vllm_omni/diffusion/model_metadata.py

@@ -56,10 +56,6 @@ _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/sana_wm/__init__.py

@@ -1,25 +0,0 @@
-# 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

@@ -1,356 +0,0 @@
-# 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

@@ -1,93 +0,0 @@
-# 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

@@ -1,942 +0,0 @@
-# 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

@@ -1,269 +0,0 @@
-# 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

@@ -1,536 +0,0 @@
-# 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/distributed_layerwise_backend.py

@@ -1034,6 +1034,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         logger.info("Realized %d loader-planned tensors as mmap views", len(loaded_names))
 
+        if getattr(pipeline, "_dlo_fp8_cache_normalized", False):
+            self._process_normalized_fp8_cache(pipeline)
+
         # Keep file handles open — _shard_and_pin will read from the mmap views.
         # They will be released after _shard_and_pin completes (when params are
         # replaced with offload placeholders).
@@ -1103,6 +1106,52 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 f"(first 5: {remaining_meta[:5]})."
             )
 
+    @staticmethod
+    def _process_normalized_fp8_cache(pipeline: nn.Module) -> None:
+        """Finalize serialized FP8 parameters after mmap replacement.
+
+        Phase-I cache files use the serialized FP8 checkpoint layout: the
+        weight is stored before the quantization method's final transpose, and
+        ``weight_scale`` is stored as an explicit cache tensor.  The ordinary
+        loader normally runs this method before DLO owns the plan; mmap loading
+        must run the same finalizer after replacing meta parameters.
+        """
+        from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
+        from vllm.model_executor.utils import replace_parameter
+
+        for module in pipeline.modules():
+            quant_method = getattr(module, "quant_method", None)
+            if not isinstance(quant_method, Fp8LinearMethod):
+                continue
+            if getattr(module, "_already_called_process_weights_after_loading", False):
+                continue
+
+            # ``Fp8LinearMethod.process_weights_after_loading`` also handles
+            # non-serialized checkpoints.  Its common path calls
+            # ``process_fp8_weight_tensor_strategy`` to requantize fused
+            # tensors, which is a CUDA-only operation.  A normalized cache
+            # already contains the serialized FP8 values and the generated
+            # scales, so only perform the layout transition and the kernel's
+            # own preprocessing here.  This keeps the warm path CPU/mmap
+            # compatible and preserves the cache bytes exactly.
+            if quant_method.block_quant:
+                raise ValueError("Phase-I normalized FP8 cache does not support block-quantized weights")
+            replace_parameter(module, "weight", module.weight.t().data)
+            # The online quantizer records one scalar.  A serialized FP8
+            # checkpoint may expose one repeated slot per logical shard in a
+            # fused module, but those slots must agree before they can be
+            # collapsed to the canonical per-tensor ``[1, 1]`` form.
+            weight_scale = module.weight_scale.reshape(-1)
+            if weight_scale.numel() > 1 and not torch.all(weight_scale == weight_scale[0]):
+                raise ValueError("Phase-I normalized FP8 cache contains non-uniform fused weight scales")
+            weight_scale = weight_scale[:1].reshape(1, 1)
+            replace_parameter(module, "weight_scale", weight_scale.data)
+            module.input_scale = None
+            if quant_method.use_marlin and hasattr(quant_method.fp8_linear, "marlin_input_dtype"):
+                quant_method.fp8_linear.marlin_input_dtype = quant_method.marlin_input_dtype
+            quant_method.fp8_linear.process_weights_after_loading(module)
+            module._already_called_process_weights_after_loading = True
+
     def _init_dp_group(self) -> None:
         """Reuse the process group initialized by parallel_state.
 

vllm_omni/diffusion/registry.py

@@ -171,11 +171,6 @@ _DIFFUSION_MODELS = {
         "pipeline_ming_imagegen",
         "MingImagePipeline",
     ),
-    "SanaWmPipeline": (
-        "sana_wm",
-        "pipeline_sana_wm",
-        "SanaWmPipeline",
-    ),
     "InternVLAA1Pipeline": (
         "internvla_a1",
         "pipeline_internvla_a1",
@@ -638,7 +633,6 @@ _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/engine/arg_utils.py

@@ -546,6 +546,7 @@ class OrchestratorArgs:
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = 0
+    dlo_fp8_cache_dir: str | None = None
     boundary_ratio: float | None = None
     flow_shift: float | None = None
     diffusion_kv_cache_dtype: str | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1046,6 +1046,7 @@ class AsyncOmniEngine:
             "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
             "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
             "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
+            "dlo_fp8_cache_dir": kwargs.get("dlo_fp8_cache_dir", None),
             "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
             "diffusion_compile_granularity": (
                 "regional"

vllm_omni/entrypoints/cli/serve.py

@@ -729,6 +729,14 @@ class OmniServeCommand(CLISubcommand):
             help="Keep this many leading main-DiT blocks resident on the device "
             "while distributed layerwise offload streams the remaining blocks.",
         )
+        omni_config_group.add_argument(
+            "--dlo-fp8-cache-dir",
+            type=str,
+            default=None,
+            help="Optional Phase-I normalized online-FP8 cache directory. "
+            "A cache miss builds the cache after the ordinary loader; a valid "
+            "cache hit uses serialized FP8 checkpoint mmap with DLO.",
+        )
         # Video model parameters (e.g., Wan2.2) - engine-level
         omni_config_group.add_argument(
             "--boundary-ratio",

vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py

@@ -10,12 +10,7 @@ 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,
-    apply_max_new_tokens,
-    conditioning_cache_salt,
-)
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, conditioning_cache_salt
 
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -109,15 +104,6 @@ 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,4 +1,3 @@
-import os
 import time
 from pathlib import Path
 
@@ -86,46 +85,3 @@ 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,13 +1,2 @@
 # 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/dit.py

@@ -216,3 +216,19 @@ 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/fm.py

@@ -3,21 +3,7 @@
 
 import torch
 
-
-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)
+from .dit import get_epss_timesteps
 
 
 def apply_sway_sampling(t, sway_sampling_coef):

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

@@ -1,56 +0,0 @@
-# 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,15 +25,14 @@ 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 resolve_model_to_local_path
-from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 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, DiT, MingAudioGenerator, build_tts_input
+from .talker_module import CFM, Aggregator, DiT, MingAudioGenerator, build_tts_input
 from .text_processing import segment_and_normalize
 from .voice_presets import VoicePresetRegistry
 
@@ -120,11 +119,7 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin)
             DiT(llm_input_dim=self.hidden_size, **config.flowmodel),
             steps=config.steps,
         )
-        # 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.aggregator = Aggregator(llm_input_dim=self.hidden_size, **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)
 
@@ -530,22 +525,27 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin)
 
     def _iter_talker_safetensors(self) -> Iterable[tuple[str, torch.Tensor]]:
         """Yield (name, tensor) pairs from talker/model*.safetensors."""
-        model_root = resolve_model_to_local_path(
-            self._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):
+        model_path = self._model_path
+        # Try local path first
+        for candidate in (os.path.join(model_path, "talker"), model_path):
             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
 
-        raise RuntimeError(f"No talker safetensors found under {model_root}. Expected talker/model*.safetensors.")
+        # HF hub fallback: download only the talker checkpoint files
+        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,7 +64,6 @@ 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 (
@@ -716,8 +715,12 @@ 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 = resolve_model_to_local_path(model_path, allow_download=True, allow_patterns=["mlp/*"])
+                local_path = download_weights_from_hf_specific(model_path, None, ["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,14 +33,9 @@ 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.cfm_solver import (
-    apply_sway_sampling,
-    get_epss_timesteps,
-    integrate_cfm_steps,
-)
-from vllm_omni.model_executor.models.common.ming.dit_blocks import CondEmbedder, DiTBlock, FinalLayer
+from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer, get_epss_timesteps
+from vllm_omni.model_executor.models.common.ming.fm import apply_sway_sampling, integrate_cfm_steps
 
 logger = init_logger(__name__)
 
@@ -463,6 +458,58 @@ 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,15 +16,16 @@ 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 resolve_model_to_local_path
-from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 
+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__)
 
 
@@ -272,19 +273,18 @@ class VoicePresetRegistry:
             if os.path.isfile(path):
                 return path, candidate
 
-        try:
-            hf_root = resolve_model_to_local_path(
-                self._model_path,
-                allow_download=True,
-                allow_patterns=["talker/data/**"],
-                cache_dir=self._download_dir,
-            )
-        except Exception as e:
-            logger.info("Could not download voice presets from HF: %s", e)
-            return None, None
-
-        candidate = os.path.join(hf_root, "talker", "data", "voice_name.json")
-        if os.path.isfile(candidate):
-            return candidate, os.path.join(hf_root, "talker")
+        if not os.path.isdir(self._model_path):
+            try:
+                hf_root = download_weights_from_hf_specific(
+                    self._model_path,
+                    self._download_dir,
+                    allow_patterns=["talker/data/**"],
+                    require_all=True,
+                )
+                candidate = os.path.join(hf_root, "talker", "data", "voice_name.json")
+                if os.path.isfile(candidate):
+                    return candidate, os.path.join(hf_root, "talker")
+            except Exception as e:  # pragma: no cover
+                logger.info("Could not download voice presets from HF: %s", e)
 
         return None, None

vllm_omni/model_executor/models/ming_tts/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_blocks import DiTBlock, FinalLayer
+from vllm_omni.model_executor.models.common.ming.dit import DiTBlock, FinalLayer
 
 
 class Aggregator(nn.Module):

vllm_omni/model_executor/models/ming_tts/flowloss_head.py

@@ -5,8 +5,9 @@
 import torch
 import torch.nn as nn
 
-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 vllm_omni.model_executor.models.common.ming.fm import Solver, build_timesteps
+
+from .fm.dit import DiT
 
 
 class FlowLoss(nn.Module):

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

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

vllm_omni/model_executor/models/ming_tts/fm/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.cfm_solver import build_timesteps
+from vllm_omni.model_executor.models.common.ming.fm import build_timesteps
 
 
 class CFMSampler(nn.Module):

vllm_omni/model_executor/models/ming_tts/fm/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_blocks import CondEmbedder, DiTBlock, FinalLayer
+from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer
 
 
 class DiT(nn.Module):

vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py

@@ -11,21 +11,15 @@ 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,
@@ -40,6 +34,7 @@ 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,
@@ -87,6 +82,8 @@ 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()
@@ -338,6 +335,8 @@ 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.model_executor.model_loader.weight_utils import resolve_model_to_local_path
+from vllm_omni.engine.stage_init_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,13 +3,16 @@
 # 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.common.ming.speaker_extractor import SpeakerEmbeddingExtractor
+from vllm_omni.model_executor.models.ming_flash_omni.spk_embedding import SpkembExtractor
 
 from .audio_prep import coerce_prompt_waveform, coerce_speaker_embeddings
 from .constants import (
@@ -26,9 +29,48 @@ logger = init_logger(__name__)
 _EXTRACTOR_LOAD_LOCK = threading.Lock()
 
 
-# Ming's bundle ships the CAM++ artifact as `campplus.onnx` in the model
-# directory, which is the shared extractor's default.
-def _load_speaker_extractor(wrapper: MingTTSForConditionalGeneration) -> SpeakerEmbeddingExtractor:
+def resolve_model_to_local_path(model):
+    if os.path.isdir(model):
+        return model
+
+    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
@@ -37,7 +79,7 @@ def _load_speaker_extractor(wrapper: MingTTSForConditionalGeneration) -> Speaker
         extractor = getattr(wrapper, "_speaker_extractor", None)
         if extractor is None:
             model_path = wrapper.vllm_config.model_config.model
-            extractor = SpeakerEmbeddingExtractor(model_path, target_sr=16000)
+            extractor = MingSpeakerEmbeddingExtractor(model_path, target_sr=16000)
             wrapper._speaker_extractor = extractor
     return extractor
 
@@ -132,6 +174,8 @@ 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,25 +58,6 @@ 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.
 
@@ -101,7 +82,6 @@ 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,7 +17,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import inspect
 import math
 import os
 import warnings
@@ -2506,17 +2505,6 @@ 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):
@@ -2529,10 +2517,6 @@ 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]
@@ -2570,18 +2554,16 @@ class MiniCPMWhisperEncoderLayer(nn.Module):
         """
         residual = hidden_states
         hidden_states = self.self_attn_layer_norm(hidden_states)
-        attention_kwargs = {
-            "hidden_states": hidden_states,
-            "attention_mask": attention_mask,
-            "layer_head_mask": layer_head_mask,
-            "output_attentions": output_attentions,
-            self._past_key_values_kwarg: past_key_values,
-        }
-        attn_out = self.self_attn(**attention_kwargs)
+        attn_out = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            layer_head_mask=layer_head_mask,
+            output_attentions=output_attentions,
+            past_key_value=past_key_values,
+        )
         hidden_states = attn_out[0]
         attn_weights = attn_out[1] if len(attn_out) > 1 else None
-        if len(attn_out) > 2:
-            past_key_values = attn_out[2]
+        past_key_values = attn_out[2] if len(attn_out) > 2 else None
         hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
         hidden_states = residual + hidden_states
 
@@ -2787,7 +2769,7 @@ class MiniCPMWhisperEncoder(WhisperEncoder):
                 past_key_values = EncoderDecoderCache(past_key_values, DynamicCache())
             else:
                 pass
-            past_key_values_length = _get_audio_cache_length(past_key_values)
+            past_key_values_length = past_key_values.self_attention_cache.get_usable_length(inputs_embeds.shape[1])
             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:, :]
@@ -4486,7 +4468,7 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
             return []
 
         if self.audio_past_key_values is not None:
-            cache_length = _get_audio_cache_length(self.audio_past_key_values)
+            cache_length = self.audio_past_key_values[0][0].shape[2]
             apm_max_len = self.apm.embed_positions.weight.shape[0]
             if cache_length + current_seq_len >= apm_max_len:
                 logger.warning(
@@ -4498,7 +4480,7 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
 
         past_len = 0
         if self.audio_past_key_values is not None:
-            past_len = _get_audio_cache_length(self.audio_past_key_values)
+            past_len = self.audio_past_key_values[0][0].shape[2]
         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,45 +128,32 @@ def _sample_token(
     do_sample: bool,
     generator: torch.Generator | None = None,
 ) -> torch.Tensor:
-    """Top-k + top-p sampling for the upstream 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.
+    """Top-k + top-p sampling (matches upstream's ``sample_token`` for the
+    inference branch).
     """
     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, top_indices = torch.topk(logits, top_k, dim=-1, sorted=True)
-        # topk is already descending: keep the following nucleus filter,
-        # softmax, and multinomial at width k instead of the full vocab.
-        logits = top_vals
-        compact_indices = top_indices
+        top_vals, _ = torch.topk(logits, top_k, dim=-1)
+        thresh = top_vals[..., -1:].expand_as(logits)
+        logits = torch.where(logits < thresh, torch.full_like(logits, float("-inf")), logits)
 
     if 0.0 < top_p < 1.0:
-        sorted_indices = None
-        if compact_indices is None:
-            logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
-        probs = F.softmax(logits, dim=-1)
+        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
+        probs = F.softmax(sorted_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
-        logits = logits.masked_fill(drop, float("-inf"))
-        if sorted_indices is not None:
-            logits = torch.full_like(logits, float("-inf")).scatter_(-1, sorted_indices, logits)
+        sorted_logits = sorted_logits.masked_fill(drop, float("-inf"))
+        logits = torch.full_like(logits, float("-inf")).scatter_(-1, sorted_idx, sorted_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,48 +52,15 @@ 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 prepare_rope_cache(
-        self,
-        max_seq_len: int,
-        device: torch.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)
-        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,
+        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],
-        )
+        position_ids = torch.arange(seq_len, device=device, dtype=torch.float32)
+        freqs = torch.einsum("s,d->sd", position_ids, self.inv_freq.to(device=device))
+        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)
 
     @staticmethod
     def _rotate_half(x: torch.Tensor) -> torch.Tensor:
@@ -233,15 +200,9 @@ 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,11 +1323,7 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
         self.mtp_hidden_size = hidden_size
         self.talker_mtp_graph_safe = True
         self.talker_mtp_output_key = ("audio_codes", "current")
-        # ``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.use_async_omni_output = False
         self.eager_omni_postprocess_before_async_output = True
         self.omni_pooler_payload_include_hidden = False
         self.postprocess_uses_multimodal_outputs = False
@@ -1470,7 +1466,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}`` state.
+        Prefill: initialise the per-request ``{is_stopping, step}`` state.
         Decode: combine the forced text token (from ``compute_logits``) with
         the audio codes the local transformer produced last step.
         """
@@ -1478,9 +1474,8 @@ 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 or is_prefill:
+        if span_len > 1 or is_first_call:
             embeds = self.model.embed_tokens(input_ids)
 
             ref_codes = (info_dict.get("codes", {}) or {}).get("ref")
@@ -1673,14 +1668,9 @@ 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
-                # Do not boolean-index a CUDA tensor here.  Even for the
-                # common single-row case, ``codes[should_continue]`` must
-                # 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
+                emitted_codes = codes[should_continue]
+                per_req_codes.append(emitted_codes)
+                have_codes = have_codes or emitted_codes.numel() > 0
             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))