Merge main into dev/vllm-align for vLLM 0.31 release

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

tests/assets/yue2/twinkle_golden.abc

@@ -0,0 +1,6 @@
+X:1
+T:Twinkle Twinkle Little Star (SheetSage2 transcription of the Wikimedia piano rendition, single voice)
+M:4/4
+L:1/16
+K:C
+C4C4G4G4|A4A4G8|F4F4E4E4|D4D4C8|G4G4F4F4|E4E4D8|G4G4F4F4|E4E4D8|C4C4G4G4|A4A4G8|F4F4E4E4|D4D4C8|

tests/config/test_config_factory.py

@@ -2928,9 +2928,17 @@ class TestPlatformOverrides:
 
         assert deploy.stages[0].engine_extras["block_size"] == 128
 
-    @pytest.mark.parametrize("platform", ["cuda", "npu"])
-    def test_recommended_native_runner_platform_defaults(self, platform):
-        filename, pipeline_key = "qwen3_tts_mrv2.yaml", "qwen3_tts"
+    @pytest.mark.parametrize(
+        "filename,pipeline_key",
+        [
+            ("qwen3_tts_mrv2.yaml", "qwen3_tts"),
+            ("moss_tts_local_mrv2.yaml", "moss_tts_local"),
+            ("moss_tts_local_mrv2_high_concurrency.yaml", "moss_tts_local"),
+            ("moss_tts_local_mrv2_low_latency.yaml", "moss_tts_local"),
+        ],
+    )
+    @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
+    def test_recommended_native_runner_platform_defaults(self, platform, filename, pipeline_key):
         deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
         deploy = _apply_platform_overrides(deploy, platform=platform)
         expect_v2 = platform == "cuda"
@@ -2946,6 +2954,60 @@ class TestPlatformOverrides:
             # v2 only engages the native plane on stages declaring support.
             assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages)
 
+    def test_moss_local_mrv2_preserves_base_profile_and_variant_scope(self):
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        base_path = Path(get_deploy_config_path(pipeline.default_deploy_config_name))
+        candidate_path = Path(get_deploy_config_path("moss_tts_local_mrv2.yaml"))
+        base = load_deploy_config(base_path)
+        candidate = load_deploy_config(candidate_path)
+        assert base.model_runner == "v1"
+        assert candidate.stages == base.stages
+        assert candidate.connectors == base.connectors
+        assert candidate.async_chunk == base.async_chunk
+        cuda = _apply_platform_overrides(candidate, platform="cuda")
+        stages = merge_pipeline_deploy(pipeline, cuda)
+        assert stages[0].yaml_engine_args["max_num_batched_tokens"] == 512
+        codec_args = stages[1].yaml_engine_args
+        assert codec_args["max_num_seqs"] == 64
+        assert codec_args["enforce_eager"] is False
+        assert codec_args["compilation_config"] == {
+            "mode": 0,
+            "cudagraph_mode": "FULL",
+            "cudagraph_capture_sizes": [1, 2, 4, 8, 16, 32, 64],
+            "cudagraph_num_of_warmups": 1,
+        }
+        for platform in ("npu", "xpu", "rocm", "musa"):
+            # Platform resolution mutates the deployment; each server loads
+            # a fresh profile before applying its own platform overrides.
+            fallback = _apply_platform_overrides(load_deploy_config(candidate_path), platform=platform)
+            original = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
+            assert fallback.stages == original.stages
+        for variant in ("moss_tts_delay", "moss_tts_realtime"):
+            assert not any(stage.supports_native_mrv2_data_plane for stage in resolve_pipeline_config(variant).stages)
+
+    def test_moss_local_high_concurrency_resolves_slot_compile_and_memory_budget(self):
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        path = Path(get_deploy_config_path("moss_tts_local_mrv2_high_concurrency.yaml"))
+        deploy = _apply_platform_overrides(load_deploy_config(path), platform="cuda")
+        talker, codec = merge_pipeline_deploy(pipeline, deploy)
+        assert talker.yaml_engine_args["max_num_seqs"] == codec.yaml_engine_args["max_num_seqs"] == 128
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_batch_prefill"] is True
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_direct_tokens"] is True
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_gpu_slot_state"] is True
+        assert talker.yaml_engine_args["max_num_batched_tokens"] == 512
+        assert talker.yaml_engine_args["kv_cache_memory_bytes"] == 32 * 1024**3
+        assert codec.yaml_engine_args["hf_overrides"]["codec_attention_backend"] == "triton_slot"
+        compilation = codec.yaml_engine_args["compilation_config"]
+        assert compilation["mode"] == 3 and compilation["backend"] == "inductor"
+        assert compilation["cudagraph_mode"] == "FULL"
+        assert compilation["cudagraph_capture_sizes"] == [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128]
+        assert compilation["inductor_compile_config"] == {"combo_kernels": False, "benchmark_combo_kernel": False}
+        base_path = Path(get_deploy_config_path("moss_tts_local.yaml"))
+        for platform in ("npu", "xpu", "rocm", "musa"):
+            actual = _apply_platform_overrides(load_deploy_config(path), platform=platform)
+            expected = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
+            assert actual.stages == expected.stages
+
     @pytest.mark.parametrize("runner,native", [("v1", False), ("v2", False), ("v2", True)])
     def test_mrv2_undeclared_transport_warns(self, monkeypatch, runner, native):
         from unittest.mock import Mock

tests/config/test_config_merge.py

@@ -35,3 +35,161 @@ def test_recursive_deploy_inheritance(tmp_path):
     merged = resolve_deploy_yaml(leaf)
     assert merged["speech_cache"] == {"nested": {"a": 1, "b": 0}}
     assert merged["async_chunk"] is False
+
+
+def test_single_stage_streaming_capability_uses_central_async_resolver():
+    from vllm_omni.config.stage_config import DeployConfig, _resolve_pipeline_async_chunk_enabled
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE
+
+    assert _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig())
+    with pytest.raises(ValueError, match="requires async_chunk=True"):
+        _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig(async_chunk=False))
+
+
+@pytest.mark.parametrize(
+    "pipeline_key,profile,stage_count",
+    [
+        ("qwen3_tts", "qwen3_tts_high_concurrency_mrv2_single_gpu.yaml", 2),
+        ("qwen3_tts_fused", "qwen3_tts_fused_single_gpu.yaml", 1),
+    ],
+)
+def test_qwen3_tts_deployment_selection_preserves_output_and_transport_contract(pipeline_key, profile, stage_count):
+    from tests.helpers.stage_config import get_deploy_config_path
+    from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+    from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
+    pipeline = resolve_pipeline_config(pipeline_key)
+    assert pipeline is not None
+    deploy = load_deploy_config(get_deploy_config_path(profile))
+    stages = merge_pipeline_deploy(pipeline, deploy)
+    assert len(stages) == stage_count
+    talker = stages[0]
+    if stage_count == 1:
+        assert talker.final_output and talker.final_output_type == "audio"
+        assert talker.yaml_engine_args["engine_output_type"] == "audio"
+        assert "custom_process_next_stage_input_func" not in talker.yaml_engine_args
+        assert not talker.yaml_extras.get("output_connectors")
+        assert talker.yaml_engine_args["supports_running_prefix_cache_reset"] is False
+        assert talker.yaml_engine_args["additional_config"]["ref_code_context_frames"] == 72
+    else:
+        assert not talker.final_output
+        assert talker.yaml_engine_args["engine_output_type"] == "latent"
+        assert talker.yaml_engine_args["custom_process_next_stage_input_func"].endswith("talker2code2wav_async_chunk")
+        assert talker.yaml_extras["output_connectors"]
+        assert stages[1].input_sources == [0] and stages[1].final_output_type == "audio"
+        assert talker.yaml_engine_args.get("supports_running_prefix_cache_reset", True)
+
+
+def test_single_stage_running_reset_capability_cannot_be_overridden():
+    from vllm_omni.config.omni_config import VllmOmniConfig
+    from vllm_omni.config.stage_config import DeployConfig, StageDeployConfig, merge_pipeline_deploy
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE
+
+    deploy = DeployConfig(
+        stages=[StageDeployConfig(stage_id=0, engine_extras={"supports_running_prefix_cache_reset": True})]
+    )
+    legacy = merge_pipeline_deploy(QWEN3_TTS_FUSED_PIPELINE, deploy)[0]
+    legacy.runtime_overrides["supports_running_prefix_cache_reset"] = True
+    assert legacy.to_omegaconf().engine_args.supports_running_prefix_cache_reset is False
+    with pytest.raises(ValueError, match="no structured config owner: supports_running_prefix_cache_reset"):
+        VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE, user_deploy_config=deploy)
+    structured = VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE)
+    assert structured.stage_by_id(0).model_config.supports_running_prefix_cache_reset is False
+
+
+@pytest.fixture
+def single_stage_tts_config(mocker):
+    from vllm.config import VllmConfig
+
+    config = mocker.Mock(spec=VllmConfig)
+    config.additional_config = {"talker_stream_decode": True}
+    config.model_config = mocker.Mock(
+        stage_connector_config={},
+        engine_output_type="audio",
+        final_output=True,
+        supports_running_prefix_cache_reset=False,
+        use_v2_model_runner=True,
+        async_chunk=True,
+    )
+    config.parallel_config = mocker.Mock(
+        tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend="uni"
+    )
+    config.cache_config = mocker.Mock(enable_prefix_caching=False)
+    config.device_config = mocker.Mock(device="cuda")
+    mocker.patch(
+        "vllm_omni.model_executor.models.qwen3_tts.stream_decode.current_omni_platform.is_cuda", return_value=True
+    )
+    return config
+
+
+def test_single_stage_tts_rejects_unsupported_decoder_configuration(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.cache_config.enable_prefix_caching = True
+    with pytest.raises(ValueError, match="enable_prefix_caching=False"):
+        talker_stream_decode_enabled(config)
+    config.cache_config.enable_prefix_caching = False
+    assert talker_stream_decode_enabled(config)
+
+
+@pytest.mark.parametrize("option_location", ["additional_config", "connector"])
+@pytest.mark.parametrize("output_type", ["latent", "audio"])
+def test_two_stage_talker_rejects_single_stage_pcm_option(single_stage_tts_config, option_location, output_type):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.model_config.engine_output_type = output_type
+    config.model_config.final_output = False
+    if option_location == "connector":
+        config.model_config.stage_connector_config = {"extra": config.additional_config}
+        config.additional_config = {}
+    with pytest.raises(ValueError, match="final audio-output Talker"):
+        talker_stream_decode_enabled(config)
+
+
+def test_two_stage_first_audio_remains_independent(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.first_audio import talker_first_audio_enabled
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.model_config.engine_output_type = "latent"
+    config.additional_config = {}
+    config.model_config.stage_connector_config = {"extra": {"talker_first_audio": True}}
+    assert not talker_stream_decode_enabled(config)
+    assert talker_first_audio_enabled(config)
+    config.model_config.stage_connector_config["extra"]["talker_stream_first_audio"] = True
+    with pytest.raises(ValueError, match="requires the single-stage"):
+        talker_stream_decode_enabled(config)
+
+
+def test_single_stage_pcm_requires_cuda_device_and_own_first_audio_option(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.device_config.device = "cpu"
+    with pytest.raises(ValueError, match="requires CUDA"):
+        talker_stream_decode_enabled(config)
+    config.device_config.device = "cuda"
+    config.additional_config["talker_first_audio"] = True
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        talker_stream_decode_enabled(config)
+    config.additional_config.pop("talker_first_audio")
+    config.additional_config["talker_stream_first_audio"] = True
+    assert talker_stream_decode_enabled(config)
+
+
+def test_single_stage_audio_talker_requires_pcm_option(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    single_stage_tts_config.additional_config = {}
+    with pytest.raises(ValueError, match="requires talker_stream_decode=True"):
+        talker_stream_decode_enabled(single_stage_tts_config)
+
+
+def test_single_stage_pcm_requires_reset_capability_declaration(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    single_stage_tts_config.model_config.supports_running_prefix_cache_reset = True
+    with pytest.raises(ValueError, match="must disable running prefix-cache reset"):
+        talker_stream_decode_enabled(single_stage_tts_config)

tests/config/test_environment_variables.py

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

tests/config/test_omni_config.py

@@ -1046,6 +1046,7 @@ def test_sub_config_fields_match_structured_scopes():
         "interleave_mm_strings",
         "media_io_kwargs",
         "final_output",
+        "supports_running_prefix_cache_reset",
         "active_stream_window",
         "session_mode",
         "duplex_max_sessions",

tests/core/sched/test_generation_batch_coalescing.py

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

tests/core/sched/test_generation_scheduler_init.py

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

tests/core/sched/test_omni_ar_scheduler_kv_transfer.py

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

tests/core/sched/test_omni_ar_scheduler_stale_drain.py

@@ -40,6 +40,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 
 # isort: on
 
@@ -104,10 +105,11 @@ 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
 
 
-def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int) -> bool:
+def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int, spec=None) -> bool:
     """Feed one model-runner frame through update_from_output.
 
     Returns True when the frame's tokens were delivered (the append ran),
@@ -115,7 +117,7 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token:
     """
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: num_scheduled}
-    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.scheduled_spec_decode_tokens = {session.request_id: spec} if spec else {}
     scheduler_output.num_invalid_spec_tokens = 0
 
     model_runner_output = MagicMock(spec=ModelRunnerOutput)
@@ -135,6 +137,27 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token:
     return sched._update_request_with_output.called
 
 
+def test_connector_prompt_replacement_drops_old_frame_and_delivers_new_frame() -> None:
+    from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+
+    session = _make_session()
+    session.num_computed_tokens = 6
+    session.num_in_flight_tokens = 1
+    session.num_output_placeholders = 1
+    sched = _make_drain_sched(session)
+    sched.chunk_transfer_adapter.replaced_streaming_prompt_ids = {session.request_id}
+    sched.chunk_transfer_adapter.requests_with_ready_chunks = {session.request_id}
+    sched.chunk_transfer_adapter.requests_num_chunks_sent = {session.external_req_id: 2}
+
+    OmniSchedulerMixin._reset_ready_async_chunk_replacements(sched)
+    assert session.num_stale_output_tokens == 1 and session.drop_stale_output
+    assert session.num_output_placeholders == 0
+    sched._release_replaced_streaming_prompt_cache.assert_called_once_with(session)
+    assert _run_step(sched, session, num_scheduled=1, token=42) is False
+    assert session.num_stale_output_tokens == 0
+    assert _run_step(sched, session, num_scheduled=1, token=43) is True
+
+
 def test_exact_drain_delivers_new_segment_frame() -> None:
     """One in-flight decode frame at replacement: its late output drains the
     counter exactly and the new segment's first frame is delivered."""
@@ -226,3 +249,80 @@ def test_in_flight_prefill_chunk_drains_exactly_without_underflow() -> None:
     assert _run_step(sched, session, num_scheduled=5, token=42) is False  # late prefill chunk dropped
     assert session.num_stale_output_tokens == 0
     assert _run_step(sched, session, num_scheduled=1, token=43) is True  # new segment survives
+
+
+@pytest.mark.parametrize("drop", [False, True])
+def test_preemption_stale_output_follows_upstream_delivery_policy(drop):
+    request = _make_session()
+    request.status = RequestStatus.PREEMPTED
+    request.num_stale_output_tokens = 1
+    request.num_in_flight_tokens = 1
+    request.drop_stale_output = drop
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=1, token=42) is (not drop)
+    assert request.num_stale_output_tokens == 0
+    if not drop:
+        sched._update_request_with_output.assert_called_once_with(request, [42], is_stale=True)
+
+
+def test_preemption_stale_spec_rejection_does_not_roll_back_resumed_counters():
+    request = _make_session()
+    request.num_stale_output_tokens = 2
+    request.num_in_flight_tokens = 2
+    request.num_computed_tokens = 3
+    request.num_output_placeholders = 2
+    request.drop_stale_output = False
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=2, token=42, spec=[41])
+    assert request.num_computed_tokens == 3
+    assert request.num_output_placeholders == 2
+
+
+@pytest.mark.parametrize(
+    "arch,supports_reset,running,reset_running,blocked",
+    [
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, True, True),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, False, True, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, False, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", True, True, True, False),
+        ("AnotherStatefulAudioModel", False, True, True, True),
+        ("AnotherStatefulAudioModel", True, True, True, False),
+    ],
+)
+def test_reset_running_streaming_codec_is_rejected_before_preemption(
+    mocker, arch, supports_reset, running, reset_running, blocked
+):
+    from vllm.config import VllmConfig
+    from vllm.v1.core.sched.scheduler import Scheduler
+
+    from vllm_omni.config.model import OmniModelConfig
+
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.vllm_config = mocker.Mock(
+        spec=VllmConfig,
+        model_config=mocker.Mock(
+            spec=OmniModelConfig,
+            model_arch=arch,
+            supports_running_prefix_cache_reset=supports_reset,
+        ),
+    )
+    request = _make_session()
+    request.num_in_flight_tokens = 1
+    request.num_output_placeholders = 1
+    scheduler.running = [request] if running else []
+
+    def reset_cache(_self, reset_running_requests=False, reset_connector=False):
+        assert reset_running_requests is reset_running and reset_connector is True
+        return True
+
+    reset = mocker.patch.object(Scheduler, "reset_prefix_cache", autospec=True, side_effect=reset_cache)
+
+    assert scheduler.reset_prefix_cache(reset_running, reset_connector=True) is (not blocked)
+    if blocked:
+        reset.assert_not_called()
+        assert scheduler.running == [request]
+        assert request.status == RequestStatus.RUNNING
+        assert request.num_in_flight_tokens == request.num_output_placeholders == 1
+        assert not request.drop_stale_output
+    else:
+        reset.assert_called_once()

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
 
@@ -176,6 +177,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}
@@ -313,6 +315,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}
@@ -401,6 +404,7 @@ def test_stale_async_frame_is_dropped_before_output_processing() -> None:
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
 
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}

tests/core/sched/test_omni_sched_deferred_free_fence.py

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

tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -81,7 +81,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/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -1048,3 +1048,83 @@ def test_minimax_music3_speech_invalid_field_values(
     online_client.send_audio_speech_http_request(
         {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
     )
+
+
+# ─── POST /v1/audio/speech · YuE2-3B ───
+
+# Text-to-music like MiniMax Music 3, with one twist: the ABC score span is
+# request-controlled via extra_params.cot/abc and the two must agree, and
+# sampling is pinned by the checkpoint preset so sampling knobs are rejected.
+_YUE2_SPEECH = [
+    pytest.param(
+        OmniServerParams(
+            model="m-a-p/YuE2-3B",
+            stage_config_path=get_deploy_config_path("yue2.yaml"),
+            server_args=_SPEECH_SERVER_ARGS,
+        ),
+        id="yue2",
+        marks=hardware_marks(res={"cuda": "H100"}),
+    ),
+]
+
+_YUE2_LYRICS = "一闪一闪亮晶晶\n满天都是小星星"
+_YUE2_CAPTION = "Chinese heavy metal, distorted electric guitars, 140 BPM, key of C."
+_YUE2_ABC = "X:1\nT:Twinkle\nM:4/4\nL:1/4\nK:C\nC C G G|A A G2|"
+# 25 frames per second; the semantic preset bounds the frame budget.
+_YUE2_MAX_FRAMES = 9000
+
+
+@pytest.mark.parametrize(
+    "overrides, err_message",
+    [
+        pytest.param({"instructions": None}, ("instructions",), id="caption_missing"),
+        pytest.param({"instructions": "   "}, ("instructions",), id="caption_blank"),
+        pytest.param({"input": "   "}, ("input",), id="lyrics_blank"),
+        pytest.param({"voice": "alloy"}, ("voice",), id="voice_rejected"),
+        pytest.param({"speed": 1.5}, ("speed",), id="speed_rejected"),
+        pytest.param(
+            {"max_new_tokens": _YUE2_MAX_FRAMES + 1},
+            ("max_new_tokens", str(_YUE2_MAX_FRAMES)),
+            id="max_new_tokens_over_cap",
+        ),
+        pytest.param({"max_new_tokens": 199}, ("max_new_tokens",), id="max_new_tokens_below_min"),
+        pytest.param(
+            {"extra_params": {"cot": "melody"}},
+            ("abc",),
+            id="cot_melody_without_abc",
+        ),
+        pytest.param(
+            {"extra_params": {"abc": _YUE2_ABC}},
+            ("cot",),
+            id="abc_without_cot",
+        ),
+        pytest.param(
+            {"extra_params": {"temperature": 0.3}},
+            ("temperature",),
+            id="temperature_rejected",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _YUE2_SPEECH, indirect=True)
+def test_yue2_speech_invalid_field_values(
+    omni_server: OmniServer,
+    online_client: OnlineOmniClient,
+    overrides: dict[str, object],
+    err_message: str | tuple[str, ...],
+) -> None:
+    body: dict[str, Any] = {
+        "model": omni_server.model,
+        "input": _YUE2_LYRICS,
+        "instructions": _YUE2_CAPTION,
+        "seed": 1,
+        "max_new_tokens": 250,
+        "response_format": "wav",
+    }
+    for key, value in overrides.items():
+        if value is None:
+            body.pop(key, None)
+        else:
+            body[key] = value
+    online_client.send_audio_speech_http_request(
+        {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )

tests/diffusion/ar_diffusion/test_capability_runner.py

@@ -376,10 +376,17 @@ def test_lru_eviction_releases_blocks_and_notifies_pipeline():
 
 def test_budget_reduced_capacity_drives_runner_lru():
     pipeline = CapablePipeline(tiny_spec(capacity=2))
-    # One tiny LingBot-like session requires 128 bytes; two require 192.
+    # One tiny LingBot-like session requires 1,808 bytes; two require 2,592.
+    #
+    # These read 128 and 192 while the paging unit was the frame. tiny_spec
+    # declares one token per frame, so it used to page a single token at a
+    # time -- a block size no attention kernel accepts. Paging at 16 makes
+    # each page sixteen times larger. The block *counts* are unchanged, and
+    # so is what this test is about: a budget that fits fewer sessions than
+    # the pipeline asked for must drive the runner's LRU.
     runner = make_runner(
         pipeline,
-        available_bytes=128,
+        available_bytes=2048,
         gpu_memory_fraction=1.0,
     )
     kv = runner.kv_cache
@@ -665,3 +672,106 @@ def test_pipeline_without_warmup_provider_is_safely_skipped(monkeypatch):
     monkeypatch.setattr(runner, "execute_model", fail_if_called)
     runner._warmup_ar_rollout()
     assert execute.called is False
+
+
+# ── paging unit vs eviction unit ────────────────────────────────────────────
+
+
+def test_a_frame_that_is_already_a_legal_block_keeps_its_paging():
+    """Every resolution that runs today must page exactly as it does today."""
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    for tokens_per_frame in (16, 320, 1440, 4096):
+        assert tokens_per_frame % 16 == 0
+        assert paging_block_size(tokens_per_frame) == tokens_per_frame
+
+
+def test_a_frame_the_kernel_would_reject_pages_at_the_finest_legal_unit():
+    """832x480 is the checkpoint's own default and the kernel rejects it.
+
+    (480/16) x (832/16) = 1560 tokens per frame, and 1560 % 16 == 8, so a
+    frame-sized block is not something FlashAttention's paged kernel accepts.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    assert 1560 % 16 == 8
+    assert paging_block_size(1560) == 16
+    # And whatever it returns is always something the kernel takes.
+    for tokens_per_frame in range(1, 200):
+        assert paging_block_size(tokens_per_frame) % 16 == 0
+
+
+# ── the kernel's constraint is data, not a constant ─────────────────────────
+
+
+def _multiple_of(base):
+    from vllm.v1.attention.backend import MultipleOf
+
+    return MultipleOf(base)
+
+
+def test_the_dispatching_module_answers_in_vllms_own_vocabulary():
+    """AR-Diffusion calls flash_attn itself, so no backend can be asked.
+
+    The answer has to come from the module that picks the kernel, and it has to
+    have the shape vLLM uses, so the policy that consumes it does not care
+    where it came from.
+    """
+    from vllm.v1.attention.backend import MultipleOf
+
+    from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import supported_kernel_block_sizes
+
+    advertised = supported_kernel_block_sizes()
+    assert advertised, "a kernel that accepts nothing cannot be paged for"
+    for entry in advertised:
+        assert isinstance(entry, int | MultipleOf)
+
+
+def test_a_kernel_that_takes_one_fixed_size_gets_that_size():
+    """hpc_attn advertises [64] -- not a multiple, a single legal value.
+
+    A frame of 1560 tokens is not 64, and 1560 % 64 == 24, so the frame cannot
+    be the block and the only thing left to page at is 64. Assuming multiples
+    of 16 here would hand the kernel a block size it rejects.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    assert 1560 % 64 == 24
+    assert paging_block_size(1560, [64]) == 64
+    assert paging_block_size(64, [64]) == 64
+    assert paging_block_size(128, [64]) == 64  # a legal multiple is still not the advertised size
+
+
+def test_a_kernel_with_large_pages_still_keeps_a_whole_frame_when_it_fits():
+    """FlashInfer offers 128-and-up pages on Blackwell and not on Hopper.
+
+    The same checkpoint on the two cards must therefore page differently, which
+    is exactly what a constant cannot express.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    hopper = [_multiple_of(16)]
+    blackwell = [_multiple_of(16), 128, 256]
+    assert paging_block_size(1440, hopper) == 1440  # a legal multiple of 16
+    assert paging_block_size(1440, blackwell) == 1440
+    assert paging_block_size(1560, hopper) == 16
+    assert paging_block_size(1560, blackwell) == 16  # finest legal unit, not the largest
+
+
+def test_whatever_comes_back_is_always_something_the_kernel_accepts():
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    for advertised, check in (
+        ([_multiple_of(16)], lambda n: n % 16 == 0),
+        ([64], lambda n: n == 64),
+        ([_multiple_of(32), 48], lambda n: n % 32 == 0 or n == 48),
+    ):
+        for tokens_per_frame in range(1, 300):
+            assert check(paging_block_size(tokens_per_frame, advertised))
+
+
+def test_a_kernel_that_advertises_nothing_is_an_error_not_a_guess():
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    with pytest.raises(ValueError, match="no legal block size"):
+        paging_block_size(1560, [])

tests/diffusion/ar_diffusion/test_kv_cache.py

@@ -625,6 +625,74 @@ def test_non_contiguous_branch_indices_rejected():
         )
 
 
+# --- eviction against the block table ---------------------------------------
+
+
+class _RecordingManager(ChunkWindowManager):
+    """Capture the block range eviction asks for, without a real block pool.
+
+    ``remove_skipped_blocks`` reads three attributes and calls one method. A
+    pool would add nothing: the bug under test is the conversion from the
+    spec's units into the block table's, and that is entirely visible here.
+    """
+
+    def __init__(self, spec):
+        self.kv_cache_spec = spec
+        self.sliding_window = spec.sliding_window
+        self.block_size = spec.block_size
+        self.freed: tuple[int, int] | None = None
+
+    def _remove_blocks_in_range(self, request_id, first_block, last_block):
+        self.freed = (first_block, last_block)
+
+
+def test_sink_is_preserved_when_a_frame_spans_many_blocks():
+    # The shipped geometry: 832x480 is 1560 tokens per frame, 1560 % 16 == 8 is
+    # not a legal block size, so paging falls back to 16 and a frame spans 98
+    # blocks. Every other eviction test uses chunk_size == BLOCK, where a frame
+    # is a block and a sink counted in frames indexes the block table correctly
+    # by accident -- which is why this went unnoticed.
+    chunk_size, sink_chunks, window_chunks = 1560, 9, 9
+    spec = ChunkWindowSpec(
+        block_size=16,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * chunk_size,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+    )
+    manager = _RecordingManager(spec)
+
+    # Far enough past sink + window that eviction has definitely started.
+    manager.remove_skipped_blocks("req", total_computed_tokens=30 * chunk_size)
+    assert manager.freed is not None
+    first_freed, last_freed = manager.freed
+
+    # The assertion is on the token boundary rather than a block count, so it
+    # stays meaningful if the paging unit changes again.
+    sink_tokens = sink_chunks * chunk_size
+    assert first_freed * spec.block_size >= sink_tokens, (
+        f"eviction starts at token {first_freed * spec.block_size}, inside the {sink_tokens}-token sink"
+    )
+    # And it must not over-protect: at most one straddling block beyond the sink.
+    assert (first_freed - 1) * spec.block_size < sink_tokens
+    # Only wholly-skipped blocks are released.
+    skipped = manager.get_num_skipped_tokens(30 * chunk_size)
+    assert last_freed * spec.block_size <= sink_tokens + skipped
+
+
+def test_sink_conversion_is_identity_when_a_frame_is_a_block():
+    # Every geometry that paged one frame per block must be untouched.
+    spec = make_spec(chunk_size=BLOCK, window_chunks=2, sink_chunks=1)
+    manager = _RecordingManager(spec)
+    manager.remove_skipped_blocks("req", total_computed_tokens=8 * BLOCK)
+    assert manager.freed is not None
+    first_freed, _ = manager.freed
+    assert first_freed == spec.sink_chunks
+
+
 def test_staged_window_reuse_matches_a_full_gather():
     """Reusing the staged history must be byte-identical to re-gathering the whole window.
 

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -23,6 +23,7 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
 )
 from vllm_omni.experimental.ar_diffusion.kv_cache import paged_attention as paged_attention_module
 from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ChunkWindowManager
 from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -36,10 +37,26 @@ NEG = "negative"
 
 
 def make_state(
-    *, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu"), reuse_history_staging=False
+    *,
+    num_layers=1,
+    window_chunks=2,
+    chunk_size=BLOCK,
+    sink_chunks=0,
+    reset_at_boundary=False,
+    dtype=torch.float32,
+    device=torch.device("cpu"),
+    reuse_history_staging=False,
 ):
+    """Build a cache. ``chunk_size`` defaults to the block size, but the two are
+    independent -- the shipped 832x480 gives 1560 tokens per frame against
+    16-token blocks, so a frame is 97.5 blocks."""
     cfg = ARDiffusionKVConfig(
-        enable=True, chunk_size=BLOCK, window_chunks=window_chunks, reuse_history_staging=reuse_history_staging
+        enable=True,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+        reuse_history_staging=reuse_history_staging,
     )
     kv = ARDiffusionKVCache(
         cfg,
@@ -109,11 +126,13 @@ def _commit_video_span(
     n_chunks: int,
     dtype: torch.dtype,
     device: torch.device,
+    chunk_size: int = BLOCK,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    ctx = st.get_kv_caches(kv_branch, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    span = n_chunks * chunk_size
+    ctx = st.get_kv_caches(kv_branch, seq_len=span, commit_current=True)[0].forward_ctx
     ctx.ensure_video_slots(device)
-    k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    k = torch.randn(1, span, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(1, span, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
     kv._k_pools[0][ctx.current_video_slot_mapping] = k[0]
     kv._v_pools[0][ctx.current_video_slot_mapping] = v[0]
     st.commit_paged_context(kv_branch)
@@ -547,6 +566,644 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(monkeypa
         torch._dynamo.reset()
 
 
+# ── a frame that is not a whole number of blocks ────────────────────────────
+
+# 24 tokens per chunk against 16-token blocks leaves 8 over, the same remainder
+# the shipped 832x480 produces (1560 % 16 == 8). Every fixture above holds
+# chunk_size == BLOCK, so nothing there can reach this path.
+RAGGED_CHUNK = 24
+
+
+@pytest.mark.parametrize("commit_current", [False, True])
+@pytest.mark.parametrize("action_len", [0, 3])
+def test_a_ragged_chunk_still_matches_the_dense_reference(commit_current, action_len):
+    """Attention reads a sequence, not a set of blocks.
+
+    With one committed chunk of 24 tokens the history stops 8 slots into its
+    last block. The committing path writes straight after it and is fine. The
+    scratch path used to restart at slot zero of a fresh region, which left
+    those 8 slots unwritten -- and since the kernel consumes the block table as
+    one contiguous run, it read them as if they were tokens and shifted every
+    token after them by 8. Shapes stayed correct throughout, so only the values
+    showed it.
+    """
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    # window_chunks=2 with one chunk of history keeps the whole sequence
+    # resident, so the window does not also round up here; that case is
+    # test_a_ragged_window_keeps_the_block_table_a_fixed_shape below.
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=RAGGED_CHUNK)
+
+    history_k, history_v = _commit_video_span(
+        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=RAGGED_CHUNK
+    )
+
+    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    assert ctx.start_offset == RAGGED_CHUNK % BLOCK == 8
+
+    current_k = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    action_k = action_v = None
+    if action_len:
+        ctx.ensure_action_slots(action_len, device)
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
+        kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
+
+    query = torch.randn(1, RAGGED_CHUNK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=action_len, query_len=query.shape[1], device=device
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+
+    new_k = torch.cat([history_k, current_k], dim=1)
+    new_v = torch.cat([history_v, current_v], dim=1)
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+
+def test_a_ragged_scratch_chunk_is_physically_continuous_with_its_history():
+    """The write targets themselves must leave no hole.
+
+    This is the property the values above depend on, asserted directly so a
+    regression names the cause rather than showing a numeric mismatch.
+    """
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=2, chunk_size=RAGGED_CHUNK)
+    _commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=torch.float32, device=device, chunk_size=RAGGED_CHUNK)
+
+    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+
+    history_tail_block = ctx.history_block_ids[-1]
+    # The chunk starts inside the history's own last block, not at a fresh one.
+    assert ctx.current_video_block_ids[0] == history_tail_block
+    first_slot = int(ctx.current_video_slot_mapping[0])
+    assert first_slot == history_tail_block * BLOCK + 8
+
+    # And the action region must not be charged for that managed block.
+    ctx.ensure_action_slots(3, device)
+    assert ctx._scratch_blocks_used == len(ctx.current_video_block_ids) - 1
+    assert ctx.action_scratch_block_ids[0] not in ctx.current_video_block_ids
+
+
+def test_a_ragged_window_keeps_the_block_table_a_fixed_shape():
+    """The table's shape must not track how full the window is.
+
+    With a sink the visible window is 2*24 + 24 = 72 tokens, which is 4.5
+    blocks. Deriving the width by flooring that would understate the capacity,
+    the max() against the live block count would take over, and the shape would
+    change as the window filled -- recompiling the graph. The same floor would
+    also let max_seq_len come out under the kv_len actually passed.
+    """
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK)
+    # The window is deliberately not a whole number of blocks.
+    max_video_tokens = kv.spec.sliding_window + kv.spec.sink_chunks * kv.spec.chunk_size
+    assert max_video_tokens % BLOCK != 0
+
+    widths: set[int] = set()
+    kv_lens: list[int] = []
+    for _ in range(6):
+        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(device)
+        block_table, _, seq_lens, _, max_seq_len = ctx.build_block_table(
+            action_len=0, query_len=RAGGED_CHUNK, device=device
+        )
+        widths.add(int(block_table.shape[1]))
+        # A bound handed to the kernel has to actually bound the sequence.
+        assert int(seq_lens[0]) <= max_seq_len
+        kv_lens.append(int(seq_lens[0]))
+        st.commit_paged_context(POS)
+
+    assert len(widths) == 1, f"block table width varied across ticks: {sorted(widths)}"
+    # The window has to have actually filled, or none of the above was tested.
+    assert max(kv_lens) > min(kv_lens)
+    assert max(kv_lens) <= ctx.max_video_blocks * BLOCK
+    # Whole blocks are read, so each of the two boundaries -- the sink's end and
+    # the recent window's start -- can contribute up to BLOCK - 1 tokens the
+    # window no longer needs.
+    assert max(kv_lens) <= max_video_tokens + 2 * (BLOCK - 1)
+
+
+def _poison_pools(kv: ARDiffusionKVCache) -> None:
+    for pool in (*kv._k_pools, *kv._v_pools):
+        pool.fill_(float("nan"))
+
+
+def _write_prepared_chunk(
+    kv: ARDiffusionKVCache, ctx, *, dtype: torch.dtype, device: torch.device
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Write fresh K/V for a prepared chunk, poisoning the rest of its blocks first.
+
+    Blocks are handed out whole and a reused block keeps whatever it held, so a
+    slot read without having been written would not look unusual. Every slot of
+    the chunk's blocks that this chunk does not write is set to NaN first --
+    except the leading slots of a block the history still occupies -- so reading
+    one turns the attention output NaN.
+    """
+    mapping = ctx.current_video_slot_mapping
+    first_slot = int(mapping[0])
+    for index, block in enumerate(ctx.current_video_block_ids):
+        start = first_slot if index == 0 else block * BLOCK
+        kv._k_pools[0][start : (block + 1) * BLOCK] = float("nan")
+        kv._v_pools[0][start : (block + 1) * BLOCK] = float("nan")
+    k = torch.randn(ctx.seq_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(ctx.seq_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][mapping] = k
+    kv._v_pools[0][mapping] = v
+    return k, v
+
+
+def _window_positions(end: int, *, sink: int, window: int, written, history: int, resident: set[int]) -> list[int]:
+    """Token positions attention must read at ``end``.
+
+    Every token of the current chunk, plus every written history token in a
+    block that is still resident and holds a sink token or one of the window's
+    most recent tokens.
+    """
+    start = end - window
+    blocks = {p // BLOCK for p in range(min(sink, end))} | {p // BLOCK for p in range(max(start, 0), end)}
+    return [p for p in sorted(written) if p >= history or (p // BLOCK in blocks and p // BLOCK in resident)]
+
+
+def _model_block_indices(kv, adapter) -> dict[int, int]:
+    """Each resident block's index in model positions, by block id.
+
+    Compaction removes the evicted gap after the sink's blocks and shifts the
+    rest of the table down by it, so a table index past the sink is
+    ``compacted_tokens`` behind the position it holds.
+    """
+    block_size = kv.block_size
+    sink_blocks = -(-(kv.spec.sink_chunks * kv.spec.chunk_size) // block_size)
+    shift = adapter.compacted_tokens // block_size
+    return {
+        int(block): index if index < sink_blocks else index + shift
+        for index, block in enumerate(kv.block_table(adapter))
+        if block != kv.null_block_id
+    }
+
+
+@pytest.mark.parametrize(
+    ("sink_chunks", "window_chunks", "reset_at_boundary"),
+    [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)],
+)
+def test_attention_after_eviction_reads_every_kept_token_and_nothing_unwritten(
+    sink_chunks, window_chunks, reset_at_boundary
+):
+    """Once the window slides, the table must still cover the window exactly.
+
+    The sink and the recent window are separate token ranges, and each can
+    straddle a block edge on its own: 24-token chunks against 16-token blocks
+    put a boundary 8 tokens past an edge on alternating chunks. Counting blocks
+    off the ends of the resident list dropped the block holding the window's
+    first tokens whenever both boundaries straddled, and the returned length
+    ran past the last written slot. Neither changes a shape, so this compares
+    values against dense attention over the tokens the window keeps -- after
+    eviction, with every unwritten slot poisoned, on both the scratch and the
+    committing forward of every chunk.
+
+    The reset case keeps only the sink across a boundary, so the recent window
+    reaches back over blocks that are no longer resident. Those must be skipped
+    by position rather than read.
+
+    Eight chunks are enough for compaction to drop the evicted gap from the
+    table, after which storage positions run behind model positions. Removing
+    a gap that is not a whole number of chunks moves where eviction snaps, and
+    in the reset case the window then keeps tokens a boundary should have
+    dropped.
+    """
+    torch.manual_seed(0)
+    device, dtype = torch.device("cpu"), torch.float32
+    kv, st = make_state(
+        device=device,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+        chunk_size=RAGGED_CHUNK,
+    )
+    _poison_pools(kv)
+    sink, window = sink_chunks * RAGGED_CHUNK, int(kv.spec.sliding_window)
+    committed: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
+    evicted = compacted = False
+
+    for chunk in range(8):
+        history = int(st.adapter(POS).absolute_num_computed_tokens)
+        end = history + RAGGED_CHUNK
+        for commit_current in (False, True):
+            ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
+            ctx.ensure_video_slots(device)
+            k, v = _write_prepared_chunk(kv, ctx, dtype=dtype, device=device)
+            written = dict(committed)
+            written.update({history + i: (k[i], v[i]) for i in range(RAGGED_CHUNK)})
+
+            query = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+            block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+                action_len=0, query_len=RAGGED_CHUNK, device=device
+            )
+            resident = set(_model_block_indices(kv, st.adapter(POS)).values())
+            positions = _window_positions(
+                end, sink=sink, window=window, written=written, history=history, resident=resident
+            )
+            label = f"chunk {chunk}, commit_current={commit_current}"
+            assert int(seq_lens[0]) == len(positions), f"{label}: kv_len does not match the written slots it covers"
+            paged = ar_diffusion_paged_attention(
+                query,
+                kv.key_cache(0),
+                kv.value_cache(0),
+                block_table=block_table,
+                query_start_loc=query_start_loc,
+                seq_lens=seq_lens,
+                max_query_len=max_query_len,
+                max_seq_len=max_seq_len,
+                softmax_scale=HEAD_DIM**-0.5,
+                causal=False,
+            )
+            ref = _dense_attention(
+                query,
+                torch.stack([written[p][0] for p in positions]).unsqueeze(0),
+                torch.stack([written[p][1] for p in positions]).unsqueeze(0),
+            )
+            torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5, msg=lambda m, label=label: f"{label}: {m}")
+        committed = written
+        st.commit_paged_context(POS)
+        compacted = compacted or st.adapter(POS).compacted_tokens > 0
+        evicted = evicted or compacted or kv.null_block_id in kv.block_table(st.adapter(POS))
+
+    assert evicted, "the window never slid, so nothing after eviction was tested"
+    assert compacted, "the table was never compacted, so storage and model positions never diverged"
+
+
+@pytest.mark.parametrize(
+    ("sink_chunks", "window_chunks", "reset_at_boundary"),
+    [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)],
+)
+def test_compaction_does_not_change_which_tokens_stay_resident(
+    monkeypatch, sink_chunks, window_chunks, reset_at_boundary
+):
+    """Compaction renumbers the block table; it must not change what eviction keeps.
+
+    Storage positions run ``compacted_tokens`` behind model positions, and
+    eviction snaps to chunk boundaries counted in storage positions. Removing a
+    gap that is not a whole number of chunks moves every later snap, so eviction
+    keeps a different set of tokens -- in the reset case, tokens a boundary
+    should have dropped, which the window then reads. Attention checked against
+    what is resident cannot see that, so the oracle here is the same rollout
+    with compaction switched off.
+    """
+
+    def resident_blocks_per_tick(compact: bool) -> tuple[list[set[int]], int]:
+        with monkeypatch.context() as patch:
+            if not compact:
+                patch.setattr(ChunkWindowManager, "compact_block_table", lambda self, request_id: 0)
+            kv, st = make_state(
+                window_chunks=window_chunks,
+                sink_chunks=sink_chunks,
+                reset_at_boundary=reset_at_boundary,
+                chunk_size=RAGGED_CHUNK,
+            )
+            ticks = []
+            for _ in range(12):
+                ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+                ctx.ensure_video_slots(torch.device("cpu"))
+                ticks.append(set(_model_block_indices(kv, st.adapter(POS)).values()))
+                st.commit_paged_context(POS)
+            return ticks, st.adapter(POS).compacted_tokens
+
+    expected, _ = resident_blocks_per_tick(compact=False)
+    actual, compacted_tokens = resident_blocks_per_tick(compact=True)
+    assert compacted_tokens > 0, "the table was never compacted, so nothing was tested"
+    for tick, (want, got) in enumerate(zip(expected, actual)):
+        assert got == want, f"tick {tick}: compaction kept blocks {sorted(got - want)} and dropped {sorted(want - got)}"
+
+
+def test_history_staging_holds_a_ragged_window_and_restages_it_whole(monkeypatch):
+    """The staging buffers must fit the table width this module pads to.
+
+    A sink or window that is not a whole number of blocks spans one more block
+    than its token count suggests, so sizing staging from tokens left it short
+    and the first prepared forward refused to run. And because such a chunk
+    shares its first block with the history's tail, the history a probe sees is
+    not left untouched by the next probe, so it is restaged rather than reused.
+    """
+    monkeypatch.setenv(KV_GATHER_ENV, "1")
+    device = torch.device("cpu")
+    kv, st = make_state(
+        device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK, reuse_history_staging=True
+    )
+    capacity = int(kv.history_staging[0][0].shape[0])
+    for chunk in range(6):
+        for probe in range(2):
+            ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=False)[0].forward_ctx
+            ctx.prepare(device=device, action_len=0, query_len=RAGGED_CHUNK)
+            assert int(ctx.max_seq_len) <= capacity
+            assert not ctx.reuse_history, f"chunk {chunk}, probe {probe}: a ragged chunk reused the staged history"
+        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+        ctx.prepare(device=device, action_len=0, query_len=RAGGED_CHUNK)
+        st.commit_paged_context(POS)
+
+
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_action_tokens_after_a_partly_written_video_block_are_refused(commit_current):
+    """Action K/V follows the video blocks in the table, and the kernel reads it as one run.
+
+    A 24-token chunk from an empty history ends 8 slots into its second 16-token
+    block, so the run would read those 8 unwritten slots as the first action
+    tokens and never reach the last ones. Every shape stays right, which is why
+    this is refused rather than computed.
+    """
+    _, st = make_state(window_chunks=2, chunk_size=RAGGED_CHUNK)
+    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
+    with pytest.raises(ValueError, match="partly written video block"):
+        ctx.build_block_table(action_len=3, query_len=RAGGED_CHUNK + 3, device=torch.device("cpu"))
+
+
+def test_the_shipped_geometry_reads_exactly_the_window_on_every_tick():
+    """832x480 with the checkpoint's own sink and window, slot by slot.
+
+    A 9-frame sink and a 9-frame window are each 8 tokens past a 16-token edge
+    but 28080 tokens together -- a whole number of blocks -- so a width taken
+    from their sum was one block short on every tick once the window had
+    filled. Twelve ticks of three frames, checking every slot the kernel reads
+    against the token that slot is supposed to hold. No attention is computed:
+    at this size the bookkeeping is the whole question.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    frame, frames_per_tick = (480 // 16) * (832 // 16), 3
+    block = paging_block_size(frame)
+    seq_len = frames_per_tick * frame
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=frame, window_chunks=9, sink_chunks=9)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=1,
+        num_kv_heads=1,
+        head_size=8,
+        dtype=torch.float16,
+        block_size=block,
+        max_model_len=1 << 20,
+        available_bytes=1 << 27,
+        kv_branches=(ARDiffusionKVBranchSpec(POS, 0),),
+        session_capacity=1,
+        frames_per_block=frames_per_tick,
+        max_scratch_tokens_per_branch=seq_len,
+        device=torch.device("cpu"),
+    )
+    st = ARDiffusionKVState(kv, "s1", {POS: kv.begin_request("r-pos")}, num_layers=1)
+    sink, window = 9 * frame, int(kv.spec.sliding_window)
+    holds: dict[int, int] = {}  # slot -> token position last written there
+    widths: set[int] = set()
+    filled = compacted = False
+
+    for tick in range(12):
+        history = int(st.adapter(POS).absolute_num_computed_tokens)
+        end = history + seq_len
+        ctx = st.get_kv_caches(POS, seq_len=seq_len, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(torch.device("cpu"))
+        for offset, slot in enumerate(ctx.current_video_slot_mapping.tolist()):
+            holds[slot] = history + offset
+        table, _, seq_lens, _, _ = ctx.build_block_table(action_len=0, query_len=seq_len, device=torch.device("cpu"))
+        widths.add(int(table.shape[1]))
+        kv_len = int(seq_lens[0])
+
+        position_of_block = {b: i * block for b, i in _model_block_indices(kv, st.adapter(POS)).items()}
+        read = []
+        for b in table[0].tolist()[: -(-kv_len // block)]:
+            assert b in position_of_block, f"tick {tick}: read block {b}, which holds no token of this session"
+            read.extend((b * block + o, position_of_block[b] + o) for o in range(block))
+        read = read[:kv_len]
+        wrong = [(slot, want) for slot, want in read if holds.get(slot) != want]
+        assert not wrong, f"tick {tick}: {len(wrong)} slots read hold a token other than the one their position implies"
+
+        start = end - window
+        kept = set(range(min(sink, end))) | set(range(max(start, 0), end))
+        missed = kept - {want for _, want in read}
+        assert not missed, f"tick {tick}: {len(missed)} tokens the window keeps were not read"
+
+        filled = filled or end > sink + window
+        st.commit_paged_context(POS)
+        compacted = compacted or st.adapter(POS).compacted_tokens > 0
+
+    assert filled, "the window never filled, so the case that fails was never reached"
+    assert compacted, "the table was never compacted, so storage and model positions never diverged"
+    assert len(widths) == 1, f"block table width varied across ticks: {sorted(widths)}"
+
+
+class _CountedBlockIds:
+    """Wraps a vLLM block manager's ``get_blocks`` and counts the block ids read.
+
+    Every path to a request's block ids goes through ``get_blocks``, so a copy of
+    the whole table is counted in full even if only a few entries are indexed.
+    """
+
+    def __init__(self, manager):
+        self.reads = 0
+        self._get_blocks = manager.get_blocks
+
+    def __call__(self, request_id):
+        from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+
+        blocks = self._get_blocks(request_id)
+        return KVCacheBlocks(tuple([_CountedBlock(block, self) for block in group] for group in blocks.blocks))
+
+
+class _CountedBlock:
+    def __init__(self, block, counter: _CountedBlockIds):
+        self._block = block
+        self._counter = counter
+
+    @property
+    def block_id(self) -> int:
+        self._counter.reads += 1
+        return self._block.block_id
+
+
+def test_reading_the_window_touches_only_the_blocks_it_can_keep(monkeypatch):
+    """However long the block table is, reading the window must not cost more.
+
+    Evicted positions stay in the table as null entries until compaction drops
+    them a whole number of chunks at a time -- up to two frames' worth of blocks
+    at 832x480 -- so the table can still be far longer than the window. Only the
+    sink's blocks and the recent window's blocks can hold a kept token, so those
+    are all the read may touch. Compaction is switched off here so the table
+    grows to many times the window, and reads are counted where the ids leave
+    vLLM's block manager.
+    """
+    monkeypatch.setattr(ChunkWindowManager, "compact_block_table", lambda self, request_id: 0)
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK)
+    reads: list[int] = []
+    table_lengths: list[int] = []
+    for _ in range(60):
+        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(device)
+        counted = _CountedBlockIds(kv.manager)
+        kv.manager.get_blocks = counted
+        try:
+            ctx.build_block_table(action_len=0, query_len=RAGGED_CHUNK, device=device)
+        finally:
+            del kv.manager.get_blocks
+        reads.append(counted.reads)
+        table_lengths.append(len(kv.block_table(st.adapter(POS))))
+        st.commit_paged_context(POS)
+
+    assert table_lengths[-1] > 10 * ctx.max_video_blocks, "the table never outgrew the window, so nothing was tested"
+    assert max(reads) <= ctx.max_video_blocks, (
+        f"read {max(reads)} table entries for a {ctx.max_video_blocks}-block window"
+    )
+
+
+def test_the_checkpoints_own_default_resolution_can_build_a_cache():
+    """832x480 is what LingBot World v2 ships as its default, and it could not run.
+
+    A frame-sized block made the resolution a kernel-compatibility question:
+    1560 tokens per frame is not a multiple of 16, so FlashAttention's paged
+    kernel rejected it and the default resolution had no realtime path at all.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    tokens_per_frame = (480 // 16) * (832 // 16)
+    assert tokens_per_frame == 1560
+    assert tokens_per_frame % 16 == 8, "the whole point is that a frame is not a legal block"
+
+    block_size = paging_block_size(tokens_per_frame)
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=tokens_per_frame, window_chunks=2)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=1,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=torch.float32,
+        block_size=block_size,
+        max_model_len=1 << 16,
+        available_bytes=1 << 28,
+        kv_branches=(ARDiffusionKVBranchSpec(POS, 0),),
+        session_capacity=1,
+        frames_per_block=2,
+        max_scratch_tokens_per_branch=block_size,
+        device=torch.device("cpu"),
+    )
+
+    assert kv.block_size % 16 == 0
+    # The eviction unit is still the frame; only the paging unit changed.
+    assert kv.spec.chunk_size == tokens_per_frame
+    assert kv.blocks_per_frame == -(-tokens_per_frame // block_size) == 98
+
+
+def test_the_shipped_resolution_geometry_matches_dense_attention():
+    """The real number, not a scaled-down stand-in: 1560 tokens per frame.
+
+    This is the case the issue reported failing at ~4e-02 on the scratch path.
+    It runs on CPU because the reference is dense attention, not a kernel.
+    """
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    tokens_per_frame = (480 // 16) * (832 // 16)
+    assert tokens_per_frame == 1560
+
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=tokens_per_frame)
+    history_k, history_v = _commit_video_span(
+        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=tokens_per_frame
+    )
+
+    # The non-committing path is the one that was wrong: four of the five
+    # forwards per generated block take it.
+    ctx = st.get_kv_caches(POS, seq_len=tokens_per_frame, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    assert ctx.start_offset == 8, "1560 % 16 == 8 is the whole reason this case exists"
+
+    current_k = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    query = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=0, query_len=tokens_per_frame, device=device
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+    ref = _dense_attention(
+        query,
+        torch.cat([history_k, current_k], dim=1),
+        torch.cat([history_v, current_v], dim=1),
+    )
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_the_shipped_resolution_geometry_on_the_real_kernel(commit_current):
+    """The shipped 832x480 geometry through FlashAttention's paged kernel.
+
+    The CPU test above proves the addressing is right against a dense
+    reference. This one proves the real kernel agrees, at the real 1560
+    tokens per frame, through the production path -- host prep followed by
+    the fused write+attend op, which is what a forward actually calls.
+    """
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.float16
+    tokens_per_frame = (480 // 16) * (832 // 16)
+    assert tokens_per_frame == 1560
+
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=tokens_per_frame)
+    history_k, history_v = _commit_video_span(
+        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=tokens_per_frame
+    )
+
+    layer_ctx = st.get_kv_caches(POS, seq_len=tokens_per_frame, commit_current=commit_current)[0]
+    ctx = layer_ctx.forward_ctx
+    current_k = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    query = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    ctx.prepare(device=device, action_len=0, query_len=tokens_per_frame)
+    assert ctx.start_offset == 8, "1560 % 16 == 8 is the whole reason this case exists"
+    assert kv.block_size == 16
+
+    inputs = layer_ctx.to_layer_inputs()
+    paged = paged_write_attn(inputs, query[0], current_k[0], current_v[0], None, None, HEAD_DIM**-0.5).unsqueeze(0)
+
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
+
+
 @pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
 @pytest.mark.parametrize("history_chunks", [0, 1, 3])
 @pytest.mark.parametrize("action_len", [0, 3])

tests/diffusion/attention/test_attention_capabilities.py

@@ -0,0 +1,359 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import sys
+from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.backends.flash_attn import (
+    FlashAttentionBackend,
+    FlashAttentionImpl,
+)
+from vllm_omni.diffusion.attention.backends.utils import fa
+from vllm_omni.diffusion.attention.capabilities import (
+    CapabilityResult,
+    CompilationMode,
+    ExecutionContext,
+    OuterBoundary,
+    ParallelStrategy,
+    SupportStatus,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_validate_fa4_head_dims = fa.validate_fa4_head_dims
+
+
+@pytest.fixture(autouse=True)
+def kernel_validator(monkeypatch):
+    # Capability unit tests do not require FA4 or a CUDA device. Real-kernel
+    # validation is exercised by test_flash_attn_compile.py.
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: True)
+
+
+def _impl(*, kernel_variant: str | None = "fa4", causal: bool = False):
+    impl = FlashAttentionImpl.__new__(FlashAttentionImpl)
+    impl._kernel_variant = kernel_variant
+    impl.causal = causal
+    return impl
+
+
+def _resolve(
+    impl,
+    *,
+    context: ExecutionContext | None = None,
+    attn_metadata: AttentionMetadata | None = None,
+    dtype=torch.bfloat16,
+):
+    tensor = torch.empty((1, 16, 8, 64), dtype=dtype)
+    return impl.resolve_execution_path(
+        context or ExecutionContext(platform="cuda"),
+        tensor,
+        tensor,
+        tensor,
+        attn_metadata,
+    )
+
+
+class _UnmigratedBackend(AttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "UNMIGRATED_TEST"
+
+    @staticmethod
+    def get_impl_cls():
+        return None
+
+    @staticmethod
+    def get_metadata_cls():
+        return None
+
+    @staticmethod
+    def get_builder_cls():
+        return None
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return []
+
+
+def test_capability_result_is_tri_state():
+    assert CapabilityResult.supported().status is SupportStatus.SUPPORTED
+    assert CapabilityResult.unsupported("not available").status is SupportStatus.UNSUPPORTED
+    assert CapabilityResult.unmigrated().status is SupportStatus.UNMIGRATED
+    with pytest.raises(ValueError, match="actionable reason"):
+        CapabilityResult(SupportStatus.UNSUPPORTED)
+
+
+def test_unmigrated_backend_default_is_conservative():
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    result = _UnmigratedBackend.resolve_capabilities(context)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+    requested = result.requested_support(context)
+    assert requested.status is SupportStatus.UNSUPPORTED
+    assert "no verified fullgraph declaration" in requested.reason
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.UNMIGRATED
+
+
+@pytest.mark.parametrize("mode", list(CompilationMode))
+def test_fullgraph_request_checks_compilation_mode(mode):
+    result = replace(_resolve(_impl()), compilation_mode=mode)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    expected = SupportStatus.UNSUPPORTED if mode is CompilationMode.EAGER_ONLY else SupportStatus.SUPPORTED
+    assert result.requested_support(context).status is expected
+    assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.SUPPORTED
+
+
+def test_request_preserves_verified_rejection_reason():
+    result = replace(_resolve(_impl()), support=CapabilityResult.unsupported("unsupported head size"))
+    assert result.requested_support(ExecutionContext(platform="cuda", require_fullgraph=True)) == result.support
+
+
+def test_exact_cuda_fa4_dense_path_is_supported():
+    context = ExecutionContext(
+        platform="cuda",
+        require_fullgraph=True,
+    )
+    result = _resolve(_impl(), context=context)
+
+    assert result.path == "fa4_dense"
+    assert result.support.status is SupportStatus.SUPPORTED
+    assert result.compilation_mode is CompilationMode.CUSTOM_OP
+    assert result.requested_support(context).status is SupportStatus.SUPPORTED
+
+
+def test_causal_fa4_path_remains_unmigrated():
+    result = _resolve(_impl(causal=True))
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_fa4_delegates_dimensions_to_kernel(monkeypatch):
+    head_dims = (80, 48)
+    validator = Mock(return_value=True)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", validator)
+    query = torch.empty((1, 16, 8, head_dims[0]), dtype=torch.bfloat16)
+    value = torch.empty((1, 16, 8, head_dims[1]), dtype=torch.bfloat16)
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, query, value, None)
+    assert result.support.status is SupportStatus.SUPPORTED
+    validator.assert_called_once_with(*head_dims, 8)
+
+
+def test_fa4_reports_kernel_rejection(monkeypatch):
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", Mock(side_effect=AssertionError("kernel dimension constraint")))
+    result = _resolve(_impl())
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert "kernel dimension constraint" in result.support.reason
+    assert "another backend" in result.support.reason
+
+
+def test_fa4_missing_validator_remains_unmigrated(monkeypatch):
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: False)
+    result = _resolve(_impl())
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+
+
+def test_fa4_rejects_mismatched_query_key_dimensions():
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    key = torch.empty((1, 16, 8, 32), dtype=torch.bfloat16)
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, key, key, None)
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert "Q and K head dimensions must match" in result.support.reason
+
+
+@pytest.mark.parametrize("arch", [90, 100, 120])
+def test_fa4_validator_adapter_uses_selected_architecture(monkeypatch, arch):
+    validator = Mock()
+    monkeypatch.setitem(
+        sys.modules,
+        "flash_attn.cute.interface",
+        SimpleNamespace(_get_device_arch=lambda: arch, _validate_head_dims=validator),
+    )
+    # Call the real adapter, not the fixture's stand-in.
+    verified = _validate_fa4_head_dims(80, 48, 8)
+    if arch == 120:
+        assert not verified
+        validator.assert_not_called()
+    else:
+        assert verified
+        validator.assert_called_once_with(80, 48, arch // 10, 8)
+
+
+def test_fa4_validator_adapter_handles_missing_private_api(monkeypatch):
+    monkeypatch.setitem(sys.modules, "flash_attn.cute.interface", SimpleNamespace())
+    assert not _validate_fa4_head_dims(64, 64, 8)
+
+
+def test_initialized_kernel_identity_overrides_caller_claim():
+    result = _resolve(
+        _impl(kernel_variant=None),
+        context=ExecutionContext(
+            platform="cuda",
+            kernel_variant="fa4",
+        ),
+    )
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_preconstruction_without_resolved_cuda_kernel_is_unmigrated():
+    result = FlashAttentionBackend.resolve_capabilities(ExecutionContext(platform="cuda"))
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.path == "unverified"
+
+
+@pytest.mark.parametrize(
+    "changes",
+    [
+        {"platform": "rocm"},
+        {"paged_kv": True},
+        {"parallel_strategy": ParallelStrategy.ULYSSES},
+        {"parallel_strategy": ParallelStrategy.RING},
+        {"parallel_strategy": ParallelStrategy.HYBRID_ULYSSES_RING},
+        {"outer_boundaries": frozenset({OuterBoundary.HSDP})},
+    ],
+)
+def test_unverified_flash_attention_variants_remain_unmigrated(changes):
+    result = _resolve(
+        _impl(),
+        context=replace(ExecutionContext(platform="cuda"), **changes),
+    )
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+
+
+def test_unverified_dtype_remains_unmigrated():
+    result = _resolve(_impl(), dtype=torch.float32)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+@pytest.mark.parametrize(
+    "attn_metadata",
+    [
+        AttentionMetadata(full_attn_spans=[[(0, 8)]]),
+        AttentionMetadata(
+            extra={
+                "cu_seqlens_q": torch.tensor([0, 16], dtype=torch.int32),
+                "cu_seqlens_k": torch.tensor([0, 16], dtype=torch.int32),
+                "max_seqlen_q": 16,
+                "max_seqlen_k": 16,
+            }
+        ),
+        AttentionMetadata(extra={"kv_cache_dtype": "fp8"}),
+    ],
+)
+def test_metadata_normalization_prevents_dense_support(attn_metadata):
+    result = _resolve(_impl(), attn_metadata=attn_metadata)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_unpublished_all_true_mask_remains_runtime_dependent():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+    )
+    result = _resolve(_impl(), attn_metadata=metadata)
+
+    assert result.path == "runtime_mask_dependent"
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_producer_published_noop_mask_resolves_dense_fa4():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+    result = _resolve(_impl(), attn_metadata=metadata)
+
+    assert result.path == "fa4_dense"
+    assert result.support.status is SupportStatus.SUPPORTED
+
+
+def test_fixed_shape_mask_semantic_change_invalidates_dense_resolution():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+    impl = _impl()
+    dense = _resolve(impl, attn_metadata=metadata)
+    assert dense.support.status is SupportStatus.SUPPORTED
+
+    metadata.attn_mask[:, 8:] = False
+    metadata.extra["attention_mask_mode"] = "padding"
+    masked = _resolve(impl, attn_metadata=metadata)
+    assert masked.support.status is SupportStatus.UNMIGRATED
+    request = ExecutionContext(platform="cuda", require_fullgraph=True)
+    assert masked.requested_support(request).status is SupportStatus.UNSUPPORTED
+
+
+def test_piecewise_dispatch_ignores_unused_incomplete_packed_metadata(monkeypatch):
+    from vllm_omni.diffusion.attention.backends import flash_attn
+    from vllm_omni.diffusion.attention.backends.utils import fa
+
+    metadata = AttentionMetadata(full_attn_spans=[[(0, 16)]], extra={"max_seqlen_q": 16})
+    impl = _impl()
+    impl.fa_deterministic = False
+    impl.softmax_scale = 0.125
+    assert _resolve(impl, attn_metadata=metadata).support.status is SupportStatus.UNMIGRATED
+
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_func", lambda *args, **kwargs: query)
+    monkeypatch.setattr(flash_attn, "piecewise_attn", lambda *args, **kwargs: query)
+    assert impl.forward_cuda(query, query, query, metadata) is query
+
+
+@pytest.mark.parametrize("input_index", [1, 2])
+@pytest.mark.parametrize("mismatch", ["dtype", "device"])
+def test_fa4_rejects_incompatible_key_value(input_index, mismatch):
+    tensors = [torch.empty((1, 16, 8, 64), dtype=torch.bfloat16) for _ in range(3)]
+    tensors[input_index] = (
+        tensors[input_index].to(dtype=torch.float16) if mismatch == "dtype" else tensors[input_index].to("meta")
+    )
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), *tensors, None)
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert f"{mismatch}s must match" in result.support.reason
+
+
+@pytest.mark.parametrize("skip_quant", [False, True])
+def test_layer_resolution_uses_effective_kv_quantization(skip_quant):
+    from vllm_omni.diffusion.attention.layer import Attention
+    from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
+
+    layer = Attention.__new__(Attention)
+    torch.nn.Module.__init__(layer)
+    layer.attention = _impl(kernel_variant=None)
+    layer._hsdp_compile_boundary_enabled = False
+    layer.skip_sequence_parallel = False
+    layer._no_parallel_strategy = NoParallelAttention()
+    layer.parallel_strategy = layer._no_parallel_strategy
+    layer.use_ring = False
+    layer.paged_kv_cache_role = None
+    layer._kv_cache_dtype = "fp8"
+    layer._disable_kv_quant = False
+    layer._kv_cache_skip_steps = None
+    layer._kv_cache_skip_layers = {0} if skip_quant else None
+    layer.layer_idx = 0
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": "fp8"})
+    result = layer.resolve_execution_path(ExecutionContext(platform="npu"), query, query, query, metadata)
+    assert result.path == ("npu_dense" if skip_quant else "npu_unverified")
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert metadata.extra == {"kv_cache_dtype": "fp8"}

tests/diffusion/attention/test_attention_state_lifetime.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Stateful contract example, independent of an optional attention kernel.
+
+The prepared module owns tensor state and passes it explicitly to an opaque op.
+This deliberately does not model opaque Python/C++ planning handles.
+"""
+
+import gc
+import weakref
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@torch.library.custom_op("vllm_omni_test::planned_attention", mutates_args=())
+def _planned_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    query_scale: torch.Tensor,
+) -> torch.Tensor:
+    return F.scaled_dot_product_attention(query * query_scale, key, value).contiguous()
+
+
+@_planned_attention.register_fake
+def _planned_attention_fake(query, key, value, query_scale):
+    return query.new_empty((*query.shape[:-1], value.shape[-1]))
+
+
+class _PreparedAttention(torch.nn.Module):
+    """Initialization/planning runs eagerly; forward only executes tensors."""
+
+    def __init__(self, query_scale: torch.Tensor):
+        super().__init__()
+        self.register_buffer("query_scale", query_scale.detach().clone())
+
+    def forward(self, query, key, value):
+        return _planned_attention(query, key, value, self.query_scale)
+
+
+def _inputs(length):
+    generator = torch.Generator().manual_seed(length)
+    return tuple(torch.randn(2, 3, length, 8, generator=generator) for _ in range(3))
+
+
+def _reference(plan, inputs):
+    query, key, value = inputs
+    return F.scaled_dot_product_attention(query * plan.query_scale, key, value)
+
+
+def test_dynamic_replay_uses_each_instances_state_without_recompilation():
+    compile_count = 0
+
+    def counting_backend(graph_module, _example_inputs):
+        nonlocal compile_count
+        compile_count += 1
+        return graph_module.forward
+
+    # Same tensor metadata, different values: state must be a graph input,
+    # never a Python instance identifier or a captured first-instance constant.
+    plans = [_PreparedAttention(torch.full((1, 3, 1, 1), scale)) for scale in (0.5, 2.0)]
+    compiled = torch.compile(
+        lambda plan, q, k, v: plan(q, k, v),
+        backend=counting_backend,
+        fullgraph=True,
+        dynamic=True,
+    )
+    for length in (5, 9, 13):
+        inputs = _inputs(length)
+        for plan in plans:
+            torch.testing.assert_close(compiled(plan, *inputs), _reference(plan, inputs))
+        assert not torch.allclose(_reference(plans[0], inputs), _reference(plans[1], inputs))
+    assert compile_count == 1
+
+
+def test_compiled_callable_retains_state_until_released():
+    plan = _PreparedAttention(torch.full((1, 3, 1, 1), 0.5))
+    plan_ref = weakref.ref(plan)
+    state_ref = weakref.ref(plan.query_scale)
+    inputs = _inputs(5)
+    expected = _reference(plan, inputs)
+    compiled = torch.compile(plan, backend="eager", fullgraph=True, dynamic=True)
+    torch.testing.assert_close(compiled(*inputs), expected)
+
+    # Deleting the caller's reference must not invalidate subsequent replay.
+    del plan
+    gc.collect()
+    assert plan_ref() is not None
+    assert state_ref() is not None
+    torch.testing.assert_close(compiled(*inputs), expected)
+
+    del compiled
+    gc.collect()
+    assert plan_ref() is None
+    assert state_ref() is None
+
+
+def test_stateful_custom_op_schema_and_fake_implementation():
+    inputs = _inputs(5)
+    # Unequal Q/K lengths and V head dimensions exercise the fake output shape.
+    query, _, _ = inputs
+    key = torch.randn(2, 3, 7, 8)
+    value = torch.randn(2, 3, 7, 4)
+    scale = torch.full((1, 3, 1, 1), 0.5)
+    torch.library.opcheck(_planned_attention, (query, key, value, scale))
+
+
+def test_stateful_attention_compiles_with_inductor():
+    plan = _PreparedAttention(torch.full((1, 3, 1, 1), 0.5))
+    compiled = torch.compile(plan, fullgraph=True, dynamic=True)
+    for length in (5, 9):
+        inputs = _inputs(length)
+        torch.testing.assert_close(compiled(*inputs), _reference(plan, inputs))

tests/diffusion/attention/test_flash_attn.py

@@ -21,6 +21,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
 from vllm_omni.diffusion.attention.backends.utils import fa  # noqa: E402
+from vllm_omni.diffusion.attention.capabilities import CompilationMode, ExecutionContext, SupportStatus
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@@ -432,6 +433,35 @@ def test_varlen_masked_routing_by_role(monkeypatch):
     assert cu_seqlens_k.tolist() == [0, 2, 5]
 
 
+@pytest.mark.cpu
+def test_published_noop_mask_uses_dense_path(monkeypatch):
+    calls = []
+
+    def fake_dense_func(q, k, v, **_kwargs):
+        calls.append("dense")
+        return q
+
+    def fail_varlen(*_args, **_kwargs):
+        raise AssertionError("published no-op mask must not use varlen attention")
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "flash_attn_func", fake_dense_func)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", fail_varlen)
+
+    impl = FlashAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5, causal=False)
+    query = torch.randn(1, 4, 2, 8)
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 4), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+
+    output = impl.forward_cuda(query, query, query, metadata)
+
+    assert torch.equal(output, query)
+    assert calls == ["dense"]
+
+
 @pytest.mark.cpu
 def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
     calls = []
@@ -550,6 +580,199 @@ def _fake_torch_npu(monkeypatch, *, npu_fusion_attention=None):
     return fake
 
 
+@pytest.fixture
+def npu_contract(monkeypatch, mocker):
+    monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: pytest.fail("NPU must not validate FA4"))
+    query = torch.zeros(1, 8, 2, 4)
+    dense = mocker.Mock(return_value=torch.zeros_like(query))
+    packed = mocker.Mock(return_value=torch.zeros_like(query.squeeze(0)))
+    _fake_mindiesd(monkeypatch, attention_forward=dense, attention_forward_varlen=packed)
+    return _npu_impl(), query, dense, packed
+
+
+def _assert_platform_path(impl, platform, query, metadata, path):
+    context = ExecutionContext(platform=platform, require_fullgraph=True)
+    result = impl.resolve_execution_path(context, query, query, query, metadata)
+    assert result.path == path
+    assert result.kernel_variant == {"npu": "mindiesd", "rocm": "aiter"}[platform]
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+    assert result.requested_support(context).status is SupportStatus.UNSUPPORTED
+
+
+def _real_pad_metadata():
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    return AttentionMetadata(
+        extra={
+            "npu_attn_varlen": True,
+            "cu_seqlens_q": cu,
+            "cu_seqlens_k": cu,
+            "max_seqlen_q": 5,
+            "max_seqlen_k": 5,
+        }
+    )
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("kv_cache_dtype", [None, "float", "auto"])
+def test_npu_contract_dense_dispatch(npu_contract, kv_cache_dtype):
+    impl, query, dense, packed = npu_contract
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": kv_cache_dtype})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_dense")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    assert dense.call_args.kwargs["attn_mask"] is None
+
+
+@pytest.mark.cpu
+def test_npu_contract_mask_layout(npu_contract):
+    impl, query, dense, packed = npu_contract
+    mask = torch.arange(8)[None] < 5
+    metadata = AttentionMetadata(attn_mask=mask)
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    torch.testing.assert_close(dense.call_args.kwargs["attn_mask"], mask[:, None, None, :].expand(1, 1, 8, 8))
+
+
+@pytest.mark.cpu
+def test_npu_contract_packed_dispatch(npu_contract):
+    impl, query, dense, packed = npu_contract
+    metadata = _real_pad_metadata()
+    _assert_platform_path(impl, "npu", query, metadata, "npu_packed_varlen")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.args[3:] == ([0, 5, 8], [0, 5, 8])
+
+
+@pytest.mark.cpu
+def test_npu_contract_prefix_dispatch(npu_contract, monkeypatch):
+    impl, query, dense, packed = npu_contract
+    metadata = _real_pad_metadata()
+    _assert_platform_path(impl, "npu", query, metadata, "npu_packed_varlen")
+    # The same metadata resolves differently after the kernel choice changes.
+    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
+    _assert_platform_path(impl, "npu", query, metadata, "npu_prefix_kv_slice")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    torch.testing.assert_close(dense.call_args.args[1], query[:, :5])
+    torch.testing.assert_close(dense.call_args.args[2], query[:, :5])
+    assert dense.call_args.kwargs["attn_mask"] is None
+
+
+@pytest.mark.cpu
+def test_npu_contract_incomplete_packing_rebuilds_mask(npu_contract):
+    impl, query, dense, packed = npu_contract
+    # CUDA-like normalization rejects this metadata, but NPU can rebuild a mask.
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True, "max_seqlen_q": 5, "valid_kv_length": 5})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked_fallback")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    mask = (torch.arange(8) < 5)[None, None, None, :].expand(1, 1, 8, 8)
+    torch.testing.assert_close(dense.call_args.kwargs["attn_mask"], mask)
+
+
+@pytest.mark.cpu
+def test_npu_contract_unusable_packing_does_not_run_unmasked(npu_contract):
+    impl, query, dense, packed = npu_contract
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked_fallback")
+    with pytest.raises(ValueError, match="refusing to run unmasked attention"):
+        impl.forward_npu(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_not_called()
+
+
+@pytest.fixture
+def rocm_contract(monkeypatch, mocker):
+    query = torch.zeros(1, 8, 2, 4)
+    dense = mocker.Mock(return_value=torch.zeros_like(query))
+    packed = mocker.Mock(side_effect=lambda *args, **kwargs: torch.zeros_like(args[0] if args else kwargs["q"]))
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "IS_AITER", True)
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_func", dense)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", packed)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: pytest.fail("ROCm must not validate FA4"))
+    return _npu_impl(), query, dense, packed
+
+
+@pytest.mark.cpu
+def test_rocm_contract_dense_dispatch(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    _assert_platform_path(impl, "rocm", query, None, "rocm_dense")
+    impl.forward_hip(query, query, query)
+    dense.assert_called_once()
+    packed.assert_not_called()
+
+
+@pytest.mark.cpu
+def test_rocm_contract_masked_unpadding(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    metadata = AttentionMetadata(attn_mask=torch.arange(8)[None] < 5, extra={"attention_mask_mode": "padding"})
+    _assert_platform_path(impl, "rocm", query, metadata, "rocm_masked_varlen")
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.args[0].shape == (5, 2, 4)
+    torch.testing.assert_close(packed.call_args.kwargs["cu_seqlens_k"], torch.tensor([0, 5], dtype=torch.int32))
+
+
+@pytest.mark.cpu
+def test_rocm_contract_multiple_packed_documents(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    cu = torch.tensor([0, 2, 4, 8], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        extra={
+            "cu_seqlens_q": cu,
+            "cu_seqlens_k": cu,
+            "max_seqlen_q": 4,
+            "max_seqlen_k": 4,
+        }
+    )
+    _assert_platform_path(impl, "rocm", query, metadata, "rocm_packed_varlen")
+    # NPU's [real, pad] path cannot consume this packing contract.
+    assert impl._resolve_packed_seq_npu(query, query, metadata.extra) is None
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.kwargs["cu_seqlens_q"] is cu
+    assert packed.call_args.kwargs["cu_seqlens_k"] is cu
+
+
+@pytest.mark.cpu
+def test_rocm_contract_unknown_mask_stays_runtime_dependent(rocm_contract, monkeypatch):
+    impl, query, dense, packed = rocm_contract
+    metadata = AttentionMetadata(attn_mask=torch.ones(1, 8, dtype=torch.bool))
+    # Resolution may inspect shape/metadata, but must not read mask values.
+    with monkeypatch.context() as patch:
+        patch.setattr(torch, "any", lambda *_args, **_kwargs: pytest.fail("resolution read mask values"))
+        _assert_platform_path(impl, "rocm", query, metadata, "rocm_runtime_mask_dependent")
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+
+
+@pytest.mark.cpu
+def test_rocm_contract_requires_initialized_aiter(monkeypatch):
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "IS_AITER", False)
+    impl = _npu_impl()
+    query = torch.zeros(1, 8, 2, 4)
+    result = impl.resolve_execution_path(
+        ExecutionContext(platform="rocm", kernel_variant="aiter"), query, query, query, None
+    )
+    assert result.path == "rocm_unverified"
+    assert result.kernel_variant is None
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
 # --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
 
 

tests/diffusion/attention/test_flash_attn_compile.py

@@ -0,0 +1,308 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import functools
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.attention.backends.flash_attn import (
+    FlashAttentionBackend,
+    FlashAttentionImpl,
+)
+from vllm_omni.diffusion.attention.backends.utils import fa
+from vllm_omni.diffusion.attention.capabilities import (
+    CompilationMode,
+    ExecutionContext,
+    OuterBoundary,
+    ParallelStrategy,
+    SupportStatus,
+)
+from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.forward_context import ForwardContext
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
+
+
+@pytest.fixture(autouse=True)
+def isolated_compiler_cache():
+    # Each case changes kernel/configuration on the same Attention.forward code
+    # object. Keep unrelated cases from exhausting its Dynamo recompile limit;
+    # retain the cache across all shapes and replays within an individual test.
+    torch.compiler.reset()
+    yield
+    torch.compiler.reset()
+
+
+@dataclass
+class _ParallelStrategyStub:
+    name: str
+    enabled: bool = True
+
+
+def _make_attention_layer(monkeypatch, *, head_size=64):
+    from vllm_omni.diffusion.attention import layer as attention_layer
+
+    monkeypatch.setattr(
+        attention_layer,
+        "get_attn_backend_for_role",
+        lambda **_kwargs: (FlashAttentionBackend, None),
+    )
+    monkeypatch.setattr(
+        attention_layer,
+        "build_parallel_attention_strategy",
+        lambda **_kwargs: NoParallelAttention(),
+    )
+    return attention_layer.Attention(num_heads=8, head_size=head_size, softmax_scale=head_size**-0.5, causal=False)
+
+
+@pytest.fixture
+def fake_fa4(monkeypatch, tmp_path):
+    marker = tmp_path / "kernel"
+    marker.write_text("loaded", encoding="utf-8")
+
+    @functools.cache
+    def cached_kernel_loader():
+        with open(marker, encoding="utf-8") as handle:
+            handle.read()
+
+    def fake_attention(query, key, value, **kwargs):
+        cached_kernel_loader()
+        return (
+            torch.nn.functional.scaled_dot_product_attention(
+                query.transpose(1, 2),
+                key.transpose(1, 2),
+                value.transpose(1, 2),
+                scale=kwargs["softmax_scale"],
+                is_causal=kwargs["causal"],
+            )
+            .transpose(1, 2)
+            .contiguous()
+        )
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", True)
+    monkeypatch.setattr(fa, "flash_attn_func", fake_attention)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: True)
+
+    return FlashAttentionImpl(
+        num_heads=8,
+        head_size=64,
+        softmax_scale=0.125,
+        causal=False,
+    )
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("noncontiguous", [False, True])
+def test_fa4_dense_dispatch_is_opaque_to_dynamic_torch_compile(fake_fa4, noncontiguous):
+    """Fast Dynamo/custom-op regression; Inductor coverage is separate."""
+    impl = fake_fa4
+    q = torch.randn(1, 16, 8, 64, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        q = q.transpose(1, 2).contiguous().transpose(1, 2)
+        assert not q.is_contiguous()
+    path = impl.resolve_execution_path(
+        ExecutionContext(
+            platform="cuda",
+            kernel_variant="fa4",
+            dtype="bfloat16",
+            require_fullgraph=True,
+        ),
+        q,
+        q,
+        q,
+        None,
+    )
+    assert path.support.status is SupportStatus.SUPPORTED
+    assert path.compilation_mode is CompilationMode.CUSTOM_OP
+
+    compile_count = 0
+
+    def counting_backend(graph_module, _example_inputs):
+        nonlocal compile_count
+        compile_count += 1
+        return graph_module.forward
+
+    compiled = torch.compile(
+        lambda query, key, value: impl.forward_cuda(query, key, value),
+        backend=counting_backend,
+        fullgraph=True,
+        dynamic=True,
+    )
+    out = compiled(q, q, q)
+    q2 = torch.randn(1, 24, 8, 64, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        q2 = q2.transpose(1, 2).contiguous().transpose(1, 2)
+    out2 = compiled(q2, q2, q2)
+    assert out.is_contiguous()
+    assert out2.is_contiguous()
+
+    assert out.shape == q.shape
+    assert out2.shape == q2.shape
+    torch.testing.assert_close(out, impl.forward_cuda(q, q, q))
+    torch.testing.assert_close(out2, impl.forward_cuda(q2, q2, q2))
+    assert compile_count == 1
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_fa4_production_attention_entry_compiles_with_inductor(
+    monkeypatch,
+    fake_fa4,
+):
+    """Compile through Attention.forward with the default Inductor backend."""
+    from vllm_omni.diffusion.attention import layer as attention_layer
+
+    layer = _make_attention_layer(monkeypatch)
+    layer.attention = fake_fa4
+    q = torch.randn(1, 16, 8, 64, device="cuda", dtype=torch.bfloat16)
+
+    resolved_contexts = []
+    resolve_execution_path = fake_fa4.resolve_execution_path
+
+    def record_execution_path(context, query, key, value, attn_metadata):
+        resolved_contexts.append(context)
+        return resolve_execution_path(
+            context,
+            query,
+            key,
+            value,
+            attn_metadata,
+        )
+
+    monkeypatch.setattr(fake_fa4, "resolve_execution_path", record_execution_path)
+    context = ExecutionContext(
+        platform="cuda",
+        kernel_variant="fa4",
+        dtype="bfloat16",
+        paged_kv=True,
+        parallel_strategy=ParallelStrategy.ULYSSES,
+        require_fullgraph=True,
+    )
+    path = layer.resolve_execution_path(context, q, q, q, None)
+    assert path.support.status is SupportStatus.SUPPORTED
+    assert path.compilation_mode is CompilationMode.CUSTOM_OP
+    assert resolved_contexts[-1].paged_kv is False
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.NONE
+
+    compiled = torch.compile(layer, fullgraph=True, dynamic=True)
+    out = compiled(q, q, q)
+
+    assert out.shape == q.shape
+    torch.testing.assert_close(out, layer(q, q, q))
+    q2 = torch.randn(1, 24, 8, 64, device="cuda", dtype=torch.bfloat16)
+    torch.testing.assert_close(compiled(q2, q2, q2), layer(q2, q2, q2))
+
+    layer.paged_kv_cache_role = "self"
+    monkeypatch.setattr(layer, "_active_paged_kv_adapter", lambda: object())
+    paged_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert resolved_contexts[-1].paged_kv is True
+    assert paged_path.support.status is SupportStatus.UNMIGRATED
+    layer.paged_kv_cache_role = None
+
+    layer.use_ring = True
+    layer.parallel_strategy = _ParallelStrategyStub(name="ulysses")
+    hybrid_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.HYBRID_ULYSSES_RING
+    assert hybrid_path.support.status is SupportStatus.UNMIGRATED
+    layer.use_ring = False
+    layer.parallel_strategy = NoParallelAttention()
+
+    layer._hsdp_compile_boundary_enabled = True
+    hsdp_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert OuterBoundary.HSDP in resolved_contexts[-1].outer_boundaries
+    assert hsdp_path.support.status is SupportStatus.UNMIGRATED
+    assert hsdp_path.compilation_mode is CompilationMode.EAGER_ONLY
+
+    layer._hsdp_compile_boundary_enabled = False
+    runtime_context = ForwardContext(
+        omni_diffusion_config=OmniDiffusionConfig(
+            parallel_config=DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=1),
+        ),
+    )
+    monkeypatch.setattr(attention_layer, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(attention_layer, "get_forward_context", lambda: runtime_context)
+    layer.use_ring = True
+    layer.parallel_strategy = _ParallelStrategyStub(name="ring")
+    runtime_hsdp_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert OuterBoundary.HSDP in resolved_contexts[-1].outer_boundaries
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.NONE
+    assert runtime_hsdp_path.support.status is SupportStatus.UNMIGRATED
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("head_dims", [(32, 32), (64, 64), (80, 48), (192, 128), (256, 256)])
+def test_real_fa4_fullgraph_matches_sdpa_across_shapes(monkeypatch, head_dims):
+    """Validate real FA4 numerics through the complete compiled attention layer."""
+    capability = current_omni_platform.get_device_capability() if current_omni_platform.is_cuda() else None
+    if capability is None or capability.major < 10:
+        pytest.skip("Requires Blackwell")
+    if not fa.IS_FLASH_ATTN_4 or fa.flash_attn_func is None:
+        pytest.skip("Requires CuTe FlashAttention-4")
+
+    from torch.nn.attention import SDPBackend, sdpa_kernel
+
+    head_size, value_head_size = head_dims
+    layer = _make_attention_layer(monkeypatch, head_size=head_size)
+    compiled = torch.compile(layer, fullgraph=True, dynamic=True)
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    # Tile boundaries, cross-attention lengths, batching, and a longer sequence.
+    for batch, q_length, kv_length in ((1, 16, 16), (1, 129, 257), (2, 257, 129), (1, 1024, 1024)):
+        query, key, value = (
+            torch.randn(batch, length, 8, dim, device="cuda", dtype=torch.bfloat16, generator=generator)
+            for length, dim in ((q_length, head_size), (kv_length, head_size), (kv_length, value_head_size))
+        )
+        path = layer.resolve_execution_path(context, query, key, value, None)
+        assert path.requested_support(context).status is SupportStatus.SUPPORTED
+        with sdpa_kernel(SDPBackend.MATH):
+            reference = torch.nn.functional.scaled_dot_product_attention(
+                query.transpose(1, 2).float(),
+                key.transpose(1, 2).float(),
+                value.transpose(1, 2).float(),
+                scale=head_size**-0.5,
+            ).transpose(1, 2)
+        actual = compiled(query, key, value)
+        assert actual.dtype == query.dtype
+        assert actual.device == query.device
+        torch.testing.assert_close(actual.float(), reference, atol=1e-2, rtol=1e-2)
+        torch.testing.assert_close(actual, layer(query, key, value))
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("head_dims", [(192, 192), (65, 64), (64, 65)])
+def test_real_fa4_dimension_rejection_matches_kernel(monkeypatch, head_dims):
+    capability = current_omni_platform.get_device_capability() if current_omni_platform.is_cuda() else None
+    if capability is None or capability.major not in (10, 11) or not fa.IS_FLASH_ATTN_4:
+        pytest.skip("Requires SM100/SM110 CuTe FlashAttention-4")
+    head_size, value_head_size = head_dims
+    layer = _make_attention_layer(monkeypatch, head_size=head_size)
+    query = torch.randn(1, 16, 8, head_size, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn(1, 16, 8, value_head_size, device="cuda", dtype=torch.bfloat16)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    path = layer.resolve_execution_path(context, query, query, value, None)
+    assert path.requested_support(context).status is SupportStatus.UNSUPPORTED
+    with pytest.raises(AssertionError) as error:
+        layer(query, query, value)
+    assert str(error.value) in path.support.reason
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("noncontiguous", [False, True])
+def test_real_fa4_custom_op_unequal_value_dimension_schema(noncontiguous):
+    if not current_omni_platform.is_cuda() or not fa.IS_FLASH_ATTN_4:
+        pytest.skip("Requires CuTe FlashAttention-4")
+    query = torch.randn(1, 17, 8, 80, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(1, 25, 8, 80, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn(1, 25, 8, 48, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        query, key, value = (t.transpose(1, 2).contiguous().transpose(1, 2) for t in (query, key, value))
+        assert not query.is_contiguous()
+    torch.library.opcheck(
+        torch.ops.vllm_omni.fa4_dense_attention.default,
+        (query, key, value, 80**-0.5, False, False),
+    )

tests/diffusion/distributed/test_cosmos3_pre_sharded.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare full and pre-sharded Cosmos loading with real two-rank FSDP."""
+
+from contextlib import ExitStack
+from datetime import timedelta
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import nn
+
+from tests.helpers.mark import hardware_test
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.core_model, pytest.mark.gpu]
+
+
+class _SDPAAttention(nn.Module):
+    """Use deterministic native attention while testing loading and FSDP."""
+
+    def __init__(self, *, causal, softmax_scale, **kwargs):
+        super().__init__()
+        self.causal = causal
+        self.scale = softmax_scale
+
+    def forward(self, query, key, value):
+        return F.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            is_causal=self.causal,
+            scale=self.scale,
+        ).transpose(1, 2)
+
+
+def _load_cosmos_worker(rank, rendezvous, checkpoint_dir, edge):
+    from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules
+
+    bootstrap_vllm_layer_custom_op_modules()
+
+    from safetensors.torch import save_file
+    from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+    from vllm.config.load import LoadConfig
+    from vllm.model_executor import parameter
+    from vllm.model_executor.layers import linear
+    from vllm.utils.torch_utils import set_default_torch_dtype
+
+    from vllm_omni.diffusion.distributed import hsdp
+    from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3, transformer_cosmos3_edge
+    from vllm_omni.platforms import current_omni_platform
+
+    device = torch.device("cuda", rank)
+    current_omni_platform.set_device(device)
+    dist.init_process_group("nccl", init_method=rendezvous, rank=rank, world_size=2, timeout=timedelta(seconds=120))
+    try:
+        with ExitStack() as stack:
+            stack.enter_context(set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cuda"))))
+            stack.enter_context(set_default_torch_dtype(torch.bfloat16))
+            # HSDP spans both workers, while every vLLM linear uses TP=1.
+            for module in (linear, parameter):
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_rank", return_value=0))
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_world_size", return_value=1))
+            for module in (transformer_cosmos3, transformer_cosmos3_edge):
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_world_size", return_value=1))
+                stack.enter_context(patch.object(module, "FrameworkAttention", _SDPAAttention))
+            stack.enter_context(patch.object(transformer_cosmos3, "_get_ulysses_state", return_value=(1, 0, None)))
+            stack.enter_context(
+                patch.object(hsdp, "get_world_group", return_value=SimpleNamespace(world_size=2, rank_in_group=rank))
+            )
+
+            model_cls = (
+                transformer_cosmos3_edge.Cosmos3EdgeVFMTransformer
+                if edge
+                else transformer_cosmos3.Cosmos3VFMTransformer
+            )
+            tf_config = dict(
+                hidden_size=16,
+                num_hidden_layers=1,
+                num_attention_heads=2,
+                num_key_value_heads=2,
+                head_dim=8,
+                intermediate_size=32,
+                vocab_size=32,
+                latent_patch_size=2,
+                latent_channel=48,
+                temporal_compression_factor=4,
+                rope_scaling={"mrope_section": [2, 1, 1]},
+            )
+            if edge:
+                tf_config.update(
+                    backbone_type=transformer_cosmos3_edge.COSMOS3_EDGE_BACKBONE_TYPE,
+                    qk_norm_for_text=False,
+                    use_und_k_norm_for_gen=True,
+                )
+            config = SimpleNamespace(
+                tf_model_config=tf_config,
+                dtype=torch.bfloat16,
+                quantization_config=None,
+                lora_path=None,
+                num_weight_load_threads=1,
+                hsdp_weight_load_strategy="full",
+                parallel_config=SimpleNamespace(
+                    use_hsdp=True, hsdp_replicate_size=1, hsdp_shard_size=2, tensor_parallel_size=1
+                ),
+            )
+
+            class Pipeline(nn.Module):
+                def __init__(self):
+                    super().__init__()
+                    self.transformer = model_cls(config)
+                    self.transformer.register_buffer("scratch", torch.tensor([7.0]), persistent=False)
+                    self.vae = nn.Linear(2, 2)
+                    self.weights_sources = [
+                        DiffusersPipelineLoader.ComponentSource(checkpoint_dir, None, None, "transformer.", False)
+                    ]
+
+                def load_weights(self, weights):
+                    state = dict(weights)
+                    self.transformer.load_state_dict(
+                        {name.removeprefix("transformer."): t for name, t in state.items()}
+                    )
+                    self.transformer.post_load_weights()
+                    self.transformer.eval()
+                    self.transformer.validate_loaded_weights(set(state))
+                    return set(state)
+
+            full = Pipeline()
+            if rank == 0:
+                torch.manual_seed(42)
+                with torch.no_grad():
+                    for name, tensor in full.transformer.named_parameters():
+                        tensor.copy_(torch.randn_like(tensor) * 0.02)
+                        if "norm" in name:
+                            tensor.fill_(1)
+                save_file(full.transformer.state_dict(), str(Path(checkpoint_dir) / "model.safetensors"))
+            dist.barrier()
+
+            def load(model):
+                loader = DiffusersPipelineLoader(LoadConfig(), config)
+                with patch.object(loader, "_init_from_load_format", return_value=model):
+                    return loader._load_model_with_hsdp(device)
+
+            full = load(full)
+            config.hsdp_weight_load_strategy = "pre_sharded"
+            pre_sharded = load(Pipeline())
+            full_parameters = dict(full.transformer.named_parameters())
+            from torch.distributed.tensor import DTensor
+
+            for name, actual in pre_sharded.transformer.named_parameters():
+                expected = full_parameters[name]
+                if name.startswith("time_embedder."):
+                    assert not isinstance(actual, DTensor)
+                    assert actual.dtype == torch.float32
+                else:
+                    assert isinstance(actual, DTensor)
+                    actual = actual.full_tensor()
+                    expected = expected.full_tensor()
+                torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+            assert pre_sharded.vae.weight.device == device
+            torch.testing.assert_close(pre_sharded.transformer.scratch, full.transformer.scratch)
+            assert all(t.device.type != "meta" for t in pre_sharded.transformer.buffers())
+
+            inputs = dict(
+                hidden_states=torch.linspace(-1, 1, 48 * 4, device=device).reshape(1, 48, 1, 2, 2),
+                timestep=torch.tensor([1.0], device=device),
+                text_ids=torch.tensor([[1, 2]], device=device),
+                text_mask=torch.ones(1, 2, dtype=torch.long, device=device),
+                video_shape=(1, 2, 2),
+                fps=24.0,
+            )
+            with torch.no_grad():
+                # Also exercise the cached UND path on the second denoising step.
+                for step in (1.0, 0.5):
+                    inputs["timestep"].fill_(step)
+                    expected = full.transformer(**inputs)
+                    actual = pre_sharded.transformer(**inputs)
+                    assert torch.isfinite(actual).all()
+                    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("edge", [False, True], ids=["cosmos3", "cosmos3_edge"])
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.accelerator.device_count() < 2 or not dist.is_nccl_available(),
+    reason="requires two CUDA GPUs and NCCL",
+)
+def test_cosmos3_pre_sharded_matches_full_hsdp(tmp_path, edge):
+    torch.multiprocessing.spawn(
+        _load_cosmos_worker,
+        args=(f"file://{tmp_path / 'rendezvous'}", str(tmp_path), edge),
+        nprocs=2,
+        join=True,
+    )

tests/diffusion/distributed/test_hsdp.py

@@ -3,11 +3,13 @@
 """Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""
 
 import gc
+from types import SimpleNamespace
 
 import pytest
 import torch
 import torch.distributed as dist
 import torch.nn as nn
+from torch.distributed.fsdp import fully_shard
 from torch.distributed.tensor import DeviceMesh, DTensor
 from vllm.utils.network_utils import get_file_store_init_method
 
@@ -18,6 +20,7 @@ from vllm_omni.diffusion.distributed.hsdp import (
     _unshardable_parameters,
     shard_model,
 )
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model]
 
@@ -338,3 +341,92 @@ class TestHSDPShardConditions:
                 matched.append(name)
         assert "blocks.0" in matched
         assert "blocks.1" in matched
+
+
+def test_pre_sharded_to_empty_materializes_fsdp_internal_storage(cpu_process_group):
+    with torch.device("meta"):
+        model = nn.Module()
+        model.transformer = nn.Sequential(nn.Linear(2, 2, bias=False))
+    model.transformer.register_buffer("scratch", torch.tensor([3.0]), persistent=False)
+    shard_model(
+        model.transformer,
+        mesh=DeviceMesh("cpu", [0]),
+        hsdp_shard_conditions=[lambda name, _module: name == "0"],
+    )
+
+    load_plan = SimpleNamespace(
+        roots=[SimpleNamespace(name="transformer", module=model.transformer)],
+        groups=[SimpleNamespace(module=model.transformer[0])],
+        binding_names=frozenset({"transformer.0.weight"}),
+    )
+    sharded_weight = model.transformer[0].weight
+
+    DiffusersPipelineLoader._materialize_pre_sharded_hsdp_state(model, load_plan, torch.device("cpu"))
+
+    fsdp_param = fully_shard.state(model.transformer[0])._fsdp_param_group.fsdp_params[0]
+    assert model.transformer[0].weight is sharded_weight
+    assert model.transformer[0].weight.to_local().device.type == "cpu"
+    assert fsdp_param._sharded_param_data.device.type == "cpu"
+    assert (
+        fsdp_param._sharded_param_data.untyped_storage().data_ptr()
+        == model.transformer[0].weight.to_local().untyped_storage().data_ptr()
+    )
+    assert torch.equal(model.transformer.scratch, torch.tensor([3.0]))
+
+
+def test_pre_sharded_load_restores_checkpoint_values(cpu_process_group, tmp_path, monkeypatch):
+    from safetensors.torch import save_file
+    from vllm.config.load import LoadConfig
+
+    from vllm_omni.diffusion.distributed import hsdp as hsdp_module
+    from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
+
+    class Transformer(nn.Module):
+        _hsdp_shard_conditions = [lambda name, module: name == "block"]
+        _hsdp_ignored_modules = ["time_embedder"]
+
+        def __init__(self):
+            super().__init__()
+            self.block = nn.Linear(3, 4, bias=False)
+            self.time_embedder = nn.Linear(2, 2, bias=False, dtype=torch.bfloat16)
+            self.root_weight = nn.Parameter(torch.zeros(2))
+            self.register_buffer("scratch", torch.tensor([7.0]), persistent=False)
+            release_module_parameters_to_meta(self.block)
+
+        def post_load_weights(self):
+            self.time_embedder.to(torch.float32)
+
+    model = nn.Module()
+    model.transformer = Transformer()
+    model.vae = nn.Linear(2, 2)
+    tensors = {
+        "block.weight": torch.arange(12, dtype=torch.float32).reshape(4, 3),
+        "root_weight": torch.tensor([2.0, 3.0]),
+        "time_embedder.weight": torch.arange(4, dtype=torch.bfloat16).reshape(2, 2),
+    }
+    save_file(tensors, str(tmp_path / "model.safetensors"))
+    model.weights_sources = [DiffusersPipelineLoader.ComponentSource(str(tmp_path), None, None, "transformer.", False)]
+    config = SimpleNamespace(
+        dtype=torch.float32,
+        quantization_config=None,
+        lora_path=None,
+        hsdp_weight_load_strategy="pre_sharded",
+        num_weight_load_threads=1,
+        parallel_config=SimpleNamespace(use_hsdp=True, hsdp_replicate_size=1, hsdp_shard_size=1),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), config)
+    monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: model)
+    monkeypatch.setattr(hsdp_module, "get_world_group", lambda: SimpleNamespace(world_size=1, rank_in_group=0))
+    monkeypatch.setattr(hsdp_module, "current_omni_platform", SimpleNamespace(device_type="cpu"))
+    loaded = loader._load_model_with_hsdp(torch.device("cpu"))
+    assert loaded is model
+    for name, expected in tensors.items():
+        actual = dict(model.transformer.named_parameters())[name]
+        if name.startswith("time_embedder."):
+            assert not isinstance(actual, DTensor)
+            assert actual.dtype == torch.float32
+            torch.testing.assert_close(actual, expected.float())
+        else:
+            assert isinstance(actual, DTensor)
+            torch.testing.assert_close(actual.to_local(), expected)
+    torch.testing.assert_close(model.transformer.scratch, torch.tensor([7.0]))

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -27,6 +27,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
     HostWeightPlanResult,
     TensorBinding,
+    build_checkpoint_binding_plan,
 )
 from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
 from vllm_omni.diffusion.models.helios import HeliosPipeline
@@ -2089,3 +2090,206 @@ def test_hsdp_broadcast_weight_load_online_quant_multiprocess():
             )
         assert os.path.exists(os.path.join(temp_dir, "rank_0_success.flag"))
         assert os.path.exists(os.path.join(temp_dir, "rank_1_success.flag"))
+
+
+def test_pre_sharded_hsdp_rejects_quantization_before_loading():
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        hsdp_weight_load_strategy="pre_sharded",
+        lora_path=None,
+        quantization_config=object(),
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    with pytest.raises(ValueError, match="does not support quantization"):
+        loader._load_model_with_hsdp(torch.device("cpu"))
+
+
+def test_pre_sharded_hsdp_strategy_dispatches_without_using_full_loader(mocker):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        hsdp_weight_load_strategy="pre_sharded",
+        lora_path=None,
+        quantization_config=None,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    discovered = PipelineModules(
+        dits=[model.transformer],
+        dit_names=["transformer"],
+        vaes=[],
+        encoders=[],
+        encoder_names=[],
+        resident_modules=[],
+        resident_names=[],
+    )
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=AssertionError("full loader must not run"))  # type: ignore[method-assign]
+    mocker.patch.object(loader_mod.ModuleDiscovery, "discover", return_value=discovered)
+    pre_sharded_load = mocker.patch.object(
+        loader,
+        "_load_model_with_pre_sharded_hsdp",
+        return_value=model,
+    )
+
+    assert loader._load_model_with_hsdp(torch.device("cpu")) is model
+    pre_sharded_load.assert_called_once()
+
+
+def test_pre_sharded_hsdp_rejects_tensor_transform():
+    def transform(tensor):
+        return tensor.transpose(-2, -1)
+
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="weight",
+                file_path="model.safetensors",
+                transform=transform,
+            )
+        }
+    )
+
+    with pytest.raises(ValueError, match="tensor transforms are unsupported"):
+        DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan)
+
+
+def test_pre_sharded_hsdp_accepts_direct_safetensors_binding():
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="weight",
+                file_path="model.safetensors",
+            )
+        }
+    )
+
+    DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan)
+
+
+def test_pre_sharded_hsdp_uses_fsdp_destination_directly():
+    target = torch.empty(8, 4, 3)
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="source_weight",
+                file_path="transformer/model.safetensors",
+            )
+        }
+    )
+
+    checkpoint_states, local_bytes = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+        load_plan,
+        {"transformer.weight": target},
+    )
+    checkpoint_target = checkpoint_states[Path("transformer")]["source_weight"]
+
+    assert checkpoint_target is target
+    assert checkpoint_target.shape == (8, 4, 3)
+    assert checkpoint_target.device.type == "cpu"
+    assert checkpoint_target.untyped_storage().data_ptr() == target.untyped_storage().data_ptr()
+    assert local_bytes == target.numel() * target.element_size()
+
+
+def test_pre_sharded_hsdp_checkpoint_writes_runtime_storage_directly():
+    runtime_target = torch.empty(2, 3)
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="source_weight",
+                file_path="transformer/model.safetensors",
+            )
+        }
+    )
+    checkpoint_tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    checkpoint_states, _ = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+        load_plan,
+        {"transformer.weight": runtime_target},
+    )
+    checkpoint_states[Path("transformer")]["source_weight"].copy_(checkpoint_tensor)
+
+    assert torch.equal(runtime_target, checkpoint_tensor)
+
+
+def test_pre_sharded_hsdp_rejects_duplicate_checkpoint_destinations():
+    binding = TensorBinding(checkpoint_key="weight", file_path="transformer/model.safetensors")
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.first": binding,
+            "transformer.second": binding,
+        }
+    )
+
+    with pytest.raises(ValueError, match="multiple runtime tensors"):
+        DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+            load_plan,
+            {"transformer.first": torch.empty(2, 3), "transformer.second": torch.empty(2, 3)},
+        )
+
+
+def test_pre_sharded_hsdp_managed_module_traversal_is_child_first_and_skips_nested_roots():
+    root = nn.Module()
+    root.left = nn.Sequential(nn.Linear(2, 2), nn.ReLU())
+    root.right = nn.Sequential(nn.Linear(2, 2), nn.ReLU())
+
+    names = {module: name or "root" for name, module in root.named_modules()}
+    modules = DiffusersPipelineLoader._hsdp_managed_modules_post_order(
+        root,
+        nested_hsdp_roots={root.left},
+    )
+
+    assert [names[module] for module in modules] == ["right.0", "right.1", "right", "root"]
+
+
+@pytest.mark.parametrize("checkpoint_key", ["renamed", "unexpected"])
+def test_hsdp_checkpoint_plan_honors_remap_and_rejects_missing(tmp_path, checkpoint_key, monkeypatch):
+    import vllm_omni.diffusion.model_loader.host_weight_plan as plan_mod
+
+    monkeypatch.setattr(plan_mod, "get_direct_mmap_adapter", lambda _model: None)
+
+    class Pipeline(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.transformer = nn.Linear(2, 2, bias=False)
+
+        @staticmethod
+        def remap_checkpoint_key(name):
+            return "transformer.weight" if name == "transformer.renamed" else name
+
+    checkpoint = tmp_path / "model.safetensors"
+    save_file({checkpoint_key: torch.ones(2, 2)}, str(checkpoint))
+    source = SimpleNamespace(
+        model_or_path=str(tmp_path),
+        subfolder=None,
+        revision=None,
+        prefix="transformer.",
+    )
+    model = Pipeline()
+    result = build_checkpoint_binding_plan(
+        model,
+        dit_modules=(("transformer", model.transformer),),
+        sources=(source,),
+        model_path=None,
+        tensor_parallel_size=1,
+        online_quantization=False,
+    )
+
+    if checkpoint_key == "renamed":
+        assert result.plan is not None
+        assert result.plan.bindings["transformer.weight"].checkpoint_key == "renamed"
+    else:
+        assert result.plan is None
+        assert "no checkpoint binding" in result.fallback_reason

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

@@ -11,10 +11,29 @@ eager off CUDA) and the tiled decode of long clips.
 import pytest
 import torch
 
