[Refactor][Perf] Qwen3-TTS: re-prefill Code Predictor with torch.compile + enable Code2Wav decoder CUDA Graph (#1617)

Signed-off-by: Sy03 <1370724210@qq.com>

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py

@@ -4,7 +4,6 @@ import os
 from collections.abc import Iterable
 from typing import Any
 
-import numpy as np
 import torch
 import torch.nn as nn
 from transformers.utils.hub import cached_file
@@ -21,7 +20,8 @@ logger = init_logger(__name__)
 
 class Qwen3TTSCode2Wav(nn.Module):
     """Stage-1 code2wav model for Qwen3-TTS (GenerationModelRunner).
-    Consumes frame-aligned codec tokens from input_ids and decodes waveform via SpeechTokenizer."""
+    Consumes frame-aligned codec tokens from input_ids and decodes waveform
+    via the SpeechTokenizer decoder directly (bypassing HF wrapper overhead)."""
 
     input_modalities = "audio"
 
@@ -34,13 +34,13 @@ class Qwen3TTSCode2Wav(nn.Module):
         self.has_preprocess = False
         self.has_postprocess = False
         self.enable_update_additional_information = True
-        # Generation-only stage (no logits / sampling).
         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._decode_upsample_rate: int | None = None
         self._output_sample_rate: int | None = None
+        self._total_upsample: int | None = None
         self._logged_codec_stats = False
 
     @staticmethod
@@ -52,18 +52,15 @@ class Qwen3TTSCode2Wav(nn.Module):
                 return buf.device
             return torch.device("cpu")
 
-    def _ensure_speech_tokenizer_loaded(self) -> Qwen3TTSTokenizer:
-        if self._speech_tokenizer is not None:
-            return self._speech_tokenizer
+    def _ensure_speech_tokenizer_loaded(self) -> None:
+        if self._decoder is not None:
+            return
 
-        # Locate speech_tokenizer dir from HF cache (or local path).
         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)
 
-        # Stage-1 only needs decode; skip HF feature extractor to avoid heavy optional deps.
-        # Still require preprocessor_config.json (use cached_file so online runs can fetch it).
         prep_cfg = cached_file(self.model_path, "speech_tokenizer/preprocessor_config.json")
         if prep_cfg is None:
             raise ValueError(
@@ -77,12 +74,10 @@ class Qwen3TTSCode2Wav(nn.Module):
             load_feature_extractor=False,
         )
 
-        # Align device with vLLM worker, then read back from module.
         if tok.model is not None:
             tok.model.to(device=self.vllm_config.device_config.device)
             tok.device = self._module_device(tok.model)
 
-        # Derive codec group count and rates from tokenizer config.
         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:
@@ -103,11 +98,39 @@ class Qwen3TTSCode2Wav(nn.Module):
         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._decode_upsample_rate = upsample
         self._output_sample_rate = out_sr
-        return tok
+        self._total_upsample = int(decoder.total_upsample)
+
+        if hasattr(decoder, "enable_cudagraph"):
+            device = self._module_device(decoder)
+            if device.type == "cuda":
+                try:
+                    capture_sizes = None
+                    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 chunk_frames > 0 and left_frames >= 0:
+                            from .cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper
+
+                            steady_window = left_frames + chunk_frames
+                            capture_sizes = sorted({*CUDAGraphDecoderWrapper.DEFAULT_CAPTURE_SIZES, steady_window})
+                    decoder.enable_cudagraph(capture_sizes=capture_sizes, device=device)
+                    logger.info("Code2Wav decoder CUDA Graph enabled")
+                except Exception:
+                    logger.warning("Failed to enable CUDA Graph for Code2Wav decoder", exc_info=True)
 
     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.
@@ -155,17 +178,19 @@ class Qwen3TTSCode2Wav(nn.Module):
         input_ids layout per request: [codec_context_frames, *flat_codes]
         where flat_codes is codebook-major [q*F].
 
-        When batched, uses forward context ubatch_slices to split the
-        concatenated input_ids and decode via a single batched forward pass.
+        Bypasses the HF Qwen3TTSTokenizer.decode() wrapper and calls the
+        decoder.chunked_decode() directly to avoid GPU->CPU->GPU round-trips.
+        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._decode_upsample_rate is not None
-        assert self._output_sample_rate is not None
+        assert self._total_upsample is not None
 
-        tok = self._speech_tokenizer
+        decoder = self._decoder
         q = int(self._num_quantizers)
-        upsample = int(self._decode_upsample_rate)
+        upsample = int(self._total_upsample)
         sr_val = int(self._output_sample_rate)
         sr_tensor = torch.tensor(sr_val, dtype=torch.int32)
         empty = torch.zeros((0,), dtype=torch.float32)
@@ -179,13 +204,9 @@ class Qwen3TTSCode2Wav(nn.Module):
         ids = input_ids.reshape(-1).to(dtype=torch.long)
         request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts"))
 
-        # Parse each request: extract left_context_size, validate, reshape codes.
-        # input_ids layout per request: [*flat_codes]
-        # where flat_codes is codebook-major [q*F].
-        parsed = []  # (left_context_size, actual_frames)
-        valid_codes = []
-        valid_indices = []
-        # Get left_context_size from runtime_additional_information (passed via kwargs)
+        parsed: list[tuple[int, int]] = []
+        valid_codes_qf: list[torch.Tensor] = []
+        valid_indices: list[int] = []
         left_context_size = [0] * len(request_ids_list)
         if runtime_additional_information is not None:
             for i, info in enumerate(runtime_additional_information):
@@ -193,8 +214,6 @@ class Qwen3TTSCode2Wav(nn.Module):
                     break
                 if "left_context_size" in info:
                     left_context_size[i] = info["left_context_size"]
-        else:
-            logger.debug("No additional_information provided to code2wav stage.")
         for i, req_ids in enumerate(request_ids_list):
             if req_ids.numel() < 1:
                 parsed.append((0, 0))
@@ -202,7 +221,6 @@ class Qwen3TTSCode2Wav(nn.Module):
             ctx_frames = left_context_size[i]
             flat = req_ids
             n = flat.numel()
-            # Warmup / dummy_run: not divisible by num_quantizers.
             if n == 0 or n % q != 0:
                 if n > 0:
                     logger.warning(
@@ -214,14 +232,14 @@ class Qwen3TTSCode2Wav(nn.Module):
                 parsed.append((0, 0))
                 continue
             frames = n // q
-            # Reshape codebook-major flat [q*F] -> [q, F] -> [F, q] for SpeechTokenizer.
-            codes_fq = flat.reshape(q, frames).transpose(0, 1).contiguous()
+            # [q*F] -> [Q, F] for direct decoder call (decoder expects [B, Q, F])
+            codes_qf = flat.reshape(q, frames)
             parsed.append((ctx_frames, frames))
-            valid_codes.append({"audio_codes": codes_fq})
+            valid_codes_qf.append(codes_qf)
             valid_indices.append(i)
 
         num_req = len(request_ids_list)
-        if not valid_codes:
+        if not valid_codes_qf:
             return OmniOutput(
                 text_hidden_states=None,
                 multimodal_outputs={
@@ -233,50 +251,55 @@ class Qwen3TTSCode2Wav(nn.Module):
         if not self._logged_codec_stats:
             self._logged_codec_stats = True
             try:
-                c = valid_codes[0]["audio_codes"]
+                c = valid_codes_qf[0]
                 logger.info(
-                    "Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] head=%s batch=%d",
-                    c.shape[0],
+                    "Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] batch=%d",
+                    c.shape[1],
                     q,
                     int(torch.unique(c).numel()),
                     int(c.min().item()),
                     int(c.max().item()),
-                    c[: min(2, c.shape[0]), : min(8, q)].cpu().tolist(),
-                    len(valid_codes),
+                    len(valid_codes_qf),
                 )
             except Exception:
                 pass
 
-        # Batched decode: single forward pass through SpeechTokenizer.
-        wavs, _ = tok.decode(valid_codes)
-        if len(wavs) != len(valid_codes):
-            raise RuntimeError(f"Code2Wav returned {len(wavs)} waveforms for {len(valid_codes)} requests")
+        # Decode directly via decoder.chunked_decode(), staying entirely on GPU.
+        # For single request: no padding needed, fast path.
+        # For multiple requests: decode each individually to avoid padding overhead.
+        wav_tensors: list[torch.Tensor] = []
+        if len(valid_codes_qf) == 1:
+            codes_bqf = valid_codes_qf[0].unsqueeze(0)  # [1, Q, F]
+            wav = decoder.chunked_decode(codes_bqf)  # [1, 1, wav_len]
+            wav_tensors.append(wav.squeeze(0).squeeze(0))  # [wav_len]
+        else:
+            for codes_qf in valid_codes_qf:
+                codes_bqf = codes_qf.unsqueeze(0)  # [1, Q, F]
+                wav = decoder.chunked_decode(codes_bqf)
+                wav_tensors.append(wav.squeeze(0).squeeze(0))
 
-        # Build per-request outputs, trimming padding and left-context.
-        audios = [empty] * num_req
+        audios: list[torch.Tensor] = [empty] * num_req
         srs = [sr_tensor] * num_req
 
         for j, idx in enumerate(valid_indices):
             ctx_frames, actual_frames = parsed[idx]
-            audio_np = wavs[j].astype(np.float32, copy=False)
-            # Trim decoder padding (output may be longer due to batch padding).
+            wav = wav_tensors[j]
             expected_len = actual_frames * upsample
-            if audio_np.shape[0] > expected_len:
-                audio_np = audio_np[:expected_len]
-            # Trim left-context waveform samples (streaming sliding window).
+            if wav.shape[0] > expected_len:
+                wav = wav[:expected_len]
             if ctx_frames > 0:
                 cut = ctx_frames * upsample
-                if cut < audio_np.shape[0]:
-                    audio_np = audio_np[cut:]
+                if cut < wav.shape[0]:
+                    wav = wav[cut:]
                 else:
                     logger.warning(
                         "Context trim %d >= decoded length %d; returning empty audio.",
                         cut,
-                        audio_np.shape[0],
+                        wav.shape[0],
                     )
                     continue
-            if audio_np.shape[0] > 0:
-                audios[idx] = torch.from_numpy(audio_np).to(dtype=torch.float32).reshape(-1)
+            if wav.shape[0] > 0:
+                audios[idx] = wav.to(dtype=torch.float32).reshape(-1)
 
         return OmniOutput(
             text_hidden_states=None,

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py

@@ -1,235 +1,217 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
-from typing import Any
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from vllm.config import VllmConfig
 from vllm.config.vllm import set_current_vllm_config
-from vllm.forward_context import set_forward_context
+from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
 )
-from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer
 from vllm.model_executor.models.utils import is_pp_missing_parameter
-from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, KVCacheTensor
-from vllm.v1.worker.gpu import attn_utils
-
-from vllm_omni.platforms import current_omni_platform
 
 from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig
 
+logger = init_logger(__name__)
 
-class _LocalPredictorKVCache:
-    """Minimal local KV cache + attention metadata for running
-    code_predictor inside one worker (independent of engine KV)."""
+
+# ===================================================================
+#  Standalone Code Predictor Layers (no vLLM paged attention)
+# ===================================================================
+#
+# These replace vLLM's Qwen3DecoderLayer for the code predictor.
+# Input is batch-major [B, seq_len, H], attention uses F.scaled_dot_product_attention.
+# Weight names match the checkpoint (self_attn.qkv_proj, mlp.gate_up_proj, etc.)
+# so load_weights works unchanged.
+
+
+class _CodePredictorAttention(nn.Module):
+    """Standalone multi-head attention for code predictor.
+
+    Uses F.scaled_dot_product_attention (SDPA) instead of vLLM's paged Attention.
+    Supports fused QKV, RoPE, q/k normalization, and native GQA via enable_gqa.
+    Input: [B, seq_len, hidden_size], output: [B, seq_len, hidden_size].
+    """
 
     def __init__(
         self,
+        config: Qwen3TTSTalkerCodePredictorConfig,
         *,
-        vllm_config: VllmConfig,
-        max_seq_len: int,
-        max_batch_size: int,
-        device: torch.device,
+        prefix: str = "",
     ) -> None:
-        self.vllm_config = vllm_config
-        self.device = device
-
-        # Collect attention layers registered in this vllm_config.
-        kv_cache_spec_by_layer = attn_utils.get_kv_cache_spec(vllm_config)
-        if not kv_cache_spec_by_layer:
-            raise RuntimeError("Local predictor KVCache requires vLLM Attention layers to be registered.")
-
-        # We only need enough blocks for a tiny per-frame sequence (<= max_seq_len).
-        any_spec = next(iter(kv_cache_spec_by_layer.values()))
-        block_size = int(any_spec.block_size)
-        blocks_per_seq = (int(max_seq_len) + block_size - 1) // block_size
-        num_blocks = max(1, int(max_batch_size) * int(blocks_per_seq))
-
-        # Allocate per-layer KV caches (small, independent).
-        kv_cache_tensors: list[KVCacheTensor] = []
-        for layer_name, spec in kv_cache_spec_by_layer.items():
-            kv_cache_tensors.append(KVCacheTensor(size=int(spec.page_size_bytes) * num_blocks, shared_by=[layer_name]))
-
-        merged_spec: KVCacheSpec = KVCacheSpec.merge(list(kv_cache_spec_by_layer.values()))
-        self.kv_cache_config = KVCacheConfig(
-            num_blocks=num_blocks,
-            kv_cache_tensors=kv_cache_tensors,
-            kv_cache_groups=[
-                KVCacheGroupSpec(layer_names=list(kv_cache_spec_by_layer.keys()), kv_cache_spec=merged_spec)
-            ],
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = getattr(
+            config,
+            "head_dim",
+            config.hidden_size // config.num_attention_heads,
         )
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self._use_gqa = self.num_kv_heads != self.num_heads
 
-        # Init backend + bind KV cache tensors to attention modules.
-        self.attn_backends, self.attn_metadata_builders = attn_utils.init_attn_backend(
-            self.kv_cache_config, vllm_config, device
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=self.hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.num_heads,
+            total_num_kv_heads=self.num_kv_heads,
+            bias=getattr(config, "attention_bias", False),
+            prefix=f"{prefix}.qkv_proj",
+            disable_tp=True,
         )
-        self.runner_kv_caches: list[torch.Tensor] = []
-        attn_utils.init_kv_cache(
-            self.runner_kv_caches,
-            vllm_config.compilation_config.static_forward_context,
-            self.kv_cache_config,
-            self.attn_backends,
-            device,
+        self.o_proj = RowParallelLinear(
+            input_size=self.num_heads * self.head_dim,
+            output_size=self.hidden_size,
+            bias=False,
+            prefix=f"{prefix}.o_proj",
+            disable_tp=True,
         )
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=config.max_position_embeddings,
+            rope_parameters=getattr(config, "rope_parameters", None),
+            dual_chunk_attention_config=None,
+        )
+        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
 
-        # Precompute a fixed block table mapping for the maximum batch.
-        self.block_size = block_size
-        self.blocks_per_seq = blocks_per_seq
-        self.max_batch_size = int(max_batch_size)
-
-        bt = torch.full((self.max_batch_size, self.blocks_per_seq), -1, dtype=torch.int32, device=device)
-        for i in range(self.max_batch_size):
-            for j in range(self.blocks_per_seq):
-                bt[i, j] = i * self.blocks_per_seq + j
-        self._block_table = bt
-
-    def build_attn_metadata(
+    def forward(
         self,
+        hidden_states: torch.Tensor,
+        position_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        bsz, seq_len, _ = hidden_states.shape
+
+        qkv, _ = self.qkv_proj(hidden_states.reshape(bsz * seq_len, -1))
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape)
+        k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape)
+
+        q, k = self.rotary_emb(position_ids, q, k)
+
+        q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        attn_out = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            scale=self.scaling,
+            is_causal=True,
+            enable_gqa=self._use_gqa,
+        )
+
+        attn_out = attn_out.transpose(1, 2).reshape(bsz * seq_len, -1)
+        output, _ = self.o_proj(attn_out)
+        return output.view(bsz, seq_len, -1)
+
+
+class _CodePredictorMLP(nn.Module):
+    """SiLU-gated MLP for code predictor, matching Qwen3MLP structure."""
+
+    def __init__(
+        self,
+        config: Qwen3TTSTalkerCodePredictorConfig,
         *,
-        num_reqs: int,
-        query_lens: torch.Tensor,  # (num_reqs,) int32 on cpu
-        seq_lens: torch.Tensor,  # (num_reqs,) int32 on cpu
-    ) -> tuple[dict[str, Any], torch.Tensor, dict[str, torch.Tensor]]:
-        """Build attention metadata, positions, and slot_mapping dict.
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=config.hidden_size,
+            output_sizes=[config.intermediate_size] * 2,
+            bias=False,
+            prefix=f"{prefix}.gate_up_proj",
+            disable_tp=True,
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=config.intermediate_size,
+            output_size=config.hidden_size,
+            bias=False,
+            prefix=f"{prefix}.down_proj",
+            disable_tp=True,
+        )
 
