[Bugfix][MiniMax-H3] Apply pre-commit formatting

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

.buildkite/amd/test-amd-merge.yml

@@ -238,13 +238,10 @@ steps:
       grade: Blocking
       timeout_in_minutes: 50
       commands:
-        - |
+        - export VLLM_LOGGING_LEVEL=DEBUG
-          timeout 40m bash -c '
+        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-            export VLLM_LOGGING_LEVEL=DEBUG
+        - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-            export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        - timeout 40m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model and cuda" --run-level "advanced_model"
-            export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-            pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model'
-          '
 
 # TODO: Bagel test on ROCm is very unstable. @tjtanaa
 # Need to debug before reneable numerical changes across large PRs

.buildkite/cuda/test-merge.yml

@@ -237,3 +237,38 @@ steps:
               pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
             "
         mirror_hardwares: h100_1
+
+      - label: "Diffusion · MiniMax-H3 L3 DLO DP2 and FastH3"
+        source_file_dependencies:
+          - tests/e2e/online_serving/minimax_h3/
+          - tests/helpers/assertions.py
+          - tests/helpers/mark.py
+          - tests/helpers/media.py
+          - tests/helpers/runtime.py
+          - vllm_omni/diffusion/models/minimax_h3/
+          - vllm_omni/diffusion/lora/
+          - vllm_omni/diffusion/offloader/
+          - vllm_omni/diffusion/distributed/
+          - vllm_omni/entrypoints/openai/
+        # Keep four server lifecycles in one Buildkite job, but never in one
+        # pytest process. DLO uses multiprocessing shared-memory queues; the
+        # parent's resource_tracker can retain a stale /psm_* entry during
+        # teardown and strand the next Ref2VA collective.
+        # A per-file process plus a hard timeout also makes a regression fail
+        # instead of leaving the merge queue stuck indefinitely.
+        timeout_in_minutes: 150
+        commands:
+          - |
+            set -eu
+            export VLLM_IMAGE_FETCH_TIMEOUT=60
+            if [ -f vllm_omni/diffusion/models/minimax_h3/lora.py ]; then
+              export VLLM_TEST_MINIMAX_H3_TURBO_LORA=$(python3 -c 'from huggingface_hub import hf_hub_download; print(hf_hub_download(repo_id="lightx2v/Minimax-h3-Turbo", filename="minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"))')
+            fi
+            if [ -f vllm_omni/diffusion/models/minimax_h3/fasth3.py ]; then
+              export VLLM_TEST_MINIMAX_H3_FASTH3_LORA=$(python3 -c 'from huggingface_hub import hf_hub_download; print(hf_hub_download(repo_id="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA", filename="dense-datafree/adapter_model.safetensors"))')
+            fi
+            timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_fl2va.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+            timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_ref2va.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+            timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_turbo_lora.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+            timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_fasth3_lora.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+        mirror_hardwares: h100_2

README.md

@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 *Latest News* 🔥
 
