Add compile_codec option for 3-4x faster batch decoding

README.md

@@ -27,6 +27,7 @@ We release **Qwen3-TTS**, a series of powerful speech generation capabilities de
   - [Released Models Description and Download](#released-models-description-and-download)
 - [Quickstart](#quickstart)
   - [Environment Setup](#environment-setup)
+  - [Performance Tips](#performance-tips)
   - [Python Package Usage](#python-package-usage)
     - [Custom Voice Generation](#custom-voice-generate)
     - [Voice Design](#voice-design)
@@ -141,6 +142,46 @@ MAX_JOBS=4 pip install -U flash-attn --no-build-isolation
 Also, you should have hardware that is compatible with FlashAttention 2. Read more about it in the official documentation of the [FlashAttention repository](https://github.com/Dao-AILab/flash-attention). FlashAttention 2 can only be used when a model is loaded in `torch.float16` or `torch.bfloat16`.
 
 
+### Performance Tips
+
+#### Compiling the Codec Decoder with `torch.compile`
+
+The speech tokenizer codec decoder contains 100+ attention modules whose Python dispatch overhead dominates waveform decoding time—profiling shows it accounts for **~47% of single-generation time and ~85% of batch generation time**. Passing `compile_codec=True` to `from_pretrained` applies `torch.compile` to the codec, fusing these modules into optimized kernels:
+
+```python
+model = Qwen3TTSModel.from_pretrained(
+    "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+    device_map="cuda:0",
+    dtype=torch.bfloat16,
+    compile_codec=True,  # ~3-4x faster batch decoding after warmup
+)
+```
+
+**What to expect:**
+- The first generation after loading incurs a one-time compilation warmup (~30–60 s depending on hardware). Subsequent calls run at full speed.
+- Batch throughput improves by **3–4x** (e.g. from ~1.3x real-time to ~4–6x real-time for batches of 12+ utterances).
+- Single-utterance latency improves by ~30% after warmup.
+- Works on NVIDIA (CUDA), AMD (ROCm), and CPU backends.
+
+You can also pass a specific `torch.compile` mode string:
+
+```python
+model = Qwen3TTSModel.from_pretrained(
+    "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+    device_map="cuda:0",
+    dtype=torch.bfloat16,
+    compile_codec="max-autotune",  # or "reduce-overhead", "default"
+)
+```
+
+If you already have a model instance, you can compile the codec after construction:
+
+```python
+model._compile_codec()  # defaults to mode="max-autotune"
+```
+
+> **Recommendation:** Use `compile_codec=True` for batch generation workloads such as audiobook production, dataset creation, or any scenario where multiple utterances are generated in sequence. For single one-off generations, the warmup cost may outweigh the benefit.
+
 ### Python Package Usage
 
 After installation, you can import `Qwen3TTSModel` to run custom voice TTS, voice design, and voice clone. The model weights can be specified either as a Hugging Face model id (recommended) or as a local directory path you downloaded. For all the `generate_*` functions below, besides the parameters shown and explicitly documented, you can also pass generation kwargs supported by Hugging Face Transformers `model.generate`, e.g., `max_new_tokens`, `top_p`, etc.

qwen_tts/inference/qwen3_tts_model.py

@@ -83,6 +83,7 @@ def __init__(self, model: Qwen3TTSForConditionalGeneration, processor, generate_
     def from_pretrained(
         cls,
         pretrained_model_name_or_path: str,
+        compile_codec: Union[bool, str] = False,
         **kwargs,
     ) -> "Qwen3TTSModel":
         """
@@ -93,10 +94,25 @@ def from_pretrained(
           2) Loads the model via AutoModel.from_pretrained(...), forwarding `kwargs` unchanged.
           3) Loads the processor via AutoProcessor.from_pretrained(model_path).
           4) Loads optional `generate_config.json` from the model directory/repo snapshot if present.
+          5) Optionally applies ``torch.compile`` to the speech tokenizer codec decoder.
 
         Args:
             pretrained_model_name_or_path (str):
                 HuggingFace repo id or local directory of the model.
+            compile_codec (bool or str):
+                If truthy, apply ``torch.compile`` to the speech tokenizer codec
+                decoder for faster waveform decoding.
+
+                * ``False`` (default) — no compilation.
+                * ``True`` — compile with ``mode="max-autotune"`` and ``dynamic=True``.
+                * A string (``"max-autotune"``, ``"reduce-overhead"``, ``"default"``) —
+                  compile with the given mode and ``dynamic=True``.
+
+                The codec decoder contains 100+ attention modules whose Python
+                dispatch overhead dominates batch decoding time.  Compilation fuses
+                these into optimized kernels and can improve batch throughput by
+                3–4x.  The first call after compilation incurs a one-time warmup
+                cost (~30–60 s depending on hardware).
             **kwargs:
                 Forwarded as-is into `AutoModel.from_pretrained(...)`.
                 Typical examples: device_map="cuda:0", dtype=torch.bfloat16, attn_implementation="flash_attention_2".
@@ -118,7 +134,33 @@ def from_pretrained(
         processor = AutoProcessor.from_pretrained(pretrained_model_name_or_path, fix_mistral_regex=True,)
 
         generate_defaults = model.generate_config
-        return cls(model=model, processor=processor, generate_defaults=generate_defaults)
+        instance = cls(model=model, processor=processor, generate_defaults=generate_defaults)
+
+        if compile_codec:
+            instance._compile_codec(compile_codec)
+
+        return instance
+
+    def _compile_codec(self, mode: Union[bool, str] = True) -> None:
+        """Apply ``torch.compile`` to the speech tokenizer codec for faster decoding.
+
+        The codec decoder contains 100+ attention modules that benefit greatly
+        from compilation, as it eliminates per-module Python dispatch overhead.
+        Profiling shows the codec accounts for ~47% of single-generation time
+        and ~85% of batch generation time.  Compilation can improve batch
+        throughput by 3–4x.
+
+        Args:
+            mode:
+                ``True`` to compile with ``mode="max-autotune"``, or a
+                ``torch.compile`` mode string (``"max-autotune"``,
+                ``"reduce-overhead"``, ``"default"``).
+        """
+        compile_mode = "max-autotune" if mode is True else str(mode)
+        codec = self.model.speech_tokenizer.model
+        self.model.speech_tokenizer.model = torch.compile(
+            codec, mode=compile_mode, dynamic=True,
+        )
 
     def _supported_languages_set(self) -> Optional[set]:
         langs = getattr(self.model, "get_supported_languages", None)