Keep baseline TTS ASR evaluation unchanged and retain failed-audio evidence

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

.buildkite/cuda/test-merge.yml

@@ -124,6 +124,30 @@ steps:
           - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model'
         mirror_hardwares: h100_2
 
+      - label: "Omni · Qwen3-Omni Async-Chunk SHM Leak"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_omni_shm_leak.py
+          - tests/helpers/runtime.py
+          - tests/helpers/stage_config.py
+          - vllm_omni/core/sched/omni_ar_scheduler.py
+          - vllm_omni/distributed/omni_connectors/connectors/shm_connector.py
+          - vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py
+          - vllm_omni/engine/orchestrator.py
+          - vllm_omni/engine/stage_engine_core_proc.py
+          - vllm_omni/engine/stage_pool.py
+          - vllm_omni/entrypoints/openai/serving_chat.py
+          - vllm_omni/model_executor/models/qwen3_omni/
+          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+          - vllm_omni/model_executor/models/common/snake_activation.py
+          - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
+          - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
+          - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
+          - vllm_omni/deploy/qwen3_omni_moe.yaml
+        timeout_in_minutes: 40
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_shm_leak.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+        mirror_hardwares: h100_2
+
       - label: "Omni · MiniCPM-o 4.5 Test"
         source_file_dependencies: omni_minicpmo_4_5_function
         timeout_in_minutes: 50

.buildkite/cuda/test-nightly.yml

@@ -156,6 +156,11 @@ steps:
       - label: ":full_moon: TTS · Function Test with L4"
         source_file_dependencies: tts_qwen3_tts_function
         timeout_in_minutes: 120
+        env:
+          VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR: .
+        artifact_paths:
+          - failed_speech_*.wav
+          - failed_speech_*.json
         commands:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
           - export VLLM_USE_DEEP_GEMM="0"

docs/configuration/environment_variables.md

@@ -214,7 +214,7 @@ collection, examples, bug reports, or logs.
 
 ## Model-specific variables
 
-The inventory includes 66 variables read by a single model or pipeline family. They are
+The inventory includes 93 variables read by a single model or pipeline family. They are
 not listed as public usage options here because doing so would turn implementation
 escape hatches into an accidental compatibility contract.
 
@@ -223,10 +223,10 @@ Every audited model-specific name has a migration disposition in the
 
 | Disposition | Count | Required outcome |
 | --- | ---: | --- |
-| Promote | 38 | Move a stable setting into typed stage or model configuration. |
+| Promote | 64 | Move a stable setting into typed stage or model configuration. |
 | Request scope | 5 | Move request-varying behavior into a declared request-option schema. |
 | External | 0 | Retain only when a supported third-party library owns the contract. |
-| Internalize | 16 | Keep a debug or diagnostic switch out of public documentation and configuration. |
+| Internalize | 17 | Keep a debug or diagnostic switch out of public documentation and configuration. |
 | Deprecate/remove | 7 | Remove a compatibility escape hatch that has no continuing contract. |
 
 The disposition is a migration target, not a statement that the existing

docs/configuration/stage_configs.md

@@ -514,3 +514,50 @@ for that comparison; this switch only controls automatic MPS management.
 
 MPS does not reserve a GPU. Use only assigned GPUs, explicitly place stages on
 the intended GPU, and warm the complete pipeline before measuring performance.
+
+## MOSS-TTS Local 1.5 with Model Runner V2
+
+`OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5` can opt into CUDA MRV2 using
+the shared runtime introduced for Qwen3-TTS:
+
+```bash
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml
+```
+
+This profile inherits the batching, codec graph buckets and 1-frame/15-frame
+chunk geometry from `moss_tts_local.yaml`, and selects V2 for both the Local
+Talker and codec stages. The Local depth predictor exposes the MRV2 `mtp`
+capabilities while retaining its V1 `talker_mtp` implementation, sampling
+defaults and explicit request-seed handling. Codec chunks use the native data
+plane; the internal Talker retains its final-only orchestrator output policy.
+
+On CUDA, the profile bounds stage-0 prefill work to 512 tokens per iteration.
+The codec retains its own CUDA Graph capture/replay, including batch buckets
+through 64, but disables Inductor compilation (`compilation_config.mode: 0`)
+of the stateful decoder to avoid lengthy compilation at startup. This is
+independent of `enforce_eager`; the profile keeps `enforce_eager: false`.
+It explicitly selects `cudagraph_mode: FULL`, which works without Inductor;
+the default `FULL_AND_PIECEWISE` would otherwise normalize to `NONE` and clear
+the codec's capture buckets when compilation is disabled.
+
+For sustained C128 serving on a large-memory CUDA GPU, select the
+separate `moss_tts_local_mrv2_high_concurrency.yaml` profile. It uses 128
+stream slots per stage, a 32 GiB Talker KV budget, and the codec's
+`triton_slot` backend with Inductor compilation and codec-owned CUDA graphs.
+The bounded KV budget leaves room for codec state and graphs; it does not
+guarantee that 128 maximum-length prompts fit simultaneously. See the
+[MOSS recipe](gh-file:recipes/OpenMOSS/MOSS-TTS.md#local-15-mrv2-and-slot-attention)
+for activation, backend comparisons, memory requirements and benchmark commands.
+
+Omitting `--deploy-config`, or selecting `moss_tts_local.yaml`, retains V1.
+NPU, XPU, ROCm and MUSA overrides also retain V1. This profile does not enable
+MRV2 for MOSS Delay, Realtime or Nano. Local 1.5 outputs 48 kHz stereo audio;
+set `VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000` and
+`VLLM_OMNI_BENCH_AUDIO_CHANNELS=2` when benchmarking raw PCM.
+
+Event-driven orchestration remains independently selectable with
+`VLLM_OMNI_EVENT_DRIVEN_ORCH=0` or `1`. Keep the runner and deployment identical
+when comparing these modes. Model-runner selection does not change the
+orchestration default or enable experimental reference encoding, chunk ramps,
+generation-output draining or MPS.

docs/user_guide/examples/offline_inference/text_to_speech.md

@@ -130,7 +130,7 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
 ### Notes
 
 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
-- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
+- H200 automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when the CUDA MPS control tool is available; other devices and environments use `cosyvoice3.yaml`. See the [serving defaults and tradeoffs](../online_serving/text_to_speech.md#cosyvoice3). Pass `--deploy-config cosyvoice3.yaml` for the previous profile or `--deploy-config <path>` to override. The example uses the resolved profile's sampling parameters.
 - `async_chunk: true` is the default; pass `--no-async-chunk` to switch to the legacy synchronous path.
 
 ---

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -155,16 +155,22 @@ vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remot
 ./cosyvoice3/run_server.sh
 ```
 
-Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.
+Streaming is on by default. On CUDA Hopper GPUs with at least 140 GiB of device memory (H200), serving automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when `nvidia-cuda-mps-control` is available on `PATH`: AR32 / codec32, standard sampling, packed Flow, cached ISTFT, HiFT decode graphs, and local CUDA MPS. No `--hf-overrides` is needed. Other devices and environments without the MPS control tool use `cosyvoice3.yaml` with RAS sampling and AR8 / codec8.
 
-Cross-request Stage-1 flow batching is opt-in. Enable it in the server environment when concurrent requests should share a flow-estimator call:
+The optimized profile is not a universal hardware configuration. Its packed kernels are Hopper-only; the talker's Model Runner V2 is unsupported on NPU/XPU, and CUDA MPS requires a local NVIDIA GPU and its control tool. Larger batches, 25-frame chunks, and graph allocations also require per-device validation on smaller or slower GPUs. Automatic selection retains the generic profile outside the validated device class; an explicit deploy configuration takes precedence.
+
+The H200 default targets streaming throughput at C32/C64. It changes sampling from RAS to standard (temperature 0.7, top-p 0.8, top-k 20, repetition penalty 1.21); output lengths and quality can differ. HiFT captures a recurring exact shape on its third use, which adds latency to that request, and its graph pool needs memory headroom beyond the engine reservations. This profile is not a claim of optimal first-audio latency at low concurrency. Pass `--deploy-config cosyvoice3.yaml` to retain the previous configuration, or `--deploy-config <path>` to select another profile.
+
+Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path; use `--deploy-config cosyvoice3.yaml` alongside it when selecting that path on H200.
+
+Cross-request Stage-1 flow batching is enabled in the H200 profile. For the generic configuration, enable it in the server environment when concurrent requests should share a flow-estimator call:
 
 ```bash
 export COSYVOICE3_BATCH_FLOW=1
 vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
 ```
 
-Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.
+Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Use `--deploy-config cosyvoice3.yaml` and leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.
 
 ### CLI client
 
@@ -194,7 +200,7 @@ The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-tex
 ### Notes
 
 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
-- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
+- Deploy config auto-loads the device-specific CosyVoice3 default described above. Pass `--deploy-config <path>` to override.
 - For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).
 
 ---

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import argparse
+import copy
 import functools
 import json
 import os
@@ -10,13 +11,11 @@ from typing import Any
 
 import numpy as np
 import soundfile as sf
-from vllm import SamplingParams
 from vllm.multimodal.media.audio import load_audio
 
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
-from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
 
 # Upstream zero-shot reference clip
 ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
@@ -59,8 +58,7 @@ def run_e2e():
         "--deploy-config",
         type=str,
         default=None,
-        help="Override the deploy config path. If unset, auto-loads "
-        "vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.",
+        help="Override the deploy config path. If unset, selects the bundled CosyVoice3 profile for the device.",
     )
     parser.add_argument("--text", type=str, default="Hello, this is a test of the CosyVoice system capability.")
     parser.add_argument(
@@ -105,8 +103,6 @@ def run_e2e():
         profiler_config=args.profiler_config,
     )
 
-    sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
-
     print("Model initialized. Preparing inputs...")
     ref_audio_path = args.ref_audio or _default_ref_audio()
     if not os.path.exists(ref_audio_path):
@@ -138,7 +134,7 @@ def run_e2e():
 
     print(f"Generating for prompt: {args.text}")
 
-    config = CosyVoice3Config()
+    config = omni.engine.stage_vllm_configs[0].model_config.hf_config
     tokenizer = get_qwen_tokenizer(
         token_path=args.tokenizer,
         skip_special_tokens=config.skip_special_tokens,
@@ -149,29 +145,10 @@ def run_e2e():
     min_len = int(base_len * config.min_token_text_ratio)
     max_len = int(base_len * config.max_token_text_ratio)
 
-    # Build SamplingParams for each stage (GPT, S2Mel, Vocoder)
-    gpt_sampling = SamplingParams(
-        temperature=1.0,
-        top_p=sampling_cfg["top_p"],
-        top_k=sampling_cfg["top_k"],
-        repetition_penalty=2.0,
-        min_tokens=min_len,
-        max_tokens=max_len,
-        stop_token_ids=[sampling_cfg["eos_token_id"]],
-        # allowed_token_ids=list(range(6561+3)),
-        detokenize=False,
-    )
-    # Not used
-    s2mel_sampling = SamplingParams(
-        temperature=1.0,
-        top_p=1.0,
-        top_k=-1,
-        repetition_penalty=2.0,
-        max_tokens=256,
-        detokenize=False,
-    )
-
-    sampling_params_list = [gpt_sampling, s2mel_sampling]
+    # Keep the deploy profile's sampling policy, penalties and required stops.
+    sampling_params_list = copy.deepcopy(omni.default_sampling_params_list)
+    sampling_params_list[0].max_tokens = max(1, min(2048, max_len))
+    sampling_params_list[0].min_tokens = min(max(1, min_len), sampling_params_list[0].max_tokens)
 
     profiler_enabled = args.profiler_config is not None
     if profiler_enabled:

examples/online_serving/text_to_speech/cosyvoice3/run_server.sh

@@ -1,11 +1,15 @@
 #!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Launch vLLM-Omni server for CosyVoice3 TTS
 #
 # Usage:
 #   ./run_server.sh
 #   CUDA_VISIBLE_DEVICES=0 ./run_server.sh
 #
-# Streaming (async-chunk) is on by default via vllm_omni/deploy/cosyvoice3.yaml.
+# Streaming (async-chunk) is on by default. H200 selects the optimized profile;
+# other devices use vllm_omni/deploy/cosyvoice3.yaml.
 # Set NO_ASYNC_CHUNK=1 to use the legacy synchronous path.
 
 set -e

recipes/OpenMOSS/MOSS-TTS.md

@@ -35,6 +35,10 @@ the requested `input` speech. Without a nonblank `ref_text`, Local uses
 audio-reference generation. Reference transcripts must match the reference
 audio; they are not style instructions.
 
+For the optional CUDA MRV2 runner, see the
+[Local 1.5 deployment profile](../../docs/configuration/stage_configs.md#moss-tts-local-15-with-model-runner-v2).
+The default Local deployment continues to use V1.
+
 ## References
 
 - Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)
@@ -152,3 +156,296 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 - No `ref_audio` required or accepted for MOSS-SoundEffect.
 - Input field maps to the `ambient_sound` parameter in the upstream processor.
 - Rate: ~12.5 tokens per second; longer descriptions produce longer audio.
+
+## Local 1.5 MRV2 and slot attention
+
+`MOSS-TTS-Local-Transformer-v1.5` supports the native CUDA MRV2 pipeline with
+an explicit deploy profile. The default Local profile continues to use V1.
+Both profiles below preserve 1-frame initial and 15-frame steady codec chunks
+and the model's sampling defaults.
+
+```bash
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml
+```
+
+This C64 profile bounds Talker prefill to 512 tokens and retains codec-owned
+CUDA graphs without compiling the codec with Inductor.
+
+For sustained high concurrency on a large-memory CUDA GPU, use:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 \
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml \
+  --stage-init-timeout 1200 --init-timeout 1500
+```
+
+The high-concurrency profile places both stages on one GPU, sets each stage's
+capacity to 128 and fixes the Talker KV budget at 32 GiB. The Talker uses
+mixed FULL CUDA graphs with both its token budget and capture limit set to
+512, so new prompt chunks remain within graph coverage. It also enables
+`hf_overrides.mrv2_gpu_slot_state`: Local hidden states, audio codes and
+continuation control stay in GPU request slots, gathered through MRV2's
+batch mapping. Prefill conditioning and explicit per-request sampling seeds
+retain their original behavior. Published audio codes own their storage, so
+slot reuse cannot overwrite an in-flight output. Set the override to `false`
+to compare with the generic Omni model state. The profile selects
+`codec_attention_backend: triton_slot`, Inductor mode 3 with combo kernels
+disabled, and codec CUDA graph buckets through 128. The configuration was
+validated on one H200. The earlier capacity-256 variant used about 110 GiB of
+sampled peak GPU memory, including loading and warmup. Reduce the stage capacities,
+graph buckets and coalescing target together when adapting it to smaller GPUs. Cold codec
+compilation can take several minutes; subsequent starts can reuse the AOT
+cache. Graph capture is still performed at startup.
+
+The high-concurrency profile also coalesces ready codec streams with
+`connectors.shm.extra.generation_min_batch_size: 32` and
+`generation_max_wait_ms: 12`. The target controls when to stop waiting; it
+neither caps execution at 32 nor guarantees every codec group has 32 rows.
+First and final chunks share the bounded window at high concurrency. When
+fewer than 32 requests remain, dispatch is immediate. Pending outputs are
+retired before waiting, with the original deadline preserved, so waiting does
+not prevent the previous batch from releasing its in-flight state. Cancellation
+and input notifications keep their ordinary scheduler bookkeeping. This option
+requires a stateful native MRV2 generation stage with TP1/PP1. Set the wait to
+`0` to disable coalescing; other profiles keep the immediate-dispatch default.
+
+The high-concurrency profile enables two additional Local model-state
+optimizations under the CUDA stage-0 `hf_overrides`, alongside
+`mrv2_gpu_slot_state: true`:
+
+- `mrv2_batch_prefill: true` reuses the batch's text embeddings and combines
+  CPU reference-code slices and their destination positions into one pinned
+  upload. Host staging is reused only after the upload event completes;
+  GPU-resident references keep the original path.
+- `mrv2_direct_tokens: true` returns the already determined text token after
+  the Local audio/stop decision. It avoids constructing and sampling the
+  full text vocabulary. Audio-code and binary-stop sampling are unchanged.
+  Requests requiring distribution metadata or token constraints, and
+  unsupported execution modes, retain the normal sampler.
+
+Other profiles leave both flags disabled. Set both flags explicitly when
+comparing their combined effect. Deployment
+inheritance replaces a stage's `hf_overrides` mapping, so an overlay must also
+retain `mrv2_gpu_slot_state: true`. Measure first-packet latency as well as
+throughput: faster prompt preparation can change competition between the
+Talker and codec sharing the GPU.
+
+For the single-H200 capacity-128 throughput baseline, keep both switches on.
+In complete Seed-TTS EN1088 runs at client concurrency 128, four combined-switch
+rounds pooled 344.11 audio-s/s versus 335.78 for batch prefill alone; direct
+tokens alone had no reliable gain. The combined path had 685 ms mean first
+audio versus 679 ms across six control rounds, with lower mean completion time.
+A fixed-seed 128-row audio sample showed mean Whisper WER of 3.68% versus
+3.78% for the control. This sample does not establish speaker similarity or
+exclude small quality changes. At concurrency 256, the combined path had
+raised mean first-audio latency by about 0.33 s, so other profiles leave the
+switches off until their latency and quality are checked.
+
+The optional connector setting `generation_coalescing_policy: idle_wait`
+waits for an inbox notification when no stream is runnable, all receivers
+are parked and registered, and no output needs retirement. The first ready
+arrival starts a fresh coalescing window, preserving the batch budget.
+Control messages, including cancellations, can wait up to two windows
+instead of one (24 ms at the 12 ms setting). The default policy remains
+`fixed`; neither policy changes the global orchestration default.
+
+The codec backends differ in state access:
+
+- `sdpa` uses PyTorch attention after gathering and updating the ring cache.
+- `triton` replaces the attention calculation, retaining the ring-cache
+  gather/copy and explicit mask construction.
+- `triton_slot` writes surviving K/V directly into request slots, attends
+  directly to the ring, and advances the active slot offsets in three
+  ordered kernels. Graph-padding rows do not advance persistent state.
+
+The slot path preserves the existing chunk-complete ring semantics, including
+retaining the final cache-capacity tokens when a chunk exceeds the ring.
+It does not introduce another cache owner or change request-slot lifetime.
+Zero-length padding rows skip the attention computation and emit zeros.
+
+The high-concurrency profile includes intermediate codec batch buckets 6, 12
+and 24. To compare with power-of-two buckets, change the stage-1
+`cudagraph_capture_sizes` to `[1, 2, 4, 8, 16, 32, 64, 128]`.
+Keep the maximum bucket equal to the state capacity to retain terminal-tail
+coalescing. Smaller buckets reduce padding work but require more graphs;
+measure complete serving runs before selecting them for a deployment.
+
+For a codec output-path comparison, set
+`connectors.shm.extra.codec_gpu_stream_output: 0` to select the synchronous
+codec output path; `1` selects GPU output snapshots on a private codec
+stream. Both configurations use MRV2 and keep streaming audio responses.
+The option changes input ownership, metadata staging and snapshot handling
+as well as transfer scheduling, so its timing difference is not just D2H
+copy time.
+
+Those chunk-complete semantics truncate the causal window: every ring holds
+exactly `context` entries and a whole chunk is written before attention runs,
+so token `i` of a `T`-token chunk sees `capacity - T + i + 1` keys. With the
+tokenizer-v2 decoder's per-layer contexts (10/10/8/4/2/1 s, i.e. 125/250/400/
+400/400/400 tokens) a 15-frame chunk is 480 tokens at the last transformer, so
+its first 80 tokens attend to nothing and no chunk sees the previous one there.
+Padding a terminal tail to 15 frames therefore decodes it differently from an
+exact-length decode; the difference is deterministic (identical in fp32 and
+fp64), not bf16 noise. The connector option `codec_ring_headroom: 1` sizes each
+ring as `context + max_chunk_frames * tokens_per_frame` (including ramp shapes); chunked streaming
+then reproduces whole-sequence decoding bit-for-bit in fp64 and padded tails
+equal exact tails. It increases state/graph memory and capture time; historical C256 experiments
+observed about a 2% throughput cost. It is off by default and its perceptual
+quality benefit has not been established. Terminal tails are padded
+into the regular 15-frame graph bucket in either mode.
+The slot kernel is specific to the CUDA tokenizer-v2 decoder; other paths
+retain their existing attention implementation. Event-driven orchestration
+remains independently selectable and its default is unchanged.
+
+### Optional progressive chunks
+
+Set `codec_chunk_ramp` under `connectors.shm.extra` to insert smaller chunks
+before steady decoding, for example:
+
+```yaml
+connectors:
+  shm:
+    extra:
+      codec_chunk_ramp: [1, 4, 15]
+```
+
+The ramp's first entry overrides `initial_codec_chunk_frames`. After the last
+entry, chunks use `codec_chunk_frames`. Each request advances independently;
+the final partial chunk is flushed and an empty terminal carries no codec
+tokens. The codec captures every ramp length, sizes the maximum execution
+step for the largest entry, and uses the ramp's first entry for its dedicated
+first-chunk graph. Additional shapes increase startup and graph memory.
+
+Ramps are opt-in. They can reduce early playback underrun while increasing
+codec calls, first-audio latency or total generation time. The measured C128
+throughput presets retain 1→15; their published figures do not establish an
+additional benefit from enabling ramps. The processor and graph-shape tests
+run without GPU or model weights:
+
+```bash
+CUDA_VISIBLE_DEVICES= PYTHONPATH=. python -m pytest -q \
+  tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py \
+  -m 'core_model and cpu' --run-level core_model
+```
+
+### Low-latency and reference-encoding options
+
+`moss_tts_local_mrv2_low_latency.yaml` keeps both stage capacities at 128 and
+uses prefill 2048, GPU slot state, batch prefill, direct tokens and eager MTP.
+Its codec first-chunk fast path uses dedicated graphs and a stream handoff
+before the regular decoder reuses the request slot. A gate limits contention
+with regular codec work. Slot waits are bounded at 30 seconds. A failed or
+stalled decode raises an error rather than replaying a partially advanced
+slot; a failed output enqueue retains owned PCM for regular delivery.
+Closing the decoder rejects new jobs. Its regular-batch limit of 16 is internal dispatch
+policy; it does not change the client concurrency or stage capacities.
+
+Reference encoding runs in the API layer, independently of MRV2. Enable
+reference graphs explicitly with `VLLM_OMNI_MOSS_REF_GRAPHS=1`; the default
+eager encoder is retained when the option is absent. The graph path uses the
+loaded tokenizer's encoder/quantizer, length buckets and optional compilation
+(`VLLM_OMNI_MOSS_REF_COMPILE=0` disables compilation). It also uses windowed
+attention; `VLLM_OMNI_MOSS_REF_ATTN=sdpa` retains the original attention.
+Compilation and attention changes need not produce bit-identical codes.
+
+For multiple API processes, `VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR` enables shared
+reference-code storage and, by default, a single encoder host with four
+workers. Use a dedicated directory per service, checkpoint and encoding
+configuration. Workers have their own graph resources; all graph captures
+complete before serving begins. `VLLM_OMNI_MOSS_REF_SHARED_ENCODER=0` retains
+separate encoders while sharing codes. More workers/graphs consume memory;
+they do not imply more GPUs. Example after selecting an available GPU and
+configuring any operator-managed MPS daemon:
+
+```bash
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \
+VLLM_OMNI_CONNECTOR_RECV_POLL_MS=1 \
+VLLM_OMNI_MOSS_REF_GRAPHS=1 \
+VLLM_OMNI_MOSS_REF_ENCODER_WORKERS=4 \
+VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS=0 \
+VLLM_OMNI_MOSS_REF_INFLIGHT=4 \
+VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS=2 \
+VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR=/dev/shm/moss-local-service \
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --api-server-count 4 \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml \
+  --stage-init-timeout 1200 --init-timeout 1500
+```
+
+This example is a configurable serving profile, not the complete launch
+configuration of a historical benchmark. In particular, the low-latency YAML
+uses synchronous stage-1 scheduling; the retained source46 experiment used
+asynchronous stage-1 scheduling, FP8 backbone, MPS and additional encoder
+environment settings. Do not infer a throughput result from the YAML alone.
+
+### Cold versus hot reference measurements
+
+A cold reference misses the encoded-reference cache: the request pays for
+reference parsing/preparation and GPU encoding before speech generation.
+A hot reference reuses reference codes but still synthesizes new speech;
+generated audio is not cached. Pre-registering a reference moves encoding
+into registration and can make subsequent synthesis hot, but registration
+latency is still part of the first-use cost.
+
+Model startup/compilation is separate from reference coldness. EN1088 contains
+repeated references, so its first pass is not uniformly cold. For an all-unique
+cold pass, use its 666 first-occurrence references, keep order, warm execution
+with disjoint references, then repeat the same 666 requests. A second pass
+through independent API-local caches is not guaranteed hot; verify cache
+coverage before comparing it with a shared-cache result. Report requests/s
+alongside audio-s/s and retain failed, empty, long and near-silent outputs.
+
+### Reproduce the serving benchmark
+
+Use the complete Seed-TTS English test set, including reference audio and text.
+Set `SEED_TTS_DATA` to the directory containing `en/meta.lst` and its 1088
+entries. Run the native benchmark after the server is ready:
+
+```bash
+export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000
+export VLLM_OMNI_BENCH_AUDIO_CHANNELS=2
+export SEED_TTS_WER_EVAL=0
+
+for concurrency in 128; do
+  for phase in warm r1 r2; do
+    vllm bench serve --omni \
+      --model OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \
+      --backend openai-audio-speech --endpoint /v1/audio/speech \
+      --dataset-name seed-tts --dataset-path "$SEED_TTS_DATA" \
+      --seed-tts-locale en --disable-shuffle --num-prompts 1088 \
+      --num-warmups 0 --ready-check-timeout-sec 0 \
+      --output-len 256 --max-concurrency "$concurrency" \
+      --request-rate inf --seed 42 \
+      --extra-body '{"task_type":"Base","max_new_tokens":256}' \
+      --save-result --save-detailed --result-dir results/moss-local \
+      --result-filename "c${concurrency}-${phase}.json"
+  done
+done
+```
+
+Discard the complete `warm` pass and combine measured runs as total generated
+audio seconds divided by total benchmark duration. Require 1088 successes,
+zero failures and 1088 nonempty-audio metric samples in each run. The output
+cap matches the benchmark protocol; success alone does not establish speech
+quality or that every sentence ended before the cap. Dataset seed 42 does
+not fix an independent sampling seed for every request.
+
+For an attention-only comparison, copy the high-concurrency YAML beside the
+original to preserve relative `base_config` resolution. Change its
+codec `compilation_config.mode` to `0` while retaining `cudagraph_mode: FULL`,
+and compare `codec_attention_backend: triton` against `triton_slot`. Keep all
+other settings, warmup and client concurrency identical. Restart the server
+between configurations and use separate result directories. Comparing the
+compiled slot profile to uncompiled Triton includes both changes.
+
+Kernel and MHA regression tests require CUDA, compatible vLLM/Triton packages,
+and no model weights:
+
+```bash
+python -m pytest -q tests/model_executor/models/moss_tts/test_slot_attention.py \
+  tests/model_executor/models/moss_tts/test_streaming_attention.py \
+  -m 'core_model and cuda' --run-level=core_model
+```

tests/config/test_config_factory.py

@@ -571,6 +571,56 @@ class TestPipelineDiscovery:
         assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
         assert "qwen3_tts" in OMNI_PIPELINES
 
+    @pytest.mark.parametrize(
+        "cuda,major,memory_gib,mps,optimized",
+        [
+            (False, None, 0, True, False),
+            (True, None, 141, True, False),
+            (True, 8, 141, True, False),
+            (True, 9, 80, True, False),
+            (True, 10, 180, True, False),
+            (True, 9, 141, False, False),
+            (True, 9, 141, True, True),
+        ],
+    )
+    def test_cosyvoice3_device_default_and_explicit_overrides(
+        self, monkeypatch, cuda, major, memory_gib, mps, optimized
+    ):
+        from vllm.platforms import current_platform
+        from vllm.platforms.interface import DeviceCapability
+
+        from vllm_omni.platforms import current_omni_platform
+
+        monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda)
+        monkeypatch.setattr(current_omni_platform, "device_name", "cuda" if cuda else "cpu")
+        capability = DeviceCapability(major, 0) if major is not None else None
+        monkeypatch.setattr(current_platform, "get_device_capability", lambda: capability)
+        monkeypatch.setattr(current_platform, "get_device_total_memory", lambda: memory_gib * 1024**3)
+        monkeypatch.setattr(
+            "vllm_omni.model_executor.models.cosyvoice3.pipeline.shutil.which",
+            lambda command: "/usr/bin/nvidia-cuda-mps-control" if mps else None,
+        )
+        pipeline = resolve_pipeline_config("cosyvoice3")
+        assert pipeline is not None
+        config = VllmOmniConfig.from_pipeline_config(pipeline)
+        talker, codec = config.stage_configs
+        assert talker.scheduler_config.max_num_seqs == codec.scheduler_config.max_num_seqs == (32 if optimized else 8)
+        assert talker.runtime_config.cuda_mps == codec.runtime_config.cuda_mps == optimized
+        assert talker.model_config.hf_overrides == ({"cosyvoice3_sampling_mode": "standard"} if optimized else None)
+        assert pipeline.stages[0].sampling_constraints["stop_token_ids"] == list(range(6561, 6761))
+        if optimized:
+            assert codec.runtime_config.env["COSYVOICE3_CACHED_ISTFT"] == "1"
+            assert codec.runtime_config.env["COSYVOICE3_HIFT_GRAPH"] == "1"
+            override = VllmOmniConfig.from_pipeline_config(
+                pipeline, cli_overrides={"hf_overrides": {"cosyvoice3_sampling_mode": "ras", "seed": 7}}
+            )
+            assert override.stage_by_id(0).model_config.hf_overrides == {"cosyvoice3_sampling_mode": "ras", "seed": 7}
+        explicit = VllmOmniConfig.from_pipeline_config(
+            pipeline, deploy_config_path=get_deploy_config_path("cosyvoice3.yaml")
+        )
+        assert explicit.stage_by_id(0).scheduler_config.max_num_seqs == 8
+        assert explicit.stage_by_id(0).model_config.hf_overrides is None
+
     def test_registry_resolver_qwen3_omni_all_stages(self):
         """Test that providing the HF config for qwen3 omni with audio enabled uses all stages."""
         pipeline = resolve_pipeline_config(
@@ -2928,9 +2978,17 @@ class TestPlatformOverrides:
 
         assert deploy.stages[0].engine_extras["block_size"] == 128
 
-    @pytest.mark.parametrize("platform", ["cuda", "npu"])
-    def test_recommended_native_runner_platform_defaults(self, platform):
-        filename, pipeline_key = "qwen3_tts_mrv2.yaml", "qwen3_tts"
+    @pytest.mark.parametrize(
+        "filename,pipeline_key",
+        [
+            ("qwen3_tts_mrv2.yaml", "qwen3_tts"),
+            ("moss_tts_local_mrv2.yaml", "moss_tts_local"),
+            ("moss_tts_local_mrv2_high_concurrency.yaml", "moss_tts_local"),
+            ("moss_tts_local_mrv2_low_latency.yaml", "moss_tts_local"),
+        ],
+    )
+    @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
+    def test_recommended_native_runner_platform_defaults(self, platform, filename, pipeline_key):
         deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
         deploy = _apply_platform_overrides(deploy, platform=platform)
         expect_v2 = platform == "cuda"
@@ -2946,6 +3004,60 @@ class TestPlatformOverrides:
             # v2 only engages the native plane on stages declaring support.
             assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages)
 
+    def test_moss_local_mrv2_preserves_base_profile_and_variant_scope(self):
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        base_path = Path(get_deploy_config_path(pipeline.default_deploy_config_name))
+        candidate_path = Path(get_deploy_config_path("moss_tts_local_mrv2.yaml"))
+        base = load_deploy_config(base_path)
+        candidate = load_deploy_config(candidate_path)
+        assert base.model_runner == "v1"
+        assert candidate.stages == base.stages
+        assert candidate.connectors == base.connectors
+        assert candidate.async_chunk == base.async_chunk
+        cuda = _apply_platform_overrides(candidate, platform="cuda")
+        stages = merge_pipeline_deploy(pipeline, cuda)
+        assert stages[0].yaml_engine_args["max_num_batched_tokens"] == 512
+        codec_args = stages[1].yaml_engine_args
+        assert codec_args["max_num_seqs"] == 64
+        assert codec_args["enforce_eager"] is False
+        assert codec_args["compilation_config"] == {
+            "mode": 0,
+            "cudagraph_mode": "FULL",
+            "cudagraph_capture_sizes": [1, 2, 4, 8, 16, 32, 64],
+            "cudagraph_num_of_warmups": 1,
+        }
+        for platform in ("npu", "xpu", "rocm", "musa"):
+            # Platform resolution mutates the deployment; each server loads
+            # a fresh profile before applying its own platform overrides.
+            fallback = _apply_platform_overrides(load_deploy_config(candidate_path), platform=platform)
+            original = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
+            assert fallback.stages == original.stages
+        for variant in ("moss_tts_delay", "moss_tts_realtime"):
+            assert not any(stage.supports_native_mrv2_data_plane for stage in resolve_pipeline_config(variant).stages)
+
+    def test_moss_local_high_concurrency_resolves_slot_compile_and_memory_budget(self):
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        path = Path(get_deploy_config_path("moss_tts_local_mrv2_high_concurrency.yaml"))
+        deploy = _apply_platform_overrides(load_deploy_config(path), platform="cuda")
+        talker, codec = merge_pipeline_deploy(pipeline, deploy)
+        assert talker.yaml_engine_args["max_num_seqs"] == codec.yaml_engine_args["max_num_seqs"] == 128
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_batch_prefill"] is True
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_direct_tokens"] is True
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_gpu_slot_state"] is True
+        assert talker.yaml_engine_args["max_num_batched_tokens"] == 512
+        assert talker.yaml_engine_args["kv_cache_memory_bytes"] == 32 * 1024**3
+        assert codec.yaml_engine_args["hf_overrides"]["codec_attention_backend"] == "triton_slot"
+        compilation = codec.yaml_engine_args["compilation_config"]
+        assert compilation["mode"] == 3 and compilation["backend"] == "inductor"
+        assert compilation["cudagraph_mode"] == "FULL"
+        assert compilation["cudagraph_capture_sizes"] == [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128]
+        assert compilation["inductor_compile_config"] == {"combo_kernels": False, "benchmark_combo_kernel": False}
+        base_path = Path(get_deploy_config_path("moss_tts_local.yaml"))
+        for platform in ("npu", "xpu", "rocm", "musa"):
+            actual = _apply_platform_overrides(load_deploy_config(path), platform=platform)
+            expected = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
+            assert actual.stages == expected.stages
+
     @pytest.mark.parametrize("runner,native", [("v1", False), ("v2", False), ("v2", True)])
     def test_mrv2_undeclared_transport_warns(self, monkeypatch, runner, native):
         from unittest.mock import Mock

tests/config/test_environment_variables.py

@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 32,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 75,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 95,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
         EnvironmentVariableCategory.INTERNAL: 7,
     }
@@ -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: 46,
+        ModelEnvironmentVariableDisposition.PROMOTE: 66,
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 17,

tests/core/sched/test_generation_batch_coalescing.py

@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Bounded readiness coalescing must not allocate or consume request state."""
+
+import queue
+from types import SimpleNamespace
+
+import pytest
+from vllm.v1.core.sched.interface import PauseState
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+from vllm_omni.outputs import OmniConnectorOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_idle_wait_wakes_on_input_without_extending_coalescing_deadline(monkeypatch):
+    events = [notification("0"), notification("1", "2", "3")]
+    s, clock = scheduler(monkeypatch, delayed=[(0.004, events[0]), (0.002, events[1])])
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.006)
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.012])
+
+
+@pytest.mark.parametrize("pending", ["in_flight", "admission", "registration", "terminal", "empty"])
+def test_idle_wait_preserves_completion_and_control_progress(monkeypatch, pending):
+    s, clock = scheduler(monkeypatch)
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    if pending == "in_flight":
+        s.requests["0"].num_in_flight_tokens = 1
+    elif pending == "admission":
+        s.requests["0"].status = RequestStatus.WAITING
+    elif pending == "registration":
+        s.input_coordinator.pending_chunk_registrations = [object()]
+    elif pending == "terminal":
+        s._pending_data_plane_terminal_req_ids = {"0"}
+    else:
+        s.requests.clear()
+    assert s._drain_omni_connector_outputs() == []
+    assert clock[0] == 10
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_idle_wait_without_arrivals_is_bounded(monkeypatch):
+    s, clock = scheduler(monkeypatch)
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    assert s._drain_omni_connector_outputs() == []
+    assert clock[0] == pytest.approx(10.012)
+
+
+def notification(*ids, terminal=()):
+    return OmniConnectorOutput(chunk_ready_req_ids=set(ids), chunk_finished_req_ids=set(terminal))
+
+
+class Inbox:
+    def __init__(self, clock, initial=(), delayed=()):
+        self.clock = clock
+        self.initial = list(initial)
+        self.delayed = list(delayed)
+        self.waits = []
+
+    def get_nowait(self):
+        if self.initial:
+            return self.initial.pop(0)
+        raise queue.Empty
+
+    def get(self, timeout):
+        self.waits.append(timeout)
+        if not self.delayed or self.delayed[0][0] > timeout:
+            self.clock[0] += timeout
+            raise queue.Empty
+        delay, value = self.delayed.pop(0)
+        self.clock[0] += delay
+        return value
+
+
+def scheduler(monkeypatch, initial=(), delayed=()):
+    clock = [10.0]
+    monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.time.monotonic", lambda: clock[0])
+    s = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
+    s._generation_max_wait_s = 0.012
+    s._generation_min_batch_size = 4
+    s._pause_state = PauseState.UNPAUSED
+    s.requests = {
+        str(i): SimpleNamespace(request_id=str(i), num_in_flight_tokens=0, is_finished=lambda: False) for i in range(8)
+    }
+    s.running = list(s.requests.values())
+    s.input_coordinator = SimpleNamespace(
+        requests_with_ready_chunks=set(), finished_requests=set(), input_terminal_req_ids=set()
+    )
+    s._latest_omni_connector_output = None
+    s._omni_connector_output_inbox = Inbox(clock, initial, delayed)
+    return s, clock
+
+
+def test_arrivals_form_target_batch_in_order(monkeypatch):
+    events = [notification("0"), notification("1", "2"), notification("3")]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1]), (0.003, events[2])])
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.005)
+    assert not s.input_coordinator.requests_with_ready_chunks
+    assert all(r.num_in_flight_tokens == 0 for r in s.requests.values())
+
+
+def test_duplicate_stale_and_finished_notifications_do_not_fill_batch(monkeypatch):
+    events = [notification("0", "gone", "1", "2"), notification("0")]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.003, events[1])])
+    s.requests["1"].is_finished = lambda: True
+    s.requests["2"].is_finished = lambda: True
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.012)
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.009])
+
+
+@pytest.mark.parametrize("reason", ["enough", "empty", "low", "paused", "off"])
+def test_immediate_dispatch(monkeypatch, reason):
+    event = notification("0")
+    s, clock = scheduler(monkeypatch, [event])
+    if reason == "enough":
+        event.chunk_ready_req_ids.update(["1", "2", "3"])
+    elif reason == "empty":
+        event.chunk_ready_req_ids.clear()
+    elif reason == "low":
+        s.requests = {"0": s.requests["0"]}
+    elif reason == "paused":
+        s._pause_state = PauseState.PAUSED_ALL
+    elif reason == "off":
+        s._generation_max_wait_s = 0
+    assert s._drain_omni_connector_outputs() == [event]
+    assert not s._omni_connector_output_inbox.waits
+    assert clock[0] == 10
+
+
+def test_terminal_arrival_keeps_original_deadline(monkeypatch):
+    events = [notification("0"), notification("1", terminal=["1"])]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1])])
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.012)
+
+
+def test_pending_ready_and_legacy_notification_preserved(monkeypatch):
+    s, _ = scheduler(monkeypatch)
+    s.input_coordinator.requests_with_ready_chunks.update(["0", "1", "2"])
+    last = notification("3")
+    s._latest_omni_connector_output = last
+    assert s._drain_omni_connector_outputs() == [last]
+    assert s._latest_omni_connector_output is None
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_deadline_is_new_for_each_batch_not_extended_by_arrivals(monkeypatch):
+    s, clock = scheduler(monkeypatch, [notification("0")], [(0.008, notification("1")), (0.008, notification("2"))])
+    assert len(s._drain_omni_connector_outputs()) == 2
+    assert clock[0] == pytest.approx(10.012)
+    s._omni_connector_output_inbox.initial.append(notification("3"))
+    assert len(s._drain_omni_connector_outputs()) == 2
+    assert clock[0] == pytest.approx(10.024)
+
+
+def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mocker):
+    event = notification("0", "cancelled")
+    event.request_metadata = {"0": {"code_predictor_codes": [7]}, "cancelled": {"code_predictor_codes": [9]}}
+    s, _ = scheduler(monkeypatch, [event])
+    s._generation_max_wait_s = 0
+    s.waiting = []
+    s.input_coordinator._async_chunk = True
+    s.input_coordinator.update_request_metadata = mocker.Mock()
+    s.input_coordinator.process_pending_chunks = mocker.Mock()
+    s._consume_pending_connector_output("generation")
+    s.input_coordinator.update_request_metadata.assert_called_once_with(
+        s.requests, {"0": {"code_predictor_codes": [7]}}, model_mode="generation"
+    )
+    s.input_coordinator.process_pending_chunks.assert_called_once_with(s.waiting, s.running, {"0"}, set())
+
+
+@pytest.mark.parametrize("target,wait", [(0, 1), (257, 1), (4, -1), (4, float("nan")), (4, float("inf"))])
+def test_invalid_config_rejected(monkeypatch, target, wait):
+    from vllm_omni.core.sched.omni_generation_scheduler import VLLMScheduler
+
+    def initialize(s, *args, **kwargs):
+        s.max_num_running_reqs = 256
+        s.vllm_config = SimpleNamespace(
+            model_config=SimpleNamespace(
+                retains_state_across_chunks=True,
+                stage_connector_config={"extra": {"generation_min_batch_size": target, "generation_max_wait_ms": wait}},
+            )
+        )
+
+    monkeypatch.setattr(VLLMScheduler, "__init__", initialize)
+    monkeypatch.setattr(OmniGenerationScheduler, "_init_omni_io_scheduling_state", lambda s: None)
+    with pytest.raises(ValueError):
+        OmniGenerationScheduler()
+
+
+@pytest.mark.parametrize("terminal", [False, True])
+def test_first_and_terminal_chunks_have_bounded_wait(monkeypatch, terminal):
+    event = notification("0", terminal=["0"] if terminal else [])
+    s, clock = scheduler(monkeypatch, [event])
+    s.running = s.running[1:]
+    assert s._drain_omni_connector_outputs() == [event]
+    assert clock[0] == pytest.approx(10.012)
+
+
+def test_retire_in_flight_output_before_waiting_without_resetting_deadline(monkeypatch):
+    events = [notification("0"), notification("1")]
+    s, clock = scheduler(monkeypatch, events[:1])
+    s.requests["7"].num_in_flight_tokens = 1
+    assert s._drain_omni_connector_outputs() == events[:1]
+    assert s._generation_defer_batch
+    assert not s._omni_connector_output_inbox.waits
+    s.input_coordinator.requests_with_ready_chunks.add("0")
+    # Simulate EngineCore retiring the previous batch after five milliseconds.
+    clock[0] += 0.005
+    s.requests["7"].num_in_flight_tokens = 0
+    s._omni_connector_output_inbox.initial.append(events[1])
+    assert s._drain_omni_connector_outputs() == events[1:]
+    assert not s._generation_defer_batch
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.007])
+    assert clock[0] == pytest.approx(10.012)
+    assert s._generation_batch_deadline is None
+
+
+def test_expired_deadline_dispatches_even_with_other_requests_in_flight(monkeypatch):
+    s, clock = scheduler(monkeypatch, [notification("0")])
+    s.requests["7"].num_in_flight_tokens = 1
+    s._generation_batch_deadline = clock[0] - 0.001
+    assert len(s._drain_omni_connector_outputs()) == 1
+    assert not s._generation_defer_batch
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_deferred_schedule_keeps_pending_request_and_allocates_nothing(monkeypatch, mocker):
+    from tests.core.sched.test_generation_scheduler_restore import _chunk_request, _make_generation_scheduler
+
+    request = _chunk_request("pending")
+    s = _make_generation_scheduler(request)
+    s._process_pending_omni_inputs = lambda **kwargs: setattr(s, "_generation_defer_batch", True)
+    allocate = mocker.spy(s.kv_cache_manager, "allocate_slots")
+    output = s.schedule()
+    assert output.total_num_scheduled_tokens == 0
+    assert not output.num_scheduled_tokens
+    allocate.assert_not_called()
+    assert list(s.waiting) == [request]
+    assert request.num_computed_tokens == request.num_in_flight_tokens == 0
+    assert s.chunk_transfer_adapter.restore_called

tests/core/sched/test_generation_scheduler_init.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise generation constructor policy after actual deploy/platform resolution."""
+
+from types import SimpleNamespace
+
+import pytest
+
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler, VLLMScheduler
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def construct_scheduler(monkeypatch, mocker):
+    # Only upstream cache allocation and connector construction are omitted.
+    # Keep the actual Omni I/O initialization, native decision and constructor.
+    mocker.patch("vllm_omni.core.sched.omni_scheduler_mixin.OmniChunkTransferAdapter")
+
+    def upstream_init(self, config):
+        self.vllm_config = config
+        self.max_num_running_reqs = config.scheduler_config.max_num_seqs
+
+    monkeypatch.setattr(VLLMScheduler, "__init__", upstream_init)
+    return OmniGenerationScheduler
+
+
+def _config(*, native=True, stateful=True, tp=1, pp=1, extras=None, capacity=128):
+    return SimpleNamespace(
+        model_config=SimpleNamespace(
+            use_v2_model_runner=native,
+            supports_native_mrv2_data_plane=True,
+            retains_state_across_chunks=stateful,
+            async_chunk=True,
+            stage_id=1,
+            stage_connector_config={"extra": extras or {}},
+        ),
+        parallel_config=SimpleNamespace(tensor_parallel_size=tp, pipeline_parallel_size=pp),
+        scheduler_config=SimpleNamespace(max_num_seqs=capacity),
+    )
+
+
+@pytest.mark.parametrize("profile", ["high_concurrency", "low_latency"])
+@pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
+def test_moss_profile_generation_constructor_after_platform_resolution(
+    construct_scheduler, monkeypatch, mocker, profile, platform
+):
+    # merge_pipeline_deploy owns platform resolution. Mock the detected device
+    # instead of applying an override first and then resolving the worker's
+    # actual platform a second time.
+    monkeypatch.setattr(current_omni_platform, "device_name", platform)
+    deploy = load_deploy_config(get_deploy_config_path(f"moss_tts_local_mrv2_{profile}.yaml"))
+    pipeline = resolve_pipeline_config("moss_tts_local")
+    codec = merge_pipeline_deploy(pipeline, deploy)[1]
+    extras = deploy.connectors["shm"]["extra"]
+    config = _config(
+        native=codec.yaml_engine_args["use_v2_model_runner"],
+        stateful=pipeline.stages[1].retains_state_across_chunks,
+        capacity=codec.yaml_engine_args["max_num_seqs"],
+        extras=extras,
+    )
+    warning = mocker.patch("vllm_omni.core.sched.omni_generation_scheduler.logger.warning")
+    scheduler = construct_scheduler(config)
+    # The raw extras persist through fallback; the live constructor must cope.
+    assert extras["generation_min_batch_size"] == 32 and extras["generation_max_wait_ms"] == 12
+    if platform == "cuda":
+        assert scheduler._native_data_plane and scheduler.input_coordinator is not None
+        assert scheduler._generation_min_batch_size == 32
+        assert scheduler._generation_max_wait_s == 0.012
+        assert scheduler._generation_max_regular_batch == (16 if profile == "low_latency" else 0)
+        warning.assert_not_called()
+        if profile == "low_latency":
+            assert extras["codec_first_chunk_fast_path"] == extras["codec_first_chunk_gate"] == 1
+    else:
+        assert not scheduler._native_data_plane and scheduler.chunk_transfer_adapter is not None
+        assert scheduler._generation_min_batch_size == 1
+        assert scheduler._generation_max_wait_s == scheduler._generation_max_regular_batch == 0
+        warning.assert_called_once()
+
+
+@pytest.mark.parametrize("stateful,tp,pp", [(False, 1, 1), (True, 2, 1), (True, 1, 2)])
+def test_native_batch_waiting_still_rejects_unsupported_state_or_topology(construct_scheduler, stateful, tp, pp):
+    config = _config(
+        stateful=stateful, tp=tp, pp=pp, extras={"generation_min_batch_size": 32, "generation_max_wait_ms": 12}
+    )
+    with pytest.raises(ValueError, match="stateful native MRV2 TP1/PP1"):
+        construct_scheduler(config)
+
+
+@pytest.mark.parametrize(
+    "extras,message",
+    [
+        ({"generation_min_batch_size": 0}, "generation_min_batch_size"),
+        ({"generation_min_batch_size": 129}, "generation_min_batch_size"),
+        ({"generation_max_wait_ms": -1}, "generation_max_wait_ms"),
+        ({"generation_max_wait_ms": float("nan")}, "generation_max_wait_ms"),
+        ({"generation_max_regular_batch": -1}, "generation_max_regular_batch"),
+        ({"generation_coalescing_policy": "unknown"}, "generation_coalescing_policy"),
+    ],
+)
+def test_non_native_fallback_still_rejects_invalid_values(construct_scheduler, extras, message):
+    with pytest.raises(ValueError, match=message):
+        construct_scheduler(_config(native=False, extras=extras))

tests/core/sched/test_omni_ar_scheduler_kv_transfer.py

@@ -1,10 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
 
 import pytest
 from vllm import SamplingParams
 from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.request import RequestStatus
 
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata
 from vllm_omni.engine.serialization import deserialize_additional_information
 
@@ -25,7 +31,7 @@ def _make_engine_request() -> EngineCoreRequest:
     )
 
 
-def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
+def _omit_flags(*, force_kv_transfer: bool) -> tuple[bool, bool, dict]:
     tagged = apply_omni_final_stage_metadata(
         _make_engine_request(),
         final_stage_id=0,
@@ -37,21 +43,182 @@ def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
         request_id="req",
         additional_information=tagged.additional_information,
     )
-    result = scheduler._request_omits_kv_transfer_to_next_stage(request)
+    omits_kv = scheduler._request_omits_kv_transfer_to_next_stage(request)
+    omits_chunk = scheduler._request_omits_chunk_transfer_to_next_stage(request)
     metadata = deserialize_additional_information(tagged.additional_information)
-    return result, metadata
+    return omits_kv, omits_chunk, metadata
 
 
 def test_stage_zero_request_omits_kv_transfer():
-    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False)
+    omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=False)
 
-    assert omits_transfer
+    assert omits_kv
+    assert omits_chunk
     assert "omni_force_kv_transfer" not in metadata
 
 
 def test_cfg_companion_forces_kv_transfer_without_downstream_payload():
-    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True)
+    omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=True)
 
-    assert not omits_transfer
+    assert not omits_kv
+    assert omits_chunk
     assert metadata["omni_final_stage_id"] == 0
     assert metadata["omni_force_kv_transfer"] is True
+
+
+class _ChunkRequest(SimpleNamespace):
+    def __hash__(self):
+        return hash(self.request_id)
+
+    def __eq__(self, other):
+        return isinstance(other, _ChunkRequest) and other.request_id == self.request_id
+
+
+def _make_chunk_request(final_stage_id: int, *, force_kv_transfer: bool = False):
+    tagged = apply_omni_final_stage_metadata(
+        _make_engine_request(),
+        final_stage_id=final_stage_id,
+        force_kv_transfer=force_kv_transfer,
+    )
+    request = _ChunkRequest(
+        request_id="req-shm-leak",
+        additional_information=tagged.additional_information,
+        num_in_flight_tokens=1,
+        num_stale_output_tokens=0,
+        sampling_params=SimpleNamespace(num_logprobs=None),
+        has_encoder_inputs=False,
+        pooling_params=None,
+        status=RequestStatus.RUNNING,
+        resumable=False,
+        num_output_placeholders=0,
+        spec_token_ids=[],
+        _output_token_ids=[],
+        num_computed_tokens=1,
+        stop_reason=None,
+        num_nans_in_logits=None,
+        client_index=0,
+        trace_headers=None,
+    )
+    request.is_finished = lambda: False
+    request.get_finished_reason = lambda: None
+    request.take_prefill_stats = lambda: None
+    request.take_events = lambda: None
+    return request
+
+
+def _run_finished_save_step(mocker, request, *, inter_stage_outputs=None):
+    adapter = mocker.MagicMock()
+    adapter._confirmed_num_computed_tokens.return_value = 1
+
+    sched = mocker.MagicMock()
+    sched.requests = {request.request_id: request}
+    sched.perf_metrics = None
+    sched.defer_block_free = False
+    sched.structured_output_manager.should_advance.return_value = False
+    sched._update_request_with_output.return_value = ([42], True)
+    sched._process_kv_transfer_trigger.return_value = False
+    sched._handle_stopped_request.return_value = True
+
+    def _free_request(req, delay_free_blocks=False):
+        sched._omits_kv_transfer_cache.pop(req.request_id, None)
+        return None, None
+
+    sched._free_request.side_effect = _free_request
+    sched.chunk_transfer_adapter = adapter
+    sched.running = [request]
+    sched.waiting_for_transfer_free = set()
+    sched.transfer_triggered_requests = set()
+    sched.active_kv_transfers = set()
+    sched.pending_stop_after_extraction = set()
+    sched.connector = None
+    sched.kv_cache_manager.take_events.return_value = None
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
+    sched.finished_req_ids_dict = {}
+    sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
+
+    scheduler_output = mocker.MagicMock()
+    scheduler_output.num_scheduled_tokens = {request.request_id: 1}
+    scheduler_output.total_num_scheduled_tokens = 1
+    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.num_invalid_spec_tokens = 0
+
+    model_runner_output = mocker.MagicMock()
+    model_runner_output.sampled_token_ids = [[42]]
+    model_runner_output.logprobs = None
+    model_runner_output.prompt_logprobs_dict = {}
+    model_runner_output.pooler_output = None
+    model_runner_output.num_nans_in_logits = None
+    model_runner_output.kv_connector_output = None
+    model_runner_output.cudagraph_stats = None
+    model_runner_output.req_id_to_index = {request.request_id: 0}
+    model_runner_output.routed_experts = None
+    model_runner_output.inter_stage_outputs = inter_stage_outputs
+
+    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+    return adapter
+
+
+@pytest.mark.parametrize(
+    ("final_stage_id", "force_kv_transfer", "expect_save"),
+    [
+        pytest.param(0, False, False, id="stage-zero-final"),
+        pytest.param(1, False, True, id="downstream-final"),
+    ],
+)
+def test_finish_saves_async_only_with_downstream_consumer(mocker, final_stage_id, force_kv_transfer, expect_save):
+    request = _make_chunk_request(final_stage_id, force_kv_transfer=force_kv_transfer)
+    adapter = _run_finished_save_step(mocker, request)
+
+    assert adapter.save_async.called is expect_save
+
+
+def test_stage_zero_final_skips_inter_stage_output(mocker):
+    warn = mocker.patch("vllm_omni.core.sched.omni_ar_scheduler.logger.warning")
+    request = _make_chunk_request(0)
+    adapter = _run_finished_save_step(mocker, request, inter_stage_outputs=[{"codes": 1}])
+
+    adapter.save_async.assert_not_called()
+    assert any("inter_stage_output is present" in str(call) for call in warn.call_args_list)
+
+
+def test_streaming_session_update_invalidates_omits_kv_transfer_cache():
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler._omits_kv_transfer_cache = {}
+    scheduler.skipped_waiting = []
+    scheduler.log_stats = False
+    tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
+    session = SimpleNamespace(
+        request_id="req",
+        additional_information=tagged.additional_information,
+        status=RequestStatus.WAITING,
+    )
+    assert scheduler._request_omits_chunk_transfer_to_next_stage(session)
+    assert "req" in scheduler._omits_kv_transfer_cache
+    update = SimpleNamespace(
+        additional_information={"omni_final_stage_id": 1},
+        model_intermediate_buffer=None,
+        arrival_time=0.0,
+        sampling_params=None,
+    )
+
+    OmniSchedulerMixin._finish_streaming_session_update(scheduler, session, update)
+
+    assert scheduler._omits_kv_transfer_cache == {}
+    assert session.additional_information == {"omni_final_stage_id": 1}
+    assert not scheduler._request_omits_chunk_transfer_to_next_stage(session)
+    assert not scheduler._request_omits_kv_transfer_to_next_stage(session)
+
+
+def test_omits_cache_misses_when_payload_object_is_replaced():
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler._omits_kv_transfer_cache = {}
+    tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
+    request = SimpleNamespace(
+        request_id="req",
+        additional_information=tagged.additional_information,
+    )
+    assert scheduler._request_omits_chunk_transfer_to_next_stage(request)
+    request.additional_information = {"omni_final_stage_id": 1}
+
+    assert not scheduler._request_omits_chunk_transfer_to_next_stage(request)

tests/core/sched/test_omni_ar_scheduler_stale_drain.py

@@ -40,6 +40,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 
 # isort: on
 
@@ -103,6 +104,7 @@ def _make_drain_sched(session: Request) -> MagicMock:
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
     return sched
 
 

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -26,6 +26,7 @@ from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARS
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 
 # isort: on
 
@@ -175,6 +176,7 @@ def _run_resumable_segment_stop(
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
 
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -312,6 +314,7 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
 
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -400,6 +403,7 @@ def test_stale_async_frame_is_dropped_before_output_processing() -> None:
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
 
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}

tests/core/sched/test_omni_sched_deferred_free_fence.py

@@ -126,6 +126,11 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request
     sched.max_num_running_reqs = 8
     sched._native_data_plane = False
     sched._retains_state_across_chunks = False
+    # Unset MagicMock attributes are truthy. Keep optional scheduling modes
+    # disabled so this fixture actually schedules its one-shot request.
+    sched._generation_defer_batch = False
+    sched._generation_first_chunk_only = False
+    sched._generation_max_regular_batch = 0
     sched._pause_state = PauseState.UNPAUSED
     sched.requests = {request.request_id: request}
     sched.running = [request]

tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 from types import SimpleNamespace
@@ -76,7 +79,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None:
     scheduler = OmniARScheduler.__new__(OmniARScheduler)
     scheduler.input_coordinator = coordinator
     scheduler.chunk_transfer_adapter = None
-    scheduler._omits_kv_transfer_cache = {"req-free": True}
+    scheduler._omits_kv_transfer_cache = {"req-free": (1, True, False)}
     scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None)
     scheduler.finished_req_ids = set()
     scheduler.finished_req_ids_dict = None

tests/diffusion/models/auk/test_auk_cudagraph.py

@@ -61,7 +61,10 @@ def test_single_request_graph_wrapper_cpu_falls_back_to_eager(cfg_strength: floa
 
     # The conditioning is prepared once per request and reused by every step.
     assert prepare_spy.call_count == 1
-    assert prepare_spy.call_args.args[-1] is (cfg_strength >= 1e-5)
+    # _prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg, timesteps)
+    assert prepare_spy.call_args.args[6] is (cfg_strength >= 1e-5)
+    # The time grid reaches prepare, so the adaLN modulations are computed once.
+    assert prepare_spy.call_args.args[7].numel() == 2
     assert step_spy.call_count == 2
     capture_spy.assert_not_called()
     assert not wrapper._cache
@@ -84,7 +87,8 @@ def test_graph_inputs_use_bounded_length_buckets() -> None:
     assert padded[4].shape == (1, 100, 4)
     assert padded[5].shape == (1, 100)
     assert [mask.sum().item() for mask in (padded[1], padded[3], padded[5])] == [65, 65, 51]
-    assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False)
+    assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False, 0)
+    assert wrapper._key(padded[0], padded[2], padded[4], False, 32) == (96, 96, 100, False, 32)
     assert wrapper.max_graphs == 32
 
 
@@ -153,18 +157,21 @@ def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_streng
             inputs["ref"],
             inputs["ref_mask"],
         )
-        key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5)
+        # sample_latents passes the grid, so the key carries its step count (two steps here).
+        key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5, 2)
         assert key in wrapper._cache
 
         # The static context was refreshed for this request's conditioning.
         entry = wrapper._cache[key]
-        fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5)
+        grid = torch.tensor([0.0, 0.4], device="cuda")
+        fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5, grid)
         for static, want in zip(entry.static_ctx.tensors(), fresh.tensors(), strict=True):
             if want is None:
                 assert static is None
             else:
                 torch.testing.assert_close(static, want)
         torch.testing.assert_close(entry.static_timestep, torch.tensor(0.4, device="cuda"))
+        assert int(entry.static_step) == 1
 
     assert len(wrapper._cache) == 1
 

tests/diffusion/models/auk/test_auk_dit_fastpaths.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The AuK DiT fast paths keep the reference numerics.
+
+Two paths are checked: the grouped conv position embedding as one batched
+GEMM, and the adaLN modulations precomputed for a whole time grid, eager and
+under the per-step CUDA graph.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.auk import auk_transformer as T
+from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, build_time_grid
+from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_dit(device: str = "cpu") -> AuKTransformer:
+    torch.manual_seed(5)
+    return (
+        AuKTransformer(
+            dim=32,
+            heads=2,
+            dim_head=16,
+            ff_mult=2,
+            latent_dim=4,
+            text_hidden_dim=8,
+            num_layers=2,
+            num_single_layers=2,
+        )
+        .eval()
+        .to(device)
+    )
+
+
+def _reference_conv_pos_embed(module: T.ConvPosEmbedding, x: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
+    """The channels-first conv stack as the reference implementation writes it."""
+    keep = None if mask is None else mask.unsqueeze(1)
+    x = x.transpose(1, 2)
+    if keep is not None:
+        x = x.masked_fill(~keep, 0.0)
+    for layer in module.conv1d:
+        x = layer(x)
+        if keep is not None and isinstance(layer, torch.nn.Conv1d):
+            x = x.masked_fill(~keep, 0.0)
+    return x.transpose(1, 2)
+
+
+@torch.inference_mode()
+def test_grouped_conv_as_batched_gemm_matches_conv1d() -> None:
+    torch.manual_seed(0)
+    conv = torch.nn.Conv1d(64, 64, 31, groups=16, padding=15)
+    x = torch.randn(2, 37, 64)
+    expected = conv(x.transpose(1, 2)).transpose(1, 2)
+    torch.testing.assert_close(T._grouped_conv1d_btc(x, conv), expected, atol=1e-5, rtol=1e-5)
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("masked", [False, True])
+def test_conv_pos_embedding_matches_the_channels_first_stack(masked: bool) -> None:
+    torch.manual_seed(1)
+    module = T.ConvPosEmbedding(64).eval()
+    x = torch.randn(2, 37, 64)
+    mask = None
+    if masked:
+        mask = torch.ones(2, 37, dtype=torch.bool)
+        mask[1, 30:] = False
+    torch.testing.assert_close(module(x, mask), _reference_conv_pos_embed(module, x, mask), atol=1e-5, rtol=1e-5)
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_infer", [False, True])
+def test_modulation_table_matches_per_step_modulation(cfg_infer: bool) -> None:
+    dit = _make_dit()
+    text = torch.randn(1, 5, 8)
+    ref = torch.randn(1, 3, 4)
+    x = torch.randn(1, 6, 4)
+    grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cpu")
+    plain = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer)
+    tabled = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer, timesteps=grid[:-1])
+    assert tabled.modulation is not None and tabled.modulation.shape[0] == 4
+    for i in range(4):
+        expected = dit.step(x, grid[i], plain)
+        torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=i), expected, atol=1e-5, rtol=1e-5)
+        index = torch.tensor(i)
+        torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=index), expected, atol=1e-5, rtol=1e-5)
+
+
+def _padded_inputs(device: str) -> dict[str, torch.Tensor]:
+    generator = torch.Generator(device="cpu").manual_seed(3)
+    text = torch.randn(1, 7, 8, generator=generator)
+    c_mask = torch.ones(1, 7, dtype=torch.bool)
+    c_mask[:, 5:] = False
+    text[:, 5:] = 0.0
+    ref = torch.randn(1, 6, 4, generator=generator)
+    ref_mask = torch.ones(1, 6, dtype=torch.bool)
+    ref_mask[:, 4:] = False
+    mask = torch.ones(1, 9, dtype=torch.bool)
+    mask[:, 7:] = False
+    inputs = {"text": text, "c_mask": c_mask, "ref": ref, "ref_mask": ref_mask, "mask": mask}
+    return {name: value.to(device) for name, value in inputs.items()}
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
+def test_graph_with_time_grid_matches_eager_with_time_grid(cfg_strength: float) -> None:
+    dit = _make_dit("cuda")
+    inputs = _padded_inputs("cuda")
+    grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cuda")
+    x = torch.randn(1, 9, 4, device="cuda")
+    wrapper = AuKCUDAGraphWrapper(dit)
+    common = dict(
+        text=inputs["text"],
+        c_mask=inputs["c_mask"],
+        ref=inputs["ref"],
+        ref_mask=inputs["ref_mask"],
+        cfg_strength=cfg_strength,
+        timesteps=grid[:-1],
+    )
+    for i in range(4):
+        graph = wrapper(x=x, timestep=grid[i], new_request=i == 0, step_index=i, **common)
+        ctx = dit.prepare(
+            inputs["text"],
+            target_len=9,
+            c_mask=inputs["c_mask"],
+            ref=inputs["ref"],
+            ref_mask=inputs["ref_mask"],
+            cfg_infer=cfg_strength >= 1e-5,
+            timesteps=grid[:-1],
+        )
+        eager = wrapper._step(x, grid[i], ctx, cfg_strength, i)
+        torch.testing.assert_close(graph, eager, atol=1e-4, rtol=1e-4)
+    # One graph serves every step of the grid.
+    assert len(wrapper._cache) == 1
+    assert next(iter(wrapper._cache))[-1] == 4

tests/diffusion/models/auk/test_auk_step_context.py

@@ -127,7 +127,9 @@ def test_hoisted_euler_loop_matches_per_step_forward(cfg_strength: float) -> Non
             v = dit(x, time=timesteps[i], **inputs)
         x = x + (timesteps[i + 1] - timesteps[i]) * v
 
-    torch.testing.assert_close(got, x, rtol=0, atol=0)
+    # sample_latents precomputes the adaLN modulations of the whole grid in one
+    # GEMM per layer, whose rows round differently from one-row GEMMs in fp32.
+    torch.testing.assert_close(got, x, rtol=0, atol=1e-6)
     assert dit.text_cond is None and dit.text_uncond is None
 
 

tests/diffusion/models/auk/test_auk_vae_fastpath.py

@@ -12,7 +12,8 @@ import pytest
 import torch
 
 from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE, LowPass, SnakeBeta, Upsample
+from vllm_omni.diffusion.models.auk import auk_vae
+from vllm_omni.diffusion.models.auk.auk_vae import AliasFreeActivation, AuKVAE, LowPass, SnakeBeta, Upsample
 from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph, plan_tiles
 from vllm_omni.platforms import current_omni_platform
 
@@ -246,6 +247,72 @@ def test_tiled_decode_matches_the_whole_decode() -> None:
         AuKVAEDecodeGraph(vae, tile_frames=60)
 
 
+@pytest.mark.cpu
+@torch.inference_mode()
+def test_fused_activation_is_skipped_off_cuda() -> None:
+    activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True))
+    x = torch.randn(1, 4, 16)
+    assert not activation._fused_eligible(x)
+    assert activation(x).shape == (1, 4, 16)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA")
+@torch.inference_mode()
+def test_fused_activation_is_skipped_without_triton(monkeypatch: pytest.MonkeyPatch) -> None:
+    activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True)).to("cuda")
+    x = torch.randn(1, 4, 16, device="cuda")
+    assert activation._fused_eligible(x)
+    monkeypatch.setattr(auk_vae, "HAS_TRITON", False)
+    assert not activation._fused_eligible(x)
+    assert activation(x).shape == (1, 4, 16)
+
+
+def _activation_input(layout: str, length: int) -> torch.Tensor:
+    """A ``[2, 6, length]`` input, contiguous or as a strided view of a larger buffer."""
+    if layout == "contiguous":
+        x = torch.randn(2, 6, length, device="cuda")
+    elif layout == "transposed":
+        x = torch.randn(2, length, 6, device="cuda").transpose(1, 2)
+    elif layout == "sliced":
+        x = torch.randn(2, 8, 2 * length, device="cuda")[:, 1:7, ::2]
+    else:
+        raise ValueError(layout)
+    assert x.shape == (2, 6, length)
+    if length > 1:  # contiguity of a single-sample view depends on the layout
+        assert x.is_contiguous() == (layout == "contiguous")
+    return x * 3
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA")
+@pytest.mark.parametrize("layout", ["contiguous", "transposed", "sliced"])
+@pytest.mark.parametrize("causal", [False, True])
+# The kernels tile the output in blocks of 1024 samples: the 2x oversampled
+# signal crosses a block boundary at 512 input samples and the decimated one at
+# 1024, so both sides of each boundary are covered, plus short and odd lengths.
+@pytest.mark.parametrize("length", [1, 7, 511, 512, 513, 1000, 1023, 1024, 1025])
+@torch.inference_mode()
+def test_fused_activation_matches_the_eager_modules(layout: str, causal: bool, length: int) -> None:
+    torch.manual_seed(0)
+    activation = AliasFreeActivation(SnakeBeta(6, alpha_logscale=True), causal=causal).to("cuda")
+    with torch.no_grad():
+        activation.act.alpha.normal_(0.0, 0.3)
+        activation.act.beta.normal_(0.0, 0.3)
+    x = _activation_input(layout, length)
+    original = x.clone()
+
+    eager = activation.downsample(activation.act(activation.upsample(x)))
+    assert activation._fused_eligible(x)
+    fused = activation(x)
+
+    # Same fp32 arithmetic, only the order in which the taps are summed differs.
+    assert fused.shape == eager.shape
+    torch.testing.assert_close(fused, eager, atol=1e-5, rtol=1e-5)
+    # Neither path writes into its (possibly strided) input.
+    torch.testing.assert_close(x, original, atol=0, rtol=0)
+
+
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
 @torch.inference_mode()

tests/diffusion/models/auk/test_pipeline_auk.py

@@ -291,6 +291,18 @@ class TestRequestParsing:
         with pytest.raises(ValueError, match="AuK max_dit_graphs must be a positive integer"):
             build_pipeline(model_config={"max_dit_graphs": max_dit_graphs})
 
+    @pytest.mark.parametrize(
+        "model_config, expected", [({}, 16), ({"auk_ref_cache_size": 0}, 0), ({"auk_ref_cache_size": 4}, 4)]
+    )
+    def test_ref_cache_size_is_taken_from_model_config(self, build_pipeline, model_config, expected):
+        pipeline, _ = build_pipeline(model_config=model_config)
+        assert pipeline._ref_cache_size == expected
+
+    @pytest.mark.parametrize("size", [-1, 1.9, True, False, "4", None])
+    def test_invalid_ref_cache_size_is_rejected(self, build_pipeline, size):
+        with pytest.raises(ValueError, match="AuK auk_ref_cache_size must be a non-negative integer"):
+            build_pipeline(model_config={"auk_ref_cache_size": size})
+
     def test_pipeline_declares_audio_output(self, build_pipeline):
         pipeline, _ = build_pipeline()
 
@@ -492,6 +504,40 @@ class TestRequestParsing:
         assert call["sample"] is True
         assert isinstance(call["generator"], torch.Generator)
 
+    def test_reference_latents_are_cached_by_content(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+        clip = _silence(1.0)
+
+        for _ in range(3):
+            pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
+        # A different clip is a new entry, the first one still hits.
+        other = (np.full_like(clip[0], 0.25), clip[1])
+        pipeline.forward(_batch(_prompt(audio=other, knobs={"gen_seconds": 1.0}), seed=1))
+        pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
+
+        assert len(pipeline.vae.encode_calls) == 2
+        assert calls[0]["ref"] is calls[1]["ref"] is calls[4]["ref"]
+
+    def test_posterior_draws_and_disabled_cache_always_encode(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+        for _ in range(2):
+            pipeline.forward(
+                _batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "vae_sample": True}), seed=1)
+            )
+        assert len(pipeline.vae.encode_calls) == 2
+
+        pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 0})
+        for _ in range(2):
+            pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0}), seed=1))
+        assert len(pipeline.vae.encode_calls) == 2
+
+    def test_reference_cache_is_bounded(self, build_pipeline):
+        pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 2})
+        for level in (0.1, 0.2, 0.3):
+            clip = (np.full(SAMPLE_RATE, level, dtype=np.float32), SAMPLE_RATE)
+            pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
+        assert len(pipeline._ref_cache) == 2
+
     def test_one_request_per_forward(self, build_pipeline):
         pipeline, _ = build_pipeline()
         batch = _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1)

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -15,6 +15,7 @@ from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
 from vllm.v1.request import Request, RequestStatus
 
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct
@@ -95,6 +96,30 @@ def test_shm_abort_reclaims_inflight_put_without_blocking(shm_sender, monkeypatc
     assert not adapter.code_prompt_token_ids
 
 
+def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender):
+    adapter, connector = shm_sender
+    ext_id = f"release_{uuid.uuid4().hex}"
+    request = _req("done", RequestStatus.RUNNING, external_req_id=ext_id)
+    sibling_key = f"{ext_id}_1_0_0"
+    assert connector.put("0", "1", sibling_key, "sibling")[0]
+    for _ in range(3):
+        adapter.save_async(None, request)
+        adapter._send_single_request(adapter._pending_save_reqs.popleft())
+    # The consumer read only the first chunk before its own request ended.
+    assert connector.get("0", "1", f"{ext_id}_0_0") is not None
+
+    adapter.release_shm_resources(ext_id)
+    # The release is queued behind sends; nothing is unlinked on the caller.
+    assert f"{ext_id}_0_1" in connector._pending_keys
+    while adapter._pending_save_reqs:
+        adapter._send_single_request(adapter._pending_save_reqs.popleft())
+
+    assert adapter._reclaimed_shm_total == 2
+    assert not any(k.startswith(f"{ext_id}_0_") for k in connector._pending_keys)
+    assert connector.get("0", "1", f"{ext_id}_0_1") is None
+    assert connector.get("0", "1", sibling_key)[0] == "sibling"
+
+
 @pytest.mark.parametrize("fail_terminal", [False, True])
 def test_shm_terminal_keeps_payload_only_after_success(shm_sender, monkeypatch, fail_terminal):
     adapter, connector = shm_sender
@@ -2620,10 +2645,9 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
     adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
     adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw))
 
-    from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
-
     scheduler = mocker.MagicMock()
     scheduler.chunk_transfer_adapter = adapter_mock
+    bind_omits_transfer_helpers(scheduler)
     scheduler.connector = None
     scheduler.perf_metrics = None
     scheduler.log_stats = False

tests/distributed/omni_connectors/test_shm_connector.py

@@ -213,7 +213,7 @@ class TestCleanup:
         connector.put("s0", "s1", "cleanup_req_42", data)
         assert "cleanup_req_42" in connector._pending_keys
 
-        connector.cleanup("cleanup_req_42")
+        assert connector.cleanup("cleanup_req_42") is True
         assert "cleanup_req_42" not in connector._pending_keys
 
         result = connector.get("s0", "s1", "cleanup_req_42", metadata=None)
@@ -224,7 +224,7 @@ class TestCleanup:
         connector.put("s0", "s1", "consumed_req_99", data)
         connector.get("s0", "s1", "consumed_req_99", metadata=None)
 
-        connector.cleanup("consumed_req_99")
+        assert connector.cleanup("consumed_req_99") is False
         assert "consumed_req_99" not in connector._pending_keys
 
     def test_close_cleans_all_pending(self, connector):
@@ -296,6 +296,24 @@ def test_consumed_bad_payload_removes_lock_file(connector, monkeypatch):
     assert key not in connector._pending_keys
 
 
+@pytest.mark.parametrize("suffix", ["_1_2", "_1", "_0", "_suffix"])
+def test_cleanup_prefix_only_matches_own_chunk_keys(connector, suffix):
+    request_id = f"prefix_{uuid.uuid4().hex}"
+    own_keys = [f"{request_id}_0_{i}" for i in range(3)]
+    sibling_key = f"{request_id}{suffix}_0_0"
+    for key in own_keys:
+        assert connector.put("0", "1", key, key)[0]
+    assert connector.put("0", "1", sibling_key, "sibling")[0]
+    # A consumed chunk is no longer counted as reclaimed.
+    assert connector.get("0", "1", own_keys[0])[0] == own_keys[0]
+
+    assert connector.cleanup_prefix(f"{request_id}_0_") == 2
+    for key in own_keys:
+        assert key not in connector._pending_keys
+        assert connector.get("0", "1", key) is None
+    assert connector.get("0", "1", sibling_key)[0] == "sibling"
+
+
 # ── Arrival wakeups ───────────────────────────────────────────────────
 
 

tests/e2e/offline_inference/test_auk_encoder_cudagraph.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""AuK encoder stage: FULL CUDA graph replay matches eager, request by request.
+
+``auk.yaml`` captures the encoder prefill as FULL CUDA graphs padded to a few
+token buckets. A replay reuses static input buffers and the graph pool, so a
+request could in principle read conditioning left behind by an earlier, longer
+prompt. This test runs the same request sequence through an eager encoder and
+through the graphed one and compares the text condition stage 0 hands to
+stage 1 (the fused thinker hidden states) for every request. The sequence
+alternates prompts of different lengths inside one capture bucket, crosses
+into a larger bucket, and then returns to the first prompt.
+
+Like ``test_auk.py`` it needs an assembled checkpoint directory:
+
+    VLLM_OMNI_AUK_MODEL_DIR=/path/to/auk-omni python -m pytest tests/e2e/offline_inference/test_auk_encoder_cudagraph.py
+"""
+
+from __future__ import annotations
+
+import bisect
+import os
+from pathlib import Path
+from typing import Any
+
+import pytest
+import torch
+import yaml
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.model_executor.stage_input_processors import auk as auk_input_processor
+from vllm_omni.model_extras.auk import auk_prompt, auk_sampling_params
+
+MODEL_DIR_ENV = "VLLM_OMNI_AUK_MODEL_DIR"
+
+# Text-only instruct prompts: the token count, and so the capture bucket,
+# follows the text length. SHORT_A and SHORT_B differ in length but share a
+# bucket; LONG lands in a larger one. The test checks these properties on the
+# recorded condition lengths, so a tokenizer change fails loudly instead of
+# silently shrinking the coverage.
+_TEMPLATE = 'Generate speech based on the following description: "{desc}". The content to speak is: "{text}".'
+SHORT_A = _TEMPLATE.format(
+    desc="A calm young woman speaking warmly.",
+    text="Welcome back, how was your day? I kept some soup warm for you, and the kettle is on if you want tea.",
+)
+SHORT_B = _TEMPLATE.format(
+    desc="An older man with a deep, slightly raspy voice, speaking slowly.",
+    text=(
+        "The train to the coast leaves at seven, so please be on time. Bring a warm coat, "
+        "some cash for the ferry, and the blue umbrella from the hall, because the forecast says rain."
+    ),
+)
+LONG = _TEMPLATE.format(
+    desc=(
+        "A cheerful radio host in her thirties, bright and energetic, with clear articulation, "
+        "a slight smile in her voice, and a brisk but relaxed pace, recorded close to the microphone"
+    ),
+    text=(
+        "Good morning and welcome to the weekend edition. Today we are talking about small gardens, "
+        "the plants that thrive on a sunny balcony, and the neighbours who turned an empty lot into "
+        "a place where the whole street now meets on Saturday afternoons to trade seedlings and stories. "
+        "Stay with us, because after the news we will open the phone lines and hear from you."
+    ),
+)
+SEQUENCE = [SHORT_A, SHORT_B, SHORT_A, LONG, SHORT_B, SHORT_A]
+
+# The DiT output is not inspected; keep stage 1 as cheap as it goes (it also
+# runs eager in both configs, see _deploy).
+GEN_SECONDS = 1.0
+NFE = 1
+
+# bf16 hidden states. Graph replay pads the prefill to the bucket size but runs
+# the same kernels on the real tokens, so the two paths agree to bf16 rounding.
+# Measured on H20: cosine >= 0.99998 and max-abs error <= 2.9e-3 of the
+# condition's max magnitude for every request; the bounds leave about 3x
+# headroom for other GPUs.
+MIN_COSINE = 0.9999
+MAX_REL_ABS = 1e-2
+
+_model_dir = os.environ.get(MODEL_DIR_ENV)
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.tts,
+    pytest.mark.skipif(not _model_dir, reason=f"set {MODEL_DIR_ENV} to an assembled AuK directory"),
+]
+
+
+def _deploy(tmp_path: Path, *, graph_encoder: bool) -> str:
+    """``auk.yaml`` with stage 1 eager, and the encoder stage eager or graphed as shipped.
+
+    Only the encoder differs between the two runs. The conditions compared here
+    are recorded before stage 1 runs, so its DiT compile, graph capture and
+    codec bucket warm-up would only add startup time; eager stage 1 skips them
+    and still takes the real stage handoff.
+    """
+    with open(get_deploy_config_path("auk.yaml")) as f:
+        deploy = yaml.safe_load(f)
+    for stage in deploy["stages"]:
+        if stage["stage_id"] == 1:
+            stage["enforce_eager"] = True
+        elif not graph_encoder:
+            stage["enforce_eager"] = True
+            stage.pop("compilation_config", None)
+    path = tmp_path / f"auk_{'graph' if graph_encoder else 'eager'}_encoder.yaml"
+    path.write_text(yaml.safe_dump(deploy))
+    return str(path)
+
+
+def _graph_capture_sizes() -> list[int]:
+    with open(get_deploy_config_path("auk.yaml")) as f:
+        deploy = yaml.safe_load(f)
+    stage = next(s for s in deploy["stages"] if s["stage_id"] == 0)
+    assert stage.get("enforce_eager") is False, "auk.yaml no longer graphs the encoder; update this test"
+    return sorted(stage["compilation_config"]["cudagraph_capture_sizes"])
+
+
+def _run_sequence(deploy_config: str, monkeypatch: pytest.MonkeyPatch) -> list[torch.Tensor]:
+    """Run SEQUENCE one request at a time and return each request's text condition."""
+    conditions: list[torch.Tensor] = []
+    original = auk_input_processor.encoder2dit
+
+    def recording_encoder2dit(*args: Any, **kwargs: Any) -> dict[str, Any] | None:
+        result = original(*args, **kwargs)
+        if result is not None:
+            conditions.append(result["prompt_embeds"].detach().to("cpu", torch.float32).clone())
+        return result
+
+    # The orchestrator runs in this process and resolves the stage input
+    # function by module attribute at startup, so the patch must precede Omni.
+    with monkeypatch.context() as patch:
+        patch.setattr(auk_input_processor, "encoder2dit", recording_encoder2dit)
+        with OmniRunner(str(Path(_model_dir).resolve()), deploy_config=deploy_config) as runner:
+            for text in SEQUENCE:
+                prompt = auk_prompt(text, None, gen_seconds=GEN_SECONDS)
+                outputs = runner.omni.generate(prompt, auk_sampling_params(nfe=NFE, cfg=2.0, seed=0))
+                assert outputs, "no outputs returned"
+    assert len(conditions) == len(SEQUENCE), f"recorded {len(conditions)} conditions for {len(SEQUENCE)} requests"
+    return conditions
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_encoder_graph_replay_matches_eager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    capture_sizes = _graph_capture_sizes()
+    eager = _run_sequence(_deploy(tmp_path, graph_encoder=False), monkeypatch)
+    graphed = _run_sequence(_deploy(tmp_path, graph_encoder=True), monkeypatch)
+
+    lengths = {text: cond.shape[0] for text, cond in zip(SEQUENCE, eager)}
+    for i, (ref, got) in enumerate(zip(eager, graphed)):
+        assert got.shape == ref.shape, f"request {i}: shape {tuple(got.shape)} != eager {tuple(ref.shape)}"
+        assert torch.isfinite(got).all(), f"request {i}: non-finite condition"
+        cosine = torch.nn.functional.cosine_similarity(got.flatten(), ref.flatten(), dim=0).item()
+        rel_abs = ((got - ref).abs().max() / ref.abs().max()).item()
+        assert cosine >= MIN_COSINE and rel_abs <= MAX_REL_ABS, (
+            f"request {i} ({lengths[SEQUENCE[i]]} tokens): cosine {cosine:.6f}, max-abs error {rel_abs:.2e} "
+            "of the condition's magnitude"
+        )
+
+    # The sequence must exercise what it claims to: two different lengths in
+    # one bucket, a second bucket, and every prompt inside the graphed range.
+    assert max(lengths.values()) <= capture_sizes[-1], f"{lengths} exceeds the largest capture size"
+    bucket = {text: capture_sizes[bisect.bisect_left(capture_sizes, n)] for text, n in lengths.items()}
+    assert lengths[SHORT_A] != lengths[SHORT_B] and bucket[SHORT_A] == bucket[SHORT_B], (lengths, bucket)
+    assert bucket[LONG] > bucket[SHORT_A], (lengths, bucket)
+
+    # Repeats of a prompt reproduce its first condition exactly, eager or graphed:
+    # nothing from the requests in between leaks into the replay.
+    for run in (eager, graphed):
+        first: dict[str, torch.Tensor] = {}
+        for text, cond in zip(SEQUENCE, run):
+            if text in first:
+                torch.testing.assert_close(cond, first[text], atol=0, rtol=0)
+            else:
+                first[text] = cond

tests/e2e/online_serving/test_qwen3_omni_shm_leak.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Live /dev/shm leak check for Qwen3-Omni ``async_chunk`` (stage-0-final).
+
+The default L2 Qwen3-Omni job uses ``ci/qwen3_omni_moe.yaml`` with
+``async_chunk: false`` plus ``--no-async-chunk``, so it never exercises the
+thinker→talker SharedMemoryConnector path.
+
+A text-only request is tagged ``omni_final_stage_id=0``, so nothing downstream
+ever ``get()``s what stage-0 ``put()``s for it. Before #7245 the scheduler still
+``put()`` a finished marker on finish and nothing unlinked it, so every such
+request left a ``SharedMemoryConnector`` segment plus lockfile in ``/dev/shm``.
+Stage-0-final finishes now skip the ``put()`` and the orchestrator reclaims any
+undrained segment on request cleanup.
+
+This module starts the production async-chunk deploy and asserts the three
+text-only request ids leave no connector lockfile or POSIX segment in
+``/dev/shm``.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from pathlib import Path
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+_USE_PD = os.environ.get("VLLM_TEST_PD_MODE", "0") == "1"
+_MODEL = os.environ.get("VLLM_OMNI_TEST_MODEL", "Qwen/Qwen3-Omni-30B-A3B-Instruct")
+_SHM_DIR = Path("/dev/shm")
+_NUM_TEXT_REQUESTS = 3
+_SETTLE_S = 15.0
+
+# CI overlay keeps seq/token limits small; flip async_chunk back on so stage-0
+# still uses SharedMemoryConnector. Production yaml is already async_chunk:true.
+_ASYNC_CHUNK_DEPLOY = modify_stage_config(
+    get_deploy_config_path("ci/qwen3_omni_moe.yaml"),
+    updates={"async_chunk": True},
+)
+
+test_params = [
+    pytest.param(
+        OmniServerParams(
+            model=_MODEL,
+            stage_config_path=_ASYNC_CHUNK_DEPLOY,
+            use_stage_cli=True,
+        ),
+        id="async_chunk",
+    )
+]
+
+
+def _system_prompt() -> dict:
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, "
+                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                    "as well as generating text and speech."
+                ),
+            }
+        ],
+    }
+
+
+def _shm_used_mib() -> str:
+    try:
+        st = os.statvfs(_SHM_DIR)
+        used = (st.f_blocks - st.f_bavail) * st.f_frsize
+        return f"{used / (1024 * 1024):.1f}MiB"
+    except OSError:
+        return "unknown"
+
+
+def _is_entry_for_request(name: str, request_id: str) -> bool:
+    """Match a segment or ``shm_`` lockfile keyed by the ``{request_id}-{suffix}`` orchestrator id."""
+    return name.removeprefix("shm_").startswith((f"{request_id}-", f"{request_id}_"))
+
+
+def _entries_for_request_ids(request_ids: list[str]) -> list[str]:
+    if not _SHM_DIR.is_dir():
+        pytest.skip("/dev/shm is not available")
+    found: list[str] = []
+    with os.scandir(_SHM_DIR) as entries:
+        for entry in entries:
+            if any(_is_entry_for_request(entry.name, request_id) for request_id in request_ids):
+                found.append(entry.name)
+    return found
+
+
+def _wait_no_entries_for_request_ids(request_ids: list[str], timeout_s: float = _SETTLE_S) -> list[str]:
+    """Reclaim is fire-and-forget; poll until these ids are gone or *timeout_s*."""
+    deadline = time.monotonic() + timeout_s
+    time.sleep(1.0)
+    leftover: list[str] = _entries_for_request_ids(request_ids)
+    while leftover:
+        if time.monotonic() >= deadline:
+            return leftover
+        time.sleep(2.0)
+        leftover = _entries_for_request_ids(request_ids)
+    return leftover
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.")
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_only_async_chunk_does_not_leak_shm(omni_server, online_client) -> None:
+    """Stage-0-final (text-only) requests must not leave connector SHM behind."""
+    print(f"[shm-leak] before requests: /dev/shm_used={_shm_used_mib()}", flush=True)
+
+    messages = dummy_messages_from_mix_data(
+        system_prompt=_system_prompt(),
+        content_text="What is the capital of China? Answer in 20 words.",
+    )
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": False,
+        "modalities": ["text"],
+        "key_words": {"text": ["beijing"]},
+    }
+    responses = online_client.send_omni_request(request_config, request_num=_NUM_TEXT_REQUESTS)
+    produced = [resp.text_content for resp in responses if getattr(resp, "text_content", None)]
+    assert len(produced) == _NUM_TEXT_REQUESTS, (
+        f"Need {_NUM_TEXT_REQUESTS} text completions to exercise stage-0 finish; "
+        f"got {len(produced)} non-empty texts. A preprocess/engine miss would "
+        f"leave /dev/shm unchanged and hide a leak."
+    )
+    request_ids = [resp.request_id for resp in responses]
+    assert all(request_ids) and len(request_ids) == _NUM_TEXT_REQUESTS, (
+        f"Need chat completion ids to assert per-request SHM cleanup; got {request_ids}"
+    )
+
+    leftover = _wait_no_entries_for_request_ids(request_ids)
+    print(
+        f"[shm-leak] after {_NUM_TEXT_REQUESTS} text-only requests: "
+        f"/dev/shm_used={_shm_used_mib()} leftover={leftover} request_ids={request_ids}",
+        flush=True,
+    )
+    assert not leftover, (
+        f"SharedMemoryConnector leaked /dev/shm entries for {leftover} after stage-0-final request(s) {request_ids}"
+    )

tests/engine/test_cfg_companion_lifecycle.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
 """Deterministic tests for the CFG-companion output lifecycle.
 
 Regression suite for the companion-output race behind the nightly Bagel
@@ -94,6 +94,9 @@ class _FakePool:
     def release_bindings(self, request_ids):
         pass
 
+    async def release_request_resources(self, request_ids):
+        pass
+
 
 def _make_orchestrator(num_stages: int = 2) -> Orchestrator:
     orch = Orchestrator.__new__(Orchestrator)

tests/engine/test_stage_engine_core_proc.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import patch
 
 import pytest
@@ -25,11 +26,15 @@ def test_first_audio_sink_preserves_first_decoder_and_opts_in_stream_decoder(
     expected,
 ):
     executor = UniProcExecutor.__new__(UniProcExecutor)
+    executor.vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)
+    )
     executor.driver_worker = mocker.Mock()
     runner = executor.driver_worker.worker.model_runner
     runner.model.first_frame_decoder = object() if first_decoder else None
     runner.model.stream_decoder = object() if stream_decoder else None
     runner.model.stream_first_audio = stream_first_audio
+    runner._omni_data_plane = None
     assert _bind_first_audio_sink(executor, mocker.Mock(), mocker.Mock()) is expected
     assert runner.model_state.set_first_audio_sink.call_count == int(expected)
 
@@ -54,3 +59,64 @@ def test_preprocess_add_request_preserves_omni_fields():
     assert current_wave == 3
     assert result.external_req_id == "external"
     assert result.additional_information == {"conditioning": "payload"}
+
+
+def test_codec_uses_generic_first_audio_binding(monkeypatch):
+    import queue
+
+    import torch
+
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
+    from vllm_omni.engine import stage_engine_core_proc as module
+
+    calls: dict[str, Any] = {}
+
+    def hook(*args):
+        return False
+
+    def bind(sink):
+        calls["sink"] = sink
+        return hook
+
+    plane = SimpleNamespace(set_first_chunk_hook=lambda value: calls.update(hook=value))
+    model = SimpleNamespace(bind_first_chunk_fast_path=bind)
+    runner = SimpleNamespace(model=model, get_model=lambda: model, _omni_data_plane=plane)
+    executor = SimpleNamespace(
+        vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)),
+        driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)),
+    )
+    monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace)
+    outputs: queue.Queue = queue.Queue()
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=2)})
+    assert module._bind_first_audio_sink(executor, outputs, scheduler)
+    assert calls["hook"] is hook
+    calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000))
+    client, batch = outputs.get_nowait()
+    assert client == 2 and FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output
+
+
+def test_first_audio_binding_preserves_talker_marker(monkeypatch):
+    import queue
+
+    import torch
+
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
+    from vllm_omni.engine import stage_engine_core_proc as module
+
+    calls: dict[str, Any] = {}
+    runner = SimpleNamespace(
+        model=SimpleNamespace(first_frame_decoder=object()),
+        model_state=SimpleNamespace(set_first_audio_sink=lambda sink: calls.update(sink=sink)),
+    )
+    executor = SimpleNamespace(
+        vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)),
+        driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)),
+    )
+    monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace)
+    outputs: queue.Queue = queue.Queue()
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=0)})
+    assert module._bind_first_audio_sink(executor, outputs, scheduler)
+    calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000))
+    assert outputs.get_nowait()[1].outputs[0].multimodal_output[FIRST_AUDIO_KEY]
+    executor.vllm_config.parallel_config.tensor_parallel_size = 2
+    assert not module._bind_first_audio_sink(executor, outputs, scheduler)

tests/engine/test_stage_pool_collective_rpc.py

@@ -154,6 +154,46 @@ def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_release_request_resources_skips_without_async_chunk():
+    async def run() -> None:
+        call = AsyncMock()
+        pool = StagePool(
+            0,
+            [SimpleNamespace(call_utility_async=call)],  # type: ignore[list-item]
+            stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=False)),
+        )
+
+        await pool.release_request_resources(["req-1"])
+        call.assert_not_awaited()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_release_request_resources_times_out_hung_replica_without_blocking_others(monkeypatch):
+    async def run() -> None:
+        async def hang(*_args):
+            await asyncio.Event().wait()
+
+        live = AsyncMock()
+        pool = StagePool(
+            0,
+            [SimpleNamespace(call_utility_async=hang), SimpleNamespace(call_utility_async=live)],  # type: ignore[list-item]
+            stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=True)),
+        )
+        monkeypatch.setattr(pool, "RELEASE_RPC_TIMEOUT_S", 0.2)
+
+        release = asyncio.create_task(pool.release_request_resources(["req-1"]))
+        await asyncio.sleep(0.05)
+        live.assert_awaited_once_with("omni_release_request_resources", ["req-1"])
+        assert not release.done()
+
+        await asyncio.wait_for(release, timeout=1.0)
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 @pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"])
 @pytest.mark.parametrize("failure", ["error", "timeout", "missing"])

tests/entrypoints/openai_api/test_moss_local_prompt.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The Local-v1.5 continuation prompt fast path equals the processor layout."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import _local_continuation_prompt
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Processor:
+    """Segment API of the model's processor with a deterministic tokenizer."""
+
+    def __init__(self):
+        self.model_config = SimpleNamespace(
+            n_vq=3,
+            im_start_token_id=1,
+            im_end_token_id=2,
+            audio_start_token_id=3,
+            audio_pad_token_id=1024,
+            audio_assistant_slot_token_id=5,
+        )
+        self.encoded = []
+
+    def _encode_text(self, text):
+        self.encoded.append(text)
+        return [100 + ord(c) % 50 for c in text]
+
+    def _user_prompt_prefix_ids(self):
+        return [1] + self._encode_text("user\n") + self._encode_text("<ref>")
+
+    def _user_prompt_after_reference_ids(self, language, fields):
+        return self._encode_text(f"lang={fields.get('language', language)}")
+
+    def _assistant_prompt_prefix_ids(self):
+        return self._encode_text("</user>") + [2] + self._encode_text("\n") + [1] + self._encode_text("assistant")
+
+    def reference(self, text, language, codes):
+        # Mirrors the processor's _build_continuation_codes + _pad (one item).
+        fields = {} if language is None else {"language": language}
+        ids = (
+            self._user_prompt_prefix_ids()
+            + self._encode_text("None")
+            + self._user_prompt_after_reference_ids(language, fields)
+            + self._encode_text(text)
+            + self._assistant_prompt_prefix_ids()
+            + [3]
+        )
+        text_rows = torch.full((len(ids), 4), 1024, dtype=torch.long)
+        text_rows[:, 0] = torch.tensor(ids)
+        audio_rows = torch.full((codes.shape[0], 4), 1024, dtype=torch.long)
+        audio_rows[:, 0] = 5
+        audio_rows[:, 1:] = codes
+        unified = torch.cat([text_rows, audio_rows])
+        return unified[:, 0].tolist(), unified[:, 1:].contiguous()
+
+
+@pytest.mark.parametrize("language", [None, "English"])
+def test_matches_processor_and_caches_constant_segments(language):
+    proc = _Processor()
+    codes = torch.randint(0, 1024, (7, 3), dtype=torch.int32)
+    kwargs = {"text": "ref words target words"}
+    if language is not None:
+        kwargs["language"] = language
+    ids, audio = _local_continuation_prompt(proc, kwargs, [codes])
+    expected_ids, expected_audio = proc.reference(kwargs["text"], language, codes.long())
+    assert ids == expected_ids
+    assert audio.dtype == torch.int64 and torch.equal(audio, expected_audio)
+
+    proc.encoded.clear()
+    _local_continuation_prompt(proc, {**kwargs, "text": "other"}, [codes])
+    assert proc.encoded == ["other"]
+
+
+def test_unsupported_layouts_fall_back():
+    proc = _Processor()
+    codes = torch.zeros((2, 3), dtype=torch.long)
+    assert _local_continuation_prompt(proc, {"text": "t", "instruction": "x"}, [codes]) is None
+    assert _local_continuation_prompt(proc, {"text": "t"}, [codes, codes]) is None
+    assert _local_continuation_prompt(proc, {"text": "t"}, ["path.wav"]) is None
+    assert _local_continuation_prompt(proc, {"text": "t"}, [torch.zeros((2, 4), dtype=torch.long)]) is None

tests/entrypoints/openai_api/test_tts_adapter.py

@@ -691,6 +691,45 @@ def test_moss_get_processor_places_audio_tokenizer_on_codec_stage_device(mocker)
     audio_tokenizer.to.assert_called_once()
 
 
+@pytest.mark.parametrize(
+    "variant,graphs_env,prepares",
+    [("local", None, False), ("local", "1", True), ("local", "0", False), ("realtime", "1", False)],
+)
+def test_moss_warmup_prepares_the_reference_encoder(variant, graphs_env, prepares, mocker, monkeypatch):
+    if graphs_env is None:
+        monkeypatch.delenv("VLLM_OMNI_MOSS_REF_GRAPHS", raising=False)
+    else:
+        monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", graphs_env)
+    adapter = _moss_adapter_with_stage_configs([], mocker)
+    adapter._moss_variant = variant
+    encoder = mocker.Mock()
+    mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
+    asyncio.run(adapter.warmup())
+    assert encoder.prepare.called is prepares
+
+
+def test_moss_warmup_failure_does_not_stop_startup(mocker, monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1")
+    adapter = _moss_adapter_with_stage_configs([], mocker)
+    encoder = mocker.Mock()
+    encoder.prepare.side_effect = RuntimeError("capture failed")
+    mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
+    asyncio.run(adapter.warmup())
+    encoder.prepare.assert_called_once_with()
+
+
+def test_moss_shared_encoder_startup_failure_stops_api_startup(mocker, monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1")
+    from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import SharedReferenceEncoderStartupError
+
+    adapter = _moss_adapter_with_stage_configs([], mocker)
+    encoder = mocker.Mock()
+    encoder.prepare.side_effect = SharedReferenceEncoderStartupError("still compiling")
+    mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
+    with pytest.raises(SharedReferenceEncoderStartupError):
+        asyncio.run(adapter.warmup())
+
+
 def test_qwen3_tts_metadata():
     assert Qwen3TTSAdapter.backend == "ar"
     assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)

tests/entrypoints/test_pd_disaggregation.py

@@ -288,6 +288,9 @@ def _setup_log_mocks(monkeypatch):
         def build_and_log_summary(self):
             return "Fake summary"
 
+        def log_timing_summary(self):
+            pass
+
     monkeypatch.setattr(
         "vllm_omni.entrypoints.omni.OrchestratorAggregator",
         _FakeOrchestratorAggregator,

tests/helpers/assertions.py

@@ -4,9 +4,11 @@
 """Assertion and response validation helpers for tests."""
 
 import base64
+import hashlib
 import io
 import json
 import math
+import os
 import tempfile
 import threading
 import wave
@@ -999,6 +1001,31 @@ def _assert_transcript_matches(
     )
 
 
+def _retain_failed_speech_audio(
+    audio_bytes: bytes | None, transcript: str, expected_text: Any, error: AssertionError
+) -> None:
+    """Retain the exact ASR input when a CI speech-content assertion fails."""
+    directory = os.environ.get("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR")
+    if not directory or not audio_bytes:
+        return
+    stem = f"failed_speech_{hashlib.sha256(audio_bytes).hexdigest()[:16]}"
+    root = Path(directory)
+    try:
+        root.mkdir(parents=True, exist_ok=True)
+        (root / f"{stem}.wav").write_bytes(audio_bytes)
+        (root / f"{stem}.json").write_text(
+            json.dumps(
+                {"transcript": transcript, "expected_text": str(expected_text), "assertion": str(error)},
+                ensure_ascii=False,
+                indent=2,
+            )
+        )
+        print(f"Retained failed speech ASR input: {root / f'{stem}.wav'}")
+    except OSError as exc:
+        # Diagnostic I/O must never replace or suppress the content failure.
+        print(f"Could not retain failed speech ASR input: {exc}")
+
+
 def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str | None = None) -> None:
     """Validate speech API results from :class:`~tests.helpers.client.OmniResponse`.
 
@@ -1063,14 +1090,21 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
             if expected_text:
                 print(f"audio content is: {transcript}")
                 print(f"input text is: {expected_text}")
-                _assert_transcript_matches(
-                    transcript,
-                    _speech_audio_for_transcription(getattr(response, "audio_bytes", None), request_config),
-                    expected_text,
-                    threshold=0.9,
-                    escalation_model=request_config.get("transcript_escalation_model"),
-                    language=request_config.get("transcript_language"),
+                transcription_audio = _speech_audio_for_transcription(
+                    getattr(response, "audio_bytes", None), request_config
                 )
+                try:
+                    _assert_transcript_matches(
+                        transcript,
+                        transcription_audio,
+                        expected_text,
+                        threshold=0.9,
+                        escalation_model=request_config.get("transcript_escalation_model"),
+                        language=request_config.get("transcript_language"),
+                    )
+                except AssertionError as exc:
+                    _retain_failed_speech_audio(transcription_audio, transcript, expected_text, exc)
+                    raise
         _assert_preset_voice_gender_from_audio(
             response.audio_bytes,
             request_config.get("voice"),

tests/helpers/client.py

@@ -97,6 +97,7 @@ def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[
 class OmniResponse:
     """Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``)."""
 
+    request_id: str | None = None
     text_content: str | None = None
     audio_data: list[str] | None = None
     audio_content: str | None = None
@@ -421,6 +422,8 @@ class OnlineOmniClient:
             text_content = ""
             audio_data = []
             for chunk in chat_completion:
+                if result.request_id is None:
+                    result.request_id = getattr(chunk, "id", None)
                 for choice in chunk.choices:
                     content = getattr(getattr(choice, "delta", None), "content", None)
                     modality = getattr(chunk, "modality", None)
@@ -453,6 +456,7 @@ class OnlineOmniClient:
         """Wall clock from *before* ``chat.completions.create`` through response parse + local decode."""
         result = OmniResponse()
         try:
+            result.request_id = getattr(chat_completion, "id", None)
             audio_data = None
             text_content = None
             for choice in chat_completion.choices:

tests/helpers/omni_scheduler.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Scheduler test helpers shared by AR / chunk-transfer unit tests."""
+
+from __future__ import annotations
+
+from types import MethodType
+
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+
+def bind_omits_transfer_helpers(sched) -> None:
+    """Bind the real omit helpers so a MagicMock scheduler does not skip save_async.
+
+    MagicMock is truthy, so an unbound ``_request_omits_chunk_transfer_to_next_stage``
+    silently skips the chunk put.
+    """
+    sched._omits_kv_transfer_cache = {}
+    sched._omni_final_stage_flags = MethodType(OmniARScheduler._omni_final_stage_flags, sched)
+    sched._request_omits_kv_transfer_to_next_stage = MethodType(
+        OmniARScheduler._request_omits_kv_transfer_to_next_stage,
+        sched,
+    )
+    sched._request_omits_chunk_transfer_to_next_stage = MethodType(
+        OmniARScheduler._request_omits_chunk_transfer_to_next_stage,
+        sched,
+    )

tests/helpers/tests/test_assertions.py

@@ -2,7 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import io
+import json
 import wave
+from contextlib import nullcontext
 from types import SimpleNamespace
 
 import numpy as np
@@ -249,6 +251,85 @@ def test_pcm_transcript_escalation_uses_same_valid_wav_and_rejects_two_wrong_tra
         _assert_wav_wraps_pcm(wav_bytes, pcm_bytes, 24_000)
 
 
+@pytest.mark.parametrize("response_format", ["wav", "pcm"])
+@pytest.mark.parametrize("escalation_model", [None, "large-v3"])
+def test_failed_speech_capture_retains_exact_asr_input_and_original_failure(
+    monkeypatch, tmp_path, response_format, escalation_model
+):
+    pcm_bytes = _pcm_sine()
+    wav = io.BytesIO()
+    with wave.open(wav, "wb") as writer:
+        writer.setnchannels(1)
+        writer.setsampwidth(2)
+        writer.setframerate(24_000)
+        writer.writeframes(pcm_bytes)
+    audio_bytes = wav.getvalue() if response_format == "wav" else pcm_bytes
+    transcribed = []
+
+    def transcribe(audio_bytes, **kwargs):
+        transcribed.append(audio_bytes)
+        return "this is unrelated speech"
+
+    directory = tmp_path / "failed_audio"
+    monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory))
+    monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", transcribe)
+    config = {
+        "input": "please read this requested sentence accurately",
+        "response_format": response_format,
+        "transcript_pcm_sample_rate": 24_000,
+        "transcript_escalation_model": escalation_model,
+    }
+    response = OmniResponse(success=True, audio_bytes=audio_bytes, audio_format=f"audio/{response_format}")
+    with pytest.raises(AssertionError, match="Transcript doesn't match input") as error:
+        assert_audio_speech_response(response, config, "full_model")
+
+    clips = list(directory.glob("*.wav"))
+    assert len(clips) == 1
+    assert all(clip == clips[0].read_bytes() for clip in transcribed)
+    _assert_wav_wraps_pcm(clips[0].read_bytes(), pcm_bytes, 24_000)
+    metadata = json.loads(clips[0].with_suffix(".json").read_text())
+    assert metadata["expected_text"] == config["input"]
+    assert metadata["transcript"] == "this is unrelated speech"
+    assert metadata["assertion"] == str(error.value)
+
+
+@pytest.mark.parametrize("capture_enabled", [False, True])
+@pytest.mark.parametrize("content_failure", [False, True])
+def test_speech_capture_is_opt_in_and_does_not_write_successes(monkeypatch, tmp_path, capture_enabled, content_failure):
+    directory = tmp_path / "failed_audio"
+    if capture_enabled:
+        monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory))
+    else:
+        monkeypatch.delenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", raising=False)
+    expected = "please read this requested sentence accurately"
+    monkeypatch.setattr(
+        assertions,
+        "convert_audio_bytes_to_text",
+        lambda *_args, **_kwargs: "this is unrelated speech" if content_failure else expected,
+    )
+    with pytest.raises(AssertionError, match="Transcript doesn't match input") if content_failure else nullcontext():
+        assert_audio_speech_response(
+            OmniResponse(success=True, audio_bytes=b"audio", audio_format="audio/wav"),
+            {"input": expected, "response_format": "wav"},
+            "full_model",
+        )
+    assert directory.exists() is (capture_enabled and content_failure)
+    assert len(list(tmp_path.rglob("*.wav"))) == int(capture_enabled and content_failure)
+
+
+def test_speech_capture_io_failure_does_not_mask_content_failure(monkeypatch, tmp_path):
+    directory = tmp_path / "blocked"
+    directory.write_text("a file cannot be an output directory")
+    monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory))
+    monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", lambda *_args, **_kwargs: "this is unrelated speech")
+    with pytest.raises(AssertionError, match="Transcript doesn't match input"):
+        assert_audio_speech_response(
+            OmniResponse(success=True, audio_bytes=b"audio", audio_format="audio/wav"),
+            {"input": "please read this requested sentence accurately", "response_format": "wav"},
+            "full_model",
+        )
+
+
 @pytest.mark.parametrize(
     "pcm_bytes",
     [

tests/metrics/test_emit_calls.py

@@ -111,7 +111,7 @@ class TestFailureCounterWiring:
         obj.request_states["req-1"] = SimpleNamespace(
             metrics=SimpleNamespace(
                 e2e_done=set(),
-                build_and_log_summary=lambda: None,
+                log_timing_summary=lambda: None,
             ),
             failure_recorded=False,
         )
@@ -127,7 +127,7 @@ class TestFailureCounterWiring:
         obj.request_states["req-1"] = SimpleNamespace(
             metrics=SimpleNamespace(
                 e2e_done=set(),
-                build_and_log_summary=lambda: None,
+                log_timing_summary=lambda: None,
             ),
             failure_recorded=False,
         )

tests/metrics/test_stats.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import pytest
@@ -155,3 +158,39 @@ def test_build_and_log_summary_multiple_requests() -> None:
     r2_stage_entry = next(e for e in summary["stage_table"] if e["request_id"] == "r2")
     assert len(r1_stage_entry["stages"]) == 2
     assert len(r2_stage_entry["stages"]) == 1
+
+
+def test_log_timing_summary_skips_tables_unless_debug(monkeypatch) -> None:
+    from vllm_omni.metrics import stats
+
+    agg = OrchestratorAggregator(num_stages=1, log_stats=True, wall_start_ts=0.0, final_stage_id_for_e2e=0)
+    agg.stage_first_ts[0] = 0.0
+    agg.stage_last_ts[0] = 0.02
+    agg.on_stage_metrics(
+        0,
+        "r1",
+        StageRequestStats(
+            batch_id=1,
+            batch_size=1,
+            num_tokens_in=3,
+            num_tokens_out=3,
+            stage_gen_time_ms=20.0,
+            rx_transfer_bytes=0,
+            rx_decode_time_ms=0.0,
+            rx_in_flight_time_ms=0.0,
+            stage_stats=StageStats(),
+        ),
+    )
+    agg.on_finalize_request(0, "r1", req_start_ts=0.0)
+    tables, lines = [], []
+
+    def record_table(*args, **kwargs):
+        tables.append(args)
+        return ""
+
+    monkeypatch.setattr(stats, "_format_table", record_table)
+    monkeypatch.setattr(stats.logger, "info", lambda msg, *args: lines.append(msg % args))
+    monkeypatch.setattr(stats.logger, "isEnabledFor", lambda level: False)
+    agg.log_timing_summary()
+    assert not tables
+    assert len(lines) == 1 and lines[0].startswith("[OmniTiming] req=r1 total=")

tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_graph.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""HiFT graphs consume live inputs and return independently owned waveforms."""
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from tests.model_executor.models.cosyvoice3.test_cosyvoice3_incremental_hift import _make_hift
+
+pytestmark = [pytest.mark.core_model]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("finalize", [False, True])
+@torch.inference_mode()
+def test_hift_graph_preserves_current_and_previous_request(finalize):
+    hift = _make_hift().cuda().eval()
+    hift.remove_weight_norm()
+    hift.f0_predictor.remove_weight_norm()
+    hift.enable_decode_graphs()
+    generator = torch.Generator(device="cuda").manual_seed(17)
+    saved: list[tuple[torch.Tensor, torch.Tensor]] = []
+    for _ in range(5):
+        mel = torch.randn(1, 80, 30, device="cuda", generator=generator)
+        source = torch.randn(1, 1, 30 * 480, device="cuda", generator=generator)
+        expected = hift._decode_eager(mel, source, finalize=finalize)
+        actual = hift.decode(mel, source, finalize=finalize)
+        torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
+        for previous, snapshot in saved:
+            assert torch.equal(previous, snapshot)
+        if hift._decode_graphs.graphs:
+            saved.append((actual, actual.clone()))
+    assert len(hift._decode_graphs.graphs) == 1
+    hift._decode_graphs.max_graphs = 1
+    mel = torch.randn(1, 80, 31, device="cuda", generator=generator)
+    source = torch.randn(1, 1, 31 * 480, device="cuda", generator=generator)
+    actual = hift.decode(mel, source, finalize=finalize)
+    expected = hift._decode_eager(mel, source, finalize=finalize)
+    assert torch.equal(actual, expected)
+    assert len(hift._decode_graphs.graphs) == 1

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -928,4 +928,51 @@ def test_talker_output_contract(monkeypatch, mrv2, packed):
         model._sampling_eps = 1e-6
         model.mrv2_custom_sampler(SimpleNamespace(penalties_state=None))
     output = model.forward(torch.ones(2, dtype=torch.long), torch.arange(2), inputs_embeds=hidden)
-    assert output is (hidden if mrv2 or packed else sentinel)
+    assert output is hidden
+
+
+def test_request_conditioning_normalizes_singletons_without_changing_schema():
+    from vllm_omni.data_entry_keys import to_struct
+    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning
+
+    token = torch.tensor([[7, 8]], dtype=torch.int32)
+    feat = torch.randn(1, 4, 80)
+    speaker = torch.randn(1, 192)
+    length = torch.tensor([[2]], dtype=torch.int32)
+    raw = {"embed": {"speech_token": token, "speech_feat": feat, "embedding": speaker, "speech_token_len": [length]}}
+    result = to_struct(_normalize_request_conditioning(raw))
+    assert result.embed.speech_token is token
+    assert result.embed.speech_feat is feat
+    assert result.embed.embedding is speaker
+    assert result.embed.speech_token_len is length
+    assert raw["embed"]["speech_token_len"] == [length]
+
+
+def test_request_conditioning_rejects_unsplit_voices():
+    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning
+
+    with pytest.raises(ValueError, match="unsplit batch"):
+        _normalize_request_conditioning({"embed": {"embedding": [torch.ones(1, 192), torch.zeros(1, 192)]}})
+
+
+def test_request_conditioning_reaches_code2wav_with_true_prompt_length():
+    model = _make_code2wav_model()
+    output = model.forward(
+        torch.tensor([1, 2]),
+        torch.arange(2),
+        model_intermediate_buffer=[
+            {
+                "embed": {
+                    "speech_token": torch.tensor([[3, 2, 0, 0]], dtype=torch.int32),
+                    "speech_feat": torch.randn(1, 8, 80),
+                    "embedding": torch.randn(1, 192),
+                    "speech_token_len": [torch.tensor([[2]], dtype=torch.int32)],
+                },
+                "meta": {"req_id": ["voice"]},
+            }
+        ],
+    )
+    call = model.code2wav.forward_calls[0]
+    assert call["prompt_token"].shape == (1, 2)
+    assert call["prompt_feat"].shape == (1, 4, 80)
+    assert output.multimodal_outputs

tests/model_executor/models/moss_tts/test_codec_async_output.py

@@ -0,0 +1,338 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Waveform ownership across graph replays and asynchronous output copies."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
+
+pytestmark = [pytest.mark.core_model]
+
+
+class CausalCodec(nn.Module):
+    downsample_rate = 1
+
+    def __init__(self, channels):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(()))
+        self.channels = channels
+        self.config = SimpleNamespace(codebook_size=1024)
+
+    def initialize_decoder_state_pool(self, capacity, scratch):
+        self.state = torch.zeros(capacity + scratch, device=self.weight.device)
+
+    def reset_decoder_state_slots(self, slots):
+        self.state[slots] = 0
+
+    def decode_streaming_batch(self, codes, lengths, slots, valid_rows):
+        audio = codes.sum(0).float().cumsum(-1) + self.state[slots, None]
+        self.state[slots] = audio[:, -1]
+        channels = torch.arange(1, self.channels + 1, device=audio.device)
+        return SimpleNamespace(audio=audio[:, None] * channels[None, :, None], audio_lengths=lengths)
+
+
+class ReplayGraph:
+    """A real CUDA graph with deliberately reused output storage."""
+
+    batch_sizes = [1]
+
+    def __init__(self, codec):
+        self.codes = torch.ones((2, 1, 3), device="cuda", dtype=torch.long)
+        self.slots = torch.zeros(1, device="cuda", dtype=torch.long)
+        self.lengths = torch.full((1,), 3, device="cuda", dtype=torch.long)
+        self.valid = torch.ones(1, device="cuda", dtype=torch.bool)
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                codec.decode_streaming_batch(self.codes, self.lengths, self.slots, self.valid)
+        torch.cuda.current_stream().wait_stream(stream)
+        self.graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.graph):
+            self.output = codec.decode_streaming_batch(self.codes, self.lengths, self.slots, self.valid)
+        codec.state.zero_()
+
+    def decode(self, codes, slots, *, allow_frame_padding):
+        if codes.shape[1:] != (1, 3):
+            return None
+        self.codes.copy_(codes)
+        self.slots.copy_(slots)
+        self.graph.replay()
+        return self.output.audio, self.output.audio_lengths, 1
+
+
+def make_decoder(*, v2=True, channels=2, device="cpu", model_type="moss_tts_local", streaming=True):
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(n_vq=2, model_type=model_type),
+            async_chunk=streaming,
+            use_v2_model_runner=v2,
+        ),
+        scheduler_config=SimpleNamespace(max_num_seqs=8),
+    )
+    decoder = MossTTSCodecDecoder(vllm_config=config)
+    decoder._codec = CausalCodec(channels).to(device)
+    decoder._n_channels = channels
+    decoder._stream_max_step_frames = 3
+    return decoder
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("ramp,frames", [(None, [1, 15]), ([1, 4, 8, 15], [1, 4, 8, 15]), ([1, 4, 4, 15], [1, 4, 15])])
+def test_ramp_chunks_are_captured_at_exact_frame_counts(ramp, frames):
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(n_vq=2, model_type="moss_tts_local"),
+            async_chunk=True,
+            use_v2_model_runner=True,
+            stage_connector_config={
+                "extra": {
+                    "initial_codec_chunk_frames": 1,
+                    "codec_chunk_frames": 15,
+                    "codec_chunk_ramp": ramp,
+                }
+            },
+        ),
+        scheduler_config=SimpleNamespace(max_num_seqs=8),
+    )
+    decoder = MossTTSCodecDecoder(vllm_config=config)
+    assert decoder._streaming_graph_frame_sizes == frames
+    assert decoder._stream_max_step_frames == 15
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "v2,streaming,model_type",
+    [
+        (False, True, "moss_tts_local"),
+        (True, False, "moss_tts_local"),
+        (True, True, "moss_tts_realtime"),
+        (True, True, "moss_tts_local"),
+    ],
+)
+def test_gpu_output_requires_local_streaming_mrv2_and_cuda(v2, streaming, model_type):
+    decoder = make_decoder(v2=v2, streaming=streaming, model_type=model_type)
+    assert decoder._gpu_stream_output == (v2 and streaming and model_type == "moss_tts_local")
+    session = decoder._ensure_stream_session()
+    slot = session.acquire()
+    assert session.step({slot: torch.ones(2, 3)})[slot].device.type == "cpu"
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("v2", [False, True])
+@pytest.mark.parametrize("channels", [1, 2])
+@pytest.mark.parametrize("graph", [False, True])
+@torch.no_grad()
+def test_replays_terminal_tails_and_slot_reuse_preserve_pending_waveforms(v2, channels, graph):
+    from vllm_omni.outputs import OmniModelRunnerOutput
+    from vllm_omni.worker_v2.omni_generation_model_runner import (
+        OmniGenerationAsyncOutput,
+        OmniGenerationModelRunner,
+        _contains_cuda_tensor,
+    )
+
+    decoder = make_decoder(v2=v2, channels=channels, device="cuda")
+    # Real Local-v1.5 weights replace the class's 24 kHz default at load time.
+    decoder._sr_tensor = torch.tensor(48_000, dtype=torch.int32)
+    session = decoder._ensure_stream_session()
+    if graph:
+        session._cudagraph_wrapper = ReplayGraph(decoder._codec)
+    copy_stream = torch.cuda.Stream()
+    pending = []
+    # Do not materialize any output until later forwards have replayed the
+    # graph, reset terminal/cancelled state, and reused those state slots.
+    for round_index in range(4):
+        value = round_index + 1
+        codes = torch.cat(
+            [
+                torch.full((14,), value, device="cuda", dtype=torch.long),
+                torch.full((6,), value + 10, device="cuda", dtype=torch.long),
+            ]
+        )
+        info = [{"request_id": rid, "meta": {"finished": rid == "tail"}} for rid in ["tail", "empty", "live"]]
+        output = decoder.forward(codes, runtime_additional_information=info, seq_token_counts=[14, 0, 6])
+        # The runner may overwrite its input buffer before codec completion.
+        codes.fill_(-100)
+        payload = output.multimodal_outputs
+        assert _contains_cuda_tensor(payload) == v2
+        assert isinstance(payload, PackedOutputSnapshot) == v2
+        if v2:
+            # Exactly one waveform slab, regardless of request/chunk count.
+            assert sum(slab.is_cuda for _, slab in payload._groups) == 1
+            result = OmniModelRunnerOutput(
+                req_ids=["tail", "empty", "live"],
+                req_id_to_index={"tail": 0, "empty": 1, "live": 2},
+                sampled_token_ids=[[], [], []],
+            )
+            pending.append(
+                OmniGenerationAsyncOutput(
+                    model_runner_output=result,
+                    multimodal_outputs=payload,
+                    num_reqs=3,
+                    main_stream=torch.cuda.current_stream(),
+                    copy_stream=copy_stream,
+                )
+            )
+        else:
+            pending.append(OmniGenerationModelRunner._build_pooler_output(output, 3))
+        decoder.on_requests_finished(["live"])
+        assert decoder._stream_req_slots == {}
+        assert not session._leased_slots
+
+    for round_index, item in enumerate(pending):
+        rows = item.get_output().multimodal_outputs if v2 else item
+        for index, (frames, value) in enumerate([(7, round_index + 1), (0, 0), (3, round_index + 11)]):
+            expected = 2 * value * torch.arange(1, frames + 1, dtype=torch.float32)
+            if channels > 1:
+                expected = expected[None] * torch.arange(1, channels + 1)[:, None]
+            torch.testing.assert_close(rows[index]["model_outputs"], expected, rtol=0, atol=0)
+            assert rows[index]["sr"].item() == 48_000
+
+
+@pytest.mark.cuda
+@torch.no_grad()
+def test_owned_input_copy_clamps_entire_ragged_batch():
+    from vllm_omni.worker_v2.omni_ar_model_runner import _async_copy_mm
+
+    decoder = make_decoder(device="cuda", channels=2)
+    codes = torch.tensor([-1, 1024, 1, -10, 2, 1025, 5000, -2], device="cuda")
+    output = decoder.forward(
+        codes,
+        runtime_additional_information=[{"request_id": name} for name in ("a", "empty", "b")],
+        seq_token_counts=[6, 0, 2],
+    )
+    codes.fill_(-100)
+    cpu = _async_copy_mm(output.multimodal_outputs, 0, pin_memory=True)
+    torch.accelerator.synchronize()
+    expected = torch.tensor([[0.0, 1025.0, 2049.0], [0.0, 2050.0, 4098.0]])
+    torch.testing.assert_close(cpu["model_outputs"][0], expected, rtol=0, atol=0)
+    assert cpu["model_outputs"][1].numel() == 0
+    torch.testing.assert_close(cpu["model_outputs"][2], torch.tensor([[1023.0], [2046.0]]), rtol=0, atol=0)
+    decoder.on_requests_finished(["a", "b"])
+
+
+@pytest.mark.cuda
+@torch.no_grad()
+def test_pinned_metadata_reuse_waits_for_dma_and_preserves_mixed_terminals():
+    decoder = make_decoder(device="cuda")
+    session = decoder._ensure_stream_session()
+    # Force reuse every call, with an outstanding transfer on the first call.
+    # This delay tests host-buffer ownership, never performance or overlap.
+    session._metadata_ring = session._metadata_ring[:1]
+    slots = [session.acquire() for _ in range(3)]
+    pending = []
+    for step in range(8):
+        order = slots[step % 3 :] + slots[: step % 3]
+        terminal = {order[1]}
+        torch.cuda._sleep(2_000_000)
+        ids, tails = session._stage_slot_ids(order, terminal)
+        pending.append((ids.clone(), tails.clone(), order, list(terminal)))
+    session.close()
+    for ids, tails, expected_ids, expected_tails in pending:
+        assert ids.cpu().tolist() == expected_ids
+        assert tails.cpu().tolist() == expected_tails
+
+
+@pytest.mark.cuda
+@torch.no_grad()
+def test_mixed_terminal_reset_does_not_reset_live_slot():
+    decoder = make_decoder(device="cuda", channels=1)
+    session = decoder._ensure_stream_session()
+    first, second = session.acquire(), session.acquire()
+    codes = torch.ones(2, 3, device="cuda")
+    initial = session.step({first: codes, second: 2 * codes}, terminal_slots={first})
+    following = session.step({second: codes, first: codes})
+    torch.testing.assert_close(initial[first].cpu(), torch.tensor([[2.0, 4.0, 6.0]]), rtol=0, atol=0)
+    torch.testing.assert_close(following[first].cpu(), torch.tensor([[2.0, 4.0, 6.0]]), rtol=0, atol=0)
+    torch.testing.assert_close(following[second].cpu(), torch.tensor([[14.0, 16.0, 18.0]]), rtol=0, atol=0)
+    session.close()
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("explicit_copy_stream", [False, True])
+@torch.no_grad()
+def test_output_copy_waits_for_private_codec_stream(explicit_copy_stream):
+    from vllm_omni.worker_v2.omni_ar_model_runner import _async_copy_mm
+
+    decoder = make_decoder(device="cuda", channels=1)
+    decoder._codec_stream = torch.cuda.Stream()
+    with torch.cuda.stream(decoder._codec_stream):
+        torch.cuda._sleep(20_000_000)
+    codes = torch.ones(6, device="cuda", dtype=torch.long)
+    output = decoder.forward(
+        codes,
+        runtime_additional_information=[{"request_id": "live"}],
+        seq_token_counts=[6],
+    )
+    codes.fill_(-100)
+    payload = output.multimodal_outputs
+    assert isinstance(payload, PackedOutputSnapshot)
+    assert payload.producer_event is not None
+    # No explicit consumer wait here: the shared copier must honor readiness.
+    copy_stream = torch.cuda.Stream()
+    with torch.cuda.stream(copy_stream):
+        cpu = _async_copy_mm(payload, 0, copy_stream=copy_stream if explicit_copy_stream else None, pin_memory=True)
+    decoder.on_requests_finished(["live"])
+    copy_stream.synchronize()
+    torch.testing.assert_close(cpu["model_outputs"][0], torch.tensor([2.0, 4.0, 6.0]), rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("channels", [1, 2])
+@torch.no_grad()
+def test_short_terminal_row_crops_each_channel_into_output_slab(channels):
+    decoder = make_decoder(device="cuda", channels=channels)
+    session = decoder._ensure_stream_session()
+    short, live = session.acquire(), session.acquire()
+    codes = {short: torch.ones(2, 2, device="cuda"), live: 3 * torch.ones(2, 5, device="cuda")}
+    slab = torch.empty(channels * 7, device="cuda")
+    buffers = {
+        short: slab[: channels * 2].view(channels, 2),
+        live: slab[channels * 2 :].view(channels, 5),
+    }
+    result = session.step(codes, terminal_slots={short}, output_buffers=buffers)
+    for slot, frames, value in [(short, 2, 1), (live, 5, 3)]:
+        expected = 2 * value * torch.arange(1, frames + 1, dtype=torch.float32)[None]
+        expected = expected * torch.arange(1, channels + 1)[:, None]
+        assert result[slot].data_ptr() == buffers[slot].data_ptr()
+        torch.testing.assert_close(result[slot].cpu(), expected, rtol=0, atol=0)
+    session.close()
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("channels", [1, 2])
+@pytest.mark.parametrize("layout", ["adjacent", "strided", "gapped"])
+@torch.no_grad()
+def test_batched_output_copy_preserves_rows_and_slab_ownership(channels, layout):
+    decoder = make_decoder(device="cuda", channels=channels)
+    session = decoder._ensure_stream_session()
+    slots = [session.acquire() for _ in range(4)]
+    slab = torch.full((6 if layout == "gapped" else 4, channels, 6), -1.0, device="cuda")
+    if layout == "gapped":
+        buffers = {slot: slab[row] for row, slot in zip([0, 1, 3, 4], slots, strict=True)}
+        frames = 6
+    elif layout == "adjacent":
+        buffers = {slot: slab.reshape(4, channels * 6)[row].view(channels, 6) for row, slot in enumerate(slots)}
+        frames = 6
+    else:
+        # Interleaved with other requests' output: the short views must not
+        # overwrite the untouched part of another row/channel.
+        buffers = {slot: slab[row, :, :3] for row, slot in enumerate(slots)}
+        frames = 3
+    codes = {slot: torch.full((2, frames), row + 1, device="cuda") for row, slot in enumerate(slots)}
+    result = session.step(codes, output_buffers=buffers)
+    for row, slot in enumerate(slots):
+        expected = 2 * (row + 1) * torch.arange(1, frames + 1, dtype=torch.float32)[None]
+        expected = expected * torch.arange(1, channels + 1)[:, None]
+        assert result[slot].data_ptr() == buffers[slot].data_ptr()
+        torch.testing.assert_close(result[slot].cpu(), expected, rtol=0, atol=0)
+    if layout == "strided":
+        assert torch.all(slab[:, :, 3:] == -1)
+    elif layout == "gapped":
+        assert torch.all(slab[[2, 5]] == -1)
+    session.close()

tests/model_executor/models/moss_tts/test_codec_ring_headroom.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Chunked streaming decode must match whole-sequence decode.
+
+``RingKVCache.complete`` writes a whole chunk before attending. With a ring of
+exactly ``context`` entries, a chunk of T tokens evicts up to T keys that its
+own queries should still see; for T > context it evicts part of the chunk
+itself. ``initialize_decoder_state_pool(chunk_frames=...)`` adds one chunk of
+headroom per layer so streaming reproduces the non-streaming decoder.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import MossAudioTokenizerModel
+from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer_v2 import (
+    MossAudioTokenizerConfig,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+TRANSFORMER = {
+    "module_type": "Transformer",
+    "d_model": 16,
+    "num_heads": 2,
+    "num_layers": 1,
+    "dim_feedforward": 32,
+    "causal": True,
+    "norm": "layer_norm",
+    "positional_embedding": "rope",
+    "max_period": 10000,
+    "gating": "none",
+    "layer_scale": 1.0,
+    "conv_layout": True,
+}
+
+
+def tiny_codec() -> MossAudioTokenizerModel:
+    # 64 Hz mono input, 4-sample code frames (16 code frames per second).
+    # Decoder: transformer @16 Hz (context 0.5 s = 8 tokens) -> x2 -> transformer
+    # @32 Hz (context 0.25 s = 8 tokens) -> x2 -> 64 Hz. A 6-frame chunk is 12
+    # tokens at the second transformer, more than its 8-token ring.
+    config = MossAudioTokenizerConfig(
+        sampling_rate=64,
+        downsample_rate=4,
+        number_channels=1,
+        enable_channel_interleave=False,
+        encoder_kwargs=[
+            {"module_type": "PatchedPretransform", "patch_size": 4},
+            {**TRANSFORMER, "input_dimension": 4, "output_dimension": 8, "context_duration": 0.5},
+        ],
+        decoder_kwargs=[
+            {**TRANSFORMER, "input_dimension": 8, "output_dimension": 16, "context_duration": 0.5},
+            {"module_type": "PatchedPretransform", "patch_size": 2},
+            {**TRANSFORMER, "input_dimension": 8, "output_dimension": 2, "context_duration": 0.25},
+            {"module_type": "PatchedPretransform", "patch_size": 2},
+        ],
+        quantizer_kwargs={
+            "input_dim": 8,
+            "rvq_dim": 8,
+            "output_dim": 8,
+            "num_quantizers": 2,
+            "codebook_size": 16,
+            "codebook_dim": 4,
+            "quantizer_type": "rlfq",
+        },
+    )
+    torch.manual_seed(0)
+    model = MossAudioTokenizerModel(config).eval()
+    with torch.no_grad():
+        for parameter in model.parameters():
+            parameter.copy_(torch.randn_like(parameter) * 0.5)
+    model.quantizer.build_decode_lut(2, dtype=torch.float32)
+    return model
+
+
+def stream_decode(model, codes, chunk_frames, *, headroom):
+    n_q, _, total = codes.shape
+    model.initialize_decoder_state_pool(1, 0, chunk_frames=chunk_frames if headroom else 0)
+    try:
+        slot = torch.zeros(1, dtype=torch.long)
+        valid = torch.ones(1, dtype=torch.bool)
+        parts = []
+        for start in range(0, total, chunk_frames):
+            chunk = codes[:, :, start : start + chunk_frames]
+            lengths = torch.full((1,), chunk.shape[-1], dtype=torch.long)
+            parts.append(model.decode_streaming_batch(chunk, lengths, slot, valid).audio)
+        return torch.cat(parts, dim=-1)
+    finally:
+        model.close_decoder_state_pool()
+
+
+@torch.no_grad()
+def test_ring_headroom_sizes_each_layer_for_one_chunk():
+    model = tiny_codec()
+    transformers = [m for m in model.decoder if hasattr(m, "tokens_per_input_frame")]
+    assert [m.tokens_per_input_frame for m in transformers] == [1, 2]
+    assert [m.transformer.layers[0].self_attn.context for m in transformers] == [8, 8]
+    model.initialize_decoder_state_pool(1, 0, chunk_frames=6)
+    try:
+        rings = [m.transformer.layers[0].self_attn._streaming_state.kv_cache.capacity for m in transformers]
+        assert rings == [8 + 6, 8 + 12]
+    finally:
+        model.close_decoder_state_pool()
+    model.initialize_decoder_state_pool(1, 0)
+    try:
+        rings = [m.transformer.layers[0].self_attn._streaming_state.kv_cache.capacity for m in transformers]
+        assert rings == [8, 8]
+    finally:
+        model.close_decoder_state_pool()
+
+
+@torch.no_grad()
+def test_chunked_streaming_matches_whole_sequence_only_with_headroom():
+    model = tiny_codec()
+    torch.manual_seed(1)
+    codes = torch.randint(0, 16, (2, 1, 18))
+    reference = model._decode_frame(codes).audio
+    with_headroom = stream_decode(model, codes, 6, headroom=True)
+    legacy = stream_decode(model, codes, 6, headroom=False)
+    assert with_headroom.shape == reference.shape == legacy.shape
+    torch.testing.assert_close(with_headroom, reference, rtol=1e-5, atol=1e-5)
+    # The legacy context-sized ring lets the 12-token chunk evict keys inside
+    # its own window, so streaming diverges from whole-sequence decoding.
+    assert not torch.allclose(legacy, reference, rtol=1e-3, atol=1e-3)
+
+
+@torch.no_grad()
+def test_padded_terminal_tail_matches_exact_tail_with_headroom():
+    model = tiny_codec()
+    torch.manual_seed(2)
+    codes = torch.randint(0, 16, (2, 1, 14))
+    history, tail = codes[:, :, :12], codes[:, :, 12:]
+    padded_tail = torch.nn.functional.pad(tail, (0, 4))  # 2 real frames padded to a 6-frame step
+    outputs = {}
+    for name, last in [("exact", tail), ("padded", padded_tail)]:
+        model.initialize_decoder_state_pool(1, 0, chunk_frames=6)
+        try:
+            slot, valid = torch.zeros(1, dtype=torch.long), torch.ones(1, dtype=torch.bool)
+            for chunk in (history[:, :, :6], history[:, :, 6:]):
+                model.decode_streaming_batch(chunk, torch.tensor([6]), slot, valid)
+            audio = model.decode_streaming_batch(last, torch.tensor([last.shape[-1]]), slot, valid).audio
+            outputs[name] = audio[..., : 2 * model.downsample_rate]
+        finally:
+            model.close_decoder_state_pool()
+    torch.testing.assert_close(outputs["padded"], outputs["exact"], rtol=1e-5, atol=1e-5)

tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py

@@ -0,0 +1,547 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The first-chunk fast path emits audio out of band; the main path resumes the stream."""
+
+import queue
+import threading
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from tests.model_executor.models.moss_tts.test_streaming_terminal_batch import session as cpu_session
+from vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path import MossFirstChunkFastPath, _SlotHandoff
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder
+from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink
+
+pytestmark = pytest.mark.core_model
+
+N_VQ, SAMPLES, CHANNELS = 2, 4, 2
+
+
+class _Session:
+    def __init__(self, capacity):
+        self.free = list(range(capacity))
+        self.reset_waits = []
+
+    def acquire(self, *, oldest=False):
+        return self.free.pop(0 if oldest else -1) if self.free else None
+
+    def order_after_reset(self, slot, stream):
+        self.reset_waits.append(slot)
+
+    def release(self, slot, **kwargs):
+        self.free.append(slot)
+
+
+class _Wrapper:
+    batch_sizes = [1, 2, 4, 8]
+
+    def __init__(self, device):
+        self.state = torch.zeros(16, device=device)
+        self.calls = []
+        self.gate = threading.Event()
+        self.gate.set()
+
+    def decode(self, codes, slots):
+        self.gate.wait()
+        self.calls.append(int(codes.shape[1]))
+        base = codes.sum(0)[:, 0].float() + 100 * slots.float()
+        channel = torch.arange(CHANNELS, device=codes.device).float()
+        audio = base[:, None, None] + channel[None, :, None] + torch.zeros(1, 1, SAMPLES, device=codes.device)
+        self.state[slots] = 1
+        return audio.to(torch.bfloat16), None, int(codes.shape[1])
+
+
+@pytest.fixture
+def cuda():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device("cuda")
+
+
+def _fast(device, capacity=4):
+    session, wrapper = _Session(capacity), _Wrapper(device)
+    fast = MossFirstChunkFastPath(
+        session,
+        wrapper,
+        n_vq=N_VQ,
+        frames=1,
+        codebook_size=1024,
+        samples_per_frame=SAMPLES,
+        n_channels=CHANNELS,
+        sample_rate=torch.tensor(24000, dtype=torch.int32),
+        device=device,
+    )
+    return fast, session, wrapper
+
+
+class _CallbackSink:
+    def __init__(self, callback):
+        self.callback = callback
+
+    def prepare(self, request_ids):
+        callback = self.callback
+
+        class Delivery:
+            routes = {rid: (3 if rid == "r0" else 0) for rid in request_ids}
+
+            def __call__(self, ids, rows, sr):
+                for rid, row in zip(ids, rows, strict=True):
+                    callback(self.routes[rid], rid, {"model_outputs": row, "sr": sr})
+
+            def fail(self, ids):
+                pass
+
+        return Delivery()
+
+
+@pytest.mark.cuda
+def test_emits_each_first_chunk_and_hands_slot_to_main_path(cuda):
+    fast, session, wrapper = _fast(cuda)
+    emitted = {}
+    done = threading.Event()
+
+    def sink(client, req_id, payload):
+        emitted[req_id] = (client, payload["model_outputs"], int(payload["sr"]))
+        if len(emitted) == 2:
+            done.set()
+
+    slots: dict[str, int] = {}
+    assert not fast.submit("r0", "k0", [1, 2], slots)  # unbound: main path owns it
+    fast.bind(_CallbackSink(sink))
+    wrapper.gate.clear()
+    assert fast.submit("r0", "k0", torch.tensor([1, 2]), slots)
+    assert fast.submit("r1", "k1", [5, 7], slots)
+    assert not fast.submit("r0", "k0", [1, 2], slots)  # already has a slot
+    assert not fast.submit("r2", "k2", [1, 2, 3, 4], slots)  # not a 1-frame chunk
+    assert slots == {"k0": 0, "k1": 1} and session.free == [2, 3]
+    wrapper.gate.set()
+    assert done.wait(10)
+    fast.close()
+
+    assert wrapper.calls in ([1, 1], [2])
+    for req_id, client, code_sum, slot in (("r0", 3, 3, 0), ("r1", 0, 12, 1)):
+        got_client, wav, sr = emitted[req_id]
+        assert got_client == client and sr == 24000 and wav.device.type == "cpu"
+        expected = (code_sum + 100 * slot + torch.arange(CHANNELS).float())[:, None].expand(CHANNELS, SAMPLES)
+        torch.testing.assert_close(wav, expected)
+    assert sorted(session.reset_waits) == [0, 1]
+
+    # The main path consumes the marker once and is ordered after the decode.
+    assert fast.take_decoded("k0") and not fast.take_decoded("k0")
+    fast.order_after(0)
+    fast.order_after(0)  # no pending handoff: no-op
+    assert wrapper.state[:2].tolist() == [1, 1]
+
+
+@pytest.mark.cuda
+def test_order_after_blocks_until_decode_finishes(cuda):
+    fast, _, wrapper = _fast(cuda)
+    fast.bind(_CallbackSink(lambda *args: None))
+    wrapper.gate.clear()
+    assert fast.submit("r0", "k0", [1, 2], {})
+    finished = threading.Event()
+
+    def wait_for_slot():
+        fast.order_after(0)
+        finished.set()
+
+    waiter = threading.Thread(target=wait_for_slot)
+    waiter.start()
+    assert not finished.wait(0.2)
+    wrapper.gate.set()
+    assert finished.wait(10)
+    waiter.join()
+    fast.close()
+
+
+@pytest.mark.cuda
+def test_no_free_slot_leaves_chunk_to_main_path(cuda):
+    fast, _, _ = _fast(cuda, capacity=0)
+    fast.bind(_CallbackSink(lambda *args: None))
+    slots: dict[str, int] = {}
+    assert not fast.submit("r0", "k0", [1, 2], slots)
+    assert slots == {} and not fast.take_decoded("k0")
+    fast.close()
+
+
+class _FakeFast:
+    frames = 1
+
+    def __init__(self, decoded, unsent=None):
+        self.decoded = set(decoded)
+        self.unsent = dict(unsent or {})
+        self.ordered = []
+        self.forgotten = []
+
+    def take_decoded(self, key):
+        if key in self.decoded:
+            self.decoded.remove(key)
+            return True
+        return False
+
+    def get_request_slot(self, key, slots):
+        return slots.get(key)
+
+    def order_after(self, slot):
+        self.ordered.append(slot)
+
+    def take_unsent_audio(self, key):
+        return self.unsent.pop(key, None)
+
+    def forget(self, key):
+        self.forgotten.append(key)
+
+    def gate(self):
+        self.gated = getattr(self, "gated", 0) + 1
+
+
+def _decoder(s, fast):
+    decoder = object.__new__(MossTTSCodecDecoder)
+    nn.Module.__init__(decoder)
+    decoder._stream_max_step_frames = 15
+    decoder._stream_state_capacity = 8
+    decoder._stream_req_slots = {}
+    decoder._ensure_stream_session = lambda: s
+    decoder._first_chunk_fast_path = fast
+    decoder._stream_session = s
+    decoder._codec_stream = None
+    return decoder
+
+
+@pytest.mark.cpu
+def test_main_path_skips_fast_decoded_chunk_and_resumes_state():
+    s = cpu_session()
+    reference = cpu_session()
+    fast = _FakeFast({"a"})
+    decoder = _decoder(s, fast)
+    # The fast path leased slot 0 and decoded frame 0 of "a" into it.
+    slot = s.acquire(oldest=True)
+    decoder._stream_req_slots["a"] = slot
+    s.step({slot: torch.full((2, 1), 3.0)})
+    ref_slot = reference.acquire()
+    reference.step({ref_slot: torch.full((2, 1), 3.0)})
+
+    out = decoder._decode_streaming_batch([(0, "a", torch.full((2, 1), 3.0), False), (1, "b", torch.ones(2, 1), False)])
+    assert 0 not in out and out[1].tolist() == [[2.0]]
+    assert fast.ordered == [slot]
+
+    out = decoder._decode_streaming_batch([(0, "a", torch.ones(2, 3), True)])
+    expected = reference.step({ref_slot: torch.ones(2, 3)})[ref_slot]
+    torch.testing.assert_close(out[0], expected)
+    assert "a" not in decoder._stream_req_slots and fast.ordered == [slot, slot]
+
+
+@pytest.mark.cpu
+def test_abort_orders_release_after_fast_decode():
+    s = cpu_session()
+    fast = _FakeFast({"a"})
+    decoder = _decoder(s, fast)
+    slot = decoder._stream_req_slots["a"] = s.acquire(oldest=True)
+    decoder.on_requests_finished({"a"})
+    assert fast.ordered == [slot] and fast.forgotten == ["a"] and decoder._stream_req_slots == {}
+    assert slot in s._free_stream_slots
+
+
+@pytest.mark.cpu
+def test_main_path_returns_unsent_first_chunk_without_decoding_twice():
+    s = cpu_session()
+    first = torch.full((1, 1), 7.0)
+    fast = _FakeFast({"a"}, {"a": first})
+    decoder = _decoder(s, fast)
+    slot = s.acquire(oldest=True)
+    decoder._stream_req_slots["a"] = slot
+    s.step({slot: torch.full((2, 1), 3.0)})
+
+    out = decoder._decode_streaming_batch([(0, "a", torch.full((2, 1), 3.0), False)])
+    torch.testing.assert_close(out[0], first)
+    assert fast.unsent == {} and fast.ordered == [slot]
+
+    # A payload containing a later frame must concatenate it after the
+    # preserved first samples, while advancing state only for the later frame.
+    first = torch.full((1, 1), 7.0)
+    fast.decoded.add("b")
+    fast.unsent["b"] = first
+    slot_b = s.acquire(oldest=True)
+    decoder._stream_req_slots["b"] = slot_b
+    s.step({slot_b: torch.full((2, 1), 3.0)})
+    out = decoder._decode_streaming_batch([(0, "b", torch.full((2, 2), 3.0), False)])
+    torch.testing.assert_close(out[0][..., :1], first)
+    assert out[0].shape[-1] == 2
+
+
+@pytest.mark.cpu
+def test_oldest_acquire_reuses_least_recently_released_slot():
+    s = cpu_session()
+    a, b = s.acquire(), s.acquire()
+    s.release(a)
+    s.release(b)
+    assert s.acquire() == b
+    assert s.acquire(oldest=True) == 7  # never-used slots were freed before any release
+    assert s._free_stream_slots[-1] == a
+
+
+@pytest.mark.cuda
+def test_failed_decode_reports_error_without_replaying_state(cuda):
+    fast, _, wrapper = _fast(cuda)
+
+    def fail(codes, slots):
+        raise RuntimeError("boom")
+
+    wrapper.decode = fail
+    fast.bind(_CallbackSink(lambda *args: pytest.fail("no output expected")))
+    slots: dict[str, int] = {}
+    assert fast.submit("r0", "k0", [1, 2], slots)
+    with pytest.raises(RuntimeError, match="first-chunk decode failed"):
+        fast.order_after(slots["k0"])
+    assert fast.take_decoded("k0")
+    fast.close()
+
+
+@pytest.mark.cuda
+def test_failed_output_handoff_retains_decoded_samples(cuda):
+    fast, _, _ = _fast(cuda)
+
+    def reject(*args):
+        raise RuntimeError("output queue closed")
+
+    fast.bind(_CallbackSink(reject))
+    slots: dict[str, int] = {}
+    assert fast.submit("r0", "k0", [1, 2], slots)
+    fast.order_after(slots["k0"])
+    assert fast.take_decoded("k0")
+    wav = fast.take_unsent_audio("k0")
+    assert wav is not None and wav.shape == (CHANNELS, SAMPLES)
+    assert fast.take_unsent_audio("k0") is None
+    fast.close()
+
+
+@pytest.mark.cuda
+def test_gate_orders_main_stream_after_inflight_fast_decode(cuda):
+    fast, _, _ = _fast(cuda)
+    side = torch.cuda.Stream()
+    with torch.cuda.stream(side):
+        torch.cuda._sleep(200_000_000)
+        pending = torch.cuda.Event()
+        pending.record(side)
+    fast._inflight = pending
+    fast.gate()  # disabled: no wait
+    unordered = torch.cuda.Event()
+    unordered.record()
+    fast._gate_main = True
+    fast.gate()
+    ordered = torch.cuda.Event()
+    ordered.record()
+    unordered.synchronize()
+    assert not pending.query() and not ordered.query()
+    ordered.synchronize()
+    assert pending.query()
+
+
+def _cpu_fast(monkeypatch, *, scheduler=None):
+    # Replace only CUDA allocations, retaining the real lifecycle and worker.
+    import vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path as module
+
+    monkeypatch.setattr(torch.Tensor, "pin_memory", lambda tensor: tensor)
+    monkeypatch.setattr(torch.cuda, "Stream", lambda **kwargs: None)
+    monkeypatch.setattr(torch.cuda, "Event", lambda: None)
+    monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(set_device=lambda device: None))
+    session = _Session(2)
+    fast = MossFirstChunkFastPath(
+        session,
+        SimpleNamespace(batch_sizes=[1, 2]),
+        n_vq=2,
+        frames=1,
+        codebook_size=1024,
+        samples_per_frame=4,
+        n_channels=1,
+        sample_rate=torch.tensor(24000),
+        device=torch.device("cpu:0"),
+        handoff_timeout_s=0.01,
+    )
+    outputs: queue.Queue = queue.Queue()
+    scheduler = scheduler or SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)})
+    sink = engine_output_queue_sink(outputs, scheduler, upstream_first_audio=False)
+    return fast, session, sink, outputs
+
+
+def _cancel_during_admission(fast, session, scheduler, sink):
+    entered, release = threading.Event(), threading.Event()
+
+    class PausedSink:
+        def prepare(self, ids):
+            delivery = sink.prepare(ids)
+            if delivery.routes:
+                entered.set()
+                assert release.wait(5)
+            return delivery
+
+    fast.bind(PausedSink())
+    decoder = _decoder(session, fast)
+    errors = []
+    finished = threading.Event()
+    finishing = threading.Event()
+
+    def submit():
+        try:
+            assert fast.submit("r", "r", [1, 2], decoder._stream_req_slots)
+        except BaseException as error:
+            errors.append(error)
+
+    def cancel():
+        try:
+            finishing.set()
+            decoder.on_requests_finished({"r"})
+            finished.set()
+        except BaseException as error:
+            errors.append(error)
+
+    producer = threading.Thread(target=submit)
+    cleaner = threading.Thread(target=cancel)
+    producer.start()
+    try:
+        assert entered.wait(5)
+        assert decoder._stream_req_slots == {}
+        # Scheduler cancellation retires the route before the runner hook.
+        scheduler.requests.pop("r")
+        cleaner.start()
+        assert finishing.wait(5)
+        assert not finished.wait(0.1)
+        release.set()
+        producer.join(5)
+        cleaner.join(5)
+        assert not producer.is_alive() and not cleaner.is_alive()
+        assert not errors, errors
+        assert finished.is_set()
+        assert decoder._stream_req_slots == {}
+        assert fast._decoded == {} and fast._handoffs == {}
+        assert sorted(session.free) == [0, 1]
+        assert not fast.submit("r", "r", [1, 2], decoder._stream_req_slots)
+
+        # A fresh request can use and return the same capacity immediately.
+        scheduler.requests["reuse"] = SimpleNamespace(client_index=3)
+        release.set()
+        assert fast.submit("reuse", "reuse", [1, 2], decoder._stream_req_slots)
+        scheduler.requests.pop("reuse")
+        decoder.on_requests_finished({"reuse"})
+        assert decoder._stream_req_slots == {} and sorted(session.free) == [0, 1]
+    finally:
+        release.set()
+        producer.join(5)
+        if cleaner.ident is not None:
+            cleaner.join(5)
+        fast.close()
+
+
+@pytest.mark.cpu
+def test_cancel_waits_for_first_chunk_admission_and_reclaims_slot(monkeypatch):
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)})
+    fast, session, sink, _ = _cpu_fast(monkeypatch, scheduler=scheduler)
+
+    # Keep the actual queue/worker/handoff lifecycle, omitting CUDA decode.
+    def decode(jobs):
+        for job in jobs:
+            job.handoff.done.set()
+
+    monkeypatch.setattr(fast, "_decode", decode)
+    fast._handoff_timeout_s = 5.0
+    _cancel_during_admission(fast, session, scheduler, sink)
+
+
+@pytest.mark.cuda
+def test_cancel_during_first_chunk_admission_orders_cuda_handoff_and_reuse(cuda):
+    fast, session, _ = _fast(cuda, capacity=2)
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)})
+    sink = engine_output_queue_sink(queue.Queue(), scheduler, upstream_first_audio=False)
+    _cancel_during_admission(fast, session, scheduler, sink)
+
+
+@pytest.mark.cpu
+def test_slot_timeout_keeps_handoff_and_prevents_slot_reuse(monkeypatch):
+    fast, _, _, _ = _cpu_fast(monkeypatch)
+    handoff = _SlotHandoff()
+    fast._handoffs[0] = handoff
+    for _ in range(2):
+        with pytest.raises(TimeoutError, match="slot 0"):
+            fast.order_after(0)
+        assert fast._handoffs[0] is handoff
+    handoff.done.set()
+    fast.order_after(0)
+    assert fast._handoffs == {}
+
+
+@pytest.mark.cpu
+def test_closed_decoder_rejects_jobs_and_cannot_restart(monkeypatch):
+    fast, session, sink, _ = _cpu_fast(monkeypatch)
+    fast.bind(sink)
+    fast.close()
+    slots: dict[str, int] = {}
+    assert not fast.submit("r", "external", [1, 2], slots)
+    assert slots == {} and session.free == [0, 1]
+    with pytest.raises(RuntimeError, match="closed"):
+        fast.bind(sink)
+
+
+@pytest.mark.cpu
+def test_missing_route_does_not_claim_slot(monkeypatch):
+    fast, session, sink, outputs = _cpu_fast(monkeypatch, scheduler=SimpleNamespace(requests={}))
+    fast.bind(sink)
+    try:
+        assert not fast.submit("gone", "external", [1, 2], {})
+        assert session.free == [0, 1] and outputs.empty()
+    finally:
+        fast.close()
+
+
+@pytest.mark.cpu
+def test_decode_failure_wakes_current_and_queued_handoffs(monkeypatch):
+    from vllm.v1.engine import FinishReason
+
+    fast, _, sink, outputs = _cpu_fast(monkeypatch)
+    entered, release = threading.Event(), threading.Event()
+
+    def fail(jobs):
+        entered.set()
+        assert release.wait(2)
+        raise RuntimeError("partial replay failed")
+
+    fast._decode = fail
+    fast.bind(sink)
+    slots: dict[str, int] = {}
+    try:
+        assert fast.submit("r", "a", [1, 2], slots)
+        assert entered.wait(2)
+        assert fast.submit("r", "b", [3, 4], slots)
+        release.set()
+        fast._thread.join(2)
+        assert not fast._thread.is_alive()
+        for slot in slots.values():
+            with pytest.raises(RuntimeError, match="decode failed"):
+                fast.order_after(slot)
+        assert all(fast.take_decoded(key) for key in slots)
+        assert not fast.submit("r", "c", [1, 2], slots)
+        assert outputs.qsize() == 2
+        while not outputs.empty():
+            assert outputs.get_nowait()[1].outputs[0].finish_reason == FinishReason.ERROR
+    finally:
+        release.set()
+        fast.close()
+
+
+@pytest.mark.cpu
+def test_platform_failure_wakes_pending_handoff(monkeypatch):
+    import vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path as module
+
+    fast, _, _, _ = _cpu_fast(monkeypatch)
+    fast._handoffs[0] = _SlotHandoff()
+
+    def fail(device):
+        raise RuntimeError("device unavailable")
+
+    monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(set_device=fail))
+    fast._run()
+    assert fast._closed
+    with pytest.raises(RuntimeError, match="decode failed"):
+        fast.order_after(0)

tests/model_executor/models/moss_tts/test_local_code_rows_snapshot.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Per-request code rows keep the runner payload layout with one host copy."""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts.local_model_state import _CodeRowsSnapshot
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_rows_map_to_batch_positions_and_copy_once():
+    codes = torch.arange(6, dtype=torch.long).reshape(2, 3)
+    snapshot = _CodeRowsSnapshot(codes, [2, 0], 4)
+    assert isinstance(snapshot, PackedOutputSnapshot) and snapshot.producer_event is None
+    rows = snapshot["codes"]["audio"]
+    assert [tuple(row.shape) for row in rows] == [(1, 3), (0, 3), (1, 3), (0, 3)]
+    assert torch.equal(rows[2], codes[:1]) and torch.equal(rows[0], codes[1:])
+
+    copies = []
+
+    def copy(tensor):
+        copies.append(tensor)
+        return tensor.clone()
+
+    host = snapshot.copy_to_cpu(copy)["codes"]["audio"]
+    assert len(copies) == 1 and copies[0] is codes
+    codes.zero_()
+    assert host[2].tolist() == [[0, 1, 2]] and host[0].tolist() == [[3, 4, 5]]
+    assert host[1].numel() == 0 and host[1].dtype == torch.long

tests/model_executor/models/moss_tts/test_local_compiled_sampler.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.tts,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+
+def test_explicit_generator_bypasses_compiled_audio_sampler():
+    from transformers import GPT2Config
+
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+        MossTTSLocalDepthTransformer,
+    )
+
+    model = MossTTSLocalDepthTransformer(GPT2Config(n_embd=80, n_head=1, n_inner=160)).cuda().bfloat16().eval()
+    heads = torch.nn.ModuleList([torch.nn.Linear(80, 32) for _ in range(2)]).cuda().bfloat16()
+    embeddings = torch.nn.ModuleList([torch.nn.Embedding(32, 80) for _ in range(2)]).cuda().bfloat16()
+    binary = torch.nn.Linear(80, 2).cuda().bfloat16()
+    hidden = torch.randn(1, 80, device="cuda", dtype=torch.bfloat16)
+    generator = torch.Generator(device="cuda").manual_seed(17)
+    expected = model.generate_frame(hidden, heads, embeddings, binary, n_vq=2, generator=generator)
+
+    def forbidden(*args, **kwargs):
+        raise AssertionError("compiled path selected")
+
+    model._compiled_audio_sampler = forbidden
+    generator.manual_seed(17)
+    actual = model.generate_frame(hidden, heads, embeddings, binary, n_vq=2, generator=generator)
+    assert all(torch.equal(a, b) for a, b in zip(expected, actual))
+    with pytest.raises(AssertionError, match="compiled path selected"):
+        model.generate_frame(hidden, heads, embeddings, binary, n_vq=2)

tests/model_executor/models/moss_tts/test_local_determined_tokens.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Determined text tokens retain the upstream sampling/counting contract."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from vllm.sampling_params import SamplingParams
+from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.states import RequestState
+
+from tests.model_executor.models.moss_tts.test_local_model_state import _batch, _state
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner
+from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("seed", [None, 17])
+@pytest.mark.parametrize("temperature", [0.0, 1.7])
+def test_matches_real_sampler_for_mixed_prefill_decode_stop_and_reordering(seed, temperature):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    device = torch.device("cuda")
+    reqs = RequestState(5, 32, 16, 0, 8, device)
+    sampler = Sampler(5, 8, device, reqs)
+    for i, length in enumerate([4, 1, 3]):
+        reqs.add_request(str(i), length, [2] * length, 0, 8)
+        slot = reqs.req_id_to_index[str(i)]
+        sampler.add_request(slot, length, SamplingParams(temperature=temperature, top_k=3, top_p=0.8, seed=seed))
+    reqs.apply_staged_writes()
+    sampler.apply_staged_writes()
+    state = _state(MossLocalModelState, device)
+    state._direct_tokens = True
+    state.model._batch_state = None
+    batch = _batch(device, [2, 4, 3], [1, 2, 1])
+    batch.seq_lens = torch.tensor([5, 2, 2], device=device, dtype=torch.int32)
+    batch.cu_num_logits_np = np.arange(4, dtype=np.int32)
+    batch.cu_num_logits = torch.tensor(batch.cu_num_logits_np, device=device)
+    batch.expanded_idx_mapping = batch.idx_mapping
+    batch.expanded_local_pos = torch.zeros(3, dtype=torch.int32, device=device)
+    batch.logits_indices = batch.query_start_loc[1:].long() - 1
+    batch.positions = torch.arange(batch.num_tokens, device=device)
+    state.model._batch_should_continue = torch.tensor([False, True, True], device=device)
+    with torch.inference_mode():
+        logits = state.model.compute_logits(torch.ones(3, 4, device=device))
+        expected = sampler(logits, batch)
+        actual = state.sample_determined_tokens(batch, sampler)
+    assert actual is not None
+    for field in ["sampled_token_ids", "num_sampled", "num_rejected"]:
+        torch.testing.assert_close(getattr(actual, field), getattr(expected, field), rtol=0, atol=0)
+    assert actual.num_sampled.tolist() == [1, 0, 1]
+    snapshot = actual.sampled_token_ids.clone()
+    state.model._batch_should_continue.logical_not_()
+    torch.testing.assert_close(actual.sampled_token_ids, snapshot, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize(
+    "constraint",
+    [
+        {"min_tokens": 2, "stop_token_ids": [1]},
+        {"allowed_token_ids": [1, 2]},
+        {"logit_bias": {1: -10.0}},
+    ],
+)
+def test_real_sampler_constraints_disable_determined_tokens(constraint):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    device = torch.device("cuda")
+    reqs = RequestState(5, 32, 16, 0, 8, device)
+    sampler = Sampler(5, 8, device, reqs)
+    reqs.add_request("constrained", 1, [2], 0, 8)
+    slot = reqs.req_id_to_index["constrained"]
+    sampler.add_request(slot, 1, SamplingParams(**constraint))
+    state = _state(MossLocalModelState, device)
+    state._direct_tokens = True
+    # Use the actual upstream request registration, including its combined
+    # min-token/allowed-ID/bias flag, rather than mocking that flag ourselves.
+    assert state.sample_determined_tokens(SimpleNamespace(idx_mapping_np=np.array([slot])), sampler) is None
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("feature", ["logprobs", "bias", "penalty", "bad_words", "thinking", "trace", "mask", "nans"])
+def test_distribution_features_fall_back_before_using_gpu(feature):
+    sampler = object.__new__(Sampler)
+    sampler.compute_nans = feature == "nans"
+    sampler.return_sampling_mask = feature == "mask"
+    sampler.trace_replay_state = object() if feature == "trace" else None
+    sampler.get_logprobs_dims = lambda rows: (1, 0) if feature == "logprobs" else None
+    sampler.logit_bias_state = SimpleNamespace(use_logit_bias=np.array([feature == "bias"]))
+    sampler.penalties_state = SimpleNamespace(use_penalty=np.array([feature == "penalty"]))
+    sampler.bad_words_state = SimpleNamespace(num_bad_words=SimpleNamespace(np=np.array([feature == "bad_words"])))
+    sampler.thinking_budget_state = SimpleNamespace(enabled=True, use_thinking_budget=np.array([feature == "thinking"]))
+    state = object.__new__(MossLocalModelState)
+    state._direct_tokens = True
+    assert state.sample_determined_tokens(SimpleNamespace(idx_mapping_np=np.array([0])), sampler) is None
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("feature", ["grammar", "sharding", "pipeline", "speculative", "empty", "declined"])
+def test_runner_retains_upstream_path_for_unsupported_execution(mocker, feature):
+    hook = mocker.Mock(return_value=None)
+
+    class State:
+        def sample_determined_tokens(self, batch, sampler):
+            return hook(batch, sampler)
+
+    runner = object.__new__(OmniARModelRunner)
+    runner.model_state = State()
+    runner.sampler = object()
+    runner.batch_sharder = object() if feature == "sharding" else None
+    runner.pp_handler = object() if feature == "pipeline" else None
+    batch = SimpleNamespace(num_draft_tokens=int(feature == "speculative"), num_reqs=0 if feature == "empty" else 1)
+    grammar = object() if feature == "grammar" else None
+    fallback = mocker.patch.object(OmniGPUModelRunner, "sample", return_value="upstream")
+    hidden = object()
+    assert runner.sample(hidden, batch, grammar) == "upstream"
+    fallback.assert_called_once_with(hidden, batch, grammar)
+    assert hook.call_count == int(feature == "declined")

tests/model_executor/models/moss_tts/test_local_eager_mtp.py

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Eager MTP emits the same frames one step earlier with identical backbone inputs."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from tests.model_executor.models.moss_tts.test_local_model_state import _admit, _batch, _state, _step
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)])
+def device(request):
+    if request.param == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device(request.param)
+
+
+def _run(state, device, schedule, stop_step):
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[3], prompt[0] = 4, 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    embeds = []
+    frames: dict[int, list[torch.Tensor]] = {3: [], 0: []}
+    for index, (slots, counts) in enumerate(schedule):
+        state.model.force_stop = index == stop_step
+        embed, payload, _ = _step(state, _batch(device, slots, counts), req_states)
+        embeds.append(embed)
+        rows = payload.get("codes", {}).get("audio", [])
+        for slot, row in zip(slots, rows):
+            if row.numel() and bool(row.ne(state.model.audio_pad_token_id).any()):
+                frames[slot].append((index, row.clone()))
+        for slot, count in zip(slots, counts):
+            computed[slot] += count
+    return embeds, frames
+
+
+def test_eager_frames_match_canonical_one_step_earlier(device):
+    schedule = [([3, 0], [2, 2]), ([3, 0], [2, 1]), ([3, 0], [1, 1]), ([0, 3], [1, 1]), ([3, 0], [1, 1])]
+    canonical, eager = (_state(MossLocalModelState, device) for _ in range(2))
+    eager._local_eager_mtp = True
+    for state in (canonical, eager):
+        _admit(state, 3, "long", 17)
+        _admit(state, 0, "short", 17)
+    # Stop both streams at the canonical step that would draw the 4th step's frames.
+    c_embeds, c_frames = _run(canonical, device, schedule, stop_step=4)
+    e_embeds, e_frames = _run(eager, device, schedule, stop_step=3)
+
+    for c, e in zip(c_embeds, e_embeds):
+        torch.testing.assert_close(c, e, rtol=0, atol=0)
+    for slot in (3, 0):
+        assert [f for _, f in c_frames[slot]] and len(c_frames[slot]) == len(e_frames[slot])
+        for (ci, cf), (ei, ef) in zip(c_frames[slot], e_frames[slot]):
+            assert ei == ci - 1
+            torch.testing.assert_close(cf, ef, rtol=0, atol=0)
+
+
+def test_stopped_stream_does_not_emit_again(device):
+    state = _state(MossLocalModelState, device)
+    state._local_eager_mtp = True
+    _admit(state, 0, "short", 17)
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[0] = 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    emitted = []
+    for index, count in enumerate([2, 1, 1, 1]):
+        state.model.force_stop = index == 1
+        _, payload, _ = _step(state, _batch(device, [0], [count]), req_states)
+        row = payload["codes"]["audio"][0]
+        emitted.append(bool(row.ne(state.model.audio_pad_token_id).any()))
+        computed[0] += count
+        state.model.force_stop = False
+    # Frame at the prefill step, stop at step 1, then asynchronous overrun steps.
+    assert emitted == [True, False, False, False]
+
+
+def test_unmarked_decode_row_ignores_stale_slot_state(device):
+    canonical, eager = (_state(MossLocalModelState, device) for _ in range(2))
+    eager._local_eager_mtp = True
+    for state in (canonical, eager):
+        _admit(state, 0, "short", 17)
+    # A decode row whose completing prefill was not observed in this state.
+    eager._keep_pool[0] = False
+    eager._eager_emb_pool[0].fill_(3.0)
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[0], computed[0] = 2, 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    for state in (canonical, eager):
+        state.intermediate_buffer.buffers[0]["audio_state"] = {"is_stopping": False}
+    c_embed, c_payload, _ = _step(canonical, _batch(device, [0], [1]), req_states)
+    e_embed, e_payload, _ = _step(eager, _batch(device, [0], [1]), req_states)
+    torch.testing.assert_close(c_embed, e_embed, rtol=0, atol=0)
+    torch.testing.assert_close(c_payload["codes"]["audio"][0], e_payload["codes"]["audio"][0], rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_first_prefill_through_runner_preserves_parent_eager_contract(monkeypatch, mocker):
+    """The Local config must not enable the parent's incompatible eager hook."""
+    from vllm_omni.worker_v2 import omni_ar_model_runner as runner_module
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+    from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+    from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner
+
+    device = torch.device("cpu")
+    template = _state(MossLocalModelState, device)
+    model = template.model
+    model.config = SimpleNamespace(mrv2_eager_mtp=True)
+
+    def initialize_parent(state, *args):
+        state.__dict__.update(template.__dict__)
+        state._eager_mtp, state._eager_rows = False, None
+        state._eager_state = EagerMTPState(state)
+
+    monkeypatch.setattr(OmniModelState, "__init__", initialize_parent)
+    state = MossLocalModelState(SimpleNamespace(), model, None, device)
+    _admit(state, 0, "req", 17)
+    batch = _batch(device, [0], [2])
+    batch.req_ids = ["req"]
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[0] = 2
+    _step(state, batch, SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed))
+    hook = mocker.spy(state, "run_eager_mtp")
+
+    runner = OmniARModelRunner.__new__(OmniARModelRunner)
+    runner.execute_model_state = SimpleNamespace(
+        input_batch=batch,
+        hidden_states=torch.zeros(2, 4),
+        finished_req_ids=set(),
+        ec_connector_output=None,
+        routed_experts=None,
+    )
+    runner._kv_extracted_req_ids = runner._last_aux_output = runner._last_multimodal_outputs = None
+    runner.is_last_pp_rank, runner.pp_handler, runner.check_ep_fault = True, None, False
+    runner.model_config = SimpleNamespace(async_chunk=False)
+    runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text"))
+    runner.model_state, runner.model = state, model
+    runner.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(gpu=torch.zeros(5, 2)),
+        num_computed_tokens=SimpleNamespace(gpu=torch.from_numpy(computed)),
+        prompt_len=SimpleNamespace(np=prompt, gpu=torch.from_numpy(prompt)),
+    )
+    runner.main_stream = runner.output_copy_stream = mocker.Mock()
+    runner.eplb = runner._finalize_native_data_plane_output = runner._reserve_native_data_plane_outputs = mocker.Mock()
+    runner.sampler = None
+    runner.sample = mocker.Mock(
+        return_value=(
+            SimpleNamespace(sampled_token_ids=torch.tensor([[2]])),
+            torch.ones(1),
+            torch.zeros(1),
+        )
+    )
+    runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=mocker.Mock(return_value={}))
+    runner.postprocess_sampled = mocker.Mock()
+    runner.kv_connector = SimpleNamespace(post_forward=mocker.Mock(return_value=None))
+    mock_out = SimpleNamespace(copy_event=None)
+    monkeypatch.setattr(runner_module, "OmniAsyncOutput", mocker.Mock(return_value=mock_out))
+
+    assert runner.sample_tokens(None) is mock_out
+    hook.assert_called_once()
+    assert state._eager_mtp is False and state._eager_rows is None
+    assert state._local_eager_mtp is True

tests/model_executor/models/moss_tts/test_local_model_state.py

@@ -0,0 +1,286 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare slot lifecycle with canonical Local hooks, including output ownership."""
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+import torch
+from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration
+from vllm_omni.worker_v2.model_states import init_omni_model_state
+from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer
+from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)])
+def device(request):
+    if request.param == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device(request.param)
+
+
+def _model(device):
+    model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+    torch.nn.Module.__init__(model)
+    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
+    model.model = torch.nn.Module()
+    weight = torch.arange(32, dtype=dtype, device=device).reshape(8, 4) / 16
+    model.model.embed_tokens = torch.nn.Embedding(8, 4, _weight=weight)
+    model.hidden_size, model.n_vq = 4, 2
+    model.audio_pad_token_id, model.text_vocab_size = 8, 8
+    model.audio_assistant_slot_token_id, model.im_end_token_id = 2, 3
+    model.talker_mtp_output_key = ("audio_codes", "current")
+    model.talker_mtp_graph_safe = False
+    model.gpu_resident_buffer_keys = {("hidden_states", "last"), ("audio_codes", "current")}
+    model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None
+    model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).to(dtype) / 8
+    model.force_stop = False
+
+    def frame(hidden, *_args, generator=None, temperature=None, top_k=None, top_p=None, **_kwargs):
+        assert (temperature, top_k, top_p) == (1.7, 25, 0.8)
+        noise = torch.randint(0, 8, (hidden.shape[0], 2), device=device, generator=generator)
+        codes = (hidden[:, :2].mul(16).long() + noise) % 8
+        return torch.full((hidden.shape[0],), not model.force_stop, device=device), codes
+
+    model.local_transformer = SimpleNamespace(generate_frame=frame)
+    return model
+
+
+def _state(cls, device):
+    state = object.__new__(cls)
+    state.model = _model(device)
+    state.device, state.dtype = device, state.model.model.embed_tokens.weight.dtype
+    state.has_preprocess = state.has_postprocess = state.have_multimodal_outputs = True
+    state.scheduler_config = SimpleNamespace(max_num_seqs=5)
+    state.vllm_config = SimpleNamespace()
+    state.intermediate_buffer = OmniIntermediateBuffer(5)
+    state._static_inputs_embeds = torch.zeros(16, 4, device=device, dtype=state.dtype)
+    state._mtp_input_ids = torch.zeros(5, device=device, dtype=torch.long)
+    state._mtp_offsets = torch.zeros(5, device=device, dtype=torch.long)
+    for name in ("_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step"):
+        setattr(state, name, torch.zeros(5, 4, device=device, dtype=state.dtype))
+    state._mtp_runner = state._mtp_sample_uniforms = None
+    state._mtp_generators = {}
+    state._stream_pos = {}
+    state._stream_decode_event = None
+    state._eager_state = EagerMTPState(state)
+    if cls is MossLocalModelState:
+        state._init_slot_buffers(5, 4, device, state.dtype)
+    return state
+
+
+def _admit(state, slot, name, seed):
+    params = SimpleNamespace(extra_args={"tts_local_seed": seed} if seed is not None else {}, seed=None)
+    with patch.object(DefaultModelState, "add_request", return_value=None):
+        state.add_request(slot, SimpleNamespace(req_id=name, mm_features=[], sampling_params=params))
+    state.intermediate_buffer.buffers[slot]["codes"] = {"ref": torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])}
+
+
+def _batch(device, slots, counts):
+    starts = np.array([0, *np.cumsum(counts)], dtype=np.int32)
+    return SimpleNamespace(
+        idx_mapping_np=np.array(slots),
+        idx_mapping=torch.tensor(slots, device=device),
+        num_reqs=len(slots),
+        num_tokens=int(starts[-1]),
+        num_scheduled_tokens=np.array(counts),
+        query_start_loc_np=starts,
+        query_start_loc=torch.tensor(starts, device=device),
+        # Upstream MRV2 uses int32 token buffers, unlike the MTP graph inputs.
+        input_ids=torch.full((int(starts[-1]),), 2, dtype=torch.int32, device=device),
+    )
+
+
+def _step(state, batch, req_states, dispatcher=None):
+    inputs = {"input_ids": batch.input_ids, "inputs_embeds": state._static_inputs_embeds[: batch.num_tokens]}
+    with (
+        patch("vllm.forward_context.set_forward_context", return_value=nullcontext()),
+        patch(
+            "vllm_omni.model_executor.models.moss_tts.local_model_state.set_forward_context", return_value=nullcontext()
+        ),
+        torch.inference_mode(),
+    ):
+        state.run_preprocess(batch, inputs, req_states, dispatcher)
+        hidden = inputs["inputs_embeds"] + 0.25
+        state.run_postprocess(hidden, batch)
+        _, payload = state.postprocess_model_output(hidden, batch, req_states)
+        last = batch.query_start_loc[1:] - 1
+        tokens = state.model.compute_logits(hidden.index_select(0, last)).argmax(-1)
+        result = inputs["inputs_embeds"].clone(), payload, tokens
+        hidden.fill_(-100)  # next graph replay can overwrite its output immediately
+    return result
+
+
+@pytest.mark.parametrize("seed", [None, 17])
+@pytest.mark.parametrize("batch_prefill", [False, True])
+def test_slot_matches_canonical_mixed_prefill_reorder_stop_and_reuse(device, seed, batch_prefill):
+    reference, candidate = (_state(cls, device) for cls in (OmniModelState, MossLocalModelState))
+    candidate._batch_prefill = batch_prefill
+    for state in (reference, candidate):
+        _admit(state, 3, "long", seed)
+        _admit(state, 0, "short", seed)
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[3], prompt[0] = 4, 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    saved = []
+    for index, (slots, counts) in enumerate(
+        [
+            ([3, 0], [2, 2]),
+            ([0, 3], [1, 1]),
+            ([3, 0], [1, 1]),
+            ([3, 0], [1, 1]),
+            ([0, 3], [1, 1]),
+            ([3, 0], [1, 1]),
+        ]
+    ):
+        if index == 4:
+            # Cancellation/removal followed by reuse must not retain hidden,
+            # reference offset, control, or the old request's RNG stream.
+            for state in (reference, candidate):
+                state.remove_request("short")
+                assert "short" not in state._mtp_generators
+                _admit(state, 0, "replacement", seed)
+            computed[0], prompt[0] = 0, 1
+        outputs = []
+        for state in (reference, candidate):
+            state.model.force_stop = index == 3
+            torch.manual_seed(91 + index)
+            outputs.append(_step(state, _batch(device, slots, counts), req_states))
+        a, b = outputs
+        torch.testing.assert_close(a[0], b[0], rtol=0, atol=0)
+        torch.testing.assert_close(a[2], b[2], rtol=0, atol=0)
+        if index == 3:
+            assert b[2].tolist() == [3, 3]
+        assert a[1].keys() == b[1].keys()
+        for ar, br in zip(a[1].get("codes", {}).get("audio", []), b[1].get("codes", {}).get("audio", [])):
+            torch.testing.assert_close(ar, br, rtol=0, atol=0)
+            saved.append((br, br.clone()))
+        for slot, count in zip(slots, counts):
+            computed[slot] += count
+            old = reference.intermediate_buffer.buffers[slot]
+            new = candidate.intermediate_buffer.buffers[slot]
+            assert old["ref_offset"] == new["ref_offset"]
+            torch.testing.assert_close(candidate._hidden_pool[slot], old["hidden_states"]["last"], rtol=0, atol=0)
+            assert "hidden_states" not in new and "audio_codes" not in new
+    candidate._codes_pool.fill_(8)
+    candidate._hidden_pool.zero_()
+    for current, snapshot in saved:
+        torch.testing.assert_close(current, snapshot, rtol=0, atol=0)
+
+
+def test_slot_honors_external_stopping_control(device):
+    states = [_state(cls, device) for cls in (OmniModelState, MossLocalModelState)]
+    for state in states:
+        _admit(state, 0, "stop", None)
+        req = SimpleNamespace(prompt_len=np.ones(5), num_computed_tokens=np.zeros(5))
+        _step(state, _batch(device, [0], [1]), req)
+        state.intermediate_buffer.buffers[0]["audio_state"]["is_stopping"] = True
+        req.num_computed_tokens[:] = 1
+        output = _step(state, _batch(device, [0], [1]), req)
+        assert output[2].tolist() == [3]
+        assert output[1]["codes"]["audio"][0].tolist() == [[8, 8]]
+
+
+@pytest.mark.cuda
+def test_slot_mtp_graph_padding_reorder_and_owned_output():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    states = [_state(cls, torch.device("cuda")) for cls in (OmniModelState, MossLocalModelState)]
+    for state in states:
+
+        def frame(hidden, *_args, **_kwargs):
+            return hidden[:, 0] < 3, hidden[:, :2].mul(16).long() % 8
+
+        state.model.local_transformer.generate_frame = frame
+        for slot in [0, 3, 4]:
+            _admit(state, slot, f"req{slot}", None)
+    candidate = states[1]
+    buffers = [
+        candidate._mtp_input_ids[:4],
+        candidate._mtp_input_embeds[:4],
+        candidate._mtp_hidden[:4],
+        candidate._mtp_text_step[:4],
+    ]
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream), torch.inference_mode():
+        for _ in range(3):
+            candidate.model.mtp(*buffers)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph), torch.inference_mode():
+        graph_output = candidate.model.mtp(*buffers)
+
+    def replay(*args, **kwargs):
+        assert all(a.data_ptr() == b.data_ptr() for a, b in zip(args, buffers))
+        assert args[0].shape[0] == 4
+        graph.replay()
+        return graph_output
+
+    candidate._mtp_runner = replay
+    candidate._is_mtp_graph_runner = lambda: True
+    from vllm.config import CUDAGraphMode
+
+    def dispatcher(bsz):
+        return SimpleNamespace(num_tokens=4, cg_mode=CUDAGraphMode.FULL)
+
+    req = SimpleNamespace(prompt_len=np.ones(5, dtype=np.int32), num_computed_tokens=np.zeros(5, dtype=np.int32))
+    retained = []
+    for slots in [[0, 3, 4], [3, 0], [4, 0, 3], [3, 4]]:
+        out = [
+            _step(
+                state,
+                _batch(torch.device("cuda"), slots, [1] * len(slots)),
+                req,
+                dispatcher if state is candidate else None,
+            )
+            for state in states
+        ]
+        torch.testing.assert_close(out[0][0], out[1][0], rtol=0, atol=0)
+        torch.testing.assert_close(out[0][2], out[1][2], rtol=0, atol=0)
+        for a, b in zip(out[0][1].get("codes", {}).get("audio", []), out[1][1].get("codes", {}).get("audio", [])):
+            torch.testing.assert_close(a, b, rtol=0, atol=0)
+            retained.append((b, b.clone()))
+        req.num_computed_tokens[slots] += 1
+    with torch.inference_mode():
+        for tensor in graph_output:
+            tensor.fill_(7)
+    for tensor, copy in retained:
+        torch.testing.assert_close(tensor, copy, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_factory_requires_explicit_class_capability_and_valid_state(monkeypatch):
+    monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None)
+    dynamic = MagicMock(has_preprocess=True, create_omni_model_state=None)
+    assert isinstance(init_omni_model_state(None, dynamic, None, torch.device("cpu")), OmniModelState)
+    dynamic.create_mrv2_model_state.assert_not_called()
+
+    class Invalid:
+        has_preprocess = True
+
+        def create_mrv2_model_state(self, *args):
+            return object()
+
+    with pytest.raises(TypeError, match="must return a ModelState"):
+        init_omni_model_state(None, Invalid(), None, torch.device("cpu"))
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("enabled", [False, True])
+def test_local_factory_default_and_option(monkeypatch, enabled):
+    monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None)
+    monkeypatch.setattr(MossLocalModelState, "__init__", lambda *args: None)
+    model = _model(torch.device("cpu"))
+    model.config = SimpleNamespace(mrv2_gpu_slot_state=enabled)
+    state = init_omni_model_state(None, model, None, torch.device("cpu"))
+    assert type(state) is (MossLocalModelState if enabled else OmniModelState)

tests/model_executor/models/moss_tts/test_local_prefill_batch.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Reference batching preserves coordinates and asynchronous source ownership."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from tests.model_executor.models.moss_tts.test_local_model_state import _batch, _state
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("flat", [False, True])
+def test_prefill_reuses_text_embeddings_and_noncontiguous_reference_offsets(mocker, flat):
+    state = _state(MossLocalModelState, torch.device("cpu"))
+    state._batch_prefill = True
+    codes = torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])
+    for slot, offset in [(3, 1), (0, 2)]:
+        state.intermediate_buffer.buffers[slot] = {
+            "req_id": str(slot),
+            "codes": {"ref": codes.flatten() if flat else codes},
+            "ref_offset": offset,
+        }
+    batch = _batch(torch.device("cpu"), [3, 0], [2, 1])
+    req = SimpleNamespace(prompt_len=np.full(5, 8), num_computed_tokens=np.zeros(5))
+    embeds = state._static_inputs_embeds[:3]
+    expected = state.model.embed_input_ids(batch.input_ids) + state.model._audio_embed(codes[[1, 2, 2]])
+    embed = mocker.spy(state.model.model.embed_tokens, "forward")
+    with torch.inference_mode():
+        state.run_preprocess(batch, {"input_ids": batch.input_ids, "inputs_embeds": embeds}, req)
+    torch.testing.assert_close(embeds, expected, rtol=0, atol=0)
+    assert embed.call_count == 1
+    assert state.intermediate_buffer.buffers[3]["ref_offset"] == 3
+    assert state.intermediate_buffer.buffers[0]["ref_offset"] == 3
+
+
+@pytest.mark.cuda
+def test_pending_reference_uploads_survive_staging_reuse_and_overflow():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    state = _state(MossLocalModelState, torch.device("cuda"))
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    outputs = []
+    with torch.cuda.stream(stream), torch.inference_mode():
+        # Keep the first H2D pending while CPU prepares all subsequent batches.
+        torch.cuda._sleep(100_000_000)
+        for value in range(1, 8):
+            embeds = torch.zeros(4, 4, device="cuda", dtype=state.dtype)
+            codes = torch.full((2, 2), value)
+            state._apply_reference_batch(embeds, [(1, codes)])
+            outputs.append((embeds, value))
+            codes.fill_(-1)
+    stream.synchronize()
+    assert len(state._prefill_staging) <= 3
+    for output, value in outputs:
+        expected = torch.zeros_like(output)
+        expected[1:3] = value / 8
+        torch.testing.assert_close(output, expected, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+def test_batched_real_audio_embedding_matches_scalar_bfloat16_reduction():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    state = _state(MossLocalModelState, torch.device("cuda"))
+    model = state.model
+    model.n_vq = 12
+    model.audio_vocab_size = model.audio_pad_token_id = 1024
+    model._stacked_audio_emb_w = torch.randn(12, 1024, 3072, device="cuda", dtype=torch.bfloat16)
+    model._audio_embed = MossTTSLocalTalkerForGeneration._audio_embed.__get__(model)
+    state._static_inputs_embeds = torch.zeros(512, 3072, device="cuda", dtype=torch.bfloat16)
+    references = [(start, torch.randint(0, 1025, (count, 12))) for start, count in [(0, 169), (201, 11), (220, 257)]]
+    with torch.inference_mode():
+        actual = torch.randn_like(state._static_inputs_embeds)
+        expected = actual.clone()
+        for start, codes in references:
+            expected[start : start + len(codes)] += model._audio_embed(codes.to("cuda"))
+        state._apply_reference_batch(actual, references)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)

tests/model_executor/models/moss_tts/test_local_short_attention.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.tts,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+
+@pytest.mark.parametrize("batch,heads,dim", [(1, 4, 64), (16, 32, 80), (128, 32, 80), (3, 8, 128)])
+def test_all_depth_positions_with_strided_cache_and_dirty_future(batch, heads, dim):
+    from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention
+
+    torch.manual_seed(43)
+    query = torch.randn(batch, heads, 1, dim, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(batch, heads, 12, dim, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn_like(key)
+    for position in range(12):
+        # Non-contiguous prefixes retain the stride of the full 12-slot cache.
+        k, v = key[:, :, : position + 1], value[:, :, : position + 1]
+        before_k, before_v = key.clone(), value.clone()
+        expected = F.scaled_dot_product_attention(query, k, v)
+        actual = local_short_attention(query, k, v)
+        assert torch.equal(key, before_k) and torch.equal(value, before_v)
+        if position == 0:
+            assert torch.equal(actual, v)
+        torch.testing.assert_close(actual, expected, atol=0.015625, rtol=0.01)
+        # Future cache contents must never affect an earlier position.
+        key[:, :, position + 1 :] += 100
+        value[:, :, position + 1 :] -= 100
+        assert torch.equal(local_short_attention(query, k, v), actual)
+        key.copy_(before_k)
+        value.copy_(before_v)
+
+
+def test_compile_and_graph_replay_reads_fresh_inputs():
+    from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention
+
+    query = torch.randn(16, 32, 1, 96, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(16, 32, 12, 96, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn_like(key)
+    compiled = torch.compile(local_short_attention, fullgraph=True)
+    for _ in range(3):
+        compiled(query, key, value)
+    torch.accelerator.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        result = compiled(query, key, value)
+    for _ in range(3):
+        query.normal_()
+        key.normal_()
+        value.normal_()
+        graph.replay()
+        torch.testing.assert_close(result, F.scaled_dot_product_attention(query, key, value), atol=0.015625, rtol=0.01)
+
+
+def test_depth_block_cache_overwrite_and_compiled_positions(monkeypatch):
+    from transformers import GPT2Config
+
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+        MossTTSLocalDepthTransformer,
+    )
+
+    monkeypatch.setenv("VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN", "1")
+    model = MossTTSLocalDepthTransformer(GPT2Config(n_embd=320, n_head=4, n_inner=640)).cuda().bfloat16().eval()
+    attn = model.h[0].attn
+    attn.prepare_rope_cache(12, torch.device("cuda"), torch.bfloat16)
+    cache = tuple(torch.empty(4, 4, 12, 80, device="cuda", dtype=torch.bfloat16) for _ in range(2))
+    forward = torch.compile(model._forward_prefix, fullgraph=True, dynamic=True)
+    with torch.inference_mode():
+        for _ in range(2):
+            inputs = torch.randn(4, 12, 320, device="cuda", dtype=torch.bfloat16)
+            # The second frame overwrites the previous frame's K/V before use.
+            for position in range(12):
+                expected = model._forward_prefix(inputs[:, : position + 1])[:, -1:]
+                actual = forward(inputs[:, position : position + 1], cache, position)
+                torch.testing.assert_close(actual, expected, atol=0.03125, rtol=0.02)

tests/model_executor/models/moss_tts/test_reference_encoder.py

@@ -6,6 +6,7 @@ single-flight, and micro-batched encoding."""
 import asyncio
 import threading
 import time
+from tempfile import TemporaryDirectory
 
 import pytest
 import torch
@@ -598,3 +599,239 @@ def test_reference_resampler_is_cached_and_matches_functional_resample():
         assert info.misses == 1 and info.hits == 1
     finally:
         _reference_resampler.cache_clear()
+
+
+# --------------------------------------------------------------------------- #
+# Inline data: URI fast path                                                  #
+# --------------------------------------------------------------------------- #
+
+
+class _DigestAudio(_FakeAudio):
+    """Resolver whose key for data: URIs is the SHA-1 digest, like serving."""
+
+    def __init__(self, *, digest_keys=True):
+        super().__init__()
+        self._digest_keys = digest_keys
+
+    async def resolve(self, ref_str):
+        wav_list, sr, _ = await super().resolve(ref_str)
+        if not self._digest_keys:
+            return wav_list, sr, "rk:" + ref_str
+        import hashlib
+
+        return wav_list, sr, hashlib.sha1(ref_str.encode("utf-8")).hexdigest()
+
+    def artifact_key(self, cache_key):
+        return None
+
+
+async def test_inline_reference_skips_repeat_resolution(make_encoder):
+    proc, audio = _FakeProcessor(), _DigestAudio()
+    ref = "data:audio/wav;base64,QUJD"
+    audio.register(ref, 5)
+    enc = make_encoder(proc)
+
+    first, key1 = await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key)
+    first.fill_(0)  # callers own their copy
+    second, key2 = await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key)
+
+    assert audio.resolve_calls == [ref]
+    assert key1 == key2
+    assert torch.equal(second, torch.full((3, _N_VQ), 5, dtype=torch.long))
+
+
+async def test_inline_index_requires_matching_resolve_key(make_encoder):
+    proc, audio = _FakeProcessor(), _DigestAudio(digest_keys=False)
+    ref = "data:audio/wav;base64,REVG"
+    audio.register(ref, 6)
+    enc = make_encoder(proc)
+
+    for _ in range(2):
+        await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key)
+    assert audio.resolve_calls == [ref, ref]
+
+
+async def test_inline_reference_shared_across_encoders(make_encoder, monkeypatch):
+    # Keep the AF_UNIX socket path independent of pytest's nested base path.
+    with TemporaryDirectory(prefix="moss-ref-") as shared_dir:
+        monkeypatch.setenv("VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR", shared_dir)
+        ref = "data:audio/wav;base64,R0hJ"
+        first_audio, second_audio = _DigestAudio(), _DigestAudio()
+        first_audio.register(ref, 9)
+        second_audio.register(ref, 9)
+        first, second = make_encoder(_FakeProcessor()), make_encoder(_FakeProcessor())
+
+        a, key_a = await first.encode(
+            ref, resolve_ref_audio=first_audio.resolve, get_artifact_key=first_audio.artifact_key
+        )
+        b, key_b = await second.encode(
+            ref, resolve_ref_audio=second_audio.resolve, get_artifact_key=second_audio.artifact_key
+        )
+
+        assert first_audio.resolve_calls == [ref] and second_audio.resolve_calls == []
+        assert key_a == key_b and torch.equal(a, b)
+
+
+# --------------------------------------------------------------------------- #
+# Reference-encoder CUDA graphs: selection and processor-equivalent prep      #
+# --------------------------------------------------------------------------- #
+
+
+class _GraphsStub:
+    def __init__(self, result):
+        self.result = result
+        self.calls: list[list[torch.Tensor]] = []
+
+    def encode(self, wavs):
+        self.calls.append(wavs)
+        return self.result
+
+
+class _LoudnessProcessor(_FakeProcessor):
+    @staticmethod
+    def loudness_normalize(wav):
+        return wav * 2.0
+
+
+async def test_encode_prepared_prefers_graphs_and_falls_back_to_processor(make_encoder):
+    proc = _LoudnessProcessor()
+    enc = make_encoder(proc)
+    graphs_codes = [torch.full((3, _N_VQ), 7, dtype=torch.long)]
+    enc._graphs_enabled, enc._graphs = True, _GraphsStub(graphs_codes)
+    prepared = [torch.ones(1, 100)]
+    assert enc._encode_prepared(prepared) is graphs_codes and proc.attempt_sizes == []
+    enc._graphs = _GraphsStub([None])  # no graph fits this clip
+    out = enc._encode_prepared(prepared)
+    assert proc.attempt_sizes == [1] and out[0].tolist() == [[1] * _N_VQ] * 3
+
+
+async def test_processor_prepare_matches_processor_channel_and_loudness_handling(make_encoder):
+    enc = make_encoder(_LoudnessProcessor())
+    mono = enc._processor_prepare(torch.ones(100))
+    assert mono.shape == (2, 100) and torch.equal(mono, torch.full((2, 100), 2.0))
+    many = enc._processor_prepare(torch.arange(3 * 4, dtype=torch.float32).reshape(3, 4))
+    assert torch.equal(many, torch.arange(8, dtype=torch.float32).reshape(2, 4) * 2.0)
+
+
+async def test_processor_prepare_resamples_to_the_tokenizer_rate(make_encoder):
+    torchaudio = pytest.importorskip("torchaudio")
+    proc = _LoudnessProcessor(sampling_rate=48000)
+    enc = make_encoder(proc)  # works at 24 kHz, the Local-v1.5 reference rate
+    wav = torch.randn(1, 2400)
+    expected = torchaudio.functional.resample(wav.repeat(2, 1), 24000, 48000) * 2.0
+    # The encoder resamples the mono clip before duplicating it. Some CPU conv
+    # kernels round a one-row batch differently from a two-row one, so the two
+    # orders agree to float32 precision, not bit for bit.
+    torch.testing.assert_close(enc._processor_prepare(wav), expected, rtol=0, atol=1e-5)
+
+
+@pytest.mark.parametrize("role", ["", "host", "client"])
+async def test_prepare_captures_graphs_or_starts_the_shared_host(make_encoder, monkeypatch, tmp_path, role):
+    from types import SimpleNamespace
+
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+
+    enc = make_encoder(_FakeProcessor())
+    enc._shares_encoder, enc._shared_codes_dir = True, str(tmp_path)
+    monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: role)
+    captured, hosts, ready = [], [], []
+    monkeypatch.setattr(enc, "_reference_graphs", lambda: captured.append(True))
+    monkeypatch.setattr(enc, "_host_workers", lambda: [])
+    helper = SimpleNamespace(wait_until_ready=lambda: ready.append(True))
+
+    def make_host(*args, **kwargs):
+        hosts.append(args)
+        return helper
+
+    monkeypatch.setattr(re_mod, "SharedReferenceEncoderHost", make_host)
+    monkeypatch.setattr(re_mod, "SharedReferenceEncoderClient", lambda *args: helper)
+    enc.prepare()
+    enc.prepare()
+    # Every API process waits for the shared host before startup completes.
+    assert (len(captured), len(hosts)) == {"": (2, 0), "host": (0, 1), "client": (0, 0)}[role]
+    assert len(ready) == (2 if role else 0)
+
+
+async def test_clips_reach_the_encoder_at_the_tokenizer_rate(make_encoder, monkeypatch):
+    torchaudio = pytest.importorskip("torchaudio")
+    proc = _LoudnessProcessor(sampling_rate=48000)
+    enc = make_encoder(proc)  # works at 24 kHz
+    sent = []
+
+    def record_prepared(prepared):
+        sent.extend(prepared)
+        return [None] * len(prepared)
+
+    monkeypatch.setattr(enc, "_encode_shared_or_local", record_prepared)
+    wav = [float(v) for v in torch.randn(2400)]
+    enc._encode_batch_sync([(wav, _SR)])
+    # Resampled where the request arrived (a shared encoder receives it this way); still one channel.
+    expected = torchaudio.functional.resample(torch.tensor(wav)[None], 24000, 48000)
+    torch.testing.assert_close(sent[0], expected, rtol=0, atol=0)
+    # The eager path is told the clip is already at the tokenizer rate.
+    calls = []
+
+    def record_encoding(wavs, sampling_rate, n_vq=None):
+        calls.append(sampling_rate)
+        return [None]
+
+    monkeypatch.setattr(proc, "encode_audios_from_wav", record_encoding)
+    enc._graphs_enabled = False
+    enc._encode_prepared(sent)
+    assert calls == [48000]
+
+
+async def test_batcher_runs_up_to_max_inflight_batches_at_once():
+    inside = threading.Barrier(2, timeout=5)
+    seen: list[int] = []
+
+    def encode_batch(payload):
+        inside.wait()  # two batches are encoding at the same time
+        seen.append(len(payload))
+        return [torch.zeros(1) for _ in payload]
+
+    batcher = _RefEncodeBatcher(encode_batch, window_ms=0.0, max_batch=8, max_inflight=2)
+    try:
+        first = asyncio.ensure_future(batcher.submit([1.0], _SR))
+        await asyncio.sleep(0.05)  # the first batch is in flight before the second clip arrives
+        await asyncio.gather(first, batcher.submit([2.0], _SR))
+    finally:
+        await batcher.aclose()
+    assert seen == [1, 1]
+
+
+def test_reference_stream_priority(monkeypatch):
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+
+    monkeypatch.delenv(re_mod._REF_ENCODE_PRIORITY_ENV, raising=False)
+    assert re_mod.reference_stream_priority() == 0
+    monkeypatch.setenv(re_mod._REF_ENCODE_PRIORITY_ENV, "high")
+    expected = torch.cuda.Stream.priority_range()[1] if torch.cuda.is_available() else 0
+    assert re_mod.reference_stream_priority() == expected
+
+
+async def test_device_prep_resamples_on_the_tokenizer_device(make_encoder, monkeypatch):
+    torchaudio = pytest.importorskip("torchaudio")
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+
+    proc = _LoudnessProcessor(sampling_rate=48000)
+    monkeypatch.setattr(proc, "audio_tokenizer", torch.nn.Linear(1, 1), raising=False)  # CPU placement
+    enc = make_encoder(proc)
+    enc._gpu_prep = True
+    wav = torch.randn(1, 2400)
+    clip = enc._encoder_input(wav)
+    assert clip.shape == (1, 2400)  # left at the working rate for the encoder's device
+    expected = torchaudio.functional.resample(wav, 24000, 48000).repeat(2, 1) * 2.0
+    torch.testing.assert_close(enc._finish_prepare(clip), expected, rtol=1e-5, atol=1e-6)
+    # The eager fallback is told the clip's own rate.
+    calls = []
+
+    def record_encoding(wavs, sampling_rate, n_vq=None):
+        calls.append(sampling_rate)
+        return [None]
+
+    monkeypatch.setattr(proc, "encode_audios_from_wav", record_encoding)
+    enc._graphs_enabled = False
+    enc._encode_prepared([clip])
+    assert calls == [24000]
+    assert re_mod._REF_ENCODE_GPU_PREP_ENV == "VLLM_OMNI_MOSS_REF_GPU_PREP"

tests/model_executor/models/moss_tts/test_reference_encoder_graphs.py

@@ -0,0 +1,320 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Reference-encoder CUDA graphs reproduce the tokenizer's eager encode."""
+
+from contextlib import nullcontext
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm_omni.model_executor.models.moss_tts.reference_encoder_graphs import (
+    MossReferenceEncoderGraphs,
+    _cached_reference_rope,
+    _reference_rope_factors,
+    _validate_unmasked_reference_encoder,
+    _windowed_attention,
+    install_windowed_attention,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+FRAME, N_VQ = 8, 3
+
+
+class _Stage(nn.Module):
+    """Causal per-frame stage: frames never mix, like a causal encoder layer."""
+
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out)
+
+    def forward(self, x, lengths):
+        return torch.tanh(self.proj(x)), lengths
+
+
+class _Quantizer(nn.Module):
+    def forward(self, hidden, lengths, n_q):
+        # hidden (B, frames, D): codes (n_q, B, frames) from per-frame argmax slices.
+        codes = torch.stack([hidden[..., q::n_q].argmax(-1) for q in range(n_q)])
+        return hidden, codes, lengths
+
+
+class _Tokenizer(nn.Module):
+    number_channels = 2
+    downsample_rate = FRAME
+    sampling_rate = 64
+
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.ModuleList([_Stage(FRAME * 2, 16), _Stage(16, 12)])
+        self.quantizer = _Quantizer()
+
+    def _flatten_channels_for_codec(self, audio, lengths):
+        batch, channels, samples = audio.shape
+        frames = samples // FRAME
+        values = audio.reshape(batch, channels, frames, FRAME).permute(0, 2, 1, 3).reshape(batch, frames, -1)
+        return values, lengths // FRAME
+
+    def _codec_inference_autocast(self):
+        return nullcontext()
+
+    def eager(self, wavs):
+        out = []
+        for wav in wavs:
+            values, lengths = self._flatten_channels_for_codec(wav[None], torch.tensor([wav.shape[-1]]))
+            for module in self.encoder:
+                values, lengths = module(values, lengths)
+            _, codes, _ = self.quantizer(values, lengths, N_VQ)
+            out.append(codes[:, 0, : wav.shape[-1] // FRAME].transpose(0, 1).long())
+        return out
+
+
+@pytest.fixture
+def cuda():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device("cuda", torch.accelerator.current_device_index())
+
+
+def test_cached_rope_reuses_longest_table_without_modifying_inputs(cuda):
+    from types import SimpleNamespace
+
+    module = SimpleNamespace(rope=SimpleNamespace(max_period=10000.0))
+    cache: dict[tuple[int, float, torch.device], tuple[torch.Tensor, torch.Tensor]] = {}
+    torch.manual_seed(17)
+    # Prime once; all subsequent sequences and batches use an immutable prefix.
+    q = torch.randn(1, 3, 128, 32, device=cuda, dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    _cached_reference_rope(module, q, k, cache)
+    stored = next(iter(cache.values()))
+    for batch, length in [(1, 17), (4, 64), (2, 128)]:
+        q = torch.randn(batch, 3, length, 32, device=cuda, dtype=torch.bfloat16)
+        k = torch.randn_like(q)
+        before = (q.clone(), k.clone())
+        got = _cached_reference_rope(module, q, k, cache)
+        cos, sin = _reference_rope_factors(length, 32, 10000.0, cuda)
+        torch.testing.assert_close(stored[0][:length], cos, rtol=0, atol=0)
+        torch.testing.assert_close(stored[1][:length], sin, rtol=0, atol=0)
+        for value, source, saved in zip(got, (q, k), before):
+            real, imag = source.float().reshape(batch, 3, length, 16, 2).unbind(-1)
+            expected = torch.stack((real * cos - imag * sin, real * sin + imag * cos), -1).to(source.dtype)
+            torch.testing.assert_close(value, expected.reshape_as(source), rtol=0, atol=0)
+            torch.testing.assert_close(source, saved, rtol=0, atol=0)
+        assert len(cache) == 1
+        assert next(iter(cache.values())) is stored
+
+
+def test_cached_rope_distinguishes_period_and_head_width(cuda):
+    from types import SimpleNamespace
+
+    cache: dict[tuple[int, float, torch.device], tuple[torch.Tensor, torch.Tensor]] = {}
+    for width, period in [(32, 10000.0), (64, 10000.0), (32, 1000.0)]:
+        module = SimpleNamespace(rope=SimpleNamespace(max_period=period))
+        q = torch.randn(1, 2, 32, width, device=cuda, dtype=torch.bfloat16)
+        _cached_reference_rope(module, q, q, cache)
+    assert len(cache) == 3
+
+
+@pytest.mark.parametrize("batched_transfer", [False, True])
+def test_graph_codes_match_eager_per_clip_and_fall_back_outside_buckets(cuda, batched_transfer):
+    torch.manual_seed(0)
+    tokenizer = _Tokenizer().to(cuda)
+    # Buckets of 2 s and 4 s at 64 Hz: 128 and 256 samples.
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer, n_vq=N_VQ, batch_sizes=(1, 2, 4), bucket_seconds=(2.0, 4.0), batched_transfer=batched_transfer
+    )
+    graphs.capture()
+    assert graphs.captured == [(1, 128), (1, 256), (2, 128), (2, 256), (4, 128), (4, 256)]
+    assert graphs.max_batch == 4
+
+    for count, frames in [(1, [5]), (3, [16, 3, 9]), (4, [32, 1, 20, 17]), (2, [16, 16])]:
+        wavs = [torch.randn(2, n * FRAME) for n in frames]
+        codes = graphs.encode(wavs)
+        expected = tokenizer.eager([w.to(cuda) for w in wavs])
+        assert len(codes) == count
+        for got, want in zip(codes, expected):
+            assert got.device.type == "cpu" and got.dtype == torch.long
+            torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+    # Longer than every bucket: left to the caller; the rest still use graphs.
+    long_clip, short_clip = torch.randn(2, 33 * FRAME), torch.randn(2, 4 * FRAME)
+    codes = graphs.encode([long_clip, short_clip])
+    assert codes[0] is None
+    torch.testing.assert_close(codes[1], tokenizer.eager([short_clip.to(cuda)])[0].cpu(), rtol=0, atol=0)
+
+    # Wider than every batch: split, longest first, and returned in input order.
+    wavs = [torch.randn(2, n * FRAME) for n in (3, 30, 7, 12, 1)]
+    for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+    # Result ownership includes the degenerate single-frame shape. Reusing
+    # the pinned output slab must not overwrite a previous caller's codes.
+    previous = graphs.encode([torch.randn(2, FRAME)])[0]
+    snapshot = previous.clone()
+    graphs.encode([torch.randn(2, 16 * FRAME) for _ in range(4)])
+    torch.testing.assert_close(previous, snapshot, rtol=0, atol=0)
+
+
+def test_batch_audio_budget_limits_captured_graphs(cuda):
+    tokenizer = _Tokenizer().to(cuda)
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer, n_vq=N_VQ, batch_sizes=(1, 2, 4), bucket_seconds=(2.0, 4.0), max_batch_seconds=8.0
+    )
+    graphs.capture()
+    # 4 x 4 s exceeds the 8 s budget; a single clip is always captured.
+    assert graphs.captured == [(1, 128), (1, 256), (2, 128), (2, 256), (4, 128)]
+    wavs = [torch.randn(2, 30 * FRAME) for _ in range(3)]
+    for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+
+def test_finer_singleton_buckets_preserve_multiclip_batching(cuda):
+    tokenizer = _Tokenizer().to(cuda)
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer,
+        n_vq=N_VQ,
+        batch_sizes=(1, 2, 4),
+        bucket_seconds=(2.0, 4.0),
+        singleton_bucket_seconds=(1.0, 3.0),
+    )
+    graphs.capture()
+    assert (1, 64) in graphs.captured and (1, 192) in graphs.captured
+    assert (2, 64) not in graphs.captured and (4, 192) not in graphs.captured
+    replay = graphs._replay
+    shapes = []
+
+    def record(entry, wavs):
+        shapes.append(tuple(entry.audio.shape))
+        return replay(entry, wavs)
+
+    graphs._replay = record
+    one = torch.randn(2, 20 * FRAME)
+    out = graphs.encode([one])
+    assert shapes == [(1, 2, 192)]
+    torch.testing.assert_close(out[0], tokenizer.eager([one.to(cuda)])[0].cpu(), rtol=0, atol=0)
+    shapes.clear()
+    wavs = [one, one[:, : 8 * FRAME], one[:, : 12 * FRAME]]
+    out = graphs.encode(wavs)
+    assert shapes == [(4, 2, 256)]
+    for got, want in zip(out, tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+
+def test_compiled_capture_matches_eager(cuda):
+    torch.manual_seed(2)
+    tokenizer = _Tokenizer().to(cuda)
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer, n_vq=N_VQ, batch_sizes=(1, 4), bucket_seconds=(2.0, 4.0), compile_core=True
+    )
+    graphs.capture()
+    assert graphs.captured == [(1, 128), (1, 256), (4, 128), (4, 256)]
+    wavs = [torch.randn(2, n * FRAME) for n in (5, 30, 12)]
+    for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+
+class _Attention(nn.Module):
+    """The attention surface the tokenizer's MHA exposes to the windowed kernel."""
+
+    def __init__(self, embed_dim=128, heads=2, context=5):
+        super().__init__()
+        self.embed_dim, self.num_heads, self.context, self.causal = embed_dim, heads, context, True
+        self.in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False)
+
+    def _project_qkv(self, x):
+        b, t, _ = x.shape
+        return self.in_proj(x).reshape(b, t, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4).unbind(0)
+
+    def _apply_dense_rope(self, q, k):
+        return q, k
+
+    def _forward_non_streaming_sdpa(self, x, input_lengths):
+        # Reference semantics: causal, distance < context, padded rows zeroed.
+        b, t, _ = x.shape
+        q, k, v = self._project_qkv(x)
+        pos = torch.arange(t, device=x.device)
+        delta = pos.view(-1, 1) - pos.view(1, -1)
+        mask = (delta >= 0) & (delta < self.context)
+        out = F.scaled_dot_product_attention(q, k, v, mask[None, None])
+        valid = (pos.view(1, t) < input_lengths.view(-1, 1)).view(b, 1, t, 1)
+        out = torch.where(valid, out, torch.zeros((), device=out.device, dtype=out.dtype))
+        return out.transpose(1, 2).reshape(b, t, self.embed_dim)
+
+
+@pytest.mark.parametrize("fa_version", [2, 3])
+def test_windowed_attention_matches_masked_sdpa(cuda, fa_version):
+    if fa_version == 3 and torch.cuda.get_device_capability(cuda)[0] != 9:
+        pytest.skip("FA3 requires Hopper")
+    torch.manual_seed(1)
+    model = nn.Module()
+    model.encoder = nn.ModuleList([_Attention().to(cuda, torch.bfloat16)])
+    attention = model.encoder[0]
+    x = torch.randn(3, 40, 128, device=cuda, dtype=torch.bfloat16)
+    lengths = torch.tensor([40, 17, 3], device=cuda)
+    expected = attention._forward_non_streaming_sdpa(x, lengths)
+    assert install_windowed_attention(model, fa_version=fa_version) == 1
+    actual = attention._forward_non_streaming_sdpa(x, lengths)
+    torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+    assert actual[1, 17:].count_nonzero() == 0 and actual[2, 3:].count_nonzero() == 0
+    # Non-bf16 inputs keep the tokenizer's own SDPA path.
+    x32 = x.float()
+    attention.float()
+    torch.testing.assert_close(
+        attention._forward_non_streaming_sdpa(x32, lengths), expected.float(), rtol=3e-2, atol=3e-2
+    )
+
+
+@pytest.mark.parametrize("fa_version", [2, 3])
+def test_unmasked_causal_queries_preserve_valid_prefix_and_ignore_future_padding(cuda, fa_version):
+    if fa_version == 3 and torch.cuda.get_device_capability(cuda)[0] != 9:
+        pytest.skip("FA3 requires Hopper")
+    torch.manual_seed(44)
+    attention = _Attention().to(cuda, torch.bfloat16)
+    x = torch.randn(4, 40, 128, device=cuda, dtype=torch.bfloat16)
+    lengths = torch.tensor([40, 17, 1, 0], device=cuda)
+    kwargs = dict(sdpa=attention._forward_non_streaming_sdpa, fa_version=fa_version)
+    with torch.no_grad():
+        masked = _windowed_attention(attention, x, lengths, **kwargs)
+        unmasked = _windowed_attention(attention, x, lengths, skip_padded_query_mask=True, **kwargs)
+        poisoned = x.clone()
+        for row, length in enumerate([40, 17, 1, 0]):
+            poisoned[row, length:] = torch.randn_like(poisoned[row, length:]) * 30
+        changed = _windowed_attention(attention, poisoned, lengths, skip_padded_query_mask=True, **kwargs)
+        for row, length in enumerate([40, 17, 1, 0]):
+            torch.testing.assert_close(unmasked[row, :length], masked[row, :length], rtol=0, atol=0)
+            torch.testing.assert_close(changed[row, :length], masked[row, :length], rtol=0, atol=0)
+            assert masked[row, length:].count_nonzero() == 0
+        assert unmasked[1, 17:].count_nonzero() > 0
+        # The FP32 fallback retains its original masking contract.
+        attention.float()
+        got = _windowed_attention(attention, x.float(), lengths, skip_padded_query_mask=True, **kwargs)
+        expected = attention._forward_non_streaming_sdpa(x.float(), lengths)
+        torch.testing.assert_close(got, expected, rtol=0, atol=0)
+
+
+def test_unmasked_queries_reject_unknown_encoder_before_patching():
+    tokenizer = _Tokenizer()
+    with pytest.raises(ValueError, match="patch-downsample"):
+        install_windowed_attention(tokenizer, skip_padded_query_mask=True)
+
+
+@pytest.mark.parametrize("bad_part", ["upsample", "noncausal", "quantizer", "convolution"])
+def test_unmasked_queries_reject_structures_that_can_read_invalid_rows(bad_part):
+    from types import SimpleNamespace
+
+    patch = type("MossAudioTokenizerPatchedPretransform", (), {})()
+    patch.is_downsample = bad_part != "upsample"
+    projected = type("MossAudioTokenizerProjectedTransformer", (), {})()
+    projected.transformer = SimpleNamespace(
+        layers=[SimpleNamespace(self_attn=SimpleNamespace(causal=bad_part != "noncausal"))]
+    )
+    quantizer_cls = type(
+        "MossAudioTokenizerResidualVQ" if bad_part != "quantizer" else "UnknownQuantizer", (nn.Module,), {}
+    )
+    quantizer = quantizer_cls()
+    quantizer.proj = nn.Conv1d(2, 2, 3 if bad_part == "convolution" else 1)
+    tokenizer = SimpleNamespace(encoder=[patch, projected], quantizer=quantizer)
+    with pytest.raises(ValueError):
+        _validate_unmasked_reference_encoder(tokenizer)

tests/model_executor/models/moss_tts/test_shared_reference_encoder.py

@@ -0,0 +1,341 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""API processes share one reference encoder through a Unix socket."""
+
+import threading
+from tempfile import TemporaryDirectory
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts import shared_reference_encoder as shared
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _codes(wav):
+    # Deterministic per clip: frames from the clip length, values from its first sample.
+    return torch.full((wav.shape[-1] // 10, 4), int(wav.reshape(-1)[0]), dtype=torch.long)
+
+
+class _Recorder:
+    def __init__(self, gate=None):
+        self.batches = []
+        self.gate = gate
+
+    def __call__(self, wavs):
+        if self.gate is not None:
+            self.gate.wait(5)
+        self.batches.append(len(wavs))
+        return [RuntimeError("bad clip") if int(w.reshape(-1)[0]) < 0 else _codes(w) for w in wavs]
+
+
+@pytest.fixture
+def host_dir(monkeypatch):
+    monkeypatch.setattr(shared, "_host_lock_handle", None)
+    # AF_UNIX limits the full socket path, including pytest directory names.
+    with TemporaryDirectory(prefix="moss-ref-") as path:
+        yield path
+
+
+def test_only_one_process_per_directory_hosts(host_dir):
+    assert shared.elect_host(host_dir) is True
+    assert shared.elect_host(host_dir) is False  # a second open file description cannot take the lock
+
+
+def test_close_releases_queued_and_active_callers(host_dir):
+    entered, release = threading.Event(), threading.Event()
+
+    def encode(wavs):
+        entered.set()
+        assert release.wait(5)
+        return [_codes(wav) for wav in wavs]
+
+    host = shared.SharedReferenceEncoderHost(host_dir, [(encode, None)])
+    results: dict[str, list[torch.Tensor | BaseException]] = {}
+    threads = [
+        threading.Thread(target=lambda key=key: results.setdefault(key, host.encode([torch.ones(1, 20)])), daemon=True)
+        for key in ("active", "queued")
+    ]
+    try:
+        threads[0].start()
+        assert entered.wait(2)
+        threads[1].start()
+        with host._ready:
+            assert host._ready.wait_for(lambda: len(host._pending) == 1, timeout=2)
+        host.close()
+        for thread in threads:
+            thread.join(1)
+            assert not thread.is_alive(), "closed encoder left a request waiting"
+        assert all(isinstance(results[key][0], RuntimeError) for key in results)
+    finally:
+        release.set()
+        host.close()
+        for thread in threads:
+            if thread.ident is not None:
+                thread.join(1)
+
+
+def test_worker_result_count_mismatch_fails_entire_batch(host_dir):
+    host = shared.SharedReferenceEncoderHost(host_dir, [(lambda wavs: [_codes(wavs[0])], None)])
+    results = []
+    caller = threading.Thread(target=lambda: results.extend(host.encode([torch.ones(1, 20)] * 2)), daemon=True)
+    try:
+        caller.start()
+        caller.join(2)
+        assert not caller.is_alive(), "missing encoder result left a request waiting"
+        assert len(results) == 2
+        assert all(isinstance(result, RuntimeError) for result in results)
+    finally:
+        host.close()
+        caller.join(1)
+
+
+def test_client_round_trip_and_per_clip_errors(host_dir):
+    host = shared.SharedReferenceEncoderHost(host_dir, [(_Recorder(), None)])
+    try:
+        client = shared.SharedReferenceEncoderClient(host_dir)
+        wavs = [torch.full((1, 40), 3.0), torch.full((1, 70), -1.0), torch.full((1, 20), 5.0)]
+        out = client.encode(wavs)
+        assert out[0].dtype == torch.long and torch.equal(out[0], _codes(wavs[0]))
+        assert isinstance(out[1], RuntimeError) and "bad clip" in str(out[1])
+        assert torch.equal(out[2], _codes(wavs[2]))
+        # The host process's own clips skip the socket.
+        assert torch.equal(host.encode([wavs[2]])[0], _codes(wavs[2]))
+    finally:
+        host.close()
+
+
+def test_host_merges_clips_from_concurrent_clients(host_dir):
+    gate = threading.Event()
+    recorder = _Recorder(gate)
+    host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, None)])
+    try:
+        results = {}
+
+        def send(name, value):
+            results[name] = shared.SharedReferenceEncoderClient(host_dir).encode([torch.full((1, 30), value)] * 2)
+
+        # Encoding is held until both clients' clips are queued.
+        warm = threading.Thread(target=lambda: host.encode([torch.full((1, 10), 1.0)]))
+        warm.start()
+        threads = [threading.Thread(target=send, args=(n, v)) for n, v in (("a", 2.0), ("b", 4.0))]
+        for thread in threads:
+            thread.start()
+        threading.Timer(0.5, gate.set).start()
+        for thread in (warm, *threads):
+            thread.join(10)
+        # Both clients' clips were encoded together, not one request at a time.
+        assert sum(recorder.batches) == 5 and max(recorder.batches) >= 4
+        assert all(int(c[0, 0]) == 2 for c in results["a"]) and all(int(c[0, 0]) == 4 for c in results["b"])
+    finally:
+        host.close()
+
+
+def test_client_without_host_raises(host_dir, monkeypatch):
+    monkeypatch.setattr(shared, "_CONNECT_TIMEOUT_S", 0.3)
+    with pytest.raises((FileNotFoundError, ConnectionRefusedError)):
+        shared.SharedReferenceEncoderClient(host_dir).encode([torch.zeros(1, 10)])
+
+
+def test_host_window_collects_across_workers_without_reserving_singletons(host_dir):
+    recorder = _Recorder()
+    host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, None)] * 3, batch_window_ms=500)
+    try:
+        barrier = threading.Barrier(3)
+        results = []
+
+        def send(value):
+            barrier.wait()
+            results.extend(host.encode([torch.full((1, 20), value)]))
+
+        threads = [threading.Thread(target=send, args=(v,)) for v in (1.0, 2.0)]
+        for thread in threads:
+            thread.start()
+        barrier.wait()
+        for thread in threads:
+            thread.join(5)
+            assert not thread.is_alive()
+        assert recorder.batches == [2]
+        assert sorted(int(x[0, 0]) for x in results) == [1, 2]
+    finally:
+        host.close()
+
+
+def test_reference_encoder_client_falls_back_to_local_encode(host_dir, monkeypatch):
+    from tests.model_executor.models.moss_tts.test_reference_encoder import _N_VQ, _SR, _FakeProcessor
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+    from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+    monkeypatch.setenv(re_mod._SHARED_CODES_DIR_ENV, host_dir)
+    monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: "client")
+    monkeypatch.setattr(shared, "_CONNECT_TIMEOUT_S", 0.3)
+    proc = _FakeProcessor()
+    enc = re_mod.MossReferenceEncoder(
+        proc, variant="local", n_vq=_N_VQ, sr_target=_SR, speaker_cache=SpeakerEmbeddingCache(max_bytes=1 << 20)
+    )
+    placed = []
+    monkeypatch.setattr(enc, "_place_tokenizer_for_local_encode", lambda: placed.append(True))
+    out = enc._encode_batch_sync([([7.0] * 50, _SR)])
+    assert placed == [True] and proc.attempt_sizes == [1] and int(out[0][0, 0]) == 7
+
+
+def test_workers_encode_concurrently_after_their_warmups(host_dir):
+    inside, release = threading.Barrier(3, timeout=5), threading.Event()
+    warmed = []
+
+    def worker(name):
+        def encode(wavs):
+            inside.wait()  # both workers are inside an encode at the same time
+            release.wait(5)
+            return [_codes(w) for w in wavs]
+
+        return encode, lambda: warmed.append(name)
+
+    host = shared.SharedReferenceEncoderHost(host_dir, [worker("a"), worker("b")])
+    try:
+        out = []
+
+        def submit(value):
+            out.append(host.encode([torch.full((1, 20), value)]))
+
+        first = threading.Thread(target=submit, args=(1.0,))
+        first.start()
+        # The first worker is busy by now, so the second clip goes to the other one.
+        threading.Event().wait(0.1)
+        second = threading.Thread(target=submit, args=(2.0,))
+        second.start()
+        inside.wait()
+        release.set()
+        for thread in (first, second):
+            thread.join(5)
+        assert sorted(warmed) == ["a", "b"] and sorted(int(o[0][0, 0]) for o in out) == [1, 2]
+    finally:
+        host.close()
+
+
+def test_client_requests_in_flight_use_their_own_connections(host_dir):
+    entered, release = threading.Semaphore(0), threading.Event()
+
+    def worker():
+        def encode(wavs):
+            entered.release()
+            release.wait(5)
+            return [_codes(w) for w in wavs]
+
+        return encode, None
+
+    host = shared.SharedReferenceEncoderHost(host_dir, [worker(), worker()])
+    try:
+        client = shared.SharedReferenceEncoderClient(host_dir)
+        out = []
+        threads = [
+            threading.Thread(target=lambda v=v: out.append(client.encode([torch.full((1, 20), v)]))) for v in (1.0, 2.0)
+        ]
+        for thread in threads:
+            thread.start()
+            # A worker holds this request, still unanswered, before the next is
+            # sent. Sent together, one free worker may take both as one batch.
+            assert entered.acquire(timeout=5)
+        release.set()  # both requests reached the host at once
+        for thread in threads:
+            thread.join(5)
+        assert sorted(int(o[0][0, 0]) for o in out) == [1, 2]
+        assert len(client._idle) == 2  # both connections are kept for reuse
+    finally:
+        host.close()
+
+
+def test_readiness_waits_for_every_worker_warmup(host_dir, monkeypatch):
+    entered, release = threading.Event(), threading.Event()
+    recorder = _Recorder()
+    warmed = []
+
+    def last_warmup():
+        entered.set()
+        assert release.wait(5)
+        warmed.append("last")
+
+    host = shared.SharedReferenceEncoderHost(
+        host_dir, [(recorder, lambda: warmed.append("first")), (recorder, last_warmup)]
+    )
+    client = shared.SharedReferenceEncoderClient(host_dir)
+    try:
+        assert entered.wait(5)
+        monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 0.03)
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            host.wait_until_ready()
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            client.wait_until_ready()
+        assert warmed == ["first"] and recorder.batches == [] and not client._idle
+        release.set()
+        monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 5)
+        host.wait_until_ready()
+        client.wait_until_ready()
+        assert warmed == ["first", "last"] and recorder.batches == []
+        assert len(client._idle) == 1
+        assert torch.equal(client.encode([torch.ones(1, 20)])[0], _codes(torch.ones(1, 20)))
+    finally:
+        release.set()
+        host.close()
+
+
+def test_slow_startup_does_not_use_encode_timeout(host_dir, monkeypatch):
+    release = threading.Event()
+    monkeypatch.setattr(shared, "_REQUEST_TIMEOUT_S", 0.05)
+    monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 5, raising=False)
+    host = shared.SharedReferenceEncoderHost(host_dir, [(_Recorder(), lambda: release.wait(5))])
+    timer = threading.Timer(0.2, release.set)
+    timer.start()
+    try:
+        # The first socket waits for readiness before starting its much shorter
+        # encode deadline. No API warmup call is required to protect this path.
+        client = shared.SharedReferenceEncoderClient(host_dir)
+        assert torch.equal(client.encode([torch.ones(1, 20)])[0], _codes(torch.ones(1, 20)))
+    finally:
+        release.set()
+        timer.join(5)
+        host.close()
+
+
+def test_failed_host_warmup_is_not_ready(host_dir):
+    def fail():
+        raise RuntimeError("warmup failed")
+
+    recorder = _Recorder()
+    host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, fail)])
+    try:
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            host.wait_until_ready()
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            shared.SharedReferenceEncoderClient(host_dir).wait_until_ready()
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            host.encode([torch.ones(1, 20)])
+        assert recorder.batches == []
+    finally:
+        host.close()
+
+
+def test_startup_failure_does_not_trigger_local_capture(host_dir, monkeypatch):
+    from unittest.mock import Mock
+
+    from tests.model_executor.models.moss_tts.test_reference_encoder import _N_VQ, _SR, _FakeProcessor
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+    from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+    monkeypatch.setenv(re_mod._SHARED_CODES_DIR_ENV, host_dir)
+    monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: "client")
+    enc = re_mod.MossReferenceEncoder(
+        _FakeProcessor(),
+        variant="local",
+        n_vq=_N_VQ,
+        sr_target=_SR,
+        speaker_cache=SpeakerEmbeddingCache(max_bytes=1 << 20),
+    )
+    enc._shared_client = Mock()
+    enc._shared_client.encode.side_effect = shared.SharedReferenceEncoderStartupError("still compiling")
+    local = Mock()
+    monkeypatch.setattr(enc, "_place_tokenizer_for_local_encode", local)
+    with pytest.raises(shared.SharedReferenceEncoderStartupError):
+        enc._encode_shared_or_local([torch.ones(1, 20)])
+    local.assert_not_called()

tests/model_executor/models/moss_tts/test_slot_attention.py

@@ -0,0 +1,301 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import copy
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import (
+    MossAudioTokenizerMultiheadAttention,
+    MossAudioTokenizerRotaryEmbedding,
+    RingKVCache,
+    StreamingExecutionContext,
+)
+from vllm_omni.model_executor.models.moss_tts.slot_attention import slot_ring_attention, slot_ring_attention_rows
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize(
+    "frames,capacity", [(1, 125), (2, 250), (15, 125), (30, 250), (32, 400), (60, 400), (480, 400)]
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("execution", ["graph", "compiled"])
+def test_empty_tile_skip_matches_static_ring_across_reuse(frames, capacity, dtype, execution):
+    torch.manual_seed(45)
+    batch, heads, states, dim = 3, 2, 5, 64
+    projected = torch.randn(batch, frames, 3, heads, dim, device="cuda", dtype=dtype)
+    q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0)
+    cache = torch.randn(2, states, heads, capacity, dim, device="cuda", dtype=dtype)
+    expected_cache = cache.clone()
+    end = torch.zeros(states, 2, device="cuda", dtype=torch.long)[:, 0]
+    expected_end = end.clone()
+    slots = torch.tensor([0, 4, 2], device="cuda")
+    lengths = torch.full((batch, 2), frames, device="cuda", dtype=torch.int32)[:, 0]
+
+    def candidate():
+        return slot_ring_attention(q, k, v, cache, end, slots, lengths, 17, skip_empty_tiles=True)
+
+    graph = None
+    if execution == "compiled":
+        run = torch.compile(candidate, fullgraph=True)
+    else:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            candidate()
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            actual = candidate()
+
+    for step, start in enumerate([0, 1, capacity - 1, capacity + 9, 0, 3 * capacity]):
+        end.fill_(start)
+        expected_end.fill_(start)
+        cache.copy_(expected_cache)
+        projected.normal_()
+        slots.copy_(torch.roll(torch.arange(states, device="cuda"), step)[:batch])
+        lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32))
+        expected = slot_ring_attention(q, k, v, expected_cache, expected_end, slots, lengths, 17)
+        if graph is not None:
+            graph.replay()
+        else:
+            actual = run()
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0)
+        torch.testing.assert_close(end, expected_end, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("frames,capacity", [(1, 125), (15, 140), (480, 400), (480, 880)])
+def test_all_padding_preserves_state_and_replay_can_become_active(dtype, frames, capacity):
+    q, k, v = [torch.randn(2, 2, frames, 64, device="cuda", dtype=dtype) for _ in range(3)]
+    cache = torch.randn(2, 4, 2, capacity, 64, device="cuda", dtype=dtype)
+    end = torch.tensor([0, capacity + 3, 9, 0], device="cuda")
+    slots = torch.tensor([1, 3], device="cuda")
+    lengths = torch.zeros(2, device="cuda", dtype=torch.int32)
+    before_cache, before_end = cache.clone(), end.clone()
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            slot_ring_attention(q, k, v, cache, end, slots, lengths, capacity)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = slot_ring_attention(q, k, v, cache, end, slots, lengths, capacity)
+    for active in [False, True, False, True]:
+        lengths.fill_(frames if active else 0)
+        expected = _ring_reference(q, k, v, before_cache, before_end, slots, lengths, capacity)
+        graph.replay()
+        torch.testing.assert_close(cache, before_cache, rtol=0, atol=0)
+        torch.testing.assert_close(end, before_end, rtol=0, atol=0)
+        torch.testing.assert_close(out, expected, rtol=0.025, atol=0.025)
+        if not active:
+            assert out.count_nonzero() == 0
+
+
+def test_dense_ring_retains_last_tokens_when_chunk_exceeds_capacity():
+    torch.manual_seed(19)
+    capacity, frames = 400, 480
+    state = RingKVCache(3, 2, 64, capacity, device=torch.device("cuda"), dtype=torch.float32)
+    slots = torch.tensor([1, 0], device="cuda")
+    context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=torch.ones(2, device="cuda", dtype=torch.bool))
+    for start in [0, frames]:
+        k = torch.randn(2, 2, frames, 64, device="cuda")
+        v = torch.randn_like(k)
+        state.complete(k, v, execution_context=context)
+        retained = torch.arange(frames - capacity, frames, device="cuda")
+        destinations = (start + retained) % capacity
+        torch.testing.assert_close(state.cache[0, slots][:, :, destinations], k[:, :, retained], rtol=0, atol=0)
+        torch.testing.assert_close(state.cache[1, slots][:, :, destinations], v[:, :, retained], rtol=0, atol=0)
+        torch.testing.assert_close(state.end_offset[slots], torch.full_like(slots, start + frames), rtol=0, atol=0)
+
+
+def _ring_reference(q, k, v, cache, end, slots, lengths, context):
+    """Independent chronological ring oracle with FP32 attention arithmetic."""
+    capacity = cache.shape[-2]
+    output = torch.zeros_like(q)
+    for row, (slot, length) in enumerate(zip(slots.tolist(), lengths.tolist())):
+        if length == 0:
+            continue
+        start = int(end[slot])
+        for token in range(max(0, length - capacity), length):
+            cache[0, slot, :, (start + token) % capacity] = k[row, :, token]
+            cache[1, slot, :, (start + token) % capacity] = v[row, :, token]
+        stop = start + length
+        physical_positions = torch.full((capacity,), -1, device=q.device, dtype=torch.long)
+        available = torch.arange(max(0, stop - capacity), stop, device=q.device)
+        physical_positions[available % capacity] = available
+        queries = torch.arange(start, stop, device=q.device)
+        distance = queries[:, None] - physical_positions[None]
+        mask = (physical_positions[None] >= 0) & (distance >= 0)
+        if context > 0:
+            mask &= distance < context
+        output[row, :, :length] = F.scaled_dot_product_attention(
+            q[row, :, :length].float(), cache[0, slot].float(), cache[1, slot].float(), mask[None]
+        ).to(q.dtype)
+        end[slot] = stop
+    return output
+
+
+@pytest.mark.parametrize("frames,capacity,head_dim", [(1, 125, 64), (15, 125, 64), (480, 400, 64), (33, 65, 32)])
+@pytest.mark.parametrize("execution", ["eager", "graph", "compiled"])
+def test_slot_attention_preserves_ring_state_and_padding(frames, capacity, head_dim, execution):
+    torch.manual_seed(29)
+    batch, heads, state_capacity = 3, 2, 5
+    projected = torch.randn(batch, frames, 3, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0)
+    cache = torch.randn(2, state_capacity, heads, capacity, head_dim, device="cuda", dtype=q.dtype)
+    # Strided metadata is allowed by the kernel contract.
+    end = torch.zeros(state_capacity, 2, device="cuda", dtype=torch.long)[:, 0]
+    slots = torch.arange(batch * 2, device="cuda")[::2]
+    lengths = torch.full((batch, 2), frames, device="cuda", dtype=torch.int32)[:, 0]
+    context = capacity
+    attention = slot_ring_attention
+    if execution == "compiled":
+        # Functionalization must preserve mutations of the persistent ring
+        # and its offsets between calls, including slot reuse.
+        attention = torch.compile(slot_ring_attention, backend="inductor", fullgraph=True)
+    if execution == "graph":
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                slot_ring_attention(q, k, v, cache, end, slots, lengths, context)
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = slot_ring_attention(q, k, v, cache, end, slots, lengths, context)
+    end.zero_()
+    expected_cache, expected_end = cache.clone(), end.clone()
+    for step in range(8):
+        slots.copy_(torch.roll(torch.arange(state_capacity, device="cuda"), step)[:batch])
+        lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32))
+        projected.normal_()
+        if step == 4:
+            end[1:3].zero_()
+            expected_end[1:3].zero_()
+        expected = _ring_reference(q, k, v, expected_cache, expected_end, slots, lengths, context)
+        if execution == "graph":
+            graph.replay()
+        else:
+            actual = attention(q, k, v, cache, end, slots, lengths, context)
+        torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0)
+        torch.testing.assert_close(end, expected_end, rtol=0, atol=0)
+        torch.testing.assert_close(actual, expected, rtol=0.025, atol=0.025)
+        padding = torch.arange(frames, device="cuda")[None] >= lengths[:, None]
+        assert actual.permute(0, 2, 1, 3)[padding].count_nonzero() == 0
+
+
+@pytest.mark.parametrize("frames,context", [(1, 9), (3, 9), (15, 125), (480, 400)])
+def test_slot_attention_integrates_with_current_mha(frames, context):
+    torch.manual_seed(9)
+    reference = MossAudioTokenizerMultiheadAttention(
+        128,
+        2,
+        causal=True,
+        context=context,
+        rope=MossAudioTokenizerRotaryEmbedding(),
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    candidate = copy.deepcopy(reference)
+    candidate._slot_attention = slot_ring_attention
+    for model in (reference, candidate):
+        model._streaming_state = model._init_streaming_state(4)
+    for step in range(8):
+        slots = torch.tensor([1, 0, 3] if step % 2 else [0, 1, 3], device="cuda")
+        valid = torch.tensor([True, True, False], device="cuda")
+        execution_context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid)
+        x = torch.randn(3, frames, 128, device="cuda", dtype=torch.bfloat16)
+        with torch.inference_mode():
+            expected = reference(x, x, x, execution_context=execution_context)
+            actual = candidate(x, x, x, execution_context=execution_context)
+        torch.testing.assert_close(actual, expected, atol=0.01, rtol=0.02)
+        cs, rs = candidate._streaming_state, reference._streaming_state
+        torch.testing.assert_close(cs.offset, rs.offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.end_offset, rs.kv_cache.end_offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.cache[:, :3], rs.kv_cache.cache[:, :3], rtol=0, atol=0)
+        if step == 4:
+            for model in (reference, candidate):
+                model._streaming_state.reset_slots(torch.tensor([0], device="cuda"))
+
+
+@pytest.mark.parametrize(
+    "frames,capacity,head_dim", [(1, 125, 64), (15, 140, 64), (40, 125, 64), (480, 400, 64), (33, 65, 32)]
+)
+@pytest.mark.parametrize("execution", ["eager", "graph", "compiled"])
+def test_fused_rows_attention_matches_three_launch_op_exactly(frames, capacity, head_dim, execution):
+    torch.manual_seed(31)
+    batch, heads, state_capacity = 3, 2, 6
+    projected = torch.randn(batch, frames, 3, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0)
+    cache = torch.randn(2, state_capacity, heads, capacity, head_dim, device="cuda", dtype=q.dtype)
+    reference_cache = cache.clone()
+    end = torch.zeros(state_capacity, device="cuda", dtype=torch.long)
+    slots = torch.zeros(batch, device="cuda", dtype=torch.long)
+    lengths = torch.zeros(batch, device="cuda", dtype=torch.int32)
+    starts = torch.zeros(batch, device="cuda", dtype=torch.long)
+    attention = slot_ring_attention_rows
+    if execution == "compiled":
+        attention = torch.compile(slot_ring_attention_rows, backend="inductor", fullgraph=True)
+    if execution == "graph":
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                slot_ring_attention_rows(q, k, v, cache, starts, slots, lengths, capacity)
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = slot_ring_attention_rows(q, k, v, cache, starts, slots, lengths, capacity)
+        cache.copy_(reference_cache)
+    for step in range(8):
+        slots.copy_(torch.roll(torch.arange(state_capacity, device="cuda"), step)[:batch])
+        lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32))
+        starts.copy_(end[slots])
+        projected.normal_()
+        expected = slot_ring_attention(q, k, v, reference_cache, end, slots, lengths, capacity)
+        if execution == "graph":
+            graph.replay()
+        else:
+            actual = attention(q, k, v, cache, starts, slots, lengths, capacity)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        torch.testing.assert_close(cache, reference_cache, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("frames,context", [(1, 9), (15, 125), (480, 400)])
+def test_fused_rows_attention_integrates_with_mha_offsets(frames, context):
+    torch.manual_seed(11)
+    reference = MossAudioTokenizerMultiheadAttention(
+        128,
+        2,
+        causal=True,
+        context=context,
+        rope=MossAudioTokenizerRotaryEmbedding(),
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    candidate = copy.deepcopy(reference)
+    reference._slot_attention = slot_ring_attention
+    candidate._slot_attention_rows = slot_ring_attention_rows
+    for model in (reference, candidate):
+        model._streaming_state = model._init_streaming_state(4)
+    for step in range(8):
+        slots = torch.tensor([1, 0, 3] if step % 2 else [0, 1, 3], device="cuda")
+        valid = torch.tensor([True, True, False], device="cuda")
+        execution_context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid)
+        x = torch.randn(3, frames, 128, device="cuda", dtype=torch.bfloat16)
+        with torch.inference_mode():
+            expected = reference(x, x, x, execution_context=execution_context)
+            actual = candidate(x, x, x, execution_context=execution_context)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        cs, rs = candidate._streaming_state, reference._streaming_state
+        torch.testing.assert_close(cs.offset, rs.offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.end_offset, rs.kv_cache.end_offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.cache, rs.kv_cache.cache, rtol=0, atol=0)
+        if step == 4:
+            for model in (reference, candidate):
+                model._streaming_state.reset_slots(torch.tensor([0], device="cuda"))

tests/model_executor/models/moss_tts/test_streaming_decoder_wrapper_concurrency.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""A second streaming-decoder graph set can replay concurrently with the first."""
+
+import threading
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import (
+    CUDAGraphStreamingDecoderWrapper,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+CAP, DIM = 8, 64
+
+
+class _Codec(nn.Module):
+    def __init__(self, device):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(DIM, DIM, device=device) / DIM)
+        self.state = torch.zeros(CAP + 16, DIM, device=device)
+
+    def decode_streaming_tensors(self, codes, lengths, slots, valid_rows):
+        x = codes.float().sum(0)[..., None].expand(-1, -1, DIM) + self.state[slots][:, None]
+        y = x @ self.weight
+        self.state.index_copy_(0, slots, y[:, -1])
+        return y.transpose(1, 2).contiguous(), lengths
+
+    def reset_decoder_state_slots(self, slots):
+        self.state.index_fill_(0, slots.to(self.state.device), 0)
+
+
+def _wrapper(codec, **kwargs):
+    return CUDAGraphStreamingDecoderWrapper(
+        codec,
+        state_capacity=CAP,
+        num_quantizers=2,
+        vllm_config=None,
+        compiled_decode=codec.decode_streaming_tensors,
+        **kwargs,
+    )
+
+
+def test_private_wrapper_captures_on_its_own_stream_and_replays_concurrently():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    device = torch.device("cuda", torch.accelerator.current_device_index())
+    codec = _Codec(device)
+    main = _wrapper(codec, batch_sizes=[4], frame_sizes=[3])
+    fast = _wrapper(codec, batch_sizes=[1, 2], frame_sizes=[1], scratch_base=CAP + 4, private_pool=True)
+    main.warmup(device)
+    fast.warmup(device)
+    assert main.is_ready and fast.is_ready
+    assert main._capture_stream is None and fast._capture_stream is not None
+
+    codes_main = torch.randint(0, 5, (2, 4, 3), device=device)
+    codes_fast = torch.randint(0, 5, (2, 2, 1), device=device)
+    errors = []
+
+    def run(wrapper, codes, slots, iters):
+        stream = torch.cuda.Stream(device=device)
+        slot_ids = torch.tensor(slots, device=device)
+        try:
+            with torch.cuda.stream(stream):
+                for _ in range(iters):
+                    codec.reset_decoder_state_slots(slot_ids)
+                    audio = wrapper.decode(codes, slot_ids)[0][: len(slots)].clone()
+                    expected = (codes.float().sum(0)[..., None].expand(-1, -1, DIM) @ codec.weight).transpose(1, 2)
+                    stream.synchronize()
+                    torch.testing.assert_close(audio, expected, rtol=1e-3, atol=1e-3)
+        except Exception as error:  # surfaced below
+            errors.append(error)
+
+    threads = [
+        threading.Thread(target=run, args=(main, codes_main, [0, 1, 2, 3], 200)),
+        threading.Thread(target=run, args=(fast, codes_fast, [4, 5], 200)),
+    ]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join(timeout=60)
+    assert not any(thread.is_alive() for thread in threads), "concurrent graph replay did not finish"
+    assert not errors, errors

tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py

@@ -54,8 +54,19 @@ def test_terminal_padding_preserves_audio_live_state_and_reused_slots():
     assert [resumed[i].item() for i in (0, 1)] == [2, 36]
 
 
-@pytest.mark.parametrize("graph_capacity, expected_calls", [(None, 5), (4, 5), (8, 3)])
-def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capacity, expected_calls):
+@pytest.mark.parametrize(
+    "graph_capacity,expected_groups",
+    [
+        # No graph, or graphs that do not cover the state capacity: exact tails.
+        (None, {3: [0], 7: [1], 8: [2], 1: [3], 15: [4]}),
+        (4, {3: [0], 7: [1], 8: [2], 1: [3], 15: [4]}),
+        # Full graph coverage: finished tails (3, 7) join the padded 8-frame
+        # bucket with the live 8-frame row; the finished 1-frame row keeps the
+        # exact 1-frame graph and the live 15-frame row stays alone.
+        (8, {8: [0, 1, 2], 1: [3], 15: [4]}),
+    ],
+)
+def test_terminal_tails_coalesce_into_padded_frame_buckets(mocker, graph_capacity, expected_groups):
     s = session()
     s._cudagraph_wrapper = (
         None
@@ -81,9 +92,16 @@ def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capa
             for i, (length, done) in enumerate(zip(lengths, finished, strict=True))
         ]
     )
-    assert step.call_count == expected_calls
+    # Slots are leased in request order, so slot == row here.
+    groups = {}
+    for call in step.call_args_list:
+        plan = call.args[0]
+        groups[max(int(codes.shape[1]) for codes in plan.values())] = sorted(plan)
+        terminal = call.kwargs["terminal_slots"] or set()
+        assert terminal == {slot for slot in plan if finished[slot]}
+    assert groups == expected_groups
+    assert step.call_count == len(expected_groups)
+    # Output is always cropped to the exact frame count, padded or not.
     assert [result[i].shape[-1] for i in range(5)] == lengths
-    if graph_capacity == 8:
-        assert list(step.call_args_list[0].args[0]) == [0, 1, 2]
-        assert step.call_args_list[0].kwargs["terminal_slots"] == {0, 1}
-        assert list(step.call_args_list[1].args[0]) == [3]  # Preserve the exact one-frame graph.
+    # Finished rows released their slots; live rows keep theirs.
+    assert set(decoder._stream_req_slots) == {"2", "4"}

tests/model_executor/models/qwen3_tts/test_qwen3_tts_prompt_text_batch.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+    Qwen3TTSConfig,
+    Qwen3TTSTalkerConfig,
+)
+from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+    PRECOMPUTED_TEXT_IDS_KEY,
+    Qwen3TTSPromptEmbedsBuilder,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_batch_preprocess_projects_new_non_streaming_texts_once():
+    torch.manual_seed(0)
+    text_embedding = torch.nn.Embedding(64, 4)
+    projection = torch.nn.Linear(4, 4)
+    projected_shapes: list[tuple[int, ...]] = []
+
+    def text_projection(x: torch.Tensor) -> torch.Tensor:
+        projected_shapes.append(tuple(x.shape))
+        return projection(x)
+
+    builder = Qwen3TTSPromptEmbedsBuilder.__new__(Qwen3TTSPromptEmbedsBuilder)
+    builder._text_embedding = text_embedding
+    builder._text_projection = text_projection
+    builder._batched_text_embeds = {}
+
+    ids_a = list(range(1, 13))
+    ids_b = list(range(20, 31))
+    # Serving stores the assistant-template ids wrapped in a one-element list.
+    serving_ids = [ids_a]
+    buf: dict[str, dict[str, Any]] = {
+        "a": {"req_id": "a", "task_type": ["CustomVoice"], PRECOMPUTED_TEXT_IDS_KEY: [ids_a]},
+        "b": {"req_id": "b", "task_type": ["VoiceDesign"], PRECOMPUTED_TEXT_IDS_KEY: [torch.tensor(ids_b)]},
+        "streaming": {
+            "req_id": "streaming",
+            "task_type": ["CustomVoice"],
+            "non_streaming_mode": [False],
+            PRECOMPUTED_TEXT_IDS_KEY: serving_ids,
+        },
+        "base": {"req_id": "base", "task_type": ["Base"], PRECOMPUTED_TEXT_IDS_KEY: serving_ids},
+        "built": {
+            "req_id": "built",
+            "task_type": ["CustomVoice"],
+            "embed": {"prefill": torch.zeros(1)},
+            PRECOMPUTED_TEXT_IDS_KEY: serving_ids,
+        },
+    }
+
+    builder.preprocess_infos_batch(req_infos=list(buf.values()), device=torch.device("cpu"))
+
+    # One embedding + projection over both requests' text tokens (template stripped).
+    assert projected_shapes == [(1, (len(ids_a) - 8) + (len(ids_b) - 8), 4)]
+    assert set(builder._batched_text_embeds) == {"a", "b"}
+    with torch.no_grad():
+        for req_id, ids in (("a", ids_a), ("b", ids_b)):
+            expected = projection(text_embedding(torch.tensor([ids[3:-5]])))
+            torch.testing.assert_close(builder._batched_text_embeds[req_id], expected)
+            assert buf[req_id][PRECOMPUTED_TEXT_IDS_KEY].tolist() == [ids]
+    for skipped in ("streaming", "base", "built"):
+        assert buf[skipped][PRECOMPUTED_TEXT_IDS_KEY] is serving_ids
+
+
+class _Tokenizer:
+    def __call__(self, text: str, **kwargs: Any) -> dict[str, torch.Tensor]:
+        return {"input_ids": torch.tensor([[2, 3, 4, 5]])}
+
+
+def _prompt_builder() -> Qwen3TTSPromptEmbedsBuilder:
+    """Use the real constructor and prompt assembly with small embedding tables."""
+    torch.manual_seed(42)
+    config = Qwen3TTSConfig(tts_bos_token_id=50, tts_eos_token_id=51, tts_pad_token_id=52)
+    talker_config = Qwen3TTSTalkerConfig(
+        codec_nothink_id=10,
+        codec_think_id=11,
+        codec_think_bos_id=12,
+        codec_think_eos_id=13,
+        codec_pad_id=14,
+        codec_bos_id=15,
+        codec_language_id={"english": 16},
+        spk_id={"vivian": 17, "serena": 18},
+    )
+    builder = Qwen3TTSPromptEmbedsBuilder(
+        config=config,
+        talker_config=talker_config,
+        model_path="",
+        text_embedding=torch.nn.Embedding(64, 4),
+        text_projection=torch.nn.Linear(4, 4),
+        codec_embed=torch.nn.Embedding(64, 4),
+        residual_code_embeddings=lambda: [],
+        speaker_encoder=torch.nn.Identity(),
+        tts_pad_embed=torch.zeros(1, 4),
+        encode_ref_audio_batch=lambda *args, **kwargs: [],
+    )
+    builder._text_tokenizer = _Tokenizer()
+    return builder
+
+
+@pytest.mark.parametrize("task_type", ["CustomVoice", "VoiceDesign"])
+@pytest.mark.parametrize("non_streaming", [False, True])
+def test_batched_prompt_consumption_matches_serial_and_clears_cache(task_type: str, non_streaming: bool):
+    batched, serial = _prompt_builder(), _prompt_builder()
+
+    def request(req_id: str, ids: list[int], speaker: str) -> dict[str, Any]:
+        return {
+            "req_id": req_id,
+            "text": ["hello"],
+            "task_type": [task_type],
+            "language": ["English"],
+            "speaker": [speaker],
+            "instruct": ["Speak calmly"],
+            "non_streaming_mode": [non_streaming],
+            PRECOMPUTED_TEXT_IDS_KEY: [ids],
+        }
+
+    with torch.inference_mode():
+        # The same request id can be used again after its earlier prompt was consumed.
+        for ids_a, ids_b in [(list(range(12)), list(range(20, 31))), (list(range(10, 24)), list(range(30, 46)))]:
+            pending = [request("a", ids_a, "Vivian"), request("b", ids_b, "Serena")]
+            expected = [request("a", ids_a, "Vivian"), request("b", ids_b, "Serena")]
+            batched.preprocess_infos_batch(req_infos=pending, device=torch.device("cpu"))
+            assert set(batched._batched_text_embeds) == ({"a", "b"} if non_streaming else set())
+            for actual_info, expected_info in zip(pending, expected, strict=True):
+                actual = batched.build_prompt_embeds(task_type=task_type, info_dict=actual_info)
+                reference = serial.build_prompt_embeds(task_type=task_type, info_dict=expected_info)
+                torch.testing.assert_close(actual[0], reference[0])
+                torch.testing.assert_close(actual[1], reference[1])
+                assert actual[2:] == reference[2:]
+                assert PRECOMPUTED_TEXT_IDS_KEY not in actual_info
+            assert not batched._batched_text_embeds
+        batched.preprocess_infos_batch(req_infos=[], device=torch.device("cpu"))
+        assert not batched._batched_text_embeds

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -13,6 +13,7 @@ from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
     text2flow_full_payload,
     text2flow_token_only,
 )
+from vllm_omni.outputs.mm_outputs import MultimodalPayload
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -121,6 +122,24 @@ def test_text2flow_full_payload_does_not_send_codec_ids():
     assert torch.equal(payload["embed"]["speech_token"], torch.tensor([[1, 2]], dtype=torch.long))
 
 
+def test_text2flow_full_payload_accepts_structured_multimodal_payload():
+    speech_token = torch.tensor([[1, 2]], dtype=torch.long)
+    payload = text2flow_full_payload(
+        None,
+        MultimodalPayload(
+            tensors={
+                "embed.speech_token": speech_token,
+                "codes.audio": torch.tensor([7, 8, 9], dtype=torch.long),
+            }
+        ),
+        SimpleNamespace(external_req_id="req-structured"),
+    )
+
+    assert payload is not None
+    assert "codes" not in payload
+    assert torch.equal(payload["embed"]["speech_token"], speech_token)
+
+
 def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
     transfer_manager = _transfer_manager()
     request = SimpleNamespace(
@@ -355,3 +374,36 @@ def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio(
     assert payload_final is not None
     assert payload_final.meta.finished.item() is True
     assert payload_final.codes.audio.tolist() == []
+
+
+def test_full_payload_retains_true_length_for_padded_conditioning():
+    from vllm_omni.data_entry_keys import to_struct
+
+    token = torch.tensor([[7, 8, 0, 0]], dtype=torch.int32)
+    feat = torch.randn(1, 8, 80)
+    speaker = torch.randn(1, 192)
+    length = torch.tensor([[2]], dtype=torch.int32)
+    result = text2flow_full_payload(
+        None,
+        {
+            "embed.speech_token": token,
+            "embed.speech_feat": feat,
+            "embed.embedding": speaker,
+            "embed.speech_token_len": length,
+        },
+        SimpleNamespace(request_id="padded"),
+    )
+    payload = to_struct(result)
+    assert payload.embed.speech_token_len is length
+
+
+def test_reference_features_are_snapshots_under_output_accumulation():
+    from vllm_omni.outputs.mm_outputs import MultimodalPayload
+    from vllm_omni.outputs.output_modality import OutputModality
+
+    first = torch.ones(1, 174, 80)
+    latest = torch.zeros(1, 1500, 80)
+    payload = MultimodalPayload.from_dict({"embed.speech_feat": first})
+    payload = payload.merged_with(MultimodalPayload.from_dict({"embed.speech_feat": latest}))
+    payload.consolidate_tensors(OutputModality.LATENT)
+    assert payload["embed.speech_feat"] is latest

tests/model_executor/stage_input_processors/test_moss_raw_async_chunk.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""MOSS raw codec rows are buffered into first/regular chunks in codebook-major order."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.moss_tts import (
+    _MOSS_AUDIO_PAD_CODE,
+    talker2codec_raw_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _manager(first=1, regular=3):
+    config = {"extra": {"codec_chunk_frames": regular, "initial_codec_chunk_frames": first}}
+    return SimpleNamespace(connector=SimpleNamespace(config=config))
+
+
+def _step(manager, frames, finished=False):
+    request = SimpleNamespace(request_id="r", external_req_id="r")
+    payload = talker2codec_raw_async_chunk(manager, {"codes": {"audio": frames}}, request, is_finished=finished)
+    if payload is not None and payload.codes is not None:
+        manager.put_req_chunk["r"] += 1
+        manager.ramp_chunk_count["r"] += 1
+    return payload
+
+
+def test_first_regular_and_terminal_chunks():
+    manager = _manager()
+    rows = torch.arange(6 * 4, dtype=torch.int32).reshape(6, 4)
+    first = _step(manager, rows[:1])
+    assert first.meta.codec_chunk_frames == 1
+    assert first.codes.audio.dtype == torch.int64
+    assert torch.equal(first.codes.audio, rows[0].long())
+
+    assert _step(manager, rows[1:2]) is None
+    padded = torch.full((1, 4), _MOSS_AUDIO_PAD_CODE, dtype=torch.int32)
+    assert _step(manager, padded) is None  # PAD rows are not audio
+    assert _step(manager, rows[2:3]) is None
+    regular = _step(manager, rows[3:4])
+    assert regular.meta.codec_chunk_frames == 3
+    assert torch.equal(regular.codes.audio, rows[1:4].long().T.reshape(-1))
+
+    tail = _step(manager, rows[4:6], finished=True)
+    assert tail.meta.codec_chunk_frames == 2
+    assert torch.equal(tail.codes.audio, rows[4:6].long().T.reshape(-1))
+    assert bool(tail.meta.finished)
+    assert "r" not in manager.code_prompt_token_ids
+
+
+def test_buffered_rows_do_not_alias_producer_storage():
+    manager = _manager(first=2, regular=2)
+    frames = torch.tensor([[7, 8, 9]])
+    assert _step(manager, frames) is None
+    frames.fill_(0)
+    chunk = _step(manager, torch.tensor([[1, 2, 3]]))
+    assert torch.equal(chunk.codes.audio, torch.tensor([7, 1, 8, 2, 9, 3]))
+
+
+@pytest.mark.parametrize("total", [1, 3, 5, 10, 13, 28, 31])
+def test_ramp_preserves_every_frame_and_terminal(total):
+    manager = _manager(first=1, regular=15)
+    manager.connector.config["extra"]["codec_chunk_ramp"] = [1, 4, 8, 15]
+    rows = torch.arange(total * 4, dtype=torch.int32).reshape(total, 4)
+    packets = []
+    for index in range(total):
+        packet = _step(manager, rows[index : index + 1], finished=index == total - 1)
+        if packet is not None:
+            packets.append(packet)
+    rebuilt = torch.cat([packet.codes.audio.reshape(4, -1).T for packet in packets])
+    assert torch.equal(rebuilt, rows.long())
+    assert all(not bool(packet.meta.finished) for packet in packets[:-1])
+    assert bool(packets[-1].meta.finished)
+    assert "r" not in manager.code_prompt_token_ids
+    if total == 31:
+        assert [packet.meta.codec_chunk_frames for packet in packets] == [1, 4, 8, 15, 3]

tests/model_executor/stage_input_processors/test_moss_raw_empty_terminal.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Empty MOSS terminals must not execute prewarmed or previous codec tokens."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.core.sched.omni_scheduling_coordinator import OmniSchedulingCoordinator
+from vllm_omni.data_entry_keys import to_dict
+from vllm_omni.distributed.omni_connectors.model_runner.omni_connector_payload_transport import (
+    _OmniConnectorPayloadTransportMixin as Transport,
+)
+from vllm_omni.model_executor.stage_input_processors.moss_tts import (
+    _MOSS_AUDIO_PAD_CODE,
+    talker2codec_raw_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _request(tokens):
+    return SimpleNamespace(
+        request_id="r",
+        external_req_id="r",
+        prompt_token_ids=list(tokens),
+        num_prompt_tokens=len(tokens),
+        _all_token_ids=list(tokens),
+        _output_token_ids=[],
+        num_computed_tokens=0,
+    )
+
+
+@pytest.mark.parametrize("after_audio", [False, True])
+def test_empty_terminal_clears_native_codec_prompt(after_audio):
+    manager = SimpleNamespace(
+        connector=SimpleNamespace(config={"extra": {"initial_codec_chunk_frames": 1, "codec_chunk_frames": 3}})
+    )
+    source = _request([0] * 88)
+    previous = [0] * 88
+    if after_audio:
+        packet = talker2codec_raw_async_chunk(manager, {"codes": {"audio": torch.ones(1, 12)}}, source)
+        manager.put_req_chunk["r"] += 1
+        previous = packet.codes.audio.tolist()
+    terminal = talker2codec_raw_async_chunk(
+        manager,
+        {"codes": {"audio": torch.full((1, 12), _MOSS_AUDIO_PAD_CODE)}},
+        source,
+        is_finished=True,
+    )
+    payload = to_dict(terminal)
+    assert not Transport._payload_is_consumable(payload)
+    metadata = Transport._extract_scheduling_metadata(payload)
+    receiver = _request(previous)
+    coordinator = OmniSchedulingCoordinator(stage_id=1)
+    coordinator.update_request_metadata({"r": receiver}, {"r": metadata}, model_mode="generation")
+    assert receiver.prompt_token_ids == []
+    assert receiver._all_token_ids == []
+    assert receiver.num_prompt_tokens == receiver.num_computed_tokens == 0
+    assert "r" in coordinator.input_terminal_req_ids
+    assert "r" not in manager.code_prompt_token_ids
+
+
+@pytest.mark.parametrize("empty", [[], torch.empty(0, dtype=torch.long)])
+def test_explicit_empty_codes_clear_generation_prompt(empty):
+    request = _request([17, 18, 19])
+    request._output_token_ids = [23]
+    request.num_computed_tokens = 3
+    coordinator = OmniSchedulingCoordinator(stage_id=1)
+    coordinator.update_request_metadata(
+        {"r": request}, {"r": {"code_predictor_codes": empty, "input_terminal": True}}, model_mode="generation"
+    )
+    assert request.prompt_token_ids == request._all_token_ids == request._output_token_ids == []
+    assert request.num_prompt_tokens == request.num_computed_tokens == 0
+    assert "r" in coordinator.input_terminal_req_ids
+
+
+@pytest.mark.parametrize("metadata", [{"input_terminal": True}, {"code_predictor_codes": None}])
+def test_absent_codes_preserve_generation_prompt(metadata):
+    request = _request([17, 18, 19])
+    coordinator = OmniSchedulingCoordinator(stage_id=1)
+    coordinator.update_request_metadata({"r": request}, {"r": metadata}, model_mode="generation")
+    assert request.prompt_token_ids == request._all_token_ids == [17, 18, 19]
+
+
+def test_ar_prompt_is_not_replaced_by_empty_audio_codes():
+    request = _request([17, 18, 19])
+    coordinator = OmniSchedulingCoordinator(stage_id=1)
+    coordinator.update_request_metadata(
+        {"r": request}, {"r": {"code_predictor_codes": [], "input_terminal": True}}, model_mode="ar"
+    )
+    assert request.prompt_token_ids == request._all_token_ids == [17, 18, 19]

tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Progressive Local codec chunks preserve frames and agree with graph shapes."""
+
+from collections import defaultdict
+from dataclasses import dataclass, field
+from unittest.mock import Mock
+
+import pytest
+import torch
+from transformers import PretrainedConfig
+from vllm.config import CompilationConfig, ModelConfig, SchedulerConfig, VllmConfig
+
+from vllm_omni.data_entry_keys import OmniPayloadStruct
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder
+from vllm_omni.model_executor.stage_input_processors.moss_tts import talker2codec_raw_async_chunk
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclass
+class _Connector:
+    config: dict[str, dict[str, object]]
+
+
+@dataclass
+class _TransferManager:
+    connector: _Connector
+    code_prompt_token_ids: dict[str, list[object]] = field(default_factory=lambda: defaultdict(list))
+    request_payload: dict[str, object] = field(default_factory=dict)
+    put_req_chunk: dict[str, int] = field(default_factory=lambda: defaultdict(int))
+    ramp_chunk_count: dict[str, int] = field(default_factory=lambda: defaultdict(int))
+
+
+@dataclass
+class _Request:
+    external_req_id: str
+
+
+def _manager(ramp: object = None, *, first: int = 1) -> _TransferManager:
+    extra = {"initial_codec_chunk_frames": first, "codec_chunk_frames": 15, "codec_chunk_ramp": ramp}
+    return _TransferManager(_Connector({"extra": extra}))
+
+
+def _emit(
+    manager: _TransferManager, req_id: str, value: int | None, finished: bool = False
+) -> OmniPayloadStruct | None:
+    output = None if value is None else {"codes": {"audio": torch.full((1, 4), value)}}
+    payload = talker2codec_raw_async_chunk(manager, output, _Request(req_id), finished)
+    if payload is not None:
+        manager.put_req_chunk[req_id] += 1
+        manager.ramp_chunk_count[req_id] += 1
+    return payload
+
+
+@pytest.mark.parametrize(
+    "ramp,first,expected",
+    [
+        (None, 1, [1, 15, 15, 9]),
+        ([1, 2, 4, 8, 15], 1, [1, 2, 4, 8, 15, 10]),
+        ([1, 4, 15], 7, [1, 4, 15, 15, 5]),
+        ([4, 8, 15], 1, [4, 8, 15, 13]),
+        ([1, 20, 15], 1, [1, 20, 15, 4]),
+        ("1,4,15", 1, [1, 4, 15, 15, 5]),
+        ([1, 0, 15], 1, [1, 15, 15, 9]),
+        ([1], 1, [1, 15, 15, 9]),
+        ("invalid", 1, [1, 15, 15, 9]),
+    ],
+)
+def test_order_and_final_flush(ramp: object, first: int, expected: list[int]) -> None:
+    manager = _manager(ramp, first=first)
+    packets = []
+    for index in range(40):
+        packet = _emit(manager, "a", index, finished=index == 39)
+        if packet is not None:
+            packets.append(packet)
+    assert [p.meta.codec_chunk_frames for p in packets] == expected
+    actual = torch.cat([p.codes.audio.reshape(4, -1).T for p in packets])
+    assert torch.equal(actual, torch.arange(40).unsqueeze(1).expand(-1, 4))
+    assert all(not bool(p.meta.finished) for p in packets[:-1])
+    assert bool(packets[-1].meta.finished)
+    assert "a" not in manager.code_prompt_token_ids
+
+
+def test_requests_progress_independently_and_empty_finish() -> None:
+    manager = _manager([1, 2, 4, 8, 15])
+    for name in ("a", "b"):
+        first = _emit(manager, name, 0)
+        assert first is not None
+        assert first.meta.codec_chunk_frames == 1
+    assert _emit(manager, "a", 1) is None
+    final_a = _emit(manager, "a", None, finished=True)
+    assert final_a is not None
+    assert final_a.meta.codec_chunk_frames == 1
+    assert bool(final_a.meta.finished)
+    final_b = _emit(manager, "b", None, finished=True)
+    assert final_b is not None
+    assert final_b.meta.codec_chunk_frames == 0
+    assert bool(final_b.meta.finished)
+    assert final_b.codes.audio.numel() == 0
+    assert not manager.code_prompt_token_ids
+
+
+@pytest.mark.parametrize(
+    "ramp,first,expected,max_step,effective_first",
+    [
+        (None, 1, [1, 15], 15, 1),
+        ([1, 2, 4, 8, 15], 1, [1, 2, 4, 8, 15], 15, 1),
+        ([1, 20, 15], 1, [1, 15, 20], 20, 1),
+        ([4, 8, 15], 1, [4, 8, 15], 15, 4),
+        ([1, 4, 15], 7, [1, 4, 15], 15, 1),
+        ([1, 0, 15], 1, [1, 15], 15, 1),
+    ],
+)
+def test_graph_shapes_and_fast_path_follow_sender(
+    ramp: object, first: int, expected: list[int], max_step: int, effective_first: int
+) -> None:
+    connector = _manager(ramp, first=first).connector.config
+    connector["extra"]["codec_first_chunk_fast_path"] = 1
+    config = Mock(
+        spec=VllmConfig,
+        model_config=Mock(
+            spec=ModelConfig,
+            hf_config=PretrainedConfig(model_type="moss_tts_local"),
+            async_chunk=True,
+            use_v2_model_runner=True,
+            enforce_eager=False,
+            stage_connector_config=connector,
+        ),
+        scheduler_config=Mock(spec=SchedulerConfig, max_num_seqs=8),
+        compilation_config=Mock(spec=CompilationConfig, cudagraph_capture_sizes=[1, 2, 4, 8]),
+    )
+    codec = MossTTSCodecDecoder(vllm_config=config)
+    assert codec._streaming_graph_frame_sizes == expected
+    assert codec._stream_max_step_frames == max_step
+    assert codec._first_chunk_fast
+    assert codec._initial_stream_chunk_frames == effective_first

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -971,12 +971,14 @@ def test_cosyvoice3_text2flow_full_payload_nested_fallback() -> None:
 
 
 def test_cosyvoice3_full_payload_replace_keys_present() -> None:
-    """Confirm _FULL_PAYLOAD_REPLACE_KEYS lists the three embed.* keys."""
+    """Reference tensors and their length travel as complete snapshots."""
     from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
         _FULL_PAYLOAD_REPLACE_KEYS,
     )
 
-    assert _FULL_PAYLOAD_REPLACE_KEYS == frozenset({"embed.speech_token", "embed.speech_feat", "embed.embedding"})
+    assert _FULL_PAYLOAD_REPLACE_KEYS == frozenset(
+        {"embed.speech_token", "embed.speech_token_len", "embed.speech_feat", "embed.embedding"}
+    )
 
 
 def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None:

tests/model_executor/test_seeded_exponential_platform.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.utils import seeded_exponential
+
+pytestmark = [pytest.mark.core_model]
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("is_nvidia", [False, True], ids=["rocm", "cuda"])
+def test_cuda_tensor_requires_nvidia_distribution_contract(monkeypatch, is_nvidia: bool) -> None:
+    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: is_nvidia)
+    # HIP tensors expose the same Torch device predicate as NVIDIA tensors.
+    q = Mock(spec=torch.Tensor)
+    q.is_cuda = True
+    q.dtype = torch.float32
+    q.dim.return_value = 2
+    q.is_contiguous.return_value = True
+    generators = {0: object(), 1: object()}
+
+    assert seeded_exponential.batched_seeded_exponential_supported(q, generators) is is_nvidia
+
+
+@pytest.mark.cpu
+def test_shared_generator_keeps_ordered_torch_draws(monkeypatch) -> None:
+    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: True)
+    q = Mock(spec=torch.Tensor)
+    q.is_cuda = True
+    q.dtype = torch.float32
+    q.dim.return_value = 2
+    q.is_contiguous.return_value = True
+    shared = object()
+
+    assert not seeded_exponential.batched_seeded_exponential_supported(q, {0: shared, 1: shared})
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+@pytest.mark.parametrize("seeded_rows", [(0, 1, 2), (0, 2)], ids=["all_seeded", "mixed"])
+@pytest.mark.parametrize("use_fp64_gumbel", [False, True])
+def test_torch_fallback_preserves_samples_and_generator_state(monkeypatch, seeded_rows, use_fp64_gumbel) -> None:
+    from vllm.v1.sample.ops import topk_topp_sampler as sampler_ops
+
+    import vllm_omni.patch  # noqa: F401
+
+    original = sampler_ops.random_sample.__wrapped__
+    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: False)
+
+    def reject_kernel(*args, **kwargs):
+        raise AssertionError("the CUDA distribution kernel must not run on the fallback path")
+
+    monkeypatch.setattr(seeded_exponential, "fill_exponential_rows", reject_kernel)
+    probs = torch.arange(1, 258, device="cuda", dtype=torch.float32).expand(3, -1).contiguous()
+    probs /= probs.sum(dim=-1, keepdim=True)
+    expected_gens = {row: torch.Generator(device="cuda").manual_seed(42 + row) for row in seeded_rows}
+    actual_gens = {row: torch.Generator(device="cuda").manual_seed(42 + row) for row in seeded_rows}
+    expected_default = torch.cuda.get_rng_state()
+    actual_default = expected_default.clone()
+    try:
+        for _ in range(5):
+            torch.cuda.set_rng_state(expected_default)
+            # FP32 sampling divides probabilities in place. Give both paths
+            # the same untouched input on every step.
+            expected = original(probs.clone(), expected_gens, use_fp64_gumbel)
+            expected_default = torch.cuda.get_rng_state()
+            torch.cuda.set_rng_state(actual_default)
+            actual = sampler_ops.random_sample(probs.clone(), actual_gens, use_fp64_gumbel)
+            actual_default = torch.cuda.get_rng_state()
+            assert torch.equal(actual, expected)
+            assert torch.equal(actual_default, expected_default)
+            for row in seeded_rows:
+                assert torch.equal(actual_gens[row].get_state(), expected_gens[row].get_state())
+    finally:
+        torch.cuda.set_rng_state(expected_default)
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+@pytest.mark.parametrize("rows", [None, [0, 2, 4]], ids=["grouped_codebooks", "selected_rows"])
+def test_direct_fill_fallback_preserves_ordered_draws(monkeypatch, rows) -> None:
+    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: False)
+
+    def reject_kernel(*args, **kwargs):
+        raise AssertionError("the NVIDIA kernel must not run for direct fallback callers")
+
+    monkeypatch.setattr(seeded_exponential, "torch_exponential_policy", reject_kernel)
+    expected = torch.ones(6, 257, device="cuda")
+    actual = expected.clone()
+    expected_gen = torch.Generator(device="cuda").manual_seed(71)
+    actual_gen = torch.Generator(device="cuda").manual_seed(71)
+    # A shared seeded generator and an unseeded request make loop order and
+    # default-state advancement observable, including grouped codebooks.
+    expected_gens = [expected_gen, None, expected_gen]
+    actual_gens = [actual_gen, None, actual_gen]
+    original_default = torch.cuda.get_rng_state()
+    expected_default = original_default.clone()
+    actual_default = original_default.clone()
+    try:
+        for _ in range(5):
+            torch.cuda.set_rng_state(expected_default)
+            targets = expected.reshape(3, -1) if rows is None else [expected[row] for row in rows]
+            for target, generator in zip(targets, expected_gens):
+                target.exponential_(generator=generator)
+            expected_default = torch.cuda.get_rng_state()
+            torch.cuda.set_rng_state(actual_default)
+            assert seeded_exponential.fill_exponential_rows(actual, actual_gens, rows) is actual
+            actual_default = torch.cuda.get_rng_state()
+            assert torch.equal(actual, expected)
+            assert torch.equal(actual_gen.get_state(), expected_gen.get_state())
+            assert torch.equal(actual_default, expected_default)
+    finally:
+        torch.cuda.set_rng_state(original_default)

tests/worker/test_omni_connector_mixin.py

@@ -39,6 +39,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class MockConnector:
+    sender_host: str | None = None
+    sender_zmq_port: int | None = None
+
     """In-memory connector for testing (mimics OmniConnectorBase)."""
 
     def __init__(self, stage_id: int = 0):
@@ -187,7 +190,7 @@ def test_init_payload_connector_ownership(role, custom_func, expected):
 
 @pytest.mark.parametrize("async_chunk", [False, True])
 @pytest.mark.parametrize("rank", [0, 1])
-def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank):
+def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank, mocker):
     host = MixinHost()
     host.init_omni_connectors(_make_model_config(async_chunk=async_chunk))
     # No background threads: deterministic interleaved polling.
@@ -196,7 +199,7 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank):
     connector = host._omni_connector
     connector.sender_host = "configured-host"
     connector.sender_zmq_port = 50051
-    connector.get = MagicMock(return_value=None)
+    get = mocker.patch.object(connector, "get", return_value=None)
     first = _make_request("r1", "external1")
     first.payload_sender_info = {"host": "producer-a", "zmq_port": "50101"}
     second = _make_request("r2", "external2")
@@ -209,14 +212,14 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank):
             for req_id in ("r2", "r1", "r2"):
                 assert not host._poll_single_request(req_id)
             if rank == 1:
-                connector.get.assert_not_called()
+                get.assert_not_called()
             else:
-                assert connector.get.call_args_list == [
+                assert get.call_args_list == [
                     unittest.mock.call("0", "1", "external2_0_0", {"source_host": "producer-b", "source_port": 51101}),
                     unittest.mock.call("0", "1", "external1_0_0", {"source_host": "producer-a", "source_port": 50101}),
                     unittest.mock.call("0", "1", "external2_0_0", {"source_host": "producer-b", "source_port": 51101}),
                 ]
-                connector.get.return_value = ({"ids": {"output": [1]}, "meta": {"finished": True}}, 1)
+                get.return_value = ({"ids": {"output": [1]}, "meta": {"finished": True}}, 1)
                 assert host._poll_single_request("r2")
                 assert "r2" not in host._pending_load_reqs
                 assert "r1" in host._pending_load_reqs
@@ -225,14 +228,14 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank):
             assert not host._pending_load_reqs
             # Reused internal ID without sender info cannot inherit either endpoint.
             replacement = _make_request("r1", "external-new")
-            connector.get.reset_mock(return_value=True)
-            connector.get.return_value = None
+            get.reset_mock(return_value=True)
+            get.return_value = None
             host.register_chunk_recv(replacement)
             assert not host._poll_single_request("r1")
             if rank == 0:
-                connector.get.assert_called_once_with("0", "1", "external-new_0_0")
+                get.assert_called_once_with("0", "1", "external-new_0_0")
             else:
-                connector.get.assert_not_called()
+                get.assert_not_called()
             assert connector.sender_host == "configured-host"
             assert connector.sender_zmq_port == 50051
             assert first.payload_sender_info == {"host": "producer-a", "zmq_port": "50101"}
@@ -1437,6 +1440,58 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
 
         host.shutdown_omni_connectors()
 
+    def test_first_chunk_hook_sees_only_chunk_zero_before_staging(self):
+        host = MixinHost()
+        host.init_omni_connectors(
+            model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="gen"),
+        )
+        host._omni_connector = MagicMock()
+        host._stage_id = 1
+        host._async_chunk = True
+        host._model_mode = "gen"
+        seen = []
+
+        def hook(req_id, handle, payload):
+            # The chunk is not yet visible to the scheduler when the hook runs.
+            seen.append((req_id, handle.external_req_id, payload["codes"]["audio"], "r1" in host._finished_load_reqs))
+
+        host.set_first_chunk_hook(hook)
+        host.register_chunk_recv(_make_request("r1", "ext-r1"))
+        chunk = {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(False)}}
+        host._omni_connector.get.return_value = (chunk, 1)
+        self.assertTrue(host._poll_single_request("r1"))
+        self.assertEqual(seen, [("r1", "ext-r1", [1, 2], False)])
+        self.assertIn("r1", host._finished_load_reqs)
+
+        host.get_omni_connector_output()
+        host._local_stage_payload_cache.clear()
+        host._local_request_metadata.clear()
+        host._omni_connector.get.return_value = ({"codes": {"audio": [3, 4]}, "meta": {}}, 1)
+        self.assertTrue(host._poll_single_request("r1"))
+        self.assertEqual(len(seen), 1)
+
+        host.shutdown_omni_connectors()
+
+    def test_first_chunk_hook_skips_terminal_first_chunk(self):
+        host = MixinHost()
+        host.init_omni_connectors(
+            model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="gen"),
+        )
+        host._omni_connector = MagicMock()
+        host._stage_id = 1
+        host._async_chunk = True
+        host._model_mode = "gen"
+        hook = MagicMock()
+        host.set_first_chunk_hook(hook)
+        host.register_chunk_recv(_make_request("r1"))
+        host._omni_connector.get.return_value = (
+            {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(True)}},
+            1,
+        )
+        self.assertTrue(host._poll_single_request("r1"))
+        hook.assert_not_called()
+        host.shutdown_omni_connectors()
+
     def test_non_ar_recv_waits_for_scheduler_handoff_before_fetching_next_chunk(self):
         host = MixinHost()
         host.init_omni_connectors(
@@ -1480,6 +1535,39 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
         host.shutdown_omni_connectors()
 
 
+@pytest.mark.parametrize("cancel_at", ["connector_get", "first_chunk_hook"])
+def test_cancelled_chunk_receive_does_not_recreate_delivery_state(mocker, cancel_at):
+    host = MixinHost()
+    host.init_omni_connectors(_make_model_config(async_chunk=True, worker_type="gen"))
+    host._stage_id = 1
+    host._omni_connector = mocker.MagicMock()
+    host.register_chunk_recv(_make_request("r1"))
+    seen = []
+    payload = {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(False)}}
+
+    def get(*args, **kwargs):
+        if cancel_at == "connector_get":
+            host.cleanup_finished_request("r1")
+        return payload, 1
+
+    def hook(req_id, handle, chunk):
+        seen.append(req_id)
+        host.cleanup_finished_request(req_id)
+
+    host._omni_connector.get.side_effect = get
+    host.set_first_chunk_hook(hook)
+    try:
+        assert not host._poll_single_request("r1")
+        assert seen == (["r1"] if cancel_at == "first_chunk_hook" else [])
+        assert host.get_local_stage_payload("r1") is None
+        assert "r1" not in host._get_req_chunk
+        assert "r1" not in host._pending_load_reqs
+        output = host.get_omni_connector_output()
+        assert output.chunk_ready_req_ids == set() and output.request_metadata == {}
+    finally:
+        host.shutdown_omni_connectors()
+
+
 class TestRankAwareKVRouting(unittest.TestCase):
     def _make_host(self, *, from_tp: int, to_tp: int, local_rank: int) -> MixinHost:
         host = MixinHost()

tests/worker_v2/test_first_audio_sender.py

@@ -235,3 +235,19 @@ def test_sender_copy_failure_reports_all_prepared_routes():
     assert client == 3
     assert [(o.request_id, o.finish_reason) for o in outputs.outputs] == [("r", FinishReason.ERROR)]
     assert output_queue.empty()
+
+
+def test_final_codec_delivery_omits_upstream_marker_and_freezes_route():
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
+
+    outputs: queue.Queue = queue.Queue()
+    scheduler = _scheduler(r=3)
+    delivery = engine_output_queue_sink(outputs, scheduler, upstream_first_audio=False).prepare(["r"])
+    scheduler.requests.clear()
+    pcm = torch.ones(4)
+    delivery(["r"], [pcm], torch.tensor(24000))
+    client, batch = outputs.get_nowait()
+    assert client == 3
+    assert FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output
+    assert batch.outputs[0].multimodal_output["model_outputs"] is pcm
+    assert batch.outputs[0].finish_reason is None

tests/worker_v2/test_omni_ar_model_runner.py

@@ -15,6 +15,7 @@ from vllm.v1.outputs import RoutedExpertsTensors
 from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors
 
 import vllm_omni.worker_v2.omni_ar_model_runner as omni_ar_model_runner
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
 from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner, OmniAsyncOutput
 from vllm_omni.worker_v2.output_snapshot import pack_output_snapshot
 
@@ -144,6 +145,7 @@ def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None:
     runner.model = SimpleNamespace()
     runner._async_mm_snapshot_slots, runner._async_mm_snapshot_events = [{}], [None]
     runner._async_mm_snapshot_pending, runner._async_mm_snapshot_cursor = [False], 0
+    runner._last_multimodal_snapshot_slot = None
     waited: list[object] = []
     runner.main_stream = SimpleNamespace(wait_event=waited.append)
     source = torch.tensor([[7, 8]], dtype=torch.long)
@@ -161,6 +163,35 @@ def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None:
     assert waited == [copy_event]
 
 
+def test_producer_snapshot_is_not_repacked() -> None:
+    runner = OmniARModelRunner.__new__(OmniARModelRunner)
+    runner.model_config = SimpleNamespace(async_chunk=True)
+    runner.model = SimpleNamespace()
+    runner._async_mm_snapshot_slots, runner._async_mm_snapshot_events = [{}], [None]
+    runner._async_mm_snapshot_pending, runner._async_mm_snapshot_cursor = [False], 0
+    runner._last_multimodal_snapshot_slot = None
+    source = torch.tensor([[7, 8]], dtype=torch.float32)
+    snapshot = pack_output_snapshot({"audio": source}, {}, max_buckets=1)
+    assert isinstance(snapshot, PackedOutputSnapshot)
+
+    retained = runner._retain_multimodal_outputs(snapshot)
+
+    assert retained is snapshot
+    assert runner._last_multimodal_snapshot_slot is None
+    assert runner._async_mm_snapshot_pending == [False]
+
+
+def test_fresh_per_step_outputs_are_not_repacked() -> None:
+    runner = OmniARModelRunner.__new__(OmniARModelRunner)
+    runner.model_config = SimpleNamespace(async_chunk=True)
+    runner.model = SimpleNamespace(mm_outputs_fresh_per_step=True)
+    runner._last_multimodal_snapshot_slot = None
+    outputs = {"model_outputs": torch.ones(2, 4)}
+
+    assert runner._retain_multimodal_outputs(outputs) is outputs
+    assert runner._last_multimodal_snapshot_slot is None
+
+
 def test_snapshot_slots_bounded_by_shape_and_packed_grouping_isolation(monkeypatch) -> None:
     slot: dict[tuple[Any, ...], torch.Tensor] = {}
     omni_ar_model_runner._copy_mm_to_snapshot_slot(torch.ones(1, 2), slot)

tests/worker_v2/test_omni_model_state.py

@@ -24,6 +24,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 class _DummyInputBatch:
     is_prefilling_np: np.ndarray
+    req_ids: list[str]
 
     input_ids: SimpleNamespace
     query_start_loc: torch.Tensor
@@ -224,6 +225,61 @@ def test_static_decode_embeddings_refresh_from_input_ids():
     assert OmniModelState._preprocess_result_needs_writeback(original, original.view_as(original)) is True
 
 
+def test_moss_local_decode_runs_depth_predictor_and_routes_eos(mocker):
+    """Exercise MRV2 dispatch through the real Local hook, output and logits."""
+    from vllm.sampling_params import SamplingParams
+    from vllm.v1.worker.gpu.states import RequestState
+
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import MossTTSLocalDepthTransformer
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSLocalTalkerForGeneration,
+    )
+
+    model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+    torch.nn.Module.__init__(model)
+    model.model = torch.nn.Module()
+    model.model.embed_tokens = torch.nn.Embedding(8, 4, _weight=torch.zeros(8, 4))
+    model.hidden_size = 4
+    model.n_vq = 2
+    model.audio_pad_token_id = 8
+    model.audio_assistant_slot_token_id = 2
+    model.im_end_token_id = 3
+    model.text_vocab_size = 8
+    model.talker_mtp_output_key = ("audio_codes", "current")
+    model.talker_mtp_graph_safe = False
+    model.gpu_resident_buffer_keys = set()
+    model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None
+    model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).float()
+    frame = mocker.Mock(return_value=(torch.tensor([True, False]), torch.tensor([[1, 2], [3, 4]])))
+    model.local_transformer = mocker.Mock(spec=MossTTSLocalDepthTransformer, generate_frame=frame)
+
+    state = _make_state(max_num_reqs=2, has_preprocess=True, have_multimodal_outputs=True)
+    state.model = model
+    for idx in range(2):
+        state.intermediate_buffer.buffers[idx] = {
+            "req_id": f"r{idx}",
+            "audio_state": {"is_stopping": False},
+            "hidden_states": {"last": torch.full((4,), float(idx + 1))},
+        }
+    batch = _DummyInputBatch([1, 0], num_computed_tokens_cpu=[1, 1])
+    inputs = {"input_ids": torch.tensor([2, 2]), "inputs_embeds": torch.zeros(2, 4)}
+    request_state = mocker.Mock(spec=RequestState, prompt_len=np.array([1, 1]), num_computed_tokens=None)
+    state.run_preprocess(batch, inputs, request_state)
+
+    frame.assert_called_once()
+    torch.testing.assert_close(frame.call_args.args[0], torch.tensor([[2.0] * 4, [1.0] * 4]))
+    assert frame.call_args.kwargs["temperature"] == 1.7
+    assert frame.call_args.kwargs["top_k"] == 25
+    assert frame.call_args.kwargs["top_p"] == 0.8
+    torch.testing.assert_close(inputs["inputs_embeds"], torch.tensor([[1.0] * 4, [0.0] * 4]))
+    _, payload = state.postprocess_model_output(torch.zeros(2, 4), batch, request_state)
+    assert [codes.tolist() for codes in payload["codes"]["audio"]] == [[[1, 2]], [[8, 8]]]
+    assert model.compute_logits(torch.zeros(2, 4)).argmax(-1).tolist() == [2, 3]
+    # An explicit local seed must reach the request-owned MRV2 generator.
+    params = SamplingParams(extra_args={"tts_local_seed": 17}, seed=99)
+    assert state._get_mtp_generator("seeded", params, torch.device("cpu")).initial_seed() == 17
+
+
 @pytest.mark.parametrize("owned", [False, True])
 def test_batched_postprocess_gpu_snapshot_writeback(owned):
     # batch=[1, 0]: last-token indices follow the reordered batch, the scalar
@@ -646,6 +702,58 @@ def test_publish_sampled_embeddings_is_opt_in() -> None:
     assert state.publish_sampled_embeddings(SimpleNamespace(num_reqs=1), torch.tensor([[1]])) is None
 
 
+def test_split_settled_rows_matches_per_row_check():
+    batch = _DummyInputBatch(np.array([3, 0, 2, 1, 4]))
+    batch.req_ids = ["c", "a", "x", "b", "e"]
+    batch.num_scheduled_tokens = [1, 1, 1, 4, 1]
+    batch.query_start_loc_np = [0, 1, 2, 3, 7]
+    req_indices = batch.idx_mapping_np.tolist()
+    # Slot 2 still names a finished request, slot 1 schedules a prefill
+    # chunk, slot 4 is not settled and slot 7 is not scheduled.
+    settled = {3: "c", 0: "a", 2: "old", 1: "b", 7: "z"}
+    split = OmniModelState._split_settled_rows(batch, req_indices, settled)
+    assert split is not None
+    settled_rows, remaining = split
+    assert settled_rows == [(0, 3, 0, "c"), (1, 0, 1, "a")]
+    assert remaining == [(2, 2), (3, 1), (4, 4)]
+    assert OmniModelState._split_settled_rows(batch, req_indices, {2: "old"}) is None
+
+
+def test_settled_row_split_preserves_mixed_batch_and_pending_replay(monkeypatch):
+    """Exercise run_preprocess, including the scalar fallback during replay."""
+
+    def prepare():
+        state = _make_eager_state()
+        _fill_buffers(state, "prefill", "decode")
+        state._eager_fastpath = True
+        state._eager_settled = {1: "decode"}
+        state._eager_ready = {1: "decode"}
+        state._eager_embeds[1].fill_(7)
+        state.model.eager_settled_text_step = lambda: torch.ones(_EAGER_DIM)
+        state.model.eager_decode_settled = lambda info: False
+        state.model.preprocess = lambda input_ids, input_embeds, **info: (input_ids, input_embeds, {})
+        batch = _EagerBatch([3, 1])
+        batch.req_ids = ["prefill", "decode"]
+        inputs = {"input_ids": torch.zeros(4, dtype=torch.long), "inputs_embeds": torch.zeros(4, _EAGER_DIM)}
+        return state, batch, inputs
+
+    vector, batch, inputs = prepare()
+    vector.run_preprocess(batch, inputs)
+    reference, ref_batch, ref_inputs = prepare()
+    monkeypatch.setattr(reference, "_split_settled_rows", lambda *args: None)
+    reference.run_preprocess(ref_batch, ref_inputs)
+    torch.testing.assert_close(inputs["inputs_embeds"], ref_inputs["inputs_embeds"], rtol=0, atol=0)
+    assert vector._eager_rows[1] == reference._eager_rows[1]
+    assert inputs["inputs_embeds"][3].tolist() == [8.0] * _EAGER_DIM
+
+    replay, replay_batch, replay_inputs = prepare()
+    replay._eager_state._restore_audio["prefill"] = [torch.ones(1)]
+    assert replay._eager_state.has_pending_replay()
+    monkeypatch.setattr(replay, "_split_settled_rows", MagicMock(side_effect=AssertionError("replay must skip split")))
+    replay.run_preprocess(replay_batch, replay_inputs)
+    torch.testing.assert_close(replay_inputs["inputs_embeds"], ref_inputs["inputs_embeds"], rtol=0, atol=0)
+
+
 @pytest.mark.parametrize("prefilling", [False, True])
 def test_identity_preprocess_skips_only_decode_rows(prefilling):
     state = _make_state(has_preprocess=True)

vllm_omni/config/environment_variable_inventory.py

@@ -148,11 +148,13 @@ _MODEL_PROMOTE = (
     "COSMOS3_SOUND_TOKENIZER_CONFIG_PATH",
     "COSMOS3_SOUND_TOKENIZER_PATH",
     "COSYVOICE3_BATCH_FLOW",
+    "COSYVOICE3_CACHED_ISTFT",
     "COSYVOICE3_ESTIMATOR_ONNX",
     "COSYVOICE3_FLOW_CUDAGRAPH",
     "COSYVOICE3_FLOW_GRAPH",
     "COSYVOICE3_FLOW_GRAPH_MAX_TOTAL",
     "COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS",
+    "COSYVOICE3_HIFT_GRAPH",
     "COSYVOICE3_PACKED_STREAMING",
     "COSYVOICE3_REFERENCE_PREFETCH",
     "COSYVOICE3_REFERENCE_WORKERS",
@@ -180,6 +182,24 @@ _MODEL_PROMOTE = (
     "VLLM_GEPARD_FIRST_CHUNK_FRAMES",
     "VLLM_GEPARD_LOOKBACK_FRAMES",
     "VLLM_OMNI_FISH_KVCACHE_ATTN",
+    "VLLM_OMNI_MOSS_LOCAL_COMPILE_AUDIO_SAMPLER",
+    "VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN",
+    "VLLM_OMNI_MOSS_REF_ATTN",
+    "VLLM_OMNI_MOSS_REF_BATCHED_TRANSFER",
+    "VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS",
+    "VLLM_OMNI_MOSS_REF_CACHE_ROPE",
+    "VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR",
+    "VLLM_OMNI_MOSS_REF_COMPILE",
+    "VLLM_OMNI_MOSS_REF_ENCODER_WORKERS",
+    "VLLM_OMNI_MOSS_REF_FA_VERSION",
+    "VLLM_OMNI_MOSS_REF_GPU_PREP",
+    "VLLM_OMNI_MOSS_REF_GRAPHS",
+    "VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS",
+    "VLLM_OMNI_MOSS_REF_INFLIGHT",
+    "VLLM_OMNI_MOSS_REF_SHARED_ENCODER",
+    "VLLM_OMNI_MOSS_REF_SINGLETON_BUCKET_MS",
+    "VLLM_OMNI_MOSS_REF_SKIP_PADDED_QUERY_MASK",
+    "VLLM_OMNI_MOSS_REF_STREAM_PRIORITY",
     "VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER",
     "VLLM_OMNI_SEEDVR2_LONG_JOB_TTL_SECONDS",
     "VLLM_OMNI_SEEDVR2_LONG_MAX_UPLOAD_BYTES",

vllm_omni/config/pipeline_registry.py

@@ -54,7 +54,7 @@ from vllm_omni.model_executor.models.cosmos3.pipeline import (
     COSMOS3_OMNI_DEPLOY_PIPELINE,
     COSMOS3_POLICY_PIPELINE,
 )
-from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE
+from vllm_omni.model_executor.models.cosyvoice3.pipeline import resolve_cosyvoice3_pipeline
 from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE
 from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE
 from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE
@@ -150,7 +150,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     # key (see deploy/cosmos3_policy_droid.yaml and deploy/cosmos3_omni.yaml).
     "cosmos3_omni_deploy": COSMOS3_OMNI_DEPLOY_PIPELINE,
     "cosmos3_policy": COSMOS3_POLICY_PIPELINE,
-    "cosyvoice3": COSYVOICE3_PIPELINE,
+    "cosyvoice3": resolve_cosyvoice3_pipeline,
     "covo_audio": COVO_AUDIO_PIPELINE,
     "dreamzero": DREAMZERO_PIPELINE,
     "fish_qwen3_omni": FISH_SPEECH_PIPELINE,

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -144,8 +144,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # Track requests that have already triggered prefill transfer to avoid duplicates
         self.transfer_triggered_requests: set[str] = set()
 
-        # Cache per-request flag to avoid repeated deserialization of additional_information
-        self._omits_kv_transfer_cache: dict[str, bool] = {}
+        # Cache per-request flags to avoid repeated deserialization of
+        # additional_information. Value is (payload_id, is_stage_zero_final,
+        # force_kv_transfer); payload_id invalidates the entry when the
+        # object is replaced.
+        self._omits_kv_transfer_cache: dict[str, tuple[int, bool, bool]] = {}
 
         # KV-wait start ts for the vllm_omni:kv_wait_s metric; see
         # _emit_kv_wait_output for the engine-core → orchestrator carry.
@@ -264,26 +267,40 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             return config.get(key, default)
         return getattr(config, key, default) if config is not None else default
 
-    def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool:
-        """True when this stage-zero-final request does not need downstream KV.
-
-        The result is cached per request to avoid repeated deserialization of
-        additional_information on every scheduler tick.
-        """
-        rid = request.request_id
-        cached = self._omits_kv_transfer_cache.get(rid)
-        if cached is not None:
-            return cached
-
+    def _omni_final_stage_flags(self, request: Request) -> tuple[bool, bool]:
+        """Return ``(is_stage_zero_final, force_kv_transfer)`` from request metadata."""
         payload = getattr(request, "additional_information", None)
         if payload is None:
-            result = False
-        else:
-            info = deserialize_additional_information(payload)
-            result = info.get("omni_final_stage_id") == 0 and not bool(info.get("omni_force_kv_transfer", False))
+            return False, False
 
-        self._omits_kv_transfer_cache[rid] = result
-        return result
+        cache = getattr(self, "_omits_kv_transfer_cache", None)
+        rid = request.request_id
+        payload_id = id(payload)
+        if cache is not None:
+            cached = cache.get(rid)
+            if cached is not None and cached[0] == payload_id:
+                return cached[1], cached[2]
+
+        info = deserialize_additional_information(payload)
+        is_final = info.get("omni_final_stage_id") == 0
+        force_kv = bool(info.get("omni_force_kv_transfer", False))
+        if cache is not None:
+            cache[rid] = (payload_id, is_final, force_kv)
+        return is_final, force_kv
+
+    def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool:
+        """True when this stage-zero-final request does not need downstream KV."""
+        is_final, force_kv = self._omni_final_stage_flags(request)
+        return is_final and not force_kv
+
+    def _request_omits_chunk_transfer_to_next_stage(self, request: Request) -> bool:
+        """True when this request has no downstream chunk consumer.
+
+        CFG companions still force KV transfer but are stage-0-final for
+        ordinary inter-stage chunks, so they omit ``save_async`` here.
+        """
+        is_final, _ = self._omni_final_stage_flags(request)
+        return is_final
 
     def _should_defer_waiting_admission(self) -> bool:
         return False
@@ -669,6 +686,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             confirmed_num_computed_tokens = None
             boundary_generation = None
+            # Only when this step might save. Read additional_information
+            # before _free_request rewrites it.
+            omits_chunk_transfer = False
+            if self.chunk_transfer_adapter is not None and (inter_stage_output is not None or stopped):
+                omits_chunk_transfer = self._request_omits_chunk_transfer_to_next_stage(request)
             # Capture before resumable stop handling can clear token history.
             output_token_ids: Any = getattr(request, "output_token_ids", None)
             if output_token_ids is None:
@@ -787,8 +809,16 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
 
-            if self.chunk_transfer_adapter is not None and (
-                inter_stage_output is not None or is_segment_finished or finished
+            if omits_chunk_transfer and inter_stage_output is not None:
+                logger.warning(
+                    "Skipping inter-stage chunk for request %s: "
+                    "omni_final_stage_id=0 but inter_stage_output is present",
+                    req_id,
+                )
+            if (
+                self.chunk_transfer_adapter is not None
+                and not omits_chunk_transfer
+                and (inter_stage_output is not None or is_segment_finished or finished)
             ):
                 save_kwargs = {
                     "new_token_ids": new_token_ids,

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -3,7 +3,9 @@
 
 from __future__ import annotations
 
+import math
 import os
+import queue
 import time
 from collections import defaultdict
 from typing import TYPE_CHECKING, Any
@@ -58,6 +60,48 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._init_omni_io_scheduling_state()
         self._retains_state_across_chunks = bool(getattr(model_config, "retains_state_across_chunks", False))
         self._pending_finish_reqs: list[Request] = []
+        connector_config = getattr(model_config, "stage_connector_config", None)
+        extra = (
+            connector_config.get("extra", connector_config)
+            if isinstance(connector_config, dict)
+            else getattr(connector_config, "extra", None)
+        ) or {}
+        self._generation_min_batch_size = int(extra.get("generation_min_batch_size", 1))
+        self._generation_max_wait_s = float(extra.get("generation_max_wait_ms", 0)) / 1000
+        self._generation_coalescing_policy = extra.get("generation_coalescing_policy", "fixed")
+        self._generation_max_regular_batch = int(extra.get("generation_max_regular_batch", 0) or 0)
+        if self._generation_max_regular_batch < 0:
+            raise ValueError("generation_max_regular_batch must be nonnegative")
+        if self._generation_coalescing_policy not in ("fixed", "idle_wait"):
+            raise ValueError("generation_coalescing_policy must be fixed or idle_wait")
+        if not 1 <= self._generation_min_batch_size <= self.max_num_running_reqs:
+            raise ValueError("generation_min_batch_size must be in [1, max_num_seqs]")
+        if not math.isfinite(self._generation_max_wait_s) or self._generation_max_wait_s < 0:
+            raise ValueError("generation_max_wait_ms must be finite and nonnegative")
+        if self._generation_max_wait_s and self._generation_min_batch_size > 1:
+            parallel = self.vllm_config.parallel_config
+            if not self._native_data_plane:
+                # CUDA-oriented profiles can inherit a platform's V1 runner.
+                # Their connector extras survive that override, but batching
+                # must revert to the fallback runner's default scheduling.
+                logger.warning(
+                    "Generation batch waiting and regular-batch limit ignored: "
+                    "requires native MRV2; using default generation scheduling."
+                )
+                self._generation_min_batch_size = 1
+                self._generation_max_wait_s = 0.0
+                self._generation_max_regular_batch = 0
+            elif not (
+                self._retains_state_across_chunks
+                and parallel.tensor_parallel_size == parallel.pipeline_parallel_size == 1
+            ):
+                raise ValueError("Generation batch waiting requires a stateful native MRV2 TP1/PP1 stage")
+            else:
+                logger.info(
+                    "Generation input coalescing: target batch=%d, max wait=%.3f ms",
+                    self._generation_min_batch_size,
+                    self._generation_max_wait_s * 1000,
+                )
         self._first_chunk_express = os.environ.get("VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS") == "1"
         self._last_step_express = False
         # Streams that already had a chunk scheduled; the rest await their first.
@@ -76,6 +120,101 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 self._express_min_slack_s,
             )
 
+    def _drain_omni_connector_outputs(self):
+        self._generation_defer_batch = False
+        outputs = super()._drain_omni_connector_outputs()
+        max_wait = getattr(self, "_generation_max_wait_s", 0)
+        target = getattr(self, "_generation_min_batch_size", 1)
+        coordinator = self.input_coordinator
+        if not max_wait or target <= 1 or coordinator is None or self._pause_state != PauseState.UNPAUSED:
+            return outputs
+
+        # Count distinct runnable streams, not notifications, padding rows, or
+        # chunks whose previous execution still owns their state. Waiting here
+        # precedes allocation and metadata replacement. Do not schedule a batch
+        # and then discard it to try to accumulate more input.
+        ready = set(coordinator.requests_with_ready_chunks)
+        terminal = set(coordinator.finished_requests) | set(coordinator.input_terminal_req_ids)
+        deadline = getattr(self, "_generation_batch_deadline", None)
+        if deadline is None:
+            deadline = time.monotonic() + max_wait
+        in_flight = any(
+            not request.is_finished() and request.num_in_flight_tokens > 0 for request in self.requests.values()
+        )
+        new_outputs = outputs
+        waited_without_input = False
+        while True:
+            for output in new_outputs:
+                ready.update(output.chunk_ready_req_ids)
+                terminal.update(output.chunk_finished_req_ids)
+                terminal.update(
+                    req_id for req_id, metadata in output.request_metadata.items() if metadata.get("input_terminal")
+                )
+            runnable = {
+                req_id
+                for req_id in ready | terminal
+                if (request := self.requests.get(req_id)) is not None
+                and not request.is_finished()
+                and request.num_in_flight_tokens == 0
+            }
+            if runnable and waited_without_input:
+                # Idle waiting must not eat the ready batch's coalescing budget.
+                # Start that budget at its first arrival, once only. The two
+                # bounded windows can delay control processing by at most 2W.
+                deadline = time.monotonic() + max_wait
+                waited_without_input = False
+            # A first/final chunk participates in the same bounded window.
+            # At high concurrency almost every batch contains one; bypassing
+            # the entire batch in that case defeats coalescing. Low concurrency
+            # still dispatches immediately, including single-request tails.
+            if not runnable:
+                # With no output to retire and all receivers already parked,
+                # a wakeable inbox wait avoids repeatedly executing empty steps.
+                # Admission/registration/terminal work must reach the runner
+                # first; waiting for a not-yet-registered receiver would stall it.
+                can_wait_idle = (
+                    getattr(self, "_generation_coalescing_policy", "fixed") == "idle_wait"
+                    and not in_flight
+                    and bool(self.requests)
+                    and not getattr(coordinator, "pending_chunk_registrations", ())
+                    and not getattr(coordinator, "pending_input_registrations", ())
+                    and not getattr(self, "_pending_data_plane_terminal_req_ids", ())
+                    and all(request.status == RequestStatus.WAITING_FOR_CHUNK for request in self.requests.values())
+                )
+                if not can_wait_idle:
+                    break
+                waited_without_input = True
+            elif len(runnable) >= target or len(self.requests) < target:
+                break
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                break
+            if in_flight:
+                # Let EngineCore retire the prior output before waiting for
+                # more input. Blocking this thread also blocks the in-flight
+                # fence that makes those streams eligible for the next batch.
+                # Emit a normal empty scheduler step (including connector and
+                # cancellation bookkeeping), never discard an allocated step.
+                self._generation_defer_batch = True
+                self._generation_batch_deadline = deadline
+                return outputs
+            try:
+                output = self._omni_connector_output_inbox.get(timeout=remaining)
+            except queue.Empty:
+                break
+            outputs.append(output)
+            new_outputs = [output]
+            # Drain a burst before re-evaluating, retaining notification order.
+            while True:
+                try:
+                    output = self._omni_connector_output_inbox.get_nowait()
+                except queue.Empty:
+                    break
+                outputs.append(output)
+                new_outputs.append(output)
+        self._generation_batch_deadline = None
+        return outputs
+
     def _build_generation_scheduler_output(
         self,
         new_reqs_data: list[OmniNewRequestData],
@@ -256,6 +395,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._drop_aborted_queued_requests()
         self._requeue_completed_native_chunks()
         self._process_pending_omni_inputs(model_mode="generation")
+        if getattr(self, "_generation_defer_batch", False):
+            token_budget = 0
+        max_regular = getattr(self, "_generation_max_regular_batch", 0)
+        if max_regular:
+            execution_batch_size = min(execution_batch_size, max_regular)
         self._drop_aborted_queued_requests()
         self._resync_streaming_input_counter()
         async_chunk_transport = self._async_chunk_transport_enabled()

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -229,6 +229,9 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
 
     def _finish_streaming_session_update(self, session: Request, update: StreamingUpdate) -> None:
         """Install payload metadata and send an updated session to admission."""
+        cache = getattr(self, "_omits_kv_transfer_cache", None)
+        if cache is not None:
+            cache.pop(session.request_id, None)
         session.additional_information = update.additional_information or None
         session.model_intermediate_buffer = getattr(
             update,
@@ -328,13 +331,8 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
             # ready-chunk marker remains until scheduler admission succeeds.
             replaced_ids.discard(request_id)
 
-    def _consume_pending_connector_output(self, model_mode: str) -> None:
-        """Drain ``self._latest_omni_connector_output`` into the coordinator.
-
-        Called at the top of every ``schedule()`` cycle.  Identical between
-        AR and generation schedulers except for the ``model_mode`` argument
-        forwarded to ``update_request_metadata``.
-        """
+    def _drain_omni_connector_outputs(self) -> list[OmniConnectorOutput]:
+        """Collect control messages in order, without changing request state."""
         connector_outputs: list[OmniConnectorOutput] = []
         inbox = getattr(self, "_omni_connector_output_inbox", None)
         if inbox is not None:
@@ -347,6 +345,11 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
         self._latest_omni_connector_output = None
         if connector_output is not None:
             connector_outputs.append(connector_output)
+        return connector_outputs
+
+    def _consume_pending_connector_output(self, model_mode: str) -> None:
+        """Drain input notifications into the coordinator on the scheduler thread."""
+        connector_outputs = self._drain_omni_connector_outputs()
         input_coordinator = getattr(self, "input_coordinator", None)
         if input_coordinator is None:
             return

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -379,14 +379,17 @@ class OmniSchedulingCoordinator:
                             )
 
             if model_mode != "ar":
-                new_ids = self._flatten_prompt_token_ids(metadata.get("code_predictor_codes"))
+                new_codes = metadata.get("code_predictor_codes")
+                new_ids = self._flatten_prompt_token_ids(new_codes)
                 runtime_seed = None
                 if "left_context_size" in metadata:
                     runtime_seed = {
                         "meta": {"left_context_size": metadata["left_context_size"]},
                     }
                 request._omni_initial_model_buffer = runtime_seed
-                if new_ids:
+                # An explicit empty snapshot clears a previous codec chunk
+                # or prewarmed placeholder. Missing codes leave it intact.
+                if new_codes is not None:
                     request.prompt_token_ids = new_ids
                     request.num_prompt_tokens = len(new_ids)
                     request._all_token_ids.clear()

vllm_omni/deploy/auk.yaml

@@ -20,14 +20,20 @@ enable_chunked_prefill: false
 
 stages:
   # Stage 0: Qwen2.5-Omni thinker as a frozen text encoder.
-  # Prefill only — one sampled token, discarded downstream. eager because the
-  # fused-hidden forward walks the decoder layers itself and so bypasses the
-  # compiled graph.
+  # Prefill only: one sampled token, discarded downstream. The fused-hidden
+  # forward walks the decoder layers itself, so torch.compile has nothing to
+  # hook (mode 0); the whole prefill is instead captured as FULL CUDA graphs
+  # padded to the sizes below, which removes most of its launch overhead
+  # (about 20 ms per request on H200). Longer prompts run eager.
   - stage_id: 0
     devices: "0"
     max_num_seqs: 1
     gpu_memory_utilization: 0.45
-    enforce_eager: true
+    enforce_eager: false
+    compilation_config:
+      mode: 0
+      cudagraph_mode: FULL
+      cudagraph_capture_sizes: [64, 128, 192, 256, 384, 512]
     max_num_batched_tokens: 32768
     # AuK takes text plus at most one clip; no image or video tower is built.
     limit_mm_per_prompt: {audio: 1}

vllm_omni/deploy/cosyvoice3_packed_streaming_optimized_standard.yaml

@@ -0,0 +1,28 @@
+# Single-Hopper streaming throughput profile: AR32 / codec32, Flow10.
+# Automatically selected on Hopper devices with at least 140 GiB (H200)
+# when nvidia-cuda-mps-control is available on PATH.
+# Standard sampling is configured below; request stream=true for first audio.
+# F0 and HiFT stay FP32; only deterministic HiFT decode is graphed.
+# HiFT graphs capture on the third use of each exact shape, adding latency then.
+# The codec graph pool is not capped by gpu_memory_utilization; leave memory headroom.
+base_config: cosyvoice3_packed_streaming.yaml
+stages:
+  - stage_id: 0
+    max_num_seqs: 32
+    hf_overrides:
+      cosyvoice3_sampling_mode: standard
+    default_sampling_params:
+      max_tokens: 2048
+      repetition_penalty: 1.21
+      temperature: 0.7
+      top_k: 20
+      top_p: 0.8
+  - stage_id: 1
+    max_num_seqs: 32
+    env:
+      COSYVOICE3_TRT: "0"
+      COSYVOICE3_BATCH_FLOW: "1"
+      COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS: "0"
+      COSYVOICE3_PACKED_STREAMING: "1"
+      COSYVOICE3_CACHED_ISTFT: "1"
+      COSYVOICE3_HIFT_GRAPH: "1"

vllm_omni/deploy/moss_tts_local.yaml

@@ -18,6 +18,9 @@ connectors:
       connector_get_max_wait: 300
       initial_codec_chunk_frames: 1   # 1 frame * 80 ms/frame = 80 ms first audio chunk
       codec_chunk_frames: 15          # 15 frames * 80 ms/frame = 1.2 s steady chunks
+      # Optional startup ramp; its first entry overrides initial chunk size.
+      # Extra graph shapes/calls trade codec cost and latency for less underrun.
+      # codec_chunk_ramp: [1, 4, 15]
       codec_left_context_frames: 0
       # Chunk ramp (opt-in): gradual chunk-size ladder replacing the
       # initial_codec_chunk_frames -> codec_chunk_frames cliff. The last

vllm_omni/deploy/moss_tts_local_mrv2.yaml

@@ -0,0 +1,30 @@
+# Opt-in CUDA MRV2 for MOSS-TTS-Local-Transformer-v1.5.
+# Keep the base profile's batching, chunk geometry and V1 default.
+base_config: moss_tts_local.yaml
+model_runner: v2
+
+platforms:
+  cuda:
+    stages:
+      - stage_id: 0
+        # Bound prefill work so new prompts do not monopolize streaming decode.
+        max_num_batched_tokens: 512
+      - stage_id: 1
+        compilation_config:
+          # The codec owns CUDA Graph capture/replay. Skip Inductor compilation
+          # of its stateful decoder; plain CUDA graphs remain enabled.
+          mode: 0
+          # FULL works without Inductor; the default FULL_AND_PIECEWISE would
+          # normalize to NONE and erase the codec's capture sizes.
+          cudagraph_mode: FULL
+          # Platform overrides replace compilation_config as a whole.
+          cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64]
+          cudagraph_num_of_warmups: 1
+  npu:
+    model_runner: v1
+  xpu:
+    model_runner: v1
+  rocm:
+    model_runner: v1
+  musa:
+    model_runner: v1

vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml

@@ -0,0 +1,53 @@
+# Opt-in CUDA profile for MOSS-TTS-Local-Transformer-v1.5 on a large-memory GPU.
+# Both stages share GPU 0. Validated with one H200 and Seed-TTS EN1088.
+base_config: moss_tts_local_mrv2.yaml
+
+connectors:
+  shm:
+    extra:
+      # Coalesce ready codec streams before submitting a small execution batch.
+      # This is a dispatch target, not a cap: execution capacity remains 128.
+      generation_min_batch_size: 32
+      generation_max_wait_ms: 12
+      # Optional; preserves the first/steady sizes while adding a middle chunk.
+      # codec_chunk_ramp: [1, 4, 15]
+
+# Optional codec attention-ring headroom (see recipes/OpenMOSS/MOSS-TTS.md):
+# sizes each ring as context + one chunk so chunked streaming matches
+# whole-sequence decoding. Requires additional state/graph memory; evaluate
+# throughput and audio quality before enabling for your context and capacity.
+# Add codec_ring_headroom: 1 under connectors.shm.extra above.
+
+platforms:
+  cuda:
+    stages:
+      - stage_id: 0
+        max_num_seqs: 128
+        max_num_batched_tokens: 512
+        hf_overrides:
+          mrv2_gpu_slot_state: true
+          # Throughput-oriented pair; measure first-packet latency for your load.
+          mrv2_batch_prefill: true
+          mrv2_direct_tokens: true
+        # Mixed prefill/decode fits the same bounded token budget as capture.
+        compilation_config:
+          mode: 3
+          cudagraph_mode: FULL
+          max_cudagraph_capture_size: 512
+          cudagraph_num_of_warmups: 1
+        # Leave room for the codec state pool and its batch-128 CUDA graphs.
+        kv_cache_memory_bytes: 34359738368  # 32 GiB
+      - stage_id: 1
+        max_num_seqs: 128
+        hf_overrides:
+          codec_attention_backend: triton_slot
+        compilation_config:
+          mode: 3
+          backend: inductor
+          # The codec owns capture/replay after Inductor compiles its decoder.
+          cudagraph_mode: FULL
+          cudagraph_capture_sizes: [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128]
+          cudagraph_num_of_warmups: 1
+          inductor_compile_config:
+            combo_kernels: false
+            benchmark_combo_kernel: false

vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml

@@ -0,0 +1,66 @@
+# Opt-in CUDA profile for MOSS-TTS-Local-Transformer-v1.5 tuned for first-packet
+# latency at client concurrency 128 on one H200 (Seed-TTS EN1088, closed loop).
+# Both stages share GPU 0. Use moss_tts_local_mrv2_high_concurrency.yaml for
+# throughput-first serving.
+#
+# Optional launch-level settings (see the full command in recipes/OpenMOSS/MOSS-TTS.md):
+#   - CUDA MPS for the GPU (stage 0 below lowers its own MPS client priority);
+#   - VLLM_OMNI_EVENT_DRIVEN_ORCH=1;
+#   - --api-server-count 4 with VLLM_WORKER_MULTIPROC_METHOD=spawn and
+#     VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR=/dev/shm/<dir> so frontends share reference codes;
+#   - VLLM_OMNI_CONNECTOR_RECV_POLL_MS=1 so Stage 1 picks up first chunks sooner.
+base_config: moss_tts_local_mrv2.yaml
+
+connectors:
+  shm:
+    extra:
+      generation_min_batch_size: 32
+      generation_max_wait_ms: 12
+      # Stage 1 decodes each first chunk from its receive thread on a small
+      # dedicated graph set and emits it directly, so first chunks no longer
+      # wait for the engine loop or a regular codec step.
+      codec_first_chunk_fast_path: 1
+      # Regular codec replays wait for an in-flight first-chunk decode.
+      codec_first_chunk_gate: 1
+      # Short regular codec steps bound contention with first-chunk decodes.
+      generation_max_regular_batch: 16
+      # Optional ramp. Both the first-chunk graph and regular graphs follow it.
+      # codec_chunk_ramp: [1, 4, 15]
+
+platforms:
+  cuda:
+    stages:
+      - stage_id: 0
+        max_num_seqs: 128
+        # Larger prefill budget admits bursts of new requests in fewer steps.
+        max_num_batched_tokens: 2048
+        hf_overrides:
+          mrv2_gpu_slot_state: true
+          mrv2_batch_prefill: true
+          mrv2_direct_tokens: true
+          # Sample each frame right after the forward that produced it.
+          mrv2_eager_mtp: true
+        compilation_config:
+          mode: 3
+          cudagraph_mode: FULL
+          max_cudagraph_capture_size: 2048
+          cudagraph_num_of_warmups: 1
+        kv_cache_memory_bytes: 34359738368  # 32 GiB
+        env:
+          # Under MPS, let codec kernels schedule ahead of the talker's.
+          CUDA_MPS_CLIENT_PRIORITY: "1"
+      - stage_id: 1
+        max_num_seqs: 128
+        # Keep synchronous scheduling; benchmark async scheduling separately.
+        async_scheduling: false
+        hf_overrides:
+          codec_attention_backend: triton_slot
+        compilation_config:
+          mode: 3
+          backend: inductor
+          cudagraph_mode: FULL
+          cudagraph_capture_sizes: [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128]
+          cudagraph_num_of_warmups: 1
+          inductor_compile_config:
+            combo_kernels: false
+            benchmark_combo_kernel: false

vllm_omni/diffusion/models/auk/alias_free_triton.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Triton kernels for the AuK codec's alias-free SnakeBeta activation.
+
+The eager module oversamples with a grouped transposed convolution, applies
+SnakeBeta, and decimates with a grouped strided convolution. Both grouped
+convolutions carry one 12-tap filter shared by every channel, and cuDNN and
+ATen serve them with generic depthwise kernels that dominate the decode (more
+than half of a 12 s decode on H200). Here the oversampling is written as its
+polyphase form fused with the activation, and the decimation as a direct
+strided FIR, each one pass over memory per row. The arithmetic is the same
+as the eager module in fp32; only the summation order of the taps differs.
+
+The ops are registered with ``torch.library.triton_op`` so the compiled
+decode buckets trace them instead of breaking the graph.
+"""
+
+from __future__ import annotations
+
+import torch
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+__all__ = ["alias_free_snake"]
+
+_BLOCK = 1024
+
+if HAS_TRITON:
+
+    @triton.jit
+    def _upsample_snake_kernel(
+        x_ptr,
+        out_ptr,
+        taps_ptr,
+        exp_alpha_ptr,
+        inv_beta_ptr,
+        channels,
+        in_len,
+        out_len,
+        ratio: tl.constexpr,
+        taps: tl.constexpr,
+        pad: tl.constexpr,
+        crop: tl.constexpr,
+        block: tl.constexpr,
+    ):
+        # out[m] = snake(ratio * sum_k taps[k] * xpad[(m + crop - k) / ratio]) over the k
+        # with (m + crop - k) divisible by ratio, where xpad is x replicate-padded by pad
+        # on both sides. That is conv_transpose1d(stride=ratio) followed by the crop.
+        row = tl.program_id(0)
+        offs = tl.program_id(1) * block + tl.arange(0, block)
+        mask = offs < out_len
+        x_row = x_ptr + row.to(tl.int64) * in_len
+        acc = tl.zeros((block,), dtype=tl.float32)
+        for k in tl.static_range(taps):
+            pos = offs + (crop - k)
+            hit = (pos % ratio) == 0
+            src = pos // ratio - pad
+            src = tl.minimum(tl.maximum(src, 0), in_len - 1)
+            w = tl.load(taps_ptr + k)
+            val = tl.load(x_row + src, mask=mask & hit, other=0.0)
+            acc += tl.where(hit, w * val, 0.0)
+        u = acc * ratio
+        channel = row % channels
+        alpha = tl.load(exp_alpha_ptr + channel)
+        inv_beta = tl.load(inv_beta_ptr + channel)
+        s = tl.sin(u * alpha)
+        u = u + inv_beta * (s * s)
+        tl.store(out_ptr + row.to(tl.int64) * out_len + offs, u, mask=mask)
+
+    @triton.jit
+    def _lowpass_kernel(
+        x_ptr,
+        out_ptr,
+        taps_ptr,
+        in_len,
+        out_len,
+        stride: tl.constexpr,
+        taps: tl.constexpr,
+        pad_left: tl.constexpr,
+        block: tl.constexpr,
+    ):
+        # out[n] = sum_k taps[k] * x[clamp(stride * n + k - pad_left)], i.e. replicate
+        # padding followed by a strided conv1d.
+        row = tl.program_id(0)
+        offs = tl.program_id(1) * block + tl.arange(0, block)
+        mask = offs < out_len
+        x_row = x_ptr + row.to(tl.int64) * in_len
+        acc = tl.zeros((block,), dtype=tl.float32)
+        for k in tl.static_range(taps):
+            src = offs * stride + (k - pad_left)
+            src = tl.minimum(tl.maximum(src, 0), in_len - 1)
+            acc += tl.load(taps_ptr + k) * tl.load(x_row + src, mask=mask, other=0.0)
+        tl.store(out_ptr + row.to(tl.int64) * out_len + offs, acc, mask=mask)
+
+    @torch.library.triton_op("vllm_omni::auk_alias_free_snake", mutates_args=())
+    def alias_free_snake(
+        x: torch.Tensor,
+        up_taps: torch.Tensor,
+        down_taps: torch.Tensor,
+        exp_alpha: torch.Tensor,
+        inv_beta: torch.Tensor,
+        up_ratio: int,
+        up_pad: int,
+        up_crop: int,
+        down_stride: int,
+        down_pad_left: int,
+        down_pad_right: int,
+    ) -> torch.Tensor:
+        """Upsample by ``up_ratio``, SnakeBeta, low-pass and decimate: ``[B, C, T] -> [B, C, T_out]``."""
+        batch, channels, length = x.shape
+        x = x.contiguous()
+        up_len = length * up_ratio
+        oversampled = torch.empty((batch, channels, up_len), device=x.device, dtype=torch.float32)
+        rows = batch * channels
+        torch.library.wrap_triton(_upsample_snake_kernel)[(rows, triton.cdiv(up_len, _BLOCK))](
+            x,
+            oversampled,
+            up_taps,
+            exp_alpha,
+            inv_beta,
+            channels,
+            length,
+            up_len,
+            ratio=up_ratio,
+            taps=up_taps.numel(),
+            pad=up_pad,
+            crop=up_crop,
+            block=_BLOCK,
+        )
+        taps = down_taps.numel()
+        out_len = (up_len + down_pad_left + down_pad_right - taps) // down_stride + 1
+        out = torch.empty((batch, channels, out_len), device=x.device, dtype=torch.float32)
+        torch.library.wrap_triton(_lowpass_kernel)[(rows, triton.cdiv(out_len, _BLOCK))](
+            oversampled,
+            out,
+            down_taps,
+            up_len,
+            out_len,
+            stride=down_stride,
+            taps=taps,
+            pad_left=down_pad_left,
+            block=_BLOCK,
+        )
+        return out
+
+else:  # pragma: no cover
+
+    def alias_free_snake(*args, **kwargs):
+        raise RuntimeError("Triton is not available")

vllm_omni/diffusion/models/auk/auk_transformer.py

@@ -153,6 +153,33 @@ class TimeEmbedding(nn.Module):
         return self.time_mlp(hidden.to(self.time_mlp[0].weight.dtype))
 
 
+def _grouped_conv1d_btc(x: torch.Tensor, conv: nn.Conv1d) -> torch.Tensor:
+    """``conv`` applied to channels-last ``x`` ``[B, T, C]``, returning ``[B, T, C]``.
+
+    A same-padded grouped convolution written as one im2col copy and one
+    batched GEMM over the groups. cuDNN runs a grouped conv of this shape as
+    one kernel per group plus two layout transposes, which at AuK's sizes
+    (1536 channels in 16 groups, kernel 31, a few hundred frames) costs about
+    ten times the arithmetic.
+    """
+    batch, frames, channels = x.shape
+    groups = conv.groups
+    kernel = conv.kernel_size[0]
+    per_group = channels // groups
+    pad = conv.padding[0]
+    padded = F.pad(x, (0, 0, pad, pad))
+    # unfold gives a [B, T, C, K] view; one copy lays it out as [B, T, G, C/G * K].
+    cols = padded.unfold(1, kernel, 1).reshape(batch * frames, groups, per_group * kernel)
+    # conv.weight [C_out, C/G, K] -> [G, C/G * K, C_out/G]; the reduction order matches the columns.
+    weight = conv.weight.reshape(groups, channels // groups, per_group * kernel).transpose(1, 2).to(cols.dtype)
+    # [G, B*T, C/G*K] @ [G, C/G*K, C_out/G] -> [G, B*T, C_out/G]
+    out = torch.bmm(cols.transpose(0, 1), weight)
+    out = out.transpose(0, 1).reshape(batch, frames, channels)
+    if conv.bias is not None:
+        out = out + conv.bias.to(out.dtype)
+    return out
+
+
 class ConvPosEmbedding(nn.Module):
     """Depthwise-grouped conv stack that adds local position information."""
 
@@ -169,17 +196,19 @@ class ConvPosEmbedding(nn.Module):
         )
 
     def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
-        keep = None if mask is None else mask.unsqueeze(1)
-        x = x.transpose(1, 2)
+        keep = None if mask is None else mask.unsqueeze(-1)
         if keep is not None:
             x = x.masked_fill(~keep, 0.0)
         for layer in self.conv1d:
-            x = layer(x)
-            # Re-zero padding after each convolution so it cannot bleed into
-            # valid frames through the kernel window.
-            if keep is not None and isinstance(layer, nn.Conv1d):
-                x = x.masked_fill(~keep, 0.0)
-        return x.transpose(1, 2)
+            if isinstance(layer, nn.Conv1d):
+                x = _grouped_conv1d_btc(x, layer)
+                # Re-zero padding after each convolution so it cannot bleed into
+                # valid frames through the kernel window.
+                if keep is not None:
+                    x = x.masked_fill(~keep, 0.0)
+            else:
+                x = layer(x)
+        return x
 
 
 class AudioEmbedding(nn.Module):
@@ -205,9 +234,12 @@ class AdaLayerNorm(nn.Module):
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
 
     def forward(
-        self, x: torch.Tensor, emb: torch.Tensor
+        self, x: torch.Tensor, emb: torch.Tensor | None, mod: torch.Tensor | None = None
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.linear(self.silu(emb)).chunk(6, dim=1)
+        """``mod`` is the precomputed ``linear(silu(emb))``; ``emb`` is then unused."""
+        if mod is None:
+            mod = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1)
         x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
         return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
 
@@ -221,8 +253,11 @@ class AdaLayerNormFinal(nn.Module):
         self.linear = nn.Linear(dim, dim * 2)
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
 
-    def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
-        scale, shift = self.linear(self.silu(emb)).chunk(2, dim=1)
+    def forward(self, x: torch.Tensor, emb: torch.Tensor | None, mod: torch.Tensor | None = None) -> torch.Tensor:
+        """``mod`` is the precomputed ``linear(silu(emb))``; ``emb`` is then unused."""
+        if mod is None:
+            mod = self.linear(self.silu(emb))
+        scale, shift = mod.chunk(2, dim=1)
         return self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
 
 
@@ -383,9 +418,11 @@ class DoubleBlock(nn.Module):
         c_rope: RopeCache,
         c_mask: torch.Tensor | None,
         bias: torch.Tensor | None = None,
+        mod_c: torch.Tensor | None = None,
+        mod_x: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, t)
-        norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, t)
+        norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, t, mod_c)
+        norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, t, mod_x)
 
         x_attn, c_attn = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope, c_mask=c_mask, bias=bias)
 
@@ -416,8 +453,9 @@ class SingleBlock(nn.Module):
         mask: torch.Tensor | None,
         rope: RopeCache,
         bias: torch.Tensor | None = None,
+        mod: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, t)
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, t, mod)
         x = x + gate_msa[:, None] * self.attn(x=norm, mask=mask, rope=rope, bias=bias)
         norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
         return x + gate_mlp[:, None] * self.ff(norm)
@@ -491,6 +529,28 @@ class AuKTransformer(nn.Module):
 
         self.norm_out = AdaLayerNormFinal(dim)
         self.proj_out = nn.Linear(dim, latent_dim)
+        # Widths of the adaLN modulations in the order modulation_table() lays them out.
+        self._modulation_widths = [module.linear.out_features for module in self._adaln_modules()]
+
+    def _adaln_modules(self) -> list[nn.Module]:
+        """Every timestep-conditioned norm, in the order :meth:`step` consumes them."""
+        modules: list[nn.Module] = []
+        for block in self.transformer_blocks:
+            modules += [block.attn_norm_c, block.attn_norm_x]
+        modules += [block.attn_norm for block in self.single_transformer_blocks]
+        modules.append(self.norm_out)
+        return modules
+
+    def modulation_table(self, timesteps: torch.Tensor) -> torch.Tensor:
+        """adaLN modulations for every timestep, ``[len(timesteps), sum of widths]``.
+
+        The modulations depend only on the timestep, and a request's time grid
+        is known before its first step. Computing them for the whole grid at
+        once reads each modulation weight once per request instead of once per
+        step, which at batch 1 is what those matrix-vector products cost.
+        """
+        emb = F.silu(self.time_embed(timesteps.reshape(-1).float()))
+        return torch.cat([module.linear(emb) for module in self._adaln_modules()], dim=1)
 
     def project_text(self, text: torch.Tensor) -> torch.Tensor:
         """Project LLM hidden states ``[B, nt, text_hidden_dim]`` to model width."""
@@ -520,6 +580,7 @@ class AuKTransformer(nn.Module):
         drop_text: bool = False,
         cfg_infer: bool = False,
         cache: bool = False,
+        timesteps: torch.Tensor | None = None,
     ) -> "AuKStepContext":
         """Everything a denoise step needs that does not depend on ``x`` or the time.
 
@@ -528,6 +589,9 @@ class AuKTransformer(nn.Module):
         prompt embedding, the padding masks and biases and the rotary tables
         are not recomputed on every step. Arguments are those of
         :meth:`forward`; ``target_len`` is the target frame count ``n``.
+        ``timesteps`` are the times the steps will run at (the grid without
+        its end point); when given, their adaLN modulations are precomputed
+        and :meth:`step` selects them by ``step_index``.
         """
         batch = text.shape[0]
         if c_mask is None:
@@ -593,25 +657,44 @@ class AuKTransformer(nn.Module):
             rope_text=_rope_cos_sin(self.rotary_embed(text_len, c_mask)),
             rope_single=_rope_cos_sin(self.rotary_embed(text_len + audio_len, single_mask)),
             branches=branches,
+            modulation=None if timesteps is None else self.modulation_table(timesteps),
         )
 
-    def step(self, x: torch.Tensor, time: torch.Tensor, ctx: "AuKStepContext") -> torch.Tensor:
+    def step(
+        self,
+        x: torch.Tensor,
+        time: torch.Tensor,
+        ctx: "AuKStepContext",
+        step_index: torch.Tensor | int | None = None,
+    ) -> torch.Tensor:
         """Predict the velocity of ``x`` at ``time`` under a :meth:`prepare` context.
 
         The target embedding does not depend on the CFG branch, so it is
-        computed once and shared by both branches.
+        computed once and shared by both branches. When the context carries a
+        modulation table, ``step_index`` selects this step's row of it and
+        ``time`` is unused; the row broadcasts over the CFG branches.
         """
-        if time.ndim == 0:
-            time = time.repeat(x.shape[0])
-        t = self.time_embed(time)
         target = self.audio_embed(x, mask=ctx.target_mask)
+        if ctx.modulation is not None and step_index is not None:
+            if isinstance(step_index, int):
+                row = ctx.modulation[step_index : step_index + 1]
+            else:
+                row = ctx.modulation.index_select(0, step_index.reshape(1))
+            mods = list(torch.split(row, self._modulation_widths, dim=1))
+            t = None
+        else:
+            if time.ndim == 0:
+                time = time.repeat(x.shape[0])
+            t = self.time_embed(time)
+            if ctx.branches == 2:
+                t = torch.cat((t, t), dim=0)
+            mods = [None] * len(self._modulation_widths)
         if ctx.branches == 2:
-            t = torch.cat((t, t), dim=0)
             target = torch.cat((target, target), dim=0)
         audio = target if ctx.prompt is None else torch.cat([ctx.prompt, target], dim=1)
         c = ctx.c
 
-        for block in self.transformer_blocks:
+        for i, block in enumerate(self.transformer_blocks):
             c, audio = block(
                 audio,
                 c,
@@ -621,14 +704,24 @@ class AuKTransformer(nn.Module):
                 c_rope=ctx.rope_text,
                 c_mask=ctx.c_mask,
                 bias=ctx.joint_bias,
+                mod_c=mods[2 * i],
+                mod_x=mods[2 * i + 1],
             )
 
         h = torch.cat([c, audio], dim=1)
-        for block in self.single_transformer_blocks:
-            h = block(h, t, mask=ctx.single_mask, rope=ctx.rope_single, bias=ctx.single_bias)
+        offset = 2 * len(self.transformer_blocks)
+        for i, block in enumerate(self.single_transformer_blocks):
+            h = block(
+                h,
+                t,
+                mask=ctx.single_mask,
+                rope=ctx.rope_single,
+                bias=ctx.single_bias,
+                mod=mods[offset + i],
+            )
 
         h = h[:, h.shape[1] - x.shape[1] :]
-        return self.proj_out(self.norm_out(h, t))
+        return self.proj_out(self.norm_out(h, t, mods[-1]))
 
     def forward(
         self,
@@ -706,6 +799,8 @@ class AuKStepContext:
     rope_text: RopeCache
     rope_single: RopeCache
     branches: int
+    #: adaLN modulations per step, ``[steps, widths]``, when the time grid was known.
+    modulation: torch.Tensor | None = None
 
     def tensors(self) -> list[torch.Tensor | None]:
         """Every tensor field in a fixed order, rope pairs flattened."""
@@ -721,6 +816,7 @@ class AuKStepContext:
             *self.rope_audio,
             *self.rope_text,
             *self.rope_single,
+            self.modulation,
         ]
 
     def copy_(self, other: "AuKStepContext") -> None:
@@ -752,6 +848,7 @@ class AuKStepContext:
             rope_text=(self.rope_text[0].clone(), self.rope_text[1].clone()),
             rope_single=(self.rope_single[0].clone(), self.rope_single[1].clone()),
             branches=self.branches,
+            modulation=_clone(self.modulation),
         )
 
 
@@ -881,6 +978,8 @@ def sample_latents(
                     timestep=t[i],
                     cfg_strength=cfg_strength,
                     new_request=i == 0,
+                    timesteps=t[:-1],
+                    step_index=i,
                 )
                 x = x + (t[i + 1] - t[i]) * velocity
         finally:
@@ -922,9 +1021,10 @@ def _sample_latents(
             ref_mask=ref_mask,
             cfg_infer=guided,
             cache=guided,
+            timesteps=timesteps[:-1],
         )
         for i in range(timesteps.shape[0] - 1):
-            v = dit.step(x, timesteps[i], ctx)
+            v = dit.step(x, timesteps[i], ctx, step_index=i)
             if guided:
                 v_cond, v_uncond = v.chunk(2, dim=0)
                 v = v_cond + (v_cond - v_uncond) * cfg_strength

vllm_omni/diffusion/models/auk/auk_vae.py

@@ -32,7 +32,9 @@ from torch import nn
 from torch.nn.utils import remove_weight_norm as _fold_weight_norm
 from torch.nn.utils import weight_norm as _apply_weight_norm
 from vllm.logger import init_logger
+from vllm.triton_utils import HAS_TRITON
 
+from vllm_omni.diffusion.models.auk.alias_free_triton import alias_free_snake
 from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta as _SharedSnakeBeta
 
 logger = init_logger(__name__)
@@ -319,8 +321,45 @@ class AliasFreeActivation(nn.Module):
         self.downsample = Downsample(down_ratio, down_kernel_size, causal=causal)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._fused_eligible(x):
+            return self._fused_forward(x)
         return self.downsample(self.act(self.upsample(x)))
 
+    def _fused_eligible(self, x: torch.Tensor) -> bool:
+        """Whether ``x`` can take the Triton kernels in alias_free_triton (CUDA fp32, AuK's filter geometry)."""
+        up = self.upsample
+        taps = up.filter.numel()
+        return (
+            x.is_cuda
+            and x.dtype == torch.float32
+            and x.dim() == 3
+            and isinstance(self.act, SnakeBeta)
+            and HAS_TRITON
+            # The kernel indexes the crop from the start of the transposed
+            # conv output, so the crop must cover the filter's reach, and the
+            # cropped length must be exactly ratio * T.
+            and up.pad_left >= taps - 1
+            and (2 * up.pad - 1) * up.ratio + taps == up.pad_left + up.pad_right
+        )
+
+    def _fused_forward(self, x: torch.Tensor) -> torch.Tensor:
+        up, low, act = self.upsample, self.downsample.lowpass, self.act
+        if not act._cached:
+            act.precompute_exp_cache()
+        return alias_free_snake(
+            x,
+            up.filter.reshape(-1),
+            low.filter.reshape(-1),
+            act._exp_alpha,
+            act._inv_beta,
+            up.ratio,
+            up.pad,
+            up.pad_left,
+            low.stride,
+            low.pad_left,
+            low.pad_right,
+        )
+
 
 class AmpBlock(nn.Module):
     """Anti-aliased multi-periodicity residual block: alternating alias-free snake and dilated convs."""

vllm_omni/diffusion/models/auk/cudagraph_wrapper.py

@@ -24,16 +24,19 @@ class _GraphEntry:
     static_ctx: AuKStepContext
     static_timestep: torch.Tensor
     static_cfg: torch.Tensor | None
+    static_step: torch.Tensor | None
     static_out: torch.Tensor
 
 
 class AuKCUDAGraphWrapper:
     """Replay one DiT denoise step and leave Euler scheduling to the caller.
 
-    The graph is keyed by the target, text and reference sequence lengths plus
-    whether the CFG branch is enabled. Timestep and the CFG strength are
+    The graph is keyed by the target, text and reference sequence lengths,
+    whether the CFG branch is enabled and, when the caller passes the time
+    grid, its length. Timestep (or the step index), and the CFG strength are
     mutable scalar buffers, so all Euler steps and CFG values within one path
-    reuse one graph.
+    reuse one graph. With the time grid, the adaLN modulations of every step
+    are part of the per-request context and the graph selects its row.
 
     Only :meth:`AuKTransformer.step` is captured. The per-request context
     (text projection, reference embedding, padding biases, rotary tables) is
@@ -85,8 +88,9 @@ class AuKCUDAGraphWrapper:
         text: torch.Tensor,
         ref: torch.Tensor,
         uses_cfg: bool,
-    ) -> tuple[int, int, int, bool]:
-        return (x.shape[1], text.shape[1], ref.shape[1], uses_cfg)
+        steps: int = 0,
+    ) -> tuple[int, int, int, bool, int]:
+        return (x.shape[1], text.shape[1], ref.shape[1], uses_cfg, steps)
 
     def _retire_graph_generation_if_full(self) -> None:
         """Retire all graphs together so none outlive shared workspaces."""
@@ -103,6 +107,7 @@ class AuKCUDAGraphWrapper:
         ref: torch.Tensor,
         ref_mask: torch.Tensor,
         uses_cfg: bool,
+        timesteps: torch.Tensor | None = None,
     ) -> AuKStepContext:
         return self.dit.prepare(
             text,
@@ -112,6 +117,7 @@ class AuKCUDAGraphWrapper:
             ref=ref,
             ref_mask=ref_mask,
             cfg_infer=uses_cfg,
+            timesteps=timesteps,
         )
 
     def _step(
@@ -120,9 +126,10 @@ class AuKCUDAGraphWrapper:
         timestep: torch.Tensor,
         ctx: AuKStepContext,
         cfg_strength: torch.Tensor | float,
+        step_index: torch.Tensor | int | None = None,
     ) -> torch.Tensor:
         """One guided (or unguided) velocity; the CFG combine is part of the graph."""
-        velocity = self.dit.step(x, timestep, ctx)
+        velocity = self.dit.step(x, timestep, ctx, step_index=step_index)
         if ctx.branches == 1:
             return velocity
         conditional, unconditional = velocity.chunk(2, dim=0)
@@ -140,6 +147,8 @@ class AuKCUDAGraphWrapper:
         timestep: torch.Tensor,
         cfg_strength: float,
         new_request: bool = True,
+        timesteps: torch.Tensor | None = None,
+        step_index: int | None = None,
     ) -> torch.Tensor:
         """Velocity of ``x`` at ``timestep``.
 
@@ -147,26 +156,33 @@ class AuKCUDAGraphWrapper:
         (``text``, ``c_mask``, ``ref``, ``ref_mask``) is then prepared anew; the
         later steps of the request pass ``False`` and reuse it. The default
         prepares on every call, which is always correct.
+
+        ``timesteps`` (the times every step of the request runs at) and
+        ``step_index`` (this step's position in them) let the adaLN
+        modulations be computed once per request instead of once per step.
         """
         uses_cfg = cfg_strength >= 1e-5
+        if timesteps is None or step_index is None:
+            timesteps, step_index = None, None
         if not self.enabled or x.device.type != "cuda" or torch.cuda.is_current_stream_capturing():
             if new_request or self._eager_ctx is None:
-                self._eager_ctx = self._prepare(x, None, text, c_mask, ref, ref_mask, uses_cfg)
-            return self._step(x, timestep, self._eager_ctx, cfg_strength)
+                self._eager_ctx = self._prepare(x, None, text, c_mask, ref, ref_mask, uses_cfg, timesteps)
+            return self._step(x, timestep, self._eager_ctx, cfg_strength, step_index)
 
         target_frames = x.shape[1]
         x, x_mask, text, c_mask, ref, ref_mask = self._bucket_inputs(x, text, c_mask, ref, ref_mask)
-        key = self._key(x, text, ref, uses_cfg)
+        steps = 0 if timesteps is None else int(timesteps.numel())
+        key = self._key(x, text, ref, uses_cfg, steps)
         entry = self._cache.get(key)
         ctx = None
         if entry is None or new_request or key != self._active_key:
-            ctx = self._prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg)
+            ctx = self._prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg, timesteps)
         if entry is None:
             # Same Retirement as #6587, details discussed in #7469
             self._retire_graph_generation_if_full()
             try:
                 assert ctx is not None
-                entry = self._capture(x, ctx, timestep, cfg_strength=cfg_strength)
+                entry = self._capture(x, ctx, timestep, cfg_strength=cfg_strength, step_index=step_index)
             except Exception:
                 self._cache.clear()
                 self._active_key = None
@@ -186,6 +202,8 @@ class AuKCUDAGraphWrapper:
 
         entry.static_x.copy_(x)
         entry.static_timestep.copy_(timestep)
+        if entry.static_step is not None:
+            entry.static_step.fill_(step_index)
         if entry.static_cfg is not None:
             entry.static_cfg.fill_(cfg_strength)
         entry.graph.replay()
@@ -198,10 +216,14 @@ class AuKCUDAGraphWrapper:
         timestep: torch.Tensor,
         *,
         cfg_strength: float,
+        step_index: int | None = None,
     ) -> _GraphEntry:
         static_x = x.clone()
         static_ctx = ctx.clone()
         static_timestep = timestep.clone()
+        static_step = None
+        if static_ctx.modulation is not None and step_index is not None:
+            static_step = torch.full((), step_index, device=static_x.device, dtype=torch.long)
         static_cfg = None
         if ctx.branches == 2:
             static_cfg = torch.empty((), device=static_x.device, dtype=torch.float32)
@@ -209,13 +231,13 @@ class AuKCUDAGraphWrapper:
         scale = static_cfg if static_cfg is not None else cfg_strength
         # Warm-up runs any lazy torch.compile outside the capture.
         for _ in range(3):
-            self._step(static_x, static_timestep, static_ctx, scale)
+            self._step(static_x, static_timestep, static_ctx, scale, static_step)
         if self._pool_handle is None:
             self._pool_handle = torch.cuda.graph_pool_handle()
         graph = torch.cuda.CUDAGraph()
         with torch.cuda.graph(graph, pool=self._pool_handle):
             # The CFG scalar stays a mutable buffer, so cfg=2 and cfg=3 share this graph.
-            static_out = self._step(static_x, static_timestep, static_ctx, scale)
+            static_out = self._step(static_x, static_timestep, static_ctx, scale, static_step)
 
         logger.info(
             "Captured AuK DiT single-step CUDA graph: target_frames=%d text_tokens=%d ref_frames=%d cfg=%s",
@@ -230,6 +252,7 @@ class AuKCUDAGraphWrapper:
             static_ctx=static_ctx,
             static_timestep=static_timestep,
             static_cfg=static_cfg,
+            static_step=static_step,
             static_out=static_out,
         )
 

vllm_omni/diffusion/models/auk/pipeline_auk.py

@@ -11,9 +11,11 @@ both, and the BigVGAN-flow decoder renders the waveform.
 
 from __future__ import annotations
 
+import hashlib
 import json
 import os
 import time
+from collections import OrderedDict
 from collections.abc import Iterable
 from typing import Any, ClassVar
 
@@ -27,7 +29,12 @@ from vllm.logger import init_logger
 from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, dit_state_dict, sample_latents
+from vllm_omni.diffusion.models.auk.auk_transformer import (
+    AuKTransformer,
+    build_time_grid,
+    dit_state_dict,
+    sample_latents,
+)
 from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE
 from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper
 from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph
@@ -59,6 +66,10 @@ _FLASH_CFG = 0.0
 _WARMUP_TARGET_FRAMES = (150, 300, 600)
 _WARMUP_REF_FRAMES = (0, 150)
 _WARMUP_TEXT_TOKENS = 96
+# Reference clips whose posterior-mean latents are kept. Voice cloning reuses a
+# small set of reference voices, so repeat requests skip the VAE encode.
+# ``model_config.auk_ref_cache_size`` overrides it; 0 disables the cache.
+_REF_CACHE_SIZE = 16
 # Decorrelates the VAE posterior draw from the initial latent for the same request seed.
 _VAE_SEED_OFFSET = 0x5EED_0A0C
 
@@ -152,6 +163,9 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
         max_dit_graphs = od_config.model_config.get("max_dit_graphs", 32)
         if isinstance(max_dit_graphs, bool) or not isinstance(max_dit_graphs, int) or max_dit_graphs < 1:
             raise ValueError("AuK max_dit_graphs must be a positive integer")
+        ref_cache_size = od_config.model_config.get("auk_ref_cache_size", _REF_CACHE_SIZE)
+        if isinstance(ref_cache_size, bool) or not isinstance(ref_cache_size, int) or ref_cache_size < 0:
+            raise ValueError("AuK auk_ref_cache_size must be a non-negative integer (0 disables the cache)")
         self.device = get_local_device()
         self.dtype = getattr(od_config, "dtype", None) or torch.bfloat16
 
@@ -213,6 +227,8 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
         if model_config.get("auk_vae_tile_frames") is not None:
             vae_decode_kwargs["tile_frames"] = int(model_config["auk_vae_tile_frames"])
         self.vae_decode = AuKVAEDecodeGraph(self.vae, enabled=not od_config.enforce_eager, **vae_decode_kwargs)
+        self._ref_cache_size = ref_cache_size
+        self._ref_cache: OrderedDict[str, torch.Tensor] = OrderedDict()
 
         logger.info(
             "AuK pipeline ready: variant=%s dtype=%s latent_dim=%d hop=%d sample_rate=%d",
@@ -267,6 +283,12 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
         text = torch.zeros(1, _WARMUP_TEXT_TOKENS, self.text_hidden_dim, device=self.device, dtype=self.dtype)
         c_mask = torch.ones(text.shape[:2], dtype=torch.bool, device=self.device)
         timestep = torch.zeros((), device=self.device, dtype=torch.float32)
+        # The graphs are keyed by the number of steps, so warm the variant's default grid.
+        if self.is_flash:
+            grid = torch.tensor(self.flash_t_grid, device=self.device, dtype=torch.float32)
+        else:
+            sway = None if self.default_sway is None else float(self.default_sway)
+            grid = build_time_grid(nfe=self.default_nfe, sway_sampling_coef=sway, t_grid=None, device=self.device)
         start = time.perf_counter()
         with torch.inference_mode(), self._dit_autocast():
             for ref_frames in _WARMUP_REF_FRAMES:
@@ -283,6 +305,8 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
                         timestep=timestep,
                         cfg_strength=cfg,
                         new_request=True,
+                        timesteps=grid[:-1],
+                        step_index=0,
                     )
         if self.device.type != "cpu":
             torch.accelerator.synchronize(self.device)
@@ -386,8 +410,24 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
 
         if audio is None:
             return torch.zeros(1, 0, self.latent_dim, device=self.device, dtype=torch.float32)
-        wav = self._prepare_waveform(_unwrap_single(audio)).to(self.device)
-        return self.vae.encode(wav, sample=sample, generator=generator)
+        wav = self._prepare_waveform(_unwrap_single(audio))
+        # The posterior mean is a pure function of the prepared waveform, so it
+        # is cached by content. A posterior draw depends on the generator and
+        # is always recomputed. Callers must not modify the returned tensor.
+        key = None
+        if not sample and self._ref_cache_size:
+            digest = hashlib.blake2b(wav.numpy().tobytes(), digest_size=16).hexdigest()
+            key = f"{wav.shape[-1]}:{digest}"
+            cached = self._ref_cache.get(key)
+            if cached is not None:
+                self._ref_cache.move_to_end(key)
+                return cached
+        latents = self.vae.encode(wav.to(self.device), sample=sample, generator=generator)
+        if key is not None:
+            self._ref_cache[key] = latents
+            while len(self._ref_cache) > self._ref_cache_size:
+                self._ref_cache.popitem(last=False)
+        return latents
 
     def _target_frames(self, gen_seconds: Any, ref_frames: int) -> int:
         """Target latent length, capped by the remaining context."""

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -307,9 +307,13 @@ class SharedMemoryConnector(OmniConnectorBase):
             self._metrics["gets"] += 1
         return result
 
-    def cleanup(self, request_id: str) -> None:
-        """Unlink the exact key passed to ``put()``, never a request-id prefix."""
+    def cleanup(self, request_id: str) -> bool:
+        """Unlink the exact key passed to ``put()``, never a request-id prefix.
+
+        Returns True when an unconsumed segment was actually unlinked.
+        """
         key = request_id
+        unlinked = False
         with self._pending_keys_lock:
             self._pending_keys.pop(key, None)
             try:
@@ -317,6 +321,7 @@ class SharedMemoryConnector(OmniConnectorBase):
                 seg.close()
                 seg.unlink()
                 logger.debug("cleanup: unlinked unconsumed SHM segment %s", key)
+                unlinked = True
             except FileNotFoundError:
                 pass
             except Exception as e:
@@ -327,6 +332,18 @@ class SharedMemoryConnector(OmniConnectorBase):
                     os.remove(lock_file)
                 except OSError:
                     pass
+        return unlinked
+
+    def cleanup_prefix(self, key_prefix: str) -> int:
+        """Unlink every tracked key of the form ``{key_prefix}{chunk_id}``.
+
+        Only keys still in ``_pending_keys`` are considered, and the suffix
+        must be a bare integer chunk id, so another request whose id merely
+        shares the prefix is never matched. Returns the unlinked count.
+        """
+        with self._pending_keys_lock:
+            keys = [k for k in self._pending_keys if k.startswith(key_prefix) and k[len(key_prefix) :].isdigit()]
+        return sum(self.cleanup(key) for key in keys)
 
     def close(self) -> None:
         """Unlink all remaining tracked SHM segments."""

vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py

@@ -1169,7 +1169,12 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                 self._payload_finished(payload_data),
             )
 
-        self._get_req_chunk[req_id] += 1
+        with self._lock:
+            if self._async_chunk and request is not None and req_id not in self._pending_load_reqs:
+                # A connector get can return after cancellation removed its
+                # receiver. Do not recreate delivery state for that request.
+                return False
+            self._get_req_chunk[req_id] += 1
 
         if self._async_chunk:
             is_finished = self._payload_finished(payload_data)
@@ -1180,8 +1185,16 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                 if not self._payload_value_has_content(new_ids) and not is_finished:
                     return False
                 payload_consumable = self._payload_is_consumable(payload_data)
+                first_chunk_hook = getattr(self, "_first_chunk_hook", None)
+                if first_chunk_hook is not None and chunk_id == 0 and not is_finished and payload_consumable:
+                    first_chunk_hook(req_id, request, payload_data)
 
             with self._lock:
+                if request is not None and req_id not in self._pending_load_reqs:
+                    # The first-chunk hook may have overlapped cancellation.
+                    # Its slot cleanup is ordered by the model; this payload
+                    # must not republish readiness after receive teardown.
+                    return False
                 if self._model_mode == "ar":
                     # Accumulation, staging, and model-side consume/ack share
                     # this lock. Keeping the transition atomic prevents the
@@ -1723,3 +1736,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         sink: Callable[[OmniConnectorOutput], None] | None,
     ) -> None:
         self._omni_connector_output_sink = sink
+
+    def set_first_chunk_hook(self, hook: Callable[[str, Any, OmniPayload], bool]) -> None:
+        """Claim eligible first chunks before publishing them to the model loop."""
+        self._first_chunk_hook = hook

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -28,6 +28,8 @@ from .base import OmniTransferAdapterBase
 
 logger = get_connector_logger(__name__)
 
+_RECLAIM_WARN_INTERVAL = 1000
+
 
 class _SenderGeneration:
     """Fence one external request generation without blocking cleanup."""
@@ -196,6 +198,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.request_payload = {}
         self.code_prompt_token_ids: dict[str, list[torch.Tensor]] = defaultdict(list)
         self.request_ids_mapping: dict[str, str] = {}
+        # Save-thread only: running count of segments reclaimed after finish.
+        self._reclaimed_shm_total = 0
 
         self.waiting_for_chunk_waiting_requests: deque[Any] = deque()
         self.waiting_for_chunk_running_requests: deque[Any] = deque()
@@ -688,6 +692,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             for chunk_id in range(num_puts):
                 self.connector.cleanup(f"{key_prefix}_{chunk_id}")
             return
+        if "release_shm" in task:
+            self._release_shm_prefix(task["release_shm"])
+            return
         request = task["request"]
         external_req_id = request.external_req_id
         sender_token = task.get("sender_token")
@@ -979,8 +986,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
         Called after a terminal chunk is sent or when the scheduler aborts the
         request before a terminal chunk can be produced. In-flight sends are
-        cancelled here and reclaim their own state in ``finally``; cleanup
-        never waits for connector I/O on the scheduler thread.
+        cancelled here and reclaim adapter state plus leftover SHM in
+        ``finally`` after ``put()`` returns; cleanup never waits for connector
+        I/O on the scheduler thread.
 
         Idempotent: calling with an already-cleaned or unknown id is safe.
         """
@@ -1042,6 +1050,42 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         if cached_ic is not None:
             cached_ic.pop(external_req_id, None)
 
+    def release_shm_resources(self, request_id: str) -> None:
+        """Queue reclaim of inter-stage segments this request left unconsumed.
+
+        Safe only once every stage has finished with the request: the producer
+        keeps writing until its own request ends, and a consumer that stopped
+        early never drains the remainder. Successful terminal chunks are kept
+        by ``_queue_shm_cleanup_locked`` for exactly that consumer, so the
+        orchestrator -- the only party that knows all stages are done -- drives
+        this (see ``Orchestrator._cleanup_request_ids``).
+
+        Like abort cleanup, the unlink runs on the save thread, so the
+        scheduler never performs SHM I/O.
+        """
+        if not isinstance(self.connector, SharedMemoryConnector):
+            return
+        external_req_id = self.request_ids_mapping.get(request_id, request_id)
+        self._pending_save_reqs.append({"release_shm": f"{external_req_id}_{self.connector.stage_id}_"})
+        with self._save_cond:
+            self._save_cond.notify()
+
+    def _release_shm_prefix(self, key_prefix: str) -> None:
+        reclaimed = self.connector.cleanup_prefix(key_prefix)
+        if not reclaimed:
+            return
+        # Non-zero means a consumer stopped before draining the producer.
+        # Common on audio requests, so warn only when the running total
+        # crosses another _RECLAIM_WARN_INTERVAL boundary.
+        prev = self._reclaimed_shm_total
+        self._reclaimed_shm_total = prev + reclaimed
+        if prev // _RECLAIM_WARN_INTERVAL != self._reclaimed_shm_total // _RECLAIM_WARN_INTERVAL:
+            logger.warning(
+                "Reclaimed %d unconsumed inter-stage segments so far; "
+                "a downstream stage finished before draining its producer",
+                self._reclaimed_shm_total,
+            )
+
     def cleanup(
         self,
         request_id: str,

vllm_omni/engine/orchestrator.py

@@ -301,6 +301,7 @@ class OrchestratorBase:
     _transfer_emitter: Any = None
     _prom_metrics: Any = None
     _stat_logger: OmniPrometheusStatLogger | None = None
+    _transfer_release_tasks: set[asyncio.Task] = set()
 
     def __init__(
         self,
@@ -348,6 +349,8 @@ class OrchestratorBase:
             self._pd_bootstrap_addr = pd_config.get("bootstrap_addr")
             self._pd_prefill_engine_id = pd_config.get("prefill_engine_id")
         self.request_states: dict[str, OrchestratorRequestState] = {}
+        # Strong refs for in-flight releases; the loop only weak-refs tasks, so dropping these risks mid-flight GC.
+        self._transfer_release_tasks: set[asyncio.Task] = set()
         self._init_metrics_state(
             stage_pools,
             running_counter,
@@ -690,6 +693,36 @@ class OrchestratorBase:
             output_msg.finished = index == last_index_by_req[output_msg.request_id]
         return abort_outputs
 
+    def _release_stage_transfer_resources(self, request_ids: list[str]) -> None:
+        """Drop each stage's inter-stage transfer resources for finished requests.
+
+        This is the only point that knows every stage is done with the request,
+        so it is the only safe place to reclaim segments a consumer never
+        drained. Scheduled rather than awaited: reclaim is best-effort and must
+        not add RPC latency to request teardown.
+        """
+        if not request_ids:
+            return
+
+        async def _run() -> None:
+            results = await asyncio.gather(
+                *(pool.release_request_resources(request_ids) for pool in self.stage_pools),
+                return_exceptions=True,
+            )
+            for result in results:
+                if isinstance(result, Exception):
+                    logger.warning("[Orchestrator] release transfer resources failed: %s", result)
+
+        try:
+            task = asyncio.get_running_loop().create_task(_run())
+            self._transfer_release_tasks.add(task)
+            task.add_done_callback(self._transfer_release_tasks.discard)
+        except RuntimeError:
+            logger.warning(
+                "[Orchestrator] no running event loop; skipped reclaim of transfer resources for %s",
+                request_ids,
+            )
+
     def _release_request_bindings(self, request_ids: list[str]) -> None:
         """Release all stage-local route bindings for the given request ids."""
         for pool in self.stage_pools:
@@ -1645,6 +1678,7 @@ class OrchestratorBase:
         if abort:
             abort_outputs = await self._abort_request_ids(cleanup_ids)
         self._release_request_bindings(cleanup_ids)
+        self._release_stage_transfer_resources(cleanup_ids)
         for request_id in cleanup_ids:
             self._pd_kv_params.pop(request_id, None)
             req_state = self.request_states.pop(request_id, None)

vllm_omni/engine/stage_engine_core_proc.py

@@ -80,9 +80,12 @@ def _signal_exit_code(signum: int) -> int:
 
 
 def _bind_first_audio_sink(model_executor: Any, output_queue: Any, scheduler: Any) -> bool:
-    """Let a TP1 in-process runner that decodes first audio deliver it as its own output."""
+    """Bind in-process first audio to the generic engine output sink."""
     if not isinstance(model_executor, UniProcExecutor):
         return False
+    parallel_config = model_executor.vllm_config.parallel_config
+    if parallel_config.tensor_parallel_size != 1 or parallel_config.pipeline_parallel_size != 1:
+        return False
     worker = getattr(getattr(model_executor, "driver_worker", None), "worker", None)
     model_runner = getattr(worker, "model_runner", None)
     model_state = getattr(model_runner, "model_state", None)
@@ -90,12 +93,23 @@ def _bind_first_audio_sink(model_executor: Any, output_queue: Any, scheduler: An
     decodes_audio = getattr(model, "first_frame_decoder", None) is not None or (
         getattr(model, "stream_decoder", None) is not None and bool(getattr(model, "stream_first_audio", False))
     )
-    if not decodes_audio or not hasattr(model_state, "set_first_audio_sink"):
-        return False
     from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink
 
-    assert model_state is not None
-    model_state.set_first_audio_sink(engine_output_queue_sink(output_queue, scheduler))
+    if decodes_audio and hasattr(model_state, "set_first_audio_sink"):
+        assert model_state is not None
+        model_state.set_first_audio_sink(engine_output_queue_sink(output_queue, scheduler))
+        return True
+    data_plane = getattr(model_runner, "_omni_data_plane", None)
+    get_model = getattr(model_runner, "get_model", None)
+    if data_plane is None or not callable(get_model):
+        return False
+    bind = getattr(get_model(), "bind_first_chunk_fast_path", None)
+    if not callable(bind):
+        return False
+    hook = bind(engine_output_queue_sink(output_queue, scheduler, upstream_first_audio=False))
+    if hook is None:
+        return False
+    data_plane.set_first_chunk_hook(hook)
     return True
 
 
@@ -131,6 +145,21 @@ class StageEngineCoreProc(EngineCoreProc):
         if _bind_native_data_plane_ready_sink(self.model_executor, self.scheduler):
             logger.info("Bound native MRv2 connector readiness directly to the scheduler inbox.")
 
+    def omni_release_request_resources(self, request_ids: list[str]) -> None:
+        """Release this stage's inter-stage transfer resources for *request_ids*.
+
+        Invoked over the UTILITY channel by the orchestrator once every stage
+        has finished with the request. Idempotent and safe for unknown ids.
+        """
+        adapter = getattr(getattr(self, "scheduler", None), "chunk_transfer_adapter", None)
+        if adapter is None:
+            return
+        for request_id in request_ids or ():
+            try:
+                adapter.release_shm_resources(request_id)
+            except Exception as e:
+                logger.debug("omni_release_request_resources(%s) failed: %s", request_id, e)
+
     def preprocess_add_request(self, request: OmniEngineCoreRequest) -> tuple[Any, int]:
         """Preserve omni payloads when vLLM builds its scheduler request."""
         scheduler_request, current_wave = super().preprocess_add_request(request)

vllm_omni/engine/stage_pool.py

@@ -99,6 +99,8 @@ class StagePool:
 
     DISPATCH_WAIT_TIMEOUT_S: float = 10.0
     DISPATCH_RETRY_INTERVAL_S: float = 0.1
+    # A replica that dies mid-release never answers; without a bound its background release never finishes.
+    RELEASE_RPC_TIMEOUT_S: float = 5.0
     # Only these EngineCore helpers may skip collective_rpc_async. A generic
     # ``{method}_async`` on AsyncMPClient must not silently drop timeout.
     _CACHE_RESET_METHODS = frozenset({"reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"})
@@ -135,6 +137,9 @@ class StagePool:
         self.clients: list[StagePoolClient | None] = list(normalized_clients)
         self._output_processor = output_processor
         self._stage_vllm_config = stage_vllm_config
+        self._has_chunk_transfer_adapter = bool(
+            getattr(getattr(stage_vllm_config, "model_config", None), "async_chunk", False)
+        )
         self._next_replica_id = 0
         self._request_bindings: dict[str, int] = {}
         self._unavailable_replicas: set[int] = set()
@@ -1304,6 +1309,28 @@ class StagePool:
 
         return abort_outputs
 
+    async def release_request_resources(self, request_ids: list[str]) -> None:
+        """Ask every live replica to drop transfer resources for *request_ids*.
+
+        Broadcast rather than binding-routed: the orchestrator releases route
+        bindings as part of the same teardown, so a binding lookup here would
+        race it. The engine-core handler is idempotent for unknown ids.
+        """
+        if not request_ids or not self._has_chunk_transfer_adapter:
+            return
+        ids = list(request_ids)
+
+        async def release(replica_id: int, call: Any) -> None:
+            try:
+                await asyncio.wait_for(call("omni_release_request_resources", ids), timeout=self.RELEASE_RPC_TIMEOUT_S)
+            except Exception as e:
+                logger.warning(
+                    "[StagePool-%s] release_request_resources on replica %s failed: %r", self.stage_id, replica_id, e
+                )
+
+        calls = [(i, getattr(self.clients[i], "call_utility_async", None)) for i in self.live_replica_ids()]
+        await asyncio.gather(*(release(i, call) for i, call in calls if call is not None))
+
     async def collective_rpc(
         self,
         replica_id: int,

vllm_omni/entrypoints/omni_base.py

@@ -441,7 +441,7 @@ class OmniBase(PDDisaggregationMixin):
             if self.log_stats:
                 # Emit per-request orchestrator timing (including e2e_total_ms)
                 # before dropping request state.
-                req_state.metrics.build_and_log_summary()
+                req_state.metrics.log_timing_summary()
         except Exception:
             logger.exception(
                 "[%s] Failed to build/log summary for req=%s",

vllm_omni/entrypoints/openai/tts_adapters/base.py

@@ -81,13 +81,24 @@ def conditioning_cache_salt(
     This replaces the uploaded data URI in the hash without changing the request.
     """
     h = hashlib.sha256()
+    ref_audio: str | list[str] | tuple[str, int] | None = request.ref_audio
+    if (
+        isinstance(ref_audio, str)
+        and len(ref_audio) > 256
+        and tts_params is not None
+        and tts_params.get("ref_audio_cache_key") is not None
+    ):
+        # The resolve key below is already a content-aware digest of this
+        # locator. Hashing a repr of a large inline data URI again dominates
+        # the per-request serving cost without adding information.
+        ref_audio = ("ref_audio_cache_key", len(ref_audio))
     for part in (
         request.input,
         request.task_type,
         request.language,
         request.voice,
         request.ref_text,
-        registered_voice if registered_voice is not None else request.ref_audio,
+        registered_voice if registered_voice is not None else ref_audio,
         request.instructions,
         request.x_vector_only_mode,
         request.speaker_embedding,

vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py

@@ -32,6 +32,59 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def _local_continuation_prompt(proc: Any, user_kwargs: dict[str, Any], reference: Any) -> tuple[list[int], Any] | None:
+    """Build the Local-v1.5 continuation prompt without the generic processor call.
+
+    Mirrors ``_build_continuation_codes`` of the model's processor: every
+    template segment is encoded separately there, so the constant segments
+    (keyed by ``language``) are cached and only the request text is
+    tokenized. Returns ``(text_ids, audio_codes)`` equal to
+    ``unified[:, 0].tolist()`` / ``unified[:, 1:]``, or ``None`` when the
+    request or processor does not match that exact layout.
+    """
+    import torch
+
+    if set(user_kwargs) - {"text", "language"}:
+        return None
+    if not isinstance(reference, list) or len(reference) != 1 or not isinstance(reference[0], torch.Tensor):
+        return None
+    cfg = getattr(proc, "model_config", None)
+    n_vq = getattr(cfg, "n_vq", None)
+    if cfg is None or not isinstance(n_vq, int) or n_vq <= 0:
+        return None
+    codes = reference[0]
+    if codes.ndim != 2 or int(codes.shape[1]) != n_vq or codes.shape[0] == 0:
+        return None
+    required = (
+        "_user_prompt_prefix_ids",
+        "_encode_text",
+        "_user_prompt_after_reference_ids",
+        "_assistant_prompt_prefix_ids",
+    )
+    if not all(callable(getattr(proc, name, None)) for name in required):
+        return None
+    language = user_kwargs.get("language")
+    cache = proc.__dict__.setdefault("_vllm_omni_continuation_segments", {})
+    segments = cache.get(language)
+    if segments is None:
+        fields = {} if language is None else {"language": language}
+        head = (
+            proc._user_prompt_prefix_ids()
+            + proc._encode_text("None")
+            + proc._user_prompt_after_reference_ids(language, fields)
+        )
+        tail = proc._assistant_prompt_prefix_ids() + [int(cfg.audio_start_token_id)]
+        segments = cache[language] = ([int(t) for t in head], [int(t) for t in tail])
+    head, tail = segments
+    text = user_kwargs.get("text")
+    prompt_ids = head + [int(t) for t in proc._encode_text("" if text is None else str(text))] + tail
+    num_audio = int(codes.shape[0])
+    text_ids = prompt_ids + [int(cfg.audio_assistant_slot_token_id)] * num_audio
+    audio_codes = torch.full((len(text_ids), n_vq), int(cfg.audio_pad_token_id), dtype=torch.int64)
+    audio_codes[len(prompt_ids) :] = codes.to(dtype=torch.int64, device="cpu")
+    return text_ids, audio_codes
+
+
 class _MossTTSAdapterBase(ARTTSAdapter):
     accumulate_nonstreaming: bool = False
 
@@ -82,6 +135,28 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         self._moss_ref_encoder = encoder
         return encoder
 
+    async def warmup(self) -> None:
+        """Set up the reference encoder at startup, not on the first cloning request.
+
+        Its CUDA graph capture takes seconds per API process (about a minute for
+        a shared encoder's workers).
+        """
+        if self._moss_variant not in ("tts", "local", "ttsd"):
+            return
+        from vllm_omni.model_executor.models.moss_tts.reference_encoder import reference_graphs_enabled
+        from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import SharedReferenceEncoderStartupError
+
+        if not reference_graphs_enabled():
+            return
+        try:
+            await asyncio.to_thread(self._get_moss_ref_encoder().prepare)
+        except SharedReferenceEncoderStartupError:
+            # Do not advertise readiness and trigger concurrent local captures
+            # while the shared host is still compiling or has failed startup.
+            raise
+        except Exception:  # noqa: BLE001 - requests still encode (and capture) on demand
+            logger.warning("MOSS reference encoder warmup failed", exc_info=True)
+
     async def _encode_moss_references(
         self,
         request: "OpenAICreateSpeechRequest",
@@ -269,8 +344,16 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
         if hasattr(proc, "audio_tokenizer"):
             device = self._resolve_ref_encoder_device()
-            proc.audio_tokenizer = proc.audio_tokenizer.to(device).eval()
-            logger.info("MOSS reference-audio encoder (audio_tokenizer) placed on %s", device)
+            proc._vllm_omni_ref_encoder_device = device
+            from vllm_omni.model_executor.models.moss_tts.reference_encoder import shared_encoder_role
+
+            if shared_encoder_role() == "client":
+                # Another API process hosts the shared encoder on the GPU.
+                proc.audio_tokenizer = proc.audio_tokenizer.eval()
+                logger.info("MOSS reference-audio encoder: using the shared encoder of another API process")
+            else:
+                proc.audio_tokenizer = proc.audio_tokenizer.to(device).eval()
+                logger.info("MOSS reference-audio encoder (audio_tokenizer) placed on %s", device)
         self._moss_processor_cache = proc
         return proc
 
@@ -405,18 +488,24 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         # Build the unified-codes prompt: (L, 1+n_vq) where col 0 is text/special
         # tokens and cols 1..n_vq are the delay-pattern audio code grid (mostly
         # audio_pad_code outside the reference block).
+        fast = None
         if v == "local" and request.ref_text and request.ref_text.strip():
             reference = user_kwargs.pop("reference")
             user_kwargs["text"] = request.ref_text.strip() + " " + (request.input or "")
-            user_msg = proc.build_user_message(**user_kwargs)
-            assistant_msg = proc.build_assistant_message(audio_codes_list=reference)
-            batch = proc(conversations=[[user_msg, assistant_msg]], mode="continuation")
+            fast = _local_continuation_prompt(proc, user_kwargs, reference)
+            if fast is None:
+                user_msg = proc.build_user_message(**user_kwargs)
+                assistant_msg = proc.build_assistant_message(audio_codes_list=reference)
+                batch = proc(conversations=[[user_msg, assistant_msg]], mode="continuation")
         else:
             user_msg = proc.build_user_message(**user_kwargs)
             batch = proc(conversations=[[user_msg]], mode="generation")
-        unified = batch["input_ids"][0]  # torch.LongTensor (L, 1+n_vq)
-        text_ids: list[int] = unified[:, 0].tolist()
-        audio_codes: torch.Tensor = unified[:, 1:].contiguous().to(torch.int64)
+        if fast is not None:
+            text_ids, audio_codes = fast
+        else:
+            unified = batch["input_ids"][0]  # torch.LongTensor (L, 1+n_vq)
+            text_ids = unified[:, 0].tolist()
+            audio_codes = unified[:, 1:].contiguous().to(torch.int64)
 
         params = {
             "prompt_token_ids": text_ids,

vllm_omni/metrics/stats.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import copy
+import logging
 import time
 from collections import defaultdict
 from collections.abc import Callable, Sequence
@@ -851,9 +852,62 @@ class OrchestratorAggregator:
             for cached_key in [k for k in self._replica_cache if k[1] == rid_key]:
                 self._replica_cache.pop(cached_key, None)
 
+    def log_timing_summary(self) -> None:
+        """Emit only the concise per-request ``[OmniTiming]`` lines.
+
+        The serving path uses this unless DEBUG tables are requested; it skips
+        the table rows that ``build_and_log_summary`` builds for its result.
+        """
+        if not self.log_stats:
+            return
+        if logger.isEnabledFor(logging.DEBUG):
+            self.build_and_log_summary()
+            return
+        for rid in sorted(set(self.stage_events.keys()) | {e.request_id for e in self.e2e_events}):
+            e2e_evt = next((e for e in self.e2e_events if e.request_id == rid), None)
+            stage_evts = sorted(
+                self.stage_events.get(rid, []),
+                key=lambda e: e.stage_id if e.stage_id is not None else -1,
+            )
+            self._log_omni_timing(rid, e2e_evt, stage_evts)
+
+    def _log_omni_timing(self, rid: str, e2e_evt: RequestE2EStats | None, stage_evts: list[StageRequestStats]) -> None:
+        pt: dict[str, float] = {}
+        if stage_evts:
+            pt = stage_evts[-1].pipeline_timings or {}
+        if pt or e2e_evt:
+            parts = [f"req={rid}"]
+            if e2e_evt:
+                parts.append(f"total={e2e_evt.e2e_total_ms / 1000.0:.2f}s")
+            if "preprocess_ms" in pt:
+                parts.append(f"preprocess={pt['preprocess_ms'] / 1000.0:.2f}s")
+            if e2e_evt:
+                engine_ms = e2e_evt.e2e_total_ms - pt.get("preprocess_ms", 0.0)
+                parts.append(f"engine={engine_ms / 1000.0:.2f}s")
+            stage_parts = []
+            for evt in stage_evts:
+                sid = evt.stage_id if evt.stage_id is not None else "?"
+                t = evt.stage_gen_time_ms / 1000.0
+                stage_parts.append(f"{sid}:{t:.2f}s")
+            if stage_parts:
+                parts.append(f"stages=[{','.join(stage_parts)}]")
+            transfer_parts = []
+            for te in self.transfer_events.values():
+                if te.request_id == rid:
+                    transfer_parts.append(f"{te.from_stage}->{te.to_stage}={te.tx_time_ms:.2f}ms")
+            if transfer_parts:
+                parts.append(f"transfers=[{','.join(transfer_parts)}]")
+            if "ar2diffusion_ms" in pt:
+                parts.append(f"ar2diffusion={pt['ar2diffusion_ms']:.2f}ms")
+            logger.info("[OmniTiming] %s", " ".join(parts))
+
     def build_and_log_summary(self) -> dict[str, Any]:
         if not self.log_stats:
             return {}
+        # Per-request tables are diagnostic. Formatting them on the serving
+        # event loop costs more than the request's own bookkeeping, so they
+        # are emitted at DEBUG; the concise [OmniTiming] line stays at INFO.
+        log_tables = logger.isEnabledFor(logging.DEBUG)
         wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0)
         e2e_avg_req = (wall_time_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
         e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / wall_time_ms) if wall_time_ms > 0 else 0.0
@@ -895,8 +949,8 @@ class OrchestratorAggregator:
             v = overall_summary.get(k, None)
             if v not in (0, 0.0, 0.000, None, ""):
                 overall_fields.append(k)
-        if overall_fields:
-            logger.info(
+        if log_tables and overall_fields:
+            logger.debug(
                 "\n%s",
                 _format_table("Overall Summary", overall_summary, overall_fields),
             )
@@ -921,8 +975,8 @@ class OrchestratorAggregator:
                         nonzero_e2e_fields.add(k)
                 value_fields_e2e = sorted(nonzero_e2e_fields)
 
-                if value_fields_e2e:
-                    logger.info(
+                if log_tables and value_fields_e2e:
+                    logger.debug(
                         "\n%s",
                         _format_table(
                             f"RequestE2EStats [request_id={rid}]",
@@ -936,34 +990,7 @@ class OrchestratorAggregator:
                 self.stage_events.get(rid, []),
                 key=lambda e: e.stage_id if e.stage_id is not None else -1,
             )
-            pt = {}
-            if stage_evts:
-                pt = stage_evts[-1].pipeline_timings or {}
-            if pt or e2e_evt:
-                parts = [f"req={rid}"]
-                if e2e_evt:
-                    parts.append(f"total={e2e_evt.e2e_total_ms / 1000.0:.2f}s")
-                if "preprocess_ms" in pt:
-                    parts.append(f"preprocess={pt['preprocess_ms'] / 1000.0:.2f}s")
-                if e2e_evt:
-                    engine_ms = e2e_evt.e2e_total_ms - pt.get("preprocess_ms", 0.0)
-                    parts.append(f"engine={engine_ms / 1000.0:.2f}s")
-                stage_parts = []
-                for evt in stage_evts:
-                    sid = evt.stage_id if evt.stage_id is not None else "?"
-                    t = evt.stage_gen_time_ms / 1000.0
-                    stage_parts.append(f"{sid}:{t:.2f}s")
-                if stage_parts:
-                    parts.append(f"stages=[{','.join(stage_parts)}]")
-                transfer_parts = []
-                for te in self.transfer_events.values():
-                    if te.request_id == rid:
-                        transfer_parts.append(f"{te.from_stage}->{te.to_stage}={te.tx_time_ms:.2f}ms")
-                if transfer_parts:
-                    parts.append(f"transfers=[{','.join(transfer_parts)}]")
-                if "ar2diffusion_ms" in pt:
-                    parts.append(f"ar2diffusion={pt['ar2diffusion_ms']:.2f}ms")
-                logger.info("[OmniTiming] %s", " ".join(parts))
+            self._log_omni_timing(rid, e2e_evt, stage_evts)
 
             # === Stage table (columns = stage_id) ===
             # if any stage has diffusion_metrics, remove postprocess_time_ms field
@@ -987,7 +1014,7 @@ class OrchestratorAggregator:
 
             result_stage_table.append({"request_id": rid, "stages": stage_rows})
 
-            if stage_rows:
+            if log_tables and stage_rows:
                 # filter out all-zero fields for logging
                 all_value_fields = set()
                 for row in stage_rows:
@@ -1006,7 +1033,7 @@ class OrchestratorAggregator:
                         value_fields_list.append(field)
 
                 if value_fields_list:
-                    logger.info(
+                    logger.debug(
                         "\n%s",
                         _format_table(
                             f"StageRequestStats [request_id={rid}]",
@@ -1027,7 +1054,7 @@ class OrchestratorAggregator:
             ]
             result_trans_table.append({"request_id": rid, "transfers": transfer_rows})
 
-            if transfer_rows:
+            if log_tables and transfer_rows:
                 # filter out all-zero fields for logging
                 all_value_fields = set()
                 for row in transfer_rows:
@@ -1046,7 +1073,7 @@ class OrchestratorAggregator:
                         value_fields_list.append(field)
 
                 if value_fields_list:
-                    logger.info(
+                    logger.debug(
                         "\n%s",
                         _format_table(
                             f"TransferEdgeStats [request_id={rid}]",

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py

@@ -966,7 +966,20 @@ class CausalHiFTGenerator(HiFTGenerator):
         self.conv_pre_look_right = conv_pre_look_right
         self.f0_predictor = f0_predictor
 
+    def enable_decode_graphs(self) -> None:
+        """Opt into exact-shape CUDA replay after frozen weights are loaded."""
+        from .hift_graph import HiFTDecodeGraphs
+
+        self.enable_cached_istft()
+        self._decode_graphs = HiFTDecodeGraphs(self._decode_eager)
+
     def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, 0), finalize: bool = True) -> torch.Tensor:
+        runner = getattr(self, "_decode_graphs", None)
+        if runner is not None:
+            return runner.run(x, s, finalize)
+        return self._decode_eager(x, s, finalize=finalize)
+
+    def _decode_eager(self, x: torch.Tensor, s: torch.Tensor, finalize: bool = True) -> torch.Tensor:
         s_stft_real, s_stft_imag = self._stft(s.squeeze(1))
         if finalize is True:
             x = self.conv_pre(x)

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hift_graph.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Bounded CUDA graph replay for HiFT's deterministic waveform decoder.
+
+F0, sine/noise generation, phase carry, and per-request streaming state remain
+outside the graph. Each replay consumes current mel and source tensors and
+returns an owned waveform, because multiple requests can reuse one shape.
+"""
+
+from collections import OrderedDict
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class _DecodeGraph:
+    graph: torch.cuda.CUDAGraph
+    mel: torch.Tensor
+    source: torch.Tensor
+    output: torch.Tensor
+
+
+class HiFTDecodeGraphs:
+    """Capture recurrent exact shapes, with at most 32 resident graphs.
+
+    ``decode`` must have warmed its ISTFT envelope before capture and remain
+    deterministic for its mel, source and finalize arguments. Uncommon shapes,
+    CPU execution and nested CUDA capture use that same eager implementation.
+    """
+
+    def __init__(self, decode: Callable[..., torch.Tensor], max_graphs: int = 32) -> None:
+        self.decode = decode
+        self.max_graphs = max_graphs
+        self.graphs: dict[tuple, _DecodeGraph] = {}
+        self.seen: OrderedDict[tuple, int] = OrderedDict()
+        self.pool = None
+
+    @torch.inference_mode()
+    def run(self, mel: torch.Tensor, source: torch.Tensor, finalize: bool) -> torch.Tensor:
+        """Replay a recurring layout without aliasing another request's output."""
+        if not mel.is_cuda or not source.is_cuda or torch.cuda.is_current_stream_capturing():
+            return self.decode(mel, source, finalize=finalize)
+        autocast = torch.is_autocast_enabled("cuda")
+        key = (
+            mel.shape,
+            source.shape,
+            mel.device,
+            mel.dtype,
+            source.dtype,
+            finalize,
+            torch.backends.cudnn.allow_tf32,
+            torch.backends.cuda.matmul.allow_tf32,
+            autocast,
+            torch.get_autocast_dtype("cuda") if autocast else None,
+        )
+        state = self.graphs.get(key)
+        if state is None:
+            if len(self.graphs) >= self.max_graphs:
+                return self.decode(mel, source, finalize=finalize)
+            count = self.seen.pop(key, 0) + 1
+            self.seen[key] = count
+            if len(self.seen) > 128:
+                self.seen.popitem(last=False)
+            if count < 3:
+                return self.decode(mel, source, finalize=finalize)
+            static_mel, static_source = mel.clone(), source.clone()
+            stream = torch.cuda.Stream(device=mel.device)
+            stream.wait_stream(torch.cuda.current_stream(mel.device))
+            with torch.cuda.stream(stream):
+                self.decode(static_mel, static_source, finalize=finalize)
+            torch.cuda.current_stream(mel.device).wait_stream(stream)
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, pool=self.pool):
+                output = self.decode(static_mel, static_source, finalize=finalize)
+            if self.pool is None:
+                self.pool = graph.pool()
+            state = _DecodeGraph(graph, static_mel, static_source, output)
+            self.graphs[key] = state
+        state.mel.copy_(mel)
+        state.source.copy_(source)
+        state.graph.replay()
+        # A later request may replay this shape before the caller consumes it.
+        return state.output.clone()

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_dit.py

@@ -170,7 +170,7 @@ class PackedDiT:
             t,
             rows,
             attention,
-            max_seqlen_q=attention.page_table.shape[1],
+            max_seqlen_q=attention.max_seqlen_q,
             block_modulation=block_modulation,
             final_modulation=final_modulation,
             rope=rope,
@@ -771,7 +771,9 @@ class FlowGraphRunner:
         attention.heads, attention.head_dim = self.heads, self.head_dim
         attention.cache_seqlens, attention.cu_seqlens_q = state.cache_seqlens, state.cu_seqlens_q
         attention.page_table = state.page_table
-        attention.max_seqlen_q = state.page_table.shape[1]
+        # Streaming queries cover one training chunk, even when their KV
+        # prefix and the packed graph allocation contain thousands of frames.
+        attention.max_seqlen_q = estimator.dit.static_chunk_size if key[0] else state.page_table.shape[1]
         mark_packed_compile_metadata(rows, attention)
         mu_cfg = torch.cat((state.mu, torch.zeros_like(state.mu)), dim=1)
         cond_cfg = torch.cat((state.cond, torch.zeros_like(state.cond)), dim=1)

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -74,6 +74,27 @@ logger = init_logger(__name__)
 _RUNTIME_COMPONENTS_CACHE: dict[str, dict] = {}
 
 
+def _normalize_request_conditioning(payload: dict) -> dict:
+    """Unwrap singleton conditioning containers at the per-request boundary.
+
+    A processor/prefix-cache passthrough can retain a one-item tensor list.
+    The code2wav schema consumes one request, so that wrapper carries no
+    batch dimension. A multi-item list here is an unsplit batch and must
+    fail instead of silently choosing another request's voice.
+    """
+    embed = payload.get("embed")
+    if not isinstance(embed, Mapping):
+        return payload
+    normalized = dict(embed)
+    for key in ("speech_token", "speech_feat", "embedding", "speech_token_len"):
+        value = normalized.get(key)
+        if isinstance(value, (list, tuple)):
+            if len(value) != 1:
+                raise ValueError(f"CosyVoice3 per-request {key} contains an unsplit batch of {len(value)} items")
+            normalized[key] = value[0]
+    return {**payload, "embed": normalized}
+
+
 def _cosyvoice3_trt_enabled() -> bool:
     """COSYVOICE3_TRT env toggle (default on) for the optional TensorRT paths.
 
@@ -1501,9 +1522,10 @@ class CosyVoice3Model(
             # [total_tokens, hidden]
             hidden_states = self.model.llm(inputs_embeds, positions)
 
-            if getattr(self, "_mrv2_tensor_output", False) or cosyvoice3_packed_inference_enabled():
-                return hidden_states
-            return self.make_omni_output(hidden_states, **kwargs)
+            # Both runners attach the live step's conditioning after replay.
+            # Capturing an OmniOutput here would retain the dummy prefill lists
+            # and re-emit them during decode, replacing the real reference.
+            return hidden_states
         elif self.model_stage == "cosyvoice3_code2wav":
             # Lazily swap the flow-decoder estimator to a TensorRT engine on the
             # first code2wav step (after weights are loaded), gated by the same
@@ -1539,7 +1561,7 @@ class CosyVoice3Model(
 
             for idx, req_ids in enumerate(request_ids_list):
                 raw = runtime_info[idx] if idx < len(runtime_info) and isinstance(runtime_info[idx], dict) else {}
-                payload = to_struct(raw)
+                payload = to_struct(_normalize_request_conditioning(raw))
                 meta = payload.meta
                 embed = payload.embed
 

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py

@@ -769,4 +769,11 @@ class CosyVoice3Code2Wav(nn.Module):
         if f0_folded == 0:
             logger.warning("HiFT F0 predictor had no weight-norm layers to fold; check config drift")
         self.hift.eval()
+        if device.type == "cuda" and cosyvoice3_packed_streaming_enabled():
+            if os.getenv("COSYVOICE3_CACHED_ISTFT", "0") == "1":
+                self.hift.enable_cached_istft()
+                logger.info("CosyVoice3 streaming HiFT: cached ISTFT enabled")
+            if os.getenv("COSYVOICE3_HIFT_GRAPH", "0") == "1":
+                self.hift.enable_decode_graphs()
+                logger.info("CosyVoice3 streaming HiFT: decode CUDA graphs enabled")
         logger.info(f"Loaded hift weights from {hift_path}")