+- [2026/08] We released [0.28.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.28.0), featuring production-ready [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) serving on GPU and NPU, a unified AR/DiT paged KV cache runtime, and enhanced realtime full-duplex serving for the [MiniCPM-o series](recipes/OpenBMB/MiniCPM-o-4_5.md).
 - [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0).
 - [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
 - [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
@@ -57,9 +58,9 @@ vLLM-Omni is flexible and easy to use with:
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
 - **Omni-modality models** (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
-- **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
+- **TTS models** (e.g. Qwen3-TTS, IndexTTS 2.5, dots.tts, CosyVoice3)
-- **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, Qwen-Image, Wan2.2, FLUX)
+- **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2)
-- **Robot-policy and action models** (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
+- **Robot-policy and action models** (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1)
 
 ## Getting Started
 
@@ -67,6 +68,7 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
 
 - [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
 - [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
+- [Nightly CUDA Docker images](https://hub.docker.com/r/vllm/vllm-omni/tags?name=nightly) built from the latest `main`
 - [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)
 - [Deployment Recipes](https://recipes.vllm.ai) for vLLM-Omni model serving
 

docs/.nav.yml

@@ -37,7 +37,7 @@ nav:
       - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
       - Speech-To-Video: user_guide/examples/online_serving/speech_to_video.md
       - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
-      - Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
+      - Text-To-Speech (Online Serving): user_guide/examples/online_serving/text_to_speech.md
       - Text-To-Video: user_guide/examples/online_serving/text_to_video.md
   - General:
     - usage/*
@@ -127,6 +127,8 @@ nav:
         - Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
         - Automatic Prefix Caching: design/feature/prefix_caching.md
         - Realtime AR-Diffusion: design/feature/realtime_ar_diffusion.md
+      - Model-Specific Capabilities:
+        - MiniMax H3 Temporal VAE Chunks: design/feature/minimax_h3_temporal_vae_chunks.md
       - Communication:
         - OmniConnector Implementations:
           - Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md

docs/README.md

@@ -57,9 +57,9 @@ vLLM-Omni is flexible and easy to use with:
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
 - Omni-modality models (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
-- TTS models (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
+- TTS models (e.g. Qwen3-TTS, IndexTTS 2.5, dots.tts, CosyVoice3)
-- Diffusion models — image, video, and audio generation (e.g. MiniMax H3, Qwen-Image, Wan2.2, FLUX)
+- Diffusion models — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2)
-- Robot-policy and action models (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
+- Robot-policy and action models (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1)
 
 For more information, checkout the following:
 

docs/configuration/README.md

@@ -2,7 +2,7 @@
 
 This section lists the most common options for running vLLM-Omni.
 
-For options within a vLLM Engine, please refer to the [vLLM 0.26 configuration guide](https://docs.vllm.ai/en/v0.26.0/configuration/index.html).
+For options within a vLLM Engine, please refer to the [vLLM 0.28 configuration guide](https://docs.vllm.ai/en/v0.28.0/configuration/index.html).
 
 Each model defines fixed topology in a registered `PipelineConfig` and runtime
 overrides in a deploy YAML.

docs/configuration/environment_variables.md

@@ -151,7 +151,7 @@ their keys only.
 ## Inherited vLLM variables
 
 vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to
-the [vLLM 0.26 environment-variable reference](https://docs.vllm.ai/en/v0.26.0/configuration/env_vars.html)
+the [vLLM 0.28 environment-variable reference](https://docs.vllm.ai/en/v0.28.0/configuration/env_vars.html)
 for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm,
 XPU, ModelScope, and FlashInfer workspace settings.
 

docs/design/architecture_overview.md

@@ -170,6 +170,10 @@ because the user-visible output is diffusion media rather than a token stream.
 The
 [MiniMax-H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md)
 contains the hardware-specific profiles and warmup requirements.
+The internal, opt-in
+[temporal VAE chunk callback](feature/minimax_h3_temporal_vae_chunks.md)
+documents the narrower released-source contract for publishing committed video
+frames during decode. It does not change the complete-output serving topology.
 
 ### Cosmos3: unified MoT reasoner and generator
 

docs/design/feature/minimax_h3_temporal_vae_chunks.md

@@ -0,0 +1,260 @@
+# MiniMax H3 temporal VAE chunk callback
+
+!!! warning "Experimental internal capability"
+
+    This feature is an opt-in model integration seam. The current MiniMax H3
+    pipeline still uses complete-output VAE decode, and no CLI option, serving
+    path, transport, or MP4 sink enables temporal chunks by default.
+
+## Status and scope
+
+MiniMax H3 decodes video latents in temporal windows, but the released VAE
+normally exposes only the complete reconstructed video. The temporal callback
+publishes frames after the released decoder has completed overlap blending and
+padding removal. This gives a future runtime consumer a stable point at which
+it can schedule representation conversion, D2H transfer, and encoding while
+the VAE continues decoding later windows.
+
+The implementation remains model-local:
+
+- `vae.py` owns VAE topology and the thin
+  `decode_latent_with_chunks()` facade;
+- `chunked_decode.py` owns the normalized callback contract, metadata
+  validation, and fixed VAE-group rendezvous points; and
+- `temporal_chunks.py` owns the source-gated compatibility planner and the
+  released temporal assembly control flow.
+
+Crop policy, typed media, transport, backpressure, encoding, and public output
+remain runtime responsibilities. The model code does not schedule requests or
+select a presentation format.
+
+## Execution flow
+
+The default and opt-in paths are intentionally separate:
+
+```text
+Default:
+latent -> MiniMaxH3VideoVAE.decode_latent() -> complete RGB video
+
+Opt-in producer seam:
+latent
+  -> MiniMaxH3VideoVAE.decode_latent_with_chunks()
+  -> group-local preflight and callback-owner validation
+  -> released temporal decode / overlap blending / padding removal
+  -> borrowed decoder-domain chunk
+  -> normalize + clone into owned RGB float32 BCTHW [0, 1]
+  -> synchronous callback
+  -> complete RGB video is still returned
+```
+
+Calling `decode_latent_with_chunks()` is the opt-in. Supplying no callback on a
+VAE peer does not select the default decoder; every rank in a VAE patch group
+must enter the chunked method, and only the output rank receives the callback.
+
+## Producer contract
+
+The normalized callback has this structural signature:
+
+```python
+def callback(
+    frames: torch.Tensor,
+    /,
+    *,
+    chunk_index: int,
+    total_chunks: int,
+    frame_start: int,
+    is_final: bool,
+) -> None:
+    ...
+```
+
+For a successful decode:
+
+- `frames` is owned, contiguous RGB float32 in `BCTHW` layout and `[0, 1]`;
+- callback storage does not alias the returned complete video;
+- the consumer may retain or mutate a chunk;
+- `chunk_index` is zero-based, contiguous, and strictly increasing;
+- `total_chunks` is positive and stable for the decode;
+- `frame_start` equals the end of the previous published range;
+- chunks are nonempty and contain only temporally committed frames;
+- exactly one chunk has `is_final=True`; and
+- concatenating chunks along time exactly reconstructs the returned video.
+
+The callback runs synchronously on the producer's current device stream. A
+consumer that reads the tensor on another stream must establish ordering, for
+example with an event, and retain the allocation for that stream's lifetime.
+A slow callback also blocks VAE production. An overlap-oriented runtime
+consumer should therefore acquire bounded capacity, enqueue side-stream work,
+and return promptly instead of encoding frames inside the callback.
+
+Chunks retain the decoded VAE canvas dimensions. A future runtime integration
+must apply the same requested-height and requested-width crop as the complete
+MiniMax H3 pipeline before declaring the media transport-ready.
+
+The raw compatibility callback is deliberately narrower. It receives borrowed
+decoder-domain tensors and must not retain or mutate them. The coordinator is
+the ownership boundary that normalizes and clones raw chunks before invoking
+the public callback.
+
+## Temporal planning and terminal publication
+
+The released configuration advances five latent tokens per temporal window and
+includes a two-token overlap, so each `_adaptive_decode` input normally contains
+seven tokens. After temporal scaling, pre-padding removal, and five-frame
+overlap blending, the normal committed unit is 17 output frames. The final
+remainder depends on the requested frame plan.
+
+Representative released frame plans include:
+
+| Output frames | Published chunks |
+| ---: | --- |
+| 124 | `7 x 17 + 5` |
+| 209 | `12 x 17 + 5` |
+| 243 | `14 x 17 + 5` |
+| 362 | `21 x 17 + 5` |
+
+The terminal callback is held until the complete frame plan passes its logical
+frame, padding, written-frame, chunk-count, and finality checks. Earlier chunks
+are semantically stable, but they do not imply that the overall request has
+succeeded; a consumer must discard or abort partial output if a later decoder
+failure terminates the request.
+
+## Source and configuration gate
+
+The released VAE has no public temporal-callback API. vLLM-Omni therefore owns
+a compatibility implementation for one exact released code contract. The
+reference revision identifies its provenance, but runtime acceptance is based
+on source content, configuration, and structure rather than repository or
+weight metadata:
+
+- reference Hugging Face revision, not independently checked at runtime:
+  `42ed227ee7df40d41602854ae760620d6eb651fe`;
+- sorted 15-file top-level `video_vae/*.py` manifest (excluding
+  `__init__.py`) SHA-256:
+  `30e64afbaa940696bf8bfe2c86003a4b1666c84b22fa3aa12403a3e6a03f705d`;
+- clip length `17`;
+- token chunk, overlap, and drop `5`, `2`, and `3`;
+- temporal ratio `4`;
+- frame pre-padding and overlap `3` and `5`; and
+- no isolated first or last frame.
+
+The adapter does not verify Hugging Face revision metadata, repository identity,
+or model weights. The manifest pin detects Python compatibility drift; it is
+not a security boundary for remote-code execution. The source manifest, listed
+temporal attributes, structural contract, plan, and dtype are validated before
+the first `_adaptive_decode` call. A mismatch disables only the opt-in chunk path.
+`decode_latent()` remains available and performs no source fingerprint or chunk
+rendezvous.
+
+## Supported models
+
+Support in this table means the internal producer contract above, not public
+serving or end-to-end transport support.
+
+| Model or component | Source contract | Status | Validation |
+| --- | --- | --- | --- |
+| MiniMax H3 FL2VA video VAE | Released manifest/listed temporal attributes above (reference `42ed227e`) | Supported, opt-in | Same-environment full-output parity and L20X/B300 feasibility measured on the implementation through PR commit `60e9bbad` |
+| MiniMax H3 Ref2VA video VAE | Same released manifest/listed temporal attributes | Source-compatible; not separately benchmarked | Ref2VA and FL2VA use byte-identical video-VAE source and configuration in the reference snapshot |
+| MiniMax H3 with a different Python manifest or listed temporal attribute | Different enforced contract | Unsupported for chunk mode | Fails closed before decoder collectives; complete decode remains available |
+| Same manifest/temporal attributes with different weights or other component metadata | Not checked by this gate | Not independently qualified | The adapter does not inspect weights, repository identity, revision metadata, or arbitrary component configuration |
+| Other video models | Any | Unsupported | No generic temporal VAE producer ABI is declared by this feature |
+
+## Feature compatibility
+
+| Dimension or combination | Status | Constraint |
+| --- | --- | --- |
+| Existing complete-output decode | Supported and default | No source check, callback object, or callback rendezvous |
+| Explicit chunk decode | Supported internally | Caller invokes `decode_latent_with_chunks()` directly |
+| VAE patch parallel size 1 | Supported by the model seam | Caller supplies the callback only on the request output owner; no WORLD rendezvous is added |
+| VAE patch parallel size equal to the full DiT group | Supported when data parallel size is 1 | Every VAE-group rank calls; exactly group rank 0 owns the callback |
+| Partial VAE patch subgroup | Unsupported | MiniMax H3 currently accepts size 1 or the full DiT group only |
+| Data parallel size greater than 1 with VAE patch parallel size 1 | Model seam only | Per-replica callback ownership is possible, but no serving/runtime consumer is wired |
+| Data parallel size greater than 1 with VAE patch parallel size greater than 1 | Unsupported | Fails before decoder collectives because no replica-local VAE patch group is defined |
+| Too few spatial tiles for the VAE group | Correctness fallback | Decodes rank-locally and restores native tiling state; slower than patch parallel decode |
+| Device memory and full-output assembly | Not reduced or bounded | The compatibility path still assembles and returns the full decode while allocating owned callback chunks; no peak-HBM reduction is claimed |
+| DiT cache, attention, TP, SP, HSDP, LoRA, or quantization features | No direct callback dependency | This seam runs after denoising and does not independently qualify combinations beyond the MiniMax H3 model's own matrix |
+| Layerwise or distributed component offload | Not separately validated | Component residency must place the video VAE before decode |
+| Step execution or request batching | Not integrated | No pipeline caller binds a runtime chunk consumer |
+| Requested-size crop and device uint8 preparation | Not implemented here | Runtime-owned follow-up |
+| Pinned D2H, SHM/ZMQ, or cross-process transport | Not implemented here | Runtime-owned follow-up |
+| Complete MP4 or fragmented MP4 sink | Not implemented here | Existing serving behavior still waits for complete output |
+| Audio VAE chunks and incremental A/V composition | Not supported | Video chunks do not imply audio or request finality |
+| Cancellation and multi-output lease ownership | Not implemented here | Requires a bounded runtime lifecycle |
+| NVIDIA CUDA | Feasibility evidence | L20X and B300 measurements apply to the implementation through `60e9bbad`; the current coordination refactor was not reprofiled |
+| AMD ROCm | Not validated for this feature | Base MiniMax H3 support is documented separately and does not imply temporal-chunk validation |
+| Ascend, Intel, or other accelerators | No feature support claim | Consult the model support table independently; this design adds no platform qualification |
+
+## Distributed and failure semantics
+
+When VAE patch parallelism is active, all VAE-group ranks run the same
+preflight before entering checkpoint collectives. Source/planner errors and the
+resolved temporal concat dtype are reconciled at fixed group rendezvous points.
+
+Callback ownership is also a group decision. Rank `r` contributes `r + 1` when
+it owns a callback and zero otherwise. A group SUM equals one only when rank 0
+is the sole owner; every invalid owner configuration therefore fails
+identically on all ranks before decode.
+
+Callback exceptions are retained on the output rank. Further publication stops,
+but every rank completes the same remaining native decoder collective schedule.
+The original exception is then raised on the owner and
+`MiniMaxH3ChunkCallbackPeerError` on peers.
+
+Exceptions raised inside the checkpoint's native distributed decode cannot in
+general be reconciled by the adapter: a peer may already be blocked inside a
+checkpoint-owned collective. Resolving such failures requires native collective
+handling or supervisor-level timeout and worker-group teardown. The process
+group must not be assumed reusable after such a failure.
+
+## Validation evidence
+
+The contract suite covers ordinary and padded frame plans, blend boundaries,
+owned storage and mutation isolation, source-manifest/listed-temporal-attribute
+drift, callback failure, and healthy subsequent requests. Real Gloo processes cover
+rank-asymmetric preflight failures, dtype disagreement, two-rank callback
+failure propagation, three-rank wrong-owner rejection, insufficient-tile
+fallback, and group recovery for failures that occur at adapter-owned
+rendezvous points.
+
+Four-device L20X and independent four- and eight-device B300 experiments on the
+implementation through PR commit `60e9bbad` showed exact serial-versus-chunked
+frame/MP4 parity within each environment and lower VAE-to-complete-MP4 latency
+when an experimental consumer was attached. Hashes are not claimed to match
+across different hardware/software environments. The later coordination and
+module-boundary refactors were revalidated with CPU and real Gloo tests but were
+not reprofiled on GPUs. These are feasibility measurements, not a promise for
+default serving, because this change does not contain the transport or encoder
+consumer. See [RFC #6872](https://github.com/vllm-project/vllm-omni/issues/6872)
+and [PR #6885](https://github.com/vllm-project/vllm-omni/pull/6885) for the
+measurement provenance and limitations.
+
+## Generalization boundary and follow-ups
+
+This feature does not define a repository-wide `SupportsVideoChunkDecode`
+interface. H3 temporal blending, Wan causal caches, LTX temporal/spatial halos,
+and pipeline-level generation chunks have different commit and failure
+boundaries. A generic producer contract should be extracted only after a second
+model proves the same semantic event without `Any` arguments or model-name
+dispatch.
+
+The reusable future boundary is the emitted, committed video-media chunk, not a
+universal latent-to-VAE method. It should extend the typed media work in
+[RFC #6541](https://github.com/vllm-project/vllm-omni/issues/6541) and
+[PR #6615](https://github.com/vllm-project/vllm-omni/pull/6615), while keeping
+model-specific halo, cache, padding, blending, distributed, and source-gating
+logic inside each model adapter.
+
+Runtime follow-ups own:
+
+- typed video-plus-audio media and request/output identity;
+- requested-size crop and device preparation;
+- bounded leases, backpressure, cancellation, and cleanup;
+- D2H/IPC transport and observability; and
+- complete-MP4 and fragmented-MP4 consumers.
+
+## Related documentation
+
+- [Diffusion model integration](../module/diffusion/diffusion_model_integration.md)
+- [VAE patch parallelism](vae_parallel.md)
+- [Async diffusion output](async_diffusion_output.md)
+- [MiniMax H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md)

docs/design/index.md

@@ -25,6 +25,10 @@ implementation contract; it is not, by itself, a general support claim.
 - [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
 - [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md)
 
+### Model-specific capabilities
+
+- [MiniMax H3 Temporal VAE Chunks](feature/minimax_h3_temporal_vae_chunks.md)
+
 ### Communication
 
 #### OmniConnector implementations

docs/design/module/diffusion/diffusion_model_integration.md

@@ -57,3 +57,10 @@ cross-stage routing.
 
 Test registry selection, checkpoint loading, minimal inference, input and output
 contracts, and every declared optional capability.
+
+## Model-specific capability documents
+
+- [MiniMax H3 temporal VAE chunk callback](../../feature/minimax_h3_temporal_vae_chunks.md)
+  defines the source-gated producer seam and its boundary with future runtime
+  media transport. It is an internal opt-in capability, not a public serving
+  feature.

docs/getting_started/installation/gpu.md

@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0; the 0.28.0 wheels are published together with the v0.28.0 release this branch produces. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, 0.26.0, and 0.28.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

docs/getting_started/installation/gpu/cuda.inc.md

@@ -5,7 +5,7 @@
 # --8<-- [end:requirements]
 # --8<-- [start:set-up-using-python]
 
-vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.26.x release line uses vLLM 0.26.x.
+vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.28.x release line uses vLLM 0.28.x.
 
 !!! note
     PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details.
@@ -95,6 +95,13 @@ uv pip install --no-build-isolation --editable .
 
 vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images and available on Docker Hub as [vllm/vllm-omni](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which release of vLLM it is based on.
 
+Nightly CUDA images built from the latest `main` are also available:
+
+- [`vllm/vllm-omni:nightly`](https://hub.docker.com/r/vllm/vllm-omni/tags?name=nightly) is the rolling multi-architecture tag.
+- [`vllm/vllm-omni:nightly-<commit>`](https://hub.docker.com/r/vllm/vllm-omni/tags?name=nightly-) provides commit-pinned multi-architecture tags. The newest 14 are retained.
+
+Use a commit-pinned tag when you need a reproducible nightly environment.
+
 Here's an example deployment command that has been verified on 2 x H100's:
 ```bash
 docker run --runtime nvidia --gpus 2 \

docs/getting_started/installation/npu/npu.inc.md

@@ -9,8 +9,8 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built
 
 ```bash
 # vLLM-Ascend image
-# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.26.0
+# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.28.0
-# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.26.0-a3
+# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.28.0-a3
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -28,11 +28,11 @@ docker run --rm \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
     -v /root/.cache:/root/.cache \
     -p 8000:8000 \
-    -it quay.io/atlas-ci/vllm-ascend:v0.26.0-a3 bash
+    -it quay.io/atlas-ci/vllm-ascend:v0.28.0-a3 bash
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.26.0 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.28.0 https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
@@ -50,21 +50,11 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
 
 # --8<-- [start:installation-main]
 
-You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
+You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.28.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU.
 
 ```bash
-# Pin vLLM and vLLM-Ascend to the v0.26 release line
+# Inside the vLLM-Ascend v0.28.0 container, install vLLM-Omni from main
-git clone -b v0.26.0 https://github.com/vllm-project/vllm.git
+cd /vllm-workspace
-cd vllm
-VLLM_TARGET_DEVICE=empty pip install -v -e .
-cd ..
-
-git clone -b releases/v0.26.0rc https://github.com/vllm-project/vllm-ascend.git
-cd vllm-ascend
-pip install -v -e .
-cd ..
-
-# Install vLLM-Omni from the latest main branch
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 pip install -v -e . --no-build-isolation
@@ -88,8 +78,8 @@ Supported images as following.
 Here's an example deployment command that has been verified on 4 x NPUs:
 
 ```bash
-# Atlas A2: quay.io/ascend/vllm-omni:v0.26.0
+# Atlas A2: quay.io/ascend/vllm-omni:v0.28.0
-# Atlas A3: quay.io/ascend/vllm-omni:v0.26.0-a3
+# Atlas A3: quay.io/ascend/vllm-omni:v0.28.0-a3
 docker run --rm \
     --name vllm-omni-a3 \
     --shm-size=4g \
@@ -114,7 +104,7 @@ docker run --rm \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
     -v ~/.cache:/root/.cache \
     -p 8091:8091 \
-    -it quay.io/ascend/vllm-omni:v0.26.0-a3 bash
+    -it quay.io/ascend/vllm-omni:v0.28.0-a3 bash
 ```
 
 !!! tip

docs/models/supported_models.md

@@ -95,7 +95,7 @@ th {
 | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
 | `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
 | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
+| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only), SenseNova-U1.5 | `SenseNova/SenseNova-U1-8B-MoT`, `sensenova/SenseNova-U1.5-8B-MoT` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.5.md) |
 | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
 | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
 | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |

docs/serving/speech_api.md

@@ -2,7 +2,7 @@
 
 vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. Supported TTS models include:
 
-- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
+- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Native output: 24 kHz; API output can be resampled to 8 kHz.
 - **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
 - **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
 - **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz.
@@ -116,6 +116,7 @@ Content-Type: application/json
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
+| `sample_rate` | integer | model native | Target output sample rate. Qwen3-TTS supports 8000 or 24000 Hz; the model remains native 24 kHz internally and is resampled before encoding. |
 | `task_type` | string | "CustomVoice" | TTS task type: CustomVoice, VoiceDesign, or Base |
 | `language` | string | "Auto" | Language (see supported languages below) |
 | `instructions` | string | "" | Voice style/emotion instructions |

docs/user_guide/diffusion_features.md

@@ -138,7 +138,7 @@ The following tables show which models support each feature:
 | **Qwen-Image-Edit**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Qwen-Image-Edit-2509** |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |        ✅ (decode)         |          ✅           |       ❌        |        ❌         |
 | **Qwen-Image-Layered**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **SenseNova-U1**         |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **SenseNova-U1 / U1.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Stable-Diffusion-XL**  |     ❌     |     ❌      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Stable-Diffusion3.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Z-Image**              |     ✅     |     ✅      |           ✅           |       ❓        |   ✅ (TP=2 only)   |          ❌          |   ✅    |             ❌             |      ✅ (decode)      |       ✅        |        ❌         |

examples/online_serving/text_to_speech/qwen3_tts/openai_speech_client.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """OpenAI-compatible client for Qwen3-TTS via /v1/audio/speech endpoint.
 
 This script demonstrates how to use the OpenAI-compatible speech API
@@ -74,6 +77,8 @@ def run_tts_generation(args) -> None:
         "voice": args.speaker,
         "response_format": args.response_format,
     }
+    if args.sample_rate is not None:
+        payload["sample_rate"] = args.sample_rate
 
     # Add optional parameters
     if args.instructions:
@@ -246,6 +251,13 @@ def parse_args():
         choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
         help="Audio output format (default: wav)",
     )
+    parser.add_argument(
+        "--sample-rate",
+        type=int,
+        default=None,
+        choices=[8000, 24000],
+        help="Output sample rate in Hz (default: model native)",
+    )
     parser.add_argument(
         "--output",
         "-o",

recipes/MiniMaxAI/MiniMax-H3.md

@@ -1142,3 +1142,4 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
 - [Supported models](../../docs/models/supported_models.md)
 - [Video API](../../docs/serving/videos_api.md)
 - [Diffusion parallelism](../../docs/user_guide/diffusion/parallelism/overview.md)
+- [Temporal VAE chunk callback design](../../docs/design/feature/minimax_h3_temporal_vae_chunks.md)

recipes/README.md

@@ -78,6 +78,7 @@ recipes/
 | [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
 | [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB / 1x AMD MI300X |
 | [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X |
+| [`SenseNova/SenseNova-U1.5.md`](./SenseNova/SenseNova-U1.5.md) | Unified image generation and understanding, with an 8-step distilled LoRA | 1x A800 80GB |
 | [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X |
 | [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
 | [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |

recipes/SenseNova/SenseNova-U1.5.md

@@ -0,0 +1,159 @@
+# SenseNova-U1.5
+
+> Unified image generation and understanding, with an official 8-step distilled LoRA
+
+## Summary
+
+- Vendor: SenseNova
+- Model: `sensenova/SenseNova-U1.5-8B-MoT`
+- LoRA: `sensenova/SenseNova-U1.5-8B-MoT-LoRAs` (`SenseNova-U1.5-8B-MoT-LoRA-8step.safetensors`)
+- Task: text2img, img2img, img2text (visual understanding), text2text (chat)
+- Mode: Offline inference, Online serving (OpenAI-compatible API)
+- Maintainer: Community
+
+## When to use this recipe
+
+U1.5 runs on the same `SenseNovaU1Pipeline` as U1 — the checkpoint keeps
+`model_type: neo_chat`, so it resolves without `--model-class-name`. Relative to U1 the
+config flips two fields, `use_pixel_head` to `true` (the flow-matching head becomes a
+`ConvDecoder`) and `noise_scale_max_value` from 8.0 to 16.0; both are already read by
+`SenseNovaU1Config`. The checkpoint is 13 shards / 50.2 GB on disk, 30.3 GB of that fp32, and
+loads to roughly 34 GB in bf16.
+
+Use this recipe for U1.5 specifically, including its distilled few-step LoRA. For U1 see
+[SenseNova-U1](SenseNova-U1.md).
+
+## Hardware Support
+
+### GPU
+
+#### 1x A800 80GB
+
+- 1x NVIDIA A800 80GB PCIe with an Intel Xeon Gold 6336Y, driver 595.84, Linux
+- Python 3.12.13, vLLM 0.28.0, torch 2.13.0+cu130, CUDA 13.0, diffusers 0.40.0,
+  transformers 5.14.1, flashinfer 0.6.16.post3, BF16, TP=1
+- Peak GPU memory: 34.3 GB at 1024x1024, 36.4 GB at 1536x2720
+
+##### Text-to-Image
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model sensenova/SenseNova-U1.5-8B-MoT \
+    --prompt "Close portrait of an elderly woman by a farmhouse window, warm natural light." \
+    --width 1024 --height 1024 \
+    --seed 42 --num-inference-steps 50 --cfg-scale 4.0 \
+    --extra-body '{"think": false, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
+    --output sensenova_u15_t2i.png
+```
+
+Think mode (`"think": true`) is recommended for higher image quality.
+
+##### Image-to-Image Editing
+
+```bash
+python examples/offline_inference/image_to_image/image_edit.py \
+    --model sensenova/SenseNova-U1.5-8B-MoT \
+    --prompt "Turn this into an oil painting" \
+    --image input.png --resolution 1024 \
+    --seed 42 --num-inference-steps 25 --cfg-scale 4.0 \
+    --extra-args '{"think": true, "img_cfg_scale": 1.0, "cfg_norm": "none", "timestep_shift": 3.0}' \
+    --output sensenova_u15_edit.png
+```
+
+##### 8-step distilled LoRA
+
+The distilled LoRA is fused into the generation tower at load time:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model sensenova/SenseNova-U1.5-8B-MoT \
+    --lora-path SenseNova-U1.5-8B-MoT-LoRA-8step.safetensors --lora-backend distill \
+    --prompt "Close portrait of an elderly woman by a farmhouse window, warm natural light." \
+    --width 1024 --height 1024 \
+    --seed 42 --num-inference-steps 8 --cfg-scale 1.0 \
+    --extra-body '{"think": false, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
+    --output sensenova_u15_lora8.png
+```
+
+**Use `--cfg-scale 1.0` with this LoRA.** It is distilled with DMD and runs without
+classifier-free guidance; the default `4.0` applies guidance twice and produces a
+blown-out, posterised image.
+
+##### Online serving
+
+```bash
+vllm serve sensenova/SenseNova-U1.5-8B-MoT --omni --port 8091
+
+python examples/online_serving/sensenova_u1/openai_chat_client.py \
+    -s http://127.0.0.1:8091 -m img2text -i input.png -p "Describe this image."
+```
+
+`-s` takes the base URL; the client appends `/v1` itself.
+
+#### Measured latency (1x A800 80GB, 25 steps, median of 3 after a warmup)
+
+| Resolution | Step latency | Total |
+| --- | --- | --- |
+| 1024x1024 | 209.5 ms | 5.24 s |
+| 1536x1536 | 458.5 ms | 11.46 s |
+
+At 1536x2720 with 50 steps, end-to-end is 43.0 s.
+
+#### Measured latency (1x NVIDIA H200 139GB, single run, no warmup)
+
+Reported by @hsliuustc0106 against PR head `29c090d`, on one reserved CUDA device.
+
+- Python 3.12.13, vLLM 0.28.0, torch 2.13.0+cu130, CUDA 13.0, diffusers 0.40.0,
+  transformers 5.14.1, flashinfer 0.6.16.post3, BF16, TP=1
+- `sensenova/SenseNova-U1.5-8B-MoT` with `SenseNova-U1.5-8B-MoT-LoRA-8step.safetensors`,
+  seed 42, 1024x1024
+
+| Case | Stage latency | Peak GPU memory |
+| --- | --- | --- |
+| 50 steps, cfg 4.0, think off | 9017 ms | 34384 MiB |
+| 8 steps, cfg 1.0, no LoRA control | 624.00 ms, 77.85 ms/step | 34376 MiB |
+| 8 steps, cfg 1.0, distilled LoRA | 623.85 ms, 77.98 ms/step | 34364 MiB |
+| 50 steps, cfg 4.0, think on, `VLLM_OMNI_SENSENOVA_PAGED_DECODE=1` | 81622 ms | 34594 MiB |
+
+Each row runs the Text-to-Image command above with `--num-inference-steps` and
+`--cfg-scale` as listed and `think` set through `--extra-body`; the LoRA row adds
+`--lora-path` and `--lora-backend distill`.
+
+Single runs without a preceding warmup, so the think-on row carries the first-request
+compile rather than a steady state. The distilled LoRA fused into 168 parameters, and its
+image differs from the same-seed no-LoRA control in 1,048,574 of 1,048,576 pixels, MAE 58.62.
+
+```bash
+pytest -m "cpu and not cuda" tests/diffusion/models/sensenova_u1/ tests/diffusion/lora/test_loader.py -q
+pytest -m "cpu and not cuda" tests/config/test_environment_variables.py tests/diffusion/test_diffusion_worker.py -q
+pytest -m cuda tests/diffusion/models/sensenova_u1/test_sensenova_u1_attention_gqa.py \
+    tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py -q --tb=short
+```
+
+71, 13 and 9 passed. Serving reached `/health`; `img2text` returned a 1921-character
+description and `text2text` the expected answer.
+
+#### Verification
+
+```bash
+pytest -q tests/diffusion/models/sensenova_u1/
+```
+
+## Notes
+
+- Both `use_pixel_head` and `noise_scale_max_value` come from `config.json`; no flag is needed.
+- The checkpoint carries no `configuration_neo_chat.py`, so the loader logs a few
+  "does not appear to have a file named configuration_neo_chat.py" errors during startup and
+  then proceeds on the in-tree config. Generation is unaffected.
+- The LoRA targets the generation tower only (`*_mot_gen`); understanding-tower weights are
+  untouched.
+- Autoregressive decode (think, text-to-text and image-to-text) runs on a paged K/V cache under
+  a captured CUDA graph. It falls back to the ordinary cache when the device or the bundled
+  `flash_attn_varlen_func` cannot support it; set `VLLM_OMNI_SENSENOVA_PAGED_DECODE=0` to force
+  that fallback.
+- The first request after startup costs about 0.7 s more than the steady state whether the paged
+  path is on or off. Measured on one A800 with the inductor, triton and vLLM compile caches all
+  cleared, median of three runs: 718 ms above steady with the path on, 679 ms with it off.
+- Each request captures its own graphs, and a think request captures twice because the sequence
+  grows past the 512 bucket, so the capture cost is paid per request rather than once at
+  startup.

tests/buildkite/test_amd_pipeline.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from pathlib import Path
+from shlex import split
+
+import pytest
+import yaml
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+AMD_MERGE_PIPELINE = Path(".buildkite/amd/test-amd-merge.yml")
+
+
+def _find_step(label: str) -> dict:
+    pipeline = yaml.safe_load(AMD_MERGE_PIPELINE.read_text(encoding="utf-8"))
+
+    def walk(steps: list[dict]) -> dict | None:
+        for step in steps:
+            if step.get("label") == label:
+                return step
+            if nested := walk(step.get("steps", [])):
+                return nested
+        return None
+
+    step = walk(pipeline.get("steps", []))
+    assert step is not None, f"missing AMD pipeline step: {label}"
+    return step
+
+
+def test_qwen3_tts_base_preserves_advanced_model_arguments() -> None:
+    step = _find_step("Qwen3-TTS Base E2E Test")
+    commands = step["commands"]
+
+    assert all("bash -c" not in command for command in commands)
+    pytest_command = next(command for command in commands if "pytest" in command)
+    argv = split(pytest_command)
+
+    marker_index = argv.index("-m")
+    run_level_index = argv.index("--run-level")
+    assert argv[marker_index + 1] == "advanced_model and cuda"
+    assert argv[run_level_index + 1] == "advanced_model"

tests/config/test_config_factory.py

@@ -2399,6 +2399,21 @@ class TestBaseConfigInheritance:
         # CI overrides max_tokens
         assert s0["max_tokens"] == 150
 
+    def test_qwen3_omni_colocate_async_bounds_only_rocm_kv_cache(self):
+        ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe_colocate_async.yaml"))
+        pipeline = resolve_pipeline_config("qwen3_omni_moe_thinker_only")
+        assert isinstance(pipeline, PipelineConfig)
+
+        cuda = _apply_platform_overrides(load_deploy_config(ci_path), platform="cuda")
+        cuda_stage = merge_pipeline_deploy(pipeline, cuda)[0]
+        assert cuda_stage.yaml_engine_args["gpu_memory_utilization"] == 0.9
+        assert "kv_cache_memory_bytes" not in cuda_stage.yaml_engine_args
+
+        rocm = _apply_platform_overrides(load_deploy_config(ci_path), platform="rocm")
+        rocm_stage = merge_pipeline_deploy(pipeline, rocm)[0]
+        assert "gpu_memory_utilization" not in rocm_stage.yaml_engine_args
+        assert rocm_stage.yaml_engine_args["kv_cache_memory_bytes"] == 2 * 1024**3
+
     def test_pure_inheritance_overlay(self, tmp_path):
         """An overlay with only ``base_config`` inherits everything."""
         base = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
@@ -2860,6 +2875,54 @@ stages:
         assert omega_config.engine_args.max_model_len == 8192
         assert omega_config.engine_args.enforce_eager is True
 
+    def test_cli_attention_shorthand_replaces_yaml_structured_default(self):
+        stage = StageConfig(
+            stage_id=0,
+            model_stage="dit",
+            stage_type=StageType.DIFFUSION,
+            yaml_engine_args={
+                "diffusion_attention_config": {
+                    "default": {"backend": "FLASH_ATTN"},
+                    "per_role": {"cross": {"backend": "TORCH_SDPA"}},
+                },
+            },
+            runtime_overrides={"diffusion_attention_backend": "SAGE_ATTN"},
+        )
+
+        engine_args = stage.to_omegaconf().engine_args
+
+        assert engine_args.diffusion_attention_backend == "SAGE_ATTN"
+        assert "default" not in engine_args.diffusion_attention_config
+        assert engine_args.diffusion_attention_config.per_role.cross.backend == "TORCH_SDPA"
+
+    def test_cli_attention_shorthand_keeps_yaml_per_role_only_config(self):
+        stage = StageConfig(
+            stage_id=0,
+            model_stage="dit",
+            stage_type=StageType.DIFFUSION,
+            yaml_engine_args={"diffusion_attention_config": {"per_role": {"cross": {"backend": "TORCH_SDPA"}}}},
+            runtime_overrides={"diffusion_attention_backend": "SAGE_ATTN"},
+        )
+
+        engine_args = stage.to_omegaconf().engine_args
+
+        assert engine_args.diffusion_attention_backend == "SAGE_ATTN"
+        assert engine_args.diffusion_attention_config == {"per_role": {"cross": {"backend": "TORCH_SDPA"}}}
+
+    def test_cli_attention_config_replaces_yaml_shorthand(self):
+        stage = StageConfig(
+            stage_id=0,
+            model_stage="dit",
+            stage_type=StageType.DIFFUSION,
+            yaml_engine_args={"diffusion_attention_backend": "TORCH_SDPA"},
+            runtime_overrides={"diffusion_attention_config": {"default": {"backend": "SAGE_ATTN"}}},
+        )
+
+        engine_args = stage.to_omegaconf().engine_args
+
+        assert "diffusion_attention_backend" not in engine_args
+        assert engine_args.diffusion_attention_config.default.backend == "SAGE_ATTN"
+
 
 class TestSentinelDefaultPrecedence:
     """Caller-typed (non-None) values win over YAML; None values fall through

tests/config/test_environment_variables.py

@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 24,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 57,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
         EnvironmentVariableCategory.INTERNAL: 2,
     }
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
     )
     assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
-        ModelEnvironmentVariableDisposition.PROMOTE: 33,
+        ModelEnvironmentVariableDisposition.PROMOTE: 34,
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 11,

tests/diffusion/lora/test_loader.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from collections import defaultdict
 
 import pytest
@@ -106,7 +109,7 @@ def make_dummy_lora_state_dict(
     """
     key_a = f"{base_key}.{lora_a_suffix}"
     key_b = f"{base_key}.{lora_b_suffix}"
-    key_bias = f"{base_key}.{lora_bias_suffix}" if bias else None
+    key_bias = f"{base_key}.{lora_bias_suffix}"
 
     if weight_a is None:
         weight_a = torch.randn(RANK, HEAD_DIM)
@@ -436,3 +439,54 @@ def test_wan_load_lora_weights_targets_both_transformers(mocker: MockerFixture):
     assert_params_equal(pipeline.transformer, original_high)
     assert_params_equal(pipeline.transformer_2, original_low)
     assert "adapter0" not in pipeline.lora_loaded
+
+
+def test_prepare_lora_delta_disambiguates_prefix_overlapping_stacked_names():
+    """A fused name can be a prefix of another fused name.
+
+    SenseNova-U1 carries both `.qkv_proj` (understanding) and `.qkv_proj_mot_gen`
+    (generation). Selecting stacked mappings by substring matched both rules for
+    the generation weight and concatenated six shards into a delta twice as tall
+    as the parameter.
+    """
+    param_to_weight_names = {
+        ".qkv_proj_mot_gen": [".q_proj_mot_gen", ".k_proj_mot_gen", ".v_proj_mot_gen"],
+        ".qkv_proj": [".q_proj", ".k_proj", ".v_proj"],
+    }
+    state_dict = {}
+    for part in ("q", "k", "v"):
+        for tower in ("_mot_gen", ""):
+            base = f"layers.0.self_attn.{part}_proj{tower}"
+            state_dict[f"{base}.lora_A.weight"] = torch.ones(RANK, HEAD_DIM)
+            state_dict[f"{base}.lora_B.weight"] = torch.ones(HEAD_DIM, RANK)
+
+    for base_key in ("layers.0.self_attn.qkv_proj_mot_gen", "layers.0.self_attn.qkv_proj"):
+        delta, _ = _prepare_lora_delta(state_dict, base_key, param_to_weight_names)
+        assert delta.shape == (3 * HEAD_DIM, HEAD_DIM), base_key
+
+
+@pytest.mark.parametrize(
+    "fused, shards",
+    [
+        # qwen_image_transformer.py stacked_params_mapping
+        (".to_qkv", (".to_q", ".to_k", ".to_v")),
+        (".add_kv_proj", (".add_q_proj", ".add_k_proj", ".add_v_proj")),
+        # wan2_2_transformer.py stacked_params_mapping
+        (".attn1.to_qkv", (".attn1.to_q", ".attn1.to_k", ".attn1.to_v")),
+        # wan2_2_s2v_transformer.py stacked_params_mapping
+        (".self_attn.to_qkv", (".self_attn.q", ".self_attn.k", ".self_attn.v")),
+    ],
+)
+def test_prepare_lora_delta_unchanged_for_existing_fused_layouts(fused, shards):
+    """None of the fused names shipped today overlap by prefix, so matching on the
+    tail must select exactly the shards substring matching selected."""
+    param_to_weight_names = {fused: list(shards)}
+    base_key = f"blocks.0{fused}"
+    state_dict = {}
+    for shard in shards:
+        state_dict[f"blocks.0{shard}.lora_A.weight"] = torch.ones(RANK, HEAD_DIM)
+        state_dict[f"blocks.0{shard}.lora_B.weight"] = torch.ones(HEAD_DIM, RANK)
+
+    delta, used_keys = _prepare_lora_delta(state_dict, base_key, param_to_weight_names)
+    assert delta.shape == (len(shards) * HEAD_DIM, HEAD_DIM)
+    assert len(used_keys) == 2 * len(shards)

tests/diffusion/models/sensenova_u1/test_sensenova_u1_attention_gqa.py

@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU tests for SenseNova-U1 attention head handling.
+
+The model used to expand K/V from 8 heads to 32 with ``repeat_interleave``
+before calling the attention backend, so the backend's own GQA dispatch never
+saw a GQA shape, and it passed an all-zeros mask during single-token decode,
+which only served to keep SDPA off its fused kernels. These pin both contracts:
+K/V must reach the backend uncompressed, decode must carry no mask, and the
+compressed path must compute the same thing the expanded one did.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
+    SenseNovaU1Attention,
+    SenseNovaU1Model,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+N_HEADS, N_KV_HEADS, HEAD_DIM, SEQ = 32, 8, 8, 6
+GROUPS = N_HEADS // N_KV_HEADS
+
+
+class _Recorder:
+    """Stands in for the Attention module and records what it was handed."""
+
+    def __init__(self):
+        self.calls: list[tuple[torch.Size, torch.Size, torch.Tensor | None]] = []
+
+    def __call__(self, query, key, value, attn_metadata):
+        mask = None if attn_metadata is None else attn_metadata.attn_mask
+        self.calls.append((query.shape, key.shape, mask))
+        return query
+
+
+class _AttnHost:
+    """Carries only the attributes ``_run_attn``/``_run_attn_bshd`` touch."""
+
+    num_kv_groups = GROUPS
+    _align_mask_dtype = staticmethod(SenseNovaU1Attention._align_mask_dtype)
+    _run_attn = SenseNovaU1Attention._run_attn
+    _run_attn_bshd = SenseNovaU1Attention._run_attn_bshd
+
+    def __init__(self):
+        self.attn = _Recorder()
+
+
+def test_bhsd_path_hands_the_backend_compressed_kv():
+    host = _AttnHost()
+    q = torch.randn(1, N_HEADS, SEQ, HEAD_DIM)
+    k = torch.randn(1, N_KV_HEADS, SEQ, HEAD_DIM)
+    host._run_attn(q, k, k.clone(), None)
+    _, key_shape, _ = host.attn.calls[-1]
+    # [B, S, H, D] after the transpose, so heads sit on dim 2.
+    assert key_shape[2] == N_KV_HEADS, f"K reached the backend with {key_shape[2]} heads"
+
+
+def test_bshd_path_hands_the_backend_compressed_kv():
+    host = _AttnHost()
+    q = torch.randn(1, SEQ, N_HEADS, HEAD_DIM)
+    k = torch.randn(1, SEQ, N_KV_HEADS, HEAD_DIM)
+    host._run_attn_bshd(q, k, k.clone(), None)
+    _, key_shape, _ = host.attn.calls[-1]
+    assert key_shape[2] == N_KV_HEADS, f"K reached the backend with {key_shape[2]} heads"
+
+
+def test_compressed_kv_computes_what_expansion_computed():
+    """The backend must return the same tensor either way, or this is a bugfix
+    dressed up as a perf change."""
+    torch.manual_seed(0)
+    q = torch.randn(1, SEQ, N_HEADS, HEAD_DIM)
+    k = torch.randn(1, SEQ, N_KV_HEADS, HEAD_DIM)
+    v = torch.randn(1, SEQ, N_KV_HEADS, HEAD_DIM)
+    impl = SDPAImpl(
+        num_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        softmax_scale=HEAD_DIM**-0.5,
+        num_kv_heads=N_KV_HEADS,
+    )
+    compressed = impl._forward_impl(q, k, v)
+    expanded = impl._forward_impl(q, k.repeat_interleave(GROUPS, dim=2), v.repeat_interleave(GROUPS, dim=2))
+    torch.testing.assert_close(compressed, expanded, atol=1e-6, rtol=1e-6)
+
+
+class _MaskProbe(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.masks: list[torch.Tensor | None] = []
+
+    def forward(self, hidden_states, **kwargs):
+        self.masks.append(kwargs["attention_mask"]["full_attention"])
+        return hidden_states
+
+
+def _fake_rope(hidden_states, position_ids):
+    """The forward builds the three RoPE tables before the layer loop; the probe
+    ignores them, so any 2-tuple will do."""
+    return hidden_states, hidden_states
+
+
+class _ModelHost:
+    """Carries only the attributes ``SenseNovaU1Model.forward`` touches."""
+
+    forward = SenseNovaU1Model.forward
+
+    def __init__(self, probe):
+        self.layers = [probe]
+        self.norm = torch.nn.Identity()
+        self.norm_mot_gen = torch.nn.Identity()
+        self.rotary_emb = _fake_rope
+        self.rotary_emb_hw = _fake_rope
+
+
+def _run_model(seq_len: int):
+    probe = _MaskProbe()
+    _ModelHost(probe).forward(
+        inputs_embeds=torch.zeros(1, seq_len, 4),
+        indexes=torch.zeros(3, seq_len, dtype=torch.long),
+    )
+    return probe.masks[-1]
+
+
+def test_single_token_decode_carries_no_mask():
+    assert _run_model(1) is None, "decode still builds the all-zeros mask"
+
+
+def test_prefill_still_gets_a_causal_mask():
+    mask = _run_model(3)
+    assert mask is not None and mask.shape == (1, 1, 3, 3)
+    assert mask[0, 0, 0, 0] == 0.0
+    assert torch.isinf(mask[0, 0, 0, 1]) and mask[0, 0, 0, 1] < 0
+    assert mask[0, 0, 2, 1] == 0.0
+
+
+# ---------------------------------------------------------------------------
+# The decode path changes which SDPA kernel runs, so it is not bit-identical to
+# what it replaced. "Different" would be a regression if the new result were
+# worse, so pin the direction against an independent float64 reference. This
+# needs a fused GQA kernel, which only exists on CUDA.
+# ---------------------------------------------------------------------------
+
+
+def _float64_attention(q, k, v):
+    q64 = q.double()
+    k64 = k.double().repeat_interleave(GROUPS, dim=1)
+    v64 = v.double().repeat_interleave(GROUPS, dim=1)
+    scores = (q64 @ k64.transpose(-1, -2)) * (q.shape[-1] ** -0.5)
+    return torch.softmax(scores, dim=-1) @ v64
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for fused GQA")
+@pytest.mark.cuda
+@pytest.mark.L4
+@pytest.mark.parametrize("kv_len", [271, 512, 2048])
+def test_decode_without_the_no_op_mask_is_closer_to_float64(kv_len):
+    """The old decode path expanded K/V and passed an all-zeros mask, which kept
+    SDPA off its fused kernels. Dropping the mask is only defensible if the
+    kernel it unlocks is at least as accurate."""
+    dev, dtype = "cuda", torch.bfloat16
+    heads, dim = 32, 128
+    new_errs: list[float] = []
+    old_errs: list[float] = []
+    trials = 8
+    for t in range(trials):
+        torch.manual_seed(2000 + t)
+        q = torch.randn(1, heads, 1, dim, device=dev, dtype=dtype)
+        k = torch.randn(1, heads // GROUPS, kv_len, dim, device=dev, dtype=dtype)
+        v = torch.randn(1, heads // GROUPS, kv_len, dim, device=dev, dtype=dtype)
+        ref = _float64_attention(q, k, v)
+
+        zeros_mask = torch.zeros(1, 1, 1, kv_len, device=dev, dtype=dtype)
+        old = torch.nn.functional.scaled_dot_product_attention(
+            q,
+            k.repeat_interleave(GROUPS, dim=1),
+            v.repeat_interleave(GROUPS, dim=1),
+            attn_mask=zeros_mask,
+        )
+        new = torch.nn.functional.scaled_dot_product_attention(q, k, v, enable_gqa=True)
+
+        new_errs.append((new.double() - ref).abs().mean().item())
+        old_errs.append((old.double() - ref).abs().mean().item())
+
+    # A single trial can tie either way depending on the GPU and the CUDA build,
+    # so gate on the aggregate: the maskless path must not be worse on average.
+    new_mean = sum(new_errs) / trials
+    old_mean = sum(old_errs) / trials
+    wins = sum(n <= o for n, o in zip(new_errs, old_errs))
+    assert new_mean <= old_mean, (
+        f"kv_len={kv_len}: maskless mean error {new_mean:.4e} is worse than masked {old_mean:.4e} "
+        f"({wins}/{trials} per-trial wins)"
+    )

tests/diffusion/models/sensenova_u1/test_sensenova_u1_distill_lora.py

@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU tests for SenseNova-U1 distilled-LoRA fusion.
+
+The checkpoints ship unfused ``q_proj_mot_gen`` / ``gate_proj`` keys while the
+pipeline runs fused layers, so each delta lands in its own row slice.
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+from safetensors.torch import save_file
+
+from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+HIDDEN, KV, INTER, RANK, ALPHA = 8, 4, 16, 2, 4.0
+
+
+class _Layer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.self_attn = nn.Module()
+        self.self_attn.qkv_proj_mot_gen = nn.Linear(HIDDEN, HIDDEN + 2 * KV, bias=False)
+        self.self_attn.o_proj_mot_gen = nn.Linear(HIDDEN, HIDDEN, bias=False)
+        self.mlp_mot_gen = nn.Module()
+        self.mlp_mot_gen.gate_up_proj = nn.Linear(HIDDEN, 2 * INTER, bias=False)
+        self.mlp_mot_gen.down_proj = nn.Linear(INTER, HIDDEN, bias=False)
+
+
+class _Stub(SenseNovaU1Pipeline):
+    """Only the parameter tree matters here, so skip the real __init__."""
+
+    def __init__(self):
+        nn.Module.__init__(self)
+        self.language_model = nn.Module()
+        self.language_model.model = nn.Module()
+        self.language_model.model.layers = nn.ModuleList([_Layer()])
+
+
+def _triple(out: int, in_: int, fill: float) -> dict[str, torch.Tensor]:
+    return {
+        "lora_down.weight": torch.full((RANK, in_), fill, dtype=torch.float32),
+        "lora_up.weight": torch.full((out, RANK), 1.0, dtype=torch.float32),
+        "alpha": torch.tensor(ALPHA),
+    }
+
+
+def _write(tmp_path, entries: dict[str, dict[str, torch.Tensor]]):
+    flat = {f"{mod}.{leaf}": t for mod, parts in entries.items() for leaf, t in parts.items()}
+    path = tmp_path / "lora.safetensors"
+    save_file(flat, str(path))
+    return str(path)
+
+
+def _prefix(i: int = 0) -> str:
+    return f"language_model.model.layers.{i}"
+
+
+def test_qkv_shards_land_in_their_own_row_slices(tmp_path):
+    pipe = _Stub()
+    qkv = pipe.language_model.model.layers[0].self_attn.qkv_proj_mot_gen.weight
+    with torch.no_grad():
+        qkv.zero_()
+    entries = {
+        f"{_prefix()}.self_attn.q_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0),
+        f"{_prefix()}.self_attn.k_proj_mot_gen": _triple(KV, HIDDEN, 2.0),
+        f"{_prefix()}.self_attn.v_proj_mot_gen": _triple(KV, HIDDEN, 3.0),
+    }
+    pipe.load_lora_weights(_write(tmp_path, entries))
+
+    scale = ALPHA / RANK
+    assert torch.allclose(qkv[:HIDDEN], torch.full((HIDDEN, HIDDEN), scale * RANK * 1.0))
+    assert torch.allclose(qkv[HIDDEN : HIDDEN + KV], torch.full((KV, HIDDEN), scale * RANK * 2.0))
+    assert torch.allclose(qkv[HIDDEN + KV :], torch.full((KV, HIDDEN), scale * RANK * 3.0))
+
+
+def test_gate_and_up_land_in_their_own_halves(tmp_path):
+    pipe = _Stub()
+    gu = pipe.language_model.model.layers[0].mlp_mot_gen.gate_up_proj.weight
+    with torch.no_grad():
+        gu.zero_()
+    entries = {
+        f"{_prefix()}.mlp_mot_gen.gate_proj": _triple(INTER, HIDDEN, 1.0),
+        f"{_prefix()}.mlp_mot_gen.up_proj": _triple(INTER, HIDDEN, 5.0),
+    }
+    pipe.load_lora_weights(_write(tmp_path, entries))
+
+    scale = ALPHA / RANK
+    assert torch.allclose(gu[:INTER], torch.full((INTER, HIDDEN), scale * RANK * 1.0))
+    assert torch.allclose(gu[INTER:], torch.full((INTER, HIDDEN), scale * RANK * 5.0))
+
+
+def test_unfused_targets_are_added_whole(tmp_path):
+    pipe = _Stub()
+    o = pipe.language_model.model.layers[0].self_attn.o_proj_mot_gen.weight
+    before = o.detach().clone()
+    entries = {f"{_prefix()}.self_attn.o_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0)}
+    pipe.load_lora_weights(_write(tmp_path, entries))
+    assert torch.allclose(o - before, torch.full((HIDDEN, HIDDEN), (ALPHA / RANK) * RANK))
+
+
+def test_lora_that_matches_nothing_raises_instead_of_no_op(tmp_path):
+    """Without a loader the worker only warned and generated without the
+    adapter, so a wrong path produced a plausible image and no error."""
+    pipe = _Stub()
+    entries = {f"{_prefix()}.self_attn.does_not_exist": _triple(HIDDEN, HIDDEN, 1.0)}
+    with pytest.raises(ValueError, match="matched no parameter"):
+        pipe.load_lora_weights(_write(tmp_path, entries))
+
+
+def test_shape_mismatch_raises(tmp_path):
+    pipe = _Stub()
+    entries = {f"{_prefix()}.self_attn.o_proj_mot_gen": _triple(HIDDEN + 1, HIDDEN, 1.0)}
+    with pytest.raises(RuntimeError, match="must match the size"):
+        pipe.load_lora_weights(_write(tmp_path, entries))
+
+
+def test_multiple_files_rejected(tmp_path):
+    pipe = _Stub()
+    path = _write(tmp_path, {f"{_prefix()}.self_attn.o_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0)})
+    with pytest.raises(ValueError, match="exactly one"):
+        pipe.load_lora_weights([path, path])
+
+
+class _ShardedLayer(nn.Module):
+    """One rank of a two-way sharded fused QKV, with a vLLM-style weight loader.
+
+    The parameter holds half of each of q, k and v. The loader is handed the whole
+    layer's tensor and keeps only this rank's slice, which is what vLLM does.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(HIDDEN // 2 + KV, HIDDEN))
+        self.seen: list[tuple[str, tuple[int, ...]]] = []
+        self.weight.weight_loader = self._weight_loader
+
+    def _weight_loader(self, param, loaded_weight, shard_id):
+        self.seen.append((shard_id, tuple(loaded_weight.shape)))
+        sizes = {"q": HIDDEN // 2, "k": KV // 2, "v": KV // 2}
+        offsets = {"q": 0, "k": HIDDEN // 2, "v": HIDDEN // 2 + KV // 2}
+        rows = sizes[shard_id]
+        param.data[offsets[shard_id] : offsets[shard_id] + rows] = loaded_weight[:rows]
+
+
+class _ShardedStub(SenseNovaU1Pipeline):
+    def __init__(self):
+        nn.Module.__init__(self)
+        self.language_model = nn.Module()
+        self.language_model.model = nn.Module()
+        layer = nn.Module()
+        layer.self_attn = nn.Module()
+        layer.self_attn.qkv_proj_mot_gen = _ShardedLayer()
+        self.language_model.model.layers = nn.ModuleList([layer])
+
+
+def test_fused_delta_is_sharded_through_the_weight_loader(tmp_path):
+    """Under tensor parallelism the parameter is a rank-local shard while the
+    checkpoint delta spans the whole layer, so the delta has to go through the
+    layer's weight loader. Adding it directly raises a shape error."""
+    pipe = _ShardedStub()
+    layer = pipe.language_model.model.layers[0].self_attn.qkv_proj_mot_gen
+    entries = {
+        f"{_prefix()}.self_attn.q_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0),
+        f"{_prefix()}.self_attn.k_proj_mot_gen": _triple(KV, HIDDEN, 2.0),
+        f"{_prefix()}.self_attn.v_proj_mot_gen": _triple(KV, HIDDEN, 3.0),
+    }
+    pipe.load_lora_weights(_write(tmp_path, entries))
+
+    # the loader saw the whole layer's tensors, not a pre-sharded slice
+    assert layer.seen == [("q", (HIDDEN, HIDDEN)), ("k", (KV, HIDDEN)), ("v", (KV, HIDDEN))]
+
+    scale = ALPHA / RANK * RANK
+    got = layer.weight
+    assert torch.allclose(got[: HIDDEN // 2], torch.full((HIDDEN // 2, HIDDEN), scale * 1.0))
+    assert torch.allclose(got[HIDDEN // 2 : HIDDEN // 2 + KV // 2], torch.full((KV // 2, HIDDEN), scale * 2.0))
+    assert torch.allclose(got[HIDDEN // 2 + KV // 2 :], torch.full((KV // 2, HIDDEN), scale * 3.0))
+
+
+def test_fusion_does_not_retain_the_adapter_tensors(tmp_path):
+    """Fusion is one-way and this pipeline has no unload path, so holding the
+    fp32 state dict would keep about 1.5 GiB alive for the shipped checkpoint."""
+    pipe = _Stub()
+    path = _write(tmp_path, {f"{_prefix()}.self_attn.o_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0)})
+    pipe.load_lora_weights(path, adapter_name="distill")
+
+    kept = pipe.lora_loaded["distill"]
+    assert not isinstance(kept, dict)
+    assert not isinstance(kept, torch.Tensor)
+
+    # the sentinel still makes a second load a no-op
+    before = pipe.language_model.model.layers[0].self_attn.o_proj_mot_gen.weight.detach().clone()
+    pipe.load_lora_weights(path, adapter_name="distill")
+    torch.testing.assert_close(pipe.language_model.model.layers[0].self_attn.o_proj_mot_gen.weight, before)

tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py

@@ -0,0 +1,800 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Tests for the paged decode cache and its CUDA graph.
+
+The cache exists so a decode step has static shapes and can be captured. That
+puts the burden on two things a normal test would not look at: the prefix must
+survive the hand-off from the ordinary cache, and every quantity that varies
+between steps must live in a device tensor. A Python-level index is correct
+eagerly and silently wrong once captured, because capture bakes it in -- which
+is exactly what the last test here pins.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from vllm.lora.layers import BaseLayerWithLoRA
+
+import vllm_omni.diffusion.models.sensenova_u1.paged_decode as paged_decode
+from vllm_omni.diffusion.models.sensenova_u1.paged_decode import (
+    BLOCK_SIZE,
+    BUCKETS,
+    TAIL_STEP,
+    PagedDecodeCache,
+    _bucket_for,
+)
+from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
+    SenseNovaU1Attention,
+)
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+LAYERS, KV_HEADS, HEAD_DIM = 2, 4, 8
+N_HEADS = KV_HEADS * 2
+
+
+class _Layer:
+    def __init__(self, keys=None, values=None):
+        self.keys, self.values = keys, values
+
+
+class _Cache:
+    """Stands in for the transformers cache: layers with [B, H, S, D] tensors."""
+
+    def __init__(self, layers):
+        self.layers = layers
+
+
+def _dyn_cache(prefix, device=None):
+    torch.manual_seed(0)
+    return _Cache(
+        [
+            _Layer(
+                torch.randn(1, KV_HEADS, prefix, HEAD_DIM, device=device),
+                torch.randn(1, KV_HEADS, prefix, HEAD_DIM, device=device),
+            )
+            for _ in range(LAYERS)
+        ]
+    )
+
+
+def test_bucket_rounds_up_and_is_block_aligned():
+    assert _bucket_for(1) == 512
+    assert _bucket_for(512) == 512
+    assert _bucket_for(513) == 1024
+    assert _bucket_for(100_000) % BLOCK_SIZE == 0
+
+
+def test_past_the_last_bucket_the_schedule_grows_in_steps():
+    """Every bucket change reallocates the cache and re-captures the graph, so
+    the tail of the schedule must not step by BLOCK_SIZE. A think edit with two
+    2048x2048 inputs runs to about 9,100 tokens, so this regime is reachable,
+    and a request decodes at most 1024 steps.
+    """
+    last = BUCKETS[-1]
+    assert _bucket_for(last + 1) == last + TAIL_STEP
+    assert _bucket_for(100_000) % BLOCK_SIZE == 0
+    for start in (last + 1, last + TAIL_STEP - 1, last + TAIL_STEP + 7):
+        bucket, grows = _bucket_for(start), 0
+        for n in range(start, start + 1025):
+            if _bucket_for(n) != bucket:
+                bucket = _bucket_for(n)
+                grows += 1
+        assert grows <= 1, f"start {start}: {grows} reallocations over 1024 decode steps"
+
+
+def test_prefix_survives_the_hand_off():
+    dyn = _dyn_cache(37)
+    paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
+    assert paged.length == 37
+    for i in range(LAYERS):
+        stored = paged.k[i].view(-1, KV_HEADS, HEAD_DIM)[:37]
+        torch.testing.assert_close(stored, dyn.layers[i].keys[0].transpose(0, 1))
+
+
+def test_write_back_restores_what_was_handed_over():
+    """Decode is paged but the stage after it reads the ordinary cache."""
+    dyn = _dyn_cache(21)
+    original = [layer.keys.clone() for layer in dyn.layers]
+    paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
+    paged.to_dynamic_cache(dyn)
+    for before, layer in zip(original, dyn.layers):
+        torch.testing.assert_close(layer.keys, before)
+
+
+def test_growing_keeps_the_tokens_already_stored():
+    dyn = _dyn_cache(500)
+    paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
+    assert paged.bucket == 512
+    before = [k.view(-1, KV_HEADS, HEAD_DIM)[:500].clone() for k in paged.k]
+    gen = paged.generation
+    assert paged.grow(513) is True
+    assert paged.bucket == 1024
+    assert paged.generation > gen, "a reallocation must invalidate captured graphs"
+    for old, new in zip(before, paged.k):
+        torch.testing.assert_close(new.view(-1, KV_HEADS, HEAD_DIM)[:500], old)
+
+
+def test_a_failed_grow_leaves_the_cache_on_its_old_buffers(monkeypatch):
+    """A half-applied resize would strand a capture on freed storage.
+
+    ``generation`` is half of the key a captured graph is stored under, so
+    publishing new K without publishing V, the block table, the bucket and the
+    generation leaves the old key selecting a graph whose recorded addresses
+    have just been freed. A later request replays it. Injecting a failure into
+    the last allocation is the cheapest way to reach that window.
+    """
+    dyn = _dyn_cache(500)
+    paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
+    before_k, before_v = list(paged.k), list(paged.v)
+    before_table, before_bucket = paged.block_table, paged.bucket
+    before_generation, before_length = paged.generation, paged.length
+
+    real_zeros = torch.zeros
+    calls = []
+
+    def flaky(*args, **kwargs):
+        calls.append(1)
+        if len(calls) == 2 * LAYERS:  # the last of the two allocations per layer
+            raise RuntimeError("injected allocation failure")
+        return real_zeros(*args, **kwargs)
+
+    monkeypatch.setattr(torch, "zeros", flaky)
+    with pytest.raises(RuntimeError, match="injected"):
+        paged.grow(513)
+    monkeypatch.undo()
+
+    assert [id(t) for t in paged.k] == [id(t) for t in before_k], "K was published on its own"
+    assert [id(t) for t in paged.v] == [id(t) for t in before_v], "V was published on its own"
+    assert paged.block_table is before_table
+    assert paged.bucket == before_bucket
+    assert paged.generation == before_generation, "the graph key moved without new buffers"
+    assert paged.length == before_length
+
+
+def test_length_lives_in_device_tensors():
+    """Both the attended length and the write slot must be tensors.
+
+    A captured graph reads them at replay; a Python int would be frozen at
+    capture time and every replay would attend the wrong span and overwrite the
+    same slot.
+    """
+    dyn = _dyn_cache(10)
+    paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
+    paged.set_length(15)
+    assert torch.is_tensor(paged.seqused) and int(paged.seqused) == 15
+    assert torch.is_tensor(paged.pos) and int(paged.pos) == 14
+
+
+# ---------------------------------------------------------------------------
+# The capture contract. These need CUDA and the bundled flash-attention kernel.
+# ---------------------------------------------------------------------------
+
+cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@cuda_only
+@pytest.mark.cuda
+@pytest.mark.L4
+def test_paged_attention_matches_sdpa_over_the_same_buffer():
+    import torch.nn.functional as F
+
+    from vllm_omni.diffusion.models.sensenova_u1.paged_decode import PagedDecodeCache
+
+    heads, kv_heads, dim, n = 8, 2, 64, 100
+    dev, dt = torch.device("cuda"), torch.bfloat16
+    torch.manual_seed(0)
+    dyn = _Cache(
+        [
+            _Layer(
+                torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
+                torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
+            )
+        ]
+    )
+    paged = PagedDecodeCache.from_dynamic_cache(dyn, 1, dev, dt)
+    paged.set_length(n)
+    q = torch.randn(1, heads, 1, dim, device=dev, dtype=dt)
+    k = torch.zeros(1, kv_heads, 1, dim, device=dev, dtype=dt)
+    scale = dim**-0.5
+    # write the last slot back as itself so attend() is a pure read
+    k[0, :, 0] = dyn.layers[0].keys[0, :, n - 1]
+    v = torch.zeros_like(k)
+    v[0, :, 0] = dyn.layers[0].values[0, :, n - 1]
+    got = paged.attend(0, q, k, v, scale)
+    flat_k = paged.k[0].view(-1, kv_heads, dim)[:n]
+    flat_v = paged.v[0].view(-1, kv_heads, dim)[:n]
+    want = F.scaled_dot_product_attention(
+        q,
+        flat_k.unsqueeze(0).transpose(1, 2).contiguous(),
+        flat_v.unsqueeze(0).transpose(1, 2).contiguous(),
+        enable_gqa=True,
+        scale=scale,
+    ).transpose(1, 2)
+    torch.testing.assert_close(got.float(), want.float(), atol=2e-3, rtol=2e-3)
+
+
+@cuda_only
+@pytest.mark.cuda
+@pytest.mark.L4
+def test_a_captured_graph_follows_a_later_length():
+    """One capture has to serve every length in the bucket.
+
+    If the attended span were a Python value the graph would keep attending the
+    capture-time length, and the run would still look fast.
+    """
+    import torch.nn.functional as F
+
+    from vllm_omni.diffusion.models.sensenova_u1.paged_decode import PagedDecodeCache
+
+    heads, kv_heads, dim, n = 8, 2, 64, 200
+    dev, dt = torch.device("cuda"), torch.bfloat16
+    torch.manual_seed(0)
+    dyn = _Cache(
+        [
+            _Layer(
+                torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
+                torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
+            )
+        ]
+    )
+    paged = PagedDecodeCache.from_dynamic_cache(dyn, 1, dev, dt)
+    q = torch.randn(1, heads, 1, dim, device=dev, dtype=dt)
+    k = torch.zeros(1, kv_heads, 1, dim, device=dev, dtype=dt)
+    v = torch.zeros_like(k)
+    scale = dim**-0.5
+
+    paged.set_length(n)
+    k[0, :, 0] = dyn.layers[0].keys[0, :, n - 1]
+    v[0, :, 0] = dyn.layers[0].values[0, :, n - 1]
+
+    side = torch.cuda.Stream()
+    side.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side):
+        for _ in range(3):
+            paged.attend(0, q, k, v, scale)
+    torch.cuda.current_stream().wait_stream(side)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = paged.attend(0, q, k, v, scale)
+
+    shorter = n - 43
+    paged.set_length(shorter)
+    graph.replay()
+    torch.accelerator.synchronize()
+
+    flat_k = paged.k[0].view(-1, kv_heads, dim)[:shorter]
+    flat_v = paged.v[0].view(-1, kv_heads, dim)[:shorter]
+    want = F.scaled_dot_product_attention(
+        q,
+        flat_k.unsqueeze(0).transpose(1, 2).contiguous(),
+        flat_v.unsqueeze(0).transpose(1, 2).contiguous(),
+        enable_gqa=True,
+        scale=scale,
+    ).transpose(1, 2)
+    torch.testing.assert_close(out.float(), want.float(), atol=2e-3, rtol=2e-3)
+
+
+@cuda_only
+@pytest.mark.cuda
+@pytest.mark.L4
+def test_a_captured_graph_writes_each_step_to_its_own_slot():
+    """The write slot has to be a device tensor, not a Python index.
+
+    ``test_a_captured_graph_follows_a_later_length`` pins the attended span.
+    The slot written is the other half: a Python index is frozen at capture, so
+    every replay would overwrite the capture-time slot and the cache would stop
+    growing -- correct eagerly, silently wrong once captured.
+    """
+    from vllm_omni.diffusion.models.sensenova_u1.paged_decode import PagedDecodeCache
+
+    heads, kv_heads, dim, n = 8, 2, 64, 100
+    dev, dt = torch.device("cuda"), torch.bfloat16
+    torch.manual_seed(0)
+    dyn = _Cache(
+        [
+            _Layer(
+                torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
+                torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
+            )
+        ]
+    )
+    paged = PagedDecodeCache.from_dynamic_cache(dyn, 1, dev, dt)
+    q = torch.randn(1, heads, 1, dim, device=dev, dtype=dt)
+    k = torch.zeros(1, kv_heads, 1, dim, device=dev, dtype=dt)
+    v = torch.zeros_like(k)
+    scale = dim**-0.5
+
+    # Capture on the step that writes slot n-1, restoring what is already there
+    # so the capture itself changes nothing.
+    paged.set_length(n)
+    k[0, :, 0] = dyn.layers[0].keys[0, :, n - 1]
+    v[0, :, 0] = dyn.layers[0].values[0, :, n - 1]
+    side = torch.cuda.Stream()
+    side.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side):
+        for _ in range(3):
+            paged.attend(0, q, k, v, scale)
+    torch.cuda.current_stream().wait_stream(side)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        paged.attend(0, q, k, v, scale)
+
+    kept = paged.k[0].view(-1, kv_heads, dim)[n - 1].clone()
+
+    # The next decode step: one token further, different K/V.
+    paged.set_length(n + 1)
+    k[0, :, 0] = 1.5
+    v[0, :, 0] = -2.5
+    graph.replay()
+    current_omni_platform.synchronize()
+
+    flat_k = paged.k[0].view(-1, kv_heads, dim)
+    flat_v = paged.v[0].view(-1, kv_heads, dim)
+    torch.testing.assert_close(flat_k[n].float(), torch.full_like(flat_k[n], 1.5).float())
+    torch.testing.assert_close(flat_v[n].float(), torch.full_like(flat_v[n], -2.5).float())
+    torch.testing.assert_close(flat_k[n - 1].float(), kept.float())
+
+
+@cuda_only
+@pytest.mark.cuda
+@pytest.mark.L4
+def test_a_whole_decode_loop_matches_the_unpaged_path():
+    """Per-step kernel equivalence is not loop equivalence.
+
+    With the switch off a single-token step attends over a cache grown by
+    ``torch.cat`` and no mask; with it on the same step reads a bucket-sized
+    buffer through ``seqused_k``. Run the same sequence of steps through both,
+    crossing a bucket boundary on the way, and require the same argmax at every
+    step off a fixed projection -- which is the token id the pipeline would emit.
+    """
+    import torch.nn.functional as F
+
+    from vllm_omni.diffusion.models.sensenova_u1.paged_decode import PagedDecodeCache
+
+    heads, kv_heads, dim = 8, 2, 64
+    prefill, steps = 500, 40  # crosses the 512 bucket at step 13
+    dev, dt = torch.device("cuda"), torch.bfloat16
+    scale = dim**-0.5
+    torch.manual_seed(7)
+
+    keys = torch.randn(1, kv_heads, prefill, dim, device=dev, dtype=dt)
+    values = torch.randn(1, kv_heads, prefill, dim, device=dev, dtype=dt)
+    paged = PagedDecodeCache.from_dynamic_cache(_Cache([_Layer(keys, values)]), 1, dev, dt)
+    ref_k, ref_v = keys.clone(), values.clone()
+    proj = torch.randn(dim * heads, 128, device=dev, dtype=torch.float32)
+
+    paged_ids, ref_ids, worst = [], [], 0.0
+    for step in range(steps):
+        q = torch.randn(1, heads, 1, dim, device=dev, dtype=dt)
+        k = torch.randn(1, kv_heads, 1, dim, device=dev, dtype=dt)
+        v = torch.randn(1, kv_heads, 1, dim, device=dev, dtype=dt)
+
+        if paged.length + 1 > paged.bucket:
+            paged.grow(paged.length + 1)
+        paged.set_length(paged.length + 1)
+        got = paged.attend(0, q, k, v, scale)
+
+        ref_k = torch.cat([ref_k, k], dim=2)
+        ref_v = torch.cat([ref_v, v], dim=2)
+        want = F.scaled_dot_product_attention(q, ref_k, ref_v, enable_gqa=True, scale=scale).transpose(1, 2)
+
+        worst = max(worst, (got.float() - want.float()).abs().max().item())
+        paged_ids.append(int((got.float().reshape(1, -1) @ proj).argmax()))
+        ref_ids.append(int((want.float().reshape(1, -1) @ proj).argmax()))
+
+    assert paged.bucket == 1024, "the loop never crossed a bucket boundary"
+    assert paged_ids == ref_ids, (
+        f"argmax diverges at step {next(i for i, (a, b) in enumerate(zip(paged_ids, ref_ids)) if a != b)}"
+    )
+    assert worst < 5e-2, f"largest per-step deviation {worst:.3e}"
+
+
+class _StubLM(torch.nn.Module):
+    """Just enough CUDA work inside the capture for a pool to be allocated."""
+
+    def __init__(self, device):
+        super().__init__()
+        self.w = torch.randn(16, 64, device=device)
+
+    def forward(self, input_ids, indexes, past_key_values=None, use_cache=False, paged_cache=None):
+        x = input_ids.float().expand(1, 16) @ self.w
+        return SimpleNamespace(logits=x.unsqueeze(0))
+
+
+@cuda_only
+@pytest.mark.cuda
+@pytest.mark.L4
+def test_the_decode_graph_captures_into_the_shared_platform_pool(monkeypatch):
+    """``_decode_context`` builds a runner per request, so a private capture pool
+    hands every request its own arena and never gives it back. Measured on a
+    real model that was about 40 MiB of device memory per request, rising
+    linearly over twelve sequential think requests while the same twelve with
+    the switch off held flat.
+    """
+    from vllm.platforms import current_platform
+
+    from vllm_omni.diffusion.models.sensenova_u1.paged_decode import DecodeGraphRunner
+
+    pools = []
+    real_graph = torch.cuda.graph
+
+    def spy(graph, pool=None, **kwargs):
+        pools.append(pool)
+        return real_graph(graph, pool=pool, **kwargs)
+
+    monkeypatch.setattr(torch.cuda, "graph", spy)
+
+    dev = torch.device("cuda")
+    cache = PagedDecodeCache.from_dynamic_cache(_dyn_cache(8, dev), LAYERS, dev, torch.float32)
+    runner = DecodeGraphRunner(_StubLM(dev), cache, dev)
+    runner.step(0, 0)
+
+    assert pools, "no graph was captured"
+    assert pools[0] is not None, "captured into a fresh private pool"
+    assert pools[0] == current_platform.get_global_graph_pool()
+    assert not hasattr(runner, "_pool"), "the runner still keeps a pool of its own"
+
+
+# ---------------------------------------------------------------------------
+# The cache is only worth anything if production actually reaches it, and the
+# kwargs it needs are not in every wheel that exports the kernel. The tests
+# above would all still pass with the `forward_und` branch deleted, so these
+# drive the real dispatch instead of the cache object.
+# ---------------------------------------------------------------------------
+
+
+class _PagedProbe:
+    def __init__(self, out):
+        self.out = out
+        self.calls: list[tuple[int, torch.Size, float]] = []
+
+    def attend(self, layer_idx, query, key, value, scaling):
+        self.calls.append((layer_idx, query.shape, scaling))
+        return self.out
+
+
+class _AttnHost:
+    """Carries only what ``forward_und`` touches on the way to the branch."""
+
+    forward_und = SenseNovaU1Attention.forward_und
+    qkv_proj = q_norm = k_norm = q_norm_hw = k_norm_hw = None
+    layer_idx = 3
+    scaling = HEAD_DIM**-0.5
+
+    def __init__(self, seq, paged_out):
+        self._qkv = (
+            torch.randn(1, N_HEADS, seq, HEAD_DIM),
+            torch.randn(1, KV_HEADS, seq, HEAD_DIM),
+            torch.randn(1, KV_HEADS, seq, HEAD_DIM),
+        )
+        self.paged = _PagedProbe(paged_out)
+        self.sdpa_calls: list[torch.Tensor | None] = []
+
+    def _project_and_rope(self, *args, **kwargs):
+        return self._qkv
+
+    def _run_attn(self, query, key, value, mask):
+        self.sdpa_calls.append(mask)
+        return torch.zeros(1, N_HEADS, query.shape[2], HEAD_DIM)
+
+    def o_proj(self, hidden_states):
+        return hidden_states * 2, None
+
+
+def test_forward_und_sends_a_single_token_decode_to_the_paged_cache():
+    """Goes through the real dispatch, so deleting the branch turns this red."""
+    out = torch.randn(1, N_HEADS, 1, HEAD_DIM)
+    host = _AttnHost(seq=1, paged_out=out)
+    got = host.forward_und(torch.zeros(1, 1, 16), None, None, position_embeddings=None, paged_cache=host.paged)
+    assert len(host.paged.calls) == 1, "the paged hook was never reached"
+    assert host.sdpa_calls == [], "the same step also ran the SDPA path"
+    assert host.paged.calls[0][0] == host.layer_idx
+    torch.testing.assert_close(got, out.reshape(1, 1, -1).contiguous() * 2)
+
+
+@pytest.mark.parametrize(
+    "seq,mask",
+    [(2, None), (1, torch.zeros(1, 1, 1, 1))],
+    ids=["multi_token", "masked"],
+)
+def test_only_an_unmasked_single_token_step_takes_the_paged_path(seq, mask):
+    """Both other shapes are uncapturable, and the cache holds one row per step."""
+    host = _AttnHost(seq=seq, paged_out=torch.zeros(1, N_HEADS, seq, HEAD_DIM))
+    host.forward_und(torch.zeros(1, seq, 16), None, mask, position_embeddings=None, paged_cache=host.paged)
+    assert host.paged.calls == [], "a step that cannot be captured reached the paged cache"
+    assert len(host.sdpa_calls) == 1
+
+
+def test_a_kernel_missing_any_kwarg_we_pass_is_not_supported(monkeypatch):
+    """The call names every argument, so the probe has to cover every name.
+
+    An older wheel exports ``flash_attn_varlen_func`` without the paged kwargs,
+    and a future one may rename the standard ones. Either way the probe is what
+    keeps that install on SDPA instead of raising TypeError at the first decode
+    step, with the feature on by default.
+    """
+
+    def no_paged_kwargs(q, k, v, max_seqlen_q, cu_seqlens_q, max_seqlen_k, cu_seqlens_k, causal=False):
+        raise AssertionError("must not be called")  # pragma: no cover
+
+    def renamed_standard_args(q, k, v, max_q, cu_q, max_k, seqused_k=None, block_table=None, causal=False):
+        raise AssertionError("must not be called")  # pragma: no cover
+
+    for kernel in (no_paged_kwargs, renamed_standard_args):
+        monkeypatch.setattr(paged_decode, "_flash_varlen", lambda kernel=kernel: kernel)
+        assert paged_decode.paged_decode_supported(torch.device("cuda"), HEAD_DIM) is False
+
+    def current_kernel(
+        q,
+        k,
+        v,
+        max_seqlen_q,
+        cu_seqlens_q,
+        max_seqlen_k,
+        seqused_k=None,
+        block_table=None,
+        softmax_scale=None,
+        causal=False,
+    ):
+        raise AssertionError("must not be called")  # pragma: no cover
+
+    monkeypatch.setattr(paged_decode, "_flash_varlen", lambda: current_kernel)
+    assert paged_decode.paged_decode_supported(torch.device("cuda"), HEAD_DIM) is True
+
+
+def test_only_q_k_and_v_reach_the_kernel_positionally(monkeypatch):
+    """A positional standard argument binds to whatever a future wheel puts in
+    that slot, and the probe above only checks that the names exist."""
+    seen = {}
+
+    def recorder(*args, **kwargs):
+        seen["args"], seen["kwargs"] = args, kwargs
+        return torch.zeros(1, N_HEADS, HEAD_DIM)
+
+    monkeypatch.setattr(paged_decode, "_flash_varlen", lambda: recorder)
+    cache = PagedDecodeCache(1, KV_HEADS, HEAD_DIM, 8, torch.device("cpu"), torch.float32)
+    cache.set_length(4)
+    cache.attend(
+        0,
+        torch.zeros(1, N_HEADS, 1, HEAD_DIM),
+        torch.zeros(1, KV_HEADS, 1, HEAD_DIM),
+        torch.zeros(1, KV_HEADS, 1, HEAD_DIM),
+        1.0,
+    )
+    assert len(seen["args"]) == 3, "only q, k and v may be positional"
+    assert set(seen["kwargs"]) == {
+        "max_seqlen_q",
+        "cu_seqlens_q",
+        "max_seqlen_k",
+        "seqused_k",
+        "block_table",
+        "softmax_scale",
+        "causal",
+    }
+
+
+def test_the_kill_switch_returns_no_cache(monkeypatch):
+    """The recipe documents `VLLM_OMNI_SENSENOVA_PAGED_DECODE=0` as the way back
+    to the ordinary cache, so the switch needs a test rather than only prose."""
+    from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import (
+        SenseNovaU1Pipeline,
+    )
+
+    monkeypatch.setenv("VLLM_OMNI_SENSENOVA_PAGED_DECODE", "0")
+    # The env check returns before anything else on the pipeline is touched.
+    host = object.__new__(SenseNovaU1Pipeline)
+    assert SenseNovaU1Pipeline._decode_context(host, object()) is None
+
+
+def _pipeline_host(monkeypatch, device="cpu", language_model=None):
+    """A bare pipeline carrying only what ``_decode_context`` touches."""
+    from vllm_omni.diffusion.models.sensenova_u1 import pipeline_sensenova_u1 as pipe_mod
+
+    monkeypatch.setenv("VLLM_OMNI_SENSENOVA_PAGED_DECODE", "1")
+    monkeypatch.setattr(pipe_mod, "paged_decode_supported", lambda dev, head_dim: True)
+    layer = SimpleNamespace(self_attn=SimpleNamespace(head_dim=HEAD_DIM))
+    host = object.__new__(pipe_mod.SenseNovaU1Pipeline)
+    # The pipeline is an nn.Module, and the language model assigned below is one
+    # too, so the bookkeeping has to exist before the assignment.
+    torch.nn.Module.__init__(host)
+    host.device = device
+    lm = language_model if language_model is not None else torch.nn.Module()
+    lm.model = SimpleNamespace(layers=[layer] * LAYERS)
+    host.language_model = lm
+    return pipe_mod, host
+
+
+def test_the_decode_context_is_reused_across_requests(monkeypatch):
+    """A capture binds the addresses it recorded, so a cache built per request
+    forces a capture per request. Two requests whose prefill fits the same
+    bucket have to come back with the same cache and the same runner, and with
+    the generation unchanged so the captures stay valid.
+    """
+    pipe_mod, host = _pipeline_host(monkeypatch)
+    monkeypatch.setattr(pipe_mod, "DecodeGraphRunner", lambda lm, cache, dev: SimpleNamespace(cache=cache))
+    ctx = pipe_mod.SenseNovaU1Pipeline._decode_context
+
+    first = ctx(host, _dyn_cache(10))
+    assert first is not None
+    second_input = _dyn_cache(20)
+    second = ctx(host, second_input)
+    assert second[0] is first[0], "the cache was rebuilt for a request that fits it"
+    assert second[1] is first[1], "the runner was rebuilt, so its captures are gone"
+    assert second[0].generation == first[0].generation, "reuse must not invalidate a capture"
+    assert second[0].length == 20
+
+    # The reused buffers must carry the new request's prefill, not the old one.
+    flat = second[0].k[0].view(-1, KV_HEADS, HEAD_DIM)[:20]
+    torch.testing.assert_close(flat, second_input.layers[0].keys[0].transpose(0, 1))
+
+    # A prefill past the bucket cannot share the buffers.
+    bigger = ctx(host, _dyn_cache(BUCKETS[0] + 1))
+    assert bigger[0] is not first[0], "a prefill past the bucket must get new buffers"
+
+
+def test_releasing_the_captures_forces_a_rebuild(monkeypatch):
+    """Sleep level 2 discards the memory a capture recorded, so what the
+    pipeline reuses across requests has to be droppable."""
+    pipe_mod, host = _pipeline_host(monkeypatch)
+    monkeypatch.setattr(pipe_mod, "DecodeGraphRunner", lambda lm, cache, dev: SimpleNamespace(cache=cache))
+    ctx = pipe_mod.SenseNovaU1Pipeline._decode_context
+
+    first = ctx(host, _dyn_cache(10))
+    assert ctx(host, _dyn_cache(10))[0] is first[0], "the fixture is not reusing to begin with"
+    pipe_mod.SenseNovaU1Pipeline.release_captured_graphs(host)
+    after = ctx(host, _dyn_cache(10))
+    assert after[0] is not first[0], "the cache survived the release"
+    assert after[1] is not first[1], "the runner survived the release"
+
+
+class _FakeLoRAWrapper(BaseLayerWithLoRA):
+    """Stands in for what ``replace_submodule`` leaves in the module tree."""
+
+    def __init__(self):
+        super().__init__()
+
+
+def test_no_capture_crosses_a_request_once_lora_wrappers_are_installed(monkeypatch):
+    """Base -> LoRA -> base. The reuse test cannot see an adapter change.
+
+    ``DiffusionLoRAManager`` replaces the decode path's linear layers with
+    ``BaseLayerWithLoRA`` and binds, rescales or resets slot 0 per request, and
+    the wrappers stay in the tree afterwards. A graph captured before an adapter
+    was bound would replay without its matmuls; one captured with it bound would
+    survive its removal. Neither moves the shape, dtype or bucket that decide
+    reuse, so the wrappers themselves have to end it.
+    """
+    pipe_mod, host = _pipeline_host(monkeypatch)
+    monkeypatch.setattr(pipe_mod, "DecodeGraphRunner", lambda lm, cache, dev: SimpleNamespace(cache=cache))
+    ctx = pipe_mod.SenseNovaU1Pipeline._decode_context
+
+    base = ctx(host, _dyn_cache(10))
+    assert ctx(host, _dyn_cache(10))[0] is base[0], "the fixture is not reusing to begin with"
+
+    host.language_model.add_module("q_proj", _FakeLoRAWrapper())
+    adapted = ctx(host, _dyn_cache(10))
+    assert adapted[0] is not base[0], "an adapted request reused the base capture"
+    assert adapted[1] is not base[1], "an adapted request reused the base runner"
+
+    back = ctx(host, _dyn_cache(10))
+    assert back[0] is not adapted[0], "a later request reused the adapted capture"
+    assert back[1] is not adapted[1], "a later request reused the adapted runner"
+    assert getattr(host, "_paged_decode", None) is None, "an adapted capture was left on the pipeline"
+
+
+@cuda_only
+@pytest.mark.cuda
+@pytest.mark.L4
+def test_a_second_request_replays_the_first_capture(monkeypatch):
+    """Reuse is only worth anything if the capture survives it."""
+    dev = torch.device("cuda")
+    pipe_mod, host = _pipeline_host(monkeypatch, device="cuda", language_model=_StubLM(dev))
+    ctx = pipe_mod.SenseNovaU1Pipeline._decode_context
+
+    cache, runner = ctx(host, _dyn_cache(10, dev))
+    cache.set_length(cache.length + 1)
+    runner.step(0, 0)
+    assert runner.captures == 1
+
+    cache2, runner2 = ctx(host, _dyn_cache(12, dev))
+    assert cache2 is cache and runner2 is runner
+    cache2.set_length(cache2.length + 1)
+    runner2.step(0, 0)
+    assert runner.captures == 1, "the second request captured the graph again"
+
+
+class _Tok:
+    """The two tokenizer calls ``_generate_text`` makes."""
+
+    EOS = 7
+
+    def convert_tokens_to_ids(self, token):
+        return self.EOS
+
+    def decode(self, ids, skip_special_tokens=False):
+        return " ".join(str(i) for i in ids)
+
+
+class _Out:
+    def __init__(self, logits, past_key_values):
+        self.logits = logits
+        self.past_key_values = past_key_values
+
+
+class _TextHost:
+    """Carries only what ``_generate_text`` touches."""
+
+    def __init__(self, context):
+        self._context = context
+        self.seen: list[object] = []
+        self.tokenizer = _Tok()
+
+    def _decode_context(self, past_key_values):
+        return self._context
+
+    def _ar_step(self, next_token, t_idx, past_key_values, decode=None):
+        self.seen.append(decode)
+        # Stop the loop: argmax of this row is the EOS id.
+        logits = torch.full((1, 1, 8), -1.0)
+        logits[0, 0, _Tok.EOS] = 1.0
+        return _Out(logits, past_key_values)
+
+
+def test_text_decoding_uses_the_paged_context_too():
+    """Image-to-text and text-to-text run their own decode loop. It used to call
+    ``_ar_step`` without the context, so the paged path was documented but never
+    reached outside think."""
+    from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import (
+        SenseNovaU1Pipeline,
+    )
+
+    context = object()
+    host = _TextHost(context)
+    prefix = torch.full((1, 1, 8), -1.0)
+    prefix[0, 0, 5] = 1.0  # not EOS, so one step runs
+    SenseNovaU1Pipeline._generate_text(host, prefix, past_key_values=None, t_idx=0)
+    assert host.seen == [context], f"text decoding ran with decode={host.seen}"
+
+
+class _WarmupReq:
+    """The engine's dummy request, which carries the reserved id."""
+
+    def __init__(self, request_id):
+        self.request_id = request_id
+        self.prompts = [{"prompt": "", "modalities": ["image"]}]
+
+
+def test_the_dummy_warmup_request_drives_one_decode_step(monkeypatch):
+    """The dummy request is a text-to-image run with think off, so it exercises
+    the prefill shape and never the decode one. Without this hook whatever the
+    deploy config compiles for decode lands on the first real request."""
+    from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import (
+        SenseNovaU1Pipeline,
+    )
+    from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID
+
+    calls: list[str] = []
+    host = object.__new__(SenseNovaU1Pipeline)
+    monkeypatch.setattr(host, "_warm_ar_decode", lambda: calls.append("warm"), raising=False)
+    monkeypatch.setattr(
+        SenseNovaU1Pipeline,
+        "_parse_request",
+        lambda self, req: (_ for _ in ()).throw(_StopForwardError()),
+        raising=False,
+    )
+
+    with pytest.raises(_StopForwardError):
+        SenseNovaU1Pipeline.forward(host, _WarmupReq(DUMMY_DIFFUSION_REQUEST_ID))
+    assert calls == ["warm"], "the dummy warmup request did not drive a decode step"
+
+    calls.clear()
+    with pytest.raises(_StopForwardError):
+        SenseNovaU1Pipeline.forward(host, _WarmupReq("a-real-request"))
+    assert calls == [], "a real request paid for the warmup"
+
+
+class _StopForwardError(Exception):
+    """Stops ``forward`` right after the warmup hook, so the test needs no model."""

tests/diffusion/models/sensenova_u1/test_sensenova_u1_rmsnorm.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU tests for SenseNova-U1 RMSNorm and the U1.5 config surface.
+
+``Qwen3RMSNorm`` was a hand-rolled fp32 up/down-cast chain; it now calls the fused
+``F.rms_norm``. These pin the mathematics against an independent reference so a
+future kernel swap cannot silently change what the layer computes, and pin the two
+config fields that distinguish U1.5 from U1.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import Qwen3RMSNorm
+from vllm_omni.transformers_utils.configs.sensenova_u1 import SenseNovaU1Config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _reference(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    """RMSNorm written out in float64 so it is independent of the implementation."""
+    x64 = x.to(torch.float64)
+    var = x64.pow(2).mean(-1, keepdim=True)
+    return (weight.to(torch.float64) * (x64 * torch.rsqrt(var + eps))).to(x.dtype)
+
+
+@pytest.mark.parametrize("hidden", [8, 128, 4096])
+@pytest.mark.parametrize(
+    "dtype,atol",
+    [(torch.float32, 1e-5), (torch.bfloat16, 3e-2), (torch.float16, 4e-3)],
+)
+def test_rmsnorm_matches_float64_reference(hidden: int, dtype: torch.dtype, atol: float):
+    torch.manual_seed(0)
+    norm = Qwen3RMSNorm(hidden, eps=1e-6).to(dtype)
+    with torch.no_grad():
+        norm.weight.copy_(torch.rand(hidden, dtype=dtype) + 0.5)
+    x = torch.randn(7, hidden, dtype=dtype)
+    got = norm(x)
+    want = _reference(x, norm.weight, norm.variance_epsilon)
+    assert got.dtype == dtype
+    assert got.shape == x.shape
+    torch.testing.assert_close(got.to(torch.float32), want.to(torch.float32), atol=atol, rtol=atol)
+
+
+def _cast_chain(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    """The implementation this PR replaced, kept so its accuracy can be compared."""
+    var = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
+    return weight * (x.to(torch.float32) * torch.rsqrt(var + eps)).to(weight.dtype)
+
+
+@pytest.mark.parametrize("hidden", [512, 4096])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_rmsnorm_is_closer_to_float64_than_the_cast_chain(hidden: int, dtype: torch.dtype):
+    """The cast chain rounds to ``dtype`` before the weight multiply and so rounds
+    twice, while ``F.rms_norm`` rounds once. Restoring the cast chain fails here."""
+    torch.manual_seed(0)
+    norm = Qwen3RMSNorm(hidden, eps=1e-6).to(dtype)
+    with torch.no_grad():
+        norm.weight.copy_(torch.rand(hidden, dtype=dtype) + 0.5)
+    x = torch.randn(64, hidden, dtype=dtype)
+
+    reference = _reference(x, norm.weight, norm.variance_epsilon).to(torch.float64)
+    scale = reference.abs().mean()
+
+    def relative_error(out: torch.Tensor) -> float:
+        return ((out.to(torch.float64) - reference).abs().mean() / scale).item()
+
+    fused = relative_error(norm(x))
+    legacy = relative_error(_cast_chain(x, norm.weight, norm.variance_epsilon))
+    assert fused < legacy, f"{dtype} hidden={hidden}: fused {fused:.3e} not better than {legacy:.3e}"
+
+
+def test_rmsnorm_is_scale_equivariant():
+    """RMSNorm divides by the RMS, so scaling the input must not change the output."""
+    torch.manual_seed(0)
+    norm = Qwen3RMSNorm(64, eps=1e-12)
+    x = torch.randn(4, 64)
+    torch.testing.assert_close(norm(x), norm(x * 7.5), atol=1e-4, rtol=1e-4)
+
+
+def test_rmsnorm_preserves_leading_dims():
+    norm = Qwen3RMSNorm(16, eps=1e-6)
+    for shape in [(16,), (3, 16), (2, 3, 16)]:
+        assert norm(torch.randn(*shape)).shape == shape
+
+
+def test_rmsnorm_weight_is_applied():
+    norm = Qwen3RMSNorm(4, eps=1e-6)
+    x = torch.randn(2, 4)
+    base = norm(x)
+    with torch.no_grad():
+        norm.weight.mul_(3.0)
+    torch.testing.assert_close(norm(x), base * 3.0, atol=1e-5, rtol=1e-5)
+
+
+def test_u15_config_fields_default_to_u1_values():
+    """U1.5 flips exactly these two; the defaults must stay U1's so old checkpoints
+    keep their behaviour when the keys are absent."""
+    cfg = SenseNovaU1Config()
+    assert cfg.use_pixel_head is False
+    assert cfg.noise_scale_max_value == 10.0
+
+
+def test_u15_config_fields_are_read_from_kwargs():
+    cfg = SenseNovaU1Config(use_pixel_head=True, noise_scale_max_value=16.0)
+    assert cfg.use_pixel_head is True
+    assert cfg.noise_scale_max_value == 16.0

tests/diffusion/test_diffusion_model_runner.py

@@ -179,6 +179,30 @@ def test_request_scoped_cache_dit_lifecycle_is_pipeline_opt_in():
     assert events == [backend]
 
 
+def test_release_captured_graphs_clears_runners_and_delegates_to_the_pipeline():
+    """Both halves in one place, so a pipeline that keeps captures is collected
+    without the caller knowing which pipelines have one."""
+    released = []
+    runner = object.__new__(DiffusionModelRunner)
+    runner.graph_runners = {"decode": object()}
+    runner.pipeline = SimpleNamespace(release_captured_graphs=lambda: released.append(True))
+
+    runner.release_captured_graphs()
+
+    assert runner.graph_runners == {}
+    assert released == [True]
+
+
+def test_release_captured_graphs_tolerates_a_pipeline_without_captures():
+    """Most pipelines keep none, and the runner may not carry `graph_runners`."""
+    runner = object.__new__(DiffusionModelRunner)
+    runner.pipeline = object()
+
+    runner.release_captured_graphs()
+
+    assert not hasattr(runner, "graph_runners")
+
+
 def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summary: bool = True):
     runner = object.__new__(DiffusionModelRunner)
     runner.vllm_config = object()

tests/diffusion/test_diffusion_worker.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 Unit tests for DiffusionWorker class.
@@ -143,6 +143,21 @@ class TestDiffusionWorkerSleep:
         assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
         assert "buffer2" in mock_gpu_worker._sleep_saved_buffers
 
+    def test_sleep_level_2_releases_captures_through_the_model_runner(self, mocker: MockerFixture, mock_gpu_worker):
+        """Level 2 discards the memory captures were recorded against, and the
+        runner owns the release. The worker asks it once and does not reach past
+        it into the pipeline."""
+        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
+        mock_platform.get_device_total_memory.return_value = 80 * 1024**3
+        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory", side_effect=[0, 0])
+        mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(return_value=[])
+
+        mock_gpu_worker.sleep(level=2)
+
+        mock_gpu_worker.model_runner.release_captured_graphs.assert_called_once_with()
+        mock_gpu_worker.model_runner.pipeline.release_captured_graphs.assert_not_called()
+
     def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker):
         """Test that sleep validates memory was actually freed."""
         mock_allocator_class = patch_cumem_allocator(mocker)

tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py

@@ -22,7 +22,7 @@ from vllm.inputs import TokensPrompt
 
 from tests.helpers.clean import wait_for_gpu_memory_to_clear
 from tests.helpers.mark import hardware_test
-from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.platforms import current_omni_platform
 
@@ -52,26 +52,7 @@ def _colocate_async_deploy() -> str:
     CI yaml plus CuMem on every stage fails engine-core startup in the L3
     merge job: Stage 1 resolves ``devices: "1"`` while only GPU 0 is visible.
     """
-    return modify_stage_config(
+    return get_deploy_config_path("ci/qwen3_omni_moe_colocate_async.yaml")
-        get_deploy_config_path("qwen3_omni_moe_thinking.yaml"),
-        updates={
-            "stages": {
-                0: {
-                    "enable_sleep_mode": True,
-                    "enforce_eager": True,
-                    "trust_remote_code": True,
-                    "tensor_parallel_size": 1,
-                    "devices": "0",
-                    "gpu_memory_utilization": 0.9,
-                    "max_num_seqs": 1,
-                    "max_model_len": 8192,
-                    "max_num_batched_tokens": 8192,
-                    "enable_prefix_caching": False,
-                    "skip_mm_profiling": True,
-                },
-            },
-        },
-    )
 
 
 @pytest.mark.advanced_model

tests/e2e/online_serving/minimax_h3/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""MiniMax online-serving tests."""

tests/e2e/online_serving/minimax_h3/_common.py

@@ -0,0 +1,369 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Shared request and server helpers for MiniMax-H3 L3 tests."""
+
+from __future__ import annotations
+
+import base64
+import concurrent.futures
+import io
+import json
+import os
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import av
+import requests
+
+from tests.helpers.assertions import assert_video_valid
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServerParams, OpenAIClientHandler
+
+# DLO starts worker processes and uses shared-memory queues.  Always select
+# spawn for these cases so a parent fork cannot inherit stale resource-tracker
+# state from a preceding server lifecycle.
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = os.environ.get("VLLM_TEST_MINIMAX_H3_MODEL", "MiniMaxAI/MiniMax-H3")
+WIDTH = 1344
+HEIGHT = 768
+FPS = 24
+NUM_INFERENCE_STEPS = 4
+# Keep a failed collective/preprocessing regression bounded well below the
+# per-file Buildkite timeout.  Healthy H100 runs complete in under two minutes
+# after startup; ten minutes leaves ample room for a cold compile.
+REQUEST_TIMEOUT_SECONDS = 600
+
+# FastH3's four-step preview is a load-time fusion adapter.  Its model-level
+# deltas cannot use the dynamic PEFT LoRA manager, so this test deliberately
+# skips on branches that predate PR #6714 instead of starting a base-H3 server.
+_REPO_ROOT = Path(__file__).resolve().parents[4]
+FASTH3_SUPPORTED = (_REPO_ROOT / "vllm_omni/diffusion/models/minimax_h3/fasth3.py").is_file()
+# The Turbo artifact is handled by the MiniMax-specific loader introduced in
+# PR #6550; the generic diffusion LoRA manager cannot parse that checkpoint.
+TURBO_SUPPORTED = (_REPO_ROOT / "vllm_omni/diffusion/models/minimax_h3/lora.py").is_file()
+FASTH3_WIDTH = 1024
+FASTH3_HEIGHT = 576
+FASTH3_DURATION = 4.4
+
+DLO_SERVER_ARGS = [
+    "--trust-remote-code",
+    "--num-gpus",
+    "2",
+    "--tensor-parallel-size",
+    "1",
+    "--data-parallel-size",
+    "2",
+    "--request-batch-max-wait-ms",
+    "500",
+    "--usp",
+    "1",
+    "--ring",
+    "1",
+    "--text-encoder-tp-size",
+    "1",
+    "--vae-patch-parallel-size",
+    "1",
+    "--vae-parallel-mode",
+    "tile",
+    "--vae-use-tiling",
+    "--enable-distributed-layerwise-offload",
+]
+
+_FL2VA_IMAGE = base64.b64decode(generate_synthetic_image(WIDTH, HEIGHT, seed=42)["base64"])
+# Keep the DLO/DP2 Ref2VA wave on the image-only path.  The separate video
+# reference path performs per-request ffmpeg/Qwen video preprocessing and has
+# repeatedly stalled this merge wave before the first pipeline log; it remains
+# covered by the Ref2VA contract/accuracy suites.
+_REF2VA_IMAGE = base64.b64decode(generate_synthetic_image(512, 288, seed=43)["base64"])
+
+
+def resolve_turbo_lora() -> str | None:
+    """Resolve the legacy Turbo LoRA artifact without network at collection."""
+    configured = os.environ.get("VLLM_TEST_MINIMAX_H3_TURBO_LORA")
+    if configured and Path(configured).is_file():
+        return configured
+
+    try:
+        from huggingface_hub import hf_hub_download
+
+        return hf_hub_download(
+            repo_id="lightx2v/Minimax-h3-Turbo",
+            filename="minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors",
+            local_files_only=True,
+        )
+    except Exception:
+        return None
+
+
+def resolve_fasth3_lora() -> str | None:
+    """Resolve the dense/data-free FastH3 artifact without network at collection."""
+    configured = os.environ.get("VLLM_TEST_MINIMAX_H3_FASTH3_LORA")
+    if configured and Path(configured).is_file():
+        return configured
+
+    try:
+        from huggingface_hub import hf_hub_download
+
+        return hf_hub_download(
+            repo_id="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA",
+            filename="dense-datafree/adapter_model.safetensors",
+            local_files_only=True,
+        )
+    except Exception:
+        # Buildkite downloads the artifact before pytest.  Local collection
+        # must stay offline-friendly when that optional cache is absent.
+        return None
+
+
+FASTH3_LORA = resolve_fasth3_lora()
+TURBO_LORA = resolve_turbo_lora()
+
+
+def _assert_audio_stream_present(video: bytes) -> None:
+    """Assert that the generated MP4 contains decodable audio samples."""
+    with av.open(io.BytesIO(video)) as container:
+        audio_streams = [stream for stream in container.streams if stream.type == "audio"]
+        assert audio_streams, "MiniMax-H3 MP4 has no audio stream"
+        audio_frame = next(container.decode(audio=0), None)
+        assert audio_frame is not None and audio_frame.samples > 0, "MiniMax-H3 MP4 audio stream is empty"
+
+
+def assert_h3_video(video: bytes, *, width: int, height: int) -> None:
+    """Check the common H3 video/audio response contract."""
+    assert_video_valid(video, width=width, height=height, fps=FPS)
+    _assert_audio_stream_present(video)
+
+
+def post_sync(
+    client: OpenAIClientHandler,
+    form_data: dict[str, str],
+    files: Any = None,
+) -> bytes:
+    """Submit one synchronous video request and return its MP4 body."""
+    response = requests.post(
+        f"{client.base_url.rstrip('/')}/v1/videos/sync",
+        data=form_data,
+        files=files,
+        headers={"Accept": "video/mp4"},
+        timeout=REQUEST_TIMEOUT_SECONDS,
+    )
+    response.raise_for_status()
+    assert response.headers.get("content-type", "").startswith("video/mp4")
+    assert response.content, "MiniMax-H3 returned an empty video body"
+    return response.content
+
+
+def _h3_form(task: str, seed: int) -> dict[str, str]:
+    return {
+        "model": MODEL,
+        "prompt": (
+            "A cinematic live-action scene with a clear subject moving naturally; "
+            "the atmosphere includes synchronized environmental sound."
+        ),
+        "width": str(WIDTH),
+        "height": str(HEIGHT),
+        "fps": str(FPS),
+        "num_inference_steps": str(NUM_INFERENCE_STEPS),
+        "flow_shift": "12",
+        "seed": str(seed),
+        "extra_params": json.dumps(
+            {
+                "task": task,
+                "duration": 4.0,
+                "aspect_ratio": "16:9",
+                "audio_flow_shift": 3.0,
+            },
+            separators=(",", ":"),
+        ),
+    }
+
+
+def run_fl2va(client: OpenAIClientHandler, seed: int) -> bytes:
+    return post_sync(
+        client,
+        _h3_form("fl2va", seed),
+        files=fl2va_files(),
+    )
+
+
+def run_ref2va(client: OpenAIClientHandler, seed: int) -> bytes:
+    return post_sync(
+        client,
+        _h3_form("ref2va", seed),
+        files=[("input_references", ("reference.jpg", io.BytesIO(_REF2VA_IMAGE), "image/jpeg"))],
+    )
+
+
+def fl2va_files() -> dict[str, tuple[str, io.BytesIO, str]]:
+    """Return a fresh multipart first-frame payload for an FL2VA request."""
+    return {"input_reference": ("first_frame.jpg", io.BytesIO(_FL2VA_IMAGE), "image/jpeg")}
+
+
+def run_dlo_wave(
+    client: OpenAIClientHandler,
+    request_fn: Callable[[OpenAIClientHandler, int], bytes],
+) -> list[bytes]:
+    """Run one complete two-request DLO/DP2 wave concurrently."""
+    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+        return list(executor.map(request_fn, (client, client), (2101, 2102)))
+
+
+def dlo_params(task_type: str) -> OmniServerParams:
+    return OmniServerParams(
+        model=MODEL,
+        server_args=[*DLO_SERVER_ARGS, "--task-type", task_type],
+        stage_init_timeout=1800,
+        init_timeout=1800,
+    )
+
+
+def turbo_params(lora_path: str | None) -> OmniServerParams:
+    """Build the two-card TP2 + DLO-no-AllGather Turbo layout."""
+    return OmniServerParams(
+        model=MODEL,
+        server_args=[
+            "--trust-remote-code",
+            "--task-type",
+            "fl2va",
+            "--num-gpus",
+            "2",
+            "--tensor-parallel-size",
+            "2",
+            "--usp",
+            "1",
+            "--ring",
+            "1",
+            "--text-encoder-tp-size",
+            "2",
+            "--vae-patch-parallel-size",
+            "2",
+            "--vae-parallel-mode",
+            "tile",
+            "--vae-use-tiling",
+            "--enable-distributed-layerwise-offload",
+            "--dlo-no-use-allgather",
+            "--lora-backend",
+            "peft",
+            "--lora-path",
+            lora_path or "",
+        ],
+        stage_init_timeout=1800,
+        init_timeout=1800,
+    )
+
+
+def turbo_form(seed: int) -> dict[str, str]:
+    """Build the official five-sigma-point Turbo FL2VA request."""
+    return {
+        "model": MODEL,
+        "prompt": (
+            "A man stands beside a yellow car at night; the car drives away "
+            "as he follows it with his eyes and begins singing sadly."
+        ),
+        "width": str(WIDTH),
+        "height": str(HEIGHT),
+        "fps": str(FPS),
+        # The public Turbo artifact has four denoiser evaluations but its API
+        # contract requests five sigma points.
+        "num_inference_steps": "5",
+        "flow_shift": "6",
+        "seed": str(seed),
+        "extra_params": json.dumps(
+            {"task": "fl2va", "duration": FASTH3_DURATION, "aspect_ratio": "16:9", "audio_flow_shift": 3.0},
+            separators=(",", ":"),
+        ),
+        "lora": json.dumps(
+            {"name": "h3-turbo-v1.0", "path": TURBO_LORA, "scale": 1.0},
+            separators=(",", ":"),
+        ),
+    }
+
+
+def fasth3_params(lora_path: str | None) -> OmniServerParams:
+    """Build the requested two-card FastH3 parallel layout.
+
+    HSDP owns DiT sharding (the model tensor-parallel degree stays at its
+    default 1), USP2 shards the sequence, text-encoder TP2 shards Qwen3-VL,
+    and VAE patch parallel 2 is the project's VPP2 setting.
+    """
+    return OmniServerParams(
+        model=MODEL,
+        server_args=[
+            "--trust-remote-code",
+            "--task-type",
+            "fl2va",
+            "--num-gpus",
+            "2",
+            "--usp",
+            "2",
+            "--ring",
+            "1",
+            "--use-hsdp",
+            "--hsdp-shard-size",
+            "2",
+            "--text-encoder-tp-size",
+            "2",
+            "--vae-patch-parallel-size",
+            "2",
+            "--vae-parallel-mode",
+            "tile",
+            "--vae-use-tiling",
+            "--lora-path",
+            lora_path or "",
+        ],
+        stage_init_timeout=1800,
+        init_timeout=1800,
+    )
+
+
+def fasth3_form(seed: int) -> dict[str, str]:
+    """Build the four-step T2VA request accepted by the FastH3 preview."""
+    return {
+        "model": MODEL,
+        "prompt": (
+            "A cinematic live-action scene with a clear subject moving naturally; "
+            "the atmosphere includes synchronized environmental sound."
+        ),
+        "width": str(FASTH3_WIDTH),
+        "height": str(FASTH3_HEIGHT),
+        "fps": str(FPS),
+        "num_inference_steps": "4",
+        "seed": str(seed),
+        # FastH3 is fused at startup, so a request-level `lora` field is both
+        # unnecessary and rejected by the load-time adapter contract.
+        "extra_params": json.dumps(
+            {
+                "task": "t2va",
+                "duration": FASTH3_DURATION,
+                "aspect_ratio": "16:9",
+            },
+            separators=(",", ":"),
+        ),
+    }
+
+
+__all__ = [
+    "FASTH3_LORA",
+    "FASTH3_SUPPORTED",
+    "FASTH3_DURATION",
+    "FASTH3_HEIGHT",
+    "FASTH3_WIDTH",
+    "FPS",
+    "MODEL",
+    "assert_h3_video",
+    "dlo_params",
+    "fasth3_form",
+    "fasth3_params",
+    "fl2va_files",
+    "post_sync",
+    "run_dlo_wave",
+    "run_fl2va",
+    "run_ref2va",
+    "TURBO_LORA",
+    "TURBO_SUPPORTED",
+    "turbo_form",
+    "turbo_params",
+]

tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_fl2va.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""MiniMax-H3 FL2VA DLO + DP2 L3 online-serving case."""
+
+from __future__ import annotations
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OpenAIClientHandler
+
+from ._common import assert_h3_video, dlo_params, run_dlo_wave, run_fl2va
+
+pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow]
+H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    [pytest.param(dlo_params("fl2va"), id="minimax_h3_dlo_dp2_fl2va", marks=H100_TWO_CARD_MARKS)],
+    indirect=True,
+)
+def test_minimax_h3_dlo_dp2_fl2va(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    """Validate a complete concurrent DLO/DP2 wave with two FL2VA jobs."""
+    for video in run_dlo_wave(openai_client, run_fl2va):
+        assert_h3_video(video, width=1344, height=768)

tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_ref2va.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""MiniMax-H3 Ref2VA DLO + DP2 L3 online-serving case."""
+
+from __future__ import annotations
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OpenAIClientHandler
+
+from ._common import assert_h3_video, dlo_params, run_dlo_wave, run_ref2va
+
+pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow]
+H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    [pytest.param(dlo_params("ref2va"), id="minimax_h3_dlo_dp2_ref2va", marks=H100_TWO_CARD_MARKS)],
+    indirect=True,
+)
+def test_minimax_h3_dlo_dp2_ref2va(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    """Validate a complete concurrent DLO/DP2 wave with two image-only Ref2VA jobs."""
+    for video in run_dlo_wave(openai_client, run_ref2va):
+        assert_h3_video(video, width=1344, height=768)

tests/e2e/online_serving/minimax_h3/test_minimax_h3_fasth3_lora.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""MiniMax-H3 FastH3 four-step LoRA L3 online-serving case."""
+
+from __future__ import annotations
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OpenAIClientHandler
+
+from ._common import (
+    FASTH3_HEIGHT,
+    FASTH3_LORA,
+    FASTH3_SUPPORTED,
+    FASTH3_WIDTH,
+    assert_h3_video,
+    fasth3_form,
+    fasth3_params,
+    post_sync,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow]
+H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+
+@pytest.mark.skipif(not FASTH3_SUPPORTED, reason="requires FastH3 support from MiniMax-H3 PR #6714")
+@pytest.mark.skipif(
+    FASTH3_LORA is None,
+    reason="set VLLM_TEST_MINIMAX_H3_FASTH3_LORA or populate the local HF cache",
+)
+@pytest.mark.parametrize(
+    "omni_server",
+    [
+        pytest.param(
+            fasth3_params(FASTH3_LORA),
+            id="minimax_h3_fasth3_lora_hsdp2_usp2_vpp2",
+            marks=H100_TWO_CARD_MARKS,
+        )
+    ],
+    indirect=True,
+)
+def test_minimax_h3_fasth3_lora_hsdp2_usp2_vpp2(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+) -> None:
+    """Run the dense/data-free FastH3 adapter on two H100-class cards.
+
+    The preview adapter is T2VA-only and must be fused through ``--lora-path``;
+    it is not a request-switchable PEFT adapter. 1024x576 is intentionally
+    below the 1344x768 H3 reference shape to stay inside an 80 GiB H100 after
+    regional compile cache growth.
+    """
+    video = post_sync(openai_client, fasth3_form(seed=3101))
+    assert_h3_video(video, width=FASTH3_WIDTH, height=FASTH3_HEIGHT)

tests/e2e/online_serving/minimax_h3/test_minimax_h3_turbo_lora.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""MiniMax-H3 Turbo LoRA + DLO L3 online-serving case."""
+
+from __future__ import annotations
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OpenAIClientHandler
+
+from ._common import (
+    TURBO_LORA,
+    TURBO_SUPPORTED,
+    assert_h3_video,
+    fl2va_files,
+    post_sync,
+    turbo_form,
+    turbo_params,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow]
+H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+
+@pytest.mark.skipif(
+    not TURBO_SUPPORTED,
+    reason="requires the MiniMax-H3 Turbo loader from PR #6550",
+)
+@pytest.mark.skipif(
+    TURBO_LORA is None,
+    reason="set VLLM_TEST_MINIMAX_H3_TURBO_LORA or populate the local HF cache",
+)
+@pytest.mark.parametrize(
+    "omni_server",
+    [
+        pytest.param(
+            turbo_params(TURBO_LORA),
+            id="minimax_h3_dlo_turbo_lora_tp2",
+            marks=H100_TWO_CARD_MARKS,
+        )
+    ],
+    indirect=True,
+)
+def test_minimax_h3_dlo_turbo_lora_fl2va(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+) -> None:
+    """Validate the official four-evaluation Turbo LoRA + DLO FL2VA request."""
+    video = post_sync(
+        openai_client,
+        turbo_form(seed=2201),
+        files=fl2va_files(),
+    )
+    assert_h3_video(video, width=1344, height=768)

tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py

@@ -138,6 +138,23 @@ def test_voice_003(omni_server, online_client) -> None:
     online_client.send_audio_speech_request(request_config)
 
 
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_sample_rate_001(omni_server, online_client) -> None:
+    """The speech API resamples Qwen3-TTS output from native 24 kHz to 8 kHz."""
+    request_config = {
+        "model": omni_server.model,
+        "input": "This response should be encoded as eight kilohertz audio.",
+        "stream": False,
+        "response_format": "wav",
+        "sample_rate": 8000,
+        "task_type": "CustomVoice",
+        "voice": "vivian",
+    }
+
+    online_client.send_audio_speech_request(request_config)
+
+
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
 def test_language_001(omni_server, online_client) -> None:

tests/engine/test_async_omni_engine_stage_init.py

@@ -11,7 +11,7 @@ import types
 
 import pytest
 
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, normalize_omni_diffusion_kwargs
 from vllm_omni.engine import async_omni_engine as async_omni_engine_module
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.stage_init_utils import (
@@ -1795,6 +1795,41 @@ def test_resolve_stage_configs_preserves_stage_diffusion_attention(monkeypatch):
     assert stage_configs[0].engine_args.diffusion_attention_config is existing_attention
 
 
+def test_resolve_stage_configs_does_not_inject_over_stage_diffusion_attention_backend(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    stage_cfg = types.SimpleNamespace(
+        stage_type="diffusion",
+        engine_args=types.SimpleNamespace(
+            diffusion_attention_backend="TORCH_SDPA",
+            diffusion_attention_config=None,
+            lora_path=None,
+            lora_scale=None,
+            enable_sleep_mode=None,
+            quantization_config=None,
+        ),
+    )
+
+    monkeypatch.setattr(
+        engine_mod,
+        "load_and_resolve_stage_configs",
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
+    )
+
+    _config_path, stage_configs = engine._resolve_stage_configs(
+        model="dummy-model",
+        kwargs={"diffusion_attention_backend": "FLASH_ATTN"},
+        trust_remote_code=False,
+    )
+
+    engine_args = stage_configs[0].engine_args
+    assert engine_args.diffusion_attention_backend == "TORCH_SDPA"
+    assert engine_args.diffusion_attention_config is None
+    normalized = normalize_omni_diffusion_kwargs(vars(engine_args))
+    assert normalized["diffusion_attention_config"].default.backend == "TORCH_SDPA"
+
+
 def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stage(monkeypatch):
     import vllm_omni.engine.async_omni_engine as engine_mod
 

tests/engine/test_stage_engine_args.py

@@ -133,7 +133,7 @@ def _stable_engine_arg_environment(monkeypatch, tmp_path):
 
     for worker_type in ("ar", "generation"):
         worker_module = types.ModuleType(f"test.{worker_type}")
-        worker_module.Worker = _TestWorker
+        setattr(worker_module, "Worker", _TestWorker)
         monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
 
 
@@ -434,6 +434,63 @@ def test_typed_diffusion_engine_args_use_structured_diffusion_config(tmp_path):
     assert typed_args["diffusion_attention_config"].per_role["cross"].backend == "TORCH_SDPA"
 
 
+def test_engine_args_consume_stage_diffusion_attention_shorthand(tmp_path):
+    """A stage-level ``diffusion_attention_backend`` shorthand is folded into the structured config."""
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    deploy.stages[2] = replace(
+        deploy.stages[2],
+        diffusion_attention_config=None,
+        diffusion_attention_backend="TORCH_SDPA",
+    )
+    legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[2], model)
+    typed_args = build_engine_args_dict_from_omni_stage_config(
+        omni_config.stage_by_id(2),
+        model,
+    )
+
+    for engine_args in (legacy_args, typed_args):
+        assert engine_args.get("diffusion_attention_backend") is None
+        assert isinstance(engine_args["diffusion_attention_config"], AttentionConfig)
+        assert engine_args["diffusion_attention_config"].default.backend == "TORCH_SDPA"
+        od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
+        assert od_config.diffusion_attention_config.default.backend == "TORCH_SDPA"
+
+
+@pytest.mark.parametrize(
+    "yaml_attention_config",
+    [
+        {"default": {"backend": "FLASH_ATTN"}, "per_role": {"cross": {"backend": "SAGE_ATTN"}}},
+        {"per_role": {"cross": {"backend": "SAGE_ATTN"}}},
+    ],
+    ids=["yaml-default", "yaml-per-role-only"],
+)
+def test_engine_args_apply_cli_attention_shorthand_over_yaml_config(tmp_path, yaml_attention_config):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    deploy.stages[2] = replace(deploy.stages[2], diffusion_attention_config=yaml_attention_config)
+    legacy_stages, omni_config = _legacy_and_typed_stages(
+        pipeline,
+        deploy,
+        model,
+        cli_overrides={"stage_2_diffusion_attention_backend": "TORCH_SDPA"},
+    )
+
+    legacy_args = build_legacy_engine_args_dict(legacy_stages[2], model)
+    typed_args = build_engine_args_dict_from_omni_stage_config(
+        omni_config.stage_by_id(2),
+        model,
+    )
+
+    for engine_args in (legacy_args, typed_args):
+        assert engine_args.get("diffusion_attention_backend") is None
+        attention_config = engine_args["diffusion_attention_config"]
+        assert attention_config.default.backend == "TORCH_SDPA"
+        assert attention_config.per_role["cross"].backend == "SAGE_ATTN"
+        od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
+        assert od_config.diffusion_attention_config.default.backend == "TORCH_SDPA"
+
+
 def test_typed_engine_args_preserve_explicit_backend_default_overrides(tmp_path):
     pipeline, deploy, model = _engine_arg_inputs(tmp_path)
     deploy.enable_prefix_caching = False

tests/entrypoints/openai_api/test_audio_format.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for audio output format handling in chat completions.
 
 Covers:
@@ -11,10 +12,15 @@ Covers:
 
 from __future__ import annotations
 
+from io import BytesIO
+
 import numpy as np
 import pytest
+import soundfile
+import torch
+import torchaudio
 
-from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
+from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin, StreamingAudioResampler
 from vllm_omni.entrypoints.openai.protocol.audio import (
     DEFAULT_AUDIO_FORMAT,
     SUPPORTED_AUDIO_FORMATS,
@@ -122,6 +128,81 @@ class TestCreateAudio:
         decoded = base64.b64decode(response.audio_data)
         assert decoded[:4] == b"RIFF"
 
+    def test_resamples_wav_to_requested_output_rate(self, mixin, audio_tensor):
+        response = mixin.create_audio(
+            CreateAudio(
+                audio_tensor=audio_tensor,
+                sample_rate=24000,
+                output_sample_rate=8000,
+                response_format="wav",
+                speed=1.0,
+                base64_encode=False,
+            )
+        )
+
+        with soundfile.SoundFile(BytesIO(response.audio_data)) as audio_file:
+            assert audio_file.samplerate == 8000
+            assert audio_file.frames == 8000
+
+
+class TestStreamingAudioResampler:
+    @staticmethod
+    def _resample(waveform):
+        resampler = StreamingAudioResampler(24000, 8000)
+        return np.concatenate(
+            (
+                resampler.process(waveform),
+                resampler.process(np.empty(0), final=True),
+            )
+        )
+
+    def test_chunk_boundaries_do_not_change_output(self):
+        waveform = np.sin(np.linspace(0, 200 * np.pi, 24000, endpoint=False)).astype(np.float32)
+
+        expected = self._resample(waveform)
+
+        chunked = StreamingAudioResampler(24000, 8000)
+        pieces = [chunked.process(chunk) for chunk in np.split(waveform, [137, 2048, 9001, 17003])]
+        pieces.append(chunked.process(np.empty(0), final=True))
+        actual = np.concatenate(pieces)
+
+        assert actual.shape == (8000,)
+        np.testing.assert_allclose(actual, expected, atol=1e-6)
+
+        reference = torchaudio.functional.resample(torch.from_numpy(waveform), 24000, 8000).numpy()
+        np.testing.assert_allclose(actual[100:-100], reference[100:-100], atol=2e-3, rtol=1e-3)
+
+    def test_preserves_passband_signal(self):
+        samples = np.arange(24000, dtype=np.float32)
+        waveform = np.sin(2 * np.pi * 3000 * samples / 24000).astype(np.float32)
+
+        output = self._resample(waveform)
+
+        input_rms = np.sqrt(np.mean(waveform**2))
+        output_rms = np.sqrt(np.mean(output[100:-100] ** 2))
+        assert output_rms == pytest.approx(input_rms, rel=0.02)
+
+    @pytest.mark.parametrize("frequency", [6000, 10000])
+    def test_attenuates_aliasing_frequencies(self, frequency):
+        samples = np.arange(24000, dtype=np.float32)
+        waveform = np.sin(2 * np.pi * frequency * samples / 24000).astype(np.float32)
+
+        output = self._resample(waveform)
+
+        input_rms = np.sqrt(np.mean(waveform**2))
+        output_rms = np.sqrt(np.mean(output[100:-100] ** 2))
+        assert output_rms < input_rms * 0.01
+
+    def test_rejects_non_integer_downsampling_ratio(self):
+        with pytest.raises(ValueError, match="integer downsampling ratio"):
+            StreamingAudioResampler(24000, 16000)
+
+    def test_rejects_multichannel_audio(self):
+        resampler = StreamingAudioResampler(24000, 8000)
+
+        with pytest.raises(ValueError, match="only supports mono audio"):
+            resampler.process(np.zeros((2, 240), dtype=np.float32))
+
 
 class TestResolveAudioFormat:
     """Test _resolve_audio_format via the serving chat class."""

tests/entrypoints/openai_api/test_serving_speech.py

@@ -2953,8 +2953,10 @@ class TestStreamingResponse:
             usage_acc=None,
             tts_params=None,
             collect=None,
+            target_sample_rate=None,
         ):
             captured["tts_params"] = tts_params
+            captured["target_sample_rate"] = target_sample_rate
             yield b"\x00\x00"
 
         monkeypatch.setattr(speech_server, "_prepare_speech_generation", prepare)
@@ -2967,6 +2969,7 @@ class TestStreamingResponse:
 
         assert response.status_code == 200
         assert captured["tts_params"] is finalized_tts_params
+        assert captured["target_sample_rate"] is None
 
     def test_sse_streaming(self, streaming_app):
         """stream_format=sse without stream=True returns audio deltas as SSE."""
@@ -3292,6 +3295,28 @@ class TestMergeBatchItem:
 
         assert merged.non_streaming_mode is False
 
+    def test_sample_rate_batch_default_used(self):
+        """Batch-level sample_rate should be used when the item omits it."""
+        batch = BatchSpeechRequest(
+            items=[SpeechBatchItem(input="hi")],
+            sample_rate=8000,
+        )
+
+        merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0])
+
+        assert merged.sample_rate == 8000
+
+    def test_sample_rate_item_override_wins(self):
+        """Per-item sample_rate should override the batch-level default."""
+        batch = BatchSpeechRequest(
+            items=[SpeechBatchItem(input="hi", sample_rate=24000)],
+            sample_rate=8000,
+        )
+
+        merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0])
+
+        assert merged.sample_rate == 24000
+
 
 def test_streaming_speech_session_config_accepts_non_streaming_mode():
     config = StreamingSpeechSessionConfig(non_streaming_mode=True)
@@ -3299,6 +3324,33 @@ def test_streaming_speech_session_config_accepts_non_streaming_mode():
     assert config.non_streaming_mode is True
 
 
+@pytest.mark.parametrize("sample_rate", [8000, 24000])
+def test_speech_requests_accept_supported_sample_rates(sample_rate):
+    assert OpenAICreateSpeechRequest(input="hello", sample_rate=sample_rate).sample_rate == sample_rate
+
+
+@pytest.mark.parametrize("sample_rate", [0, -8000])
+def test_speech_request_rejects_non_positive_sample_rate(sample_rate):
+    with pytest.raises(ValidationError):
+        OpenAICreateSpeechRequest(input="hello", sample_rate=sample_rate)
+
+
+def test_sample_rate_uses_adapter_capabilities():
+    server = object.__new__(OmniOpenAIServingSpeech)
+    request = OpenAICreateSpeechRequest(input="hello", sample_rate=8000)
+
+    server._tts_model_type = "qwen3_tts"
+    assert server._validate_speech_sample_rate(request) is None
+
+    request.sample_rate = 16000
+    assert server._validate_speech_sample_rate(request) == (
+        "sample_rate=16000 is not supported by the current TTS model; supported rates: 8000, 24000"
+    )
+
+    server._tts_model_type = "fish_speech"
+    assert server._validate_speech_sample_rate(request) == ("sample_rate is not supported by the current TTS model")
+
+
 def test_streaming_speech_session_config_scopes_native_speed_to_http():
     with pytest.raises(
         ValidationError,
@@ -5027,6 +5079,78 @@ class TestTTSAsyncOffloading:
         assert "partial_audio" not in payload["error"]
         assert "action" not in payload["error"]
 
+    @pytest.mark.asyncio
+    async def test_generate_audio_chunks_resamples_pcm_to_8khz(self, qwen3_tts_server):
+        async def pcm_generator():
+            output = create_mock_audio_output_for_test()
+            output.multimodal_output.update(
+                audio=torch.linspace(-0.5, 0.5, 2400, dtype=torch.float32),
+                sr=24000,
+            )
+            yield output
+
+        chunks = [
+            chunk
+            async for chunk in qwen3_tts_server._generate_audio_chunks(
+                pcm_generator(),
+                "req-8khz",
+                target_sample_rate=8000,
+            )
+        ]
+
+        assert len(chunks) >= 2
+        assert len(b"".join(chunks)) == 800 * 2
+
+    @pytest.mark.asyncio
+    async def test_generate_audio_chunks_wav_header_uses_target_sample_rate(self, qwen3_tts_server):
+        async def pcm_generator():
+            output = create_mock_audio_output_for_test()
+            output.multimodal_output.update(
+                audio=torch.zeros(2400, dtype=torch.float32),
+                sr=24000,
+            )
+            yield output
+
+        chunks = [
+            chunk
+            async for chunk in qwen3_tts_server._generate_audio_chunks(
+                pcm_generator(),
+                "req-8khz-wav",
+                response_format="wav",
+                target_sample_rate=8000,
+            )
+        ]
+
+        assert struct.unpack("<I", chunks[0][24:28])[0] == 8000
+
+    @pytest.mark.asyncio
+    async def test_streaming_resampler_retains_first_chunk_sample_rate_metadata(self, qwen3_tts_server):
+        async def pcm_generator():
+            first = create_mock_audio_output_for_test()
+            first.multimodal_output.update(
+                audio=torch.zeros(10, dtype=torch.float32),
+                sr=24000,
+            )
+            yield first
+
+            second = create_mock_audio_output_for_test()
+            second.multimodal_output.update(audio=torch.zeros(230, dtype=torch.float32))
+            second.multimodal_output.pop("sr", None)
+            yield second
+
+        chunks = [
+            chunk
+            async for chunk in qwen3_tts_server._generate_audio_chunks(
+                pcm_generator(),
+                "req-retained-sr",
+                response_format="wav",
+                target_sample_rate=8000,
+            )
+        ]
+
+        assert chunks[0][:4] == b"RIFF"
+        assert struct.unpack("<I", chunks[0][24:28])[0] == 8000
+
     @pytest.mark.asyncio
     async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server):
         async def pcm_generator():

tests/helpers/assertions.py

@@ -873,6 +873,13 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
     elif req_fmt == "wav" and response.audio_format:
         assert req_fmt in response.audio_format
 
+    expected_sample_rate = request_config.get("sample_rate")
+    if expected_sample_rate is not None and req_fmt == "wav" and response.audio_bytes:
+        info = sf.info(io.BytesIO(response.audio_bytes))
+        assert info.samplerate == int(expected_sample_rate), (
+            f"Expected sample_rate={expected_sample_rate}, got {info.samplerate}"
+        )
+
     if run_level in {"advanced_model", "full_model"}:
         if req_fmt == "pcm" and response.audio_bytes:
             min_hnr_db = float(request_config.get("min_hnr_db", _MIN_PCM_SPEECH_HNR_DB))

tests/helpers/client.py

@@ -1461,6 +1461,7 @@ class OnlineOmniClient:
             "seed",
             "instructions",
             "speed",
+            "sample_rate",
             "stream_format",
             "x_vector_only_mode",
         ):

tests/helpers/stage_config.py

@@ -381,6 +381,40 @@ _CI_OVERLAYS: dict[str, dict[str, Any]] = {
             },
         },
     },
+    "qwen3_omni_moe_colocate_async": {
+        "base_config": "qwen3_omni_moe_thinking.yaml",
+        "stages": [
+            {
+                "stage_id": 0,
+                "enable_sleep_mode": True,
+                "enforce_eager": True,
+                "trust_remote_code": True,
+                "tensor_parallel_size": 1,
+                "devices": "0",
+                "max_num_seqs": 1,
+                "max_model_len": 8192,
+                "max_num_batched_tokens": 8192,
+                "enable_prefix_caching": False,
+                "skip_mm_profiling": True,
+            },
+        ],
+        "platforms": {
+            "rocm": {
+                "stages": [
+                    {
+                        "stage_id": 0,
+                        # This control-plane test needs only 8K tokens. On MI300,
+                        # percentage sizing attempted a ~125 GiB KV allocation
+                        # after the 64.5 GiB model load and then failed while
+                        # allocating the cache. Keep CUDA's existing sizing and
+                        # give only ROCm a bounded cache with ample headroom.
+                        "gpu_memory_utilization": None,
+                        "kv_cache_memory_bytes": 2 * 1024**3,
+                    },
+                ],
+            },
+        },
+    },
     "qwen3_omni_moe_multi_replicas_4gpu": {
         "base_config": "qwen3_omni_moe.yaml",
         "async_chunk": True,

vllm_omni/config/environment_variable_inventory.py

@@ -167,6 +167,7 @@ _MODEL_PROMOTE = (
     "VLLM_GEPARD_LOOKBACK_FRAMES",
     "VLLM_OMNI_FISH_KVCACHE_ATTN",
     "VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER",
+    "VLLM_OMNI_SENSENOVA_PAGED_DECODE",
 )
 
 _MODEL_REQUEST_SCOPE = (

vllm_omni/config/omni_config.py

@@ -47,6 +47,7 @@ from vllm_omni.config.stage_config import (
     build_stage_runtime_overrides,
     load_deploy_config,
     normalize_pipeline_cli_overrides,
+    reconcile_diffusion_attention_overrides,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 
@@ -1207,6 +1208,9 @@ def _stage_engine_values(
     if topology.omni_kv_config:
         engine["omni_kv_config"] = _copy_value(topology.omni_kv_config)
     if stage_cli_overrides:
+        if topology.execution_type == StageExecutionType.DIFFUSION:
+            # Mirror StageConfig.to_omegaconf so both projections resolve alike.
+            reconcile_diffusion_attention_overrides(engine, stage_cli_overrides)
         engine.update(_copy_value(stage_cli_overrides))
     _validate_stage_engine_override_ownership(
         topology.stage_id,

vllm_omni/config/stage_config.py

@@ -7,7 +7,7 @@ from __future__ import annotations
 import functools
 import re
 import warnings
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from dataclasses import asdict, dataclass, field, fields
 from enum import Enum
 from pathlib import Path
@@ -161,6 +161,28 @@ def _apply_diffusion_parallel_runtime_overrides(
         engine_args["parallel_config"] = parallel_config_dict
 
 
+def reconcile_diffusion_attention_overrides(
+    engine_args: dict[str, Any],
+    runtime_overrides: Mapping[str, Any],
+) -> None:
+    """Apply CLI precedence across the two diffusion attention representations.
+
+    ``diffusion_attention_backend`` and ``diffusion_attention_config.default``
+    express the same setting and are rejected downstream when both are set, so
+    a CLI value in one form replaces the YAML value in the other.
+    """
+    if runtime_overrides.get("diffusion_attention_backend") is not None:
+        yaml_config = engine_args.get("diffusion_attention_config")
+        if isinstance(yaml_config, Mapping) and yaml_config.get("default") is not None:
+            remaining = {k: v for k, v in yaml_config.items() if k != "default"}
+            if remaining:
+                engine_args["diffusion_attention_config"] = remaining
+            else:
+                engine_args.pop("diffusion_attention_config")
+    if runtime_overrides.get("diffusion_attention_config") is not None:
+        engine_args.pop("diffusion_attention_backend", None)
+
+
 class StageType(str, Enum):
     """Type of processing stage in the Omni pipeline."""
 
@@ -1092,6 +1114,7 @@ class StageConfig:
 
         if StageType(self.stage_type) == StageType.DIFFUSION:
             _apply_diffusion_parallel_runtime_overrides(engine_args, runtime_overrides)
+            reconcile_diffusion_attention_overrides(engine_args, runtime_overrides)
 
         # CLI overrides take precedence over YAML defaults
         for key, value in runtime_overrides.items():

vllm_omni/diffusion/lora/loader.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import os
 from collections import OrderedDict, defaultdict
 from collections.abc import Callable
@@ -41,7 +44,7 @@ def _prepare_lora_delta(
     lora_b_suffix: str = "lora_B.weight",
     lora_bias_suffix: str = "bias",
 ):
-    used_keys = set()
+    used_keys: set[str] = set()
     # stacked_params_mapping                       param_to_weight_names
     # [(".to_qkv", ".to_q.", "q")
     # (".to_qkv", ".to_k.", "k")  ========> {".to_qkv": [".to_q", ".to_k", ".to_v"]}
@@ -54,7 +57,10 @@ def _prepare_lora_delta(
     if param_to_weight_names is None:
         param_to_weight_names = defaultdict(list)
     for param_name, weight_names in param_to_weight_names.items():
-        if param_name not in base_key:
+        # Fused names can overlap by prefix (`.qkv_proj` is a substring of
+        # `.qkv_proj_mot_gen`), and a substring match stacks both mappings into
+        # one delta of twice the height.
+        if not base_key.endswith(param_name):
             continue
         is_stacked_param = True
         # handle lora_a_key and lora_b_key together
@@ -78,7 +84,7 @@ def _prepare_lora_delta(
                 stacked_deltas.append(delta)
                 used_keys.add(lora_a_key)
                 used_keys.add(lora_b_key)
-        continue
+        break
 
     if is_stacked_param:
         return torch.concat(stacked_deltas), used_keys
@@ -200,6 +206,10 @@ def _apply_diffusers_lora_alpha_scaling(state_dict: dict[str, torch.Tensor]) ->
 class LoraLoaderMixin:
     transformer_name = "transformer"
 
+    # Set by the pipeline this mixin is combined with; annotated, not assigned.
+    transformer: torch.nn.Module
+    _lora_loaded: dict[str | None, dict[str, torch.Tensor]]
+
     # Lazy initialization to avoid MRO issues: __init__ may not be called
     # when mixin with nn.Module
     @property
@@ -373,6 +383,7 @@ class QwenImageLoraLoaderMixin(LoraLoaderMixin):
 
 class WanLoraLoaderMixin(LoraLoaderMixin):
     transformer_2_name = "transformer_2"
+    has_transformer_2: bool
 
     def load_lora_weights(
         self,

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -2223,7 +2223,7 @@ class Cosmos3OmniDiffusersPipeline(
         if prompt_suffix:
             prompt = f"{prompt.rstrip()} {prompt_suffix.lstrip()}".strip()
         if _is_rank_zero():
-            logger.info("Final prompt: '%s'", prompt)
+            logger.debug("Final prompt: '%s'", prompt)
 
         if negative_metadata_mode == "none":
             negative_dur_tmpl = None

vllm_omni/diffusion/models/minimax_h3/chunked_decode.py

@@ -0,0 +1,391 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Model-local coordination for MiniMax H3 temporal chunk decoding.
+
+This boundary owns released-code compatibility and reconciles preflight and
+callback failures at fixed VAE-group rendezvous points. Cropping,
+representation conversion, transport, and encoding remain runtime
+responsibilities.
+"""
+
+from __future__ import annotations
+
+from contextlib import AbstractContextManager
+from typing import Protocol
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+
+from .temporal_chunks import (
+    MiniMaxH3TemporalChunkCompatibilityError,
+    MiniMaxH3TemporalDecodePlan,
+    _TemporalChunkCallbackSignature,
+    decode_minimax_h3_temporal_chunks_compat,
+    prepare_minimax_h3_temporal_chunks_compat,
+)
+
+
+class MiniMaxH3VideoChunkCallback(_TemporalChunkCallbackSignature, Protocol):
+    """Consume one ordered, temporally committed H3 video chunk.
+
+    ``frames`` is an owned, contiguous float32 RGB tensor in ``BCTHW`` layout
+    and the ``[0, 1]`` value range. The callback runs synchronously on the
+    producer's current device stream. Before reading it on another stream, the
+    consumer must establish ordering from the producer stream (for example via
+    an event) and retain the allocation for that stream's lifetime. The
+    consumer may retain or mutate the tensor; it does not alias the returned
+    complete decode.
+    """
+
+
+class MiniMaxH3ChunkedDecodeUnsupportedError(RuntimeError):
+    """Raised when loaded MiniMax H3 code cannot honor the chunk contract."""
+
+
+class MiniMaxH3ChunkCallbackPeerError(RuntimeError):
+    """Raised on peer VAE ranks when the output-rank callback fails."""
+
+
+class _MiniMaxH3ChunkDecodeHost(Protocol):
+    """Narrow VAE capabilities required by the chunk coordinator."""
+
+    model: nn.Module
+
+    def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor: ...
+
+    def _decode_tiling_context(
+        self,
+        latent: torch.Tensor,
+    ) -> AbstractContextManager: ...
+
+    def _normalize_decoded_frames(
+        self,
+        decoded: torch.Tensor,
+    ) -> torch.Tensor: ...
+
+
+class _MiniMaxH3ChunkDecodeCoordinator:
+    """Run source-gated chunk decode with group-local failure semantics.
+
+    The coordinator is a plain Python object. It caches only successful source
+    validation; request callbacks, tensors, metadata, and process groups remain
+    local to :meth:`decode`.
+
+    Errors raised inside the checkpoint decoder after native collectives begin
+    cannot be reconciled here: another rank may already be blocked inside that
+    collective. Resolving such a failure requires native collective handling or
+    supervisor-level timeout and worker-group teardown; the process group cannot
+    be assumed reusable.
+    """
+
+    def __init__(self) -> None:
+        self._sources_validated = False
+
+    @staticmethod
+    def _build_execution_signature(
+        latent: torch.Tensor,
+        plan: MiniMaxH3TemporalDecodePlan,
+    ) -> tuple[int, ...]:
+        """Describe every shape and plan field used by the temporal schedule."""
+
+        dtype_codes: dict[torch.dtype | None, int] = {
+            None: 0,
+            torch.float16: 1,
+            torch.bfloat16: 2,
+            torch.float32: 3,
+            torch.float64: 4,
+        }
+        return (
+            # Keep the layout versioned so adding a field cannot silently make
+            # an older and newer worker appear compatible.
+            1,
+            *[int(size) for size in latent.shape],
+            dtype_codes.get(latent.dtype, -1),
+            *[int(size) for size in plan.latent.shape],
+            dtype_codes.get(plan.latent.dtype, -1),
+            int(plan.num_chunks),
+            int(plan.total_frames),
+            int(plan.pad_frames),
+            int(plan.output_frames),
+            len(plan.chunk_frames),
+            *[int(frame_count) for frame_count in plan.chunk_frames],
+            dtype_codes.get(plan.temporal_dtype, -1),
+        )
+
+    @staticmethod
+    def _is_output_rank(
+        chunk_callback: MiniMaxH3VideoChunkCallback | None,
+        *,
+        group: dist.ProcessGroup | None,
+        device: torch.device,
+    ) -> bool:
+        if group is None:
+            return chunk_callback is not None
+
+        group_rank = dist.get_rank(group)
+        # A valid configuration has one contribution: rank 0 contributes 1.
+        # A non-zero owner contributes at least 2, and multiple owners sum to
+        # at least 3, so every rank makes the same decision from one reduction.
+        callback_owner = torch.tensor(
+            [group_rank + 1 if chunk_callback is not None else 0],
+            dtype=torch.int64,
+            device=device,
+        )
+        dist.all_reduce(callback_owner, op=dist.ReduceOp.SUM, group=group)
+        if int(callback_owner.item()) != 1:
+            raise ValueError("MiniMax-H3 chunk decode requires a callback on exactly VAE group rank 0")
+        return group_rank == 0
+
+    @staticmethod
+    def _synchronize_callback_error(
+        callback_error: BaseException | None,
+        *,
+        group: dist.ProcessGroup | None,
+        device: torch.device,
+    ) -> None:
+        if group is None:
+            if callback_error is not None:
+                raise callback_error
+            return
+
+        is_output_rank = dist.get_rank(group) == 0
+        failed = torch.tensor(
+            [int(is_output_rank and callback_error is not None)],
+            dtype=torch.int32,
+            device=device,
+        )
+        dist.broadcast(
+            failed,
+            src=dist.get_global_rank(group, 0),
+            group=group,
+        )
+        if int(failed.item()) == 0:
+            return
+        if is_output_rank:
+            assert callback_error is not None
+            raise callback_error
+        raise MiniMaxH3ChunkCallbackPeerError("MiniMax-H3 video chunk callback failed on this VAE group's output rank")
+
+    @staticmethod
+    def _synchronize_preflight_error(
+        preflight_error: Exception | None,
+        temporal_dtype: torch.dtype | None,
+        execution_signature: tuple[int, ...] | None,
+        *,
+        group: dist.ProcessGroup | None,
+        device: torch.device,
+    ) -> None:
+        if group is None:
+            if preflight_error is not None:
+                raise preflight_error
+            return
+        failed = torch.tensor(
+            [int(preflight_error is not None)],
+            dtype=torch.int32,
+            device=device,
+        )
+        dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=group)
+        if int(failed.item()) == 0:
+            dtype_codes = {
+                None: 0,
+                torch.float16: 1,
+                torch.bfloat16: 2,
+                torch.float32: 3,
+            }
+            code = torch.tensor(
+                [dtype_codes.get(temporal_dtype, -1)],
+                dtype=torch.int32,
+                device=device,
+            )
+            lowest = code.clone()
+            highest = code.clone()
+            dist.all_reduce(lowest, op=dist.ReduceOp.MIN, group=group)
+            dist.all_reduce(highest, op=dist.ReduceOp.MAX, group=group)
+            if int(lowest.item()) == int(highest.item()) and int(lowest.item()) >= 0:
+                if execution_signature is None:
+                    raise RuntimeError("MiniMax-H3 temporal decode preflight produced no execution signature")
+
+                # The plan length is shape-dependent, so first agree on a
+                # common tensor size and then compare the complete signatures.
+                # This keeps all ranks in tensor collectives and catches a
+                # schedule mismatch before entering the decoder's collectives.
+                signature_length = torch.tensor(
+                    [len(execution_signature)],
+                    dtype=torch.int64,
+                    device=device,
+                )
+                max_signature_length = signature_length.clone()
+                dist.all_reduce(
+                    max_signature_length,
+                    op=dist.ReduceOp.MAX,
+                    group=group,
+                )
+                padded_signature = torch.full(
+                    (int(max_signature_length.item()),),
+                    -2,
+                    dtype=torch.int64,
+                    device=device,
+                )
+                padded_signature[: len(execution_signature)] = torch.tensor(
+                    execution_signature,
+                    dtype=torch.int64,
+                    device=device,
+                )
+                gathered_signatures = [torch.empty_like(padded_signature) for _ in range(dist.get_world_size(group))]
+                dist.all_gather(gathered_signatures, padded_signature, group=group)
+                if any(not torch.equal(padded_signature, peer) for peer in gathered_signatures):
+                    raise MiniMaxH3ChunkedDecodeUnsupportedError(
+                        "MiniMax-H3 VAE ranks built incompatible temporal decode plans"
+                    )
+                return
+            raise MiniMaxH3ChunkedDecodeUnsupportedError(
+                "MiniMax-H3 VAE ranks selected incompatible temporal concat dtypes"
+            )
+        if preflight_error is not None:
+            raise preflight_error
+        raise MiniMaxH3ChunkedDecodeUnsupportedError("A peer VAE rank rejected the MiniMax-H3 temporal chunk preflight")
+
+    def decode(
+        self,
+        host: _MiniMaxH3ChunkDecodeHost,
+        latent: torch.Tensor,
+        chunk_callback: MiniMaxH3VideoChunkCallback | None,
+        *,
+        group: dist.ProcessGroup | None,
+    ) -> torch.Tensor:
+        """Decode one latent while publishing committed normalized chunks."""
+
+        plan: MiniMaxH3TemporalDecodePlan | None = None
+        preflight_error: Exception | None = None
+        if latent.ndim != 5:
+            preflight_error = MiniMaxH3ChunkedDecodeUnsupportedError(
+                f"MiniMax-H3 video latent must be rank 5, got {tuple(latent.shape)}"
+            )
+        else:
+            try:
+                denormalized = host._denormalize_latent(latent)
+                plan = prepare_minimax_h3_temporal_chunks_compat(
+                    host.model,
+                    denormalized,
+                    validate_sources=not self._sources_validated,
+                )
+                self._sources_validated = True
+                del denormalized
+            except MiniMaxH3TemporalChunkCompatibilityError as exc:
+                preflight_error = MiniMaxH3ChunkedDecodeUnsupportedError(str(exc))
+            except Exception as exc:  # noqa: BLE001
+                # Every VAE rank must reach the preflight reduction even when
+                # one node's private planner or environment fails locally.
+                exc.__traceback__ = None
+                preflight_error = exc
+        self._synchronize_preflight_error(
+            preflight_error,
+            plan.temporal_dtype if plan is not None else None,
+            self._build_execution_signature(latent, plan) if plan is not None else None,
+            group=group,
+            device=latent.device,
+        )
+        assert plan is not None
+        is_output_rank = self._is_output_rank(
+            chunk_callback,
+            group=group,
+            device=latent.device,
+        )
+        callback_error: BaseException | None = None
+        next_chunk_index = 0
+        expected_total_chunks: int | None = None
+        next_frame_start = 0
+        saw_final = False
+
+        def publish(
+            decoded_chunk: torch.Tensor,
+            *,
+            chunk_index: int,
+            total_chunks: int,
+            frame_start: int,
+            is_final: bool,
+        ) -> None:
+            nonlocal callback_error, expected_total_chunks
+            nonlocal next_chunk_index, next_frame_start, saw_final
+            if not is_output_rank or callback_error is not None:
+                return
+            try:
+                if expected_total_chunks is None:
+                    expected_total_chunks = total_chunks
+                if (
+                    saw_final
+                    or total_chunks <= 0
+                    or total_chunks != expected_total_chunks
+                    or chunk_index != next_chunk_index
+                    or chunk_index >= total_chunks
+                    or frame_start != next_frame_start
+                    or is_final != (chunk_index + 1 == total_chunks)
+                ):
+                    raise RuntimeError(
+                        "MiniMax-H3 temporal backend emitted discontinuous "
+                        "chunk metadata: "
+                        f"index={chunk_index}/{next_chunk_index}, "
+                        f"total={total_chunks}/{expected_total_chunks}, "
+                        f"frame_start={frame_start}/{next_frame_start}, "
+                        f"final={is_final}, after_final={saw_final}"
+                    )
+                frames = host._normalize_decoded_frames(decoded_chunk)
+                frames = frames.clone(memory_format=torch.contiguous_format)
+                frame_count = int(frames.shape[2])
+                if frame_count <= 0:
+                    raise RuntimeError("MiniMax-H3 temporal backend emitted an empty chunk")
+                next_chunk_index += 1
+                next_frame_start += frame_count
+                saw_final = is_final
+                assert chunk_callback is not None
+                chunk_callback(
+                    frames,
+                    chunk_index=chunk_index,
+                    total_chunks=total_chunks,
+                    frame_start=frame_start,
+                    is_final=is_final,
+                )
+            except BaseException as exc:  # noqa: BLE001
+                # Keep a failed GPU chunk and its callback locals out of the
+                # traceback while peer ranks finish the remaining collectives.
+                exc.__traceback__ = None
+                callback_error = exc
+
+        # Do not catch and reduce exceptions from this region. A peer may
+        # already be blocked in one of the checkpoint's native tile
+        # collectives, in which case a new adapter-level collective would also
+        # deadlock. A native backend or supervisor must handle timeout and
+        # worker-group teardown.
+        with host._decode_tiling_context(latent):
+            decoded = decode_minimax_h3_temporal_chunks_compat(
+                host.model,
+                plan,
+                publish,
+            )
+        frames = host._normalize_decoded_frames(decoded)
+        if (
+            is_output_rank
+            and callback_error is None
+            and (not saw_final or next_chunk_index != expected_total_chunks or next_frame_start != int(frames.shape[2]))
+        ):
+            callback_error = RuntimeError(
+                "MiniMax-H3 temporal backend did not reconstruct the full "
+                "decode: "
+                f"final={saw_final}, "
+                f"chunks={next_chunk_index}/{expected_total_chunks}, "
+                f"frames={next_frame_start}/{frames.shape[2]}"
+            )
+        self._synchronize_callback_error(
+            callback_error,
+            group=group,
+            device=frames.device,
+        )
+        return frames
+
+
+__all__ = [
+    "MiniMaxH3ChunkCallbackPeerError",
+    "MiniMaxH3ChunkedDecodeUnsupportedError",
+    "MiniMaxH3VideoChunkCallback",
+]

vllm_omni/diffusion/models/minimax_h3/temporal_chunks.py

@@ -0,0 +1,414 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Source-gated temporal chunk compatibility for the released MiniMax H3 VAE."""
+
+from __future__ import annotations
+
+import hashlib
+import importlib
+import inspect
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Protocol
+
+import torch
+import torch.nn as nn
+
+MINIMAX_H3_RELEASED_REVISION = "42ed227ee7df40d41602854ae760620d6eb651fe"
+# Compatibility pins detect upstream code drift; they are not a trust or
+# security boundary for remote code execution.
+MINIMAX_H3_RELEASED_BUNDLE_SHA256 = "30e64afbaa940696bf8bfe2c86003a4b1666c84b22fa3aa12403a3e6a03f705d"
+
+_RELEASED_TEMPORAL_CONFIG = {
+    "clip_length": 17,
+    "tokens_chunk_size": 5,
+    "token_overlap": 2,
+    "token_drop": 3,
+    "vae_ratio_t": 4,
+    "frame_pre_padding": 3,
+    "frame_overlap": 5,
+    "isolated_first_frame": False,
+    "isolated_last_frame": False,
+}
+
+
+class MiniMaxH3TemporalChunkCompatibilityError(RuntimeError):
+    """Raised when remote H3 code cannot use the local compatibility path."""
+
+
+class _TemporalChunkCallbackSignature(Protocol):
+    """Shared callable shape; each producer defines tensor semantics."""
+
+    def __call__(
+        self,
+        frames: torch.Tensor,
+        /,
+        *,
+        chunk_index: int,
+        total_chunks: int,
+        frame_start: int,
+        is_final: bool,
+    ) -> None: ...
+
+
+@dataclass(slots=True)
+class MiniMaxH3TemporalDecodePlan:
+    latent: torch.Tensor
+    num_chunks: int
+    total_frames: int
+    pad_frames: int
+    output_frames: int
+    chunk_frames: tuple[int, ...]
+    temporal_dtype: torch.dtype | None
+
+
+class _RawChunkAssembly:
+    def __init__(
+        self,
+        *,
+        callback: _TemporalChunkCallbackSignature,
+        plan: MiniMaxH3TemporalDecodePlan,
+    ) -> None:
+        self.callback = callback
+        self.plan = plan
+        self.frames: torch.Tensor | None = None
+        self.write_pos = 0
+        self.logical_frames = 0
+        self.dropped_frames = 0
+        self.chunk_index = 0
+        self.pending_final: tuple[torch.Tensor, int, int] | None = None
+        self.contract_error: RuntimeError | None = None
+
+    def write(self, part: torch.Tensor) -> None:
+        part_frames = int(part.shape[2])
+        if part_frames <= 0:
+            return
+        self.logical_frames += part_frames
+        remaining = self.plan.output_frames - self.write_pos
+        copy_frames = min(part_frames, max(0, remaining))
+        if copy_frames <= 0:
+            self.dropped_frames += part_frames
+            return
+        if self.chunk_index >= len(self.plan.chunk_frames):
+            self.contract_error = RuntimeError("MiniMax-H3 compatibility decode emitted an extra chunk")
+        else:
+            expected_frames = self.plan.chunk_frames[self.chunk_index]
+            if copy_frames != expected_frames and self.contract_error is None:
+                self.contract_error = RuntimeError(
+                    "MiniMax-H3 compatibility chunk size mismatch: "
+                    f"chunk={self.chunk_index}, frames={copy_frames}/{expected_frames}"
+                )
+
+        emitted = part[:, :, :copy_frames]
+        if self.frames is None:
+            output_shape = list(emitted.shape)
+            output_shape[2] = self.plan.output_frames
+            self.frames = torch.empty(
+                output_shape,
+                dtype=emitted.dtype,
+                device=emitted.device,
+            )
+        frame_start = self.write_pos
+        frame_stop = frame_start + copy_frames
+        self.frames[:, :, frame_start:frame_stop].copy_(emitted)
+        if frame_stop == self.plan.output_frames:
+            self.pending_final = (emitted, self.chunk_index, frame_start)
+        elif self.contract_error is None:
+            self.callback(
+                emitted,
+                chunk_index=self.chunk_index,
+                total_chunks=len(self.plan.chunk_frames),
+                frame_start=frame_start,
+                is_final=False,
+            )
+        self.chunk_index += 1
+        self.write_pos = frame_stop
+        self.dropped_frames += part_frames - copy_frames
+
+    def finish(self) -> torch.Tensor:
+        if self.contract_error is not None:
+            raise self.contract_error
+        if self.frames is None or self.pending_final is None:
+            raise RuntimeError("MiniMax-H3 compatibility decode produced no terminal chunk")
+        if (
+            self.logical_frames != self.plan.total_frames
+            or self.dropped_frames != self.plan.pad_frames
+            or self.write_pos != self.plan.output_frames
+            or self.chunk_index != len(self.plan.chunk_frames)
+        ):
+            raise RuntimeError(
+                "MiniMax-H3 compatibility frame plan mismatch: "
+                f"logical={self.logical_frames}/{self.plan.total_frames}, "
+                f"dropped={self.dropped_frames}/{self.plan.pad_frames}, "
+                f"written={self.write_pos}/{self.plan.output_frames}, "
+                f"chunks={self.chunk_index}/{len(self.plan.chunk_frames)}"
+            )
+
+        final_frames, final_index, final_start = self.pending_final
+        self.callback(
+            final_frames,
+            chunk_index=final_index,
+            total_chunks=len(self.plan.chunk_frames),
+            frame_start=final_start,
+            is_final=True,
+        )
+        return self.frames
+
+
+def _source_bundle_sha256(model: nn.Module) -> tuple[str | None, str | None]:
+    try:
+        source = inspect.getsourcefile(type(model))
+    except TypeError:
+        return None, None
+    if source is None:
+        return None, None
+    package_dir = Path(source).parent
+    files = sorted(path for path in package_dir.glob("*.py") if path.name != "__init__.py")
+    digest = hashlib.sha256()
+    try:
+        for path in files:
+            digest.update(path.name.encode())
+            digest.update(b"\0")
+            digest.update(path.read_bytes())
+            digest.update(b"\0")
+    except OSError:
+        return str(package_dir), None
+    return str(package_dir), digest.hexdigest()
+
+
+def _validate_released_sources(model: nn.Module) -> None:
+    source_dir, actual = _source_bundle_sha256(model)
+    if actual != MINIMAX_H3_RELEASED_BUNDLE_SHA256:
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            "The vLLM-owned MiniMax-H3 chunk adapter supports only remote code "
+            f"from MiniMaxAI/MiniMax-H3@{MINIMAX_H3_RELEASED_REVISION}: "
+            f"bundle={actual!r} ({source_dir}). "
+            "The normal complete-output decode remains available."
+        )
+
+
+def _validate_released_config(model: nn.Module) -> None:
+    actual = {name: getattr(model, name, None) for name in _RELEASED_TEMPORAL_CONFIG}
+    if actual != _RELEASED_TEMPORAL_CONFIG:
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            "The vLLM-owned MiniMax-H3 chunk adapter requires the released temporal "
+            f"configuration {_RELEASED_TEMPORAL_CONFIG}, found {actual}. "
+            "The normal complete-output decode remains available."
+        )
+
+
+def _validate_compat_contract(model: nn.Module, latent: torch.Tensor) -> None:
+    if latent.ndim != 5 or int(latent.shape[2]) <= 0:
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            f"MiniMax-H3 video latent must have positive rank-5 time, got {tuple(latent.shape)}"
+        )
+    if not bool(getattr(model, "use_3d_conv", False)):
+        raise MiniMaxH3TemporalChunkCompatibilityError("MiniMax-H3 temporal chunk decode requires the 3D decoder")
+    if bool(getattr(model, "training", False)):
+        raise MiniMaxH3TemporalChunkCompatibilityError("MiniMax-H3 temporal chunk decode is inference-only")
+    required_methods = (
+        "_adaptive_decode",
+        "_decode_temporal_output_frame_plan",
+        "_decode_temporal_pad_frames",
+        "blend",
+    )
+    missing = [name for name in required_methods if not callable(getattr(model, name, None))]
+    if missing:
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            "Loaded MiniMax-H3 VAE lacks required released-code methods: " + ", ".join(missing)
+        )
+
+
+def _truncate_chunk_frames(
+    logical_chunk_frames: list[int],
+    output_frames: int,
+) -> tuple[int, ...]:
+    remaining = output_frames
+    result: list[int] = []
+    for logical_frames in logical_chunk_frames:
+        emitted_frames = min(logical_frames, max(0, remaining))
+        if emitted_frames > 0:
+            result.append(emitted_frames)
+            remaining -= emitted_frames
+    if remaining != 0:
+        raise RuntimeError(f"MiniMax-H3 compatibility chunk plan cannot cover output: remaining={remaining}")
+    return tuple(result)
+
+
+def _build_decode_plan(
+    model: nn.Module,
+    latent: torch.Tensor,
+    temporal_dtype: torch.dtype | None,
+) -> MiniMaxH3TemporalDecodePlan:
+    pseudo_total_tokens = int(latent.shape[2]) + int(model.token_drop)  # type: ignore[attr-defined]
+    tokens_chunk_size = int(model.tokens_chunk_size)  # type: ignore[attr-defined]
+    remainder = pseudo_total_tokens % tokens_chunk_size
+    pad_tokens = 0 if remainder == 0 else tokens_chunk_size - remainder
+    pseudo_total_tokens += pad_tokens
+    num_chunks = pseudo_total_tokens // tokens_chunk_size - int(model.token_drop > 0)  # type: ignore[attr-defined]
+    if num_chunks <= 0:
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            f"MiniMax-H3 temporal chunk plan is empty for latent T={latent.shape[2]}"
+        )
+
+    temporal_latent = latent
+    if pad_tokens:
+        pad = temporal_latent[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)
+        temporal_latent = torch.cat((temporal_latent, pad), dim=2)
+    total_frames, pad_frames, output_frames = model._decode_temporal_output_frame_plan(  # type: ignore[attr-defined]
+        temporal_latent,
+        None,
+        None,
+        num_chunks,
+        pad_tokens,
+    )
+    if output_frames <= 0:
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            f"MiniMax-H3 temporal chunk plan has {output_frames} output frames"
+        )
+
+    chunk_decoded_frames = tokens_chunk_size * int(model.vae_ratio_t)  # type: ignore[attr-defined]
+    split_count = int(model.token_drop > 0) + 1  # type: ignore[attr-defined]
+    logical_chunk_frames: list[int] = []
+    final_overlap_frames = 0
+    for temporal_index in range(num_chunks):
+        token_start = temporal_index * tokens_chunk_size
+        token_stop = token_start + tokens_chunk_size + int(model.token_overlap)  # type: ignore[attr-defined]
+        clip_tokens = max(
+            0,
+            min(token_stop, int(temporal_latent.shape[2])) - min(token_start, int(temporal_latent.shape[2])),
+        )
+        clip_frames = clip_tokens * int(model.vae_ratio_t)  # type: ignore[attr-defined]
+        for split_index in range(split_count):
+            frame_start = split_index * chunk_decoded_frames
+            frame_stop = min(frame_start + chunk_decoded_frames, clip_frames)
+            part_frames = max(
+                0,
+                frame_stop - frame_start - int(model.frame_pre_padding),  # type: ignore[attr-defined]
+            )
+            if split_index == 0:
+                if part_frames > 0:
+                    logical_chunk_frames.append(part_frames)
+            else:
+                final_overlap_frames = part_frames
+    if final_overlap_frames > 0:
+        logical_chunk_frames.append(final_overlap_frames)
+    if sum(logical_chunk_frames) != int(total_frames):
+        raise RuntimeError(
+            "MiniMax-H3 compatibility planner disagrees with released remote code: "
+            f"frames={sum(logical_chunk_frames)}/{total_frames}"
+        )
+    chunk_frames = _truncate_chunk_frames(logical_chunk_frames, int(output_frames))
+    return MiniMaxH3TemporalDecodePlan(
+        latent=temporal_latent,
+        num_chunks=num_chunks,
+        total_frames=int(total_frames),
+        pad_frames=int(pad_frames),
+        output_frames=int(output_frames),
+        chunk_frames=chunk_frames,
+        temporal_dtype=temporal_dtype,
+    )
+
+
+def resolve_minimax_h3_temporal_cat_dtype(model: nn.Module) -> torch.dtype | None:
+    module = importlib.import_module(type(model).__module__)
+    resolver = getattr(module, "_resolve_temporal_cat_dtype", None)
+    if resolver is None:
+        return None
+    if not callable(resolver):
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            "Loaded MiniMax-H3 VAE has a non-callable temporal dtype resolver"
+        )
+    try:
+        dtype = resolver()
+    except ValueError as exc:
+        raise MiniMaxH3TemporalChunkCompatibilityError(str(exc)) from None
+    if dtype is not None and not isinstance(dtype, torch.dtype):
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            f"Loaded MiniMax-H3 VAE returned invalid temporal dtype {dtype!r}"
+        )
+    return dtype
+
+
+def decode_minimax_h3_temporal_chunks_compat(
+    model: nn.Module,
+    plan: MiniMaxH3TemporalDecodePlan,
+    callback: _TemporalChunkCallbackSignature,
+) -> torch.Tensor:
+    """Mirror the released Apache-2.0 H3 temporal assembler with callbacks.
+
+    Keep this implementation and the source/config fingerprints above in sync
+    with ``MiniMaxAI/MiniMax-H3@42ed227e``. Unknown remote code fails closed
+    before this function is called.
+
+    Callback tensors are borrowed decoder-domain views and must not be retained
+    or mutated. The high-level adapter owns normalization and storage isolation.
+    """
+
+    chunk_decoded_frames = int(model.tokens_chunk_size) * int(model.vae_ratio_t)  # type: ignore[attr-defined]
+    split_count = int(model.token_drop > 0) + 1  # type: ignore[attr-defined]
+    assembly = _RawChunkAssembly(callback=callback, plan=plan)
+    overlap: torch.Tensor | None = None
+
+    for temporal_index in range(plan.num_chunks):
+        token_start = temporal_index * int(model.tokens_chunk_size)  # type: ignore[attr-defined]
+        token_stop = token_start + int(model.tokens_chunk_size) + int(model.token_overlap)  # type: ignore[attr-defined]
+        chunk_latent = plan.latent[:, :, token_start:token_stop]
+        decoded = model._adaptive_decode(chunk_latent)  # type: ignore[attr-defined]
+        if plan.temporal_dtype is not None and decoded.dtype != plan.temporal_dtype:
+            decoded = decoded.to(plan.temporal_dtype)
+        if decoded.device != plan.latent.device:
+            decoded = decoded.to(plan.latent.device)
+
+        for split_index in range(split_count):
+            frame_start = split_index * chunk_decoded_frames
+            frame_stop = min(frame_start + chunk_decoded_frames, int(decoded.shape[2]))
+            part = decoded[:, :, frame_start:frame_stop]
+            part = part[:, :, int(model.frame_pre_padding) :]  # type: ignore[attr-defined]
+            if split_index == 0:
+                if overlap is not None:
+                    part = model.blend(  # type: ignore[attr-defined]
+                        overlap,
+                        part,
+                        int(model.frame_overlap),  # type: ignore[attr-defined]
+                        dim=-3,
+                    )
+                    overlap = None
+                assembly.write(part)
+            else:
+                overlap = part.contiguous()
+            del part
+
+        if temporal_index == plan.num_chunks - 1 and overlap is not None:
+            assembly.write(overlap)
+            overlap = None
+        del decoded, chunk_latent
+
+    return assembly.finish()
+
+
+def prepare_minimax_h3_temporal_chunks_compat(
+    model: nn.Module,
+    latent: torch.Tensor,
+    *,
+    validate_sources: bool = True,
+) -> MiniMaxH3TemporalDecodePlan:
+    """Build a validated plan without entering decoder collectives."""
+
+    if validate_sources:
+        _validate_released_sources(model)
+    _validate_released_config(model)
+    _validate_compat_contract(model, latent)
+    temporal_dtype = resolve_minimax_h3_temporal_cat_dtype(model)
+    return _build_decode_plan(model, latent, temporal_dtype)
+
+
+__all__ = [
+    "MINIMAX_H3_RELEASED_REVISION",
+    "MINIMAX_H3_RELEASED_BUNDLE_SHA256",
+    "MiniMaxH3TemporalDecodePlan",
+    "MiniMaxH3TemporalChunkCompatibilityError",
+    "decode_minimax_h3_temporal_chunks_compat",
+    "prepare_minimax_h3_temporal_chunks_compat",
+    "resolve_minimax_h3_temporal_cat_dtype",
+]

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -21,12 +21,22 @@ from vllm.logger import init_logger
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
     DistributedVaeMixin,
 )
-from vllm_omni.diffusion.distributed.parallel_state import get_world_group
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_data_parallel_world_size,
+    get_world_group,
+    model_parallel_is_initialized,
+)
 from vllm_omni.diffusion.offloader.module_residency import (
     BoundedAllocatorCache,
     PinnedModuleStager,
 )
 
+from .chunked_decode import (
+    MiniMaxH3ChunkCallbackPeerError,
+    MiniMaxH3ChunkedDecodeUnsupportedError,
+    MiniMaxH3VideoChunkCallback,
+    _MiniMaxH3ChunkDecodeCoordinator,
+)
 from .ops import install_h3_vae_optimizations
 from .packed_tokens import minimax_h3_patchify_video_latent
 
@@ -153,6 +163,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         self.use_slicing = False
         self.parallel_size = 1
         self.device_module = torch.get_device_module()
+        self._chunk_decode_coordinator: _MiniMaxH3ChunkDecodeCoordinator | None = None
 
     def load_to_device(self) -> None:
         if self._stager is not None:
@@ -265,6 +276,59 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
     def is_distributed_enabled(self) -> bool:
         return self.parallel_size > 1 and dist.is_initialized()
 
+    def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        channels = int(self.config_dict["latent_channels"])
+        mean = torch.tensor(
+            self.config_dict["latents_mean"],
+            device=latent.device,
+            dtype=latent.dtype,
+        ).view(1, channels, 1, 1, 1)
+        std = torch.tensor(
+            self.config_dict["latents_std"],
+            device=latent.device,
+            dtype=latent.dtype,
+        ).view(1, channels, 1, 1, 1)
+        return latent * std + mean
+
+    def _decode_tiling_context(self, latent: torch.Tensor) -> AbstractContextManager:
+        # The checkpoint hands rank r the tiles ``range(r, num_tiles, sp_size)``
+        # and then rejects an empty share inside the gather. A rank with no
+        # tiles raises and leaves the collective while the others block in it
+        # forever, so too few tiles hangs the whole stage rather than failing
+        # it. Tile count depends only on the latent shape, so every rank takes
+        # this branch together.
+        num_tiles = self._decoder_tile_count(latent)
+        if self.parallel_size > 1 and num_tiles < self.parallel_size:
+            logger.warning_once(
+                "MiniMax-H3 VAE decode splits into %d tile(s) but the tile group has "
+                "%d ranks; decoding rank-locally for this shape instead, which is "
+                "slower but avoids ranks without tiles hanging the collective.",
+                num_tiles,
+                self.parallel_size,
+            )
+            return self._rank_local_tiling()
+        return nullcontext()
+
+    def _normalize_decoded_frames(self, decoded: torch.Tensor) -> torch.Tensor:
+        frames = self.model.processor.revert_tensor(decoded)
+        if frames.ndim == 4:
+            frames = frames.unsqueeze(0).transpose(1, 2)
+        if frames.ndim != 5:
+            raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}")
+        return frames.float()
+
+    def _chunk_process_group(self) -> dist.ProcessGroup | None:
+        if not self.is_distributed_enabled():
+            return None
+        if model_parallel_is_initialized() and get_data_parallel_world_size() > 1:
+            raise MiniMaxH3ChunkedDecodeUnsupportedError(
+                "MiniMax-H3 temporal chunks do not yet support combining data parallelism with VAE patch parallelism"
+            )
+        group = self._native_parallel_state().get("sp_process_group")
+        if group is None or dist.get_world_size(group) != self.parallel_size:
+            raise RuntimeError("MiniMax-H3 VAE chunk decode has an invalid spatial-parallel group")
+        return group
+
     @torch.inference_mode()
     def encode_image(self, image: Image.Image) -> torch.Tensor:
         previous_parallel = self.model.parallel_tiling
@@ -357,44 +421,40 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
 
     @torch.inference_mode()
     def decode_latent(self, latent: torch.Tensor) -> torch.Tensor:
-        channels = int(self.config_dict["latent_channels"])
+        with self._decode_tiling_context(latent):
-        mean = torch.tensor(
+            decoded = self.model.decode_base(self._denormalize_latent(latent))
-            self.config_dict["latents_mean"],
+        return self._normalize_decoded_frames(decoded)
-            device=latent.device,
-            dtype=latent.dtype,
-        ).view(1, channels, 1, 1, 1)
-        std = torch.tensor(
-            self.config_dict["latents_std"],
-            device=latent.device,
-            dtype=latent.dtype,
-        ).view(1, channels, 1, 1, 1)
-        # The checkpoint hands rank r the tiles ``range(r, num_tiles, sp_size)``
-        # and then rejects an empty share inside the gather. A rank with no
-        # tiles raises and leaves the collective while the others block in it
-        # forever, so too few tiles hangs the whole stage rather than failing
-        # it. Tile count depends only on the latent shape, so every rank takes
-        # this branch together.
-        num_tiles = self._decoder_tile_count(latent)
-        if self.parallel_size > 1 and num_tiles < self.parallel_size:
-            logger.warning_once(
-                "MiniMax-H3 VAE decode splits into %d tile(s) but the tile group has "
-                "%d ranks; decoding rank-locally for this shape instead, which is "
-                "slower but avoids ranks without tiles hanging the collective.",
-                num_tiles,
-                self.parallel_size,
-            )
-            tiling_context: AbstractContextManager = self._rank_local_tiling()
-        else:
-            tiling_context = nullcontext()
 
-        with tiling_context:
+    @torch.inference_mode()
-            decoded = self.model.decode_base(latent * std + mean)
+    def decode_latent_with_chunks(
-        frames = self.model.processor.revert_tensor(decoded)
+        self,
-        if frames.ndim == 4:
+        latent: torch.Tensor,
-            frames = frames.unsqueeze(0).transpose(1, 2)
+        chunk_callback: MiniMaxH3VideoChunkCallback | None,
-        if frames.ndim != 5:
+    ) -> torch.Tensor:
-            raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}")
+        """Decode a video while publishing committed temporal chunks.
-        return frames.float()
+
+        Every rank collaborating in the configured H3 VAE tile group must call
+        this method, with ``chunk_callback`` supplied only on group rank 0. A
+        callback failure is deferred until the compatibility decoder finishes its
+        remaining collectives, then propagated to every group rank. With
+        ``parallel_size=1`` the runtime designates each request's output owner
+        by supplying a callback there and ``None`` on replica peers.
+
+        This is a model-local producer seam before the typed media boundary in
+        RFC #6541 / PR #6615. Representation conversion, transport,
+        presentation policy, and encoding remain runtime-owned. Chunks retain
+        the VAE canvas dimensions; a runtime consumer must apply the same
+        requested-size crop as the complete H3 pipeline output.
+        """
+
+        if self._chunk_decode_coordinator is None:
+            self._chunk_decode_coordinator = _MiniMaxH3ChunkDecodeCoordinator()
+        return self._chunk_decode_coordinator.decode(
+            self,
+            latent,
+            chunk_callback,
+            group=self._chunk_process_group(),
+        )
 
 
 class MiniMaxH3AudioVAE(nn.Module):
@@ -516,4 +576,10 @@ class MiniMaxH3AudioVAE(nn.Module):
         return waveform.permute(1, 0, 2).contiguous().float()
 
 
-__all__ = ["MiniMaxH3AudioVAE", "MiniMaxH3VideoVAE"]
+__all__ = [
+    "MiniMaxH3AudioVAE",
+    "MiniMaxH3ChunkCallbackPeerError",
+    "MiniMaxH3ChunkedDecodeUnsupportedError",
+    "MiniMaxH3VideoChunkCallback",
+    "MiniMaxH3VideoVAE",
+]

vllm_omni/diffusion/models/sensenova_u1/paged_decode.py

@@ -0,0 +1,368 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Paged KV cache for SenseNova-U1 autoregressive decode, and a CUDA graph over it.
+
+Autoregressive decode dominates a think request -- 80% of the wall clock, at
+about 20 ms per token for 42 layers -- and almost all of the GPU idle in that
+window is Python/ATen dispatch rather than any CUDA call. A CUDA graph retires
+those dispatches, but capture needs static shapes, and the default cache grows
+K/V with ``torch.cat`` on every step.
+
+Padding the cache to a bucket and masking the tail does give static shapes, but
+the mask is what costs: measured on an A800 at real decode shapes, a masked
+bucket attention runs 7.96-11.55 ms per step against 1.03 ms for an exact-length
+unmasked one, which is more than the graph saves.
+
+The way out is a paged cache. ``flash_attn_varlen_func`` takes the used length
+as a *tensor* (``seqused_k``) alongside a ``block_table``, so the buffers stay
+bucket-sized and capturable while the kernel reads only the valid prefix. One
+captured graph then serves every length in the bucket.
+
+Scope, and when to delete this. The cache is model-local on purpose:
+``DiffusionKVCacheManager`` reserves once per scheduler request, and
+``ARDiffusionModelRunner`` lives under ``vllm_omni/experimental``. SenseNova runs
+its whole autoregressive loop inside one pipeline ``forward``, so the decode
+steps never surface to the scheduler and it cannot admit, pool, evict, reuse a
+prefix for, or continuously batch them. What is here is therefore one set of
+buffers behind an identity block table, reused by whichever request fits them,
+not general paged-KV support: it holds while the pipeline serves one sequence
+per forward, and is released when sleep discards the memory it captured
+against. Once the decode loop is
+driven by a scheduler-visible runner, or the pipeline handles more than one
+sequence per forward, delete this file and use the manager instead.
+"""
+
+from __future__ import annotations
+
+import inspect
+
+import torch
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+BLOCK_SIZE = 16
+# Growth schedule. Each new bucket costs one capture, so keep the list short and
+# let the last entry cover the tail rather than doubling forever.
+BUCKETS = (512, 1024, 2048, 4096, 8192)
+# Past the last bucket the schedule grows by this much at a time. A request
+# decodes at most `max_think_tokens` (1024) or `max_tokens` (512) steps, so a
+# step this size costs at most one reallocation per request. Doubling the last
+# bucket instead would also cost at most one, but a 9,100-token request would
+# then reserve 766 MiB it never uses, against 118 MiB here.
+TAIL_STEP = 2048
+
+
+# Every argument to the kernel is passed by keyword, so the probe covers all of
+# them. An older wheel can export `flash_attn_varlen_func` without the paged
+# ones, or under different names, and a version string would not say which;
+# probe the signature and fall back to SDPA when any is missing.
+_REQUIRED_KWARGS = (
+    "max_seqlen_q",
+    "cu_seqlens_q",
+    "max_seqlen_k",
+    "seqused_k",
+    "block_table",
+    "softmax_scale",
+    "causal",
+)
+
+
+def _flash_varlen():
+    """Import lazily: the bundled kernel is CUDA-only and optional."""
+    from vllm.vllm_flash_attn import flash_attn_varlen_func
+
+    return flash_attn_varlen_func
+
+
+def _accepts_paged_kwargs(fn) -> bool:
+    try:
+        params = inspect.signature(fn).parameters
+    except (TypeError, ValueError):  # pragma: no cover - no introspectable signature
+        return False
+    return all(name in params for name in _REQUIRED_KWARGS)
+
+
+def paged_decode_supported(device: torch.device, head_dim: int) -> bool:
+    if device.type != "cuda":
+        return False
+    if head_dim % 8 or head_dim > 256:
+        return False
+    try:
+        fn = _flash_varlen()
+    except Exception as exc:  # pragma: no cover - depends on the wheel
+        logger.debug("Paged decode unavailable: %s", exc)
+        return False
+    if not _accepts_paged_kwargs(fn):
+        logger.debug("Paged decode unavailable: flash_attn_varlen_func lacks %s", _REQUIRED_KWARGS)
+        return False
+    return True
+
+
+def dynamic_lora_wrappers_present(module) -> bool:
+    """True once ``DiffusionLoRAManager`` has wrapped layers under ``module``.
+
+    The manager replaces linear layers with ``BaseLayerWithLoRA`` and then
+    binds, rescales or resets slot 0 per request; the wrappers stay in the tree
+    afterwards. Capture records the module tree and the branches that ran during
+    it, so a graph taken before an adapter was bound would replay without the
+    adapter's matmuls and one taken with it bound would survive its removal --
+    both silently, and neither visible in the shape, dtype and bucket that decide
+    reuse. Nothing captured may outlive a request once these exist.
+
+    The distilled LoRA is not affected: it is fused one-way into the weights and
+    leaves no wrapper behind.
+    """
+    try:
+        from vllm.lora.layers import BaseLayerWithLoRA
+    except ImportError:  # pragma: no cover - depends on the wheel
+        return False
+    return any(isinstance(m, BaseLayerWithLoRA) for m in module.modules())
+
+
+def _bucket_for(length: int) -> int:
+    for b in BUCKETS:
+        if length <= b:
+            return b
+    # Past the schedule, grow in TAIL_STEP-sized steps. Rounding to a whole
+    # number of blocks here would reallocate and re-capture every BLOCK_SIZE
+    # tokens: a think edit with two 2048x2048 inputs runs to about 9,100 tokens
+    # and paid 41 reallocations and 42 captures in one request.
+    over = length - BUCKETS[-1]
+    return BUCKETS[-1] + ((over + TAIL_STEP - 1) // TAIL_STEP) * TAIL_STEP
+
+
+class PagedDecodeCache:
+    """Per-layer paged K/V for single-token decode.
+
+    Layout follows what ``flash_attn_varlen_func`` accepts with a block table:
+    ``(num_blocks, BLOCK_SIZE, kv_heads, head_dim)``. ``seqused`` holds the
+    number of valid tokens and is the only thing that changes between decode
+    steps, which is what lets a captured graph be replayed unchanged.
+    """
+
+    def __init__(self, num_layers, kv_heads, head_dim, length, device, dtype):
+        self.kv_heads = kv_heads
+        self.head_dim = head_dim
+        self.device = device
+        self.dtype = dtype
+        self.bucket = _bucket_for(length)
+        nblocks = self.bucket // BLOCK_SIZE
+        shape = (nblocks, BLOCK_SIZE, kv_heads, head_dim)
+        self.k = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(num_layers)]
+        self.v = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(num_layers)]
+        self.block_table = torch.arange(nblocks, device=device, dtype=torch.int32).unsqueeze(0)
+        self.cu_seqlens_q = torch.tensor([0, 1], device=device, dtype=torch.int32)
+        self.seqused = torch.zeros(1, device=device, dtype=torch.int32)
+        # The write slot has to live in a tensor, not a Python int. Under graph
+        # capture a Python index is baked in at capture time, so every replay
+        # would overwrite the same slot -- correct eagerly, silently wrong once
+        # captured.
+        self.pos = torch.zeros(1, device=device, dtype=torch.int64)
+        self._length = 0
+        # Bumped whenever the buffers are reallocated, so a graph runner can
+        # tell that its capture no longer points at live memory.
+        self.generation = 0
+
+    @property
+    def length(self) -> int:
+        return self._length
+
+    def set_length(self, n: int) -> None:
+        self._length = n
+        self.seqused.fill_(n)
+        self.pos.fill_(n - 1)
+
+    def reusable_for(self, num_layers, kv_heads, head_dim, dtype, length) -> bool:
+        """Can this cache serve another request without reallocating?
+
+        Reusing the buffers is what lets captured graphs survive across
+        requests: a capture binds the addresses it recorded, so a fresh cache
+        forces a fresh capture.
+        """
+        return (
+            len(self.k) == num_layers
+            and self.kv_heads == kv_heads
+            and self.head_dim == head_dim
+            and self.dtype == dtype
+            and _bucket_for(length) <= self.bucket
+        )
+
+    def load_prefix(self, cache) -> bool:
+        """Copy a prefill cache into the paged buffers.
+
+        Only the first ``seqused`` rows are ever read, so the tail is left as
+        whatever the previous request wrote. Returns False for an empty cache.
+        """
+        layer0 = cache.layers[0]
+        if layer0.keys is None:
+            return False
+        prefix = layer0.keys.shape[2]
+        for i in range(len(self.k)):
+            keys, values = cache.layers[i].keys, cache.layers[i].values
+            # [B, H, S, D] -> [S, H, D], then into the paged blocks
+            self._write_prefix(i, keys[0].transpose(0, 1), values[0].transpose(0, 1))
+        self.set_length(prefix)
+        return True
+
+    @classmethod
+    def from_dynamic_cache(cls, cache, num_layers, device, dtype, min_length=0):
+        """Allocate buffers for a prefill cache and copy it in."""
+        layer0 = cache.layers[0]
+        if layer0.keys is None:
+            return None
+        _, kv_heads, prefix, head_dim = layer0.keys.shape
+        obj = cls(num_layers, kv_heads, head_dim, max(prefix + 1, min_length), device, dtype)
+        obj.load_prefix(cache)
+        return obj
+
+    def _write_prefix(self, layer_idx, keys_shd, values_shd):
+        n = keys_shd.shape[0]
+        flat_k = self.k[layer_idx].view(-1, self.kv_heads, self.head_dim)
+        flat_v = self.v[layer_idx].view(-1, self.kv_heads, self.head_dim)
+        flat_k[:n].copy_(keys_shd)
+        flat_v[:n].copy_(values_shd)
+
+    def grow(self, length: int) -> bool:
+        """Move to the next bucket, keeping the tokens already stored.
+
+        Everything is built into locals first and published in one step. A
+        resize that failed part way -- an allocation for a later layer, say --
+        would otherwise leave new K beside old V under the ``generation`` a
+        captured graph is keyed on, so the next request would keep selecting
+        that graph and replay it over storage that had just been freed.
+        """
+        new_bucket = _bucket_for(length)
+        if new_bucket == self.bucket:
+            return False
+        nblocks = new_bucket // BLOCK_SIZE
+        shape = (nblocks, BLOCK_SIZE, self.kv_heads, self.head_dim)
+        old_len = self._length
+        new_k, new_v = [], []
+        for old_k, old_v in zip(self.k, self.v):
+            k = torch.zeros(shape, device=self.device, dtype=self.dtype)
+            v = torch.zeros(shape, device=self.device, dtype=self.dtype)
+            k.view(-1, self.kv_heads, self.head_dim)[:old_len].copy_(
+                old_k.view(-1, self.kv_heads, self.head_dim)[:old_len]
+            )
+            v.view(-1, self.kv_heads, self.head_dim)[:old_len].copy_(
+                old_v.view(-1, self.kv_heads, self.head_dim)[:old_len]
+            )
+            new_k.append(k)
+            new_v.append(v)
+        block_table = torch.arange(nblocks, device=self.device, dtype=torch.int32).unsqueeze(0)
+
+        # Publish. New buffers mean any graph captured against the old ones is
+        # stale, so the generation moves in the same step that hands them over.
+        self.k, self.v, self.block_table, self.bucket = new_k, new_v, block_table, new_bucket
+        self.generation += 1
+        self.set_length(old_len)
+        logger.debug("Paged decode grew to bucket %d at length %d (generation %d)", new_bucket, length, self.generation)
+        return True
+
+    def attend(self, layer_idx, query_bhsd, key_bhsd, value_bhsd, softmax_scale):
+        """Append this step's K/V and attend over the whole valid prefix.
+
+        Writes at ``length - 1`` because the caller has already advanced the
+        length for this step. Both the slot and the attended length come from
+        device tensors, so a captured graph reads whatever they hold at replay
+        rather than whatever they held at capture.
+        """
+        flat_k = self.k[layer_idx].view(-1, self.kv_heads, self.head_dim)
+        flat_v = self.v[layer_idx].view(-1, self.kv_heads, self.head_dim)
+        flat_k.index_copy_(0, self.pos, key_bhsd[0, :, 0].unsqueeze(0))
+        flat_v.index_copy_(0, self.pos, value_bhsd[0, :, 0].unsqueeze(0))
+        heads, head_dim = query_bhsd.shape[1], query_bhsd.shape[3]
+        q = query_bhsd[0].transpose(0, 1).reshape(1, heads, head_dim)
+        out = _flash_varlen()(
+            q,
+            self.k[layer_idx],
+            self.v[layer_idx],
+            max_seqlen_q=1,
+            cu_seqlens_q=self.cu_seqlens_q,
+            max_seqlen_k=self.bucket,
+            seqused_k=self.seqused,
+            block_table=self.block_table,
+            softmax_scale=softmax_scale,
+            causal=False,
+        )
+        return out.reshape(1, 1, heads, head_dim)
+
+    def to_dynamic_cache(self, cache) -> None:
+        """Write the decoded K/V back so the generation stage sees one cache.
+
+        Only decode is paged; the DiT stage that follows reads the ordinary
+        cache, so the two are reconciled once at the hand-off rather than kept
+        in sync on every step.
+        """
+        n = self._length
+        for i, layer in enumerate(cache.layers):
+            keys = self.k[i].view(-1, self.kv_heads, self.head_dim)[:n]
+            values = self.v[i].view(-1, self.kv_heads, self.head_dim)[:n]
+            layer.keys = keys.transpose(0, 1).unsqueeze(0).contiguous()
+            layer.values = values.transpose(0, 1).unsqueeze(0).contiguous()
+
+
+class DecodeGraphRunner:
+    """One captured decode step per KV bucket, replayed for every token.
+
+    The whole point of the paged cache is that everything the step reads which
+    varies -- the token, its position, and how much of the cache is live -- sits
+    in device tensors. So a single capture serves every step in a bucket: fill
+    the tensors, replay, read the logits out of the static output.
+
+    A capture is invalidated when the cache reallocates (tracked by
+    ``PagedDecodeCache.generation``), which happens once per bucket boundary.
+    """
+
+    def __init__(self, language_model, cache, device):
+        self.lm = language_model
+        self.cache = cache
+        self.device = device
+        self.input_ids = torch.zeros(1, 1, dtype=torch.long, device=device)
+        self.indexes = torch.zeros(3, 1, dtype=torch.long, device=device)
+        self._graphs: dict[tuple[int, int], tuple[torch.cuda.CUDAGraph, torch.Tensor]] = {}
+        self.captures = 0
+
+    def _forward(self):
+        return self.lm(
+            input_ids=self.input_ids,
+            indexes=self.indexes,
+            past_key_values=None,
+            use_cache=False,
+            paged_cache=self.cache,
+        ).logits
+
+    def _capture(self, key):
+        # Warm up on a side stream first so cuBLAS workspaces and any lazy
+        # allocation happen outside the capture. The warm-up recomputes this
+        # same step, writing the same K/V to the same slot, so it is idempotent.
+        side = torch.cuda.Stream()
+        side.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(side):
+            for _ in range(3):
+                self._forward()
+        torch.cuda.current_stream().wait_stream(side)
+
+        graph = torch.cuda.CUDAGraph()
+        # A private pool would hand every runner its own arena and never give it
+        # back, and a runner is rebuilt whenever the cache reallocates or a LoRA
+        # wrapper appears: measured, that grew device memory by about 40 MiB per
+        # request and never levelled off. The platform pool is shared, which is
+        # what every other captured path in the tree uses.
+        with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
+            logits = self._forward()
+        self.captures += 1
+        logger.debug("Captured decode graph for bucket=%d generation=%d", self.cache.bucket, self.cache.generation)
+        self._graphs[key] = (graph, logits)
+        return self._graphs[key]
+
+    def step(self, token, t_index):
+        """Run one decode step. Returns the static logits tensor."""
+        self.input_ids[0, 0] = token
+        self.indexes[0, 0] = t_index
+        key = (self.cache.bucket, self.cache.generation)
+        entry = self._graphs.get(key) or self._capture(key)
+        entry[0].replay()
+        return entry[1]

vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """SenseNova-U1 Pipeline for vLLM-Omni.
 
 SenseNova-U1 is a unified Qwen3-based model that uses Mixture-of-Tokenizers
@@ -35,6 +35,13 @@ from vllm.logger import init_logger
 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.lora.loader import (
+    LoraLoaderMixin,
+    _apply_diffusers_lora_alpha_scaling,
+    _load_lora_state_dict,
+    _prepare_lora_delta,
+    _remap_state_dict_keys,
+)
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
@@ -44,7 +51,14 @@ from vllm_omni.transformers_utils.configs.sensenova_u1 import (
     SenseNovaU1Config,
 )
 
+from .paged_decode import (
+    DecodeGraphRunner,
+    PagedDecodeCache,
+    dynamic_lora_wrappers_present,
+    paged_decode_supported,
+)
 from .sensenova_u1_transformer import (
+    SenseNovaU1CausalLMOutput,
     SenseNovaU1ForCausalLM,
     clear_flash_kv_cache,
     create_block_causal_mask,
@@ -446,7 +460,13 @@ class SenseNovaU1DenoisingAdapter(nn.Module):
         return self.language_model(*args, **kwargs)
 
 
-class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin, CFGParallelMixin):
+class SenseNovaU1Pipeline(
+    nn.Module,
+    SupportsComponentDiscovery,
+    DiffusionPipelineProfilerMixin,
+    CFGParallelMixin,
+    LoraLoaderMixin,
+):
     """SenseNova-U1 text-to-image and image-to-image pipeline for vllm-omni.
 
     Builds the full model graph internally:
@@ -709,8 +729,71 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         sigma = shift * sigma / (1 + (shift - 1) * sigma)
         return 1 - sigma
 
-    def _ar_step(self, next_token, t_idx, past_key_values):
+    def _decode_context(self, past_key_values):
+        """Paged cache plus a captured graph for the AR loop, or ``None``.
+
+        Decode is where a think request spends most of its time, and every step
+        issues the same work at the same shapes, so it is worth capturing. The
+        cache has to be paged for that: it keeps the buffers a fixed size while
+        the kernel reads only ``seqused_k`` of them, which is what lets one
+        capture serve every step until the next bucket.
+
+        Both live on the pipeline and are reused across requests unless a LoRA
+        adapter is being served dynamically. A capture binds the addresses it
+        recorded, so a cache built per request forces a capture per request:
+        measured, that left about 40 MiB of device memory behind on every
+        request and re-captured every time. Reuse rests on the same thing the
+        paged path already rests on -- one sequence in flight per pipeline
+        forward.
+        """
+        if os.environ.get("VLLM_OMNI_SENSENOVA_PAGED_DECODE", "1") != "1":
+            return None
+        head_dim = self.language_model.model.layers[0].self_attn.head_dim
+        if not paged_decode_supported(torch.device(self.device), head_dim):
+            return None
+        if past_key_values is None or not past_key_values.layers:
+            return None
+        layer0 = past_key_values.layers[0]
+        if layer0.keys is None:
+            return None
+        _, kv_heads, prefix, kv_head_dim = layer0.keys.shape
+        device = torch.device(self.device)
+        num_layers = len(self.language_model.model.layers)
+        # Reuse is off once the LoRA manager holds the decode path: the test
+        # below sees shape, dtype and bucket, none of which move when an adapter
+        # is bound, rescaled or reset between requests.
+        reuse = not dynamic_lora_wrappers_present(self.language_model)
+        decode = getattr(self, "_paged_decode", None) if reuse else None
+        if decode is None or not decode[0].reusable_for(
+            num_layers, kv_heads, kv_head_dim, layer0.keys.dtype, prefix + 1
+        ):
+            cache = PagedDecodeCache(num_layers, kv_heads, kv_head_dim, prefix + 1, device, layer0.keys.dtype)
+            decode = (cache, DecodeGraphRunner(self.language_model, cache, device))
+            self._paged_decode = decode if reuse else None
+        if not decode[0].load_prefix(past_key_values):
+            return None
+        return decode
+
+    def release_captured_graphs(self) -> None:
+        """Drop the reused paged cache and the graphs captured against it.
+
+        Sleep level 2 discards the memory a capture recorded, so anything held
+        across requests has to go with it. The next request rebuilds both.
+        """
+        self._paged_decode = None
+
+    def _ar_step(self, next_token, t_idx, past_key_values, decode=None):
         """One autoregressive token step. Constructs single-token `indexes` explicitly."""
+        if decode is not None:
+            cache, runner = decode
+            if cache.length + 1 > cache.bucket:
+                cache.grow(cache.length + 1)
+            cache.set_length(cache.length + 1)
+            logits = runner.step(int(next_token), t_idx + 1)
+            # The caller reassigns its cache handle from this field; the paged
+            # tokens are merged into it once, at the hand-off below.
+            return SenseNovaU1CausalLMOutput(logits=logits, past_key_values=past_key_values)
+
         indexes = torch.tensor([[t_idx + 1], [0], [0]], dtype=torch.long, device=self.device)
         outputs = self.language_model(
             input_ids=next_token.unsqueeze(0),
@@ -725,24 +808,28 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         think_end_token_id = self.tokenizer.convert_tokens_to_ids("</think>")
         think_token_ids = []
         next_token = torch.argmax(prefix_outputs.logits[:, -1, :], dim=-1)
+        decode = self._decode_context(past_key_values)
 
         for _ in range(max_think_tokens):
             token_item = next_token.item()
             if token_item == eos_token_id:
                 break
             if token_item == think_end_token_id:
-                outputs = self._ar_step(next_token, t_idx, past_key_values)
+                outputs = self._ar_step(next_token, t_idx, past_key_values, decode)
                 past_key_values = outputs.past_key_values
                 t_idx += 1
                 think_token_ids.append(token_item)
                 break
 
             think_token_ids.append(token_item)
-            outputs = self._ar_step(next_token, t_idx, past_key_values)
+            outputs = self._ar_step(next_token, t_idx, past_key_values, decode)
             past_key_values = outputs.past_key_values
             t_idx += 1
             next_token = torch.argmax(outputs.logits[:, -1, :], dim=-1)
 
+        if decode is not None:
+            decode[0].to_dynamic_cache(past_key_values)
+
         # Append "\n\n<img>" tokens to cache
         append_ids = self.tokenizer(
             "\n\n<img>",
@@ -797,6 +884,10 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         """Autoregressive text decoding seeded from prefix logits."""
         eos_token_id = self.tokenizer.convert_tokens_to_ids("<|im_end|>")
         generated_ids: list[int] = []
+        # Same paged decode as the think loop. No write-back at the end: the
+        # caller returns the decoded string and never reads the cache again,
+        # unlike the think path, which hands it to the generation stage.
+        decode = self._decode_context(past_key_values)
 
         logits = prefix_logits[:, -1, :]
         for _ in range(max_tokens):
@@ -810,7 +901,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
             if token_id == eos_token_id:
                 break
             generated_ids.append(token_id)
-            outputs = self._ar_step(next_token, t_idx, past_key_values)
+            outputs = self._ar_step(next_token, t_idx, past_key_values, decode)
             past_key_values = outputs.past_key_values
             logits = outputs.logits[:, -1, :]
             t_idx += 1
@@ -1174,8 +1265,37 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
             layer.values = layer.values.expand(batch_size, *layer.values.shape[1:])
         prepare_flash_kv_cache(kv, current_len=token_hw, batch_size=batch_size)
 
+    @staticmethod
+    def _is_warmup_request(req) -> bool:
+        is_dummy_run = getattr(req, "is_dummy_run", None)
+        if callable(is_dummy_run):
+            return bool(is_dummy_run())
+        return OmniDiffusionRequest.is_dummy_run_request_id(getattr(req, "request_id", None))
+
+    def _warm_ar_decode(self) -> None:
+        """Run one single-token decode at startup so its compiled region is built.
+
+        The engine's dummy request is a text-to-image run with think off. It
+        exercises the prefill shape but never the decode one, so whatever the
+        deploy config compiles for decode lands on the first real request
+        instead of on readiness. Warmup is best effort and must not fail startup.
+        """
+        try:
+            lm = self.language_model
+            device = torch.device(self.device)
+            ids = torch.zeros(1, 2, dtype=torch.long, device=device)
+            idx = torch.zeros(3, 2, dtype=torch.long, device=device)
+            prefill = lm(input_ids=ids, indexes=idx, use_cache=True)
+            nxt = torch.zeros(1, 1, dtype=torch.long, device=device)
+            nidx = torch.tensor([[2], [0], [0]], dtype=torch.long, device=device)
+            lm(input_ids=nxt, indexes=nidx, past_key_values=prefill.past_key_values, use_cache=True)
+        except Exception as exc:  # pragma: no cover - warmup is best effort
+            logger.warning("Autoregressive decode warmup skipped: %s", exc)
+
     @torch.inference_mode()
     def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+        if self._is_warmup_request(req):
+            self._warm_ar_decode()
         p = self._parse_request(req)
 
         input_images = self._extract_input_images(p.first_prompt)
@@ -1408,21 +1528,129 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
     # Weight loading
     # -----------------------------------------------------------------------
 
+    # Shard name -> fused parameter, used by weight loading and LoRA loading.
+    # More specific _mot_gen patterns FIRST to avoid substring ambiguity
+    # (e.g. `.q_proj` is a substring of `.q_proj_mot_gen`).
+    stacked_params_mapping: ClassVar[list[tuple[str, str, str | int]]] = [
+        (".qkv_proj_mot_gen", ".q_proj_mot_gen", "q"),
+        (".qkv_proj_mot_gen", ".k_proj_mot_gen", "k"),
+        (".qkv_proj_mot_gen", ".v_proj_mot_gen", "v"),
+        (".qkv_proj", ".q_proj", "q"),
+        (".qkv_proj", ".k_proj", "k"),
+        (".qkv_proj", ".v_proj", "v"),
+        # MLP gate/up fused into MergedColumnParallelLinear
+        (".gate_up_proj", ".gate_proj", 0),
+        (".gate_up_proj", ".up_proj", 1),
+    ]
+
+    def load_lora_weights(
+        self,
+        pretrained_model_name_or_path: str | list[str],
+        adapter_name: str | None = None,
+    ) -> None:
+        """Fuse a distilled few-step LoRA into the weights.
+
+        The checkpoints ship kohya `lora_down`/`lora_up`/`alpha` names, renamed
+        here to the Diffusers names before `load_lora_into_module` routes each
+        delta into its slice of the fused projections.
+        """
+        if adapter_name in self.lora_loaded:
+            return
+
+        paths = (
+            [pretrained_model_name_or_path]
+            if isinstance(pretrained_model_name_or_path, str)
+            else list(pretrained_model_name_or_path)
+        )
+        if len(paths) != 1:
+            raise ValueError(f"Expected exactly one distilled LoRA file, got {len(paths)}: {paths}")
+
+        state_dict = _load_lora_state_dict(paths[0])
+        state_dict = _remap_state_dict_keys(
+            state_dict,
+            [(".lora_down.weight", ".lora_A.weight"), (".lora_up.weight", ".lora_B.weight")],
+        )
+        # Fold alpha and the low-rank product in fp32; scaling B in bf16 would
+        # round once before the matmul and once after.
+        state_dict = {k: v.to(torch.float32) for k, v in state_dict.items()}
+        state_dict = _apply_diffusers_lora_alpha_scaling(state_dict)
+
+        applied = self._add_lora_deltas(state_dict)
+        if not applied:
+            raise ValueError(f"Distilled LoRA {paths[0]} matched no parameter of this pipeline.")
+        logger.info("Fused distilled LoRA %s into %d parameters", paths[0], applied)
+        # Fusion is one-way and there is no unload path, so keep only a sentinel;
+        # the fp32 state dict is ~1.5 GiB and would stay resident for the process.
+        self.lora_loaded[adapter_name] = paths[0]
+
+    def _fused_shards(self, base_key: str) -> list[tuple[str, str | int]] | None:
+        """Unfused source names and shard ids for a fused parameter, else None."""
+        for param_name, weight_name, shard_id in self.stacked_params_mapping:
+            if base_key.endswith(param_name):
+                return [
+                    (base_key[: -len(param_name)] + w, sid)
+                    for p, w, sid in self.stacked_params_mapping
+                    if p == param_name
+                ]
+        return None
+
+    def _add_lora_deltas(self, state_dict: dict[str, torch.Tensor]) -> int:
+        """Add every matching delta into its parameter, sharding through vLLM.
+
+        The delta read from the checkpoint spans the whole layer. A parameter of
+        a parallel layer only holds this rank's slice of it, so the delta goes
+        through the layer's own weight loader -- into a zeroed copy of the
+        parameter, which is then added -- instead of being added directly.
+
+        Each parameter is fused in place as it is reached, so a failure part way
+        through leaves the earlier ones fused and the pipeline has to be
+        reloaded. Deferring the writes would mean holding every delta at once,
+        and the deltas are full rank: 15.09 GiB for the 8B checkpoint, against
+        the 192 MiB one parameter costs here. Fusion runs at load time, before
+        the pipeline serves anything, so a partially fused model is never
+        reachable from a request.
+        """
+        applied = 0
+        for name, param in self.named_parameters():
+            if not name.endswith(".weight"):
+                continue
+            base_key = name[: -len(".weight")]
+            shards = self._fused_shards(base_key)
+
+            deltas: list[tuple[torch.Tensor, str | int | None]] = []
+            for src, shard_id in shards or [(base_key, None)]:
+                delta, _ = _prepare_lora_delta(state_dict, src, None)
+                if delta is None:
+                    deltas = []
+                    break
+                deltas.append((delta, shard_id))
+            if not deltas:
+                continue
+
+            weight_loader = getattr(param, "weight_loader", None)
+            buffer = torch.zeros_like(param.data)
+            if weight_loader is None:
+                # Not a parallel layer: the delta is already rank-local.
+                buffer = torch.cat([d for d, _ in deltas]).to(device=param.device, dtype=param.dtype)
+            else:
+                saved = param.data
+                try:
+                    param.data = buffer
+                    for delta, shard_id in deltas:
+                        delta = delta.to(device=param.device, dtype=param.dtype)
+                        if shard_id is None:
+                            weight_loader(param, delta)
+                        else:
+                            weight_loader(param, delta, shard_id)
+                finally:
+                    param.data = saved
+            with torch.no_grad():
+                param.data.add_(buffer)
+            applied += 1
+        return applied
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
+        stacked_params_mapping = self.stacked_params_mapping
-            # More specific _mot_gen patterns FIRST to avoid substring
-            # ambiguity (e.g. `.q_proj` is a substring of `.q_proj_mot_gen`).
-            (".qkv_proj_mot_gen", ".q_proj_mot_gen", "q"),
-            (".qkv_proj_mot_gen", ".k_proj_mot_gen", "k"),
-            (".qkv_proj_mot_gen", ".v_proj_mot_gen", "v"),
-            (".qkv_proj", ".q_proj", "q"),
-            (".qkv_proj", ".k_proj", "k"),
-            (".qkv_proj", ".v_proj", "v"),
-            # MLP gate/up fused into MergedColumnParallelLinear
-            (".gate_up_proj", ".gate_proj", 0),
-            (".gate_up_proj", ".up_proj", 1),
-        ]
-        self.stacked_params_mapping = stacked_params_mapping
 
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()

vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Qwen3 LLM with Mixture-of-Tokenizers (MoT) for SenseNova-U1.
 
 Ported from the sensenova_u1 package with vllm tensor-parallel support:
@@ -16,6 +16,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from cache_dit import ForwardPattern
 from transformers.cache_utils import DynamicCache
 from vllm.logger import init_logger
@@ -163,6 +164,18 @@ class Qwen3RotaryEmbedding(nn.Module):
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
 
+def _build_3d_rope(config) -> tuple[Qwen3RotaryEmbedding, Qwen3RotaryEmbedding]:
+    """The t and h/w rotary modules: t uses half head_dim, h/w a quarter each."""
+    t_config = copy.deepcopy(config)
+    t_config.head_dim = config.head_dim // 2
+
+    hw_config = copy.deepcopy(config)
+    hw_config.head_dim = config.head_dim // 4
+    hw_config.rope_theta = config.rope_theta_hw
+    hw_config.max_position_embeddings = config.max_position_embeddings_hw
+    return Qwen3RotaryEmbedding(config=t_config), Qwen3RotaryEmbedding(config=hw_config)
+
+
 # ---------------------------------------------------------------------------
 # Norms
 # ---------------------------------------------------------------------------
@@ -175,11 +188,9 @@ class Qwen3RMSNorm(nn.Module):
         self.variance_epsilon = eps
 
     def forward(self, hidden_states):
-        input_dtype = hidden_states.dtype
+        # F.rms_norm keeps the fp32 accumulation internal and rounds once; the
-        hidden_states = hidden_states.to(torch.float32)
+        # cast chain it replaces rounded to bf16 before the weight multiply.
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        return F.rms_norm(hidden_states, self.weight.shape, self.weight, self.variance_epsilon)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states.to(input_dtype)
 
 
 # ---------------------------------------------------------------------------
@@ -303,23 +314,16 @@ class SenseNovaU1Attention(nn.Module):
         self.q_norm_hw_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)
         self.k_norm_hw_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)
 
-        # 3D RoPE: t uses half head_dim, h/w each use quarter head_dim
-        t_config = copy.deepcopy(config)
-        t_config.head_dim = config.head_dim // 2
-        self.rotary_emb = Qwen3RotaryEmbedding(config=t_config)
-
-        hw_config = copy.deepcopy(config)
-        hw_config.head_dim = config.head_dim // 4
-        hw_config.rope_theta = config.rope_theta_hw
-        hw_config.max_position_embeddings = config.max_position_embeddings_hw
-        self.rotary_emb_hw = Qwen3RotaryEmbedding(config=hw_config)
-
         self.attn = Attention(
             num_heads=self.num_heads,
             head_size=self.head_dim,
             causal=False,
             softmax_scale=self.scaling,
-            num_kv_heads=self.num_heads,
+            # Keep K/V at their real head count. The SDPA backend already picks
+            # a fused GQA kernel when the runtime shape allows one and expands
+            # K/V itself when it does not; expanding here hides the GQA shape
+            # from that check, so it never fires.
+            num_kv_heads=self.num_kv_heads,
             prefix=f"{prefix}.attn",
         )
         self.attn.attention = self.attn.sdpa_fallback
@@ -339,10 +343,6 @@ class SenseNovaU1Attention(nn.Module):
         attention_mask: torch.Tensor | None,
     ) -> torch.Tensor:
         """Run unified attention with [B, H, S, D] inputs. Returns [B, S, H, D]."""
-        if self.num_kv_groups > 1:
-            n = self.num_kv_groups
-            key_bhsd = key_bhsd.repeat_interleave(n, dim=1)
-            value_bhsd = value_bhsd.repeat_interleave(n, dim=1)
         q = query_bhsd.transpose(1, 2).contiguous()
         k = key_bhsd.transpose(1, 2).contiguous()
         v = value_bhsd.transpose(1, 2).contiguous()
@@ -358,16 +358,16 @@ class SenseNovaU1Attention(nn.Module):
         attention_mask: torch.Tensor | None,
     ) -> torch.Tensor:
         """Run unified attention with [B, S, H, D] inputs. Returns [B, S, H, D]."""
-        if self.num_kv_groups > 1:
-            n = self.num_kv_groups
-            key_bshd = key_bshd.repeat_interleave(n, dim=2)
-            value_bshd = value_bshd.repeat_interleave(n, dim=2)
         attention_mask = self._align_mask_dtype(attention_mask, query_bshd)
         attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
         return self.attn(query_bshd, key_bshd, value_bshd, attn_metadata)
 
-    def _project_and_rope(self, hidden_states, indexes, qkv_proj, q_norm, k_norm, q_norm_hw, k_norm_hw):
+    def _project_and_rope(self, hidden_states, position_embeddings, qkv_proj, q_norm, k_norm, q_norm_hw, k_norm_hw):
-        """Project Q/K/V via the given QKVParallelLinear and apply 3D RoPE."""
+        """Project Q/K/V via the given QKVParallelLinear and apply 3D RoPE.
+
+        The three tables come from the model, which builds them once per forward
+        because every layer is handed the same `indexes`.
+        """
         input_shape = hidden_states.shape[:-1]  # (B, S)
         qkv, _ = qkv_proj(hidden_states)
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
@@ -376,9 +376,7 @@ class SenseNovaU1Attention(nn.Module):
         k = k.view(*input_shape, self.num_kv_heads, self.head_dim)
         v = v.view(*input_shape, self.num_kv_heads, self.head_dim)
 
-        cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0))
+        (cos_t, sin_t), (cos_h, sin_h), (cos_w, sin_w) = position_embeddings
-        cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0))
-        cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0))
 
         value_states = v.transpose(1, 2)  # [B, H, S, D]
         try:
@@ -433,13 +431,23 @@ class SenseNovaU1Attention(nn.Module):
         input_shape = hidden_states.shape[:-1]
         query_states, key_states, value_states = self._project_and_rope(
             hidden_states,
-            indexes,
+            kwargs["position_embeddings"],
             self.qkv_proj,
             self.q_norm,
             self.k_norm,
             self.q_norm_hw,
             self.k_norm_hw,
         )
+        # Single-token decode with a paged cache: append this step's K/V and let
+        # the kernel read the valid prefix via `seqused_k`. Shapes stay fixed,
+        # which is what makes the step capturable as a CUDA graph.
+        paged = kwargs.get("paged_cache")
+        if paged is not None and attention_mask is None and query_states.shape[2] == 1:
+            attn_output = paged.attend(self.layer_idx, query_states, key_states, value_states, self.scaling)
+            attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+            attn_output, _ = self.o_proj(attn_output)
+            return attn_output
+
         update_cache = kwargs.get("update_cache", True)
         if past_key_values is not None:
             if update_cache:
@@ -460,7 +468,7 @@ class SenseNovaU1Attention(nn.Module):
         input_shape = hidden_states.shape[:-1]
         query_states, key_states, value_states = self._project_and_rope(
             hidden_states,
-            indexes,
+            kwargs["position_embeddings"],
             self.qkv_proj_mot_gen,
             self.q_norm_mot_gen,
             self.k_norm_mot_gen,
@@ -659,6 +667,10 @@ class SenseNovaU1Model(nn.Module):
         )
         self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.norm_mot_gen = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # Every layer receives the same `indexes`, so the tables are the same for
+        # all of them: build them once per forward here and thread them down,
+        # rather than giving every layer its own pair.
+        self.rotary_emb, self.rotary_emb_hw = _build_3d_rope(config)
 
     def forward(
         self,
@@ -690,15 +702,25 @@ class SenseNovaU1Model(nn.Module):
             past_len = past_key_values.get_seq_length() if past_key_values else 0
             seq_len = inputs_embeds.shape[1]
             total_len = past_len + seq_len
-            mask = torch.zeros(1, 1, seq_len, total_len, device=inputs_embeds.device)
             if seq_len > 1:
+                mask = torch.zeros(1, 1, seq_len, total_len, device=inputs_embeds.device)
                 causal = torch.tril(torch.ones(seq_len, seq_len, device=inputs_embeds.device))
                 mask[:, :, :, past_len:] = torch.where(causal == 1, 0.0, float("-inf"))
+            else:
+                # A single query token attends to every cached key, so the mask
+                # this branch used to build was all zeros. Passing it changed no
+                # result, only which SDPA kernel was reachable.
+                mask = None
             causal_mask_mapping = {"full_attention": mask}
         else:
             causal_mask_mapping = attention_mask
 
         hidden_states = inputs_embeds
+        position_embeddings = (
+            self.rotary_emb(hidden_states, indexes[0].unsqueeze(0)),
+            self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0)),
+            self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0)),
+        )
         for layer in self.layers:
             hidden_states = layer(
                 hidden_states,
@@ -706,6 +728,7 @@ class SenseNovaU1Model(nn.Module):
                 exist_und=exist_und,
                 exist_gen=exist_gen,
                 indexes=indexes,
+                position_embeddings=position_embeddings,
                 attention_mask=causal_mask_mapping,
                 past_key_values=past_key_values,
                 use_cache=use_cache,

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -519,6 +519,23 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         if self.prompt_embed_cache is not None:
             self.prompt_embed_cache.clear()
 
+    def release_captured_graphs(self) -> None:
+        """Drop every CUDA graph held for this model, wherever it is kept.
+
+        Sleep level 2 discards the memory a capture was recorded against, so a
+        graph that outlives it replays over freed storage. The runner owns
+        compilation and execution resources, so it owns the release: a pipeline
+        that captures graphs of its own implements ``release_captured_graphs``
+        and is collected here, instead of every caller having to know which
+        pipelines have one.
+        """
+        runners = getattr(self, "graph_runners", None)
+        if runners is not None:
+            runners.clear()
+        release = getattr(self.pipeline, "release_captured_graphs", None)
+        if callable(release):
+            release()
+
     def get_prompt_embed_cache_stats(self) -> dict | None:
         """Return hit/miss statistics for the prompt-embedding cache, if enabled.
 

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -786,9 +786,8 @@ class DiffusionWorker:
         usage_before = allocator.get_current_usage()
 
         if level == 2 and self.model_runner is not None:
-            if hasattr(self.model_runner, "graph_runners"):
+            self.model_runner.release_captured_graphs()
-                self.model_runner.graph_runners.clear()
+            logger.info(f"[Worker {self.rank}] CUDA Graphs cleared.")
-                logger.info(f"[Worker {self.rank}] CUDA Graphs cleared.")
             model = self.model_runner.pipeline
             self._sleep_saved_buffers = {name: buffer.cpu().clone() for name, buffer in model.named_buffers()}
 

vllm_omni/engine/async_omni_engine.py

@@ -1349,8 +1349,8 @@ class AsyncOmniEngine:
                     or kwargs.get("fastvideo_vsa_topk") is not None
                 ):
                     has_stage_attention = (
-                        hasattr(cfg.engine_args, "diffusion_attention_config")
+                        getattr(cfg.engine_args, "diffusion_attention_config", None) is not None
-                        and cfg.engine_args.diffusion_attention_config is not None
+                        or getattr(cfg.engine_args, "diffusion_attention_backend", None) is not None
                     )
                     if not has_stage_attention:
                         cfg.engine_args.diffusion_attention_config = parse_attention_config(

vllm_omni/engine/stage_init_utils.py

@@ -1253,9 +1253,19 @@ def _finalize_engine_args_dict(
     if is_diffusion:
         from vllm_omni.diffusion.data import parse_attention_config
 
-        if engine_args_dict.get("diffusion_attention_config") is not None:
+        # Fold the attention shorthand into the structured config so only one
+        # representation reaches OmniDiffusionConfig.from_kwargs.
+        attention_backend = engine_args_dict.pop("diffusion_attention_backend", None)
+        fastvideo_vsa_topk = engine_args_dict.pop("fastvideo_vsa_topk", None)
+        if (
+            engine_args_dict.get("diffusion_attention_config") is not None
+            or attention_backend is not None
+            or fastvideo_vsa_topk is not None
+        ):
             engine_args_dict["diffusion_attention_config"] = parse_attention_config(
-                engine_args_dict["diffusion_attention_config"],
+                engine_args_dict.get("diffusion_attention_config"),
+                attention_backend=attention_backend,
+                fastvideo_vsa_topk=fastvideo_vsa_topk,
             )
     else:
         resolve_worker_cls(engine_args_dict)

vllm_omni/entrypoints/openai/audio_utils_mixin.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from io import BytesIO
 
 import numpy as np
@@ -15,6 +18,88 @@ except ImportError:
 logger = init_logger(__name__)
 
 
+class StreamingAudioResampler:
+    """Stateful resampler for streaming mono audio.
+
+    Only integer downsampling ratios are supported so every input chunk can
+    produce output without buffering the complete stream.
+    """
+
+    def __init__(self, source_rate: int, target_rate: int):
+        if source_rate <= 0 or target_rate <= 0:
+            raise ValueError("Audio sample rates must be positive")
+        self.source_rate = source_rate
+        self.target_rate = target_rate
+        if source_rate < target_rate or source_rate % target_rate != 0:
+            raise ValueError(
+                "Streaming audio resampling requires an integer downsampling ratio, "
+                f"got {source_rate} Hz to {target_rate} Hz"
+            )
+
+        self._ratio = source_rate // target_rate
+        self._buffer = np.empty((0,), dtype=np.float32)
+        self._buffer_start = 0
+        self._total_samples = 0
+        self._next_center = 0
+
+        if self._ratio == 1:
+            self._taps = np.ones((1,), dtype=np.float32)
+        else:
+            # Use a Kaiser-windowed sinc and leave a small transition band
+            # below the target Nyquist frequency.
+            num_taps = 20 * self._ratio + 1
+            half = num_taps // 2
+            offsets = np.arange(num_taps, dtype=np.float64) - half
+            cutoff = 0.95 / self._ratio
+            taps = cutoff * np.sinc(cutoff * offsets) * np.kaiser(num_taps, 5.0)
+            self._taps = (taps / taps.sum()).astype(np.float32)
+
+    def process(self, audio: np.ndarray, *, final: bool = False) -> np.ndarray:
+        chunk = np.asarray(audio, dtype=np.float32)
+        if chunk.ndim != 1:
+            raise ValueError(f"Streaming audio resampling only supports mono audio, got shape {chunk.shape}")
+
+        if self._ratio == 1:
+            return chunk
+
+        if chunk.size:
+            self._buffer = np.concatenate((self._buffer, chunk))
+            self._total_samples += int(chunk.size)
+
+        half = self._taps.size // 2
+        max_center = self._total_samples - 1 if final else self._total_samples - 1 - half
+        if self._next_center > max_center:
+            return np.empty((0,), dtype=np.float32)
+
+        centers = np.arange(self._next_center, max_center + 1, self._ratio, dtype=np.int64)
+        first_center = int(centers[0])
+        last_center = int(centers[-1])
+        window_start = first_center - half
+        window_end = last_center + half
+        segment = np.zeros((window_end - window_start + 1,), dtype=np.float32)
+
+        copy_start = max(window_start, self._buffer_start, 0)
+        copy_end = min(window_end + 1, self._buffer_start + self._buffer.size)
+        if copy_end > copy_start:
+            dst_start = copy_start - window_start
+            src_start = copy_start - self._buffer_start
+            segment[dst_start : dst_start + copy_end - copy_start] = self._buffer[
+                src_start : src_start + copy_end - copy_start
+            ]
+
+        windows = np.lib.stride_tricks.sliding_window_view(segment, self._taps.size)[:: self._ratio]
+        output = windows @ self._taps[::-1]
+        self._next_center = last_center + self._ratio
+
+        keep_from = max(0, self._next_center - half)
+        drop = min(max(keep_from - self._buffer_start, 0), self._buffer.size)
+        if drop:
+            self._buffer = self._buffer[drop:]
+            self._buffer_start += int(drop)
+
+        return np.asarray(output, dtype=np.float32)
+
+
 class AudioMixin:
     """Mixin class to add audio-related utilities."""
 
@@ -44,6 +129,10 @@ class AudioMixin:
 
         audio_tensor, sample_rate = self._apply_speed_adjustment(audio_tensor, speed, sample_rate)
 
+        if audio_obj.output_sample_rate is not None and audio_obj.output_sample_rate != sample_rate:
+            audio_tensor = self._resample_audio(audio_tensor, sample_rate, audio_obj.output_sample_rate)
+            sample_rate = audio_obj.output_sample_rate
+
         supported_formats = {
             "wav": ("WAV", "audio/wav", {}),
             "pcm": ("RAW", "audio/pcm", {"subtype": "PCM_16"}),
@@ -69,6 +158,19 @@ class AudioMixin:
 
         return AudioResponse(audio_data=audio_data, media_type=media_type)
 
+    @staticmethod
+    def _resample_audio(audio_tensor: np.ndarray, source_rate: int, target_rate: int) -> np.ndarray:
+        """Resample complete audio while preserving soundfile's channels-last layout."""
+        if source_rate == target_rate:
+            return audio_tensor
+
+        audio_array = np.asarray(audio_tensor)
+        if not np.issubdtype(audio_array.dtype, np.floating):
+            audio_array = audio_array.astype(np.float32)
+        waveform = torch.from_numpy(audio_array.T.copy() if audio_array.ndim == 2 else audio_array.copy())
+        resampled = torchaudio.functional.resample(waveform, source_rate, target_rate).cpu().numpy()
+        return resampled.T if audio_array.ndim == 2 else resampled
+
     def _apply_speed_adjustment(self, audio_tensor: np.ndarray, speed: float, sample_rate: int):
         """Apply speed adjustment to the audio tensor while preserving pitch.
 

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import math
-from typing import Any, Literal
+from typing import Annotated, Any, Literal
 
 import numpy as np
 from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator
@@ -15,6 +15,7 @@ _MAX_EMBEDDING_DIM = 8192
 SUPPORTED_AUDIO_FORMATS: frozenset[str] = frozenset({"wav", "pcm", "flac", "mp3", "opus"})
 SUPPORTED_CHAT_AUDIO_FORMATS: frozenset[str] = SUPPORTED_AUDIO_FORMATS | {"pcm16"}
 DEFAULT_AUDIO_FORMAT: str = "wav"
+SpeechSampleRate = Annotated[int, Field(gt=0)]
 
 
 def _normalize_ref_audio_value(value):
@@ -80,6 +81,10 @@ class OpenAICreateSpeechRequest(BaseModel):
         description="Instructions for voice style/emotion (maps to 'instruct' for Qwen3-TTS)",
     )
     response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
+    sample_rate: SpeechSampleRate | None = Field(
+        default=None,
+        description="Target sample rate of the returned audio. If omitted, use the model's native sample rate.",
+    )
     speed: float | None = Field(
         default=1.0,
         ge=0.25,
@@ -396,6 +401,7 @@ class OpenAICreateAudioGenerateRequest(BaseModel):
 class CreateAudio(BaseModel):
     audio_tensor: np.ndarray
     sample_rate: int = 24000
+    output_sample_rate: int | None = None
     response_format: str = "wav"
     speed: float = 1.0
     base64_encode: bool = True
@@ -420,6 +426,7 @@ class SpeechBatchItem(BaseModel):
     voice: str | None = Field(default=None, validation_alias=AliasChoices("voice", "speaker"))
     instructions: str | None = None
     response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] | None = None
+    sample_rate: SpeechSampleRate | None = None
     speed: float | None = Field(default=None, ge=0.25, le=4.0)
     task_type: Literal["CustomVoice", "VoiceDesign", "Base"] | None = None
     language: str | None = None
@@ -440,6 +447,7 @@ class BatchSpeechRequest(BaseModel):
     voice: str | None = Field(default=None, validation_alias=AliasChoices("voice", "speaker"))
     instructions: str | None = None
     response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
+    sample_rate: SpeechSampleRate | None = None
     speed: float | None = Field(default=1.0, ge=0.25, le=4.0)
     task_type: Literal["CustomVoice", "VoiceDesign", "Base"] | None = None
     language: str | None = None

vllm_omni/entrypoints/openai/serving_speech.py

@@ -37,7 +37,7 @@ from vllm.utils import random_uuid
 from vllm.utils.async_utils import make_async
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
-from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
+from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin, StreamingAudioResampler
 from vllm_omni.entrypoints.openai.protocol.audio import (
     AudioResponse,
     BatchSpeechRequest,
@@ -1296,6 +1296,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
     def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
         """Validate TTS request parameters. Returns error message or None."""
+        sample_rate_error = self._validate_speech_sample_rate(request)
+        if sample_rate_error is not None:
+            return sample_rate_error
+
         adapter = self._get_tts_adapter()
         if adapter is not None:
             return adapter.validate(request)
@@ -1310,6 +1314,22 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         )
         return adapter_cls(ctx).validate(request)
 
+    def _validate_speech_sample_rate(self, request: OpenAICreateSpeechRequest) -> str | None:
+        if request.sample_rate is None:
+            return None
+
+        adapter_cls = resolve_adapter(self._tts_model_type)
+        supported_rates = adapter_cls.supported_output_sample_rates if adapter_cls is not None else frozenset()
+        if request.sample_rate not in supported_rates:
+            if supported_rates:
+                rates = ", ".join(str(rate) for rate in sorted(supported_rates))
+                return (
+                    f"sample_rate={request.sample_rate} is not supported by the current TTS model; "
+                    f"supported rates: {rates}"
+                )
+            return "sample_rate is not supported by the current TTS model"
+        return None
+
     def _voxcpm2_encode(self, text: str) -> list[int]:
         """Tokenize text for VoxCPM2, splitting multichar Chinese tokens."""
         from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
@@ -2191,6 +2211,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         usage_acc: SpeechOutputTokenCounter | None = None,
         tts_params: dict[str, Any] | None = None,
         collect: dict | None = None,
+        target_sample_rate: int | None = None,
     ):
         """Generate audio chunks for streaming response.
 
@@ -2214,6 +2235,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         first_audio_chunk_s: float | None = None
         stream_start_s = request_start_s if request_start_s is not None else time.perf_counter()
         artifact_ready = False
+        source_sample_rate: int | None = None
+        resampler: StreamingAudioResampler | None = None
 
         # SSE supplies an accumulator for usage output. Raw-audio and WebSocket
         # streams retain terminal metrics only when their model adapter needs
@@ -2239,7 +2262,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 if sr_raw is not None:
                     sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
                     sample_rate_val = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
-
+                    if source_sample_rate is not None and sample_rate_val != source_sample_rate:
+                        raise ValueError(
+                            "Audio sample rate changed during streaming: "
+                            f"{source_sample_rate} Hz to {sample_rate_val} Hz"
+                        )
                 audio_val = audio_output[audio_key]
                 if isinstance(audio_val, list):
                     # Cumulative mode: each update grows the list; emit only new tail.
@@ -2253,12 +2280,25 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     else:
                         new_chunks = []
 
+                if new_chunks and source_sample_rate is None:
+                    if response_format == "wav" and sr_raw is None:
+                        raise ValueError("First audio output must include sample rate metadata for WAV streaming")
+                    source_sample_rate = sample_rate_val
+                    if target_sample_rate is not None and target_sample_rate != source_sample_rate:
+                        resampler = StreamingAudioResampler(source_sample_rate, target_sample_rate)
+
+                output_sample_rate = target_sample_rate or sample_rate_val
+
                 for chunk_tensor in new_chunks:
                     chunk_np = (
                         chunk_tensor.float().detach().cpu().numpy() if hasattr(chunk_tensor, "float") else chunk_tensor
                     )
                     if chunk_np.ndim > 1:
                         chunk_np = chunk_np.squeeze()
+                    if resampler is not None:
+                        chunk_np = resampler.process(chunk_np)
+                        if chunk_np.size == 0:
+                            continue
                     if self._tts_model_type in _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES and int(np.size(chunk_np)) == 0:
                         # Zero-size chunks must not emit a WAV header or count
                         # as first audio; the post-loop guard below needs to
@@ -2266,14 +2306,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         continue
                     # For WAV format, emit header before first audio chunk
                     if response_format == "wav" and first_chunk:
-                        # Assert that sample rate has been set from chunk metadata (not just default)
-                        # This ensures the WAV header contains the correct sample rate
-                        assert sr_raw is not None, (
-                            "First audio chunk must include sample rate metadata for WAV streaming"
-                        )
                         num_channels = _infer_audio_num_channels(np.asarray(chunk_np))
                         wav_header = _create_wav_header(
-                            sample_rate=sample_rate_val,
+                            sample_rate=output_sample_rate,
                             num_channels=num_channels,
                             bits_per_sample=16,
                         )
@@ -2283,7 +2318,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     # Convert audio to PCM bytes
                     audio_obj = CreateAudio(
                         audio_tensor=chunk_np,
-                        sample_rate=sample_rate_val,
+                        sample_rate=output_sample_rate,
                         response_format="pcm",
                         speed=1.0,
                         base64_encode=False,
@@ -2292,7 +2327,33 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         first_audio_chunk_s = time.perf_counter()
                     audio_bytes = self.create_audio(audio_obj).audio_data
                     if include_sample_rate:
-                        yield audio_bytes, sample_rate_val
+                        yield audio_bytes, output_sample_rate
+                    else:
+                        yield audio_bytes
+
+            if resampler is not None:
+                final_chunk = resampler.process(np.empty((0,), dtype=np.float32), final=True)
+                if final_chunk.size:
+                    output_sample_rate = target_sample_rate or sample_rate_val
+                    if response_format == "wav" and first_chunk:
+                        yield _create_wav_header(
+                            sample_rate=output_sample_rate,
+                            num_channels=1,
+                            bits_per_sample=16,
+                        )
+                        first_chunk = False
+                    audio_obj = CreateAudio(
+                        audio_tensor=final_chunk,
+                        sample_rate=output_sample_rate,
+                        response_format="pcm",
+                        speed=1.0,
+                        base64_encode=False,
+                    )
+                    if first_audio_chunk_s is None:
+                        first_audio_chunk_s = time.perf_counter()
+                    audio_bytes = self.create_audio(audio_obj).audio_data
+                    if include_sample_rate:
+                        yield audio_bytes, output_sample_rate
                     else:
                         yield audio_bytes
             if self._tts_model_type in _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES and first_audio_chunk_s is None:
@@ -2398,6 +2459,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 request_start_s=request_start_s,
                 usage_acc=usage_acc,
                 tts_params=tts_params,
+                target_sample_rate=request.sample_rate if request is not None else None,
             ):
                 payload = {
                     "type": "speech.audio.delta",
@@ -3003,6 +3065,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         if self.engine_client.errored:
             raise self.engine_client.dead_error
 
+        sample_rate_error = self._validate_speech_sample_rate(request)
+        if sample_rate_error is not None:
+            raise ValueError(sample_rate_error)
+
         request_id = request_id or f"speech-{random_uuid()}"
         qwen3_ref_audio_warmup_artifact_key: str | None = None
 
@@ -3151,6 +3217,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         include_sample_rate: bool = False,
         tts_params: dict[str, Any] | None = None,
         collect: dict | None = None,
+        target_sample_rate: int | None = None,
     ):
         """Yield raw PCM byte chunks from the engine generator.
 
@@ -3166,6 +3233,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             include_sample_rate=include_sample_rate,
             tts_params=tts_params,
             collect=collect,
+            target_sample_rate=target_sample_rate,
         ):
             yield chunk
 
@@ -3173,7 +3241,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         """Yield raw PCM bytes for a speech request as soon as chunks are decoded."""
         request_id, generator, tts_params = await self._prepare_speech_generation(request)
         try:
-            async for chunk in self._generate_pcm_chunks(generator, request_id, tts_params=tts_params):
+            async for chunk in self._generate_pcm_chunks(
+                generator,
+                request_id,
+                tts_params=tts_params,
+                target_sample_rate=request.sample_rate,
+            ):
                 yield chunk
         finally:
             self._discard_ref_audio_artifact_warmup(request_id)
@@ -3323,6 +3396,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             audio_obj = CreateAudio(
                 audio_tensor=audio_tensor,
                 sample_rate=sample_rate,
+                output_sample_rate=request.sample_rate,
                 response_format=request.response_format or "wav",
                 speed=self._audio_encode_speed(request),
                 base64_encode=base64_encode,
@@ -3453,6 +3527,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             audio_obj = CreateAudio(
                 audio_tensor=audio_tensor,
                 sample_rate=sample_rate,
+                output_sample_rate=request.sample_rate,
                 response_format=request.response_format or "wav",
                 speed=self._audio_encode_speed(request),
                 base64_encode=False,
@@ -3523,6 +3598,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         - ref_audio: Reference audio for voice cloning (Base task)
         - ref_text: Transcript of reference audio (Base task)
         - x_vector_only_mode: Use speaker embedding only (Base task)
+        - sample_rate: Target output sample rate (8000 or 24000 Hz)
 
         Streaming is supported via the ``stream=True`` switch or ``stream_format='sse'``,
         which return OpenAI ``speech.audio.*`` SSE events. ``stream_format='audio'``
@@ -3534,6 +3610,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             if _is_default_voice(request.voice.lower(), self._get_available_speakers()):
                 request.voice = None
 
+        sample_rate_error = self._validate_speech_sample_rate(request)
+        if sample_rate_error is not None:
+            return self.create_error_response(sample_rate_error)
+
         if self._diffusion_mode:
             return await self._create_diffusion_speech(request)
 
@@ -3580,6 +3660,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         raw_request=raw_request,
                         request_start_s=request_start_s,
                         tts_params=raw_tts_params,
+                        target_sample_rate=request.sample_rate,
                     ),
                     media_type=media_type,
                 )
@@ -3730,6 +3811,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             voice=_pick("voice"),
             instructions=_pick("instructions"),
             response_format=_pick("response_format") or "wav",
+            sample_rate=_pick("sample_rate"),
             speed=picked_speed if picked_speed is not None else 1.0,
             stream=False,
             task_type=_pick("task_type"),

vllm_omni/entrypoints/openai/tts_adapters/base.py

@@ -166,6 +166,9 @@ class TTSModelAdapter(ABC):
     validates_generation: ClassVar[bool] = False
     #: Whether the model consumes ``request.speed`` in its native parameters.
     native_speed_control: ClassVar[bool] = False
+    #: Target sample rates validated for this adapter's output path. An empty
+    #: set means that the adapter does not expose per-request resampling.
+    supported_output_sample_rates: ClassVar[frozenset[int]] = frozenset()
 
     max_new_tokens_min = 1
 

vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py

@@ -42,6 +42,7 @@ class Qwen3TTSAdapter(ARTTSAdapter):
     validates_generation = True
     stage_keys = frozenset({"qwen3_tts"})
     name = "qwen3_tts"
+    supported_output_sample_rates = frozenset({8000, 24000})
 
     def normalize(self, request: "OpenAICreateSpeechRequest") -> None:
         """Qwen3-TTS normalization (Base-task inference, voice lowercasing) is