[CI] Incorporate upstream Server VAD test quarantine

Merge upstream #7314 for the known Qwen3-Omni Server VAD input_backpressure failure tracked in #7279. Preserve all MOSS-TTS runtime changes and existing validation coverage.

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

.buildkite/cuda/test-ready.yml

@@ -93,7 +93,7 @@ steps:
     commands:
       - |
         timeout 15m bash -c "
-                  pytest -sv tests/engine/ tests/model_executor/ -m 'core_model and cuda' --run-level 'core_model'
+                  pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py -m 'core_model and cuda' --run-level 'core_model'
         "
     mirror_hardwares: l4_1
 

examples/online_serving/text_to_speech/README.md

@@ -259,21 +259,25 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
 0.6B DualAR TTS at 44.1 kHz, 11 languages, zero-shot voice cloning.
 
 ### Prerequisites
+
 None beyond the base install: the neural audio codec is implemented in tree and
 its weights (`codec.pth`) ship with the checkpoint.
 
 ### Launch
+
 ```bash
 vllm serve Audio8/Audio8-TTS-Preview-0.6b --omni --port 8092
 # or:
 ./audio8_tts/run_server.sh
 ```
+
 The deploy config auto-loads from `vllm_omni/deploy/audio8_tts.yaml` (HF
 `model_type` is `arktts`). Do **not** pass `--trust-remote-code`: vllm-omni
 registers its own `arktts` config, and transformers would otherwise prefer the
 checkpoint's remote code and bypass it.
 
 ### Text-only synthesis
+
 ```bash
 curl -X POST http://localhost:8092/v1/audio/speech \
     -H "Content-Type: application/json" \
@@ -284,6 +288,7 @@ curl -X POST http://localhost:8092/v1/audio/speech \
 ```
 
 ### Voice cloning
+
 ```bash
 curl -X POST http://localhost:8092/v1/audio/speech \
     -H "Content-Type: application/json" \
@@ -293,18 +298,21 @@ curl -X POST http://localhost:8092/v1/audio/speech \
         "ref_text": "The exact transcript of the reference recording."
     }' --output cloned.wav
 ```
+
 `ref_audio` also accepts a `data:audio/wav;base64,...` URL. `ref_text` is
 mandatory whenever `ref_audio` is present and must match the recording.
 Uploading a voice via `POST /v1/audio/voices` lets the server reuse the encoded
 reference codes across requests (`voice: "<name>"`).
 
 ### Streaming
+
 ```bash
 python audio8_tts/speech_client.py --text "Welcome to Audio8 TTS." --stream --output out.pcm
 ffplay -f s16le -ar 44100 -ac 1 out.pcm
 ```
 
 ### Gradio demo
+
 ```bash
 ./audio8_tts/run_gradio_demo.sh          # server + demo
 python audio8_tts/gradio_demo.py --api-base http://localhost:8092   # demo only
@@ -328,7 +336,7 @@ these as a lower bound. 20 prompts from `seed_tts_smoke/en`, mean audio 4.5 s,
 deploy defaults:
 
 | | HF reference (batch=1) | vllm-omni c=1 | c=4 | c=8 |
-|---|---|---|---|---|
+| --- | --- | --- | --- | --- |
 | RTF (mean) | 1.008 | **0.19** | 0.30 | 0.54 |
 | Throughput (req/s) | 0.23 | 1.18 | 2.84 | **2.93** |
 | E2E latency (mean, ms) | 4375 | 850 | 1349 | 2361 |
@@ -347,7 +355,7 @@ job holding ~80% SM, no repeats), so read them as directional, not as measured
 speedups. Same setup, varying stage 0's `max_num_seqs` (stage 1 stays at 1):
 
 | stage-0 `max_num_seqs` | client concurrency | req/s | RTF | TTFP mean / p99 (ms) | underrun p99 (s) |
-|---|---|---|---|---|---|
+| --- | --- | --- | --- | --- | --- |
 | 4 (default) | 1 | 1.18 | 0.19 | 63 / 67 | 0.00 |
 | 4 | 4 | 2.84 | 0.30 | 105 / 148 | 0.00 |
 | 4 | 8 | 2.93 | 0.54 | 1124 / 1642 | 0.00 |
@@ -376,6 +384,7 @@ underrun. Raise it to 8 if first-packet latency under load matters more than a
 rare sub-300 ms gap.
 
 ### Notes
+
 - Output: 44.1 kHz mono; ~21.5 codec frames per second.
 - No built-in speaker presets. Omit `voice` for a random timbre, or clone one.
 - `max_new_tokens` caps generated codec frames (1 frame ~= 46 ms).
@@ -553,6 +562,46 @@ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.
 
 ---
 
+## MOSS-TTS Local Transformer v1.5
+
+For a single H200, the optional
+[`moss_tts_local_h200.yaml`](../../../vllm_omni/deploy/moss_tts_local_h200.yaml)
+deployment places both the talker and codec on logical GPU 0:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \
+    vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \
+    --omni --trust-remote-code \
+    --deploy-config vllm_omni/deploy/moss_tts_local_h200.yaml \
+    --disable-log-stats
+```
+
+Run from the repository root and select an available physical GPU with
+`CUDA_VISIBLE_DEVICES`. The command enables event-driven orchestration and
+disables detailed per-request statistics logging to reduce CPU overhead at
+high concurrency. Remove `--disable-log-stats` when those statistics are needed;
+keep these settings identical when comparing performance.
+This preset configures 256 request slots per stage,
+a 32 GiB talker KV cache, talker CUDA Graph buckets through 512 scheduled tokens,
+and codec batch buckets through 256. It requires more than 80 GiB of GPU memory;
+request capacity also depends on input and generated lengths. The default
+[`moss_tts_local.yaml`](../../../vllm_omni/deploy/moss_tts_local.yaml) remains
+available for smaller deployments.
+
+The preset enables the optional codec backend with
+`hf_overrides.codec_attention_backend: triton` on stage 1. It preserves the
+streaming ring-cache mask and uses BF16 attention with 64-dimensional heads;
+other attention shapes use PyTorch SDPA. Set the backend to `sdpa` to use the
+default attention implementation. Codec terminal tails share execution only
+when they already map to the same padded CUDA Graph; returned audio retains
+each request's actual length.
+
+Voice cloning requests use `ref_audio` and `ref_text`. For streaming output,
+set `stream: true`, `stream_format: "audio"`, and `response_format: "pcm"`;
+the native PCM format is 48 kHz, stereo, signed 16-bit little-endian. Compute
+audio throughput as `PCM bytes / (48000 * 2 * 2) / elapsed seconds`, and compare
+the same requests, concurrency, and reference-cache state.
+
 ## MOSS-TTS-Nano
 
 Single-stage 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec at 48 kHz mono. Every request must include `ref_audio`; there are no built-in speaker presets.

recipes/OpenMOSS/MOSS-TTS.md

@@ -16,7 +16,7 @@ Use this recipe for 24 kHz multilingual TTS with voice cloning (20 languages inc
 Chinese and English). Choose a variant based on your latency and quality requirements:
 
 | Model | Params | Use case |
-|---|---|---|
+| --- | --- | --- |
 | MOSS-TTS | 8B | General TTS, highest quality |
 | MOSS-TTS-v1.5 | 8B | General TTS upgrade of 1.0: 31 languages, steadier cloning, `[pause Xs]` markers (set `language` for best results); same `MossTTSDelay` API |
 | MOSS-TTS-Realtime | 1.7B | Lowest latency (TTFB ~180 ms), streaming-first |
@@ -24,9 +24,17 @@ Chinese and English). Choose a variant based on your latency and quality require
 | MOSS-SoundEffect | 8B | Sound effect synthesis from text description |
 | MOSS-VoiceGenerator | 1.7B | Zero-shot voice design |
 
-All variants share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
+The variants above share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
 output 24 kHz mono audio.
 
+MOSS-TTS-Local-Transformer-v1.5 uses MOSS-Audio-Tokenizer-v2 and outputs 48 kHz
+stereo audio. For Local voice cloning through `/v1/audio/speech`, provide an
+accurate `ref_text` transcript alongside `ref_audio`. The adapter uses the
+reference transcript and audio as a continuation prefix, then generates only
+the requested `input` speech. Without a nonblank `ref_text`, Local uses
+audio-reference generation. Reference transcripts must match the reference
+audio; they are not style instructions.
+
 ## References
 
 - Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)

tests/benchmarks/test_seed_tts_dataset_variants.py

@@ -26,6 +26,7 @@ _MODULE_NAME = "vllm_omni.benchmarks.data_modules.seed_tts_dataset"
 
 if _MODULE_NAME not in sys.modules:
     _spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
+    assert _spec is not None and _spec.loader is not None
     _mod = importlib.util.module_from_spec(_spec)
     sys.modules[_MODULE_NAME] = _mod
     _spec.loader.exec_module(_mod)
