[Refactor][Qwen3-TTS] Construct Code2Wav decoder natively (#2341)

Signed-off-by: Nick Cao <ncao@redhat.com> Co-authored-by: Claude <noreply@anthropic.com>
2026-05-05 15:26:30 +00:00
1 parent b64ab05d09
commit fde38ca54f
@@ -2,18 +2,25 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 from types import SimpleNamespace
+from unittest.mock import patch

 import pytest
 import torch
 import torch.nn as nn

-from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import Qwen3TTSCode2Wav
+from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import (
+    Qwen3TTSCode2Wav,
+)

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

+_NUM_QUANTIZERS = 2
+_TOTAL_UPSAMPLE = 4
+_OUTPUT_SAMPLE_RATE = 24000
+

 class _FakeDecoder(nn.Module):
-    def __init__(self, total_upsample: int = 4):
+    def __init__(self, total_upsample: int = _TOTAL_UPSAMPLE):
        super().__init__()
        self.total_upsample = total_upsample

@@ -24,18 +31,35 @@ class _FakeDecoder(nn.Module):
        return wav.view(1, 1, -1)


-def _make_model() -> Qwen3TTSCode2Wav:
-    model = Qwen3TTSCode2Wav(
-        vllm_config=SimpleNamespace(
-            model_config=SimpleNamespace(model="unused"),
-            device_config=SimpleNamespace(device=torch.device("cpu")),
-        )
+def _fake_dec_config():
+    return SimpleNamespace(
+        num_quantizers=_NUM_QUANTIZERS,
+        sliding_window=0,
    )
-    model._decoder = _FakeDecoder()
-    model._num_quantizers = 2
-    model._output_sample_rate = 24000
-    model._total_upsample = 4
-    model._ensure_speech_tokenizer_loaded = lambda: None
+
+
+def _make_model() -> Qwen3TTSCode2Wav:
+    dec_config = _fake_dec_config()
+    tok_config = SimpleNamespace(
+        decoder_config=dec_config,
+        output_sample_rate=_OUTPUT_SAMPLE_RATE,
+    )
+    with (
+        patch(
+            "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav.Qwen3TTSTokenizerV2Config.from_pretrained",
+            return_value=tok_config,
+        ),
+        patch(
+            "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav.Qwen3TTSTokenizerV2Decoder._from_config",
+            return_value=_FakeDecoder(),
+        ),
+    ):
+        model = Qwen3TTSCode2Wav(
+            vllm_config=SimpleNamespace(
+                model_config=SimpleNamespace(model="unused"),
+                device_config=SimpleNamespace(device=torch.device("cpu")),
+            )
+        )
    return model


@@ -1,19 +1,24 @@
 from __future__ import annotations

-import os
 from collections.abc import Iterable
 from typing import Any

 import torch
 import torch.nn as nn
-from transformers.utils.hub import cached_file
 from vllm.config import VllmConfig
 from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.logger import init_logger
+from vllm.model_executor.model_loader import DefaultModelLoader
+from vllm.model_executor.models.utils import AutoWeightsLoader

 from vllm_omni.model_executor.models.output_templates import OmniOutput

-from .qwen3_tts_tokenizer import Qwen3TTSTokenizer
+from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2Config,
+)
+from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2Decoder,
+)

 logger = init_logger(__name__)

@@ -36,128 +41,24 @@ class Qwen3TTSCode2Wav(nn.Module):
        self.enable_update_additional_information = True
        self.requires_raw_input_tokens = True

-        self._speech_tokenizer: Qwen3TTSTokenizer | None = None
-        self._decoder: nn.Module | None = None
-        self._num_quantizers: int | None = None
-        self._output_sample_rate: int | None = None
-        self._total_upsample: int | None = None
-        self._decoder_sliding_window: int | None = None
        self._decode_chunk_frames = 300
        self._decode_left_context_frames = 25
        self._logged_codec_stats = False

