diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index c93d027d7..eee751e16 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/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, diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index d78e16b9c..7f3fdfd7a 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/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