@@ -237,7 +238,8 @@ def test_seed_tts_design_dataset_rejects_missing_description(seed_tts_design_roo
         assert req.seed_tts_utterance_id == "ok"
 
 
-def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
+@pytest.mark.parametrize("audio_seconds", [0.1, 61])
+def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch, audio_seconds):
     from vllm_omni.benchmarks.data_modules import seed_tts_eval
 
     calls = {}
@@ -252,8 +254,9 @@ def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
             return self
 
     class FakeProcessor:
-        def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False):
+        def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False, **kwargs):
             calls["return_attention_mask"] = return_attention_mask
+            calls["processor_kwargs"] = kwargs
             assert sampling_rate == 16000
             assert return_tensors == "pt"
             assert len(wav) > 0
@@ -262,9 +265,10 @@ def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
                 attention_mask=FakeTensor("mask") if return_attention_mask else None,
             )
 
-        def get_decoder_prompt_ids(self, *, language, task):
+        def get_decoder_prompt_ids(self, *, language, task, no_timestamps=True):
             assert language == "english"
             assert task == "transcribe"
+            calls["no_timestamps"] = no_timestamps
             return [(1, 2)]
 
         def batch_decode(self, predicted_ids, *, skip_special_tokens):
@@ -283,10 +287,18 @@ def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
     monkeypatch.setattr(seed_tts_eval, "_en_model", FakeModel())
     monkeypatch.setattr(seed_tts_eval, "_device", "cuda:1")
 
-    text = seed_tts_eval._transcribe_en_f32_16k(np.ones(1600, dtype=np.float32))
+    text = seed_tts_eval._transcribe_en_f32_16k(np.ones(int(audio_seconds * 16000), dtype=np.float32))
 
     assert text == "hello"
     assert calls["return_attention_mask"] is True
     assert calls["input_features"].device == "cuda:1"
     assert calls["generate_kwargs"]["attention_mask"].device == "cuda:1"
     assert calls["generate_kwargs"]["forced_decoder_ids"] == [(1, 2)]
+    if audio_seconds > 30:
+        assert calls["processor_kwargs"] == {"truncation": False, "padding": "longest"}
+        assert calls["generate_kwargs"]["return_timestamps"] is True
+        assert calls["no_timestamps"] is False
+    else:
+        assert calls["processor_kwargs"] == {}
+        assert "return_timestamps" not in calls["generate_kwargs"]
+        assert calls["no_timestamps"] is True

tests/entrypoints/openai_api/test_tts_adapter.py

@@ -166,6 +166,56 @@ def test_moss_tts_seed_falls_back_to_stage_default(adapter_cls, mocker):
     assert prepared.tts_params["seed"] == [42]
 
 
+@pytest.mark.parametrize(
+    "variant,ref_text,mode",
+    [
+        ("local", " Reference. ", "continuation"),
+        ("local", None, "generation"),
+        ("local", "  ", "generation"),
+        ("tts", "Reference.", "generation"),
+    ],
+)
+def test_moss_reference_transcript_mode(variant, ref_text, mode, mocker):
+    import torch
+
+    reference = [torch.ones((3, 12), dtype=torch.int64)]
+    unified = torch.arange(52, dtype=torch.int64).reshape(1, 4, 13)
+    processor = mocker.Mock(return_value={"input_ids": unified})
+    server = mocker.Mock(_moss_variant="local", uploaded_speakers={"speaker": {}})
+    server._voice_created_at.return_value = 123
+    request = OpenAICreateSpeechRequest(
+        input="Target.", ref_text=ref_text, language="English", voice="Speaker", seed=0, max_new_tokens=2048
+    )
+    adapter = MossTTSAdapter(SpeechServingContext(server=server))
+
+    adapter._moss_variant = variant
+    mocker.patch.object(adapter, "_get_moss_processor", return_value=processor)
+    mocker.patch.object(
+        adapter, "_encode_moss_references", new=mocker.AsyncMock(return_value=(reference, {0: "reference-key"}))
+    )
+
+    prepared = asyncio.run(adapter.build(request, [], has_inline_ref_audio=False))
+
+    if mode == "continuation":
+        processor.build_user_message.assert_called_once_with(text="Reference. Target.", language="English")
+        processor.build_assistant_message.assert_called_once_with(audio_codes_list=reference)
+        conversation = [processor.build_user_message.return_value, processor.build_assistant_message.return_value]
+    else:
+        processor.build_user_message.assert_called_once_with(text="Target.", language="English", reference=reference)
+        processor.build_assistant_message.assert_not_called()
+        conversation = [processor.build_user_message.return_value]
+    processor.assert_called_once_with(conversations=[conversation], mode=mode)
+    assert prepared.prompt["prompt_token_ids"] == unified[0, :, 0].tolist()
+    assert torch.equal(prepared.tts_params["codes"]["ref"], unified[0, :, 1:])
+    assert prepared.tts_params["max_new_frames"] == [2048]
+    assert prepared.tts_params["ref_audio_cache_key"] == "reference-key"
+    assert prepared.tts_params["voice_name"] == ["speaker"]
+    assert prepared.tts_params["voice_created_at"] == [123]
+    assert prepared.tts_params["seed"] == [0]
+    assert prepared.prompt["cache_salt"]
+    assert request.input == "Target."
+
+
 def test_qwen3_tts_metadata():
     assert Qwen3TTSAdapter.backend == "ar"
     assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)

tests/entrypoints/test_omni_entrypoints.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 queue
@@ -97,9 +100,11 @@ class FakeAsyncOmniEngine:
     ) -> None:
         self.model = model
         self.config_path = None
-        self.stage_configs: list[Any] = []
         self.stage_metadata = stage_metadata or [THREE_STAGE_META[-1]]
         self.num_stages = len(self.stage_metadata)
+        self.stage_configs = [
+            StageConfig(stage_id=i, model_stage="dummy-model").to_omegaconf() for i in range(self.num_stages)
+        ]
         self.default_sampling_params_list = default_sampling_params_list or [
             SamplingParams(max_tokens=8) for _ in range(self.num_stages)
         ]
@@ -222,6 +227,25 @@ def test_resolve_sampling_params_list_preserves_stage_constraints():
     assert caller_params.stop_token_ids == [7]
 
 
+@pytest.mark.parametrize("use_defaults", [False, True])
+def test_moss_local_output_policy_preserves_codec_streaming(use_defaults):
+    from vllm_omni.model_executor.models.moss_tts.pipeline import MOSS_TTS_LOCAL_PIPELINE
+
+    base = _make_base()
+    base.engine.num_stages = 2
+    base.sampling_constraints_list = [stage.sampling_constraints for stage in MOSS_TTS_LOCAL_PIPELINE.stages]
+    base.default_sampling_params_list = [
+        base._apply_sampling_constraints(SamplingParams(), constraints)
+        for constraints in base.sampling_constraints_list
+    ]
+    caller = [SamplingParams(output_kind=RequestOutputKind.DELTA) for _ in range(2)]
+    result = base.resolve_sampling_params_list(None if use_defaults else caller, allow_delta_coercion=True)
+
+    assert [params.output_kind for params in result] == [RequestOutputKind.FINAL_ONLY, RequestOutputKind.DELTA]
+    assert all(params.output_kind == RequestOutputKind.DELTA for params in caller)
+    assert base.default_sampling_params_list[0].output_kind == RequestOutputKind.FINAL_ONLY
+
+
 def _stage_spec(
     stage_id: int,
     *,

tests/model_executor/models/moss_tts/test_local_depth_rope_cache.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
 """CPU parity tests for MOSS-TTS Local Depth RoPE caching."""
 
 import pytest
@@ -42,3 +42,16 @@ def test_rope_cache_matches_reference_and_reuses_storage() -> None:
     # Smaller slices retain the fixed allocation used by CUDA Graph.
     attn.prepare_rope_cache(6, device, torch.float32)
     assert attn._rope_cos_cache.data_ptr() == cache_ptr
+
+
+def test_frame_local_kv_matches_full_prefix_and_reuses_slots() -> None:
+    model = MossTTSLocalDepthTransformer(_test_config()).eval()
+    model.h[0].attn.prepare_rope_cache(12, torch.device("cpu"), torch.float32)
+    cache = (torch.empty(2, 4, 12, 8), torch.empty(2, 4, 12, 8))
+    with torch.inference_mode():
+        for _ in range(2):
+            inputs = torch.randn(2, 12, 32)
+            for position in range(12):
+                expected = model._forward_prefix(inputs[:, : position + 1])[:, -1:]
+                actual = model._forward_prefix(inputs[:, position : position + 1], cache, position)
+                torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5)

tests/model_executor/models/moss_tts/test_streaming_attention.py

