diff --git a/tests/e2e/online_serving/test_qwen3_tts_customvoice.py b/tests/e2e/online_serving/test_qwen3_tts_customvoice.py index 2aab84840..62528c272 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_customvoice.py +++ b/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 diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_tokenizer_mask_cache.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_tokenizer_mask_cache.py new file mode 100644 index 000000000..aeb538adc --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index 447823abb..9d6eb7aaa 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/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, )