-    @staticmethod
-    def _module_device(module: nn.Module) -> torch.device:
-        try:
-            return next(module.parameters()).device
-        except StopIteration:
-            for _, buf in module.named_buffers(recurse=True):
-                return buf.device
-            return torch.device("cpu")
-
-    def _ensure_speech_tokenizer_loaded(self) -> None:
-        if self._decoder is not None:
-            return
-
-        cfg_path = cached_file(self.model_path, "speech_tokenizer/config.json")
-        if cfg_path is None:
-            raise ValueError(f"{self.model_path}/speech_tokenizer/config.json not found")
-        speech_tokenizer_dir = os.path.dirname(cfg_path)
-
-        prep_cfg = cached_file(self.model_path, "speech_tokenizer/preprocessor_config.json")
-        if prep_cfg is None:
-            raise ValueError(
-                f"{self.model_path}/speech_tokenizer/preprocessor_config.json not found. "
-                "Please make sure the checkpoint contains the required HF preprocessing files."
-            )
-
-        tok = Qwen3TTSTokenizer.from_pretrained(
-            speech_tokenizer_dir,
-            torch_dtype=torch.float32,
-            load_feature_extractor=False,
+        # Construct decoder from config so it is visible to vLLM's
+        # memory profiler at startup.  Weights are loaded later in
+        # load_weights().
+        tok_config = Qwen3TTSTokenizerV2Config.from_pretrained(
+            self.model_path,
+            subfolder="speech_tokenizer",
        )
-
-        if tok.model is not None:
-            tok.model.to(device=self.vllm_config.device_config.device)
-            tok.device = self._module_device(tok.model)
-
-        dec_cfg = getattr(tok.model.config, "decoder_config", None)
-        num_q = getattr(dec_cfg, "num_quantizers", None) if dec_cfg is not None else None
-        if num_q is None:
-            raise ValueError("speech_tokenizer decoder_config.num_quantizers not found")
-        num_q = int(num_q)
-        if num_q <= 0:
-            raise ValueError(f"Invalid speech_tokenizer num_quantizers={num_q}")
-
-        try:
-            upsample = int(tok.get_decode_upsample_rate())
-        except Exception as e:
-            raise ValueError(f"Failed to get decode upsample rate: {e}") from e
-        if upsample <= 0:
-            raise ValueError(f"Invalid decode upsample rate: {upsample}")
-
-        try:
-            out_sr = int(tok.get_output_sample_rate())
-        except Exception as e:
-            raise ValueError(f"Failed to get output sample rate: {e}") from e
-
-        decoder = tok.model.decoder
-        decoder.eval()
-
-        self._speech_tokenizer = tok
-        self._decoder = decoder
-        self._num_quantizers = num_q
-        self._output_sample_rate = out_sr
-        self._total_upsample = int(decoder.total_upsample)
-        self._decoder_sliding_window = int(getattr(dec_cfg, "sliding_window", 0) or 0)
-
-        # Precompute SnakeBeta exp caches (benefits both Triton and eager paths)
-        if hasattr(decoder, "precompute_snake_caches"):
-            decoder.precompute_snake_caches()
-
-        chunk_frames = 0
-        left_frames = 0
-        model_cfg = getattr(self.vllm_config, "model_config", None)
-        connector_cfg = getattr(model_cfg, "stage_connector_config", None)
-        extra_cfg = (
-            connector_cfg.get("extra", connector_cfg)
-            if isinstance(connector_cfg, dict)
-            else getattr(connector_cfg, "extra", None)
-        )
-        if isinstance(extra_cfg, dict):
-            chunk_frames = int(extra_cfg.get("codec_chunk_frames") or 0)
-            left_frames = int(extra_cfg.get("codec_left_context_frames") or 0)
-            if getattr(model_cfg, "async_chunk", False) and chunk_frames > 0:
-                self._decode_chunk_frames = chunk_frames
-                self._decode_left_context_frames = left_frames if left_frames > 0 else 0
-
-        if hasattr(decoder, "enable_cudagraph"):
-            device = self._module_device(decoder)
-            if device.type == "cuda":
-                try:
-                    if (
-                        chunk_frames > 0
-                        and left_frames > 0
-                        and self._decoder_sliding_window
-                        and left_frames < self._decoder_sliding_window
-                    ):
-                        logger.warning(
-                            "Qwen3-TTS streaming codec_left_context_frames=%d is smaller than "
-                            "decoder sliding_window=%d; chunk-boundary distortion may occur. "
-                            "Increase codec_left_context_frames to at least %d for streaming.",
-                            left_frames,
-                            self._decoder_sliding_window,
-                            self._decoder_sliding_window,
-                        )
-
-                    decoder.enable_cudagraph(
-                        device=device,
-                        codec_chunk_frames=chunk_frames,
-                        codec_left_context_frames=left_frames,
-                    )
-                    logger.info("Code2Wav decoder CUDA Graph enabled")
-                except Exception:
-                    logger.warning("Failed to enable CUDA Graph for Code2Wav decoder", exc_info=True)
+        dec_config = tok_config.decoder_config
+        self.decoder = Qwen3TTSTokenizerV2Decoder._from_config(dec_config)
+        self.decoder.eval()
+        self._num_quantizers = int(dec_config.num_quantizers)
+        self._output_sample_rate = int(tok_config.output_sample_rate)
+        self._total_upsample = int(self.decoder.total_upsample)
+        self._decoder_sliding_window = int(getattr(dec_config, "sliding_window", 0) or 0)

    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
        # This stage ignores token embeddings. Keep a stable dummy embedding for vLLM runner.
@@ -210,12 +111,7 @@ class Qwen3TTSCode2Wav(nn.Module):
        Length management is done here instead of relying on HF's padding=-1
        sentinel logic.
        """
-        self._ensure_speech_tokenizer_loaded()
-        assert self._decoder is not None
-        assert self._num_quantizers is not None
-        assert self._total_upsample is not None
-
-        decoder = self._decoder
+        decoder = self.decoder
        q = int(self._num_quantizers)
        upsample = int(self._total_upsample)
        sr_val = int(self._output_sample_rate)
@@ -362,6 +258,80 @@ class Qwen3TTSCode2Wav(nn.Module):
        )

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        # SpeechTokenizer weights live under `speech_tokenizer/` and are loaded
-        # lazily from that directory. Ignore main checkpoint weights.
-        return set()
+        # The primary weights iterator contains no Code2Wav parameters.
+        # Drain it so callers don't hang on an unconsumed generator.
+        for _ in weights:
+            pass
+
+        # Load decoder weights from the speech_tokenizer/ subfolder
+        # via vLLM's weight loader (handles sharded safetensors, index
+        # files, and all load formats).  AutoWeightsLoader matches
+        # "decoder.*" weights to self.decoder and skips encoder weights.
+        model_loader = DefaultModelLoader(self.vllm_config.load_config)
+        source = DefaultModelLoader.Source(
+            model_or_path=self.model_path,
+            revision=self.vllm_config.model_config.revision,
+            subfolder="speech_tokenizer",
+        )
+        subfolder_weights = model_loader._get_weights_iterator(source)
+        loaded = AutoWeightsLoader(
+            self,
+            skip_prefixes=["encoder."],
+        ).load_weights(subfolder_weights)
+
+        device = self.vllm_config.device_config.device
+        self.decoder.to(device=device, dtype=torch.float32)
+
+        # Precompute SnakeBeta exp caches (benefits both Triton and eager paths)
+        if hasattr(self.decoder, "precompute_snake_caches"):
+            self.decoder.precompute_snake_caches()
+
+        # Read chunk config from stage connector and update decode params
+        chunk_frames = 0
+        left_frames = 0
+        model_cfg = getattr(self.vllm_config, "model_config", None)
+        connector_cfg = getattr(model_cfg, "stage_connector_config", None)
+        extra_cfg = (
+            connector_cfg.get("extra", connector_cfg)
+            if isinstance(connector_cfg, dict)
+            else getattr(connector_cfg, "extra", None)
+        )
+        if isinstance(extra_cfg, dict):
+            chunk_frames = int(extra_cfg.get("codec_chunk_frames") or 0)
+            left_frames = int(extra_cfg.get("codec_left_context_frames") or 0)
+            if getattr(model_cfg, "async_chunk", False) and chunk_frames > 0:
+                self._decode_chunk_frames = chunk_frames
+                self._decode_left_context_frames = left_frames if left_frames > 0 else 0
+
+        if hasattr(self.decoder, "enable_cudagraph") and device.type == "cuda":
+            try:
+                if (
+                    chunk_frames > 0
+                    and left_frames > 0
+                    and self._decoder_sliding_window
+                    and left_frames < self._decoder_sliding_window
+                ):
+                    logger.warning(
+                        "Qwen3-TTS streaming codec_left_context_frames=%d "
+                        "is smaller than decoder sliding_window=%d; "
+                        "chunk-boundary distortion may occur. "
+                        "Increase codec_left_context_frames to at least "
+                        "%d for streaming.",
+                        left_frames,
+                        self._decoder_sliding_window,
+                        self._decoder_sliding_window,
+                    )
+
+                self.decoder.enable_cudagraph(
+                    device=device,
+                    codec_chunk_frames=chunk_frames,
+                    codec_left_context_frames=left_frames,
+                )
+                logger.info("Code2Wav decoder CUDA Graph enabled")
+            except Exception:
+                logger.warning(
+                    "Failed to enable CUDA Graph for Code2Wav decoder",
+                    exc_info=True,
+                )
+
+        return loaded