@@ -0,0 +1,58 @@
+# 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,
+    StreamingExecutionContext,
+)
+from vllm_omni.model_executor.models.moss_tts.streaming_attention import masked_attention
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize(
+    "batch,heads,queries,keys", [(2, 20, 1, 125), (4, 20, 15, 125), (1, 12, 480, 400), (4, 12, 32, 250)]
+)
+def test_masked_attention_matches_sdpa_with_strides_and_invalid_rows(batch, heads, queries, keys):
+    torch.manual_seed(7)
+    q = torch.randn(batch, queries, 3, heads, 64, device="cuda", dtype=torch.bfloat16)[:, :, 0].permute(0, 2, 1, 3)
+    k = torch.randn(batch, heads, keys, 64, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn_like(k)
+    positions = torch.arange(keys, device="cuda").roll(17)
+    query_positions = torch.arange(queries, device="cuda") + keys - queries
+    mask = (positions[None, :] <= query_positions[:, None])[None, None].expand(batch, 1, -1, -1).clone()
+    mask[0, :, 0] = False  # Cover an all-masked row alongside valid rows.
+    actual = masked_attention(q, k, v, mask)
+    expected = F.scaled_dot_product_attention(q, k, v, mask)
+    assert bool(actual.isfinite().all())
+    torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.01)
+    assert actual[0, :, 0].count_nonzero() == 0
+
+
+def test_streaming_attention_preserves_ring_wrap_reordering_and_slot_reset():
+    torch.manual_seed(9)
+    reference = MossAudioTokenizerMultiheadAttention(
+        128, 2, causal=True, context=9, device="cuda", dtype=torch.bfloat16
+    )
+    candidate = copy.deepcopy(reference)
+    candidate._streaming_attention = masked_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")
+        context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid)
+        x = torch.randn(3, 3, 128, device="cuda", dtype=torch.bfloat16)
+        with torch.inference_mode():
+            expected = reference(x, x, x, execution_context=context)
+            actual = candidate(x, x, x, execution_context=context)
+        torch.testing.assert_close(actual, expected, atol=0.01, rtol=0.02)
+        torch.testing.assert_close(candidate._streaming_state.offset, reference._streaming_state.offset, 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_terminal_batch.py

@@ -0,0 +1,89 @@
+# 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.moss_tts.modeling_moss_tts_codec import (
+    MossTTSCodecDecoder,
+    _MossCodecStreamSession,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class CausalCodec(nn.Module):
+    downsample_rate = 1
+
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(()))
+
+    def initialize_decoder_state_pool(self, capacity, scratch):
+        self.state = torch.zeros(capacity + scratch)
+
+    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]
+        return SimpleNamespace(audio=audio[:, None], audio_lengths=lengths)
+
+
+def session():
+    return _MossCodecStreamSession(CausalCodec(), state_capacity=8, n_vq=2, vllm_config=None)
+
+
+def test_terminal_padding_preserves_audio_live_state_and_reused_slots():
+    s = session()
+    assert [s.acquire(), s.acquire()] == [0, 1]
+    plan = {0: torch.ones(2, 3), 1: torch.ones(2, 15)}
+    with pytest.raises(ValueError, match="Only terminal"):
+        s.step(plan)
+    assert s._codec.state.count_nonzero() == 0
+    s.step({0: torch.ones(2, 1), 1: torch.full((2, 1), 2)})
+    actual = s.step(plan, terminal_slots={0})
+    torch.testing.assert_close(actual[0], torch.tensor([[4.0, 6.0, 8.0]]))
+    torch.testing.assert_close(actual[1], torch.arange(6.0, 36.0, 2.0)[None])
+    s.release(0, state_already_reset=True)
+    assert s.acquire() == 0
+    resumed = s.step({0: torch.ones(2, 1), 1: torch.ones(2, 1)})
+    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):
+    s = session()
+    s._cudagraph_wrapper = (
+        None
+        if graph_capacity is None
+        else SimpleNamespace(
+            batch_sizes=[1, graph_capacity],
+            _select_frame_size=lambda size, allow_padding: next((b for b in [1, 8, 15] if b >= size), None),
+            decode=lambda *args, **kwargs: None,
+        )
+    )
+    step = mocker.spy(s, "step")
+    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
+    lengths = [3, 7, 8, 1, 15]
+    finished = [True, True, False, True, False]
+    result = decoder._decode_streaming_batch(
+        [
+            (i, str(i), torch.ones(2, length), done)
+            for i, (length, done) in enumerate(zip(lengths, finished, strict=True))
+        ]
+    )
+    assert step.call_count == expected_calls
+    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.