-from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE, LowPass, SnakeBeta, Upsample
+from tests.helpers.mark import hardware_test
+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
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model]
+
+
+@pytest.fixture(autouse=True)
+def _bound_cpu_decode_threads(request):
+    if request.node.get_closest_marker("cpu") is None:
+        yield
+        return
+    # These narrow convolutions run in shared CI pods. A host-sized Torch
+    # pool spends more time coordinating workers than decoding small clips.
+    previous_threads = torch.get_num_threads()
+    torch.set_num_threads(1)
+    print(f"AuK CPU decode threads: {previous_threads} -> {torch.get_num_threads()}")
+    try:
+        yield
+    finally:
+        torch.set_num_threads(previous_threads)
 
 
 def _small_vae() -> AuKVAE:
@@ -45,6 +64,23 @@ def _reference_snake(module: SnakeBeta, x: torch.Tensor) -> torch.Tensor:
     return x + (1.0 / (beta + 1e-9)) * torch.sin(x * alpha).pow(2)
 
 
+def _assert_plain_graph_matches_eager(replay: torch.Tensor, eager: torch.Tensor) -> None:
+    assert replay.shape == eager.shape and replay.dtype == eager.dtype
+    assert torch.isfinite(replay).all() and torch.isfinite(eager).all()
+    error = (replay - eager).abs()
+    rms = error.square().mean().sqrt().item()
+    print(f"AuK plain graph: max_abs={error.max().item():.9g}, rms={rms:.9g}")
+    if current_omni_platform.is_rocm():
+        # MI300 build 13040 measured max_abs <= 1.64e-7 and RMS <= 4.13e-8.
+        # Bound both peak and aggregate FP32 error; near-zero samples must
+        # not hide behind a relative tolerance. NVIDIA retains bit parity.
+        torch.testing.assert_close(replay, eager, atol=1e-6, rtol=0.0)
+        assert rms <= 1e-7, rms
+    else:
+        assert torch.equal(replay, eager)
+
+
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_window_padding_decodes_to_raw_zero() -> None:
     """Bucket padding is the normalized latent whose denormalized value is zero."""
