perf(qwen3-tts): cache decoder masks and compile pre-transformer (#5107)

Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
Co-authored-by: Sy03 <1370724210@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/e2e/online_serving/test_qwen3_tts_customvoice.py

@@ -12,15 +12,21 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 import pytest
+from vllm.platforms import current_platform
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServerParams
-from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from tests.helpers.stage_config import (
+    get_deploy_config_path,
+    get_deploy_config_stage,
+    modify_stage_config,
+)
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
 
+_DEFAULT_STAGE_CONFIG = get_deploy_config_path("qwen3_tts.yaml")
 _STAGE_CONFIG = modify_stage_config(
-    get_deploy_config_path("qwen3_tts.yaml"),
+    _DEFAULT_STAGE_CONFIG,
     updates={"stages": {0: {"default_sampling_params.max_tokens": 500}}},
 )
 
@@ -50,6 +56,39 @@ tts_server_params = [
     )
 ]
 
+default_tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=_DEFAULT_STAGE_CONFIG,
+            server_args=["--trust-remote-code"],
+        ),
+        id="async_chunk",
+    )
+]
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA Graph startup test requires CUDA")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", default_tts_server_params, indirect=True)
+def test_default_cuda_graph_startup(omni_server) -> None:
+    """Verify both stages start with the shipped CUDA Graph configuration.
+
+    The fixture reaching this test is the smoke assertion: it waits for the
+    server to become ready after both stages finish model initialization and
+    CUDA Graph capture. The regression covered here exited during stage 1
+    capture, before fixture setup could complete.
+    """
+    for stage_id in (0, 1):
+        stage = get_deploy_config_stage("qwen3_tts.yaml", stage_id=stage_id)
+        assert stage.get("enforce_eager", False) is False
+
+    assert omni_server.proc is not None
+    assert omni_server.proc.poll() is None
+
 
 @pytest.mark.core_model
 @pytest.mark.advanced_model

tests/model_executor/models/qwen3_tts/test_qwen3_tts_tokenizer_mask_cache.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import inspect
+
+import pytest
+import torch
+from transformers.masking_utils import create_sliding_window_causal_mask
+
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2DecoderConfig,
+)
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2DecoderTransformerModel,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_config(attention_implementation: str) -> Qwen3TTSTokenizerV2DecoderConfig:
+    config = Qwen3TTSTokenizerV2DecoderConfig(
+        hidden_size=32,
+        latent_dim=16,
+        max_position_embeddings=512,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_quantizers=2,
+        decoder_dim=32,
+        upsample_rates=(2,),
+        upsampling_ratios=(2,),
+    )
+    config._attn_implementation = attention_implementation
+    return config
+
+
+def _make_model(
+    attention_implementation: str,
+    dtype: torch.dtype = torch.float32,
+) -> Qwen3TTSTokenizerV2DecoderTransformerModel:
+    model = Qwen3TTSTokenizerV2DecoderTransformerModel._from_config(_make_config(attention_implementation))
+    return model.to(dtype=dtype).eval()
+
+
+def _create_reference_mask(
+    model: Qwen3TTSTokenizerV2DecoderTransformerModel,
+    inputs_embeds: torch.Tensor,
+) -> torch.Tensor | None:
+    sequence_length = inputs_embeds.shape[1]
+    mask_kwargs = {
+        "config": model.config,
+        "attention_mask": None,
+        "past_key_values": None,
+        "position_ids": None,
+    }
+    sig = inspect.signature(create_sliding_window_causal_mask)
+    if "input_embeds" in sig.parameters:
+        mask_kwargs["input_embeds"] = inputs_embeds[:1]
+        mask_kwargs["cache_position"] = torch.arange(sequence_length)
+    else:
+        mask_kwargs["inputs_embeds"] = inputs_embeds[:1]
+    return create_sliding_window_causal_mask(**mask_kwargs)
+
+
+@pytest.mark.parametrize("sequence_length", [17, 72, 73, 169, 325])
+@pytest.mark.parametrize(
+    ("attention_implementation", "dtype"),
+    [("sdpa", torch.float32), ("eager", torch.float32), ("eager", torch.float16)],
+)
+def test_cached_sliding_mask_matches_transformers(
+    sequence_length: int,
+    attention_implementation: str,
+    dtype: torch.dtype,
+):
+    model = _make_model(attention_implementation, dtype)
+    inputs_embeds = torch.empty(2, sequence_length, model.config.hidden_size, dtype=dtype)
+
+    actual = model._get_sliding_attention_mask(inputs_embeds)
+    expected = _create_reference_mask(model, inputs_embeds)
+
+    if expected is None:
+        assert actual is None
+        assert model._get_sliding_attention_mask(inputs_embeds) is None
+        return
+
+    assert actual is not None
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+    assert actual.shape == (1, 1, sequence_length, sequence_length)
+    assert actual.is_contiguous()
+    assert model._get_sliding_attention_mask(inputs_embeds) is actual
+
+
+def test_cached_eager_mask_tracks_compute_dtype():
+    model = _make_model("eager")
+
+    fp32_inputs = torch.empty(1, 73, model.config.hidden_size)
+    fp32_mask = model._get_sliding_attention_mask(fp32_inputs)
+    assert fp32_mask is not None
+    assert fp32_mask.dtype == torch.float32
+
+    model.half()
+    fp16_inputs = fp32_inputs.half()
+    fp16_mask = model._get_sliding_attention_mask(fp16_inputs)
+    assert fp16_mask is not None
+    assert fp16_mask.dtype == torch.float16
+    assert torch.isfinite(fp16_mask).all()
+    assert fp16_mask.min() == torch.finfo(torch.float16).min
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "message"),
+    [
+        ({"use_cache": True}, "use_cache=True"),
+        ({"past_key_values": object()}, "past_key_values"),
+        ({"cache_position": torch.arange(1, 18)}, "zero-based cache_position"),
+        ({"position_ids": torch.arange(1, 18).unsqueeze(0)}, "zero-based position_ids"),
+    ],
+)
+def test_cached_mask_rejects_unsupported_cache_and_positions(kwargs: dict, message: str):
+    model = _make_model("sdpa")
+    inputs_embeds = torch.randn(1, 17, model.config.latent_dim)
+
+    with pytest.raises(ValueError, match=message):
+        model(inputs_embeds=inputs_embeds, **kwargs)

vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py

@@ -25,12 +25,9 @@ from torch import nn
 from torch.nn import functional as F
 from transformers import MimiConfig, MimiModel
 from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache, DynamicCache
+from transformers.cache_utils import Cache
 from transformers.integrations import use_kernel_forward_from_hub
-from transformers.masking_utils import (
-    create_causal_mask,
-    create_sliding_window_causal_mask,
-)
+from transformers.masking_utils import create_sliding_window_causal_mask
 from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
 from transformers.modeling_layers import GradientCheckpointingLayer
 from transformers.modeling_outputs import BaseModelOutputWithPast
@@ -523,15 +520,64 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
         self.norm = Qwen3TTSTokenizerV2DecoderRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.rotary_emb = Qwen3TTSTokenizerV2DecoderRotatoryEmbedding(config=config)
         self.gradient_checkpointing = False
-        self.has_sliding_layers = "sliding_attention" in self.config.layer_types
         self.window_size = config.sliding_window
 
         self.input_proj = nn.Linear(config.latent_dim, config.hidden_size)
         self.output_proj = nn.Linear(config.hidden_size, config.latent_dim)
 
+        # Code2Wav is cache-free and every decoder layer uses sliding attention.
+        # Cache an exact, compact mask per warmup shape instead of retaining a
+        # max_position_embeddings square buffer. CUDA Graph warmup populates
+        # these entries before capture, and dtype/device are part of the key.
+        self._sliding_attention_mask_cache: dict[tuple[int, torch.dtype, torch.device, str], torch.Tensor | None] = {}
+
         # Initialize weights and apply final processing
         self.post_init()
 
+    def _apply(self, fn, recurse=True):
+        # Cached masks are keyed by dtype/device and should not retain storage
+        # from the model's previous placement.
+        self._sliding_attention_mask_cache.clear()
+        return super()._apply(fn, recurse=recurse)
+
+    def _get_sliding_attention_mask(self, inputs_embeds: torch.Tensor) -> torch.Tensor | None:
+        sequence_length = int(inputs_embeds.shape[1])
+        max_position_embeddings = int(self.config.max_position_embeddings)
+        if sequence_length > max_position_embeddings:
+            raise ValueError(
+                f"Input length {sequence_length} exceeds max_position_embeddings {max_position_embeddings}"
+            )
+
+        attention_implementation = str(self.config._attn_implementation)
+        cache_key = (
+            sequence_length,
+            inputs_embeds.dtype,
+            inputs_embeds.device,
+            attention_implementation,
+        )
+        if cache_key not in self._sliding_attention_mask_cache:
+            # A batch-one mask broadcasts across requests and avoids caching a
+            # duplicate tensor for every batch size.
+            mask_inputs = inputs_embeds[:1]
+            mask_kwargs = {
+                "config": self.config,
+                "attention_mask": None,
+                "past_key_values": None,
+                "position_ids": None,
+            }
+            sig = inspect.signature(create_sliding_window_causal_mask)
+            if "input_embeds" in sig.parameters:
+                mask_kwargs["input_embeds"] = mask_inputs
+                mask_kwargs["cache_position"] = torch.arange(sequence_length, device=inputs_embeds.device)
+            else:
+                mask_kwargs["inputs_embeds"] = mask_inputs
+            sliding_attention_mask = create_sliding_window_causal_mask(**mask_kwargs)
+            if sliding_attention_mask is not None:
+                sliding_attention_mask = sliding_attention_mask.contiguous()
+            self._sliding_attention_mask_cache[cache_key] = sliding_attention_mask
+
+        return self._sliding_attention_mask_cache[cache_key]
+
     # Note: @check_model_inputs decorator removed for vLLM compatibility
     # The decorator causes "unexpected keyword argument 'inputs_embeds'" error
     @auto_docstring
