diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 7940e2176..2ea871890 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -2625,6 +2625,15 @@ class TestPlatformOverrides: assert stages[1].yaml_engine_args["dtype"] == dtype + def test_moss_tts_rocm_disables_codec_cudagraph(self): + deploy_path = Path(get_deploy_config_path("moss_tts.yaml")) + + base = load_deploy_config(deploy_path) + assert base.stages[1].enforce_eager is False + + rocm = _apply_platform_overrides(base, platform="rocm") + assert rocm.stages[1].enforce_eager is True + def test_higgs_audio_v3_rocm_uses_triton_attention(self): deploy_path = Path(get_deploy_config_path("higgs_multimodal_qwen3.yaml")) diff --git a/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py b/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py index 487a6bfa6..b86db369e 100644 --- a/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py +++ b/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py @@ -1,19 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E offline inference test for MOSS-TTS-v1.5 (MossTTSDelay-8B). MOSS-TTS-v1.5 is a continued-training upgrade of MOSS-TTS 1.0 with the SAME -``MossTTSDelay`` architecture and API. The 8B checkpoint is H100-gated; small-GPU -MossTTSDelay coverage stays in ``test_moss_tts_expansion.py``. +``MossTTSDelay`` architecture and API. The 8B checkpoint runs on H100 or MI325; +small-GPU MossTTSDelay coverage stays in ``test_moss_tts_expansion.py``. """ from __future__ import annotations +import gc import os import urllib.request import pytest import torch +from transformers import AutoProcessor from vllm import SamplingParams from tests.helpers.mark import hardware_test @@ -28,14 +30,10 @@ _OMNI_RUNNER_PARAM = ( {"stage_init_timeout": 600, "trust_remote_code": True}, ) -_SKIP_ISSUE_6417 = pytest.mark.skip( - reason="https://github.com/vllm-project/vllm-omni/issues/6417", -) - pytestmark = [ + pytest.mark.full_model, pytest.mark.slow, pytest.mark.tts, - _SKIP_ISSUE_6417, pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), ] @@ -69,15 +67,19 @@ def ref_audio_path(tmp_path_factory) -> str: def _build_request(text: str, ref_audio_path: str) -> dict: + processor = AutoProcessor.from_pretrained(MODEL, trust_remote_code=True) + message = processor.build_user_message( + text=text, + reference=[ref_audio_path], + language="English", + ) + unified = processor(conversations=[[message]], mode="generation")["input_ids"][0] + del processor + gc.collect() + return { - "prompt": "<|im_start|>assistant\n", - "additional_information": { - "task_type": ["voice_clone"], - "text": [text], - "mode": ["voice_clone"], - "prompt_audio_path": [ref_audio_path], - "seed": [42], - }, + "prompt_token_ids": unified[:, 0].tolist(), + "additional_information": {"codes": {"ref": unified[:, 1:].contiguous().to(torch.int64)}}, } @@ -112,7 +114,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor raise AssertionError("No stage outputs received") -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) def test_moss_tts_v15_voice_clone(omni_runner: OmniRunner, ref_audio_path) -> None: """MOSS-TTS-v1.5 (8B): voice_clone produces non-empty 24 kHz audio.""" req = _build_request("Hello, this is a MOSS-TTS v1.5 voice cloning test.", ref_audio_path) diff --git a/tests/model_executor/models/moss_tts/test_audio_tokenizer_v1_streaming.py b/tests/model_executor/models/moss_tts/test_audio_tokenizer_v1_streaming.py new file mode 100644 index 000000000..c03526307 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_audio_tokenizer_v1_streaming.py @@ -0,0 +1,199 @@ +# 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.audio_tokenizer import ( + MossAudioTokenizerConfig, + MossAudioTokenizerModel, + _Attention, + _StreamingExecutionContext, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_codec() -> MossAudioTokenizerModel: + torch.manual_seed(0) + config = MossAudioTokenizerConfig( + sampling_rate=100, + downsample_rate=1, + causal_transformer_context_duration=0.08, + encoder_kwargs=[{"module_type": "PatchedPretransform", "patch_size": 1}], + decoder_kwargs=[ + { + "module_type": "Transformer", + "input_dimension": 4, + "output_dimension": 1, + "d_model": 4, + "num_heads": 2, + "num_layers": 2, + "dim_feedforward": 8, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10_000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + } + ], + quantizer_kwargs={ + "input_dim": 4, + "rvq_dim": 4, + "output_dim": 4, + "num_quantizers": 2, + "codebook_size": 8, + "codebook_dim": 2, + "quantizer_type": "rlfq", + }, + ) + return MossAudioTokenizerModel(config).eval() + + +def _codes(values: list[list[int]]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.long) + + +def _reference_audio(codec: MossAudioTokenizerModel, codes: torch.Tensor) -> torch.Tensor: + output = codec.batch_decode([codes], num_quantizers=int(codes.shape[0])) + assert output.audio is not None + return output.audio[0] + + +def _stream( + codec: MossAudioTokenizerModel, + codes: torch.Tensor, + *, + slot: int, +) -> torch.Tensor: + output = codec.decode_streaming_batch( + codes.unsqueeze(1), + torch.tensor([codes.shape[1]], dtype=torch.long), + torch.tensor([slot], dtype=torch.long), + torch.tensor([True]), + ) + assert output.audio is not None + return output.audio[0] + + +def test_streaming_chunks_match_full_decode() -> None: + codec = _make_codec() + codes = _codes([[1, 2, 3, 4], [4, 3, 2, 1]]) + expected = _reference_audio(codec, codes) + + codec.initialize_decoder_state_pool(state_capacity=2) + first = _stream(codec, codes[:, :2], slot=1) + second = _stream(codec, codes[:, 2:], slot=1) + + torch.testing.assert_close(torch.cat([first, second], dim=-1), expected, atol=1e-5, rtol=1e-5) + + +def test_streaming_attention_ring_wrap_matches_full_attention() -> None: + torch.manual_seed(0) + attention = _Attention( + embed_dim=8, + num_heads=2, + causal=True, + max_period=10_000, + context=8, + ).eval() + hidden_states = torch.randn(1, 12, 8) + expected = attention(hidden_states) + + attention.initialize_streaming_state(torch.zeros(1, dtype=torch.long)) + execution_context = _StreamingExecutionContext( + state_slot_ids=torch.tensor([0]), + valid_rows=torch.tensor([True]), + ) + chunks = [ + attention(hidden_states[:, start : start + 3], execution_context) + for start in range(0, hidden_states.shape[1], 3) + ] + + torch.testing.assert_close(torch.cat(chunks, dim=1), expected, atol=1e-6, rtol=1e-6) + + +def test_streaming_slots_keep_independent_state() -> None: + codec = _make_codec() + first_codes = _codes([[1, 2, 3, 4], [4, 3, 2, 1]]) + second_codes = _codes([[5, 6, 7, 1], [1, 7, 6, 5]]) + first_expected = _reference_audio(codec, first_codes) + second_expected = _reference_audio(codec, second_codes) + + codec.initialize_decoder_state_pool(state_capacity=3) + chunks: list[tuple[torch.Tensor, torch.Tensor]] = [] + for start in (0, 2): + batch = torch.stack([first_codes[:, start : start + 2], second_codes[:, start : start + 2]], dim=1) + output = codec.decode_streaming_batch( + batch, + torch.tensor([2, 2], dtype=torch.long), + torch.tensor([2, 0], dtype=torch.long), + torch.tensor([True, True]), + ) + assert output.audio is not None + chunks.append((output.audio[0], output.audio[1])) + + torch.testing.assert_close( + torch.cat([chunk[0] for chunk in chunks], dim=-1), + first_expected, + atol=1e-5, + rtol=1e-5, + ) + torch.testing.assert_close( + torch.cat([chunk[1] for chunk in chunks], dim=-1), + second_expected, + atol=1e-5, + rtol=1e-5, + ) + + +def test_reset_slot_removes_previous_request_state() -> None: + codec = _make_codec() + old_codes = _codes([[1, 2], [4, 3]]) + new_codes = _codes([[7, 6], [2, 5]]) + expected = _reference_audio(codec, new_codes) + + codec.initialize_decoder_state_pool(state_capacity=1) + _stream(codec, old_codes, slot=0) + codec.reset_decoder_state_slots(torch.tensor([0], dtype=torch.long)) + + torch.testing.assert_close(_stream(codec, new_codes, slot=0), expected, atol=1e-5, rtol=1e-5) + + +def test_invalid_padding_row_does_not_advance_state() -> None: + codec = _make_codec() + codes = _codes([[1, 2], [4, 3]]) + expected = _reference_audio(codec, codes) + + codec.initialize_decoder_state_pool(state_capacity=1, scratch_capacity=1) + padding_output = codec.decode_streaming_batch( + codes.unsqueeze(1), + torch.tensor([0], dtype=torch.long), + torch.tensor([1], dtype=torch.long), + torch.tensor([False]), + ) + assert padding_output.audio_lengths is not None + assert int(padding_output.audio_lengths[0]) == 0 + + torch.testing.assert_close(_stream(codec, codes, slot=1), expected, atol=1e-5, rtol=1e-5) + + +def test_streaming_requires_initialized_state_pool() -> None: + codec = _make_codec() + codes = _codes([[1], [2]]) + + with pytest.raises(RuntimeError, match="state pool is not initialized"): + _stream(codec, codes, slot=0) + + +def test_closed_state_pool_can_be_reinitialized() -> None: + codec = _make_codec() + codes = _codes([[1], [2]]) + + codec.initialize_decoder_state_pool(state_capacity=1) + codec.close_decoder_state_pool() + codec.initialize_decoder_state_pool(state_capacity=1) + + assert _stream(codec, codes, slot=0).numel() > 0 diff --git a/vllm_omni/deploy/moss_tts.yaml b/vllm_omni/deploy/moss_tts.yaml index 92bd2f3ba..229377431 100644 --- a/vllm_omni/deploy/moss_tts.yaml +++ b/vllm_omni/deploy/moss_tts.yaml @@ -57,3 +57,11 @@ stages: max_tokens: 65536 seed: 42 repetition_penalty: 1.0 + +platforms: + rocm: + stages: + - stage_id: 1 + # MIOPEN_FIND_MODE=FAST changes solver selection, not HIP graph safety. + # Keep the Conv1d codec outside the outer vLLM graph on ROCm. + enforce_eager: true diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py index 99dcb33e4..a8fd5b25d 100644 --- a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py +++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.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"); @@ -8,7 +11,6 @@ # # Vendored from OpenMOSS-Team/MOSS-Audio-Tokenizer (configuration_moss_audio_tokenizer.py # and modeling_moss_audio_tokenizer.py). Simplified for inference-only use: -# - Streaming KV-cache infrastructure removed (single-pass batch decode only). # - Training-only methods (forward, encode, decode) removed. # - Dead branches removed: gating="none" always, weights_per_step=0 always, # positional_embedding="rope" always, norm="layer_norm" always in default config. @@ -155,6 +157,29 @@ class MossAudioTokenizerDecoderOutput(ModelOutput): audio_lengths: torch.Tensor | None = None +@dataclass(frozen=True) +class _StreamingExecutionContext: + state_slot_ids: torch.Tensor + valid_rows: torch.Tensor + + def validate(self, *, batch_size: int, state_capacity: int, device: torch.device) -> None: + if self.state_slot_ids.shape != (batch_size,): + raise ValueError(f"Expected state_slot_ids shape ({batch_size},), got {tuple(self.state_slot_ids.shape)}") + if self.valid_rows.shape != (batch_size,): + raise ValueError(f"Expected valid_rows shape ({batch_size},), got {tuple(self.valid_rows.shape)}") + if self.state_slot_ids.device != device or self.valid_rows.device != device: + raise ValueError("Streaming execution metadata must be on the same device as decoder inputs.") + if self.state_slot_ids.dtype != torch.long or self.valid_rows.dtype != torch.bool: + raise TypeError("state_slot_ids must be int64 and valid_rows must be bool.") + if device.type == "cpu" and self.state_slot_ids.numel() > 0: + min_slot = int(self.state_slot_ids.min()) + max_slot = int(self.state_slot_ids.max()) + if min_slot < 0 or max_slot >= state_capacity: + raise ValueError( + f"Decoder state slots must be in [0, {state_capacity}), got {self.state_slot_ids.tolist()}" + ) + + # --------------------------------------------------------------------------- # Building blocks # --------------------------------------------------------------------------- @@ -169,7 +194,12 @@ class _LayerScale(nn.Module): return self.scale * x -def _apply_rope(q: torch.Tensor, k: torch.Tensor, max_period: float = 10_000) -> tuple[torch.Tensor, torch.Tensor]: +def _apply_rope( + q: torch.Tensor, + k: torch.Tensor, + offsets: torch.Tensor | None = None, + max_period: float = 10_000, +) -> tuple[torch.Tensor, torch.Tensor]: """Rotary position embedding over sequence dimension (B, H, T, D). Matches upstream MossAudioTokenizer's ``apply_rope``: pair the last dim @@ -181,7 +211,10 @@ def _apply_rope(q: torch.Tensor, k: torch.Tensor, max_period: float = 10_000) -> half = D // 2 ds = torch.arange(half, device=q.device, dtype=torch.float32) freqs = torch.exp(ds * (-math.log(max_period) * 2 / D)) - ts = torch.arange(T, device=q.device, dtype=torch.float32).view(1, 1, -1, 1) # (1, 1, T, 1) + if offsets is None: + offsets = torch.zeros(B, device=q.device, dtype=torch.long) + ts = offsets.float().view(B, 1) + torch.arange(T, device=q.device, dtype=torch.float32) + ts = ts.view(B, 1, T, 1) rotr = torch.cos(freqs * ts) # (1, 1, T, D/2) roti = torch.sin(freqs * ts) @@ -201,8 +234,80 @@ def _apply_rope(q: torch.Tensor, k: torch.Tensor, max_period: float = 10_000) -> return qo, ko +class _RingKVCache: + def __init__( + self, + state_capacity: int, + num_heads: int, + context: int, + head_dim: int, + *, + device: torch.device, + dtype: torch.dtype, + ) -> None: + self.context = context + self.cache = torch.zeros( + (2, state_capacity, num_heads, context, head_dim), + device=device, + dtype=dtype, + ) + + def complete( + self, + k: torch.Tensor, + v: torch.Tensor, + offsets: torch.Tensor, + execution_context: _StreamingExecutionContext, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size, num_heads, num_frames, head_dim = k.shape + slots = execution_context.state_slot_ids + row_cache = self.cache.index_select(1, slots) + cache_indexes = torch.arange(self.context, device=offsets.device, dtype=torch.long) + last_offset = offsets.view(-1, 1) - 1 + end_index = last_offset % self.context + delta = cache_indexes - end_index + cached_positions = torch.where( + delta <= 0, + last_offset + delta, + last_offset + delta - self.context, + ) + cached_positions = torch.where( + (cache_indexes < offsets.view(-1, 1)) & execution_context.valid_rows.view(-1, 1), + cached_positions, + torch.full_like(cached_positions, -1), + ) + current_positions = offsets.view(-1, 1) + torch.arange(num_frames, device=offsets.device, dtype=offsets.dtype) + current_positions = torch.where( + execution_context.valid_rows.view(-1, 1), + current_positions, + torch.full_like(current_positions, -1), + ) + + # Keep this chunk separate from the pre-write cache so a ring wrap does + # not overwrite keys that early queries in the chunk still need. + keys = torch.cat((row_cache[0], k), dim=2) + values = torch.cat((row_cache[1], v), dim=2) + positions = torch.cat((cached_positions, current_positions), dim=1) + + cache_frames = min(num_frames, self.context) + cache_start = offsets.view(-1, 1) + num_frames - cache_frames + indexes = torch.arange(cache_frames, device=offsets.device, dtype=offsets.dtype) + indexes = (indexes + cache_start) % self.context + scatter_indexes = indexes.view(batch_size, 1, cache_frames, 1).expand(-1, num_heads, -1, head_dim) + row_cache[0].scatter_(2, scatter_indexes, k[:, :, -cache_frames:]) + row_cache[1].scatter_(2, scatter_indexes, v[:, :, -cache_frames:]) + self.cache.index_copy_(1, slots, row_cache) + return keys, values, positions + + +@dataclass +class _AttentionStreamingState: + offsets: torch.Tensor + kv_cache: _RingKVCache + + class _Attention(nn.Module): - """Causal multi-head self-attention with RoPE, no streaming KV cache.""" + """Causal multi-head self-attention with request-addressed streaming state.""" def __init__( self, @@ -220,6 +325,7 @@ class _Attention(nn.Module): self.causal = causal self.max_period = max_period self.context = context + self._streaming_state: _AttentionStreamingState | None = None kw = {"device": device, "dtype": dtype} self.in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False, **kw) self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False, **kw) @@ -240,12 +346,58 @@ class _Attention(nn.Module): dst = src.replace("in_projs.0.", "in_proj.").replace("out_projs.0.", "out_proj.") state_dict[dst] = state_dict.pop(src) - def forward(self, x: torch.Tensor) -> torch.Tensor: + def initialize_streaming_state(self, offsets: torch.Tensor) -> None: + if self._streaming_state is not None: + raise RuntimeError("MOSS Audio Tokenizer attention is already streaming.") + weight = self.in_proj.weight + head_dim = self.embed_dim // self.num_heads + self._streaming_state = _AttentionStreamingState( + offsets=offsets, + kv_cache=_RingKVCache( + offsets.shape[0], + self.num_heads, + self.context or 1024, + head_dim, + device=weight.device, + dtype=weight.dtype, + ), + ) + + def close_streaming_state(self) -> None: + self._streaming_state = None + + def forward( + self, + x: torch.Tensor, + execution_context: _StreamingExecutionContext | None = None, + ) -> torch.Tensor: B, T, _ = x.shape H, D = self.num_heads, self.embed_dim // self.num_heads qkv = self.in_proj(x).reshape(B, T, 3, H, D).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] # each (B, H, T, D) - q, k = _apply_rope(q, k, self.max_period) + state = self._streaming_state + if execution_context is None: + offsets = torch.zeros(B, device=x.device, dtype=torch.long) + else: + if state is None: + raise RuntimeError("MOSS Audio Tokenizer decoder state pool is not initialized.") + offsets = state.offsets.index_select(0, execution_context.state_slot_ids) + q, k = _apply_rope(q, k, offsets, self.max_period) + + if execution_context is not None: + assert state is not None + k, v, key_positions = state.kv_cache.complete(k, v, offsets, execution_context) + query_positions = offsets.view(-1, 1, 1) + torch.arange(T, device=x.device).view(1, -1, 1) + delta = query_positions - key_positions.view(B, 1, -1) + mask = (key_positions.view(B, 1, -1) >= 0) & (delta >= 0) + if self.context is not None: + mask &= delta < self.context + out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask.unsqueeze(1)) + next_offsets = torch.where(execution_context.valid_rows, offsets + T, offsets) + state.offsets.index_copy_(0, execution_context.state_slot_ids, next_offsets) + out = out.transpose(1, 2).reshape(B, T, self.embed_dim) + return self.out_proj(out) + if self.context is not None and self.context < T: # Local-windowed causal attention: query i may only see keys in # [i - context + 1, i]. Matches upstream's per-stage receptive @@ -310,8 +462,12 @@ class _TransformerLayer(nn.Module): self.ls1 = nn.Identity() self.ls2 = nn.Identity() - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = x + self.ls1(self.attn(self.norm1(x))) + def forward( + self, + x: torch.Tensor, + execution_context: _StreamingExecutionContext | None = None, + ) -> torch.Tensor: + x = x + self.ls1(self.attn(self.norm1(x), execution_context)) x = x + self.ls2(self.ff2(F.gelu(self.ff1(self.norm2(x))))) return x @@ -340,9 +496,13 @@ class _Transformer(nn.Module): ] ) - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward( + self, + x: torch.Tensor, + execution_context: _StreamingExecutionContext | None = None, + ) -> torch.Tensor: for layer in self.layers: - x = layer(x) + x = layer(x, execution_context) return x @@ -360,9 +520,14 @@ class _ProjectedTransformer(nn.Module): nn.Linear(d_model, output_dimension, bias=False) if output_dimension != d_model else nn.Identity() ) - def forward(self, x: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def forward( + self, + x: torch.Tensor, + lengths: torch.Tensor, + execution_context: _StreamingExecutionContext | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: x = self.in_proj(x.transpose(1, 2)) # (B, D, T) → (B, T, d_model) - x = self.transformer(x) + x = self.transformer(x, execution_context) x = self.out_proj(x).transpose(1, 2) # (B, T, D) → (B, D, T) return x, lengths @@ -544,7 +709,7 @@ def _build_modules( class MossAudioTokenizerModel(PreTrainedModel): - """MOSS Audio Tokenizer — inference-only (batch_encode / batch_decode).""" + """MOSS Audio Tokenizer inference model with stateful decoder streaming.""" config_class = MossAudioTokenizerConfig base_model_prefix = "" @@ -579,8 +744,83 @@ class MossAudioTokenizerModel(PreTrainedModel): kw = dict(config.quantizer_kwargs) self.quantizer = _ResidualQ(**kw) + self._decoder_streaming_attentions: list[_Attention] = [] + self._decoder_slot_offsets: torch.Tensor | None = None + self._decoder_state_capacity = 0 self.post_init() + def initialize_decoder_state_pool(self, state_capacity: int, scratch_capacity: int = 0) -> None: + if state_capacity <= 0 or scratch_capacity < 0: + raise ValueError(f"Invalid decoder state capacities: state={state_capacity}, scratch={scratch_capacity}.") + if self._decoder_streaming_attentions: + raise RuntimeError("MOSS Audio Tokenizer decoder state pool is already initialized.") + attentions = [module for module in self.decoder.modules() if isinstance(module, _Attention)] + if not attentions: + raise RuntimeError("MOSS Audio Tokenizer decoder has no streaming attention layers.") + total_capacity = state_capacity + scratch_capacity + offsets = torch.zeros( + (len(attentions), total_capacity), + device=next(self.parameters()).device, + dtype=torch.long, + ) + for index, attention in enumerate(attentions): + attention.initialize_streaming_state(offsets[index]) + self._decoder_streaming_attentions = attentions + self._decoder_slot_offsets = offsets + self._decoder_state_capacity = total_capacity + + def close_decoder_state_pool(self) -> None: + for attention in self._decoder_streaming_attentions: + attention.close_streaming_state() + self._decoder_streaming_attentions = [] + self._decoder_slot_offsets = None + self._decoder_state_capacity = 0 + + def reset_decoder_state_slots(self, state_slot_ids: torch.Tensor) -> None: + if not self._decoder_streaming_attentions or self._decoder_slot_offsets is None: + raise RuntimeError("MOSS Audio Tokenizer decoder state pool is not initialized.") + if state_slot_ids.numel() == 0: + return + if state_slot_ids.device.type == "cpu": + min_slot = int(state_slot_ids.min()) + max_slot = int(state_slot_ids.max()) + if min_slot < 0 or max_slot >= self._decoder_state_capacity: + raise ValueError( + f"Decoder state slots must be in [0, {self._decoder_state_capacity}), got {state_slot_ids.tolist()}" + ) + slots = state_slot_ids.to(device=self._decoder_slot_offsets.device, dtype=torch.long) + # A zero offset marks every old ring-cache entry invalid, so the large + # KV buffers do not need to be cleared when a slot is reused. + self._decoder_slot_offsets.index_fill_(1, slots, 0) + + def decode_streaming_batch( + self, + codes: torch.Tensor, + codes_lengths: torch.Tensor, + state_slot_ids: torch.Tensor, + valid_rows: torch.Tensor, + ) -> MossAudioTokenizerDecoderOutput: + if not self._decoder_streaming_attentions: + raise RuntimeError("MOSS Audio Tokenizer decoder state pool is not initialized.") + execution_context = _StreamingExecutionContext(state_slot_ids, valid_rows) + execution_context.validate( + batch_size=int(codes.shape[1]), + state_capacity=self._decoder_state_capacity, + device=codes.device, + ) + audio, audio_lengths = self._decode_tensors(codes, codes_lengths, execution_context) + return MossAudioTokenizerDecoderOutput(audio=audio, audio_lengths=audio_lengths) + + def decode_streaming_tensors( + self, + codes: torch.Tensor, + codes_lengths: torch.Tensor, + state_slot_ids: torch.Tensor, + valid_rows: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + execution_context = _StreamingExecutionContext(state_slot_ids, valid_rows) + return self._decode_tensors(codes, codes_lengths, execution_context) + @torch.no_grad() def batch_encode( self, @@ -688,13 +928,26 @@ class MossAudioTokenizerModel(PreTrainedModel): self, codes: torch.Tensor, lengths: torch.Tensor, + execution_context: _StreamingExecutionContext | None = None, ) -> MossAudioTokenizerDecoderOutput: + audio, audio_lengths = self._decode_tensors(codes, lengths, execution_context) + return MossAudioTokenizerDecoderOutput(audio=audio, audio_lengths=audio_lengths) + + def _decode_tensors( + self, + codes: torch.Tensor, + lengths: torch.Tensor, + execution_context: _StreamingExecutionContext | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: z = self.quantizer.decode_codes(codes) d, d_len = z, lengths for m in self.decoder: - d, d_len = m(d, d_len) + if isinstance(m, _ProjectedTransformer): + d, d_len = m(d, d_len, execution_context) + else: + d, d_len = m(d, d_len) d, d_len = self._restore_channels_from_codec(d, d_len) - return MossAudioTokenizerDecoderOutput(audio=d, audio_lengths=d_len) + return d, d_len __all__ = [ diff --git a/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py index 52815dd3b..4d277f44e 100644 --- a/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py +++ b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.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"). @@ -30,11 +33,10 @@ logger = init_logger(__name__) } ) class _MossStreamingDecodeCompileAdapter(nn.Module): - """Expose the BF16 codec streaming hot path to vLLM compile. + """Expose the codec streaming hot path to vLLM compile. - The dtype is part of the class identity intentionally: vLLM's AOT cache - drops runtime guards, so an artifact traced with the old FP32 decoder - weights must never be reused after the decoder is materialized as BF16. + The wrapper preserves the codec dtype: v1 decodes in FP32 to match its + quantizer output, while v2 uses BF16. """ def __init__(self, codec: nn.Module, *, vllm_config: VllmConfig) -> None: 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 c27949310..fbefe3a6a 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 @@ -767,7 +767,8 @@ class MossTTSCodecDecoder(nn.Module): ) codec.eval() - if device.type != "cpu": + # The v1 quantizer emits FP32 tensors, so its decoder must remain FP32. + if device.type != "cpu" and isinstance(codec, MossAudioTokenizerV2Model): codec.decoder.to(dtype=torch.bfloat16) attention_backend = getattr(self.vllm_config.model_config.hf_config, "codec_attention_backend", "sdpa") if attention_backend != "sdpa":