@@ -57,6 +93,7 @@ def test_window_padding_decodes_to_raw_zero() -> None:
     torch.testing.assert_close(raw, torch.zeros_like(raw), atol=1e-6, rtol=0.0)
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_eager_snake_matches_the_reference_formula() -> None:
     module = SnakeBeta(6, alpha_logscale=True)
@@ -69,6 +106,7 @@ def test_eager_snake_matches_the_reference_formula() -> None:
     assert torch.equal(module(x), _reference_snake(module, x))
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_filter_cache_leaves_the_filters_and_output_unchanged() -> None:
     lowpass = LowPass(cutoff=0.25, half_width=0.3, stride=2, kernel_size=12, causal=True)
@@ -88,6 +126,7 @@ def test_filter_cache_leaves_the_filters_and_output_unchanged() -> None:
         assert module._expanded.shape == (3, 1, 12)
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_decode_fast_paths_reproduce_the_plain_decode() -> None:
     vae = _small_vae()
@@ -103,6 +142,7 @@ def test_decode_fast_paths_reproduce_the_plain_decode() -> None:
     assert all(module._cached for module in vae.modules() if isinstance(module, SnakeBeta))
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_graph_wrapper_falls_back_to_eager_off_cuda(mocker) -> None:
     vae = _small_vae()
@@ -119,6 +159,7 @@ def test_graph_wrapper_falls_back_to_eager_off_cuda(mocker) -> None:
     assert not wrapper._cache and not wrapper._compiled
 
 
+@pytest.mark.cpu
 def test_compiled_bucket_is_the_smallest_captured_one_that_fits() -> None:
     wrapper = AuKVAEDecodeGraph(_small_vae(), compile_shapes=(16, 8, 8))
     assert wrapper.compile_shapes == [8, 16]
@@ -131,6 +172,7 @@ def test_compiled_bucket_is_the_smallest_captured_one_that_fits() -> None:
     assert wrapper.compiled_bucket(6) == 16
 
 
+@pytest.mark.cpu
 def test_plan_tiles_covers_the_clip_once_with_full_context() -> None:
     assert plan_tiles(40, 64, 10, 3) == [(0, 40, 0, 40)]
     windows = plan_tiles(1000, 512, 53, 9)
@@ -149,10 +191,16 @@ def test_plan_tiles_covers_the_clip_once_with_full_context() -> None:
     # 512 and a 256 window instead of two 512s.
     assert plan_tiles(600, 512, 53, 9, sizes=(128, 256, 512)) == [(0, 512, 0, 503), (344, 256, 503, 600)]
     assert plan_tiles(1000, 512, 53, 9, sizes=(128, 256, 512))[-1] == (872, 128, 953, 1000)
+    narrow = plan_tiles(200, 64, 53, 9)
+    assert len(narrow) == 69
+    assert [emit for _, _, emit, _ in narrow[:3]] == [0, 55, 57]
+    assert all(a[3] == b[2] for a, b in zip(narrow, narrow[1:]))
+    assert narrow[-1][3] == 200
     with pytest.raises(ValueError, match="must exceed"):
         plan_tiles(100, 60, 53, 9)
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_decode_context_bounds_the_measured_receptive_field() -> None:
     vae = _small_vae()
@@ -170,25 +218,26 @@ def test_decode_context_bounds_the_measured_receptive_field() -> None:
     assert differing and all(index < left or end - start - index <= right for index in differing)
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_tiled_decode_matches_the_whole_decode() -> None:
     vae = _small_vae()
     # Off the accelerator every window decodes eagerly, so this isolates the stitching.
-    wrapper = AuKVAEDecodeGraph(vae, compile_shapes=(64,), tile_frames=64)
-    assert wrapper.tile_frames == 64 and wrapper.context_frames == (53, 9)
-    for frames in (64, 65, 200, 331):
+    wrapper = AuKVAEDecodeGraph(vae, compile_shapes=(128,), tile_frames=128)
+    assert wrapper.tile_frames == 128 and wrapper.context_frames == (53, 9)
+    for frames in (128, 129, 200, 331):
         latents = torch.randn(1, frames, vae.latent_dim)
         whole = vae.decode(latents)
         tiled = wrapper(latents)
         assert tiled.shape == whole.shape
         torch.testing.assert_close(tiled, whole, atol=1e-5, rtol=0.0)
-        assert wrapper.last_mode == ("eager" if frames == 64 else "tiled")
+        assert wrapper.last_mode == ("eager" if frames == 128 else "tiled")
     # The streaming form yields the same audio in order.
     latents = torch.randn(1, 200, vae.latent_dim)
     pieces = [(start, chunk.clone()) for start, chunk in wrapper.decode_tiles(latents)]
-    # With a 64-frame tile and a 62-frame halo each interior tile adds two frames.
-    assert [start for start, _ in pieces] == [emit for _, _, emit, _ in plan_tiles(200, 64, 53, 9)]
-    assert [start for start, _ in pieces][:3] == [0, 55, 57]
+    # A 128-frame tile avoids repeatedly decoding the same 62-frame halo.
+    assert [start for start, _ in pieces] == [emit for _, _, emit, _ in plan_tiles(200, 128, 53, 9)]
+    assert [start for start, _ in pieces][:3] == [0, 119, 185]
     torch.testing.assert_close(
         torch.cat([chunk for _, chunk in pieces], dim=1), vae.decode(latents), atol=1e-5, rtol=0.0
     )
@@ -198,20 +247,90 @@ def test_tiled_decode_matches_the_whole_decode() -> None:
         AuKVAEDecodeGraph(vae, tile_frames=60)
 
 
-@pytest.mark.cuda
+@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()
 def test_graph_replay_matches_eager_per_length() -> None:
     vae = _small_vae().to("cuda")
     wrapper = AuKVAEDecodeGraph(vae, max_graphs=2)
     pools = []
+    saved_outputs: list[tuple[torch.Tensor, torch.Tensor]] = []
     for frames in (6, 9, 6, 12, 6):
         latents = torch.randn(1, frames, vae.latent_dim, device="cuda")
         eager = vae.decode(latents)
         replay = wrapper(latents)
-        # Exact-length graphs replay the very same kernels: bit-identical,
-        # including after the generation was retired.
-        assert torch.equal(replay, eager), frames
+        _assert_plain_graph_matches_eager(replay, eager)
+        # Returned waveforms own their storage and survive later replays
+        # and retirement of the generation that produced them.
+        for previous, snapshot in saved_outputs:
+            assert torch.equal(previous, snapshot)
+        saved_outputs.append((replay, replay.clone()))
         pools.append(wrapper._plain_pool)
     # The third distinct length found the cache full, so the whole generation
     # (6 and 9) was retired together with its pool rather than one graph at a
@@ -220,7 +339,7 @@ def test_graph_replay_matches_eager_per_length() -> None:
     assert pools[0] is pools[2] and pools[3] is pools[4] and pools[2] is not pools[3]
 
 
-@pytest.mark.cuda
+@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()
 def test_bucketed_graph_only_disturbs_the_tail() -> None:
@@ -240,7 +359,7 @@ def test_bucketed_graph_only_disturbs_the_tail() -> None:
     torch.testing.assert_close(replay, eager, atol=0.1, rtol=0.0)
 
 
-@pytest.mark.cuda
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="torch.compile + CUDA graph capture requires CUDA")
 @torch.inference_mode()
 def test_compiled_buckets_replay_within_fusion_tolerance_and_leave_longer_clips_to_plain_graphs() -> None:
@@ -266,11 +385,12 @@ def test_compiled_buckets_replay_within_fusion_tolerance_and_leave_longer_clips_
     # a longer clip gets its own plain graph.
     assert wrapper.tile_frames == 0
     long = torch.randn(1, 12, vae.latent_dim, device="cuda")
-    assert torch.equal(wrapper(long), vae.decode(long))
+    replay, eager = wrapper(long), vae.decode(long)
+    _assert_plain_graph_matches_eager(replay, eager)
     assert wrapper.last_mode == "graph" and list(wrapper._cache) == [12]
 
 
-@pytest.mark.cuda
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="torch.compile + CUDA graph capture requires CUDA")
 @torch.inference_mode()
 def test_tiles_replay_the_compiled_bucket_for_long_clips() -> None:

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/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -868,6 +868,93 @@ def test_qwen_gen_mlp_remains_gated() -> None:
     assert hasattr(model.gen_layers[0].mlp, "gate_proj")
 
 
+@pytest.mark.parametrize("edge", [False, True])
+@pytest.mark.parametrize(
+    "use_hsdp,strategy", [(False, "full"), (False, "pre_sharded"), (True, "full"), (True, "pre_sharded")]
+)
+def test_cosmos3_releases_blocks_only_for_pre_sharded_hsdp(edge, use_hsdp, strategy, monkeypatch):
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3, transformer_cosmos3_edge
+    from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
+
+    released = []
+
+    def release_and_check(block):
+        parameters = dict(block.named_parameters())
+        metadata = {name: dict(parameter.__dict__) for name, parameter in parameters.items()}
+        release_module_parameters_to_meta(block)
+        for name, parameter in block.named_parameters():
+            assert parameter is parameters[name]
+            assert parameter.device.type == "meta"
+            assert parameter.__dict__ == metadata[name]
+        released.append(block)
+
+    monkeypatch.setattr(transformer_cosmos3, "release_module_parameters_to_meta", release_and_check)
+    monkeypatch.setattr(transformer_cosmos3_edge, "release_module_parameters_to_meta", release_and_check)
+    model_cls = (
+        transformer_cosmos3_edge.Cosmos3EdgeVFMTransformer if edge else transformer_cosmos3.Cosmos3VFMTransformer
+    )
+    config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config
+    model = model_cls(
+        SimpleNamespace(
+            tf_model_config=config(num_hidden_layers=1),
+            dtype=torch.float32,
+            parallel_config=SimpleNamespace(use_hsdp=use_hsdp),
+            hsdp_weight_load_strategy=strategy,
+        )
+    )
+    blocks = [*model.language_model.layers, *model.gen_layers]
+    expected_meta = use_hsdp and strategy == "pre_sharded"
+    assert len(released) == (len(blocks) if expected_meta else 0)
+    assert all((p.device.type == "meta") == expected_meta for block in blocks for p in block.parameters())
+    assert model.time_embedder.linear_1.weight.device.type != "meta"
+    assert all(buffer.device.type != "meta" for buffer in model.buffers())
+
+
+@pytest.mark.parametrize("edge", [False, True])
+def test_cosmos3_pre_sharded_meta_processing_accepts_real_unquantized_blocks(edge, monkeypatch):
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+    from vllm_omni.diffusion.models.cosmos3 import Cosmos3EdgeVFMTransformer, Cosmos3VFMTransformer
+
+    model_cls = Cosmos3EdgeVFMTransformer if edge else Cosmos3VFMTransformer
+    config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config
+    model = model_cls(
+        SimpleNamespace(
+            tf_model_config=config(num_hidden_layers=1),
+            dtype=torch.float32,
+            parallel_config=SimpleNamespace(use_hsdp=True),
+            hsdp_weight_load_strategy="pre_sharded",
+        )
+    )
+    groups = [SimpleNamespace(module=block) for block in (*model.language_model.layers, *model.gen_layers)]
+    plan = SimpleNamespace(roots=[SimpleNamespace(module=model)], groups=groups)
+    processed = []
+    # Exercise model selection with real vLLM linears, without invoking the
+    # CPU-specific weight packing hook on meta tensors in CPU-only CI.
+    monkeypatch.setattr(
+        UnquantizedLinearMethod, "process_weights_after_loading", lambda self, layer: processed.append(layer)
+    )
+    DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    expected = [module for group in groups for module in group.module.modules() if hasattr(module, "quant_method")]
+    assert expected
+    assert processed == expected
+
+    model._hsdp_pre_sharded_meta_post_load = False
+    with pytest.raises(ValueError, match="has not declared"):
+        DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    model._hsdp_pre_sharded_meta_post_load = True
+
+    class UnsupportedMethod(UnquantizedLinearMethod):
+        pass
+
+    expected[-1].quant_method = UnsupportedMethod()
+    processed.clear()
+    with pytest.raises(ValueError, match="UnsupportedMethod"):
+        DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    assert not processed
+
+
 def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 

tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py

@@ -148,7 +148,8 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
     assert array.shape == (2, 4, 4, 3)
     video = post({"video": frames}, SimpleNamespace(output_type="pt"))
     assert set(video) == {"video"}
-    assert video["video"] is frames
+    assert isinstance(video["video"], np.ndarray)
+    np.testing.assert_array_equal(video["video"], frames.numpy())
     image = post(
         {"image": torch.ones(4, 4, 3)},
         SimpleNamespace(output_type="pt"),
@@ -162,6 +163,40 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
     assert latent_output["image"] is latent
 
 
+@pytest.mark.parametrize("output_type", [None, "pt", "np"])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_postprocess_preserves_video_range_for_serving(output_type: str | None, dtype: torch.dtype) -> None:
+    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
+    from vllm_omni.entrypoints.openai.video_api_utils import _coerce_video_to_frames
+
+    # Include black and midtones: all-white frames cannot expose double normalization.
+    frames = torch.tensor([0.0, 0.25, 0.5, 0.75, 1.0], dtype=dtype)
+    frames = frames.reshape(5, 1, 1, 1).expand(5, 2, 4, 3).clone()
+    expected = frames.float().numpy().copy()
+    sampling = None if output_type is None else OmniDiffusionSamplingParams(output_type=output_type)
+    post = get_lingbot_video_post_process_func(None)
+
+    processed = post({"video": frames}, sampling)
+    encoded_frames = np.stack(_coerce_video_to_frames(processed["video"]))
+
+    assert set(processed) == {"video"}
+    assert encoded_frames.dtype == np.float32
+    np.testing.assert_array_equal(encoded_frames, expected)
+    np.testing.assert_array_equal(frames.float().numpy(), expected)
+
+
+@pytest.mark.parametrize("output_key", ["image", "video"])
+def test_postprocess_preserves_latent_payload(output_key: str) -> None:
+    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
+
+    latents = torch.tensor([-2.0, 0.0, 2.0]).reshape(1, 3, 1, 1, 1)
+    sampling = OmniDiffusionSamplingParams(output_type="latent")
+    processed = get_lingbot_video_post_process_func(None)({output_key: latents}, sampling)
+
+    assert set(processed) == {output_key}
+    assert processed[output_key] is latents
+
+
 def test_postprocess_preserves_image_contract_for_serving():
     from vllm_omni.diffusion.data import DiffusionOutput
     from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
@@ -181,7 +216,9 @@ def test_postprocess_preserves_image_contract_for_serving():
         sampling_params=sampling,
         request_id="lingbot-image-contract-test",
     )
-    raw_output = {"image": torch.ones(8, 12, 3)}
+    image = torch.zeros(8, 12, 3)
+    image[4:] = 1.0
+    raw_output = {"image": image}
     processed = get_lingbot_video_post_process_func(SimpleNamespace())(raw_output, sampling)
     normalized = normalize_diffusion_postprocess_output(processed)
 
@@ -199,6 +236,8 @@ def test_postprocess_preserves_image_contract_for_serving():
     endpoint_images = _extract_images_from_result(result)
     assert len(endpoint_images) == 1
     assert endpoint_images[0].size == (12, 8)
+    assert endpoint_images[0].getpixel((0, 0)) == (0, 0, 0)
+    assert endpoint_images[0].getpixel((0, 7)) == (255, 255, 255)
 
 
 def test_forward_resolves_t2v_sampling_and_flow_shift_alias():

tests/diffusion/offloader/test_diffusion_layerwise_offload.py

@@ -13,7 +13,6 @@ import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
 from tests.helpers.mark import hardware_marks
-from tests.helpers.monitor import DeviceMemoryMonitor
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -55,6 +54,7 @@ IMAGE_VIDEO_MODELS_PARAMS: dict[str, dict[str, Any]] = {
 }
 
 _OFFLOAD_STATE_PROBE = "vllm_omni_layerwise_offload_test"
+_OFFLOAD_MEMORY_PROBE = "vllm_omni_layerwise_offload_memory_test"
 # A second CI build can reverse the same exact-head measurement without a
 # source edit; the default keeps the existing baseline-first workload.
 _MEASUREMENT_ORDER_ENV = "VLLM_OMNI_OFFLOAD_TEST_ORDER"
@@ -65,7 +65,7 @@ _MEASUREMENT_ORDERS: dict[str, tuple[bool, bool]] = {
 
 
 class OffloadStateProbe:
-    """Test-only worker extension exposing the configured offloader state."""
+    """Test-only worker extension exposing offloader and allocator state."""
 
     model_runner: Any
 
@@ -86,6 +86,21 @@ class OffloadStateProbe:
             "block_count": sum(group_sizes),
         }
 
+    def reset_peak_memory_for_test(self) -> dict[str, str]:
+        """Start a process-local inference peak after model initialization."""
+        current_omni_platform.synchronize()
+        torch.accelerator.reset_peak_memory_stats()
+        return {"probe": _OFFLOAD_MEMORY_PROBE, "operation": "reset"}
+
+    def get_peak_memory_for_test(self) -> dict[str, str | float]:
+        """Return this worker's allocator peak without counting other jobs."""
+        current_omni_platform.synchronize()
+        return {
+            "probe": _OFFLOAD_MEMORY_PROBE,
+            "operation": "snapshot",
+            "peak_allocated_mb": torch.accelerator.max_memory_allocated() / (1024**2),
+        }
+
 
 def check_audio_determinism(audio1: np.ndarray, audio2: np.ndarray, atol: float = 1e-2) -> bool:
     if not np.allclose(audio1, audio2, atol=atol):