@@ -554,57 +600,56 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
             raise ValueError("input_ids is not expected")
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+        if use_cache:
+            raise ValueError("Code2Wav sliding-mask caching does not support use_cache=True")
+        if past_key_values is not None:
+            raise ValueError("Code2Wav sliding-mask caching does not support past_key_values")
+        if attention_mask is not None and not isinstance(attention_mask, dict):
+            raise ValueError("Code2Wav cached masks require attention_mask to be None or a prepared mask mapping")
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
 
         inputs_embeds = self.input_proj(inputs_embeds)
 
-        if use_cache and past_key_values is None:
-            past_key_values = DynamicCache(config=self.config)
+        sequence_length = inputs_embeds.shape[1]
+        expected_cache_position = torch.arange(sequence_length, device=inputs_embeds.device)
 
         if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-            cache_position = torch.arange(
-                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
-            )
+            cache_position = expected_cache_position
+        elif attention_mask is None and (
+            cache_position.shape != expected_cache_position.shape
+            or cache_position.device != expected_cache_position.device
+            or not torch.equal(cache_position, expected_cache_position)
+        ):
+            raise ValueError("Cached Code2Wav masks require contiguous zero-based cache_position")
 
         if position_ids is None:
             position_ids = cache_position.unsqueeze(0)
-
-        # It may already have been prepared by e.g. `generate`
-        if not isinstance(causal_mask_mapping := attention_mask, dict):
-            # Prepare mask arguments
-            mask_kwargs = {
-                "config": self.config,
-                "attention_mask": attention_mask,
-                "past_key_values": past_key_values,
-                "position_ids": position_ids,
-            }
-            # Handle API changes across transformers versions
-            sig = inspect.signature(create_causal_mask)
-            if "input_embeds" in sig.parameters:
-                mask_kwargs["input_embeds"] = inputs_embeds
-                mask_kwargs["cache_position"] = cache_position
-            else:
-                mask_kwargs["inputs_embeds"] = inputs_embeds
-            # Create the masks
-            causal_mask_mapping = {
-                "full_attention": create_causal_mask(**mask_kwargs),
-            }
-            # The sliding window alternating layers are not always activated depending on the config
-            if self.has_sliding_layers:
-                causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
+        elif attention_mask is None:
+            expected_position_ids = expected_cache_position.unsqueeze(0)
+            if (
+                position_ids.ndim != 2
+                or position_ids.shape[-1] != sequence_length
+                or position_ids.device != expected_position_ids.device
+                or not torch.equal(position_ids, expected_position_ids.expand_as(position_ids))
+            ):
+                raise ValueError("Cached Code2Wav masks require contiguous zero-based position_ids")
 
         hidden_states = inputs_embeds
 
+        if isinstance(attention_mask, dict):
+            sliding_attention_mask = attention_mask["sliding_attention"]
+        else:
+            sliding_attention_mask = self._get_sliding_attention_mask(inputs_embeds)
+
         # create position embeddings to be shared across the decoder layers
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
 
         for decoder_layer in self.layers[: self.config.num_hidden_layers]:
             hidden_states = decoder_layer(
                 hidden_states,
-                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+                attention_mask=sliding_attention_mask,
                 position_ids=position_ids,
                 past_key_values=past_key_values,
                 use_cache=use_cache,
@@ -617,7 +662,7 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
         hidden_states = self.output_proj(hidden_states)
         return BaseModelOutputWithPast(
             last_hidden_state=hidden_states,
-            past_key_values=past_key_values if use_cache else None,
+            past_key_values=None,
         )