-        Returns:
-            (attn_metadata, positions, slot_mappings_by_layer)
-            - attn_metadata: per-layer attention metadata for attn backends.
-            - positions: (num_tokens,) position IDs on device.
-            - slot_mappings_by_layer: {layer_name: slot_mapping_tensor} for
-              set_forward_context so that unified_kv_cache_update can write
-              the KV cache correctly.
-        """
-        num_reqs = int(num_reqs)
-        if num_reqs <= 0:
-            return {}, torch.empty((0,), dtype=torch.int64, device=self.device), {}
-        if num_reqs > self.max_batch_size:
-            raise ValueError(f"num_reqs={num_reqs} exceeds local predictor max_batch_size={self.max_batch_size}")
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_up, _ = self.gate_up_proj(x)
+        gate, up = gate_up.chunk(2, dim=-1)
+        x = F.silu(gate) * up
+        x, _ = self.down_proj(x)
+        return x
 
-        query_lens_i32 = query_lens.to(dtype=torch.int32, device="cpu")
-        seq_lens_i32 = seq_lens.to(dtype=torch.int32, device="cpu")
 
-        # query_start_loc: prefix sums of query_lens.
-        qsl = torch.zeros((num_reqs + 1,), dtype=torch.int32, device="cpu")
-        qsl[1:] = torch.cumsum(query_lens_i32, dim=0)
-        num_tokens = int(qsl[-1].item())
-        if num_tokens <= 0:
-            return {}, torch.empty((0,), dtype=torch.int64, device=self.device), {}
+class _CodePredictorDecoderLayer(nn.Module):
+    """Transformer decoder layer for code predictor (SDPA, no KV cache)."""
 
-        # Convert tensors to Python lists once to avoid per-element .item()
-        # calls (each .item() on a GPU tensor forces a GPU→CPU sync).
-        query_lens_list = query_lens_i32[:num_reqs].tolist()
-        seq_lens_list = seq_lens_i32[:num_reqs].tolist()
-        block_table_cpu = self._block_table[:num_reqs].cpu().tolist()
+    def __init__(
+        self,
+        config: Qwen3TTSTalkerCodePredictorConfig,
+        *,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.self_attn = _CodePredictorAttention(config, prefix=f"{prefix}.self_attn")
+        self.mlp = _CodePredictorMLP(config, prefix=f"{prefix}.mlp")
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
-        # positions: for each request i, emit positions [seq_len-query_len .. seq_len-1]
-        pos_list: list[torch.Tensor] = []
-        for i in range(num_reqs):
-            ql = query_lens_list[i]
-            sl = seq_lens_list[i]
-            start = sl - ql
-            pos_list.append(torch.arange(start, sl, dtype=torch.int64))
-        positions_cpu = torch.cat(pos_list, dim=0)
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(hidden_states, position_ids)
+        hidden_states = residual + hidden_states
 
-        # slot_mapping: map each query token to a physical slot in the paged KV cache.
-        # We allocate per-request contiguous blocks; slot = base + position.
-        slot_mapping = torch.empty((num_tokens,), dtype=torch.int64, device="cpu")
-        cursor = 0
-        for i in range(num_reqs):
-            ql = query_lens_list[i]
-            sl = seq_lens_list[i]
-            start = sl - ql
-            for p in range(start, sl):
-                block_idx = p // self.block_size
-                offset = p % self.block_size
-                block_id = block_table_cpu[i][block_idx]
-                slot_mapping[cursor] = block_id * self.block_size + offset
-                cursor += 1
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
 
-        max_seq_len = int(seq_lens_i32[:num_reqs].max().item())
-        query_start_loc_gpu = qsl.to(device=self.device)
-        seq_lens_gpu = seq_lens_i32.to(device=self.device)
-        block_table = self._block_table[:num_reqs].contiguous()
-        slot_mapping_gpu = slot_mapping.to(device=self.device)
 
-        # FIXME(gcanlin): Refactor build_attn_metadata to avoid special-casing NPU backends here.
-        if current_omni_platform.is_npu():
-            # NPU requires AscendCommonAttentionMetadata with extra attributes
-            from vllm_ascend.worker.v2 import attn_utils as attn_utils_npu
-
-            max_query_len = int(query_lens_i32[:num_reqs].max().item())
-            # NPU version expects slot_mappings as a stacked tensor, not a list
-            slot_mappings_tensor = slot_mapping_gpu.unsqueeze(0)
-            attn_metadata = attn_utils_npu.build_attn_metadata(
-                attn_metadata_builders=self.attn_metadata_builders,
-                num_reqs=num_reqs,
-                num_tokens=num_tokens,
-                query_start_loc_gpu=query_start_loc_gpu,
-                query_start_loc_cpu=qsl,
-                max_query_len=max_query_len,
-                seq_lens=seq_lens_gpu,
-                max_seq_len=max_seq_len,
-                block_tables=[block_table],
-                slot_mappings=slot_mappings_tensor,
-                kv_cache_config=self.kv_cache_config,
-            )
-        else:
-            attn_metadata = attn_utils.build_attn_metadata(
-                self.attn_metadata_builders,
-                num_reqs=num_reqs,
-                num_tokens=num_tokens,
-                query_start_loc_gpu=query_start_loc_gpu,
-                query_start_loc_cpu=qsl,
-                seq_lens=seq_lens_gpu,
-                max_seq_len=max_seq_len,
-                block_tables=[block_table],
-                slot_mappings=[slot_mapping_gpu],
-                kv_cache_config=self.kv_cache_config,
-            )
-
-        # Build slot_mappings_by_layer for set_forward_context.
-        # Fix for vllm 0.15.0
-        slot_mappings_by_layer: dict[str, torch.Tensor] = {}
-        for kv_cache_group in self.kv_cache_config.kv_cache_groups:
-            for layer_name in kv_cache_group.layer_names:
-                slot_mappings_by_layer[layer_name] = slot_mapping_gpu
-
-        return attn_metadata, positions_cpu.to(device=self.device), slot_mappings_by_layer
+# ===================================================================
+#  Code Predictor Transformer Model
+# ===================================================================
 
 
 class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module):
+    """Transformer model for the code predictor (re-prefill, no KV cache)."""
+
     def __init__(
         self,
         config: Qwen3TTSTalkerCodePredictorConfig,
         *,
         talker_hidden_size: int | None = None,
-        cache_config=None,
-        quant_config=None,
         prefix: str = "",
     ) -> None:
         super().__init__()
         self.config = config
-        self.quant_config = quant_config
 
         self.layers = nn.ModuleList(
-            [
-                Qwen3DecoderLayer(
-                    config, cache_config=cache_config, quant_config=quant_config, prefix=f"{prefix}.layers.{i}"
-                )
-                for i in range(config.num_hidden_layers)
-            ]
+            [_CodePredictorDecoderLayer(config, prefix=f"{prefix}.layers.{i}") for i in range(config.num_hidden_layers)]
         )
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
-        # Official code_predictor uses one embedding table per residual group.
-        # Some Qwen3-TTS checkpoints store codec embeddings in the talker hidden
-        # space, even when `code_predictor_config.hidden_size` is smaller.
-        # We keep the embedding dim aligned with the checkpoint and project down
-        # via `small_to_mtp_projection` in the wrapper module.
+        # Codec embeddings: one per residual group. Stored in talker hidden dim
+        # (some checkpoints use talker_hidden_size != code_predictor hidden_size).
         emb_dim = int(talker_hidden_size) if talker_hidden_size is not None else int(config.hidden_size)
         self.codec_embedding = nn.ModuleList(
             [nn.Embedding(config.vocab_size, emb_dim) for _ in range(config.num_code_groups - 1)]
@@ -238,18 +220,18 @@ class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module):
     def get_input_embeddings(self) -> nn.ModuleList:
         return self.codec_embedding
 
-    def forward(self, positions: torch.Tensor, inputs_embeds: torch.Tensor) -> torch.Tensor:
-        # Token-major: [num_tokens, hidden]
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        position_ids: torch.Tensor,
+    ) -> torch.Tensor:
         hidden_states = inputs_embeds
-        residual = None
         for layer in self.layers:
-            hidden_states, residual = layer(positions, hidden_states, residual)
-        hidden_states, _ = self.norm(hidden_states, residual)
+            hidden_states = layer(hidden_states, position_ids)
+        hidden_states = self.norm(hidden_states)
         return hidden_states
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        # Match vLLM Qwen2/Qwen3 packing conventions: q_proj/k_proj/v_proj -> qkv_proj,
-        # gate_proj/up_proj -> gate_up_proj.
         stacked_params_mapping = [
             ("qkv_proj", "q_proj", "q"),
             ("qkv_proj", "k_proj", "k"),
@@ -262,13 +244,6 @@ class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module):
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                loaded_params.add(scale_name)
-                continue
 
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
@@ -309,8 +284,31 @@ class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module):
         return loaded_params
 
 
+# ===================================================================
+#  Code Predictor Wrapper (optimized re-prefill + torch.compile)
+# ===================================================================
+
+
 class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
-    """vLLM-native code_predictor used by the AR talker (residual codebooks)."""
+    """vLLM-native code_predictor for the AR talker (residual codebooks).
+
+    Re-prefill approach: each AR step forwards the full growing sequence
+    through the 5-layer transformer. No KV cache needed. This trades
+    ~O(T^2) extra attention FLOPs (negligible for T=16, 5 layers) for
+    zero KV cache management overhead and a simpler execution model.
+
+    Optimizations over baseline:
+      1. torch.compile on model forward -- fuses 60+ small kernel launches per step
+         into fewer fused kernels (4x speedup on model_fwd, ~75% of total time).
+      2. Pre-allocated embedding buffer [B, max_seq, H] -- no torch.cat per step.
+      3. Projection caching -- each token projected once and cached, avoids O(T^2)
+         redundant projections.
+      4. Pre-allocated position_ids [max_seq] -- no torch.arange per step.
+      5. Inline sampling -- no custom op / forward_context overhead.
+      6. No context managers in forward().
+      7. Cached module references -- bypass nn.Module.__call__ and ModuleList indexing.
+      8. Pre-allocated output tensor.
+    """
 
     def __init__(
         self,
@@ -325,16 +323,12 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
         self.config = config
         self.talker_config = talker_config
 
-        # Keep module/weight names aligned with official checkpoint (talker.code_predictor.model.*).
         self.model = Qwen3TTSTalkerCodePredictorModelVLLM(
             config,
             talker_hidden_size=int(talker_config.hidden_size),
-            cache_config=vllm_config.cache_config,
-            quant_config=vllm_config.quant_config,
             prefix=f"{prefix}.model",
         )
 
-        # One head per residual group.
         self.lm_head = nn.ModuleList(
             [nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(config.num_code_groups - 1)]
         )
@@ -344,14 +338,21 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
         else:
             self.small_to_mtp_projection = nn.Identity()
 
-        self._kv_cache: _LocalPredictorKVCache | None = None
+        self._num_groups = int(config.num_code_groups)
+        self._talker_hidden = int(talker_config.hidden_size)
+        self._cp_hidden = int(config.hidden_size)
+
+        # Pre-allocated buffers (lazily initialized on first forward).
+        self._proj_buf: torch.Tensor | None = None
+        self._pos_ids: torch.Tensor | None = None
+
+        # torch.compile: fuse small kernels in the 5-layer transformer.
+        self._compiled_model_fwd: object | None = None
 
     def get_input_embeddings(self) -> nn.ModuleList:
         return self.model.get_input_embeddings()
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        # Ensure all vLLM custom layers consult the predictor vllm_config
-        # (esp. for Attention static_forward_context).
         with set_current_vllm_config(self._vllm_config):
             loaded: set[str] = set()
             model_weights: list[tuple[str, torch.Tensor]] = []
@@ -373,95 +374,52 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
                 loaded.add(name)
             return loaded
 
-    def _maybe_init_kv_cache(self, device: torch.device) -> None:
-        if self._kv_cache is not None:
+    # ------------------------------------------------------------------
+    #  Pre-allocated buffer management
+    # ------------------------------------------------------------------
+
+    def _ensure_buffers(self, bsz: int, device: torch.device, dtype: torch.dtype) -> None:
+        max_seq = self._num_groups + 1
+        if (
+            self._proj_buf is not None
+            and self._proj_buf.shape[0] >= bsz
+            and self._proj_buf.device == device
+            and self._proj_buf.dtype == dtype
+        ):
             return
-        with set_current_vllm_config(self._vllm_config):
-            max_seq_len = int(getattr(self.config, "num_code_groups", 16) or 16)
-            # Upper bound on batch size: vLLM scheduler max_num_seqs (fallback 8).
-            max_batch = int(getattr(self._vllm_config.scheduler_config, "max_num_seqs", 8) or 8)
-            max_batch = max(1, max_batch)
-            self._kv_cache = _LocalPredictorKVCache(
-                vllm_config=self._vllm_config,
-                max_seq_len=max_seq_len,
-                max_batch_size=max_batch,
-                device=device,
-            )
-
-    @torch.inference_mode()
-    def reset_cache(self) -> None:
-        # We reuse a fixed kv cache buffer and overwrite starting at slot 0.
-        # No action required here (seq_lens controls what is read).
-        return
-
-    @torch.inference_mode()
-    def prefill_logits(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
-        """Prefill with 2 tokens: [past_hidden, layer0_embed]. Returns logits for residual group 0."""
-        self._maybe_init_kv_cache(inputs_embeds.device)
-        assert self._kv_cache is not None
-
-        bsz = int(inputs_embeds.shape[0])
-        qlen = 2
-        # Flatten to token-major.
-        model_dtype = next(self.model.parameters()).dtype
-        hs = inputs_embeds.to(dtype=model_dtype).reshape(bsz * qlen, -1)
-        hs = self.small_to_mtp_projection(hs)
-
-        query_lens = torch.full((bsz,), qlen, dtype=torch.int32)
-        seq_lens = query_lens.clone()
-        attn_metadata, positions, slot_mappings = self._kv_cache.build_attn_metadata(
-            num_reqs=bsz, query_lens=query_lens, seq_lens=seq_lens
+        self._proj_buf = torch.zeros(
+            bsz,
+            max_seq,
+            self._cp_hidden,
+            dtype=dtype,
+            device=device,
+        )
+        self._pos_ids = torch.arange(
+            max_seq,
+            dtype=torch.long,
+            device=device,
         )
 
-        with (
-            set_current_vllm_config(self._vllm_config),
-            set_forward_context(
-                attn_metadata,
-                self._vllm_config,
-                num_tokens=int(hs.shape[0]),
-                slot_mapping=slot_mappings,
-            ),
-        ):
-            out = self.model(positions=positions, inputs_embeds=hs)
+    def _setup_compile(self) -> None:
+        """Lazily set up torch.compiled model forward for kernel fusion.
 