@@ -96,13 +111,6 @@ def check_audio_determinism(audio1: np.ndarray, audio2: np.ndarray, atol: float
     return True
 
 
-def _device_used_mb(device_index: int) -> float:
-    current_omni_platform.synchronize()
-    with current_omni_platform.device(device_index):
-        free_bytes, total_bytes = current_omni_platform.mem_get_info()
-    return (total_bytes - free_bytes) / (1024**2)
-
-
 def _extract_audio(output: Any) -> np.ndarray | None:
     if not output:
         return None
@@ -117,17 +125,29 @@ def _extract_audio(output: Any) -> np.ndarray | None:
     return np.asarray(audio)
 
 
-def _collect_offload_states(value: Any) -> list[dict[str, Any]]:
+def _collect_probe_results(value: Any, probe: str) -> list[dict[str, Any]]:
     if isinstance(value, dict):
-        return [value] if value.get("probe") == _OFFLOAD_STATE_PROBE else []
+        return [value] if value.get("probe") == probe else []
     if isinstance(value, (list, tuple)):
-        states: list[dict[str, Any]] = []
+        results: list[dict[str, Any]] = []
         for item in value:
-            states.extend(_collect_offload_states(item))
-        return states
+            results.extend(_collect_probe_results(item, probe))
+        return results
     return []
 
 
+def _collect_offload_states(value: Any) -> list[dict[str, Any]]:
+    return _collect_probe_results(value, _OFFLOAD_STATE_PROBE)
+
+
+def _collect_memory_results(value: Any, operation: str) -> list[dict[str, Any]]:
+    return [
+        result
+        for result in _collect_probe_results(value, _OFFLOAD_MEMORY_PROBE)
+        if result.get("operation") == operation
+    ]
+
+
 def _measurement_order() -> tuple[bool, bool]:
     """Return the requested order without changing the default CI workload."""
     configured = os.environ.get(_MEASUREMENT_ORDER_ENV, "baseline-first")
@@ -144,8 +164,6 @@ def run_inference(
     num_inference_steps: int = 3,
 ) -> dict[str, Any]:
     current_omni_platform.empty_cache()
-    device_index = current_omni_platform.current_device()
-    initial_used_mb = _device_used_mb(device_index)
 
     if model_name in AUDIO_MODEL:
         params = AUDIO_MODEL_PARAMS
@@ -166,52 +184,45 @@ def run_inference(
         if not offload_states:
             raise AssertionError("The offload-state worker probe returned no diffusion worker results")
 
-        # Measure steady-state inference memory, not model construction. Enabling
-        # layerwise offload first loads the model and then replaces each block's
-        # device storage with CPU-backed weights.  Monitoring that transition
-        # captures both the original model and temporary staging allocations,
-        # which is not representative of layerwise-offloaded inference.
-        monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
-        current_omni_platform.reset_peak_memory_stats()
-        monitor.start()
+        # Reset and read allocator peaks in the diffusion worker process. A
+        # device-wide mem_get_info() sample includes unrelated sibling jobs on
+        # the same GPU and produced negative "savings" in shared ROCm CI.
+        reset_results = _collect_memory_results(
+            runner.omni.engine.collective_rpc(method="reset_peak_memory_for_test", timeout=60),
+            "reset",
+        )
+        if not reset_results:
+            raise AssertionError("The worker peak-memory probe returned no reset acknowledgements")
 
-        try:
-            # Refer to tests/e2e/offline_inference/test_wan22.py
-            # Use minimal settings for testing
-            output = runner.omni.generate(
-                "A cat sitting on a table",
-                OmniDiffusionSamplingParams(
-                    generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
-                    guidance_scale=1.0,
-                    num_inference_steps=num_inference_steps,
-                    **params["sampler_params"],
-                ),
-            )
-        finally:
-            monitor.stop()
+        # Refer to tests/e2e/offline_inference/test_wan22.py.
+        # Use minimal settings for testing.
+        output = runner.omni.generate(
+            "A cat sitting on a table",
+            OmniDiffusionSamplingParams(
+                generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+                guidance_scale=1.0,
+                num_inference_steps=num_inference_steps,
+                **params["sampler_params"],
+            ),
+        )
+        memory_results = _collect_memory_results(
+            runner.omni.engine.collective_rpc(method="get_peak_memory_for_test", timeout=60),
+            "snapshot",
+        )
+        if not memory_results:
+            raise AssertionError("The worker peak-memory probe returned no snapshots")
 
         audio = _extract_audio(output)
         del output
 
-    # DeviceMemoryMonitor reports absolute device usage. Subtract this run's
-    # starting usage. Each mode runs in its own spawned process below, so model
-    # objects and process-local allocator/compiler/backend workspace state cannot
-    # carry into the other measurement. Shared filesystem caches can persist;
-    # _MEASUREMENT_ORDER_ENV enables a reversed-order validation run for that.
-    peak_used_mb = monitor.peak_used_mb
-    incremental_peak_mb = max(0.0, peak_used_mb - initial_used_mb)
-
     del runner
     gc.collect()
     cleanup_dist_env_and_memory()
     current_omni_platform.empty_cache()
-    post_cleanup_used_mb = _device_used_mb(device_index)
 
     return {
-        "initial_used_mb": initial_used_mb,
-        "peak_used_mb": peak_used_mb,
-        "incremental_peak_mb": incremental_peak_mb,
-        "post_cleanup_used_mb": post_cleanup_used_mb,
+        "peak_allocated_mb": max(float(result["peak_allocated_mb"]) for result in memory_results),
+        "memory_results": memory_results,
         "audio": audio,
         "offload_states": offload_states,
     }
@@ -259,12 +270,8 @@ def _assert_audio_outputs(
 
 
 def _print_measurement(label: str, measurement: dict[str, Any]) -> None:
-    print(
-        f"{label}: initial={measurement['initial_used_mb']:.1f} MB, "
-        f"peak={measurement['peak_used_mb']:.1f} MB, "
-        f"incremental_peak={measurement['incremental_peak_mb']:.1f} MB, "
-        f"post_cleanup={measurement['post_cleanup_used_mb']:.1f} MB"
-    )
+    print(f"{label}: worker_peak_allocated={measurement['peak_allocated_mb']:.1f} MB")
+    print(f"{label} worker memory: {measurement['memory_results']}")
     print(f"{label} offload state: {measurement['offload_states']}")
 
 
@@ -298,6 +305,33 @@ def test_measurement_order_rejects_unknown_value(monkeypatch) -> None:
         _measurement_order()
 
 
+@pytest.mark.diffusion
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_collect_memory_results_filters_nested_worker_responses() -> None:
+    value = [
+        None,
+        [
+            {"probe": _OFFLOAD_MEMORY_PROBE, "operation": "reset"},
+            {
+                "probe": _OFFLOAD_MEMORY_PROBE,
+                "operation": "snapshot",
+                "peak_allocated_mb": 123.0,
+            },
+        ],
+        {"probe": _OFFLOAD_STATE_PROBE, "operation": "snapshot"},
+    ]
+
+    assert _collect_memory_results(value, "reset") == [{"probe": _OFFLOAD_MEMORY_PROBE, "operation": "reset"}]
+    assert _collect_memory_results(value, "snapshot") == [
+        {
+            "probe": _OFFLOAD_MEMORY_PROBE,
+            "operation": "snapshot",
+            "peak_allocated_mb": 123.0,
+        }
+    ]
+
+
 @pytest.mark.diffusion
 @pytest.mark.core_model
 @pytest.mark.cpu
@@ -366,8 +400,8 @@ def test_layerwise_offload_diffusion_model(model_name: str):
 
     # Verify that layerwise offloading significantly reduces memory usage
     # Passes only if the actual savings meets the expected savings
-    no_offload_peak_memory = no_offload["incremental_peak_mb"]
-    layerwise_offload_peak_memory = layerwise_offload["incremental_peak_mb"]
+    no_offload_peak_memory = no_offload["peak_allocated_mb"]
+    layerwise_offload_peak_memory = layerwise_offload["peak_allocated_mb"]
     actual_saved_memory = no_offload_peak_memory - layerwise_offload_peak_memory
     assert layerwise_offload_peak_memory + expected_saved_memory <= no_offload_peak_memory, (
         f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB "

tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py

@@ -77,6 +77,7 @@ class _PipelineBackedEngine:
     def __init__(self, pipeline: _StepStreamingPipeline | _FailingStreamingPipeline) -> None:
         self.pipeline = pipeline
         self.executor = SimpleNamespace(
+            is_dead=False,
             register_failure_callback=MagicMock(),
             check_health=MagicMock(),
         )
@@ -127,9 +128,9 @@ def _streaming_diffusion_output(
 
 def _streaming_payload_and_metadata(output: DiffusionOutput) -> tuple[dict, dict]:
     envelope = output.output if isinstance(output.output, dict) else {}
-    payload = envelope.get("payload") if isinstance(envelope.get("payload"), dict) else {}
-    metadata = envelope.get("metadata") if isinstance(envelope.get("metadata"), dict) else {}
-    return payload, metadata
+    payload = envelope.get("payload")
+    metadata = envelope.get("metadata")
+    return payload if isinstance(payload, dict) else {}, metadata if isinstance(metadata, dict) else {}
 
 
 class TestPipelineStreamingOutputToStageDiffusionClient:
@@ -230,7 +231,7 @@ class TestPipelineStreamingOutputToStageDiffusionClient:
         client._response_socket.bind(response_address)
         client._encoder = OmniMsgpackEncoder()
         client._decoder = OmniMsgpackDecoder()
-        client._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
+        client._output_queue = asyncio.Queue[OmniRequestOutput]()
         client._rpc_results = {}
         client._pending_rpcs = set()
         client._tasks = {}

tests/diffusion/test_attention_compile.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
 
 from types import SimpleNamespace
 
@@ -14,6 +14,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 def test_attention_uses_compile_boundary_for_hsdp(monkeypatch):
     attention = object.__new__(Attention)
+    attention._hsdp_compile_boundary_enabled = False
     calls = []
 
     def _boundary(query, key, value, attn_metadata=None):
@@ -38,6 +39,7 @@ def test_attention_uses_compile_boundary_for_hsdp(monkeypatch):
 
 def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch):
     attention = object.__new__(Attention)
+    attention._hsdp_compile_boundary_enabled = False
     calls = []
 
     def _boundary(query, key, value, attn_metadata=None):
@@ -60,8 +62,10 @@ def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch):
     assert calls == ["impl"]
 
 
-def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch):
+@pytest.mark.parametrize("initialized_with_hsdp", [False, True])
+def test_attention_compile_boundary_without_diffusion_config(monkeypatch, initialized_with_hsdp):
     attention = object.__new__(Attention)
+    attention._hsdp_compile_boundary_enabled = initialized_with_hsdp
     calls = []
 
     def _boundary(query, key, value, attn_metadata=None):
@@ -80,4 +84,4 @@ def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch):
     with set_forward_context(vllm_config=None):
         assert Attention.forward(attention, query, query, query) is query
 
-    assert calls == ["impl"]
+    assert calls == (["boundary"] if initialized_with_hsdp else ["impl"])

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -16,6 +16,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
@@ -96,6 +97,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
@@ -2663,11 +2688,10 @@ 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.aux_output_connector = None
     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/features/yue2/worker_extension.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Observe YuE2 lifecycle contracts inside a real engine worker."""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+
+class Yue2LifecycleWorkerExtension:
+    model_runner: "GPUARModelRunner"
+
+    def start_yue2_probe(self):
+        from unittest.mock import patch
+
+        from vllm.compilation.cuda_graph import CUDAGraphWrapper
+        from vllm.config import CUDAGraphMode
+        from vllm.forward_context import get_forward_context, is_forward_context_available
+
+        from vllm_omni.model_executor.models.yue2.yue2 import ABC_END, HOLD_TOKEN, MUSIC_END
+
+        model = self.model_runner.get_model()
+        self.yue2_probe = {
+            "preemptions": 0,
+            "rollbacks": 0,
+            "checked_rows": 0,
+            "mixed_steps": 0,
+            "synthesis_started": False,
+            "ar_full_replays": 0,
+            "ar_piecewise_replays": 0,
+        }
+        graph_call = CUDAGraphWrapper.__call__
+
+        def count_ar_replays(wrapper, *args, **kwargs):
+            if is_forward_context_available():
+                context = get_forward_context()
+                mode = context.cudagraph_runtime_mode
+                entry = wrapper.concrete_cudagraph_entries.get(context.batch_descriptor)
+                if mode == wrapper.runtime_mode and entry is not None and entry.cudagraph is not None:
+                    if mode == CUDAGraphMode.FULL:
+                        self.yue2_probe["ar_full_replays"] += 1
+                    elif mode == CUDAGraphMode.PIECEWISE:
+                        self.yue2_probe["ar_piecewise_replays"] += 1
+            return graph_call(wrapper, *args, **kwargs)
+
+        # This worker is created for this test and exits with the engine.
+        # Count vLLM AR replays, separately from model-owned sampler/NAR graphs.
+        self.yue2_graph_patch = patch.object(CUDAGraphWrapper, "__call__", count_ar_replays)
+        self.yue2_graph_patch.start()
+        self.yue2_cancel_job = None
+        previous: dict[str, int] = {}
+        prepare = model.prepare_runner_inputs
+        reconcile = model._reconcile_history
+        queue = model._queue_synthesis
+        sample = model.sample
+        resolve = model._resolve_pending
+        emitted: dict[str, int] = {}
+
+        def resolve_pending():
+            if model._pending is not None:
+                event, host, states = model._pending
+                event.synchronize()
+                for state, token in zip(states, host[0].tolist()):
+                    end = ABC_END if state.constants.phase == "abc" else MUSIC_END
+                    expected = HOLD_TOKEN if token == end and not state.constants.skip_synthesis else token
+                    assert emitted[state.request_id] == expected, "engine token and model history disagree"
+            resolve()
+
+        def sample_tokens(logits, sampling_metadata):
+            result = sample(logits, sampling_metadata)
+            phases = {model._states[rid].constants.phase for rid, _, _ in model._step_rows if rid in model._states}
+            if len(phases) == 2:
+                self.yue2_probe["mixed_steps"] += 1
+            for (rid, _, _), row in zip(model._step_rows, result.sampled_token_ids.tolist()):
+                emitted[rid] = row[0]
+            return result
+
+        def prepare_inputs(**kwargs):
+            for rid, computed in zip(kwargs["req_ids"], kwargs["num_computed_tokens"]):
+                if int(computed) < previous.get(rid, 0):
+                    self.yue2_probe["preemptions"] += 1
+                previous[rid] = int(computed)
+            return prepare(**kwargs)
+
+        def reconcile_history(state, accepted):
+            if len(state.history) + int(state.end_drawn) > accepted:
+                self.yue2_probe["rollbacks"] += 1
+            reconcile(state, accepted)
+            assert len(state.history) + int(state.end_drawn) <= accepted
+            if not state.finish_ready and state.job is None and not state.finished:
+                assert len(state.history) == accepted
+                assert not state.truncated and not state.end_drawn
+            self.yue2_probe["checked_rows"] += 1
+
+        def queue_synthesis(state):
+            queue(state)
+            if "cancel-synthesis" in state.request_id:
+                job = state.job
+                assert job.started and job.work is not None and job.keep
+                assert job.event is None, "test must intercept an active synthesis job"
+                # Keep an actual NAR job active until the driver cancels it.
+                # Its first work units ran normally; stop pumping later units
+                # so cancellation cannot race against a short song finishing.
+                self.yue2_cancel_advance = job.advance
+                job.advance = lambda: None
+                self.yue2_cancel_job = job
+                self.yue2_probe["synthesis_started"] = True
+
+        model.prepare_runner_inputs = prepare_inputs
+        model._reconcile_history = reconcile_history
+        model._queue_synthesis = queue_synthesis
+        model.sample = sample_tokens
+        model._resolve_pending = resolve_pending
+        return True
+
+    def get_yue2_probe(self):
+        result = dict(self.yue2_probe)
+        job = self.yue2_cancel_job
+        if job is not None:
+            if job.cancelled:
+                # Exercise the original pump after cleanup: it must not submit
+                # further units even though this test retains the job object.
+                self.yue2_cancel_advance()
+            model = self.model_runner.get_model()
+            result["cancelled"] = job.cancelled
+            result["released"] = job._released and job.work is None and not job.keep and job.stream is None
+            result["removed"] = (
+                job.request_id not in model._states
+                and model._synthesis.active is not job
+                and job not in model._synthesis.waiting
+            )
+        return result

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/offline_inference/test_yue2.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""End-to-end regressions for YuE2 offline WAV export."""
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import soundfile as sf
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import get_model_prefix
+from vllm_omni.transformers_utils.repo_utils import hf_api
+
+pytestmark = [
+    pytest.mark.full_model,
+    pytest.mark.tts,
+    *hardware_marks(res={"cuda": ["L4", "H100", "B200"]}, num_cards=1),
+]
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+
+
+@pytest.fixture(scope="module")
+def vae_path() -> Path:
+    path = Path(os.environ.get("YUE2_VAE", get_model_prefix() + "m-a-p/YuE2-Vae"))
+    if not path.is_dir():
+        path = Path(hf_api().snapshot_download(str(path), allow_patterns=["*.json", "*.safetensors"]))
+    return path
+
+
+@pytest.mark.parametrize("max_frames,truncated", [(200, True), (9000, False)], ids=["budget", "natural-eos"])
+def test_offline_example_saves_audio(vae_path: Path, tmp_path: Path, max_frames: int, truncated: bool) -> None:
+    model = Path(os.environ.get("YUE2_MODEL_DIR", get_model_prefix() + "m-a-p/YuE2-3B"))
+    if not model.is_dir():
+        model = Path(
+            hf_api().snapshot_download(str(model), allow_patterns=["*.json", "*.safetensors", "qwen.tiktoken"])
+        )
+    output = tmp_path / "song.wav"
+    env = os.environ.copy()
+    env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+    env["PYTHONPATH"] = str(REPO_ROOT) + os.pathsep + env.get("PYTHONPATH", "")
+    completed = subprocess.run(
+        [
+            sys.executable,
+            str(REPO_ROOT / "examples/offline_inference/yue2/end2end.py"),
+            "--model",
+            str(model),
+            "--vae",
+            str(vae_path),
+            "--style",
+            "gentle acoustic folk, clear English vocals, 90 BPM",
+            "--lyrics",
+            "[Verse]\nMorning light across the bay\nCarry all my dreams away\n[Outro]",
+            "--cot",
+            "off",
+            "--seed",
+            "831001",
+            "--max-frames",
+            str(max_frames),
+            "--output",
+            str(output),
+        ],
+        cwd=REPO_ROOT,
+        env=env,
+        capture_output=True,
+        text=True,
+        timeout=900,
+    )
+    assert completed.returncode == 0, completed.stdout + completed.stderr
+    assert f"truncated={truncated}" in completed.stdout
+    audio, sample_rate = sf.read(output, always_2d=True)
+    assert sample_rate == 48000 and audio.shape[1] == 2
+    assert np.isfinite(audio).all() and np.abs(audio).max() > 0
+    duration = len(audio) / sample_rate
+    if truncated:
+        assert duration == pytest.approx(max_frames / 25, abs=0.1)
+    else:
+        assert 8 <= duration < max_frames / 25

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/e2e/online_serving/test_yue2.py

@@ -0,0 +1,434 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""E2E online tests for YuE2-3B text-to-music on ``/v1/audio/speech``.
+
+Scenario: heavy-metal "Twinkle Twinkle Little Star" — Chinese lyrics plus a
+metal style caption plus an ABC score of the melody (``extra_params.cot =
+"melody"``). The golden score is a SheetSage2 transcription of the Wikimedia
+piano rendition, reduced to a single voice (``tests/assets/yue2``).
+
+Two programmatic content checks, both calibrated against a same-seed control
+request with ``cot="off"`` (no ABC conditioning):
+
+- Style: CLAP zero-shot probability of "heavy metal" must clear 0.9.
+- Melody: the generated song is transcribed back with SheetSage2 and its
+  pitch-class sequence is LCS-matched against the golden score (max over
+  major/minor variants and 12 transpositions, tolerating pedal-tone
+  insertions). SheetSage2's transcription of distorted metal vocals carries
+  roughly ±0.05 noise, so the bar sits at 0.80; the ABC-conditioned request
+  must also beat the control.
+
+The content checks run in a dedicated verification interpreter because
+SheetSage2 wants a dependency set (mir_eval, pretty_midi, numpy<2) that must
+not leak into the serving environment. They are enabled by three env vars:
+
+- ``YUE2_E2E_VERIFY_PYTHON``: python with sheetsage2 + CLAP deps installed
+- ``YUE2_E2E_CLAP_DIR``: local laion/clap-htsat-unfused snapshot (~0.6 GB)
+- ``YUE2_E2E_SHEETSAGE2_DIR``: local m-a-p/SheetSage2 snapshot (~0.25 GB,
+  plus its MERT-v2 backbone ~2.5 GB from the HF cache)
+
+Without them the test still asserts the structural floor (non-empty payload
+above the byte floor) and passes on that alone, following the MiniMax Music 3
+precedent.
+
+Weekly rather than per-PR: each request is a 16 s song and the verification
+checkpoints add ~5.5 GB of downloads. The async/graph implementation was
+validated on H200; the hardware mark targets the CI fleet's H100. The older
+RTX 4090 result does not establish this implementation's memory requirements.
+"""
+
+import concurrent.futures
+import io
+import json
+import os
+import subprocess
+import wave
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import httpx
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "m-a-p/YuE2-3B"
+DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 900.0
+
+# 400 frames = 16 s of song. The byte floor sits far below the observed
+# payload (~3.0 MB of 48 kHz stereo 16-bit) to leave room for an early EOS.
+_FRAMES = 400
+# Short concurrent requests exercise row re-indexing as budgets differ.
+_CONCURRENT_FRAMES = 250
+_MIN_BYTES = 1_000_000
+
+LYRICS = "一闪一闪亮晶晶\n满天都是小星星\n挂在天空放光明\n好像千万小眼睛"
+CAPTION = (
+    "Chinese heavy metal, distorted electric guitars, aggressive double-kick drums, powerful vocals, 140 BPM, key of C"
+)
+SEED = 831001
+
+_VERIFY_PYTHON = os.environ.get("YUE2_E2E_VERIFY_PYTHON")
+_CLAP_DIR = os.environ.get("YUE2_E2E_CLAP_DIR")
+_SHEETSAGE2_DIR = os.environ.get("YUE2_E2E_SHEETSAGE2_DIR")
+_CONTENT_CHECKS_ENABLED = all([_VERIFY_PYTHON, _CLAP_DIR, _SHEETSAGE2_DIR])
+_VERIFY_HELPER = os.path.join(os.path.dirname(__file__), "yue2_metal_verify.py")
+
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("yue2.yaml"),
+            server_args=["--disable-log-stats"],
+        ),
+        id="yue2",
+    )
+]
+
+
+def _http_speech(omni_server, *, seed: int, frames: int) -> bytes:
+    """Raw-HTTP speech request returning (and structurally checking) the WAV.
+
+    The OpenAI SDK client swallows response headers, so the truncation flag is
+    cross-checked here: ``X-Audio-Truncated`` must agree with the audio — the
+    full frame budget (a 25 frames/s song of exactly frames/25 seconds) means
+    truncated, a shorter song means a natural end token.
+    """
+    response = httpx.post(
+        f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech",
+        json={
+            "model": omni_server.model,
+            "input": LYRICS,
+            "instructions": CAPTION,
+            "seed": seed,
+            "max_new_tokens": frames,
+            "stream": False,
+            "response_format": "wav",
+            "extra_params": {"cot": "off"},
+        },
+        timeout=DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+    )
+    assert response.status_code == 200, response.text[:500]
+    truncated = response.headers.get("X-Audio-Truncated")
+    assert truncated in ("true", "false"), f"X-Audio-Truncated missing/invalid: {truncated!r}"
+    audio = response.content
+    with wave.open(io.BytesIO(audio)) as wav:
+        assert wav.getnchannels() == 2
+        assert wav.getframerate() == 48000
+        duration = wav.getnframes() / wav.getframerate()
+    hit_budget = duration >= frames / 25 - 0.5
+    assert truncated == ("true" if hit_budget else "false"), (
+        f"X-Audio-Truncated={truncated} but duration={duration:.2f}s vs {frames / 25:.1f}s budget"
+    )
+    return audio
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None:
+    """Heavy-metal Twinkle: ABC-conditioned request vs same-seed control.
+
+    Deploy Setting: yue2.yaml
+    Input Modal: text (lyrics) + text (style caption) + ABC score (extra_params)
+    Output Modal: audio, 48 kHz stereo
+    Input Setting: stream=False, seed pinned, 400 frames
+    """
+    golden_abc = get_asset_path("yue2/twinkle_golden.abc").read_text()
+
+    def _request(extra_params: dict) -> bytes:
+        responses = openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": LYRICS,
+                "instructions": CAPTION,
+                "seed": SEED,
+                "max_new_tokens": _FRAMES,
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_BYTES,
+                # Sung metal vocals are not whisper-transcribable; skip the
+                # full_model transcript check (our content checks do the work).
+                "transcript_expected_text": "",
+                "extra_params": extra_params,
+            }
+        )
+        audio = responses[0].audio_bytes
+        assert audio, "empty audio payload"
+        return audio
+
+    metal_path = tmp_path / "metal.wav"
+    control_path = tmp_path / "control.wav"
+    metal_bytes = _request({"cot": "melody", "abc": golden_abc})
+    control_bytes = _request({"cot": "off"})
+    metal_path.write_bytes(metal_bytes)
+    control_path.write_bytes(control_bytes)
+
+    # CI runs without the verification assets, so the structural checks below
+    # are the only guard against the two real bugs this PR fixed during review:
+    # mono/half-speed output (WAV header + duration bounds) and the int32
+    # crash under concurrent row re-indexing (simultaneous requests).
+    # Duration bounds: the end token is masked for the first 200 steps
+    # (min_tokens), and the 400-frame budget caps the song, so a healthy
+    # request lands in [8, 16.5] s; the half-speed bug produced ~32 s.
+    for label, audio in (("metal", metal_bytes), ("control", control_bytes)):
+        with wave.open(io.BytesIO(audio)) as wav:
+            assert wav.getnchannels() == 2, f"{label}: {wav.getnchannels()} channels, want stereo"
+            assert wav.getframerate() == 48000, f"{label}: {wav.getframerate()} Hz, want 48 kHz"
+            duration = wav.getnframes() / wav.getframerate()
+            assert _FRAMES / 50 <= duration <= _FRAMES / 25 + 0.5, (
+                f"{label}: {duration:.2f}s outside [{_FRAMES / 50}, {_FRAMES / 25 + 0.5}]s"
+            )
+
+    # A pinned-seed request must flag truncation consistently with the audio
+    # (see _http_speech); the OpenAI SDK response does not expose headers.
+    _http_speech(omni_server, seed=SEED, frames=_FRAMES)
+
+    # Different budgets exercise row re-indexing as shorter requests finish.
+    # Compare the same-budget pair so length cannot hide an ignored seed.
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as pool:
+        futures = [
+            pool.submit(_http_speech, omni_server, seed=SEED + i, frames=frames)
+            for i, frames in enumerate((_CONCURRENT_FRAMES, _FRAMES, _CONCURRENT_FRAMES))
+        ]
+        audios = [f.result() for f in futures]
+    assert audios[0] != audios[2], "different seeds produced identical audio at the same frame budget"
+
+    # Without the verification assets the structural assertions above
+    # (non-empty payload, byte floor) are the whole test — pass on them.
+    if not _CONTENT_CHECKS_ENABLED:
+        return
+    assert _VERIFY_PYTHON is not None
+    assert _CLAP_DIR is not None
+    assert _SHEETSAGE2_DIR is not None
+
+    proc = subprocess.run(
+        [
+            _VERIFY_PYTHON,
+            _VERIFY_HELPER,
+            "--clap-dir",
+            _CLAP_DIR,
+            "--sheetsage2-dir",
+            _SHEETSAGE2_DIR,
+            "--work-dir",
+            str(tmp_path),
+            str(metal_path),
+            str(control_path),
+        ],
+        capture_output=True,
+        text=True,
+        timeout=1800,
+    )
+    assert proc.returncode == 0, f"verification helper failed:\n{proc.stderr}"
+    report = json.loads(proc.stdout.strip().splitlines()[-1])
+
+    metal_prob = report["metal_probability"]
+    assert metal_prob >= 0.9, f"style check: P(heavy metal)={metal_prob:.3f} < 0.9"
+
+    melody_score = report["melody_score"]
+    control_score = report["control_score"]
+    assert melody_score >= 0.80, f"melody check: score={melody_score:.3f} < 0.80"
+    assert melody_score > control_score, (
+        f"calibration: ABC-conditioned {melody_score:.3f} not above control {control_score:.3f}"
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_yue2_abort_then_next_request_succeeds(omni_server) -> None:
+    """Serving smoke check after a client timeout, not proof of job cleanup."""
+    # The chosen seed and budget time out on the tested cards. Require the
+    # timeout so this cannot silently become another successful request;
+    # whether it lands in AR or synthesis depends on the card.
+    with pytest.raises(httpx.TimeoutException):
+        httpx.post(
+            f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech",
+            json={
+                "model": omni_server.model,
+                "input": LYRICS,
+                "instructions": CAPTION,
+                "seed": SEED + 2,
+                "max_new_tokens": 3000,
+                "stream": False,
+                "response_format": "wav",
+                "extra_params": {"cot": "off"},
+            },
+            timeout=1.0,
+        )
+    _http_speech(omni_server, seed=SEED + 3, frames=_CONCURRENT_FRAMES)
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("chunked", [True, False], ids=["chunked", "unchunked"])
+@pytest.mark.parametrize("enforce_eager", [True, False], ids=["eager", "ar-graphs"])
+async def test_yue2_preemption_and_synthesis_abort(chunked, enforce_eager, monkeypatch) -> None:
+    """Force real preemption and abort only after NAR work has been submitted."""
+    import asyncio
+    from pathlib import Path
+
+    import torch
+    from vllm import SamplingParams
+
+    from tests.helpers.runtime import get_model_prefix
+    from tests.helpers.stage_config import modify_stage_config
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.model_executor.models.yue2.yue2 import (
+        ABC_END,
+        EOD,
+        KEY_MAX_AUDIO_FRAMES,
+        KEY_MAX_HOLD_STEPS,
+        KEY_MIN_TOKENS,
+        KEY_PHASE,
+        KEY_PREFIX_IDS,
+        KEY_SEED,
+        KEY_SKIP_SYNTHESIS,
+        STOP_TOKEN_IDS,
+    )
+    from vllm_omni.tokenizers.yue2_prompt import abc_prefix_ids, semantic_frames, semantic_prefix_ids
+    from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer
+    from vllm_omni.transformers_utils.repo_utils import hf_api
+
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    model_path = Path(os.environ.get("YUE2_MODEL_DIR", get_model_prefix() + MODEL))
+    if not model_path.is_dir():
+        model_path = Path(hf_api().snapshot_download(str(model_path)))
+    tokenizer = YuE2TextTokenizer(model_path / "qwen.tiktoken")
+    prefix = semantic_prefix_ids(tokenizer.encode, CAPTION, LYRICS, "off")
+    frames = 256
+    # Each song fits alone, but four full songs cannot fit in 64 KV blocks.
+    # Chunked recompute also replays output tokens past the original prompt.
+    deploy = modify_stage_config(
+        get_deploy_config_path("yue2.yaml"),
+        updates={
+            "stages": {
+                0: {
+                    "enforce_eager": enforce_eager,
+                    "compilation_config": {
+                        "cudagraph_mode": "NONE" if enforce_eager else "FULL_AND_PIECEWISE",
+                    },
+                    "enable_chunked_prefill": chunked,
+                    "async_scheduling": True,
+                    "max_num_batched_tokens": 256 if chunked else 1024,
+                    "max_model_len": 1000,
+                    "num_gpu_blocks_override": 64,
+                    "block_size": 16,
+                    "max_num_seqs": 4,
+                    "tokenizer": str(model_path),
+                }
+            }
+        },
+    )
+    engine = AsyncOmni(
+        model=str(model_path),
+        deploy_config=deploy,
+        worker_extension_cls="tests.e2e.features.yue2.worker_extension.Yue2LifecycleWorkerExtension",
+        stage_init_timeout=600,
+        init_timeout=900,
+    )
+
+    async def probe():
+        return (await engine.collective_rpc("get_yue2_probe", stage_ids=[0]))[0][0]
+
+    async def song(request_id, seed):
+        params = SamplingParams(
+            max_tokens=1000 - len(prefix),
+            stop_token_ids=STOP_TOKEN_IDS,
+            detokenize=False,
+            extra_args={
+                KEY_PREFIX_IDS: prefix,
+                KEY_SEED: seed,
+                KEY_MAX_AUDIO_FRAMES: frames,
+                KEY_MIN_TOKENS: frames,
+                KEY_MAX_HOLD_STEPS: 1000 - len(prefix) - frames - 1,
+            },
+        )
+        audio = []
+        tokens = []
+        async for output in engine.generate(
+            {"prompt_token_ids": prefix}, request_id=request_id, sampling_params_list=[params]
+        ):
+            tokens = list(output.outputs[0].token_ids)
+            multimodal = getattr(output.outputs[0], "multimodal_output", None)
+            if multimodal is not None and "audio" in multimodal:
+                data = multimodal["audio"]
+                audio.extend(data if isinstance(data, list) else [data])
+        assert audio, "request completed without audio"
+        waveform = torch.cat([part.detach().cpu().reshape(-1) for part in audio])
+        assert torch.isfinite(waveform).all() and waveform.abs().max() > 0
+        assert len(semantic_frames(tokens)) == frames
+        # The VAE convolution trims a small number of boundary samples.
+        assert waveform.numel() == pytest.approx(frames * 1920 * 2, abs=256)
+
+    async def plan():
+        abc_prefix = abc_prefix_ids(tokenizer.encode, CAPTION, LYRICS, "melody")
+        params = SamplingParams(
+            max_tokens=64,
+            stop_token_ids=STOP_TOKEN_IDS,
+            detokenize=False,
+            extra_args={
+                KEY_PHASE: "abc",
+                KEY_PREFIX_IDS: abc_prefix,
+                KEY_SEED: SEED,
+                KEY_SKIP_SYNTHESIS: True,
+            },
+        )
+        tokens = []
+        async for output in engine.generate(
+            {"prompt_token_ids": abc_prefix}, request_id="abc-plan", sampling_params_list=[params]
+        ):
+            tokens = list(output.outputs[0].token_ids)
+        assert tokens and all(t < EOD or t == ABC_END for t in tokens)
+
+    tasks = []
+    try:
+        assert await engine.collective_rpc("start_yue2_probe", stage_ids=[0]) == [[True]]
+        tasks = [asyncio.create_task(plan())]
+        tasks.extend(asyncio.create_task(song(f"preempt-{i}", SEED + i)) for i in range(4))
+        await asyncio.wait_for(asyncio.gather(*tasks), timeout=600)
+        observations = await probe()
+        assert observations["preemptions"] > 0, observations
+        assert observations["rollbacks"] > 0, observations
+        assert observations["checked_rows"] >= 4 * frames, observations
+        assert observations["mixed_steps"] > 0, observations
+        if enforce_eager:
+            assert observations["ar_full_replays"] == 0, observations
+        else:
+            assert observations["ar_full_replays"] > 0, observations
+
+        cancelled = asyncio.create_task(song("cancel-synthesis", SEED + 7))
+        tasks.append(cancelled)
+        deadline = asyncio.get_running_loop().time() + 120
+        while not (await probe())["synthesis_started"]:
+            assert not cancelled.done(), "song finished before the active-job cancellation probe"
+            assert asyncio.get_running_loop().time() < deadline, "NAR job never started"
+            await asyncio.sleep(0.05)
+        cancelled.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await cancelled
+        deadline = asyncio.get_running_loop().time() + 30
+        while True:
+            observations = await probe()
+            if observations.get("removed"):
+                assert observations["cancelled"] and observations["released"], observations
+                break
+            assert asyncio.get_running_loop().time() < deadline, observations
+            await asyncio.sleep(0.05)
+        await asyncio.wait_for(song("after-cancel", SEED + 8), timeout=120)
+        observations = await probe()
+        assert observations["released"]
+        print("YuE2 lifecycle:", observations)
+    finally:
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        engine.shutdown()

tests/e2e/online_serving/yue2_metal_verify.py

@@ -0,0 +1,185 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Content-gate helper for the YuE2 heavy-metal Twinkle e2e test.
+
+Run as ``__main__`` by a dedicated verification interpreter (see
+``YUE2_E2E_VERIFY_PYTHON`` in test_yue2.py), NOT by the serving environment:
+SheetSage2 wants its own dependency set (mir_eval, pretty_midi, numpy<2),
+which must not leak into the serving venv.
+
+Usage:
+    python yue2_metal_verify.py --clap-dir DIR --sheetsage2-dir DIR \
+        metal.wav control.wav
+
+Prints one JSON object on stdout:
+    {"metal_probability": float, "melody_score": float, "control_score": float}
+
+Gates (applied by the caller):
+- metal_probability >= 0.9  (CLAP zero-shot heavy-metal style)
+- melody_score >= 0.80 and melody_score > control_score
+  (SheetSage2 re-transcription, pitch-class LCS vs the golden score, max over
+  major/minor variants and 12 transpositions, pedal-tone tolerant)
+"""
+
+import argparse
+import json
+import re
+import wave
+from typing import Any
+
+import numpy as np
+
+_NOTE_BASE = {"C": 0, "D": 2, "E": 4, "F": 5, "G": 7, "A": 9, "B": 11}
+_GOLDEN_MAJOR = "C4C4G4G4|A4A4G8|F4F4E4E4|D4D4C8|G4G4F4F4|E4E4D8|G4G4F4F4|E4E4D8|C4C4G4G4|A4A4G8|F4F4E4E4|D4D4C8|"
+# Metal arrangements commonly minor-ize the melody; accept both modes.
+_GOLDEN_MINOR = (
+    "C4C4G4G4|_A4_A4G8|F4F4_E4_E4|D4D4C8|G4G4F4F4|_E4_E4D8|G4G4F4F4|_E4_E4D8|C4C4G4G4|_A4_A4G8|F4F4_E4_E4|D4D4C8|"
+)
+
+_STYLE_CANDIDATES = [
+    "heavy metal music with distorted guitars and aggressive drums",
+    "a soft children's lullaby",
+    "classical piano music",
+    "pop music",
+]
+
+
+def _abc_note_to_midi(token):
+    m = re.match(r"([\^_=]*)([A-Ga-g])([,']*)(\d*)", token)
+    if not m:
+        return None
+    acc, letter, oct_marks, _ = m.groups()
+    pc = _NOTE_BASE[letter.upper()] + acc.count("^") - acc.count("_")
+    octave = 5 if letter.isupper() else 6
+    octave += oct_marks.count("'") - oct_marks.count(",")
+    return 12 * octave + pc
+
+
+def _collapse(seq):
+    out: list[Any] = []
+    for x in seq:
+        if not out or x != out[-1]:
+            out.append(x)
+    return out
+
+
+def _golden_events(body):
+    return _collapse([_abc_note_to_midi(t) for t in re.findall(r"[\^_=]*[A-G]\d*", body)])
+
+
+def _melody_events(score_abc):
+    """Note events of the voice carrying the most notes in a SheetSage2 score."""
+    voices: dict[str | None, list[int]] = {}
+    current = None
+    for line in score_abc.splitlines():
+        line = line.strip()
+        if line.startswith("V:"):
+            parts = line.split()
+            current = parts[1] if len(parts) > 1 else None
+            continue
+        if current is None or not line or line.startswith(("%", "X:", "T:", "M:", "L:", "Q:", "K:")):
+            continue
+        for tok in re.findall(r"[\^_=]*[A-Ga-g][,']*\d*", line.replace("|", " ")):
+            midi = _abc_note_to_midi(tok)
+            if midi is not None:
+                voices.setdefault(current, []).append(midi)
+    return _collapse(max(voices.values(), key=len, default=[]))
+
+
+def _lcs(a, b):
+    dp = [0] * (len(b) + 1)
+    for x in a:
+        prev = 0
+        for j, y in enumerate(b, 1):
+            tmp = dp[j]
+            dp[j] = prev + 1 if x == y else max(dp[j], dp[j - 1])
+            prev = tmp
+    return dp[-1]
+
+
+def melody_score(score_abc):
+    gen = [m % 12 for m in _melody_events(score_abc)]
+    if len(gen) < 4:
+        return 0.0
+    best = 0
+    for body in (_GOLDEN_MAJOR, _GOLDEN_MINOR):
+        gold = [m % 12 for m in _golden_events(body)]
+        best = max(best, max(_lcs(gen, [(p + r) % 12 for p in gold]) for r in range(12)))
+    return best / len(_golden_events(_GOLDEN_MAJOR))
+
+
+def _wav_to_mono(path):
+    with wave.open(path) as wf:
+        sr = wf.getframerate()
+        channels = wf.getnchannels()
+        frames = wf.readframes(wf.getnframes())
+    data = np.frombuffer(frames, dtype=np.int16).astype(np.float32) / 32768.0
+    if channels > 1:
+        data = data.reshape(-1, channels).mean(axis=1)
+    return data, sr
+
+
+def metal_probability(wav_path, clap_dir):
+    import torch
+    from transformers import ClapModel, ClapProcessor
+
+    model = ClapModel.from_pretrained(clap_dir)
+    proc = ClapProcessor.from_pretrained(clap_dir)
+    audio, sr = _wav_to_mono(wav_path)
+    if sr != 48000:
+        n = int(len(audio) * 48000 / sr)
+        audio = np.interp(np.linspace(0, len(audio), n), np.arange(len(audio)), audio).astype(np.float32)
+    audio = audio[: 48000 * 10]
+    inputs = proc(
+        text=_STYLE_CANDIDATES,
+        audios=[audio] * len(_STYLE_CANDIDATES),
+        sampling_rate=48000,
+        return_tensors="pt",
+        padding=True,
+    )
+    with torch.no_grad():
+        probs = model(**inputs).logits_per_audio[0].softmax(dim=-1)
+    return float(probs[0])
+
+
+def transcribe_melody_abc(wav_path, sheetsage2_dir, output_dir):
+    import torch
+    from transformers import AutoModel
+
+    audio, sr = _wav_to_mono(wav_path)
+    model = AutoModel.from_pretrained(sheetsage2_dir, trust_remote_code=True)
+    result = model.transcribe(
+        torch.from_numpy(audio[None, :]),
+        sampling_rate=sr,
+        output_dir=str(output_dir),
+        melody_only=True,
+        dtype="fp32",
+    )
+    abc = result.get("abc")
+    if not abc:
+        raise RuntimeError(f"SheetSage2 produced no melody ABC: {result.get('abc_error')}")
+    return abc
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--clap-dir", required=True)
+    parser.add_argument("--sheetsage2-dir", required=True)
+    parser.add_argument("metal_wav")
+    parser.add_argument("control_wav")
+    parser.add_argument("--work-dir", default=None)
+    args = parser.parse_args()
+
+    work = args.work_dir or "."
+    report = {"metal_probability": metal_probability(args.metal_wav, args.clap_dir)}
+    report["melody_score"] = melody_score(
+        transcribe_melody_abc(args.metal_wav, args.sheetsage2_dir, f"{work}/metal_score")
+    )
+    report["control_score"] = melody_score(
+        transcribe_melody_abc(args.control_wav, args.sheetsage2_dir, f"{work}/control_score")
+    )
+    print(json.dumps(report))
+
+
+if __name__ == "__main__":
+    main()

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_orchestrator_first_audio.py

@@ -114,10 +114,10 @@ def _raw(*, first=False, required=False, terminal=False, samples=(1, 1, 1, 1)):
     )
 
 
-async def _codec_output(obj, raw):
-    replica = obj.stage_pools[1].replica_id
-    processed = await obj._process_llm_stage_outputs(1, replica, raw, set())
-    await obj._handle_processed_outputs(1, replica, processed)
+async def _codec_output(obj, raw, stage=1):
+    replica = obj.stage_pools[stage].replica_id
+    processed = await obj._process_llm_stage_outputs(stage, replica, raw, set())
+    await obj._handle_processed_outputs(stage, replica, processed)
 
 
 def _messages(obj):
@@ -132,6 +132,57 @@ def _audio(output):
     return torch.cat(audio) if isinstance(audio, list) else audio
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("kind", list(RequestOutputKind))
+@pytest.mark.parametrize("suffix_overtook", [False, True])
+async def test_single_stage_first_audio_obeys_output_kind_and_terminal_order(kind, suffix_overtook):
+    obj = _orchestrator(kind, final_stage_id=0)
+    obj.stage_pools = [obj.stage_pools[1]]
+    suffix = _raw(required=True, samples=(20, 21))
+    terminal = _raw(required=True, terminal=True, samples=(30, 31))
+    if suffix_overtook:
+        await _codec_output(obj, suffix, stage=0)
+        await _codec_output(obj, terminal, stage=0)
+        assert obj.output_async_queue.empty()
+    await obj._route_upstream_first_audio(0, 2, _raw(first=True, samples=(10, 11)))
+    if not suffix_overtook:
+        await _codec_output(obj, suffix, stage=0)
+        await _codec_output(obj, terminal, stage=0)
+    messages = _messages(obj)
+    expected = {
+        RequestOutputKind.DELTA: [[10, 11], [20, 21], [30, 31]],
+        RequestOutputKind.CUMULATIVE: [[10, 11], [10, 11, 20, 21], [10, 11, 20, 21, 30, 31]],
+        RequestOutputKind.FINAL_ONLY: [[10, 11, 20, 21, 30, 31]],
+    }
+    assert [_audio(message.engine_outputs).tolist() for message in messages] == expected[kind]
+    assert all((message.stage_id, message.replica_id) == (0, 2) for message in messages)
+    assert messages[-1].finished
+    assert not obj.request_states["r"].pending_first_audio_outputs
+
+
+@pytest.mark.asyncio
+async def test_single_stage_cancellation_discards_overtaking_terminal_and_late_first_audio():
+    obj = _orchestrator(final_stage_id=0)
+    obj.stage_pools = [obj.stage_pools[1]]
+    await _codec_output(obj, _raw(required=True, terminal=True), stage=0)
+    assert obj.request_states["r"].pending_first_audio_outputs
+    await obj._cleanup_request_ids(["r"], abort=True)
+    await obj._route_upstream_first_audio(0, 2, _raw(first=True))
+    assert not obj.request_states and obj.output_async_queue.empty()
+
+
+@pytest.mark.asyncio
+async def test_single_stage_first_audio_processing_failure_releases_waiting_terminal():
+    obj = _orchestrator(final_stage_id=0)
+    obj.stage_pools = [obj.stage_pools[1]]
+    await _codec_output(obj, _raw(required=True, terminal=True), stage=0)
+    obj.stage_pools[0].process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed"))
+    await obj._route_upstream_first_audio(0, 2, _raw(first=True))
+    error = obj.output_async_queue.get_nowait()
+    assert isinstance(error, ErrorMessage) and error.stage_id == 0
+    assert not obj.request_states
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("kind", list(RequestOutputKind))
 @pytest.mark.parametrize("codec_overtook", [False, True])

tests/engine/test_stage_engine_core_proc.py

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

tests/engine/test_stage_pool_collective_rpc.py

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

tests/entrypoints/openai_api/test_audio_format.py

@@ -346,3 +346,19 @@ def test_chat_audio_metadata_survives_response_serialization(stream, fmt):
             assert audio_file.samplerate == 22050
     else:
         assert len(raw) == frame_count * 2
+
+
+def test_float32_pcm_fast_path_matches_soundfile():
+    from vllm_omni.entrypoints.openai.audio_utils_mixin import _float32_to_pcm16_bytes
+
+    rng = np.random.default_rng(0)
+    audio = np.concatenate(
+        [
+            rng.uniform(-1.2, 1.2, 4096),
+            np.array([-1.0, 1.0, 0.0, -0.0, 32767 / 32768, -1 / 32768, 0.5 / 32768, -0.5 / 32768]),
+        ]
+    ).astype(np.float32)
+    with BytesIO() as buffer:
+        soundfile.write(buffer, audio, 24000, format="RAW", subtype="PCM_16")
+        expected = buffer.getvalue()
+    assert _float32_to_pcm16_bytes(audio) == expected

tests/entrypoints/openai_api/test_moss_local_prompt.py

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

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1105,6 +1105,38 @@ class TestTTSMethods:
         assert isinstance(audio_obj, CreateAudio)
         assert audio_obj.speed == 1.0
 
