[Model] Optimize MOSS-TTS Local batched execution and streaming codec

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

@@ -425,6 +425,40 @@ 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 serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \
+    --omni --trust-remote-code \
+    --deploy-config vllm_omni/deploy/moss_tts_local_h200.yaml
+```
+
+Run from the repository root and select an available physical GPU with
+`CUDA_VISIBLE_DEVICES`. 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.

tests/model_executor/models/moss_tts/test_streaming_attention.py

@@ -0,0 +1,61 @@
+# 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", [(1, 20, 1, 125), (4, 20, 15, 125), (1, 12, 480, 400), (4, 12, 32, 250)]
+)
+@pytest.mark.parametrize("empty_mask", [False, True])
+def test_masked_attention_matches_sdpa_with_strides_and_invalid_rows(batch, heads, queries, keys, empty_mask):
+    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()
+    if empty_mask:
+        mask.zero_()
+    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)
+    if empty_mask:
+        assert actual.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,118 @@
+# 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 _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=3, n_vq=2, vllm_config=None)
+
+
+@pytest.mark.parametrize("tail_length", [1, 3, 7, 14])
+def test_terminal_padding_preserves_audio_live_state_and_reused_slots(tail_length):
+    batched, separate = session(), session()
+    for s in (batched, separate):
+        assert [s.acquire(), s.acquire()] == [0, 1]
+        s.step({0: torch.ones(2, 15), 1: torch.full((2, 15), 3)})
+    short = torch.arange(2 * tail_length).reshape(2, tail_length)
+    long = torch.full((2, 15), 5)
+    expected_short = separate.step({0: short}, terminal_slots={0})[0]
+    expected_long = separate.step({1: long})[1]
+    actual = batched.step({0: short, 1: long}, terminal_slots={0})
+    torch.testing.assert_close(actual[0], expected_short, rtol=0, atol=0)
+    torch.testing.assert_close(actual[1], expected_long, rtol=0, atol=0)
+    for s in (batched, separate):
+        s.release(0, state_already_reset=True)
+        assert s.acquire() == 0
+    next_step = {0: torch.ones(2, 1), 1: torch.ones(2, 1)}
+    a, b = batched.step(next_step), separate.step(next_step)
+    for slot in (0, 1):
+        torch.testing.assert_close(a[slot], b[slot], rtol=0, atol=0)
+
+
+def test_nonterminal_row_cannot_silently_advance_through_padding():
+    s = session()
+    s.acquire()
+    s.acquire()
+    with pytest.raises(ValueError, match="Only terminal"):
+        s.step({0: torch.ones(2, 3), 1: torch.ones(2, 15)})
+    assert s._codec.state.count_nonzero() == 0
+
+
+@pytest.mark.parametrize("graph_capacity, expected_calls", [(None, 5), (4, 5), (8, 3)])
+def test_only_existing_compatible_graph_buckets_are_coalesced(graph_capacity, expected_calls):
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder
+
+    calls = []
+    leased: list[int] = []
+
+    def acquire():
+        slot = len(leased)
+        leased.append(slot)
+        return slot
+
+    def step(plan, *, terminal_slots):
+        calls.append((dict(plan), set(terminal_slots)))
+        return {slot: torch.zeros(1, codes.shape[1]) for slot, codes in plan.items()}
+
+    graph = (
+        None
+        if graph_capacity is None
+        else SimpleNamespace(
+            batch_sizes=[1, graph_capacity],
+            frame_sizes=[1, 8, 15],
+            _select_frame_size=lambda size, allow_padding: next((b for b in [1, 8, 15] if b >= size), None),
+        )
+    )
+    fake_session = SimpleNamespace(
+        _cudagraph_wrapper=graph,
+        acquire=acquire,
+        step=step,
+        release=lambda slot, **kwargs: None,
+    )
+    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: fake_session
+    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 len(calls) == expected_calls
+    assert [result[i].shape[-1] for i in range(5)] == lengths
+    if graph_capacity == 8:
+        assert set(calls[0][0]) == {0, 1, 2}
+        assert calls[0][1] == {0, 1}
+        assert set(calls[1][0]) == {3}  # Preserve the exact one-frame graph.

tests/worker/test_batched_omni_output.py

@@ -0,0 +1,94 @@
+# 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.cpu
+@pytest.mark.parametrize("kinds", [("emit", "stop", "prefill"), ("prefill", "emit"), ("stop",), ("prefill",)])
+def test_local_stop_flags_and_output_rows_stay_request_aligned(kinds):
+    model = object.__new__(MossTTSLocalTalkerForGeneration)
+    torch.nn.Module.__init__(model)
+    model.n_vq = 12
+    model.audio_pad_token_id = 1024
+    infos = []
+    originals = []
+    for index, kind in enumerate(kinds):
+        codes = torch.full((1, 12), 1024 if kind == "stop" else index, dtype=torch.long)
+        originals.append(codes)
+        infos.append({} if kind == "prefill" else {"audio_codes": {"current": codes}})
+    result = model.make_omni_output(torch.ones((len(kinds), 4)), runtime_additional_information=infos)
+    if all(kind == "prefill" for kind in kinds):
+        assert model._batch_should_continue is None
+        assert result.multimodal_outputs == {}
+        return
+    assert model._batch_should_continue.tolist() == [kind != "stop" for kind in kinds]
+    rows = result.multimodal_outputs["codes"]["audio"]
+    assert len(rows) == len(kinds)
+    for index, kind in enumerate(kinds):
+        if kind == "prefill":
+            assert rows[index].shape == (0, 12)
+        else:
+            torch.testing.assert_close(rows[index], originals[index], rtol=0, atol=0)
+    assert all("current" not in info.get("audio_codes", {}) for info in infos)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("as_tuple", [False, True])
+def test_async_packed_snapshot_survives_source_reuse(as_tuple):
+    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().clone() for v in values]
+    payload = tuple(values) if as_tuple else values
+    snapshot = _snapshot_tensor_payload_to_cpu_async(
+        {"codes": payload, "metadata": "kept"}, copy_stream=torch.cuda.Stream(), pin_memory=True
+    )
+    for v in values:
+        v.fill_(-99)
+    snapshot.wait()
+    actual = snapshot.payload["codes"]
+    assert isinstance(actual, tuple if as_tuple else list)
+    assert snapshot.payload["metadata"] == "kept"
+    for out, ref in zip(actual, expected, strict=True):
+        torch.testing.assert_close(out, ref, rtol=0, atol=0)
+    # A second output must not alias the first output's host snapshot either.
+    later = _snapshot_tensor_payload_to_cpu_async(values, copy_stream=torch.cuda.Stream(), pin_memory=True)
+    later.wait()
+    for out, ref in zip(actual, expected, strict=True):
+        torch.testing.assert_close(out, ref, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("offsets", [[0, 1], [1, 4]])
+def test_owned_code_batch_and_embeddings_survive_next_graph_output(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")}
+    original = runner.talker_mtp
+    returned = []
+
+    def mtp(*args, **kwargs):
+        result = original(*args, **kwargs)
+        returned.append(result)
+        return result
+
+    runner.talker_mtp = mtp
+    embeddings = torch.full((6, 4), -10.0)
+    runner._talker_mtp_forward(["r1", "r2"], embeddings, offsets)
+    assert torch.equal(embeddings[offsets], torch.ones(2, 4))
+    for i in set(range(6)) - set(offsets):
+        assert torch.equal(embeddings[i], torch.full((4,), -10.0))
+    returned[0][1].fill_(999)
+    assert runner.model_intermediate_buffer["r1"]["codes"]["audio"].item() == 0
+    assert runner.model_intermediate_buffer["r2"]["codes"]["audio"].item() == 1

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/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_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,16 @@ 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)
+                indices = torch.tensor(should_indices, device=hidden.device, dtype=torch.long)
+                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/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/worker/gpu_ar_model_runner.py

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

vllm_omni/worker/gpu_model_runner.py

@@ -1968,11 +1968,25 @@ 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:
+            offsets = torch.tensor(start_offsets, device=inputs_embeds.device, dtype=torch.long)
+            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,