From 251ef4f48fe2b9f461759f99f38668ffc03feede Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Mon, 7 Sep 2026 16:20:49 +0800 Subject: [PATCH] [Model] Optimize MOSS-TTS Local batched execution and streaming codec Signed-off-by: Sy03 <1370724210@qq.com> --- .buildkite/cuda/test-ready.yml | 2 +- .../online_serving/text_to_speech/README.md | 34 +++++ .../moss_tts/test_streaming_attention.py | 61 +++++++++ .../moss_tts/test_streaming_terminal_batch.py | 118 +++++++++++++++++ tests/worker/test_batched_omni_output.py | 94 ++++++++++++++ vllm_omni/deploy/moss_tts_local_h200.yaml | 64 ++++++++++ .../models/moss_tts/audio_tokenizer_v2.py | 14 ++- .../moss_tts/modeling_moss_tts_codec.py | 42 +++++-- .../moss_tts/modeling_moss_tts_talker.py | 25 ++-- .../models/moss_tts/streaming_attention.py | 119 ++++++++++++++++++ vllm_omni/worker/gpu_ar_model_runner.py | 39 +++++- vllm_omni/worker/gpu_model_runner.py | 24 +++- 12 files changed, 608 insertions(+), 28 deletions(-) create mode 100644 tests/model_executor/models/moss_tts/test_streaming_attention.py create mode 100644 tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py create mode 100644 tests/worker/test_batched_omni_output.py create mode 100644 vllm_omni/deploy/moss_tts_local_h200.yaml create mode 100644 vllm_omni/model_executor/models/moss_tts/streaming_attention.py diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index d7e14ed64..f54a58fa4 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.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 diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index e1cdd58dd..eae3dee27 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/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. diff --git a/tests/model_executor/models/moss_tts/test_streaming_attention.py b/tests/model_executor/models/moss_tts/test_streaming_attention.py new file mode 100644 index 000000000..fc77232a6 --- /dev/null +++ b/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")) diff --git a/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py b/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py new file mode 100644 index 000000000..585c3da49 --- /dev/null +++ b/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. diff --git a/tests/worker/test_batched_omni_output.py b/tests/worker/test_batched_omni_output.py new file mode 100644 index 000000000..86e598718 --- /dev/null +++ b/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 diff --git a/vllm_omni/deploy/moss_tts_local_h200.yaml b/vllm_omni/deploy/moss_tts_local_h200.yaml new file mode 100644 index 000000000..6fb5e010d --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py index 380096af5..4dbedbd19 100644 --- a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py +++ b/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) diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index 231e33857..c27949310 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/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 diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index 0495d79d8..7d743603f 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/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={}) diff --git a/vllm_omni/model_executor/models/moss_tts/streaming_attention.py b/vllm_omni/model_executor/models/moss_tts/streaming_attention.py new file mode 100644 index 000000000..1acce7710 --- /dev/null +++ b/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) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index d293f6312..6618ba920 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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, diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index eca15c16e..ae8de37b8 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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,