+    @pytest.mark.asyncio
+    async def test_audio_synthesis_error_flag_raises_tts_generation_error(
+        self,
+        speech_server,
+        mocker: MockerFixture,
+    ):
+        """A model-flagged synthesis failure surfaces as a non-retryable
+        TTSGenerationError, so create_speech answers 500 instead of shipping a
+        zero-length WAV or crashing on an unpacked Response."""
+
+        async def mock_generate():
+            yield create_mock_audio_output_for_test()
+
+        mocker.patch.object(
+            speech_server,
+            "_prepare_speech_generation",
+            new=mocker.AsyncMock(return_value=("speech-synth-err", mock_generate(), {})),
+        )
+        adapter = mocker.MagicMock()
+        adapter.collect_response_metadata = lambda _audio_output, collect: collect.__setitem__(
+            "audio_synthesis_error", True
+        )
+        mocker.patch.object(speech_server, "_get_tts_adapter", return_value=adapter)
+
+        with pytest.raises(TTSGenerationError, match="failed to synthesize audio") as exc_info:
+            await speech_server._generate_audio_bytes(
+                OpenAICreateSpeechRequest(input="Hello"),
+                collect={},
+            )
+
+        assert exc_info.value.retryable is False
+
     def test_is_tts_detection_with_tts_stage(self, mocker: MockerFixture):
         """Test TTS model detection when TTS stage exists."""
         mock_engine_client = mocker.MagicMock()

tests/entrypoints/openai_api/test_serving_video_stream.py

@@ -7,10 +7,12 @@ from __future__ import annotations
 import asyncio
 import base64
 import gc
+import hashlib
 import io
 import json
 import threading
 import weakref
+from types import SimpleNamespace
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock
 
@@ -26,7 +28,9 @@ from vllm_omni.entrypoints.openai.serving_video_stream import (
     QwenOmniStreamingVideoHandler,
     StreamingVideoSessionConfig,
 )
-from vllm_omni.entrypoints.openai.video_stream_base import OmniStreamingVideoHandler
+from vllm_omni.entrypoints.openai.video_stream_base import (
+    OmniStreamingVideoHandler,
+)
 from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -506,17 +510,9 @@ async def test_video_frames_consumed_excludes_bad_frame_in_query_snapshot(
         pytest.param(2, ["f0", "f1", "f2"], True, [0, 2], id="repeated-image"),
     ],
 )
-async def test_video_frame_selection_preserves_uncached_frames(
-    monkeypatch, num_frames, frame_ids, repeat_image, expected_indices
-):
-    """Prewarm misses still use image URLs, with metadata paired by buffer position."""
+async def test_video_frame_selection_pairs_metadata_by_position(num_frames, frame_ids, repeat_image, expected_indices):
+    """Queries wait for prewarmed PIL frames, with metadata paired by buffer position."""
     engine_prompts = []
-    pending_decode = asyncio.Event()
-
-    async def blocked_to_thread(*args, **kwargs):
-        await pending_decode.wait()
-
-    monkeypatch.setattr(asyncio, "to_thread", blocked_to_thread)
 
     class OneOutputEngine:
         async def generate(self, **kwargs):
@@ -550,13 +546,13 @@ async def test_video_frame_selection_preserves_uncached_frames(
     assert not [message for message in ws.sent if message["type"] == "error"]
     assert len(engine_prompts) == 1
     content = engine_prompts[0]["messages"][-1]["content"]
-    assert content == [
-        *[
-            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{frames[index]}"}}
-            for index in expected_indices
-        ],
-        {"type": "text", "text": "describe"},
-    ]
+    assert content[-1] == {"type": "text", "text": "describe"}
+    image_parts = content[:-1]
+    assert [part["type"] for part in image_parts] == ["image_pil"] * len(expected_indices)
+    for part, index in zip(image_parts, expected_indices):
+        raw = base64.b64decode(frames[index])
+        assert part["image_pil"].tobytes() == Image.open(io.BytesIO(raw)).convert("RGB").tobytes()
+        assert part["uuid"] == hashlib.md5(raw, usedforsecurity=False).hexdigest()
     consumed_events = [message for message in ws.sent if message["type"] == "video.frames.consumed"]
     if not any(frame_ids):
         assert consumed_events == []
@@ -909,6 +905,9 @@ async def test_interrupted_queries_wait_for_abort_without_fixed_delay(monkeypatc
 
     class BlockingEngine:
         async def generate(self, *, request_id, **kwargs):
+            submitted = kwargs.get("submitted")
+            if submitted is not None and not submitted.done():
+                submitted.set_result(None)
             if request_ids:
                 assert aborted[-1] == request_ids[-1]
             request_ids.append(request_id)
@@ -1004,7 +1003,7 @@ async def test_video_done_waits_for_in_flight_query():
 
 
 @pytest.mark.asyncio
-async def test_frame_prewarm_does_not_block_following_query(monkeypatch):
+async def test_query_waits_for_frame_prewarm(monkeypatch):
     decode_started = threading.Event()
     release_decode = threading.Event()
     query_started = asyncio.Event()
@@ -1035,14 +1034,80 @@ async def test_frame_prewarm_does_not_block_following_query(monkeypatch):
     assert decode_started.is_set()
 
     ws.put({"type": "video.query", "text": "describe"})
-    await asyncio.wait_for(query_started.wait(), timeout=2.0)
+    await asyncio.sleep(0.05)
+    assert not query_started.is_set()
 
     release_decode.set()
+    await asyncio.wait_for(query_started.wait(), timeout=2.0)
     ws.put({"type": "video.done"})
     await asyncio.wait_for(task, timeout=2.0)
     assert "session.done" in ws.sent_types()
 
 
+@pytest.mark.asyncio
+async def test_legacy_query_does_not_wait_for_unselected_frame(monkeypatch):
+    """APC off and num_frames=1: a blocked older frame must not delay the query."""
+    frame_a = _make_jpeg(1, 0, 0)
+    frame_b = _make_jpeg(0, 1, 0)
+    decode_a_started = threading.Event()
+    release_a = threading.Event()
+    decode_b_done = threading.Event()
+    query_started = asyncio.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        if raw_bytes == frame_a:
+            decode_a_started.set()
+            release_a.wait(timeout=5.0)
+        image = Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+        if raw_bytes == frame_b:
+            decode_b_done.set()
+        return image
+
+    class RecordingHandler(QwenOmniStreamingVideoHandler):
+        async def _process_query(self, *args, **kwargs):
+            query_started.set()
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    ws = TimedWebSocket()
+    handler = RecordingHandler(chat_service=object(), engine_client=object(), idle_timeout=5.0)
+    task = asyncio.create_task(handler.handle_session(ws))
+    try:
+        ws.put(
+            {
+                "type": "session.config",
+                "model": "test",
+                "modalities": ["text"],
+                "num_frames": 1,
+                "enable_frame_filter": False,
+            }
+        )
+        await asyncio.sleep(0)
+        ws.put({"type": "video.frame", "data": _b64(frame_a), "frame_id": "A"})
+        for _ in range(100):
+            if decode_a_started.is_set():
+                break
+            await asyncio.sleep(0.01)
+        assert decode_a_started.is_set()
+        ws.put({"type": "video.frame", "data": _b64(frame_b), "frame_id": "B"})
+        for _ in range(200):
+            if decode_b_done.is_set():
+                break
+            await asyncio.sleep(0.01)
+        assert decode_b_done.is_set()
+
+        ws.put({"type": "video.query", "text": "describe"})
+        await asyncio.wait_for(query_started.wait(), timeout=2.0)
+        assert not release_a.is_set()
+        release_a.set()
+        ws.put({"type": "video.done"})
+        await asyncio.wait_for(task, timeout=2.0)
+    finally:
+        release_a.set()
+        if not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("max_frames", [1, 2], ids=["evicted", "retained"])
 async def test_frame_prewarm_only_keeps_retained_images(monkeypatch, max_frames):
@@ -1272,3 +1337,452 @@ def test_build_messages_keeps_recent_history_text_only():
     assert messages[1] == {"role": "assistant", "content": "recent answer"}
     assert messages[2] == user_message
     assert user_message["content"][-1] == {"type": "text", "text": "current question"}
+
+
+# ---------------------------------------------------------------------------
+# Incremental prefill (warmup) lifecycle
+# ---------------------------------------------------------------------------
+
+
+class RecordingReuseEngine:
+    """Engine stub with stage-0 prefix caching on; records generate() and abort()."""
+
+    def __init__(self, warmup_delay: float = 0.0, query_delay: float = 0.0):
+        self.requests: list[dict[str, Any]] = []
+        self.aborted: list[str] = []
+        self._warmup_delay = warmup_delay
+        self._query_delay = query_delay
+        self.engine = SimpleNamespace(
+            stage_vllm_configs=[SimpleNamespace(cache_config=SimpleNamespace(enable_prefix_caching=True))]
+        )
+
+    def warmup_ids(self) -> list[str]:
+        return [r["request_id"] for r in self.requests if r["request_id"].startswith("video-warmup-")]
+
+    def query_ids(self) -> list[str]:
+        return [r["request_id"] for r in self.requests if not r["request_id"].startswith("video-warmup-")]
+
+    def generate(self, **kwargs):
+        self.requests.append(kwargs)
+        submitted = kwargs.get("submitted")
+        if submitted is not None and not submitted.done():
+            submitted.set_result(None)
+        is_warmup = kwargs["request_id"].startswith("video-warmup-")
+        delay = self._warmup_delay if is_warmup else self._query_delay
+
+        async def _gen():
+            if delay:
+                await asyncio.sleep(delay)
+            yield _text_result("ok")
+
+        return _gen()
+
+    async def abort(self, request_id: str) -> None:
+        self.aborted.append(request_id)
+
+
+class PromptRecordingHandler(QwenOmniStreamingVideoHandler):
+    """Records each query turn's user-message content and processor kwargs."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.query_contents: list[list[dict[str, Any]]] = []
+        self.warmup_contents: list[list[dict[str, Any]]] = []
+        self.mm_processor_kwargs: list[Any] = []
+
+    async def _preprocess_to_engine_prompt(self, request):
+        self.mm_processor_kwargs.append(getattr(request, "mm_processor_kwargs", None))
+        return {"prompt": "p"}
+
+    def build_engine_prompt(
+        self, config, frame_buffer, audio_buffer, message_history, query_text, prewarmed_frames, *, frame_indices=None
+    ):
+        messages, user_message = super().build_engine_prompt(
+            config,
+            frame_buffer,
+            audio_buffer,
+            message_history,
+            query_text,
+            prewarmed_frames,
+            frame_indices=frame_indices,
+        )
+        self.query_contents.append(list(user_message["content"]))
+        return messages, user_message
+
+    def build_engine_prompt_prefix(self, config, frame_buffer, message_history, prewarmed_frames):
+        messages = super().build_engine_prompt_prefix(config, frame_buffer, message_history, prewarmed_frames)
+        if messages:
+            self.warmup_contents.append(list(messages[-1]["content"]))
+        return messages
+
+
+async def _poll(predicate, timeout: float = 3.0) -> bool:
+    for _ in range(int(timeout / 0.01)):
+        if predicate():
+            return True
+        await asyncio.sleep(0.01)
+    return bool(predicate())
+
+
+@pytest.mark.asyncio
+async def test_warmup_fires_on_frame_and_query_cancels_it():
+    engine = RecordingReuseEngine(warmup_delay=10.0, query_delay=0.3)
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+
+    assert await _poll(lambda: engine.warmup_ids())
+    warmup_id = engine.warmup_ids()[0]
+    assert engine.requests[0]["sampling_params"].max_tokens == 1
+
+    ws.put({"type": "video.query", "text": "describe"})
+    assert await _poll(lambda: engine.query_ids())
+    # in-flight warmup was cancelled + aborted, and not restarted mid-query
+    assert warmup_id in engine.aborted
+    assert len(engine.warmup_ids()) == 1
+    assert handler.mm_processor_kwargs[:2] == [
+        {"use_audio_in_video": True},
+        {"use_audio_in_video": True},
+    ]
+
+    assert await _poll(lambda: "response.text.done" in ws.sent_types())
+    # the committed turn changed the context -> a fresh warmup fires
+    assert await _poll(lambda: len(engine.warmup_ids()) == 2)
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_cancelled_warmup_reaches_stage0_before_abort():
+    """Cancel during submit must not abort a warmup the engine never received."""
+
+    class _HoldSubmitEngine(RecordingReuseEngine):
+        def __init__(self):
+            super().__init__()
+            self.started = asyncio.Event()
+            self.release_submit = asyncio.Event()
+
+        def generate(self, **kwargs):
+            submitted = kwargs.get("submitted")
+
+            async def _gen():
+                self.started.set()
+                await self.release_submit.wait()
+                self.requests.append(kwargs)
+                if submitted is not None and not submitted.done():
+                    submitted.set_result(None)
+                yield _text_result("ok")
+
+            return _gen()
+
+    engine = _HoldSubmitEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    assert await _poll(engine.started.is_set)
+    ws.put({"type": "video.query", "text": "describe"})
+    await asyncio.sleep(0.05)
+    assert engine.requests == []
+    assert engine.aborted == []
+    engine.release_submit.set()
+    assert await _poll(lambda: engine.aborted)
+    assert engine.warmup_ids()
+    assert engine.aborted[0] == engine.warmup_ids()[0]
+    assert await _poll(lambda: engine.query_ids())
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_cancelled_query_reaches_stage0_before_abort():
+    """Cancel during query submit must not abort a request the engine never received."""
+
+    class _HoldQueryEngine(RecordingReuseEngine):
+        def __init__(self):
+            super().__init__()
+            self.query_started = asyncio.Event()
+            self.release_submit = asyncio.Event()
+
+        def generate(self, **kwargs):
+            if str(kwargs.get("request_id", "")).startswith("video-warmup-"):
+                return super().generate(**kwargs)
+            submitted = kwargs.get("submitted")
+
+            async def _gen():
+                self.query_started.set()
+                await self.release_submit.wait()
+                self.requests.append(kwargs)
+                if submitted is not None and not submitted.done():
+                    submitted.set_result(None)
+                yield _text_result("ok")
+
+            return _gen()
+
+    engine = _HoldQueryEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    assert await _poll(lambda: engine.warmup_ids())
+    ws.put({"type": "video.query", "text": "first"})
+    assert await _poll(engine.query_started.is_set)
+    ws.put({"type": "video.query", "text": "second"})
+    await asyncio.sleep(0.05)
+    assert engine.query_ids() == []
+    assert engine.aborted == []
+    engine.release_submit.set()
+    assert await _poll(lambda: engine.query_ids())
+    first_query = engine.query_ids()[0]
+    assert await _poll(lambda: first_query in engine.aborted)
+    assert await _poll(lambda: len(engine.query_ids()) == 2)
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_audio_session_never_starts_warmup():
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text", "audio"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg(9, 9, 9))})
+    await asyncio.sleep(0.3)
+
+    assert engine.warmup_ids() == []
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_warmup_starts_only_after_frames_ready(monkeypatch):
+    decode_started = threading.Event()
+    release_decode = threading.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        decode_started.set()
+        release_decode.wait(timeout=5.0)
+        return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    assert await _poll(decode_started.is_set)
+    await asyncio.sleep(0.1)
+    assert engine.warmup_ids() == []
+
+    release_decode.set()
+    assert await _poll(lambda: engine.warmup_ids())
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_query_cancelling_warmup_keeps_prewarm_alive(monkeypatch):
+    """Cancel warmup during prefill must not drop a later frame still decoding."""
+    frame1 = _b64(_make_jpeg(1, 1, 1))
+    frame2 = _b64(_make_jpeg(2, 2, 2))
+    raw2 = base64.b64decode(frame2)
+    decode_started = threading.Event()
+    release_decode = threading.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        if raw_bytes == raw2:
+            decode_started.set()
+            release_decode.wait(timeout=5.0)
+        return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    engine = RecordingReuseEngine(warmup_delay=10.0)
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": frame1})
+    assert await _poll(lambda: engine.warmup_ids())
+
+    ws.put({"type": "video.frame", "data": frame2})
+    assert await _poll(decode_started.is_set)
+
+    ws.put({"type": "video.query", "text": "describe"})
+    await asyncio.sleep(0.1)
+    release_decode.set()
+
+    assert await _poll(lambda: "response.text.done" in ws.sent_types())
+    assert len(handler.query_contents) == 1
+    image_parts = [p for p in handler.query_contents[0] if p.get("type") == "image_pil"]
+    assert len(image_parts) == 2
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_query_keeps_frames_evicted_while_pinned(monkeypatch):
+    """Regression: frames FIFO-evicted while an in-flight query awaits prewarm
+    must keep their PIL entries and stay in the query prompt."""
+    frame_a = _b64(_make_jpeg(1, 1, 1))
+    frame_a2 = _b64(_make_jpeg(2, 2, 2))
+    frame_b = _b64(_make_jpeg(3, 3, 3))
+    raw_a = base64.b64decode(frame_a)
+    decoded: list[bytes] = []
+    release_others = threading.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        if raw_bytes != raw_a:
+            release_others.wait(timeout=5.0)
+        decoded.append(raw_bytes)
+        return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put(
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": ["text"],
+            "max_frames": 2,
+            "enable_frame_filter": False,
+        }
+    )
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": frame_a})
+    assert await _poll(lambda: raw_a in decoded)
+    ws.put({"type": "video.frame", "data": frame_a2})  # decode gated -> query will wait
+
+    ws.put({"type": "video.query", "text": "describe"})  # pins [a, a2]
+    await asyncio.sleep(0.1)
+    ws.put({"type": "video.frame", "data": frame_b})  # buffer full -> evicts pinned a
+    await asyncio.sleep(0.1)
+    release_others.set()
+
+    assert await _poll(lambda: "response.text.done" in ws.sent_types())
+    assert len(handler.query_contents) == 1
+    uuids = [p["uuid"] for p in handler.query_contents[0] if p.get("type") == "image_pil"]
+    expected = [hashlib.md5(base64.b64decode(f), usedforsecurity=False).hexdigest() for f in (frame_a, frame_a2)]
+    assert uuids == expected
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+async def _query_image_uuids(modalities: list[str], frames: list[str], *, max_frames: int) -> list[str]:
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put(
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": modalities,
+            "max_frames": max_frames,
+            "num_frames": max_frames,
+            "enable_frame_filter": False,
+        }
+    )
+    await asyncio.sleep(0)
+    for frame in frames:
+        ws.put({"type": "video.frame", "data": frame})
+    ws.put({"type": "video.query", "text": "describe"})
+    assert await _poll(lambda: handler.query_contents)
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+    return [part["uuid"] for part in handler.query_contents[0] if part.get("type") == "image_pil"]
+
+
+def _frame_uuid(frame_b64: str) -> str:
+    return hashlib.md5(base64.b64decode(frame_b64), usedforsecurity=False).hexdigest()
+
+
+@pytest.mark.asyncio
+async def test_below_max_frames_incremental_matches_legacy_window():
+    frames = [_b64(_make_jpeg(i, 0, 0)) for i in range(3)]
+    expected = [_frame_uuid(frame) for frame in frames]
+    assert await _query_image_uuids(["text"], frames, max_frames=4) == expected
+    assert await _query_image_uuids(["text", "audio"], frames, max_frames=4) == expected
+
+
+@pytest.mark.asyncio
+async def test_incremental_full_window_drops_oldest_half():
+    frames = [_b64(_make_jpeg(i, 1, 0)) for i in range(5)]
+    kept = await _query_image_uuids(["text"], frames, max_frames=4)
+    legacy = await _query_image_uuids(["text", "audio"], frames, max_frames=4)
+    assert kept == [_frame_uuid(frame) for frame in frames[2:]]
+    assert legacy == [_frame_uuid(frame) for frame in frames[1:]]
+
+
+async def _run_frame_audio_query(modalities: list[str], frame: str, pcm: bytes) -> PromptRecordingHandler:
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put(
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": modalities,
+            "enable_frame_filter": False,
+        }
+    )
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": frame})
+    if "audio" not in modalities:
+        assert await _poll(lambda: engine.warmup_ids())
+    else:
+        await asyncio.sleep(0.05)
+        assert engine.warmup_ids() == []
+    ws.put({"type": "audio.chunk", "data": _b64(pcm)})
+    ws.put({"type": "video.query", "text": "describe"})
+    assert await _poll(lambda: engine.query_ids())
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+    return handler
+
+
+@pytest.mark.asyncio
+async def test_reuse_query_with_audio_matches_legacy_prompt():
+    """Incremental and legacy query prompts match; both set the processor kwarg."""
+    frame = _b64(_make_jpeg())
+    pcm = b"\x01\x00\x02\x00"
+    reuse = await _run_frame_audio_query(["text"], frame, pcm)
+    legacy = await _run_frame_audio_query(["text", "audio"], frame, pcm)
+
+    assert reuse.warmup_contents
+    assert all(part.get("type") != "input_audio" for part in reuse.warmup_contents[0])
+    assert [part.get("type") for part in reuse.warmup_contents[0]] == ["image_pil"]
+
+    assert reuse.query_contents and legacy.query_contents
+    reuse_q = reuse.query_contents[0]
+    legacy_q = legacy.query_contents[0]
+    assert [part.get("type") for part in reuse_q] == ["image_pil", "input_audio", "text"]
+    assert [part.get("type") for part in legacy_q] == ["image_pil", "input_audio", "text"]
+    assert reuse_q[0]["uuid"] == legacy_q[0]["uuid"]
+    assert reuse_q[1] == legacy_q[1]
+    assert reuse_q[2] == legacy_q[2] == {"type": "text", "text": "describe"}
+
+    assert reuse.mm_processor_kwargs[0] == {"use_audio_in_video": True}
+    assert reuse.mm_processor_kwargs[1] == {"use_audio_in_video": True}
+    assert legacy.mm_processor_kwargs[0] == {"use_audio_in_video": True}

tests/entrypoints/openai_api/test_tts_adapter.py

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

tests/entrypoints/openai_api/test_tts_detection.py

@@ -183,6 +183,7 @@ _PIPELINE_STAGES = [
     "token2wav",
     "tts",
     "encoder",
+    "yue2",
 ]
 
 _STAGES = [*_PIPELINE_STAGES, None, "vae", "not_a_real_stage"]

tests/entrypoints/openai_api/test_video_audio_delta.py

@@ -174,7 +174,9 @@ async def test_real_producer_deltas_reach_video_client(
         produced.append(OmniRequestOutput.from_stage_output(result, final_output_type="audio"))
 
     class Engine:
-        async def generate(self, *, prompt, request_id, output_modalities, sampling_params_list=None):
+        async def generate(self, *, prompt, request_id, output_modalities, sampling_params_list=None, submitted=None):
+            if submitted is not None and not submitted.done():
+                submitted.set_result(None)
             if explicit_sampling:
                 assert sampling_params_list is not None
                 assert sampling_params_list[0].output_kind == RequestOutputKind.DELTA

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -11,6 +11,7 @@ from pydantic import ValidationError
 from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig, extract_diffusion_stage_config_kwargs
 from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config
+from vllm_omni.config.stage_config import build_stage_runtime_overrides
 from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
 from vllm_omni.engine import stage_init_utils
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
@@ -945,3 +946,33 @@ def test_generic_diffusion_structured_stage_reaches_standard_startup(mocker):
     assert launched["model"] == "generic-diffusion"
     assert launched["stage_id"] == 0
     assert runtime.stage_pools[0].clients == [client]
+
+
+def test_default_stage_config_preserves_pre_sharded_hsdp_weight_load() -> None:
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"hsdp_weight_load_strategy": "pre_sharded"})[0]
+
+    assert stage_cfg["engine_args"]["hsdp_weight_load_strategy"] == "pre_sharded"
+
+
+@pytest.mark.parametrize("prefix", ["", "stage_0_"])
+def test_pre_sharded_hsdp_weight_load_survives_stage_overrides(prefix: str) -> None:
+    overrides = build_stage_runtime_overrides(0, {f"{prefix}hsdp_weight_load_strategy": "pre_sharded"})
+
+    assert overrides["hsdp_weight_load_strategy"] == "pre_sharded"
+
+
+@pytest.mark.parametrize("strategy", ["full", "pre_sharded"])
+def test_typed_default_stage_preserves_hsdp_weight_load_strategy(strategy: str) -> None:
+    config = StageConfigFactory.create_typed_default_diffusion(
+        "generic-diffusion", {"hsdp_weight_load_strategy": strategy}
+    )
+
+    assert config.stage_configs[0].diffusion_config.hsdp_weight_load_strategy == strategy
+
+
+@pytest.mark.parametrize("strategy", ["memory_limited", "invalid"])
+def test_removed_hsdp_loading_strategies_are_rejected(strategy):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    with pytest.raises(ValueError, match="hsdp_weight_load_strategy"):
+        OmniDiffusionConfig(hsdp_weight_load_strategy=strategy)

tests/entrypoints/test_pd_disaggregation.py

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

tests/helpers/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/metrics/test_emit_calls.py

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

tests/metrics/test_stats.py

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

tests/model_executor/models/common/test_fused_sampling.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
+]
+
+
+def _reference(logits, uniforms, top_k, top_p):
+    if top_k > 0:
+        threshold = logits.topk(top_k, dim=-1).values[:, -1:]
+        logits = logits.masked_fill(logits < threshold, -torch.inf)
+    if top_p < 1:
+        values, indices = logits.sort(dim=-1, descending=True)
+        probs = values.softmax(dim=-1, dtype=torch.float32)
+        values[(probs.cumsum(-1) - probs) >= top_p] = -torch.inf
+        logits = values.scatter(1, indices, values)
+    return (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
+
+
+@pytest.mark.parametrize("batch,vocab", [(1, 128), (7, 2048), (64, 2048)])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("top_k,top_p", [(-1, 0.8), (0, 1.0), (50, 0.8), (50, 0.01)])
+def test_fixed_uniform_sampling_and_graph(batch, vocab, dtype, top_k, top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    torch.manual_seed(173)
+    logits = torch.randn(batch, vocab, device="cuda", dtype=dtype)
+    # The MTP caller supplies a noncontiguous row slice of [B, 15, vocab].
+    uniforms = torch.rand(batch, 15, vocab, device="cuda")[:, 3].clamp_min_(1e-7)
+    expected = _reference(logits, uniforms, top_k, top_p)
+    actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    logits.add_(torch.randn_like(logits) * 0.1)
+    uniforms.uniform_(1e-7, 1 - 1e-7)
+    graph.replay()
+    torch.testing.assert_close(captured, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("top_k,top_p", [(3, 0.5), (3, 1.0), (0, 0.8)])
+def test_threshold_ties_and_degenerate_rows(top_k, top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    logits = torch.tensor(
+        [[2, 2, 2, 2, 2, 1, 0, -1], [1] * 8, [-torch.inf] * 8, [torch.nan, 1, 2, 3, 4, 5, 6, 7]],
+        device="cuda",
+        dtype=torch.float32,
+    )
+    uniforms = torch.tensor([[0.4, 0.7, 0.8, 0.9, 0.2, 0.3, 0.6, 0.5]], device="cuda").expand(4, -1)
+    actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    torch.testing.assert_close(actual, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_native_nucleus_boundaries(dtype):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    torch.manual_seed(1923)
+    logits = torch.randn(1, 2048, device="cuda", dtype=dtype)
+    values, indices = logits.sort(descending=True)
+    probabilities = values.softmax(-1, dtype=torch.float32)
+    before = probabilities.cumsum(-1) - probabilities
+    for rank in [1, 19, 49, 1000, 2000]:
+        boundary = before[0, rank]
+        uniforms = torch.full_like(logits, 0.001, dtype=torch.float32)
+        uniforms[0, indices[0, rank]] = 1 - 1e-7
+        for direction in [-torch.inf, torch.inf]:
+            top_p = torch.nextafter(boundary, boundary.new_tensor(direction)).item()
+            actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=0, top_p=top_p)
+            torch.testing.assert_close(actual, _reference(logits, uniforms, 0, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("top_p", [0.8, 1.0])
+def test_strided_logits_and_non_power_of_two_vocabulary(top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    for seed in range(16):
+        torch.manual_seed(seed)
+        logits = torch.randn(128, 2050, device="cuda")[:, ::2]
+        uniforms = torch.rand(128, 3, 1025, device="cuda")[:, 1].clamp_min_(1e-7)
+        actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=-1, top_p=top_p)
+        torch.testing.assert_close(actual, _reference(logits, uniforms, -1, top_p), rtol=0, atol=0)

tests/model_executor/models/common/test_short_kv_attention.py

@@ -16,7 +16,7 @@ pytestmark = [
 @pytest.mark.parametrize("heads,kv_heads,dim", [(16, 8, 128), (8, 8, 64)])
 @pytest.mark.parametrize("positions", [[0, 1], [2], [8], [15]])
 def test_short_attention_causal_gqa_and_graph(batch, heads, kv_heads, dim, positions):
-    from vllm_omni.model_executor.models.qwen3_tts.short_kv_attention import short_kv_attention
+    from vllm_omni.model_executor.models.common.short_kv_attention import short_kv_attention
 
     with torch.inference_mode():
         torch.manual_seed(42)

tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_graph.py

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

tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py

@@ -13,6 +13,7 @@ from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import
     HiggsAudioV3TalkerForConditionalGeneration,
 )
 from vllm_omni.model_executor.models.higgs_audio_v3.model_state import HiggsModelState
+from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
 from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -34,6 +35,8 @@ def make_state():
         _audio_continuation_id=99,
         update_decode_step_metadata=Mock(),
     )
+    state._stream_pos = {}
+    state._eager_state = EagerMTPState(state)
     return state
 
 

tests/model_executor/models/moss_tts/test_codec_async_output.py

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

tests/model_executor/models/moss_tts/test_codec_ring_headroom.py

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

tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py

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

tests/model_executor/models/moss_tts/test_local_code_rows_snapshot.py

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

tests/model_executor/models/moss_tts/test_local_compiled_sampler.py

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

tests/model_executor/models/moss_tts/test_local_determined_tokens.py

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

tests/model_executor/models/moss_tts/test_local_eager_mtp.py

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

tests/model_executor/models/moss_tts/test_local_model_state.py

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

tests/model_executor/models/moss_tts/test_local_prefill_batch.py

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

tests/model_executor/models/moss_tts/test_local_short_attention.py

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

tests/model_executor/models/moss_tts/test_reference_encoder.py

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

tests/model_executor/models/moss_tts/test_reference_encoder_graphs.py

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

tests/model_executor/models/moss_tts/test_shared_reference_encoder.py

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

tests/model_executor/models/moss_tts/test_slot_attention.py

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

tests/model_executor/models/moss_tts/test_streaming_decoder_wrapper_concurrency.py

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

tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py

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

tests/model_executor/models/qwen3_omni/test_first_frame_cuda.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("rows", [16, 19])
+@torch.inference_mode()
+def test_first_frame_graph_outputs_own_storage(dtype, rows):
+    class Decoder(torch.nn.Linear):
+        config = SimpleNamespace(num_quantizers=1)
+
+        def forward(self, codes):
+            return codes.to(dtype) * self.weight
+
+    model = Decoder(1, 1, bias=False, device="cuda", dtype=dtype)
+    model.weight.fill_(1)
+    decoder = Qwen3OmniFirstFrameDecoder(model, sample_rate=24000)
+    decoder.capture()
+    codes = torch.arange(rows, device="cuda").reshape(-1, 1)
+    first = decoder.decode(codes)
+    second = decoder.decode(codes + 32)
+    assert first.dtype == second.dtype == torch.float32
+    torch.testing.assert_close(first, codes.float(), rtol=0, atol=0)
+    torch.testing.assert_close(second, (codes + 32).float(), rtol=0, atol=0)

tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_cudnn.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
+]
+
+
+def _codec(enabled):
+    config = Qwen3OmniMoeCode2WavConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        head_dim=8,
+        decoder_dim=32,
+        upsample_rates=[2, 2],
+        upsampling_ratios=[2],
+        codebook_size=16,
+        num_quantizers=2,
+    )
+    config._attn_implementation = "sdpa"
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=config, stage_connector_config={"extra": {"codec_cudnn_benchmark": enabled}}
+        )
+    )
+    model = Qwen3OmniMoeCode2Wav(vllm_config=vllm_config).cuda().eval()
+    model.precompute_snake_caches()
+    return model
+
+
+def _flags():
+    c = torch.backends.cudnn
+    return c.enabled, c.benchmark, c.benchmark_limit, c.deterministic, c.allow_tf32
+
+
+@torch.inference_mode()
+def test_autotuned_codec_graph_replay_restores_flags():
+    torch.manual_seed(21)
+    model = _codec(True)
+    codes = torch.randint(0, 16, (2, 2, 8), device="cuda")
+    original = _flags()
+    for _ in range(3):
+        model(codes)
+    assert _flags() == original
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = model(codes)
+    assert _flags() == original
+    codes.random_(0, 16)
+    graph.replay()
+    expected = model(codes)
+    torch.testing.assert_close(captured, expected)
+    assert _flags() == original
+
+
+@pytest.mark.parametrize("enabled", [False, True])
+@torch.inference_mode()
+def test_codec_failure_restores_cudnn_settings(mocker, enabled):
+    model = _codec(enabled)
+    codes = torch.randint(0, 16, (1, 2, 8), device="cuda")
+    original = _flags()
+
+    def fail(hidden):
+        if enabled:
+            assert torch.backends.cudnn.benchmark
+            assert torch.backends.cudnn.benchmark_limit == 10
+        else:
+            assert _flags() == original
+        raise RuntimeError("decode failed")
+
+    mocker.patch.object(model, "_decode_waveform", side_effect=fail)
+    with pytest.raises(RuntimeError, match="decode failed"):
+        model(codes)
+    assert _flags() == original

tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_snake.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import copy
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+    Qwen3OmniMoeCode2WavDecoderBlock,
+    Qwen3OmniMoeSnakeBeta,
+)
+
+from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav, use_fused_snake
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _blocks():
+    config = Qwen3OmniMoeCode2WavConfig(decoder_dim=32, upsample_rates=[2, 2])
+    torch.manual_seed(0)
+    blocks = torch.nn.ModuleList(Qwen3OmniMoeCode2WavDecoderBlock(config, i) for i in range(2)).eval()
+    with torch.no_grad():
+        for module in blocks.modules():
+            if isinstance(module, Qwen3OmniMoeSnakeBeta):
+                module.alpha.uniform_(-0.5, 0.5)
+                module.beta.uniform_(-0.5, 0.5)
+    return blocks
+
+
+def test_decoder_blocks_swap_every_hf_snake_and_keep_weight_names():
+    blocks = _blocks()
+    hf_count = sum(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules())
+    keys = set(blocks.state_dict())
+
+    assert use_fused_snake(blocks) == hf_count > 0
+    assert not any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules())
+    assert sum(isinstance(m, SnakeBeta) for m in blocks.modules()) == hf_count
+    assert set(blocks.state_dict()) == keys  # checkpoint weights load unchanged
+
+
+def test_fused_snake_blocks_match_hf_blocks():
+    reference = _blocks()
+    fused = copy.deepcopy(reference)
+    use_fused_snake(fused)
+    fused.load_state_dict(reference.state_dict())
+    x = torch.randn(2, 32, 20)
+    with torch.no_grad():
+        expected = x
+        actual = x
+        for ref_block, fused_block in zip(reference, fused):
+            expected = ref_block(expected)
+            actual = fused_block(actual)
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize(
+    ("option", "device", "cuda", "fused"),
+    [
+        (False, "cuda", True, False),
+        (True, "cuda", True, True),
+        (True, "cpu", True, False),
+        (True, "cuda", False, False),
+    ],
+)
+def test_decoder_block_fusion_requires_cuda_opt_in(monkeypatch, option, device, cuda, fused):
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: cuda)
+    config = Qwen3OmniMoeCode2WavConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        decoder_dim=32,
+    )
+    model = Qwen3OmniMoeCode2Wav(
+        vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(
+                hf_config=config, stage_connector_config={"extra": {"codec_fused_snake": option}}
+            ),
+            device_config=SimpleNamespace(device=device),
+        )
+    )
+    assert any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in model.decoder.modules()) is (not fused)

tests/model_executor/models/qwen3_omni/test_qwen3_omni_first_audio.py

@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import (
+    Qwen3OmniMoeCode2Wav,
+    plan_decode_groups,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_Q = 4
+_CODEBOOK = 16
+
+
+def _code2wav(quantizers=_Q) -> Qwen3OmniMoeCode2Wav:
+    torch.manual_seed(0)
+    config = Qwen3OmniMoeCode2WavConfig(
+        codebook_size=_CODEBOOK,
+        num_quantizers=quantizers,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        decoder_dim=32,
+        sliding_window=8,
+    )
+    model = Qwen3OmniMoeCode2Wav(vllm_config=SimpleNamespace(model_config=SimpleNamespace(hf_config=config)))
+    return model.eval()
+
+
+def test_first_frame_pcm_equals_code2wav_streaming_chunk0():
+    code2wav = _code2wav()
+    decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000)
+    codes = torch.randint(0, _CODEBOOK, (3, _Q))
+    with torch.inference_mode():
+        for row in range(codes.shape[0]):
+            pcm = decoder.decode(codes[row : row + 1])
+            chunk0 = code2wav.chunked_decode_streaming(
+                codes[row].reshape(1, _Q, 1), left_context_size=[0], seq_token_counts=[_Q]
+            )[0]
+            assert pcm.dtype == torch.float32
+            assert torch.equal(pcm[0], chunk0.reshape(-1).float())
+            assert 0 < pcm.shape[-1] < int(code2wav.total_upsample)
+        batched = decoder.decode(codes)
+    assert batched.shape == (3, pcm.shape[-1])
+
+
+def test_first_frame_is_a_prefix_of_a_longer_first_chunk():
+    code2wav = _code2wav()
+    decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000)
+    codes = torch.randint(0, _CODEBOOK, (_Q, 4))
+    with torch.inference_mode():
+        first = decoder.decode(codes[:, :1].T)[0]
+        chunk0 = code2wav.chunked_decode_streaming(
+            codes.reshape(1, _Q, 4), left_context_size=[0], seq_token_counts=[4 * _Q]
+        )[0]
+    torch.testing.assert_close(chunk0.reshape(-1)[: first.numel()].float(), first)
+
+
+def test_chunk_ramp_adds_exact_code2wav_graph_sizes(monkeypatch):
+    import vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper as wrapper_module
+
+    real = wrapper_module.CUDAGraphDecoderWrapper
+    wrapper = MagicMock(compute_capture_sizes=real.compute_capture_sizes)
+    monkeypatch.setattr(wrapper_module, "CUDAGraphDecoderWrapper", wrapper)
+    extra = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "codec_chunk_ramp": [1, 2, 4, 8, 16, 25]}
+    code2wav = SimpleNamespace(config=SimpleNamespace(num_quantizers=_Q))
+    code2wav.enable_cudagraph = lambda **kw: Qwen3OmniMoeCode2Wav.enable_cudagraph(
+        code2wav, device=torch.device("cuda"), **kw
+    )
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.code2wav = code2wav
+    model.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(async_chunk=True, enforce_eager=False, stage_connector_config={"extra": extra})
+    )
+    for batches in ([], [2, 4, 8]):
+        extra["codec_graph_batch_sizes"] = batches
+        Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model)
+        args = wrapper.call_args.kwargs
+        sizes = args["capture_sizes"]
+        assert {1, 3, 7, 15, 31, 50} <= set(sizes)
+        assert set(real.compute_capture_sizes(codec_chunk_frames=25, codec_left_context_frames=25)) <= set(sizes)
+        assert set(args["extra_capture_shapes"]) == {
+            (b, size) for b in batches for size in sizes if size <= 50 and b * size <= 325
+        }
+    assert code2wav._streaming_batch_sizes[50] == [1, 2, 4]
+    extra.pop("codec_chunk_ramp")
+    extra.pop("codec_graph_batch_sizes")
+    Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model)
+    assert wrapper.call_args.kwargs["capture_sizes"] is None
+
+
+_SIZES = [1, 2, 3, 4, 7, 8, 15, 16, 25, 31, 32, 50, 64]
+
+
+def _bucket(length):
+    return next((size for size in _SIZES if size >= length), None)
+
+
+def test_decode_groups_keep_a_uniform_batch_whole_and_split_by_the_largest_graph():
+    assert plan_decode_groups([50] * 6, _bucket, lambda size: [1, 2, 4, 8]) == [(list(range(6)), 50)]
+    calls = plan_decode_groups([25] * 20, _bucket, lambda size: [1, 2, 4, 8, 16])
+    assert [len(rows) for rows, _length in calls] == [16, 4]
+    calls = plan_decode_groups([1] * 8 + [50] * 8, _bucket, lambda size: [1, 2, 4, 8] if size < 50 else [1, 2, 4])
+    assert [(len(rows), length) for rows, length in calls] == [(8, 1), (4, 50), (4, 50)]
+    assert plan_decode_groups([], _bucket, lambda size: [1, 2]) == []
+
+
+@pytest.mark.parametrize("frames", [65, 129])
+def test_decode_groups_beyond_largest_graph_use_eager(monkeypatch, frames):
+    from vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper
+
+    wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper)
+    wrapper.enabled = wrapper._warmed_up = True
+    wrapper._bucket_sizes = [8, 32, 64]
+    wrapper._compiled_shapes = {(1, 64), (8, 64)}
+    wrapper._compiled_graphs = wrapper.graphs = dict.fromkeys(wrapper._compiled_shapes, MagicMock())
+    wrapper.decoder = MagicMock(side_effect=lambda codes: codes.float() * 2)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    groups = plan_decode_groups([frames, frames], wrapper._get_padded_size, lambda _: (1,))
+    assert groups == [([0], frames), ([1], frames)]
+    codes = torch.arange(2 * frames).reshape(2, 1, frames)
+    for rows, width in groups:
+        actual = wrapper.decode(codes[rows, :, :width])
+        torch.testing.assert_close(actual, codes[rows].float() * 2, rtol=0, atol=0)
+    assert wrapper.decoder.call_count == 2
+    for graph in (*wrapper.graphs.values(), *wrapper._compiled_graphs.values()):
+        graph.replay.assert_not_called()
+
+
+class _PaddingGraphs:
+    def __init__(self, model, batch_sizes):
+        self.model = model
+        self.batch_sizes = sorted(batch_sizes)
+        self.calls = []
+
+    def _get_padded_size(self, length):
+        return _bucket(length)
+
+    def decode(self, codes):
+        rows, _q, frames = codes.shape
+        size = _bucket(frames)
+        batch = next(b for b in self.batch_sizes if b >= rows)
+        self.calls.append((rows, frames))
+        padded = torch.zeros(batch, codes.shape[1], size, dtype=codes.dtype)
+        padded[:rows, :, :frames] = codes
+        out = self.model(padded)
+        return out[:rows, :, : out.shape[-1] - (size - frames) * int(self.model.total_upsample)].clone()
+
+
+def test_grouped_streaming_decode_matches_one_padded_batch():
+    code2wav = _code2wav()
+    lengths = [1, 3, 50, 7, 31]
+    lefts = [0, 1, 25, 3, 15]
+    codes = torch.zeros(len(lengths), _Q, max(lengths), dtype=torch.long)
+    for row, length in enumerate(lengths):
+        codes[row, :, :length] = torch.randint(0, _CODEBOOK, (_Q, length))
+    counts = [length * _Q for length in lengths]
+    with torch.inference_mode():
+        reference = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts)
+        graphs = _PaddingGraphs(code2wav, [1, 2, 4])
+        code2wav._cudagraph_wrapper, code2wav._cudagraph_enabled = graphs, True
+        code2wav._streaming_batch_sizes = {size: [1, 2, 4] for size in _SIZES}
+        grouped = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts)
+
+    for want, got in zip(reference, grouped, strict=True):
+        assert want.shape == got.shape
+        torch.testing.assert_close(got, want, rtol=1e-4, atol=1e-5)
+
+
+@pytest.mark.parametrize(("frames", "quantizers"), [(1, 16), (4, 16), (4, 4)])
+def test_delivered_first_frame_is_removed_from_code2wav_output(frames, quantizers):
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY
+
+    codec = _code2wav(quantizers)
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.model_stage, model.code2wav, model.code2wav_config = "code2wav", codec, codec.config
+    model.generate_audio = lambda codes, left, counts: codec.chunked_decode_streaming(
+        codes, left_context_size=left, seq_token_counts=counts
+    )
+    codes = torch.randint(0, _CODEBOOK, (2, quantizers, frames))
+    with torch.inference_mode():
+        reference = model.generate_audio(codes, [0, 0], [quantizers * frames] * 2)
+        first = Qwen3OmniFirstFrameDecoder(codec, sample_rate=24000).decode(codes[0, :, :1].T)[0]
+        out = model.forward(
+            input_ids=codes.flatten(),
+            positions=None,
+            seq_token_counts=[quantizers * frames] * 2,
+            runtime_additional_information=[
+                {"meta": {"left_context_size": 0, "first_audio": True}},
+                {"meta": {"left_context_size": 0, "first_audio": False}},
+            ],
+        ).multimodal_outputs
+    torch.testing.assert_close(torch.cat([first, out["model_outputs"][0].flatten()]), reference[0].flatten())
+    torch.testing.assert_close(out["model_outputs"][1], reference[1].reshape(1, -1))
+    assert [bool(flag) for flag in out[FIRST_AUDIO_REQUIRED_KEY]] == [True, False]

tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrv2.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_H = 4
+_Q = 3
+
+
+class _Talker(nn.Module):
+    num_code_groups = _Q
+
+    def __init__(self):
+        super().__init__()
+        self.text_projection = nn.Identity()
+        self.code_predictor = SimpleNamespace(
+            _top_k=50,
+            _top_p=0.8,
+            set_sampling_params=lambda **kwargs: setattr(self, "sampling", kwargs),
+        )
+
+
+def _talker(*, async_chunk: bool = True) -> Qwen3OmniMoeForConditionalGeneration:
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    nn.Module.__init__(model)
+    model.model_stage = "talker"
+    model.talker = _Talker()
+    model.vllm_config = SimpleNamespace(model_config=SimpleNamespace(async_chunk=async_chunk, use_v2_model_runner=True))
+    model.tts_eos_embed = torch.full((_H,), -1.0)
+    model.tts_pad_embed = torch.full((_H,), -2.0)
+    model._codec_codebook_size = 2048
+    return model
+
+
+def _rows(*values: float) -> torch.Tensor:
+    return torch.tensor([[value] * _H for value in values])
+
+
+def _step(model, payload: dict) -> torch.Tensor:
+    update: dict = {}
+    text_step = model._thinker_decode_to_talker_decode(payload, torch.device("cpu"), update)
+    for key, value in update.items():
+        payload.setdefault(key, {}).update(value)
+    return text_step.reshape(-1)
+
+
+def test_mtp_frame_valid_rejects_codec_eos_and_special_ids():
+    layer0 = torch.tensor([0, 2047, 2048, 2150, -1])
+    assert _talker().mtp_frame_valid(layer0).tolist() == [True, True, False, False, False]
+
+
+def test_decode_rows_are_consumed_one_per_step_in_arrival_order():
+    model = _talker()
+    # Two rows were already pending (cached at prefill); two more arrive at once.
+    payload = {"embed": {"cached_decode": _rows(1, 2), "decode": _rows(3, 4)}, "meta": {}}
+    steps = [_step(model, payload)[0].item()]
+    for arrival in (None, _rows(5), None, None, None, None):
+        payload["embed"]["decode"] = arrival
+        steps.append(_step(model, payload)[0].item())
+    assert steps == [1, 2, 3, 4, 5, -1, -2]
+    assert payload["embed"]["cached_decode"] is None
+
+
+@pytest.mark.parametrize(("v2", "resumable"), [(True, True), (False, False), (False, True)])
+def test_v1_and_resumable_requests_keep_the_indexed_cache_path(v2, resumable):
+    model = _talker()
+    model.vllm_config.model_config.use_v2_model_runner = v2
+    payload = {"embed": {"cached_decode": _rows(1, 2, 3), "decode": None}, "meta": {"resumable": resumable}}
+    payload["meta"]["num_processed_tokens"] = 2
+    assert _step(model, payload)[0].item() == 3
+
+
+def test_make_omni_output_aligns_codes_to_request_spans():
+    model = _talker()
+    hidden = torch.zeros((8, _H))  # 5 real tokens + graph padding
+    decode_codes = torch.tensor([[5, 6, 7]])
+    buffers = [
+        {"req_id": "_warmup_0_"},  # no preprocess, no codes
+        {"req_id": "a", "codes": {"audio": torch.zeros((3, _Q), dtype=torch.long)}},
+        {"req_id": "b", "codes": {"audio": decode_codes}},
+    ]
+    out = model.make_omni_output(
+        hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 1), (1, 4), (4, 5)]
+    )
+    codes = out.multimodal_outputs["codes"]["audio"]
+    assert codes.shape == (5, _Q)
+    assert codes[4].tolist() == [5, 6, 7]
+    assert out.text_hidden_states.shape[0] == 5
+
+
+def test_make_omni_output_without_spans_keeps_v1_concat():
+    model = _talker()
+    buffers = [{"codes": {"audio": torch.tensor([[1, 2, 3]])}}, {"codes": {"audio": torch.tensor([[4, 5, 6]])}}]
+    out = model.make_omni_output(torch.zeros((4, _H)), model_intermediate_buffer=buffers)
+    assert out.multimodal_outputs["codes"]["audio"].tolist() == [[1, 2, 3], [4, 5, 6]]
+    assert out.text_hidden_states.shape[0] == 2
+
+
+def test_eager_output_is_token_major_with_invalid_rows():
+    model = _talker()
+    model.eager_frames_active = True
+    hidden = torch.zeros((6, _H))
+    buffers = [{"req_id": "a"}, {"req_id": "b"}]
+    out = model.make_omni_output(hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 3), (3, 4)])
+    mm = out.multimodal_outputs
+    assert mm["codes"]["audio"].shape == (4, _Q)
+    assert mm["meta"]["codec_frame_valid"].tolist() == [0, 0, 0, 0]
+    assert out.text_hidden_states is hidden

tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py

@@ -104,6 +104,7 @@ def test_factory_declares_only_captures_from_earlier_ranks(monkeypatch, start, e
     if staged:
         model_config.model_stage = "thinker"
     vllm_config = SimpleNamespace(model_config=model_config, with_hf_config=lambda *args, **kwargs: None)
+    vllm_config.speculative_config = None
     thinker = nn.Module()
     thinker.weight = nn.Parameter(torch.ones(1))
     thinker.language_model = SimpleNamespace(model=SimpleNamespace(start_layer=start))

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -569,6 +569,27 @@ class TestCodePredictorPerRowGenerators:
             )
 
 
+@pytest.mark.parametrize("fused_requested", [False, True])
+@pytest.mark.parametrize("invalid_input", ["uniforms", "generators"])
+def test_sampling_contract_rejected_before_dispatch(mocker, loaded_target_classes, fused_requested, invalid_input):
+    predictor, talker_config = TestCodePredictorPerRowGenerators()._make_predictor(mocker, loaded_target_classes)
+    predictor._fused_requested = fused_requested
+    predictor._fused = mocker.Mock()
+    setup = mocker.patch.object(predictor, "_setup_compile")
+    uniforms = torch.full((2, 3, 1 if invalid_input == "uniforms" else 64), 0.5)
+    generators = [torch.Generator()] if invalid_input == "generators" else None
+    with pytest.raises(ValueError, match="sample_uniforms must have shape|one entry per row"):
+        predictor(
+            torch.zeros(2, dtype=torch.long),
+            torch.randn(2, talker_config.hidden_size),
+            torch.randn(2, talker_config.hidden_size),
+            sample_uniforms=uniforms,
+            generators=generators,
+        )
+    setup.assert_not_called()
+    predictor._fused.assert_not_called()
+
+
 class TestCodePredictorModelDtype:
     """Test the inner model forward with different dtypes."""
 
@@ -663,6 +684,8 @@ class TestCodePredictorGraphReplay:
 
         predictor = object.__new__(code_predictor_wrapper)
         torch.nn.Module.__init__(predictor)
+        predictor._fused_requested = False
+        predictor._fused = None
         predictor._num_groups = 3
         predictor._model_dtype = torch.float32
         predictor._setup_compile = mocker.Mock()
@@ -1432,3 +1455,24 @@ def test_shared_predictor_ignores_tts_fast_path_options(mocker, loaded_target_cl
     expected = (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
     actual = predictor._sample_per_call(logits, 1.0, 0, None, uniforms)
     torch.testing.assert_close(actual, expected)
+
+
+def test_tts_cached_predictor_keeps_its_own_step_dispatch(mocker: MockerFixture, loaded_target_classes) -> None:
+    """The TTS cache must not enter Omni's uninitialized incremental loop."""
+    _, _, wrapper, _, _ = loaded_target_classes
+    common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+    mocker.patch.object(common_mod.current_omni_platform, "is_cuda", return_value=True)
+    mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+    mocker.patch.object(wrapper, "_stage_connector_extra_config", return_value={"code_predictor_kv_cache": True})
+    cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+    predictor = wrapper(vllm_config=_make_vllm_config(mocker), config=cp_config, talker_config=talker_config)
+    assert predictor._kv_requested
+    predictor._model_dtype = torch.float32
+    predictor._codec_embeds_list = list(predictor.model.codec_embedding)
+    predictor._bucket_sizes = [1]
+    mocker.patch.object(predictor, "_setup_compile")
+    step = mocker.patch.object(predictor, "_predict_step_logits", return_value=torch.ones(1, cp_config.vocab_size))
+    codes = predictor(torch.zeros(1, dtype=torch.long), torch.zeros(1, 32), torch.zeros(1, 32), do_sample=False)
+    assert step.call_count == cp_config.num_code_groups - 1
+    assert codes.shape == (1, cp_config.num_code_groups)
+    assert predictor._kv_buf is None

tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Supplied-uniform samples must match the existing ATen path exactly."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -12,6 +14,79 @@ pytestmark = [
 ]
 
 
+@pytest.mark.parametrize("top_k", [0, 1, 50, 2048])
+@torch.inference_mode()
+def test_fused_sample_kernel_preserves_signed_zero_cutoff_ties(top_k):
+    from vllm_omni.model_executor.models.qwen3_tts.fused_code_predictor import _cp_sample_kernel
+
+    vocab, hidden, groups = 2048, 64, 4
+    logits = torch.zeros(2, vocab, device="cuda", dtype=torch.bfloat16)
+    logits[0, 1] = -0.0
+    logits[1, ::2] = -0.0
+    uniforms = torch.full((2, groups - 1, vocab), 0.01, device="cuda")[:, 1, :]
+    uniforms[0, 1] = uniforms[1, 2] = 0.99
+    codes = torch.full((2, groups), -1, device="cuda", dtype=torch.int64)
+    table = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16)
+    next_input = torch.empty(2, hidden, device="cuda", dtype=table.dtype)
+
+    def sample():
+        _cp_sample_kernel[(2,)](
+            logits, uniforms, uniforms.stride(0), codes, 1, groups, table, next_input, 1.0,
+            V=vocab, TOPK=top_k, HID=hidden, HAS_NEXT=True, num_warps=4,
+        )  # fmt: skip
+
+    scaled = logits.float()
+    if top_k > 0:
+        scaled = scaled.masked_fill(scaled < scaled.topk(top_k).values[:, -1:], -torch.inf)
+    expected = (scaled - torch.log(-torch.log(uniforms))).argmax(-1)
+    assert expected.tolist() == [1, 2]
+    sample()
+    torch.testing.assert_close(codes[:, 1], expected)
+    torch.testing.assert_close(next_input, table[expected])
+    assert (codes[:, [0, 2, 3]] == -1).all()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        sample()
+    codes.fill_(-1)
+    graph.replay()
+    torch.testing.assert_close(codes[:, 1], expected)
+    torch.testing.assert_close(next_input, table[expected])
+
+
+@torch.inference_mode()
+def test_fused_predictor_warmup_with_capacity_one():
+    from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+        Qwen3TTSTalkerCodePredictorConfig,
+        Qwen3TTSTalkerConfig,
+    )
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM,
+    )
+
+    cp = Qwen3TTSTalkerCodePredictorConfig(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1,
+        num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4,
+    )  # fmt: skip
+    talker = Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4)
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(stage_connector_config={}),
+        additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True},
+        scheduler_config=SimpleNamespace(max_num_seqs=1),
+    )
+    predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(
+        vllm_config=config, config=cp, talker_config=talker
+    ).to(device="cuda", dtype=torch.bfloat16)
+    predictor._setup_compile()
+    assert predictor._fused is not None
+    hidden = torch.randn(1, 1, 64, device="cuda", dtype=torch.bfloat16)
+    result = predictor(
+        torch.zeros(1, 1, device="cuda", dtype=torch.long), hidden, hidden,
+        sample_uniforms=torch.full((1, 3, 64), 0.5, device="cuda"),
+    )  # fmt: skip
+    assert result.shape == (1, 4)
+    assert ((result >= 0) & (result < 64)).all()
+
+
 @pytest.mark.parametrize("batch", [1, 64])
 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
 @pytest.mark.parametrize("top_k", [-1, 1, 50, 2048])
