[Test] Align CosyVoice3 hardware mocks across platform layers

Signed-off-by: Sy03 <1370724210@qq.com>

.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

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

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(

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: 77,
         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: 48,
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 17,

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_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_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/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/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

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",

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_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,

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/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/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

@@ -131,6 +131,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/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_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}")

vllm_omni/model_executor/models/cosyvoice3/pipeline.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
 """CosyVoice3 pipeline topology (frozen).
 
 Stage 0: Talker   — text prompt → speech tokens (LLM autoregressive).
@@ -12,11 +12,17 @@ Stage 1: Code2Wav — flow-matching decoder → acoustic features → waveform.
     through the shared-memory connector.
 """
 
+import shutil
+from dataclasses import replace
+
+from transformers import PretrainedConfig
+
 from vllm_omni.config.stage_config import (
     PipelineConfig,
     StageExecutionType,
     StagePipelineConfig,
 )
+from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
 
 _PROC = "vllm_omni.model_executor.stage_input_processors.cosyvoice3"
 
@@ -35,8 +41,9 @@ COSYVOICE3_PIPELINE = PipelineConfig(
             async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
             custom_process_next_stage_input_func=f"{_PROC}.text2flow_full_payload",
             sampling_constraints={
-                # merged speech stop token (logsumexp of all 200 stop logits)
-                "stop_token_ids": [6562],
+                # Standard sampling can emit any of the 200 control tokens.
+                # RAS merges their logits into 6562, which is also in this range.
+                "stop_token_ids": list(range(6561, 6761)),
             },
         ),
         StagePipelineConfig(
@@ -52,3 +59,28 @@ COSYVOICE3_PIPELINE = PipelineConfig(
         ),
     ),
 )
+
+
+def resolve_cosyvoice3_pipeline(hf_config: PretrainedConfig | None = None) -> PipelineConfig | None:
+    """Select a streaming default compatible with the device and runtime.
+
+    Query the platform through NVML without initializing CUDA in the parent.
+    Explicit deploy configs still take precedence over this pipeline default.
+    """
+    from vllm.platforms import current_platform
+
+    if hf_config is not None and not isinstance(hf_config, CosyVoice3Config):
+        return None
+    if current_platform.is_cuda():
+        capability = current_platform.get_device_capability()
+        if (
+            capability is not None
+            and capability.major == 9
+            and current_platform.get_device_total_memory() >= 140 * 1024**3
+            and shutil.which("nvidia-cuda-mps-control") is not None
+        ):
+            return replace(
+                COSYVOICE3_PIPELINE,
+                default_deploy_config_name="cosyvoice3_packed_streaming_optimized_standard.yaml",
+            )
+    return COSYVOICE3_PIPELINE