-        # Gather last token per request.
-        last_idx = torch.arange(qlen - 1, bsz * qlen, step=qlen, device=out.device, dtype=torch.long)
-        last_h = out.index_select(0, last_idx)
-        logits = self.lm_head[0](last_h)
-        return logits
-
-    @torch.inference_mode()
-    def decode_logits(self, input_ids: torch.Tensor, *, generation_step: int, past_seq_len: int) -> torch.Tensor:
-        """Decode one new token for residual group `generation_step` (1..Q-1)."""
-        self._maybe_init_kv_cache(input_ids.device)
-        assert self._kv_cache is not None
-        bsz = int(input_ids.shape[0])
-        if generation_step <= 0:
-            raise ValueError("generation_step must be >= 1 for decode_logits")
-
-        embed_idx = generation_step - 1
-        hs = self.model.get_input_embeddings()[embed_idx](input_ids.to(dtype=torch.long).reshape(bsz, 1))
-        hs = self.small_to_mtp_projection(hs.reshape(bsz, -1))
-
-        query_lens = torch.ones((bsz,), dtype=torch.int32)
-        seq_lens = torch.full((bsz,), int(past_seq_len) + 1, dtype=torch.int32)
-        attn_metadata, positions, slot_mappings = self._kv_cache.build_attn_metadata(
-            num_reqs=bsz, query_lens=query_lens, seq_lens=seq_lens
+        Uses ``mode="default"`` so Inductor performs operator fusion without
+        capturing its own CUDA graphs.  This avoids conflicts with vLLM's
+        ``CUDAGraphWrapper`` which manages CUDA graphs for the main Talker
+        model on the default stream.
+        """
+        if self._compiled_model_fwd is not None:
+            return
+        self._compiled_model_fwd = torch.compile(
+            self.model.forward,
+            mode="default",
+            dynamic=True,
         )
+        logger.info("code_predictor: torch.compile enabled (mode=default)")
 
-        with (
-            set_current_vllm_config(self._vllm_config),
-            set_forward_context(
-                attn_metadata,
-                self._vllm_config,
-                num_tokens=int(hs.shape[0]),
-                slot_mapping=slot_mappings,
-            ),
-        ):
-            out = self.model(positions=positions, inputs_embeds=hs)
-
-        logits = self.lm_head[generation_step](out)
-        return logits
+    # ------------------------------------------------------------------
+    #  Optimized forward: re-prefill + torch.compile + projection cache
+    # ------------------------------------------------------------------
 
     @torch.inference_mode()
     def forward(
@@ -470,62 +428,71 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
         layer0_embed: torch.Tensor,
         last_talker_hidden: torch.Tensor,
         do_sample: bool = True,
-        temperature: float = 1.0,
+        temperature: float = 0.9,
         top_k: int = 50,
         top_p: float = 1.0,
     ) -> torch.Tensor:
-        """Full autoregressive prediction of residual codebooks 1..Q-1.
+        """Predict residual codebooks 1..Q-1 autoregressively via re-prefill.
 
-        Args:
-            layer0_code: [B, 1] first-layer codec token ids.
-            layer0_embed: [B, 1, H] embedding of layer0_code.
-            last_talker_hidden: [B, 1, H] hidden state from the talker.
-            do_sample: whether to sample or take argmax.
-            temperature: sampling temperature.
-            top_k: top-k filtering.
-            top_p: top-p (nucleus) filtering.
+        torch.compile fuses the ~60 small kernel launches per step into fewer
+        fused kernels, reducing kernel launch overhead by ~75%.
 
-        Returns:
-            audio_codes: [B, Q] all codebook tokens (layer0 + residuals).
+        Projection caching: each token is projected once via small_to_mtp_projection
+        and cached in _proj_buf, avoiding redundant re-projection of past tokens.
         """
         bsz = int(layer0_code.shape[0])
-        num_groups = int(self.config.num_code_groups)
-        max_steps = num_groups - 1
+        num_groups = self._num_groups
+        device = layer0_code.device
+        dtype = layer0_embed.dtype
 
-        # Reset KV cache for a fresh sequence.
-        self.reset_cache()
+        all_codes = torch.empty(bsz, num_groups, dtype=torch.long, device=device)
+        all_codes[:, 0] = layer0_code.reshape(bsz)
 
-        # Prefill: feed [last_talker_hidden, layer0_embed] → logits for group 1.
-        prefill_input = torch.cat([last_talker_hidden, layer0_embed], dim=1)  # [B, 2, H]
-        logits = self.prefill_logits(prefill_input)  # [B, vocab]
+        self._ensure_buffers(bsz, device, dtype)
+        self._setup_compile()
 
-        all_codes = [layer0_code.reshape(bsz, 1)]
-        past_seq_len = 2
+        proj_buf = self._proj_buf
+        pos_ids = self._pos_ids
+
+        projection = self.small_to_mtp_projection
+        model_fwd = self._compiled_model_fwd
+        lm_heads = list(self.lm_head)
+        codec_embeds = list(self.model.codec_embedding)
+
+        proj_buf[:bsz, 0, :] = projection(last_talker_hidden.reshape(bsz, 1, -1)).reshape(bsz, -1)
+        proj_buf[:bsz, 1, :] = projection(layer0_embed.reshape(bsz, 1, -1)).reshape(bsz, -1)
+
+        use_sampling = do_sample and temperature > 0
+        inv_temperature = 1.0 / max(temperature, 1e-6) if use_sampling else 0.0
+        if use_sampling and top_p != 1.0:
+            raise NotImplementedError(
+                "top_p sampling is not implemented for the vLLM-native code predictor; please set top_p=1.0."
+            )
 
         for step in range(1, num_groups):
-            # Sample or argmax from logits.
-            if do_sample and temperature > 0:
-                scaled = logits.float() / temperature
+            seq_len = step + 1
+
+            projected = proj_buf[:bsz, :seq_len, :]
+            step_pos_ids = pos_ids[:seq_len] if bsz == 1 else pos_ids[:seq_len].repeat(bsz)
+
+            hidden_out = model_fwd(projected, step_pos_ids)
+
+            logits = lm_heads[step - 1](hidden_out[:, -1, :])
+
+            if use_sampling:
+                scaled = logits * inv_temperature
                 if top_k > 0:
                     topk_vals, _ = scaled.topk(top_k, dim=-1)
                     scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
-                probs = torch.softmax(scaled, dim=-1)
-                # Guard against NaN/Inf from lower-precision dtypes.
-                probs = probs.clamp(min=0.0)
-                row_sums = probs.sum(dim=-1, keepdim=True)
-                probs = torch.where(row_sums > 0, probs / row_sums, torch.ones_like(probs) / probs.shape[-1])
-                next_ids = torch.multinomial(probs, num_samples=1)  # [B, 1]
+                probs = F.softmax(scaled, dim=-1)
+                next_ids = torch.multinomial(probs, num_samples=1)
             else:
-                next_ids = logits.argmax(dim=-1, keepdim=True)  # [B, 1]
-            all_codes.append(next_ids)
+                next_ids = logits.argmax(dim=-1, keepdim=True)
 
-            # If not the last step, decode one more token.
-            if step < max_steps:
-                logits = self.decode_logits(
-                    next_ids.reshape(bsz),
-                    generation_step=step,
-                    past_seq_len=past_seq_len,
-                )
-                past_seq_len += 1
+            all_codes[:, step] = next_ids.reshape(bsz)
 
-        return torch.cat(all_codes, dim=1)  # [B, Q]
+            if step < num_groups - 1:
+                new_embed = codec_embeds[step - 1](next_ids)
+                proj_buf[:bsz, step + 1, :] = projection(new_embed.reshape(bsz, 1, -1)).reshape(bsz, -1)
+
+        return all_codes