@@ -66,3 +141,47 @@ def test_topk_gumbel_accepts_strided_supplied_uniforms(top_k):
     expected = (scaled.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
     actual = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9)
     torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.parametrize("batch", [1, 4])
+@pytest.mark.parametrize("top_k", [1, 50])
+@torch.inference_mode()
+def test_fused_predictor_matches_full_reference_with_fixed_uniforms(mocker, batch, top_k):
+    from vllm.config import VllmConfig
+
+    from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+        Qwen3TTSTalkerCodePredictorConfig,
+        Qwen3TTSTalkerConfig,
+    )
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM,
+    )
+
+    torch.manual_seed(42)
+    cp = Qwen3TTSTalkerCodePredictorConfig(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1,
+        num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4,
+    )  # fmt: skip
+    config = mocker.Mock(
+        spec=VllmConfig,
+        model_config=mocker.Mock(stage_connector_config={}),
+        additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True},
+        scheduler_config=mocker.Mock(max_num_seqs=4),
+    )
+    predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(
+        vllm_config=config, config=cp, talker_config=Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4)
+    ).to(device="cuda", dtype=torch.bfloat16)
+    predictor._setup_compile()
+    assert predictor._fused is not None
+    ids = torch.randint(0, 64, (batch, 1), device="cuda")
+    embeds = torch.randn(batch, 1, 64, device="cuda", dtype=torch.bfloat16)
+    hidden = torch.randn_like(embeds)
+    uniforms = torch.rand(batch, 3, 64, device="cuda").clamp_(1e-6, 1 - 1e-6)
+    actual = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms).clone()
+    # Use identical parameters through the full re-prefill and ATen sampling
+    # path, bypassing both the fused predictor and its frame-local KV cache.
+    predictor._fused_requested = False
+    predictor._fused = None
+    predictor._frame_cache = None
+    expected = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)

tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Prefix graphs must follow per-request first-audio delivery, not a capture hint."""
 
+import math
 from types import SimpleNamespace
 
 import pytest
@@ -18,6 +19,38 @@ pytestmark = [pytest.mark.core_model]
 DEVICE = torch.device("cuda:0")
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("capacity,model_type", [(1, "base"), (128, "base"), (128, "custom_voice")])
+def test_stream_priming_graph_workspace_respects_token_budget(monkeypatch, capacity, model_type):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import streaming_decoder
+
+    captures = []
+
+    def capture(stream, sizes, frames=1):
+        captures.append((sizes, frames))
+        return SimpleNamespace(sizes=sizes, frames=frames)
+
+    monkeypatch.setattr(streaming_decoder, "StreamingDecodeGraphs", capture)
+    model = SimpleNamespace(
+        stream_decoder=SimpleNamespace(max_frames=33),
+        stream_graphs=None,
+        stream_prime_graphs=None,
+        stream_ref_context_frames=25,
+        stream_chunk_frames=25,
+        config=SimpleNamespace(tts_model_type=model_type),
+        vllm_config=SimpleNamespace(scheduler_config=SimpleNamespace(max_num_batched_tokens=512)),
+    )
+    sizes = [size for size in [1, 2, 4, 8, 16, 32, 64, 128] if size <= capacity]
+    Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes)
+    expected = [(sizes, 1)]
+    if model_type == "base":
+        expected.append(([size for size in sizes if size * 25 <= 512], 25))
+    assert captures == expected
+    Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes)
+    assert captures == expected
+
+
 @pytest.mark.cpu
 @pytest.mark.parametrize("skip_flags", [(False, False), (True, True), (False, True), (True, False)])
 @pytest.mark.parametrize("initial_frames", [1, 3])
@@ -124,3 +157,134 @@ def test_xvec_graphs_follow_request_first_audio_flags(initial_frames, time_major
             torch.testing.assert_close(got, want, atol=1e-4, rtol=1e-4)
             for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"):
                 torch.testing.assert_close(cache[key], expected_cache[key], atol=1e-5, rtol=1e-5)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.inference_mode()
+def test_stream_decoder_preemption_restores_all_state_into_a_different_slot():
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import StreamingCodecDecoder
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    stream = StreamingCodecDecoder(decoder, num_slots=2, dtype=torch.bfloat16)
+    codes = torch.randint(0, 32, (1, 8, 2), device=DEVICE)
+    slot = torch.tensor([0], device=DEVICE, dtype=torch.int32)
+    for pos in range(7):
+        stream(codes[:, pos : pos + 1], slot, torch.tensor([pos], device=DEVICE, dtype=torch.int32))
+    saved = stream.save_slot(0)
+    expected = stream(codes[:, 7:], slot, torch.tensor([7], device=DEVICE, dtype=torch.int32)).clone()
+    # Both the freed slot and the destination can hold an unrelated request.
+    other_codes = (codes[:, :3] + 1) % 32
+    for pos in range(3):
+        stream(
+            other_codes[:, pos : pos + 1].expand(2, -1, -1).contiguous(),
+            torch.tensor([0, 1], device=DEVICE, dtype=torch.int32),
+            torch.full((2,), pos, device=DEVICE, dtype=torch.int32),
+        )
+    without_restore = stream(
+        codes[:, 7:],
+        torch.tensor([1], device=DEVICE, dtype=torch.int32),
+        torch.tensor([7], device=DEVICE, dtype=torch.int32),
+    ).clone()
+    assert not torch.equal(without_restore, expected)
+    stream.restore_slot(1, saved)
+    actual = stream(
+        codes[:, 7:],
+        torch.tensor([1], device=DEVICE, dtype=torch.int32),
+        torch.tensor([7], device=DEVICE, dtype=torch.int32),
+    )
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("use_graph", [False, True])
+@torch.inference_mode()
+def test_batched_reference_priming_matches_serial_decoder_continuation(use_graph):
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import (
+        StreamingCodecDecoder,
+        StreamingDecodeGraphs,
+    )
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    stream = StreamingCodecDecoder(decoder, num_slots=5, dtype=torch.bfloat16)
+    graphs = StreamingDecodeGraphs(stream, [2], frames=25) if use_graph else None
+    model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.stream_decoder = stream
+    model.stream_prime_graphs = graphs
+    model.stream_chunk_frames = 25
+    refs = [torch.randint(0, 32, (length, 2), device=DEVICE) for length in [25, 3, 51, 25]]
+    slots = [2, 0, 3, 4]
+    frame = torch.randint(0, 32, (4, 1, 2), device=DEVICE)
+    slot_tensor = torch.tensor(slots, device=DEVICE, dtype=torch.int32)
+    positions = torch.tensor([len(ref) for ref in refs], device=DEVICE, dtype=torch.int32)
+    for slot, ref in zip(slots, refs, strict=True):
+        for t0 in range(0, len(ref), 25):
+            stream(
+                ref[None, t0 : t0 + 25].to(torch.int32).contiguous(),
+                torch.tensor([slot], device=DEVICE, dtype=torch.int32),
+                torch.tensor([t0], device=DEVICE, dtype=torch.int32),
+            )
+    expected = stream(frame, slot_tensor, positions).clone()
+    untouched = stream.save_slot(1)
+    # Reuse dirty slots; pos=0 must reset every causal layer for each group.
+    model.prime_stream_decoder(list(zip(slots, refs, strict=True)))
+    actual = stream(frame, slot_tensor, positions)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=3e-5)
+    for got, want in zip(stream.save_slot(1), untouched, strict=True):
+        torch.testing.assert_close(got, want, rtol=0, atol=0)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("use_graph", [False, True])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@torch.inference_mode()
+def test_stream_frames_match_fp32_exact_decoder(use_graph, dtype):
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import (
+        StreamingCodecDecoder,
+        StreamingDecodeGraphs,
+    )
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=dtype)
+    # The tiny HF default initialization shrinks the deep conv stack to
+    # near-zero PCM. Keep the oracle signal large enough to catch lost state.
+    for module in decoder.modules():
+        if isinstance(module, (torch.nn.Conv1d, torch.nn.ConvTranspose1d, torch.nn.Linear)):
+            torch.nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))
+            if module.bias is not None:
+                module.bias.uniform_(-0.01, 0.01)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    reference = _make_decoder().to(device=DEVICE, dtype=torch.float32)
+    reference.load_state_dict(decoder.state_dict())
+    reference.precompute_snake_caches()
+    # Cross both the attention window and KV ring wrap, checking every frame.
+    codes = torch.randint(0, 32, (2, 136, 2), device=DEVICE)
+    expected = reference._forward_exact(codes.transpose(1, 2))[:, 0]
+    assert expected.square().sum() > 1e-3
+    stream = StreamingCodecDecoder(decoder, num_slots=3, dtype=dtype)
+    decode = StreamingDecodeGraphs(stream, [2]) if use_graph else stream
+    slots = torch.tensor([2, 0], device=DEVICE, dtype=torch.int32)
+    spf = int(stream.spf)
+    for frame in range(codes.shape[1]):
+        positions = torch.full((2,), frame, device=DEVICE, dtype=torch.int32)
+        actual = decode(codes[:, frame : frame + 1].contiguous().to(torch.int32), slots, positions).reshape(2, -1)
+        exact_frame = expected[:, frame * spf : (frame + 1) * spf]
+        relative_rms = (actual - exact_frame).square().mean(1).sqrt() / exact_frame.square().mean(1).sqrt()
+        # BF16 kernels are approximate; compare every frame to the independent
+        # FP32 full decoder, not only one streaming implementation to itself.
+        tolerance = 0.05 if dtype == torch.bfloat16 else 5e-4
+        assert torch.all(relative_rms < tolerance), (frame, relative_rms)

tests/model_executor/models/qwen3_tts/test_streaming_col2im_cuda.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Streaming col2im preserves PCM inputs and ping-pong slot state."""
+
+import pytest
+import torch
+from vllm.triton_utils import triton
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import _col2im_kernel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.tts]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("batch,frames,channels,rate", [(1, 1, 5, 2), (3, 3, 129, 3), (9, 1, 64, 4)])
+def test_col2im_partial_tiles_and_reordered_slot_state(batch: int, frames: int, channels: int, rate: int):
+    torch.manual_seed(23)
+    capacity = batch + 2
+    z = torch.randn(batch, frames, 2 * rate, channels).to(torch.bfloat16)
+    bias = torch.randn(channels).to(torch.bfloat16)
+    state = torch.randn(2, capacity, rate, channels).to(torch.bfloat16)
+    expected_state = state.clone()
+    slots = torch.arange(batch - 1, -1, -1, dtype=torch.int32)
+    positions = torch.arange(batch, dtype=torch.int32)
+    parity = torch.arange(capacity, dtype=torch.int32) % 2
+    expected = torch.empty(batch, frames, rate, channels, dtype=torch.bfloat16)
+    for b in range(batch):
+        slot = int(slots[b])
+        side = int(parity[slot])
+        for t in range(frames):
+            previous = z[b, t - 1, rate:] if t else state[side, slot] if positions[b] else torch.zeros(rate, channels)
+            expected[b, t] = (z[b, t, :rate].float() + previous.float() + bias.float()).to(torch.bfloat16)
+        expected_state[1 - side, slot] = z[b, -1, rate:]
+
+    z_gpu, bias_gpu, state_gpu = z.cuda(), bias.cuda(), state.cuda()
+    slots_gpu, positions_gpu, parity_gpu = slots.cuda(), positions.cuda(), parity.cuda()
+    actual = torch.empty_like(expected, device="cuda")
+    width = min(128, triton.next_power_of_2(channels))
+    rows = batch * frames
+    _col2im_kernel[(rows, rate, triton.cdiv(channels, width))](
+        z_gpu,
+        bias_gpu,
+        state_gpu,
+        slots_gpu,
+        positions_gpu,
+        parity_gpu,
+        actual,
+        frames,
+        capacity,
+        R=rate,
+        C=channels,
+        BC=width,
+        num_warps=4,
+    )
+    torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0)
+    torch.testing.assert_close(state_gpu.cpu(), expected_state, rtol=0, atol=0)

tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py

@@ -4,6 +4,7 @@
 
 import pytest
 import torch
+from vllm.platforms import current_platform
 
 from tests.helpers.mark import hardware_test
 from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import sample_logits
@@ -16,6 +17,10 @@ VOCAB = 2052
 EOS = 2051
 
 
+def _captured_generator_state_supported() -> bool:
+    return current_platform.is_cuda() and hasattr(torch.cuda.CUDAGraph, "register_generator_state")
+
+
 def _logits(count: int, *, device: str, seed: int = 100, vocab: int = VOCAB) -> list[torch.Tensor]:
     source = torch.Generator(device="cpu").manual_seed(seed)
     values = torch.randn((count, vocab), generator=source, dtype=torch.float32) * 3
@@ -164,6 +169,14 @@ def test_cuda_replay_preserves_tokens_rng_and_owned_outputs(monkeypatch) -> None
                     right[codebook].exponential_(generator=live_expected[index])
                 torch.testing.assert_close(left, right, atol=0, rtol=0)
             _assert_states(actual_generators, expected_generators)
+    if not _captured_generator_state_supported():
+        # ROCm deliberately uses the eager sampler because CUDA graph generator-state registration is an
+        # NVIDIA-only path. The loop above still verifies GPU tokens and every request-owned RNG stream.
+        assert not sampler._graphs
+        for value, expected in retained:
+            torch.testing.assert_close(value, expected, atol=0, rtol=0)
+        torch.testing.assert_close(torch.cuda.get_rng_state(), global_state, atol=0, rtol=0)
+        return
     # B1/3/16/17/31/32 must use exactly buckets 1/4/16/32.
     assert len(sampler._graphs) == 4
     rows = _logits(3, device="cuda", seed=2026)