tests/worker/test_batched_omni_output.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration
+from vllm_omni.worker.gpu_ar_model_runner import _snapshot_tensor_payload_to_cpu_async
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+def test_async_packed_snapshot_survives_source_reuse():
+    values = [
+        torch.arange(12, device="cuda").view(1, 12),
+        torch.empty((0, 12), device="cuda", dtype=torch.long),
+        torch.arange(24, device="cuda").view(2, 12),
+    ]
+    expected = [v.cpu() for v in values]
+    snapshot = _snapshot_tensor_payload_to_cpu_async(
+        {"list": values, "tuple": tuple(values)}, copy_stream=torch.cuda.Stream(), pin_memory=True
+    )
+    for v in values:
+        v.fill_(-99)
+    snapshot.wait()
+    torch.testing.assert_close(snapshot.payload, {"list": expected, "tuple": tuple(expected)}, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("offsets", [[0, 1], [1, 4]])
+def test_mtp_preserves_rows_and_snapshot_without_cuda_sync(monkeypatch, offsets):
+    import vllm_omni.worker.gpu_model_runner as mod
+    from tests.worker.test_omni_gpu_model_runner import _make_runner, _noop_forward_context
+
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+    runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
+    runner.model.gpu_resident_buffer_keys = {("codes", "audio")}
+    for name in ("talker_mtp_input_ids", "talker_mtp_inputs_embeds", "last_talker_hidden", "text_step"):
+        buffer = getattr(runner, name)
+        buffer.gpu = buffer.gpu.to("cuda")
+    embeds = torch.ones((2, 4), device="cuda")
+    codes = torch.arange(2, device="cuda").view(2, 1)
+    runner.talker_mtp = lambda *args, **kwargs: (embeds, codes)
+    output = torch.full((5, 4), -10.0, device="cuda")
+    previous = torch.cuda.get_sync_debug_mode()
+    try:
+        torch.cuda.set_sync_debug_mode("error")
+        runner._talker_mtp_forward(["r1", "r2"], output, offsets)
+    finally:
+        torch.cuda.set_sync_debug_mode(previous)
+    expected = torch.full((5, 4), -10.0)
+    expected[offsets] = 1.0
+    torch.testing.assert_close(output.cpu(), expected, rtol=0, atol=0)
+    codes.fill_(999)
+    assert runner.model_intermediate_buffer["r1"]["codes"]["audio"].item() == 0
+    assert runner.model_intermediate_buffer["r2"]["codes"]["audio"].item() == 1
+
+
+def test_mixed_stop_flags_do_not_synchronize_cuda():
+    model = object.__new__(MossTTSLocalTalkerForGeneration)
+    torch.nn.Module.__init__(model)
+    model.n_vq = 12
+    model.audio_pad_token_id = 1024
+    valid = torch.arange(12, device="cuda").view(1, 12)
+    stopped = torch.full_like(valid, 1024)
+    infos = [{"audio_codes": {"current": valid}}, {}, {"audio_codes": {"current": stopped}}]
+    hidden = torch.ones((3, 4), device="cuda")
+    previous = torch.cuda.get_sync_debug_mode()
+    try:
+        torch.cuda.set_sync_debug_mode("error")
+        output = model.make_omni_output(hidden, runtime_additional_information=infos)
+    finally:
+        torch.cuda.set_sync_debug_mode(previous)
+    assert model._batch_should_continue.tolist() == [True, True, False]
+    torch.testing.assert_close(output.multimodal_outputs["codes"]["audio"], [valid, valid[:0], stopped])
+    assert all("current" not in info.get("audio_codes", {}) for info in infos)
+    prefill = model.make_omni_output(hidden[:1], runtime_additional_information=[{}])
+    assert model._batch_should_continue is None and prefill.multimodal_outputs == {}

vllm_omni/benchmarks/data_modules/seed_tts_eval.py

@@ -464,25 +464,33 @@ def _transcribe_en_f32_16k(wav_f32: np.ndarray) -> str:
         return ""
     with _lock:
         assert _en_processor is not None and _en_model is not None and _device is not None
+        # Whisper's default feature extraction truncates at 30 seconds. Keep
+        # abnormal tails so WER evaluates the complete generated response.
+        long_audio = len(wav_f32) > 30 * 16000
+        processor_kwargs = {"truncation": False, "padding": "longest"} if long_audio else {}
         try:
             inputs = _en_processor(
                 wav_f32,
                 sampling_rate=16000,
                 return_tensors="pt",
                 return_attention_mask=True,
+                **processor_kwargs,
             )
         except TypeError:
-            inputs = _en_processor(wav_f32, sampling_rate=16000, return_tensors="pt")
+            inputs = _en_processor(wav_f32, sampling_rate=16000, return_tensors="pt", **processor_kwargs)
         input_features = inputs.input_features.to(_device)
         attention_mask = getattr(inputs, "attention_mask", None)
         if attention_mask is None and isinstance(inputs, dict):
             attention_mask = inputs.get("attention_mask")
         generate_kwargs: dict[str, Any] = {}
+        if long_audio:
+            generate_kwargs["return_timestamps"] = True
         if attention_mask is not None:
             generate_kwargs["attention_mask"] = attention_mask.to(_device)
         with torch.no_grad():
             try:
-                forced = _en_processor.get_decoder_prompt_ids(language="english", task="transcribe")
+                prompt_kwargs = {"no_timestamps": False} if long_audio else {}
+                forced = _en_processor.get_decoder_prompt_ids(language="english", task="transcribe", **prompt_kwargs)
                 predicted_ids = _en_model.generate(input_features, forced_decoder_ids=forced, **generate_kwargs)
             except Exception:
                 predicted_ids = _en_model.generate(
@@ -548,14 +556,13 @@ def _expand_seed_tts_turn_outputs(
         turn_pcm = getattr(output, "tts_turn_pcm_bytes", None)
         session_pcm = getattr(output, "tts_output_pcm_bytes", None)
         for turn_index, turn in enumerate(turns):
-            expanded_requests.append(
-                replace(
-                    request,
-                    prompt=turn.target_text,
-                    seed_tts_utterance_id=turn.utterance_id,
-                    seed_tts_turns=(),
-                )
+            turn_request = replace(
+                request,
+                seed_tts_utterance_id=turn.utterance_id,
+                seed_tts_turns=(),
             )
+            turn_request.prompt = turn.target_text
+            expanded_requests.append(turn_request)
             turn_output = copy(output)
             if isinstance(turn_pcm, list) and turn_index < len(turn_pcm):
                 turn_output.tts_output_pcm_bytes = turn_pcm[turn_index]

vllm_omni/deploy/moss_tts_local_h200.yaml

@@ -0,0 +1,64 @@
+# MOSS-TTS Local Transformer v1.5: both stages on one H200.
+# Select one free physical GPU with CUDA_VISIBLE_DEVICES; devices below are logical.
+# Capacity and graph capture require more than 80 GiB of GPU memory.
+async_chunk: true
+
+connectors:
+  shm:
+    name: SharedMemoryConnector
+    extra:
+      connector_get_sleep_s: 0.005
+      connector_get_max_wait_first_chunk: 1000
+      connector_get_max_wait: 300
+      initial_codec_chunk_frames: 1
+      codec_chunk_frames: 15
+      codec_left_context_frames: 0
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 256
+    gpu_memory_utilization: 0.60
+    kv_cache_memory_bytes: 34359738368  # 32 GiB; leave room for codec state and graphs.
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 65536
+    max_model_len: 4096
+    devices: "0"
+    compilation_config:
+      # Capture mixed prefill/decode batches by total scheduled tokens.
+      cudagraph_capture_sizes: [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 160, 192, 224, 256, 320, 384, 448, 512]
+      max_cudagraph_capture_size: 512
+    output_connectors:
+      to_stage_1: shm
+    default_sampling_params:
+      temperature: 1.7
+      top_p: 0.8
+      top_k: 25
+      max_tokens: 4096
+      seed: 42
+      repetition_penalty: 1.0
+
+  - stage_id: 1
+    max_num_seqs: 256
+    gpu_memory_utilization: 0.30
+    enforce_eager: false
+    enable_prefix_caching: false
+    async_scheduling: true
+    compilation_config:
+      # Codec execution-batch buckets; frame counts come from the connector.
+      cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64, 128, 160, 192, 256]
+      cudagraph_num_of_warmups: 1
+    max_num_batched_tokens: 65536
+    max_model_len: 65536
+    devices: "0"
+    hf_overrides:
+      codec_attention_backend: triton
+    input_connectors:
+      from_stage_0: shm
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.0

vllm_omni/entrypoints/omni_base.py

@@ -339,6 +339,7 @@ class OmniBase(PDDisaggregationMixin):
         allow_delta_coercion: bool = False,
     ) -> Sequence[Any]:
         """Resolve request parameters; pipeline sampling constraints override caller values."""
+        normalized: Sequence[Any]
         if sampling_params_list is None:
             normalized = self.default_sampling_params_list
             # Set the output kind to delta since no params were specified
@@ -354,7 +355,9 @@ class OmniBase(PDDisaggregationMixin):
         if len(normalized) != self.num_stages:
             raise ValueError(f"Expected {self.num_stages} sampling params, got {len(normalized)}")
 
-        if sampling_params_list is not None:
+        # Streaming coercion may also change a constrained output kind in
+        # the defaults (for example, an internal TTS stage's FINAL_ONLY).
+        if sampling_params_list is not None or allow_delta_coercion:
             normalized = [
                 self._apply_sampling_constraints(params, constraints)
                 for params, constraints in zip(normalized, self.sampling_constraints_list, strict=True)
@@ -364,7 +367,7 @@ class OmniBase(PDDisaggregationMixin):
     @staticmethod
     def _get_sampling_constraints_list(stage_configs: Sequence[Any]) -> list[dict[str, Any]]:
         """Extract each stage's required sampling settings from runtime configs."""
-        constraints_list = []
+        constraints_list: list[dict[str, Any]] = []
         for stage_config in stage_configs:
             constraints = getattr(stage_config, "sampling_constraints", {})
             if not isinstance(constraints, Mapping):
@@ -398,7 +401,7 @@ class OmniBase(PDDisaggregationMixin):
         req_state = self.request_states.get(request_id)
         prom = getattr(self, "prom_metrics", None)
         metrics = getattr(req_state, "metrics", None)
-        if metrics is None or prom is None:
+        if req_state is None or metrics is None or prom is None:
             return
         if str(request_id) in metrics.e2e_done or getattr(req_state, "failure_recorded", False):
             return

vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py

@@ -6,7 +6,7 @@ Both variants share the same build/validate flow (``_build_moss_tts_params``
 handles each); they are registered under distinct model-type names.
 """
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 from vllm.inputs import tokens_input
 
@@ -61,7 +61,7 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         # the process-wide speaker cache.
         encoder = build_reference_encoder(
             self._get_moss_processor(),
-            variant=self._moss_variant,
+            variant=cast(str, self._moss_variant),
             speaker_cache=self._speaker_cache,
         )
         self._moss_ref_encoder = encoder
@@ -101,7 +101,7 @@ class _MossTTSAdapterBase(ARTTSAdapter):
 
         return await encode_request_references(
             self._get_moss_ref_encoder(),
-            request.ref_audio,
+            cast(str, request.ref_audio),
             request.ref_audio_2 if two_speaker else None,
             resolve_ref_audio=self._resolve_ref_audio,
             get_artifact_key=self._get_resolved_ref_audio_artifact_key,
@@ -179,15 +179,16 @@ class _MossTTSAdapterBase(ARTTSAdapter):
 
         v = self._moss_variant
 
+        params: dict[str, Any]
         # ---- Legacy nano path (unchanged) ----
         if v is None:  # moss_tts_nano
-            params: dict[str, Any] = {
+            params = {
                 "text": [request.input or ""],
                 "mode": ["voice_clone"],
             }
             if request.max_new_tokens is not None:
                 params["max_new_frames"] = [request.max_new_tokens]
-            wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio)
+            wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio))
             params["prompt_audio_array"] = [[wav_list, sr]]
             params["ref_audio_cache_key"] = cache_key
             return params
@@ -201,13 +202,13 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         # talker reads for short prompts; full Realtime support needs a
         # separate processor.from_module path which we don't wire here.
         if v == "realtime":
-            params: dict[str, Any] = {
+            params = {
                 "text": [request.input or ""],
                 "mode": ["voice_clone"],
             }
             if request.max_new_tokens is not None:
                 params["max_new_frames"] = [request.max_new_tokens]
-            wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio)
+            wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio))
             params["prompt_audio_array"] = [[wav_list, sr]]
             params["ref_audio_cache_key"] = cache_key
             return params
@@ -254,13 +255,20 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         # Build the unified-codes prompt: (L, 1+n_vq) where col 0 is text/special
         # tokens and cols 1..n_vq are the delay-pattern audio code grid (mostly
         # audio_pad_code outside the reference block).
-        user_msg = proc.build_user_message(**user_kwargs)
-        batch = proc(conversations=[[user_msg]], mode="generation")
+        if v == "local" and request.ref_text and request.ref_text.strip():
+            reference = user_kwargs.pop("reference")
+            user_kwargs["text"] = request.ref_text.strip() + " " + (request.input or "")
+            user_msg = proc.build_user_message(**user_kwargs)
+            assistant_msg = proc.build_assistant_message(audio_codes_list=reference)
+            batch = proc(conversations=[[user_msg, assistant_msg]], mode="continuation")
+        else:
+            user_msg = proc.build_user_message(**user_kwargs)
+            batch = proc(conversations=[[user_msg]], mode="generation")
         unified = batch["input_ids"][0]  # torch.LongTensor (L, 1+n_vq)
         text_ids: list[int] = unified[:, 0].tolist()
         audio_codes: torch.Tensor = unified[:, 1:].contiguous().to(torch.int64)
 
-        params: dict[str, Any] = {
+        params = {
             "prompt_token_ids": text_ids,
             "codes": {"ref": audio_codes},
         }

vllm_omni/entrypoints/utils.py

@@ -278,9 +278,9 @@ def maybe_coerce_to_message_type(params: SamplingParams, is_streaming: bool):
     if params.output_kind == target_type:
         return params
     elif is_streaming and params.output_kind == RequestOutputKind.FINAL_ONLY:
-        logger.warning("Request appears to be streaming, but got request type final only!")
+        logger.debug("Coercing FINAL_ONLY output to DELTA for streaming")
     elif not is_streaming and params.output_kind == RequestOutputKind.DELTA:
-        logger.warning("Request appears to not be streaming, but got request type delta!")
+        logger.debug("Coercing DELTA output to FINAL_ONLY for non-streaming")
 
     if not params.skip_clone:
         params = params.clone()

vllm_omni/metrics/stats.py

@@ -48,7 +48,7 @@ class StageRequestStats:
     finish_reason: str | None = None
     request_id: str | None = None
     postprocess_time_ms: float = 0.0
-    diffusion_metrics: dict[str, float] = None
+    diffusion_metrics: dict[str, float] | None = None
     audio_generated_frames: int = 0
     audio_sample_rate: int = 0
     audio_duration_s: float = 0.0
@@ -181,11 +181,11 @@ class OrchestratorAggregator:
         self.e2e_total_ms = 0.0
         self.e2e_total_tokens = 0
         self.e2e_count = 0
-        self.e2e_done = set()
+        self.e2e_done: set[str] = set()
         self.wall_start_ts = float(wall_start_ts)
         self.last_finish_ts = float(wall_start_ts)
-        self.stage_first_ts = [None for _ in range(self.num_stages)]
-        self.stage_last_ts = [None for _ in range(self.num_stages)]
+        self.stage_first_ts: list[float | None] = [None for _ in range(self.num_stages)]
+        self.stage_last_ts: list[float | None] = [None for _ in range(self.num_stages)]
         self.accumulated_gen_time_ms: defaultdict[str, defaultdict[int, float]] = defaultdict(
             lambda: defaultdict(float)
         )  # {request_id: {stage_id:accumulated_gen_time_ms}}
@@ -618,6 +618,7 @@ class OrchestratorAggregator:
         final_output_type: str | None = None,
     ) -> None:
         stats = self._as_stage_request_stats(stage_id, req_id, metrics, final_output_type)
+        assert stats.stage_id is not None
         self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_out)
         if stats.stage_id == 0:
             self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_in)
@@ -774,10 +775,8 @@ class OrchestratorAggregator:
             final_stage_id_map = self.final_stage_id_for_e2e
 
         stage_wall_time_ms = [
-            ((self.stage_last_ts[i] - self.stage_first_ts[i]) * 1000.0)
-            if (self.stage_first_ts[i] is not None and self.stage_last_ts[i] is not None)
-            else 0.0
-            for i in range(self.num_stages)
+            ((last - first) * 1000.0) if first is not None and last is not None else 0.0
+            for first, last in zip(self.stage_first_ts, self.stage_last_ts, strict=True)
         ]
 
         overall_summary = {

vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -801,7 +804,16 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule):
         else:
             attn_bias = None
 
-        x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
+        streaming_attention = getattr(self, "_streaming_attention", None)
+        if (
+            streaming_attention is not None
+            and attn_bias is not None
+            and q.dtype == torch.bfloat16
+            and q.shape[-1] == 64
+        ):
+            x = streaming_attention(q, k, v, attn_bias)
+        else:
+            x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
         x = x.transpose(1, 2).reshape(B, T, self.embed_dim)
         x = apply_weights_per_step(self.out_projs, self.weights_per_step_schedule, x, offset_cpu)
 

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py

@@ -145,15 +145,22 @@ class _MossCodecStreamSession:
         terminal_slots = terminal_slots or set()
         if not terminal_slots.issubset(slot_codes):
             raise ValueError("terminal_slots must be a subset of the decode batch slots.")
-        step_lengths = {int(codes.shape[1]) for codes in slot_codes.values()}
-        if len(step_lengths) != 1:
-            raise ValueError(f"MOSS codec streaming step needs uniform T, got {sorted(step_lengths)}")
-        (step_t,) = step_lengths
+        step_t = max(int(codes.shape[1]) for codes in slot_codes.values())
+        if any(int(codes.shape[1]) != step_t and slot not in terminal_slots for slot, codes in slot_codes.items()):
+            raise ValueError("Only terminal codec rows may be padded to a larger step length")
         slots = list(slot_codes)
         if any(slot not in self._leased_slots for slot in slots):
             raise RuntimeError(f"Streaming decode references an unleased state slot: {slots}")
         codes_step = torch.stack(
-            [slot_codes[slot].to(device=self._device, dtype=torch.long) for slot in slots],
+            [
+                torch.nn.functional.pad(
+                    slot_codes[slot].to(device=self._device, dtype=torch.long),
+                    (0, step_t - int(slot_codes[slot].shape[1])),
+                )
+                if int(slot_codes[slot].shape[1]) != step_t
+                else slot_codes[slot].to(device=self._device, dtype=torch.long)
+                for slot in slots
+            ],
             dim=1,
         )
         state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long)
@@ -197,9 +204,9 @@ class _MossCodecStreamSession:
             self._reset_slot_ids(terminal_slot_ids)
 
         audio = audio_tensor.detach().to("cpu", torch.float32)
-        audio_length = step_t * self._samples_per_frame
         out: dict[int, torch.Tensor] = {}
         for row, slot in enumerate(slots):
+            audio_length = int(slot_codes[slot].shape[1]) * self._samples_per_frame
             out[slot] = audio[row, ..., :audio_length].contiguous()
         return out
 
@@ -518,6 +525,8 @@ class MossTTSCodecDecoder(nn.Module):
         outputs: dict[int, torch.Tensor] = {}
         grouped: dict[int, list[tuple[int, str, int, torch.Tensor, bool]]] = {}
         max_step_frames = max(1, int(self._stream_max_step_frames))
+        graph = session._cudagraph_wrapper
+        coalesce_tails = graph is not None and max(graph.batch_sizes, default=0) >= self._stream_state_capacity
 
         for output_index, request_id, codes_nq_t, finished in items:
             slot = self._stream_req_slots.get(request_id)
@@ -548,9 +557,13 @@ class MossTTSCodecDecoder(nn.Module):
                     )
                 continue
 
-            grouped.setdefault(int(codes_nq_t.shape[1]), []).append(
-                (output_index, request_id, slot, codes_nq_t, finished)
-            )
+            frame_count = int(codes_nq_t.shape[1])
+            # Only combine tails that already used the same padded graph.
+            # Preserve exact-size graphs (including T=1) and the eager path.
+            group_frames = frame_count
+            if coalesce_tails and finished:
+                group_frames = graph._select_frame_size(frame_count, allow_padding=True) or frame_count
+            grouped.setdefault(group_frames, []).append((output_index, request_id, slot, codes_nq_t, finished))
 
         for group in grouped.values():
             plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group}
@@ -756,6 +769,17 @@ class MossTTSCodecDecoder(nn.Module):
         codec.eval()
         if device.type != "cpu":
             codec.decoder.to(dtype=torch.bfloat16)
+        attention_backend = getattr(self.vllm_config.model_config.hf_config, "codec_attention_backend", "sdpa")
+        if attention_backend != "sdpa":
+            if attention_backend != "triton" or device.type != "cuda":
+                raise ValueError(f"Unsupported codec attention backend/device: {attention_backend}/{device.type}")
+            from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import MossAudioTokenizerMultiheadAttention
+            from vllm_omni.model_executor.models.moss_tts.streaming_attention import masked_attention
+
+            for module in codec.decoder.modules():
+                if isinstance(module, MossAudioTokenizerMultiheadAttention):
+                    module._streaming_attention = masked_attention
+            logger.info("Enabled Triton masked attention for the streaming codec decoder")
         build_decode_lut = getattr(codec.quantizer, "build_decode_lut", None)
         if callable(build_decode_lut):
             lut_dtype = torch.bfloat16 if device.type == "cuda" else torch.float32

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License").
@@ -101,8 +103,13 @@ class _MossTTSLocalAttention(nn.Module):
         odd = x[..., 1::2]
         return torch.stack((-odd, even), dim=-1).reshape_as(x)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        """``hidden_states``: ``(B, S, H)``. Re-prefills with a fresh causal mask over ``[0, S)`` every call."""
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
+        position: int = 0,
+    ) -> torch.Tensor:
+        """Run a causal prefix, or one new token using frame-local K/V."""
         batch_size, seq_len, _ = hidden_states.shape
         qkv = self.c_attn(hidden_states)
         query, key, value = qkv.split(self.embed_dim, dim=-1)
@@ -110,14 +117,22 @@ class _MossTTSLocalAttention(nn.Module):
         key = key.view(batch_size, seq_len, self.n_head, self.head_dim)
         value = value.view(batch_size, seq_len, self.n_head, self.head_dim)
 
-        cos, sin = self._rope_cos_sin(seq_len, hidden_states.device, hidden_states.dtype)
+        cos, sin = self._rope_cos_sin(position + seq_len, hidden_states.device, hidden_states.dtype)
+        cos, sin = cos[:, position:], sin[:, position:]
         query = query * cos + self._rotate_half(query) * sin
         key = key * cos + self._rotate_half(key) * sin
 
         query = query.transpose(1, 2)
         key = key.transpose(1, 2)
         value = value.transpose(1, 2)
-        attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=True)
+        if kv_cache is not None:
+            assert seq_len == 1, "Cached Local Depth execution consumes one token at a time"
+            k_cache, v_cache = kv_cache
+            k_cache[:, :, position : position + 1].copy_(key)
+            v_cache[:, :, position : position + 1].copy_(value)
+            key = k_cache[:, :, : position + 1]
+            value = v_cache[:, :, : position + 1]
+        attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=kv_cache is None)
         attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim)
         return self.c_proj(attn_output)
 
@@ -141,8 +156,13 @@ class _MossTTSLocalBlock(nn.Module):
         self.ln_2 = nn.LayerNorm(hidden_size, eps=eps)
         self.mlp = _MossTTSLocalMLP(hidden_size, inner_size)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = hidden_states + self.attn(self.ln_1(hidden_states))
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
+        position: int = 0,
+    ) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(self.ln_1(hidden_states), kv_cache, position)
         hidden_states = hidden_states + self.mlp(self.ln_2(hidden_states))
         return hidden_states
 
@@ -156,11 +176,9 @@ class MossTTSLocalDepthTransformer(nn.Module):
         and codebook-0's head (``audio_lm_heads[0]``) simultaneously.
       - codebooks 1..n_vq-1 are sampled sequentially: each sampled code is
         re-embedded (``audio_embeddings[c]``) and appended as the next
-        position, re-prefilling the block over the growing (<=n_vq) sequence
-        with a fresh causal mask each call -- mathematically identical to
-        incremental KV-cache decoding since attention is strictly causal and
-        only the last position is ever read, but avoids any cache plumbing
-        given the trivially short sequence length.
+        position. K/V for earlier positions are reused within the frame, so
+        each position's projections and MLP run once. A new frame overwrites
+        every cache position before reading it, including CUDA Graph replay.
     """
 
     def __init__(self, gpt2_config, hidden_size: int | None = None) -> None:
@@ -174,11 +192,13 @@ class MossTTSLocalDepthTransformer(nn.Module):
         self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps)
         self._compiled_forward_prefix = None
 
-    def _forward_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor:
-        hidden_states = seq_embeds
-        for block in self.h:
-            hidden_states = block(hidden_states)
-        return self.ln_f(hidden_states)
+    def _forward_prefix(
+        self,
+        seq_embeds: torch.Tensor,
+        kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
+        position: int = 0,
+    ) -> torch.Tensor:
+        return self.ln_f(self.h[0](seq_embeds, kv_cache, position))
 
     def setup_compile(self) -> None:
         if self._compiled_forward_prefix is not None:
@@ -189,14 +209,18 @@ class MossTTSLocalDepthTransformer(nn.Module):
             return
         self._compiled_forward_prefix = torch.compile(
             self._forward_prefix,
-            dynamic=False,
+            # Share batch/position shapes instead of exhausting Dynamo's
+            # recompilation budget while warming request-size graph buckets.
+            dynamic=True,
             options={"epilogue_fusion": False},
         )
-        logger.info("MOSS-TTS local depth prefix enabled with torch.compile")
+        logger.info("MOSS-TTS local depth frame-local KV execution enabled with torch.compile")
 
-    def _run_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor:
+    def _run_prefix(
+        self, seq_embeds: torch.Tensor, kv_cache: tuple[torch.Tensor, torch.Tensor], position: int
+    ) -> torch.Tensor:
         forward_prefix = self._compiled_forward_prefix or self._forward_prefix
-        return forward_prefix(seq_embeds)
+        return forward_prefix(seq_embeds, kv_cache, position)
 
     @torch.no_grad()
     def generate_frame(
@@ -240,9 +264,13 @@ class MossTTSLocalDepthTransformer(nn.Module):
         for block in self.h:
             block.attn.prepare_rope_cache(n_vq, backbone_last_hidden.device, dtype)
 
-        embeds = backbone_last_hidden.new_zeros((batch_size, n_vq, self.hidden_size), dtype=dtype)
-        embeds[:, 0, :] = backbone_last_hidden.to(dtype)
-        hidden = self._run_prefix(embeds[:, :1, :])
+        attn = self.h[0].attn
+        cache_shape = (batch_size, attn.n_head, n_vq, attn.head_dim)
+        kv_cache = (
+            backbone_last_hidden.new_empty(cache_shape, dtype=dtype),
+            backbone_last_hidden.new_empty(cache_shape, dtype=dtype),
+        )
+        hidden = self._run_prefix(backbone_last_hidden[:, None, :].to(dtype), kv_cache, 0)
         local_hidden = hidden[:, 0, :]
 
         binary_logits = local_text_lm_head(local_hidden).float()
@@ -292,9 +320,9 @@ class MossTTSLocalDepthTransformer(nn.Module):
             codes[:, channel_index] = channel_token
 
             if channel_index + 1 < n_vq:
-                embeds[:, channel_index + 1, :] = audio_embeddings[channel_index](channel_token).to(dtype)
-                hidden = self._run_prefix(embeds[:, : channel_index + 2, :])
-                local_hidden = hidden[:, channel_index + 1, :]
+                embeds = audio_embeddings[channel_index](channel_token).to(dtype)[:, None, :]
+                hidden = self._run_prefix(embeds, kv_cache, channel_index + 1)
+                local_hidden = hidden[:, 0, :]
 
         return should_continue, codes
 

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License").
@@ -1671,20 +1674,19 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
 
         self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
         self._batch_state_spans = kwargs.get("request_token_spans")
-        should_values: list[torch.Tensor] = []
-        have_mtp_output = False
+        should_codes: list[torch.Tensor] = []
+        should_indices: list[int] = []
         per_req_codes: list[torch.Tensor] = []
         have_codes = False
-        for info in info_dicts:
+        for index, info in enumerate(info_dicts):
             audio_codes = (info.get("audio_codes", {}) or {}) if isinstance(info, dict) else {}
             current = audio_codes.pop("current", None)
             if isinstance(current, torch.Tensor) and current.numel() > 0:
                 codes = current.to(device=hidden.device, dtype=torch.long)
                 if codes.dim() == 1:
                     codes = codes.unsqueeze(0)
-                should_continue = codes.ne(self.audio_pad_token_id).any(dim=-1)
-                should_values.append(should_continue.reshape(-1)[:1])
-                have_mtp_output = True
+                should_codes.append(codes[:1])
+                should_indices.append(index)
                 # Do not boolean-index a CUDA tensor here.  Even for the
                 # common single-row case, ``codes[should_continue]`` must
                 # discover a dynamic output size on the host and therefore
@@ -1694,9 +1696,19 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
                 per_req_codes.append(codes)
                 have_codes = True
             else:
-                should_values.append(torch.ones((1,), device=hidden.device, dtype=torch.bool))
                 per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
-        self._batch_should_continue = torch.cat(should_values, dim=0) if have_mtp_output and should_values else None
+        self._batch_should_continue = None
+        if should_codes:
+            active = torch.cat(should_codes, dim=0).ne(self.audio_pad_token_id).any(dim=-1)
+            if len(should_codes) == len(info_dicts):
+                self._batch_should_continue = active
+            else:
+                self._batch_should_continue = torch.ones(len(info_dicts), device=hidden.device, dtype=torch.bool)
+                # Keep mixed prefill/decode output packing asynchronous.
+                indices = torch.tensor(should_indices, device="cpu", dtype=torch.long).to(
+                    hidden.device, non_blocking=True
+                )
+                self._batch_should_continue.index_copy_(0, indices, active)
 
         if not have_codes:
             return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})

vllm_omni/model_executor/models/moss_tts/pipeline.py

@@ -1,8 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License").
 """Pipeline topology for all MOSS-TTS variants (2-stage: talker → codec)."""
 
+from vllm.sampling_params import RequestOutputKind
+
 from vllm_omni.config.stage_config import (
     PipelineConfig,
     StageExecutionType,
@@ -100,6 +105,10 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig(
             sampling_constraints={
                 "detokenize": False,
                 "stop_token_ids": [151645],
+                # The worker connector streams codes directly to the codec.
+                # This internal stage only needs to publish its terminal
+                # result; codec audio output remains incremental.
+                "output_kind": RequestOutputKind.FINAL_ONLY,
             },
         ),
         StagePipelineConfig(

vllm_omni/model_executor/models/moss_tts/streaming_attention.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Masked streaming attention preserving ring-cache position semantics."""
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _attention(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    mask_ptr,
+    out_ptr,
+    qs0: tl.constexpr,
+    qs1: tl.constexpr,
+    qs2: tl.constexpr,
+    qs3: tl.constexpr,
+    ks0: tl.constexpr,
+    ks1: tl.constexpr,
+    ks2: tl.constexpr,
+    ks3: tl.constexpr,
+    vs0: tl.constexpr,
+    vs1: tl.constexpr,
+    vs2: tl.constexpr,
+    vs3: tl.constexpr,
+    ms0: tl.constexpr,
+    ms2: tl.constexpr,
+    ms3: tl.constexpr,
+    num_heads: tl.constexpr,
+    q_len: tl.constexpr,
+    kv_len: tl.constexpr,
+    head_dim: tl.constexpr,
+    block_m: tl.constexpr,
+    block_n: tl.constexpr,
+):
+    rows = tl.program_id(0) * block_m + tl.arange(0, block_m)
+    bh = tl.program_id(1)
+    batch, head = bh // num_heads, bh % num_heads
+    dims = tl.arange(0, head_dim)
+    q = tl.load(
+        q_ptr + batch * qs0 + head * qs1 + rows[:, None] * qs2 + dims[None, :] * qs3,
+        mask=rows[:, None] < q_len,
+        other=0,
+    )
+    acc = tl.full((block_m, head_dim), 0.0, tl.float32)
+    maximum = tl.full((block_m,), -float("inf"), tl.float32)
+    denominator = tl.full((block_m,), 0.0, tl.float32)
+    for first in range(tl.cdiv(kv_len, block_n)):
+        cols = first * block_n + tl.arange(0, block_n)
+        k = tl.load(
+            k_ptr + batch * ks0 + head * ks1 + cols[None, :] * ks2 + dims[:, None] * ks3,
+            mask=cols[None, :] < kv_len,
+            other=0,
+        )
+        allowed = tl.load(
+            mask_ptr + batch * ms0 + rows[:, None] * ms2 + cols[None, :] * ms3,
+            mask=(rows[:, None] < q_len) & (cols[None, :] < kv_len),
+            other=0,
+        )
+        score = tl.dot(q, k).to(tl.float32) * (head_dim**-0.5)
+        score = tl.where(allowed & (cols[None, :] < kv_len), score, -float("inf"))
+        next_max = tl.maximum(maximum, tl.max(score, axis=1))
+        safe_max = tl.where(next_max == -float("inf"), 0.0, next_max)
+        rescale = tl.exp(maximum - safe_max)
+        p = tl.exp(score - safe_max[:, None])
+        denominator = denominator * rescale + tl.sum(p, axis=1)
+        v = tl.load(
+            v_ptr + batch * vs0 + head * vs1 + cols[:, None] * vs2 + dims[None, :] * vs3,
+            mask=cols[:, None] < kv_len,
+            other=0,
+        )
+        acc = acc * rescale[:, None] + tl.dot(p.to(v.dtype), v)
+        maximum = next_max
+    out = acc / tl.where(denominator > 0, denominator, 1.0)[:, None]
+    tl.store(
+        out_ptr + ((batch * num_heads + head) * q_len + rows[:, None]) * head_dim + dims[None, :],
+        out,
+        mask=rows[:, None] < q_len,
+    )
+
+
+@torch.library.custom_op("moss_streaming::masked_attention", mutates_args=())
+def masked_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    b, h, t, d = q.shape
+    out = torch.empty(q.shape, device=q.device, dtype=q.dtype)
+    # Small query tiles expose parallelism for short streaming chunks. Larger
+    # tiles reuse K/V for steady chunks. Keep 64-key softmax blocks to retain
+    # the same accumulation granularity across both paths.
+    bm = 16 if t <= 32 else 64
+    num_warps = 8 if k.shape[2] <= 256 else 4
+    _attention[(triton.cdiv(t, bm), b * h)](
+        q,
+        k,
+        v,
+        mask,
+        out,
+        *q.stride(),
+        *k.stride(),
+        *v.stride(),
+        mask.stride(0),
+        mask.stride(2),
+        mask.stride(3),
+        h,
+        t,
+        k.shape[2],
+        d,
+        bm,
+        64,
+        num_warps=num_warps,
+        num_stages=2,
+    )
+    return out
+
+
+@masked_attention.register_fake
+def _(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    return torch.empty(q.shape, device=q.device, dtype=q.dtype)

vllm_omni/outputs/__init__.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -52,7 +55,7 @@ class OmniModelRunnerOutput(ModelRunnerOutput):
     """
 
     multimodal_outputs: list[dict[str, object]] | None = None
-    inter_stage_outputs: list[dict[str, object]] | None = None
+    inter_stage_outputs: list[dict[str, Any] | None] | None = None
     # IDs of requests whose KV cache has been extracted from GPU/NPU to CPU.
     # The Scheduler can safely free the block tables for these requests.
     kv_extracted_req_ids: list[str] | None = None

vllm_omni/worker/gpu_ar_model_runner.py

@@ -15,7 +15,7 @@ from collections.abc import Callable, Sequence
 from contextlib import nullcontext
 from copy import copy
 from dataclasses import replace
-from typing import Any, NamedTuple
+from typing import Any, NamedTuple, cast
 
 import numpy as np
 import torch
@@ -87,14 +87,39 @@ def _clone_cuda_tensor_payload(value: Any, sources: list[torch.Tensor]) -> Any:
         return value.detach().clone()
     if isinstance(value, dict):
         return {k: _clone_cuda_tensor_payload(v, sources) for k, v in value.items()}
-    if isinstance(value, list):
-        return [_clone_cuda_tensor_payload(v, sources) for v in value]
-    if isinstance(value, tuple):
-        return tuple(_clone_cuda_tensor_payload(v, sources) for v in value)
+    if isinstance(value, (list, tuple)):
+        if (
+            value
+            and all(
+                isinstance(v, torch.Tensor)
+                and v.is_cuda
+                and v.dtype == value[0].dtype
+                and v.device == value[0].device
+                and v.layout == torch.strided
+                and v.is_contiguous()
+                for v in value
+            )
+            # Keep frame/state packing bounded; large payloads use individual snapshots.
+            and sum(v.numel() * v.element_size() for v in value) <= 1024 * 1024
+        ):
+            packed = torch.cat([v.detach().reshape(-1) for v in value])
+            sources.append(packed)
+            return _PackedTensorPayload(packed, [v.shape for v in value], isinstance(value, tuple))
+        items = [_clone_cuda_tensor_payload(v, sources) for v in value]
+        return tuple(items) if isinstance(value, tuple) else items
     return value
 
 
 def _copy_tensor_payload_to_cpu(value: Any, pin_memory: bool) -> Any:
+    if isinstance(value, _PackedTensorPayload):
+        flat = _copy_tensor_payload_to_cpu(value.tensor, pin_memory)
+        items = []
+        offset = 0
+        for shape in value.shapes:
+            count = shape.numel()
+            items.append(flat[offset : offset + count].view(shape))
+            offset += count
+        return tuple(items) if value.is_tuple else items
     if isinstance(value, torch.Tensor):
         if value.device.type != "cuda":
             return value
@@ -110,6 +135,12 @@ def _copy_tensor_payload_to_cpu(value: Any, pin_memory: bool) -> Any:
     return value
 
 
+class _PackedTensorPayload(NamedTuple):
+    tensor: torch.Tensor
+    shapes: list[torch.Size]
+    is_tuple: bool
+
+
 class _AsyncCPUPayloadSnapshot:
     def __init__(
         self,
@@ -181,7 +212,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
         if kwargs:
             raise TypeError(f"Unexpected OmniAsyncGPUModelRunnerOutput kwargs: {sorted(kwargs)}")
 
-        self._model_runner_output = None
+        self._model_runner_output: OmniModelRunnerOutput | None = None
         self._invalid_req_indices = invalid_req_indices
 
         self.async_copy_ready_event = torch.Event()
@@ -211,7 +242,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
             self._num_nans_cpu = self._num_nans.to("cpu", non_blocking=True) if self._num_nans is not None else None
             self.async_copy_ready_event.record()
 
-        self._model_runner_output_builder = model_runner_output_builder
+        self._model_runner_output_builder: Callable[[], OmniModelRunnerOutput] | None = model_runner_output_builder
         self._background_exception: BaseException | None = None
         self._background_thread: threading.Thread | None = None
         self._cuda_device = cuda_device
@@ -226,6 +257,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
         if self._model_runner_output is not None:
             return
         with record_function_or_nullcontext("omni_async_output:get_output/build_model_runner_output"):
+            assert self._model_runner_output_builder is not None
             self._model_runner_output = self._model_runner_output_builder()
         self._model_runner_output_builder = None
 
@@ -321,6 +353,9 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
     outputs per request while keeping Async output semantics.
     """
 
+    execute_model_state: ExecuteModelState | None
+    kv_extracted_req_ids: list[str] | None
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
@@ -436,6 +471,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         if not isinstance(info, dict):
             return None
         val = info.get("omni_final_stage_id")
+        if val is None:
+            return None
         try:
             return int(val)
         except (TypeError, ValueError):
@@ -1485,7 +1522,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                     )
             selected_indices.append(idx)
             selected_req_ids.append(req_id)
-            selected_req_infos.append(req_info)
+            selected_req_infos.append(cast(dict[str, Any], req_info))
 
         if not selected_indices:
             return multimodal_outputs
@@ -1552,7 +1589,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
 
     def _build_multimodal_outputs(
         self,
-        per_req_payloads: list[dict[str, object] | None] | None,
+        per_req_payloads: Sequence[dict[str, object] | None] | None,
     ) -> list[dict[str, torch.Tensor] | None] | None:
         """Build per-request multimodal output payloads (dedicated channel).
 
@@ -1648,8 +1685,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
 
     def _build_omni_step_outputs(
         self,
-        pooler_inter: list[dict[str, object]],
-        pooler_client: list[dict[str, object]],
+        pooler_inter: Sequence[dict[str, object] | None] | None,
+        pooler_client: Sequence[dict[str, object] | None] | None,
         *,
         defer_full_payload_d2h: bool,
     ) -> tuple[
@@ -1935,6 +1972,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                     )
                     pooler_output.append(flatten_payload(payload))
 
+        pooler_inter: Sequence[dict[str, object] | None] | None
+        pooler_client: Sequence[dict[str, object] | None] | None
         pooler_output = pooler_output or []
         if self._async_chunk and stage_sends_async_output(self.model_config):
             pooler_inter, pooler_client = partition_payload_list(pooler_output)

vllm_omni/worker/gpu_model_runner.py

@@ -78,6 +78,8 @@ def _filter_mrope_kwargs_for_model(model: object, kwargs: dict[str, Any]) -> dic
 
 
 class OmniGPUModelRunner(GPUModelRunner):
+    intermediate_tensors: IntermediateTensors | None
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.model_intermediate_buffer: dict[str, dict[str, Any]] = {}
@@ -842,7 +844,9 @@ class OmniGPUModelRunner(GPUModelRunner):
             return None
 
     @torch.inference_mode()
-    def extract_multimodal_outputs(self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput) -> dict:
+    def extract_multimodal_outputs(
+        self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput
+    ) -> tuple[Any, Any]:
         if (
             hasattr(self.model, "have_multimodal_outputs")
             and self.model.have_multimodal_outputs
@@ -1407,14 +1411,14 @@ class OmniGPUModelRunner(GPUModelRunner):
     def _process_additional_information_updates(
         self,
         hidden_states: torch.Tensor,
-        multimodal_outputs: object,
+        multimodal_outputs: Any,
         num_scheduled_tokens_np: np.ndarray,
         scheduler_output: "SchedulerOutput",
         combined_hidden_states: dict[str, torch.Tensor] | None = None,
-        combined_multimodal_outputs: dict[str, object] | None = None,
+        combined_multimodal_outputs: dict[str, Any] | None = None,
         req_ids_filter: set[str] | None = None,
         req_ids: list[str] | None = None,
-        query_start_loc_cpu: object | None = None,
+        query_start_loc_cpu: Any = None,
     ) -> None:
         """Process model-provided per-request updates and merge into model_intermediate_buffer."""
         req_ids = req_ids if req_ids is not None else self.input_batch.req_ids
@@ -1476,9 +1480,9 @@ class OmniGPUModelRunner(GPUModelRunner):
     def _collect_additional_information_for_prefill(
         self,
         num_scheduled_tokens_np: np.ndarray,
-    ) -> dict[str, dict]:
+    ) -> None:
         """Overlay per-request prompt_embeds for the prefill portion and collect
-        additional_information slices for this step. Returns a map req_id -> dict."""
+        additional_information slices for this step."""
         for req_index, req_id in enumerate(self.input_batch.req_ids):
             req_state = self.requests[req_id]
             pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
@@ -1707,9 +1711,9 @@ class OmniGPUModelRunner(GPUModelRunner):
 
             # Overlay custom prompt_embeds per request for the prompt portion;
             # collect additional_information (tensor/list) for prefill portion only
-            decode_req_ids = []
-            decode_start_offsets = []
-            decode_batch_items = []
+            decode_req_ids: list[str] = []
+            decode_start_offsets: list[int] = []
+            decode_batch_items: list[tuple[str, int, dict[str, Any]]] = []
             batch_decode_preprocess = getattr(self.model, "preprocess_decode_batch", None)
 
             def flush_decode_batch() -> None:
@@ -1717,6 +1721,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                 if not decode_batch_items:
                     return
 
+                assert callable(batch_decode_preprocess)
                 req_ids_b = [item[0] for item in decode_batch_items]
                 start_offsets_b = [item[1] for item in decode_batch_items]
                 req_infos_b = [item[2] for item in decode_batch_items]
@@ -1761,7 +1766,7 @@ class OmniGPUModelRunner(GPUModelRunner):
 
                 # mimo-audio check
                 req_state = self.requests.get(req_id)
-                req_infos = self._maybe_attach_mimo_audio_req_infos(req_state, req_infos, req_id)
+                req_infos = cast(dict[str, Any], self._maybe_attach_mimo_audio_req_infos(req_state, req_infos, req_id))
 
                 start_offset = int(self.query_start_loc.cpu[req_index])
                 sched_tokens = int(num_scheduled_tokens_np[req_index])
@@ -1968,11 +1973,29 @@ class OmniGPUModelRunner(GPUModelRunner):
         if start_offsets is None:
             id_to_index = self.input_batch.req_id_to_index
             start_offsets = [int(self.query_start_loc.cpu[id_to_index[req_id]]) for req_id in decode_req_ids]
-        for idx, (req_id, start_offset) in enumerate(zip(decode_req_ids, start_offsets, strict=True)):
-            inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
-            if code_predictor_codes is not None:
-                update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
-                self._update_intermediate_buffer(req_id, update_dict)
+        if start_offsets == list(range(decode_batch_size)):
+            inputs_embeds[:decode_batch_size].copy_(req_embeds[:decode_batch_size])
+        else:
+            # A device tensor constructor waits for the preceding MTP graph.
+            # Enqueue the small host index transfer without synchronizing it.
+            offsets = torch.tensor(start_offsets, device="cpu", dtype=torch.long).to(
+                inputs_embeds.device, non_blocking=True
+            )
+            inputs_embeds.index_copy_(0, offsets, req_embeds[:decode_batch_size])
+        if code_predictor_codes is not None:
+            if out_key in getattr(self.model, "gpu_resident_buffer_keys", set()):
+                # One owned batch snapshot protects all rows from graph reuse.
+                owned_codes = code_predictor_codes[:decode_batch_size].detach().clone()
+                for req_id, row in zip(decode_req_ids, owned_codes.split(1), strict=True):
+                    req_state = self.requests.get(req_id)
+                    if req_state is not None:
+                        existing = self.model_intermediate_buffer.setdefault(req_id, {})
+                        existing.setdefault(out_key[0], {})[out_key[1]] = row
+                        req_state.additional_information_cpu = existing
+            else:
+                for idx, req_id in enumerate(decode_req_ids):
+                    update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
+                    self._update_intermediate_buffer(req_id, update_dict)
 
     def _model_forward(
         self,