@@ -189,6 +202,7 @@ def test_cuda_replay_preserves_tokens_rng_and_owned_outputs(monkeypatch) -> None
 @torch.inference_mode()
 def test_cuda_parameters_vocab_and_missing_graph_api(monkeypatch) -> None:
     sampler = BreezeFirstCodeSampler()
+    captures_generator_state = _captured_generator_state_supported()
     actual_generators = _generators([42, 43, 44], "cuda")
     expected_generators = _generators([42, 43, 44], "cuda")
     cases = ((0.9, 50, 1.0), (0.7, 100, 0.8), (1.2, 0, 0.95), (1.0, 1, 1.0), (1.0, 0, 0.5))
@@ -201,12 +215,12 @@ def test_cuda_parameters_vocab_and_missing_graph_api(monkeypatch) -> None:
         expected = _reference(rows, parameters, expected_generators)
         _assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows)
         _assert_states(actual_generators, expected_generators)
-    assert len(sampler._graphs) == len(cases)
+    assert len(sampler._graphs) == (len(cases) if captures_generator_state else 0)
     rows = _logits(3, device="cuda", vocab=17)
     expected = _reference(rows, PARAMETERS, expected_generators)
     _assert_outputs(sampler.sample(rows, PARAMETERS, actual_generators), expected, rows)
     _assert_states(actual_generators, expected_generators)
-    assert len(sampler._graphs) == len(cases) + 1
+    assert len(sampler._graphs) == (len(cases) + 1 if captures_generator_state else 0)
 
     class UnavailableGraph:
         def __init__(self, *_args, **_kwargs):
@@ -243,7 +257,7 @@ def test_cuda_shared_generator_uses_sequential_eager_fallback(monkeypatch) -> No
     # The duplicate-generator guard must precede lookup of an existing graph.
     sampler.sample(rows, PARAMETERS, _generators([100, 101, 102], "cuda"))
     cached = dict(sampler._graphs)
-    assert len(cached) == 1
+    assert len(cached) == (1 if _captured_generator_state_supported() else 0)
     actual = torch.Generator(device="cuda").manual_seed(2**64 - 23)
     reference = torch.Generator(device="cuda").manual_seed(2**64 - 23)
     multinomial = torch.multinomial
@@ -270,6 +284,7 @@ def test_cuda_shared_generator_uses_sequential_eager_fallback(monkeypatch) -> No
 @torch.inference_mode()
 def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch) -> None:
     sampler = BreezeFirstCodeSampler()
+    captures_generator_state = _captured_generator_state_supported()
     rows = _logits(1, device="cuda")
     actual_generators = _generators([42], "cuda")
     expected_generators = _generators([42], "cuda")
@@ -279,7 +294,7 @@ def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch)
         _assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows)
         _assert_states(actual_generators, expected_generators)
     cached = dict(sampler._graphs)
-    assert len(cached) == 32
+    assert len(cached) == (32 if captures_generator_state else 0)
     parameters = (0.9, 33, 1.0)
     expected = _reference(rows, parameters, expected_generators)
     multinomial = torch.multinomial
@@ -298,6 +313,11 @@ def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch)
     assert all(sampler._graphs[key] is value for key, value in cached.items())
     expected = _reference(rows, (0.9, 1, 1.0), expected_generators)
 
+    if not captures_generator_state:
+        _assert_outputs(sampler.sample(rows, (0.9, 1, 1.0), actual_generators), expected, rows)
+        _assert_states(actual_generators, expected_generators)
+        return
+
     def unexpected_multinomial(*_args, **_kwargs):
         raise AssertionError("Cache exhaustion must retain existing captured entries")
 

tests/model_executor/models/test_breeze_tts_2_graphs.py

@@ -9,7 +9,11 @@ from transformers.models.t5gemma2.modeling_t5gemma2 import T5Gemma2TextEncoder
 from vllm.platforms import current_platform
 
 from tests.helpers.mark import hardware_test
-from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import BreezeDepthDecoder, sample_logits
+from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import (
+    BreezeDepthDecoder,
+    _depth_layer_compile_options,
+    sample_logits,
+)
 from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze import BreezeForConditionalGeneration
 from vllm_omni.model_executor.models.breeze_tts_2.text_encoder_graph import (
     BreezeTextEncoderCompiled,
@@ -19,6 +23,17 @@ from vllm_omni.model_executor.models.breeze_tts_2.text_encoder_graph import (
 pytestmark = [pytest.mark.core_model]
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize(("is_rocm", "expected"), [(True, False), (False, True)])
+def test_depth_layer_compile_options(monkeypatch, is_rocm: bool, expected: bool) -> None:
+    monkeypatch.setattr(current_platform, "is_rocm", lambda: is_rocm)
+
+    assert _depth_layer_compile_options() == {
+        "epilogue_fusion": False,
+        "max_autotune": expected,
+    }
+
+
 @pytest.fixture
 def full_precision_matmul():
     # Different padded shapes can select different TF32 kernels. Compare
@@ -241,6 +256,7 @@ def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_m
     seeds[0], seeds[-1] = 2**63 + 17, 2**64 - 1
     generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in seeds]
     reference_generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in seeds]
+    captures_generator_state = current_platform.is_cuda() and hasattr(torch.cuda.CUDAGraph, "register_generator_state")
 
     expected_greedy = torch.cat(
         [
@@ -314,7 +330,7 @@ def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_m
             branches = 2 if guidance_scale != 1 else 1
             graph_key = (request_bucket * branches, branches, temperature == 0)
             with monkeypatch.context() as rng_guard:
-                if temperature > 0 and graph_key in depth._graphs:
+                if temperature > 0 and graph_key in depth._graphs and captures_generator_state:
                     # A warmed replay must generate noise within the CUDA
                     # graph, without dispatching new RNG kernels from Python.
                     rng_guard.setattr(torch.Tensor, "exponential_", reject_eager_noise)
@@ -349,7 +365,7 @@ def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_m
     expected_keys.update((batch, 2, greedy) for batch in (2, 8, 16, 32) for greedy in (False, True))
     assert set(depth._graphs) == expected_keys
     for (batch_bucket, branches, greedy), entry in depth._graphs.items():
-        if greedy:
+        if greedy or not captures_generator_state:
             assert entry.noise_generators is None
         else:
             assert entry.noise_generators is not None

tests/model_executor/models/test_talker_first_audio.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import talker_first_audio_enabled
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    ("env", "extra", "async_chunk", "v2", "expected"),
+    [
+        (None, {}, True, True, False),
+        ("1", {}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, True),
+        ("0", {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [4, 8]}, True, True, False),
+        (None, {"talker_first_audio": True, "initial_codec_chunk_frames": 1}, True, True, True),
+        (None, {"talker_first_audio": True, "initial_codec_chunk_frames": 4}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1]}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, False, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, False, False),
+    ],
+)
+def test_talker_first_audio_default_and_conditions(monkeypatch, env, extra, async_chunk, v2, expected):
+    if env is None:
+        monkeypatch.delenv("VLLM_OMNI_TALKER_FIRST_AUDIO", raising=False)
+    else:
+        monkeypatch.setenv("VLLM_OMNI_TALKER_FIRST_AUDIO", env)
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            async_chunk=async_chunk, use_v2_model_runner=v2, stage_connector_config={"extra": extra}
+        ),
+        device_config=SimpleNamespace(device="cuda"),
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend=None
+        ),
+        cache_config=SimpleNamespace(enable_prefix_caching=False),
+    )
+    assert talker_first_audio_enabled(config) is expected
+    for section, field, value in [
+        ("device_config", "device", "cpu"),
+        ("parallel_config", "tensor_parallel_size", 2),
+        ("parallel_config", "pipeline_parallel_size", 2),
+        ("parallel_config", "distributed_executor_backend", "mp"),
+        ("cache_config", "enable_prefix_caching", True),
+    ]:
+        with monkeypatch.context() as patch:
+            patch.setattr(getattr(config, section), field, value)
+            assert not talker_first_audio_enabled(config)
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
+    assert not talker_first_audio_enabled(config)
+
+
+@pytest.mark.parametrize("option", [False, True])
+def test_tts_requires_connector_opt_in(monkeypatch, option):
+    from vllm_omni.model_executor.models.qwen3_tts import first_audio
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor,
+    )
+
+    monkeypatch.setattr(Predictor, "_stage_connector_extra_config", lambda _: {"talker_first_audio": option})
+    monkeypatch.setattr(first_audio, "supports_talker_first_audio", lambda _: True)
+    assert first_audio.talker_first_audio_enabled(None) is option

tests/model_executor/stage_input_processors/test_fish_speech_async_chunk.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from collections import defaultdict
 from types import SimpleNamespace
 
@@ -18,7 +21,7 @@ def _make_transfer_manager(
     left_context: int = 1,
     initial_chunk_frames: int = 0,
     single_initial_chunk: bool = False,
-    backlog_chunk_frames: int = 0,
+    backlog_chunk_frames: int | str = 0,
     backlog_load_threshold: float = 0.75,
     max_num_seqs: int = 1,
 ):
@@ -227,3 +230,77 @@ def test_invalid_backlog_config_fails_loudly():
             {"audio_codes": torch.tensor([[1, 2, 3]])},
             request,
         )
+
+
+@pytest.mark.parametrize("single_initial_chunk", [False, True])
+@pytest.mark.parametrize("initial_chunk_frames", [0, 4, 10])
+@pytest.mark.parametrize("tensor_payload", [False, True])
+@pytest.mark.parametrize("pattern", ["steady", "drop_at_finish", "rise_midstream", "oscillate"])
+def test_backlog_changes_preserve_every_frame_once(single_initial_chunk, initial_chunk_frames, tensor_payload, pattern):
+    tm = _make_transfer_manager(
+        tensor_payload=tensor_payload,
+        chunk_frames=25,
+        left_context=25,
+        initial_chunk_frames=initial_chunk_frames,
+        single_initial_chunk=single_initial_chunk,
+        backlog_chunk_frames=50,
+        max_num_seqs=64,
+    )
+    request = _make_request()
+    emitted = []
+    for frame in range(1, 141):
+        high = (
+            pattern == "steady"
+            or (pattern == "drop_at_finish" and frame < 140)
+            or (pattern == "rise_midstream" and frame >= 37)
+            or (pattern == "oscillate" and (frame // 13) % 2 == 0)
+        )
+        others = 63 if high else 1
+        tm.code_prompt_token_ids = defaultdict(list, {"req": tm.code_prompt_token_ids["req"]})
+        for idx in range(others):
+            tm.code_prompt_token_ids[f"other-{idx}"] = [None] * 50
+        out = slow_ar_to_dac_decoder_async_chunk(
+            tm,
+            {"audio_codes": torch.tensor([[frame, frame + 1000]]), "audio_code_valid": torch.tensor([True])},
+            request,
+        )
+        if out is not None:
+            codes = out.codes.audio.reshape(2, -1)
+            emitted.extend(codes[0, out.meta.left_context_size :].tolist())
+            tm.put_req_chunk["req"] += 1
+    # EOS may arrive in a separate callback after the last valid frame.
+    out = slow_ar_to_dac_decoder_async_chunk(tm, {}, request, is_finished=True)
+    if out is not None and out.codes.audio.numel():
+        emitted.extend(out.codes.audio.reshape(2, -1)[0, out.meta.left_context_size :].tolist())
+    assert emitted == list(range(1, 141))
+    assert slow_ar_to_dac_decoder_async_chunk(tm, {}, request, is_finished=True) is None
+
+
+def test_final_backlog_drop_flushes_entire_pending_tail():
+    tm = _make_transfer_manager(
+        tensor_payload=True,
+        chunk_frames=25,
+        left_context=25,
+        initial_chunk_frames=4,
+        single_initial_chunk=True,
+        backlog_chunk_frames=50,
+        max_num_seqs=64,
+    )
+    request = _make_request()
+    emitted = []
+    for frame in range(1, 41):
+        if frame < 40:
+            for idx in range(63):
+                tm.code_prompt_token_ids[f"other-{idx}"] = [None] * 30
+        else:
+            for idx in range(32, 63):
+                tm.code_prompt_token_ids.pop(f"other-{idx}")
+        out = slow_ar_to_dac_decoder_async_chunk(
+            tm,
+            {"audio_codes": torch.tensor([[frame]])},
+            request,
+            is_finished=frame == 40,
+        )
+        if out is not None:
+            emitted.extend(out.codes.audio[0, out.meta.left_context_size :].tolist())
+    assert emitted == list(range(1, 41))

tests/model_executor/stage_input_processors/test_moss_raw_async_chunk.py

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

tests/model_executor/stage_input_processors/test_moss_raw_empty_terminal.py

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

tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py

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

tests/model_executor/stage_input_processors/test_qwen3_omni_chunk_ramp.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3
+from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ramp_decode_windows
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_Q = 2
+_RAMP = [1, 2, 4, 8, 16, 25]
+
+
+def _manager(**extra):
+    config = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "initial_codec_chunk_frames": 4, **extra}
+    return SimpleNamespace(
+        code_prompt_token_ids=defaultdict(list),
+        put_req_chunk=defaultdict(int),
+        ramp_chunk_count=defaultdict(int),
+        connector=SimpleNamespace(config={"extra": config}),
+    )
+
+
+def _request(request_id="req"):
+    return SimpleNamespace(external_req_id=request_id, sampling_params=SimpleNamespace(stop_token_ids=[]))
+
+
+def _frame(index: int) -> dict:
+    return {
+        "codes": {"audio": torch.tensor([[index, 1000 + index]], dtype=torch.long)},
+        "meta": {"codec_frame_valid": torch.ones(1, dtype=torch.int8)},
+    }
+
+
+def _stream(manager, num_frames, finish_after_last=True):
+    request = _request()
+    chunks = []
+
+    def emit(payload):
+        if payload is None:
+            return
+        codes = payload.codes.audio.reshape(_Q, -1)
+        chunks.append((codes[0].tolist(), int(payload.meta.left_context_size), bool(payload.meta.finished)))
+        # The transport counts every enqueued chunk.
+        manager.put_req_chunk[request.external_req_id] += 1
+        manager.ramp_chunk_count[request.external_req_id] += 1
+
+    for index in range(num_frames):
+        last = finish_after_last and index == num_frames - 1
+        emit(q3.talker2code2wav_async_chunk(manager, _frame(index), request, is_finished=last))
+    return chunks
+
+
+def test_ramp_decode_windows_include_left_context():
+    assert ramp_decode_windows(_RAMP, 25) == [1, 3, 7, 15, 31, 50]
+    assert ramp_decode_windows([1, 2], 0) == [1, 2]
+
+
+def test_ramp_chunks_grow_and_resend_left_context():
+    chunks = _stream(_manager(codec_chunk_ramp=_RAMP), 56 + 25, finish_after_last=False)
+
+    new_frames = [len(window) - left for window, left, _finished in chunks]
+    assert new_frames == [1, 2, 4, 8, 16, 25, 25]
+    assert [left for _window, left, _finished in chunks] == [0, 1, 3, 7, 15, 25, 25]
+    assert [len(window) for window, _left, _finished in chunks] == [1, 3, 7, 15, 31, 50, 50]
+    emitted = 0
+    for window, left, _finished in chunks:
+        # Each window ends at the newest frame and re-sends `left` earlier ones.
+        assert window == list(range(emitted - left, emitted + len(window) - left))
+        emitted += len(window) - left
+    # The first chunk is the stream's first frame alone (no context).
+    assert chunks[0][0] == [0]
+
+
+def test_ramp_flushes_partial_tail_with_left_context():
+    chunks = _stream(_manager(codec_chunk_ramp=_RAMP), 10)
+
+    assert [len(window) - left for window, left, _finished in chunks] == [1, 2, 4, 3]
+    window, left, finished = chunks[-1]
+    assert finished and left == 7 and window == list(range(0, 10))
+
+
+def test_ramp_finish_on_chunk_boundary_leaves_the_marker_to_the_transport():
+    manager = _manager(codec_chunk_ramp=_RAMP)
+    request = _request()
+    for index in range(3):
+        payload = q3.talker2code2wav_async_chunk(manager, _frame(index), request)
+        if payload is not None:
+            manager.ramp_chunk_count[request.external_req_id] += 1
+    # Frames 0 and 1-2 went out as chunks 0 and 1; the finish step has no frame.
+    finish = {"codes": {"audio": torch.zeros((1, _Q), dtype=torch.long)}, "meta": {"codec_frame_valid": torch.zeros(1)}}
+    assert q3.talker2code2wav_async_chunk(manager, finish, request, is_finished=True) is None
+
+
+def test_without_ramp_the_initial_chunk_setting_is_unchanged():
+    chunks = _stream(_manager(), 29 + 3)
+
+    assert [(len(window), left) for window, left, _finished in chunks] == [(4, 0), (29, 4), (28, 25)]
+
+
+def test_first_audio_marker_only_reaches_initial_codec_chunk():
+    manager = _manager(codec_chunk_ramp=[1, 1])
+    for index in range(2):
+        frame = _frame(index)
+        frame["meta"]["first_audio"] = torch.tensor([True])
+        payload = q3.talker2code2wav_async_chunk(manager, frame, _request())
+        assert bool(payload.meta.first_audio) is (index == 0)
+        manager.put_req_chunk["req"] += 1
+        manager.ramp_chunk_count["req"] += 1

tests/model_executor/stage_input_processors/test_qwen3_omni_first_token_handoff.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+ACCEPT = 24
+HIDDEN = 8
+PROMPT = [151644, 872, 5, 6, 151645, 151644, 77091, 198]
+EOS = 151645
+
+
+class _Manager:
+    def __init__(self) -> None:
+        self.put_req_chunk: defaultdict[str, int] = defaultdict(int)
+        self.request_payload: dict[str, dict] = {}
+        self.config = SimpleNamespace(
+            hf_config=SimpleNamespace(talker_config=SimpleNamespace(accept_hidden_layer=ACCEPT))
+        )
+
+    def process(self, output, request, is_finished=False):
+        payload = q3.thinker2talker_async_chunk(self, output, request, is_finished=is_finished)
+        if payload is not None:
+            self.put_req_chunk[request.external_req_id] += 1
+        return payload
+
+
+def _request(output_token_ids, *, history=True, max_tokens=64):
+    # The native data plane keeps token history only while chunk 0 is pending.
+    return SimpleNamespace(
+        external_req_id="r",
+        prompt_token_ids=list(PROMPT) if history else [],
+        output_token_ids=list(output_token_ids) if history else [],
+        all_token_ids=list(PROMPT) + list(output_token_ids) if history else [],
+        output_token_count=len(output_token_ids),
+        last_output_token_id=output_token_ids[-1] if output_token_ids else None,
+        resumable=False,
+        additional_information=None,
+        sampling_params=SimpleNamespace(all_stop_token_ids={EOS}, max_tokens=max_tokens),
+    )
+
+
+def _embed(token_id: int) -> torch.Tensor:
+    return torch.full((1, HIDDEN), float(token_id), dtype=torch.bfloat16)
+
+
+def _output(token_ids, sampled=None):
+    layers = {
+        0: torch.cat([_embed(t) for t in token_ids]),
+        ACCEPT: torch.randn(len(token_ids), HIDDEN).to(torch.bfloat16),
+    }
+    embed = {"tts_bos": _embed(1), "tts_eos": _embed(2), "tts_pad": _embed(3)}
+    if sampled is not None:
+        embed["sampled"] = sampled
+    return {"hidden_states": {"layers": layers}, "embed": embed}
+
+
+def _run_turn(tokens, *, publish, max_tokens=64, prefill_chunks=(len(PROMPT),)):
+    manager = _Manager()
+    rows, chunks, start = [], 0, 0
+    for i, size in enumerate(prefill_chunks):
+        last = i == len(prefill_chunks) - 1
+        sampled = _embed(tokens[0]) if last else torch.empty(0, dtype=torch.bfloat16)
+        request = _request(tokens[:1] if last else [], max_tokens=max_tokens)
+        payload = manager.process(_output(PROMPT[start : start + size], sampled if publish else None), request)
+        start += size
+        if payload is not None:
+            chunks += 1
+            rows.extend(payload.embed.prefill[len(PROMPT) :])
+    for step in range(1, len(tokens)):
+        history = manager.put_req_chunk["r"] == 0
+        request = _request(tokens[: step + 1], history=history, max_tokens=max_tokens)
+        output = _output([tokens[step - 1]], _embed(tokens[step]) if publish else None)
+        payload = manager.process(output, request)
+        if payload is None:
+            continue
+        chunks += 1
+        rows.extend(payload.embed.prefill[len(PROMPT) :] if payload.embed.prefill is not None else payload.embed.decode)
+    return [int(row[0]) for row in rows], chunks, manager
+
+
+@pytest.mark.parametrize(
+    ("tokens", "max_tokens"),
+    [
+        ([11, 12, 13, EOS], 64),  # ends on a stop token
+        ([11, 12, 13], 3),  # ends on the max_tokens-th token
+        ([11, EOS], 64),
+    ],
+)
+def test_sampled_stream_matches_captured_stream(tokens, max_tokens):
+    late, _chunks, _manager = _run_turn(tokens, publish=False, max_tokens=max_tokens)
+    early, _chunks, manager = _run_turn(tokens, publish=True, max_tokens=max_tokens)
+    # The final token is never processed, so neither path feeds it to the Talker.
+    assert early == late == tokens[:-1]
+    assert q3._SAMPLED_TEXT_STREAM in manager.request_payload["r"]
+
+
+def test_sampled_stream_sends_each_token_one_step_earlier():
+    manager = _Manager()
+    chunk0 = manager.process(_output(PROMPT, sampled=_embed(11)), _request([11]))
+    assert chunk0 is not None  # at the prefill step, not after the next decode step
+    assert torch.equal(chunk0.embed.prefill, torch.cat([_embed(t) for t in [*PROMPT, 11]]))
+    assert chunk0.hidden_states.output.shape[0] == len(PROMPT)
+    assert chunk0.ids.all == [*PROMPT, 11]
+    # Decode step 1 processes 11 and samples 12: it carries 12.
+    chunk1 = manager.process(_output([11], sampled=_embed(12)), _request([11, 12], history=False))
+    assert torch.equal(chunk1.embed.decode, _embed(12))
+
+
+def test_chunked_prefill_uses_sample_of_final_chunk_only():
+    early, chunks, _manager = _run_turn([11, 12, EOS], publish=True, prefill_chunks=(5, 3))
+    assert early == [11, 12] and chunks == 2
+
+
+def test_terminal_first_token_keeps_captured_stream():
+    manager = _Manager()
+    assert manager.process(_output(PROMPT, sampled=_embed(EOS)), _request([EOS])) is None
+    assert q3._SAMPLED_TEXT_STREAM not in manager.request_payload["r"]
+
+
+@pytest.mark.parametrize("prefill_rows", [0, 1, 3])
+def test_missing_sample_on_sampled_stream(prefill_rows):
+    manager = _Manager()
+    manager.process(_output(PROMPT, sampled=_embed(11)), _request([11]))
+    chunk = manager.process(_output([11], sampled=_embed(12)), _request([11, 12], history=False))
+    assert torch.equal(chunk.embed.decode, _embed(12))
+    if not prefill_rows:
+        with pytest.raises(RuntimeError, match="sample embedding missing"):
+            manager.process(_output([12]), _request([11, 12], history=False))
+    else:
+        before = manager.put_req_chunk["r"]
+        for tokens in (PROMPT[:prefill_rows], [11]):
+            partial = _output(tokens, sampled=torch.empty(0))
+            assert manager.process(partial, _request([11, 12], history=False)) is None
+        assert manager.put_req_chunk["r"] == before
+        chunk = manager.process(_output([12], sampled=_embed(13)), _request([11, 12, 13], history=False))
+        assert torch.equal(chunk.embed.decode, _embed(13))

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -1054,3 +1054,80 @@ def test_qwen2_5_omni_thinker2talker_full_payload_noop() -> None:
 
     payload = thinker2talker_full_payload(None, {"any": "thing"}, None)
     assert payload is None
+
+
+def _codec_transfer_manager(request_id: str, *, initial_frames: int = 2, chunk_frames: int = 4):
+    return SimpleNamespace(
+        code_prompt_token_ids=defaultdict(list),
+        put_req_chunk=defaultdict(int, {request_id: 0}),
+        connector=SimpleNamespace(
+            config={
+                "extra": {
+                    "initial_codec_chunk_frames": initial_frames,
+                    "codec_chunk_frames": chunk_frames,
+                    "codec_left_context_frames": 4,
+                }
+            }
+        ),
+    )
+
+
+def _codec_request(request_id: str) -> SimpleNamespace:
+    return SimpleNamespace(
+        external_req_id=request_id,
+        sampling_params=SimpleNamespace(stop_token_ids=[2150], stop_token_id=None),
+    )
+
+
+def test_talker2code2wav_async_chunk_reads_last_valid_row_of_multi_row_span() -> None:
+    # MRv2 eager frames: a prefill span carries its first frame in its last row.
+    request_id = "eager_prefill"
+    manager = _codec_transfer_manager(request_id)
+    codes = torch.zeros((5, 3), dtype=torch.long)
+    codes[-1] = torch.tensor([11, 12, 13])
+    valid = torch.tensor([0, 0, 0, 0, 1], dtype=torch.int8)
+
+    payload = q3.talker2code2wav_async_chunk(
+        manager, {"codes": {"audio": codes}, "meta": {"codec_frame_valid": valid}}, _codec_request(request_id)
+    )
+
+    assert payload is None  # below the initial chunk size
+    assert [frame.tolist() for frame in manager.code_prompt_token_ids[request_id]] == [[[11, 12, 13]]]
+
+
+def test_talker2code2wav_async_chunk_skips_invalid_frame_and_flushes_on_finish() -> None:
+    request_id = "eager_eos"
+    manager = _codec_transfer_manager(request_id)
+    manager.code_prompt_token_ids[request_id] = [torch.tensor([[7, 8, 9]])]
+    # The step that sampled codec EOS carries an invalid (placeholder) frame.
+    output = {
+        "codes": {"audio": torch.tensor([[2150, 0, 0]])},
+        "meta": {"codec_frame_valid": torch.tensor([0], dtype=torch.int8)},
+    }
+
+    payload = q3.talker2code2wav_async_chunk(manager, output, _codec_request(request_id), is_finished=True)
+
+    assert len(manager.code_prompt_token_ids[request_id]) == 1
+    assert payload is not None
+    assert payload.codes.audio.tolist() == [7, 8, 9]
+    assert payload.meta.finished.item() is True
+
+
+def test_talker2code2wav_async_chunk_without_validity_keeps_one_row_contract() -> None:
+    # V1 / MRv2 deferred frames: one row per step, zero rows are prefill
+    # placeholders and a stop-token CB0 marks the post-EOS step.
+    request_id = "legacy_rows"
+    manager = _codec_transfer_manager(request_id)
+    request = _codec_request(request_id)
+
+    assert (
+        q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.zeros((4, 3), dtype=torch.long)}}, request)
+        is None
+    )
+    assert q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[2150, 1, 1]])}}, request) is None
+    assert manager.code_prompt_token_ids[request_id] == []
+    q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[5, 6, 7]])}}, request)
+    payload = q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[8, 9, 10]])}}, request)
+
+    assert payload is not None  # two frames fill the initial chunk
+    assert payload.codes.audio.tolist() == [5, 8, 6, 9, 7, 10]

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -594,6 +594,40 @@ def test_streaming_ref_code_context_is_bounded_for_batchable_shapes():
     torch.testing.assert_close(frames[:3], ref_code[-3:])
 
 
+@pytest.mark.parametrize("tensor_frames", [False, True])
+def test_async_chunks_preserve_codebook_order_with_reference_and_short_tail(tensor_frames):
+    tm = _tm(left_context=3, initial_chunk_frames=1)
+    rid = "r-codebook-order"
+    quantizers = 16
+    reference = [[1500 + f * quantizers + q for q in range(quantizers)] for f in range(5)]
+    tm.request_payload[rid] = torch.tensor(reference, dtype=torch.long)
+    generated = [[1 + f * quantizers + q for q in range(quantizers)] for f in range(30)]
+
+    for start, end, finished in [(0, 1, False), (1, 26, False), (26, 30, True)]:
+        rows = generated[:end]
+        if tensor_frames:
+            # Non-contiguous int32 rows exercise conversion independently of
+            # the native producer's contiguous CPU-long representation.
+            storage = torch.tensor(rows, dtype=torch.int32).repeat_interleave(2, dim=1)
+            tm.code_prompt_token_ids[rid] = list(storage[:, ::2].unbind())
+        else:
+            tm.code_prompt_token_ids[rid] = rows
+        payload = talker2code2wav_async_chunk(
+            transfer_manager=tm,
+            multimodal_output=None if finished else {},
+            request=_req(rid, finished=finished),
+            is_finished=finished,
+        )
+        assert payload is not None
+        expected = (reference[-3:] if start == 0 else []) + generated[start:end]
+        codes = payload.codes.audio
+        assert codes.dtype == torch.long and codes.device.type == "cpu" and codes.is_contiguous()
+        torch.testing.assert_close(codes.reshape(quantizers, -1).T, torch.tensor(expected))
+        assert payload.meta.left_context_size == (3 if start == 0 else 0)
+        assert payload.meta.finished.item() == finished
+        tm.put_req_chunk[rid] += 1
+
+
 def test_ref_code_context_can_be_buffered_before_first_emit():
     tm = _tm()
     rid = "r-ref-buffered"

tests/worker/test_gpu_ar_model_runner.py

@@ -15,6 +15,8 @@ import numpy as np
 import pytest
 import torch
 from vllm.sampling_params import SamplingParams
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.utils import CpuGpuBuffer
 from vllm.v1.worker import gpu_input_batch
 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 
@@ -103,6 +105,8 @@ def test_post_sample_talker_mtp_uses_current_sample_and_hidden() -> None:
     )
 
     assert received["req_ids"] == ["ready"]
+    assert isinstance(received["input_ids"], torch.Tensor)
+    assert isinstance(received["hidden_states"], torch.Tensor)
     assert received["input_ids"].tolist() == [101]
     assert received["hidden_states"].tolist() == [[1.0, 2.0]]
     assert received["req_infos"] == [{"duplex": {"data_plane": True}}]
@@ -139,6 +143,8 @@ def test_post_sample_talker_mtp_uses_gpu_token_with_async_scheduling() -> None:
     )
 
     assert received["req_ids"] == ["ready"]
+    assert isinstance(received["input_ids"], torch.Tensor)
+    assert isinstance(received["hidden_states"], torch.Tensor)
     assert received["input_ids"].tolist() == [101]
     assert received["hidden_states"].tolist() == [[1.0, 2.0]]
     assert received["req_infos"] == [{"duplex": {"data_plane": True}}]
@@ -275,7 +281,12 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc
     )
     input_batch.sampling_metadata = input_batch._make_sampling_metadata()
 
-    received = []
+    received: list[SamplingMetadata] = []
+
+    def sample(logits: torch.Tensor, metadata: SamplingMetadata) -> str:
+        received.append(metadata)
+        return "model-sampler"
+
     runner = object.__new__(GPUARModelRunner)
     runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
     runner._pooler_payload_include_hidden_flag = True
@@ -283,7 +294,7 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc
     runner.model = SimpleNamespace(
         prefer_model_sampler=True,
         skips_model_sampler_output_token_history=True,
-        sample=lambda logits, metadata: received.append(metadata) or "model-sampler",
+        sample=sample,
     )
     runner.sampler = SimpleNamespace()
     logits = torch.zeros((1, 4))
@@ -1670,7 +1681,6 @@ class TestMergeModelKvTransferMetadata:
         # ...while the engine-shared originals are untouched.
         assert original["r1"]["custom_metadata"] is original_meta
         assert original_meta == {"a": 1}
-        assert "talker_codes" not in original["r1"]["custom_metadata"]
 
     def test_merge_without_model_meta_passes_entry_through(self):
         class _Model:
@@ -1713,7 +1723,7 @@ class TestDownstreamPayloadMemoization:
         return runner
 
     def test_missing_marker_defaults_true_without_memoizing(self):
-        stages = {"r1": None}
+        stages: dict[str, int | None] = {"r1": None}
         runner = self._runner(stages)
 
         assert runner._request_needs_downstream_stage_payload("r1") is True
@@ -1861,6 +1871,7 @@ class TestPreferModelSamplerNoneFallback:
             "minicpmo_4_5",
             "minimax_music3",
             "nemotron_voicechat",
+            "yue2",  # Always supplies SamplerOutput, including empty/prefill steps.
         }
         assert declarers == expected, (
             "The set of models declaring `prefer_model_sampler` changed:\n"
@@ -1875,3 +1886,63 @@ class TestPreferModelSamplerNoneFallback:
             "least tolerates -- that fallback, then add its directory name to "
             "`expected` above. If you REMOVED one, drop its name."
         )
+
+
+@pytest.mark.parametrize("accepts_extra_args", [False, True])
+def test_prepare_hook_keeps_legacy_signature_and_orders_opt_in_metadata(accepts_extra_args: bool) -> None:
+    runner = object.__new__(GPUARModelRunner)
+    runner.model = torch.nn.Module()
+    runner.model.accepts_runner_sampling_extra_args = accepts_extra_args
+    runner.requests = {
+        rid: CachedRequestState(
+            req_id=rid,
+            prompt_token_ids=[1],
+            mm_features=[],
+            sampling_params=params,
+            generator=None,
+            block_ids=([],),
+            num_computed_tokens=0,
+            output_token_ids=[],
+        )
+        for rid, params in (
+            ("a", SamplingParams(extra_args={"seed": 7})),
+            ("b", None),
+            ("unused", SamplingParams(extra_args={"seed": 9})),
+        )
+    }
+    runner.discard_request_mask = CpuGpuBuffer(3, dtype=torch.bool, device=torch.device("cpu"), pin_memory=False)
+    runner.discard_request_mask.np[:] = [True, False, True]
+    ids, positions = torch.tensor([3, 4]), torch.tensor([0, 1])
+    step_inputs = dict(
+        input_ids=ids,
+        positions=positions,
+        inputs_embeds=None,
+        num_computed_tokens=np.array([0, 1]),
+        num_scheduled_tokens=np.array([1, 1]),
+        input_ids_buffer=ids,
+    )
+    seen = {}
+
+    def legacy_hook(
+        *, input_ids, positions, inputs_embeds, req_ids, num_computed_tokens, num_scheduled_tokens, input_ids_buffer
+    ):
+        seen["req_ids"] = req_ids
+        return input_ids, positions
+
+    def opt_in_hook(*, sampling_extra_args, discard_mask, **kwargs):
+        seen["extra_args"] = sampling_extra_args
+        seen["discard"] = discard_mask.tolist()
+        return legacy_hook(**kwargs)
+
+    out = runner._call_prepare_runner_inputs(
+        opt_in_hook if accepts_extra_args else legacy_hook,
+        req_ids=["b", "a"],
+        **step_inputs,
+    )
+    assert out[0] is ids and out[1] is positions
+    assert seen["req_ids"] == ["b", "a"]
+    if accepts_extra_args:
+        assert seen["extra_args"] == [{}, {"seed": 7}]
+        assert seen["discard"] == [True, False]
+    else:
+        assert set(seen) == {"req_ids"}