[Feature] : Support disaggregated inference pipeline for Qwen3_TTS (#1161)

Signed-off-by: Sy03 <1370724210@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-21 07:36:44 +08:00
1 parent 5295f0f81b
commit efbe4116ff
@@ -176,7 +176,7 @@ steps:
  #                 type: DirectoryOrCreate

  - label: "Omni Model Test"
-    timeout_in_minutes: 15
+    timeout_in_minutes: 20
    depends_on: image-build
    commands:
      - export VLLM_LOGGING_LEVEL=DEBUG
@@ -4,18 +4,21 @@ Provides single and batch sample inputs for CustomVoice, VoiceDesign, and Base
 tasks, then runs Omni generation and saves output wav files.
 """

+import logging
 import os
-from typing import NamedTuple
+from typing import Any, NamedTuple

 import soundfile as sf
+import torch

 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

-from vllm import SamplingParams
 from vllm.utils.argparse_utils import FlexibleArgumentParser

 from vllm_omni import Omni

+logger = logging.getLogger(__name__)
+

 class QueryResult(NamedTuple):
    """Container for a prepared Omni request."""
@@ -24,6 +27,44 @@ class QueryResult(NamedTuple):
    model_name: str


+def _estimate_prompt_len(
+    additional_information: dict[str, Any],
+    model_name: str,
+    _cache: dict[str, Any] = {},
+) -> int:
+    """Estimate prompt_token_ids placeholder length for the Talker stage.
+
+    The AR Talker replaces all input embeddings via ``preprocess``, so the
+    placeholder values are irrelevant but the **length** must match the
+    embeddings that ``preprocess`` will produce.
+    """
+    try:
+        from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
+        from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+            Qwen3TTSTalkerForConditionalGeneration,
+        )
+
+        if model_name not in _cache:
+            from transformers import AutoTokenizer
+
+            tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left")
+            cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
+            _cache[model_name] = (tok, getattr(cfg, "talker_config", None))
+
+        tok, tcfg = _cache[model_name]
+        task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
+        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
+            additional_information=additional_information,
+            task_type=task_type,
+            tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
+            codec_language_id=getattr(tcfg, "codec_language_id", None),
+            spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
+        )
+    except Exception as exc:
+        logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc)
+        return 2048
+
+
 def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
    """Build CustomVoice sample inputs.

@@ -34,6 +75,7 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
        QueryResult with Omni inputs and the CustomVoice model path.
    """
    task_type = "CustomVoice"
+    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
    if use_batch_sample:
        texts = ["其实我真的有发现，我是一个特别善于观察别人情绪的人。", "She said she would be here by noon."]
        instructs = ["", "Very happy."]
@@ -41,18 +83,18 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
        speakers = ["Vivian", "Ryan"]
        inputs = []
        for text, instruct, language, speaker in zip(texts, instructs, languages, speakers):
-            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+            additional_information = {
+                "task_type": [task_type],
+                "text": [text],
+                "instruct": [instruct],
+                "language": [language],
+                "speaker": [speaker],
+                "max_new_tokens": [2048],
+            }
            inputs.append(
                {
-                    "prompt": prompt,
-                    "additional_information": {
-                        "task_type": [task_type],
-                        "text": [text],
-                        "instruct": [instruct],
-                        "language": [language],
-                        "speaker": [speaker],
-                        "max_new_tokens": [2048],
-                    },
+                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+                    "additional_information": additional_information,
                }
            )
    else:
@@ -60,21 +102,21 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
        language = "Chinese"
        speaker = "Vivian"
        instruct = "用特别愤怒的语气说"
-        prompts = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+        additional_information = {
+            "task_type": [task_type],
+            "text": [text],
+            "language": [language],
+            "speaker": [speaker],
+            "instruct": [instruct],
+            "max_new_tokens": [2048],
+        }
        inputs = {
-            "prompt": prompts,
-            "additional_information": {
-                "task_type": [task_type],
-                "text": [text],
-                "language": [language],
-                "speaker": [speaker],
-                "instruct": [instruct],
-                "max_new_tokens": [2048],
-            },
+            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+            "additional_information": additional_information,
        }
    return QueryResult(
        inputs=inputs,
-        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+        model_name=model_name,
    )


@@ -88,6 +130,7 @@ def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
        QueryResult with Omni inputs and the VoiceDesign model path.
    """
    task_type = "VoiceDesign"
+    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
    if use_batch_sample:
        texts = [
            "哥哥，你回来啦，人家等了你好久好久了，要抱抱！",
@@ -100,39 +143,39 @@ def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
        languages = ["Chinese", "English"]
        inputs = []
        for text, instruct, language in zip(texts, instructs, languages):
-            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
-            inputs.append(
-                {
-                    "prompt": prompt,
-                    "additional_information": {
-                        "task_type": [task_type],
-                        "text": [text],
-                        "language": [language],
-                        "instruct": [instruct],
-                        "max_new_tokens": [2048],
-                        "non_streaming_mode": [True],
-                    },
-                }
-            )
-    else:
-        text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！"
-        instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。"
-        language = "Chinese"
-        prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
-        inputs = {
-            "prompt": prompt,
-            "additional_information": {
+            additional_information = {
                "task_type": [task_type],
                "text": [text],
                "language": [language],
                "instruct": [instruct],
                "max_new_tokens": [2048],
                "non_streaming_mode": [True],
-            },
+            }
+            inputs.append(
+                {
+                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+                    "additional_information": additional_information,
+                }
+            )
+    else:
+        text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！"
+        instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。"
+        language = "Chinese"
+        additional_information = {
+            "task_type": [task_type],
+            "text": [text],
+            "language": [language],
+            "instruct": [instruct],
+            "max_new_tokens": [2048],
+            "non_streaming_mode": [True],
+        }
+        inputs = {
+            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+            "additional_information": additional_information,
        }
    return QueryResult(
        inputs=inputs,
-        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
+        model_name=model_name,
    )


@@ -147,6 +190,7 @@ def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> Que
        QueryResult with Omni inputs and the Base model path.
    """
    task_type = "Base"
+    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
    ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
    ref_audio_single = ref_audio_path_1
    ref_text_single = (
@@ -163,38 +207,38 @@ def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> Que
        syn_lang_batch = ["Chinese", "English"]
        inputs = []
        for text, language in zip(syn_text_batch, syn_lang_batch):
-            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
-            inputs.append(
-                {
-                    "prompt": prompt,
-                    "additional_information": {
-                        "task_type": [task_type],
-                        "ref_audio": [ref_audio_single],
-                        "ref_text": [ref_text_single],
-                        "text": [text],
-                        "language": [language],
-                        "x_vector_only_mode": [x_vector_only_mode],
-                        "max_new_tokens": [2048],
-                    },
-                }
-            )
-    else:
-        prompt = f"<|im_start|>assistant\n{syn_text_single}<|im_end|>\n<|im_start|>assistant\n"
-        inputs = {
-            "prompt": prompt,
-            "additional_information": {
+            additional_information = {
                "task_type": [task_type],
                "ref_audio": [ref_audio_single],
                "ref_text": [ref_text_single],
-                "text": [syn_text_single],
-                "language": [syn_lang_single],
+                "text": [text],
+                "language": [language],
                "x_vector_only_mode": [x_vector_only_mode],
                "max_new_tokens": [2048],
-            },
+            }
+            inputs.append(
+                {
+                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+                    "additional_information": additional_information,
+                }
+            )
+    else:
+        additional_information = {
+            "task_type": [task_type],
+            "ref_audio": [ref_audio_single],
+            "ref_text": [ref_text_single],
+            "text": [syn_text_single],
+            "language": [syn_lang_single],
+            "x_vector_only_mode": [x_vector_only_mode],
+            "max_new_tokens": [2048],
+        }
+        inputs = {
+            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+            "additional_information": additional_information,
        }
    return QueryResult(
        inputs=inputs,
-        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
+        model_name=model_name,
    )


@@ -223,30 +267,22 @@ def main(args):
        stage_init_timeout=args.stage_init_timeout,
    )

-    sampling_params = SamplingParams(
-        temperature=0.9,
-        top_p=1.0,
-        top_k=50,
-        max_tokens=2048,
-        seed=42,
-        detokenize=False,
-        repetition_penalty=1.05,
-    )
-
-    sampling_params_list = [
-        sampling_params,
-    ]
-
    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
    os.makedirs(output_dir, exist_ok=True)

-    omni_generator = omni.generate(query_result.inputs, sampling_params_list)
+    omni_generator = omni.generate(query_result.inputs, sampling_params_list=None)
    for stage_outputs in omni_generator:
        for output in stage_outputs.request_output:
            request_id = output.request_id
-            audio_tensor = output.outputs[0].multimodal_output["audio"]
+            audio_data = output.outputs[0].multimodal_output["audio"]
+            # async_chunk mode returns a list of chunks; concatenate them.
+            if isinstance(audio_data, list):
+                audio_tensor = torch.cat(audio_data, dim=-1)
+            else:
+                audio_tensor = audio_data
            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
-            audio_samplerate = output.outputs[0].multimodal_output["sr"].item()
+            sr_val = output.outputs[0].multimodal_output["sr"]
+            audio_samplerate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val[-1])
            # Convert to numpy array and ensure correct format
            audio_numpy = audio_tensor.float().detach().cpu().numpy()

@@ -310,10 +310,12 @@ class TestTTSMethods:
        speech_server.engine_client.stage_list = [mock_stage]
        assert speech_server._is_tts_model() is True

-    def test_build_tts_prompt(self, speech_server):
-        """Test TTS prompt format."""
-        prompt = speech_server._build_tts_prompt("Hello")
-        assert prompt == "<|im_start|>assistant\nHello<|im_end|>\n<|im_start|>assistant\n"
+    def test_estimate_prompt_len_fallback(self, speech_server):
+        """Test prompt length estimation falls back to 2048 when model is unavailable."""
+        tts_params = {"text": ["Hello"], "task_type": ["CustomVoice"]}
+        result = speech_server._estimate_prompt_len(tts_params)
+        # Without a real model, it should fall back to 2048.
+        assert result == 2048

    def test_validate_tts_request_basic(self, speech_server):
        """Test basic validation cases."""
@@ -69,6 +69,7 @@ def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
    runner.text_step = DummyBuffer(torch.zeros((bsz, hidden_size), dtype=torch.float32))

    runner.talker_mtp = DummyTalkerMTP()
+    runner.model = SimpleNamespace(talker_mtp_output_key="code_predictor_codes")
    runner.vllm_config = object()

    # Provide a minimal implementation that returns the expected 4-tuple.
@@ -58,6 +58,7 @@ class OmniModelConfig(ModelConfig):
        }
    )
    omni_kv_config: dict | None = None
+    codec_frame_rate_hz: float | None = None

    @property
    def registry(self):
@@ -128,6 +129,21 @@ class OmniModelConfig(ModelConfig):
            video_pruning_rate=video_pruning_rate,
        )

+        # Qwen3-TTS: infer codec frame rate from the model config for online serving.
+        if self.codec_frame_rate_hz is None and self.model_arch == "Qwen3TTSTalkerForConditionalGenerationARVLLM":
+            talker_cfg = getattr(self.hf_config, "talker_config", None)
+            if isinstance(talker_cfg, dict):
+                pos_per_sec = talker_cfg.get("position_id_per_seconds")
+            else:
+                pos_per_sec = getattr(talker_cfg, "position_id_per_seconds", None)
+            if pos_per_sec is not None:
+                try:
+                    fps = float(pos_per_sec)
+                except Exception:
+                    fps = None
+                if fps is not None and fps > 0:
+                    self.codec_frame_rate_hz = fps
+
        # Override hf_text_config with omni-specific logic for multi-stage models
        # (e.g., thinker_config, talker_config)
        new_hf_text_config = self.draw_hf_text_config()
@@ -62,7 +62,6 @@ class OmniGenerationScheduler(VLLMScheduler):
        while req_index < len(self.running) and token_budget > 0:
            request = self.running[req_index]
            # OMNI: Skip requests that are not in self.requests
-            # This can happen when connector marks request as finished and it's removed from requests
            if request.request_id not in self.requests or (
                self.chunk_transfer_adapter is None and request.status == RequestStatus.FINISHED_STOPPED
            ):
@@ -71,7 +70,25 @@ class OmniGenerationScheduler(VLLMScheduler):
                continue

            num_computed_tokens = request.num_computed_tokens
-            required_tokens = max(len(request.prompt_token_ids) - num_computed_tokens, 1)
+            required_tokens = len(request.prompt_token_ids) - num_computed_tokens
+            # async_chunk: don't schedule placeholder tokens when no new chunk is available.
+            if required_tokens <= 0:
+                if (
+                    self.chunk_transfer_adapter is not None
+                    and request.request_id in self.chunk_transfer_adapter.finished_requests
+                ):
+                    request.status = RequestStatus.FINISHED_STOPPED
+                    # Upstream may finish with no terminal tokens; append one pad token so we can emit FINISHED.
+                    if len(request.prompt_token_ids) <= num_computed_tokens:
+                        request.prompt_token_ids.append(0)
+                        try:
+                            request._all_token_ids.append(0)  # type: ignore[attr-defined]
+                        except Exception:
+                            pass
+                    required_tokens = len(request.prompt_token_ids) - num_computed_tokens
+                else:
+                    req_index += 1
+                    continue
            num_new_tokens = min(required_tokens, token_budget)
            new_blocks = self.kv_cache_manager.allocate_slots(
                request,
@@ -109,6 +126,20 @@ class OmniGenerationScheduler(VLLMScheduler):
                self.waiting.pop_request()
                continue

+            # async_chunk: wait for the first upstream chunk (don't start with placeholders).
+            if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0:
+                if request.request_id in self.chunk_transfer_adapter.finished_requests:
+                    request.status = RequestStatus.FINISHED_STOPPED
+                    request.prompt_token_ids.append(0)
+                    try:
+                        request._all_token_ids.append(0)  # type: ignore[attr-defined]
+                    except Exception:
+                        pass
+                else:
+                    self.waiting.pop_request()
+                    skipped_waiting_requests.prepend_request(request)
+                    continue
+
            # Uniformly treat as diffusion. A feature flag can be added later
            # via config or request tag.

@@ -145,11 +176,13 @@ class OmniGenerationScheduler(VLLMScheduler):

        # If fast path scheduled none, fall back to the original scheduling
        if not num_scheduled_tokens:
-            res = super().schedule()
            if self.chunk_transfer_adapter:
+                # Don't fall back: base scheduler doesn't handle async_chunk
+                # requests with empty prompt_token_ids.
                self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
-                self.chunk_transfer_adapter.postprocess_scheduler_output(res)
-            return res
+            else:
+                res = super().schedule()
+                return res

        # Compute common prefix blocks (aligned with v1)
        num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
@@ -318,9 +351,8 @@ class OmniGenerationScheduler(VLLMScheduler):
                continue
            request = self.requests.get(req_id)
            if request is None or request.is_finished():
-                # The request is already finished. This can happen if the
-                # request is aborted while the model is executing it (e.g.,
-                # in pipeline parallelism or async scheduling).
+                # Request may already be finished (e.g., aborted during
+                # execution / pipeline parallelism / async scheduling).
                continue

            req_index = model_runner_output.req_id_to_index[req_id]
@@ -360,8 +392,14 @@ class OmniGenerationScheduler(VLLMScheduler):
            routed_experts = None

            # Diffusion request: completes in one step; mark finished and free resources
-            if request.status == RequestStatus.FINISHED_STOPPED or (
-                self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens
+            if (
+                request.status == RequestStatus.FINISHED_STOPPED
+                or (self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens)
+                or (
+                    self.chunk_transfer_adapter is not None
+                    and request.request_id in self.chunk_transfer_adapter.finished_requests
+                    and request.num_computed_tokens >= len(request.prompt_token_ids)
+                )
            ):
                request.status = RequestStatus.FINISHED_STOPPED
                # Optional: set a stop_reason for front-end clarity
@@ -375,15 +413,11 @@ class OmniGenerationScheduler(VLLMScheduler):
                finished = self._handle_stopped_request(request)
                if finished:
                    kv_transfer_params = self._free_request(request)
-                if status_before_stop == RequestStatus.RUNNING:
-                    stopped_running_reqs.add(request)
-                elif status_before_stop == RequestStatus.WAITING_FOR_CHUNK:
-                    # In async chunk mode, request may be in either queue.
-                    # Remove from both to avoid stale queue entries.
+                if status_before_stop == RequestStatus.WAITING_FOR_CHUNK:
                    stopped_running_reqs.add(request)
                    stopped_preempted_reqs.add(request)
                else:
-                    stopped_preempted_reqs.add(request)
+                    stopped_running_reqs.add(request)

            # Extract sample logprobs if needed.
            if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
@@ -179,11 +179,17 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
            else:
                if payload_data.get("finished"):
                    self.finished_requests.add(req_id)
-                    req.status = RequestStatus.FINISHED_STOPPED

-                req.prompt_token_ids = payload_data.get("code_predictor_codes", [])
+                # req.prompt_token_ids = payload_data.get("code_predictor_codes", [])
+                # req.num_computed_tokens = 0
+                new_ids = payload_data.get("code_predictor_codes", [])
+                req.prompt_token_ids = new_ids
                req.num_computed_tokens = 0

+                # Empty chunk with more data expected: keep polling.
+                if not new_ids and not payload_data.get("finished"):
+                    return
+
            # Mark as finished for consumption
            with self.lock:
                self._finished_load_reqs.add(req_id)
@@ -308,7 +314,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                    # of schedule, but have not scheduled
                    continue
                if request.request_id in self.finished_requests:
-                    request.additional_information = {}
                    continue
                # Requests that waiting for chunk
                self.load_async(request)
@@ -390,7 +390,8 @@ class AsyncOmni(OmniBase):
                    submit_flag = False
                    prompt_token_ids = engine_outputs.prompt_token_ids
                    engine_input = copy.deepcopy(prompt)
-                    engine_input["prompt_token_ids"] = [0] * compute_talker_prompt_ids_length(prompt_token_ids)
+                    next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
+                    engine_input["prompt_token_ids"] = [0] * next_prompt_len
                    engine_input["multi_modal_data"] = engine_input["mm_processor_kwargs"] = None
                    for i in range(1, len(self.stage_list)):
                        task = {
@@ -1,6 +1,14 @@
 import asyncio
+import base64
+import io
+import ipaddress
+import socket
 from typing import Any
+from urllib.parse import urlparse
+from urllib.request import urlopen

+import numpy as np
+import soundfile as sf
 from fastapi import Request
 from fastapi.responses import Response
 from vllm.entrypoints.openai.engine.serving import OpenAIServing
@@ -17,6 +25,19 @@ from vllm_omni.outputs import OmniRequestOutput

 logger = init_logger(__name__)

+_REF_AUDIO_TIMEOUT_S = 15
+_REF_AUDIO_MAX_BYTES = 50 * 1024 * 1024  # 50 MB
+_REF_AUDIO_BLOCKED_NETWORKS = [
+    ipaddress.ip_network("127.0.0.0/8"),
+    ipaddress.ip_network("10.0.0.0/8"),
+    ipaddress.ip_network("172.16.0.0/12"),
+    ipaddress.ip_network("192.168.0.0/16"),
+    ipaddress.ip_network("169.254.0.0/16"),
+    ipaddress.ip_network("::1/128"),
+    ipaddress.ip_network("fc00::/7"),
+    ipaddress.ip_network("fe80::/10"),
+]
+
 # TTS Configuration (currently supports Qwen3-TTS)
 _TTS_MODEL_STAGES: set[str] = {"qwen3_tts"}
 _TTS_LANGUAGES: set[str] = {
@@ -43,6 +64,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
        # Load supported speakers
        self.supported_speakers = self._load_supported_speakers()
        logger.info(f"Loaded {len(self.supported_speakers)} supported speakers: {sorted(self.supported_speakers)}")
+        self._tts_tokenizer = None

    def _load_supported_speakers(self) -> set[str]:
        """Load supported speakers (case-insensitive) from the model configuration."""
@@ -62,6 +84,36 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):

        return set()

+    def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int:
+        """Estimate prompt length so the placeholder matches model-side embeddings."""
+        try:
+            from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+                Qwen3TTSTalkerForConditionalGeneration,
+            )
+
+            if self._tts_tokenizer is None:
+                from transformers import AutoTokenizer
+
+                model_name = self.engine_client.model_config.model
+                self._tts_tokenizer = AutoTokenizer.from_pretrained(
+                    model_name,
+                    trust_remote_code=True,
+                    padding_side="left",
+                )
+            hf_config = self.engine_client.model_config.hf_config
+            talker_config = hf_config.talker_config
+            task_type = (tts_params.get("task_type") or ["CustomVoice"])[0]
+            return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
+                additional_information=tts_params,
+                task_type=task_type,
+                tokenize_prompt=lambda t: self._tts_tokenizer(t, padding=False)["input_ids"],
+                codec_language_id=getattr(talker_config, "codec_language_id", None),
+                spk_is_dialect=getattr(talker_config, "spk_is_dialect", None),
+            )
+        except Exception as e:
+            logger.warning("Failed to estimate TTS prompt length, using fallback 2048: %s", e)
+            return 2048
+
    def _is_tts_model(self) -> bool:
        """Check if the current model is a supported TTS model."""
        stage_list = getattr(self.engine_client, "stage_list", None)
@@ -125,9 +177,44 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):

        return None

-    def _build_tts_prompt(self, text: str) -> str:
-        """Build TTS prompt from input text."""
-        return f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+    @staticmethod
+    async def _resolve_ref_audio(ref_audio_str: str) -> tuple[list[float], int]:
+        """Resolve ref_audio URL/base64 to (wav_samples, sample_rate)."""
+        parsed = urlparse(ref_audio_str)
+
+        def _check_ssrf(url: str) -> None:
+            host = urlparse(url).hostname
+            if not host:
+                raise ValueError("ref_audio URL must include a hostname")
+            for info in socket.getaddrinfo(host, None):
+                ip_str = str(info[4][0]).split("%", 1)[0]
+                addr = ipaddress.ip_address(ip_str)
+                if any(addr in net for net in _REF_AUDIO_BLOCKED_NETWORKS):
+                    raise ValueError(f"ref_audio URL resolves to blocked address: {addr}")
+
+        def _fetch_sync() -> tuple[np.ndarray, int]:
+            if parsed.scheme in ("http", "https"):
+                _check_ssrf(ref_audio_str)
+                with urlopen(ref_audio_str, timeout=_REF_AUDIO_TIMEOUT_S) as resp:
+                    data = resp.read(_REF_AUDIO_MAX_BYTES + 1)
+                    if len(data) > _REF_AUDIO_MAX_BYTES:
+                        raise ValueError(f"ref_audio URL exceeds {_REF_AUDIO_MAX_BYTES} bytes")
+                buf = io.BytesIO(data)
+            elif ref_audio_str.startswith("data:"):
+                b64 = ref_audio_str
+                if "," in b64:
+                    b64 = b64.split(",", 1)[1]
+                buf = io.BytesIO(base64.b64decode(b64))
+            else:
+                raise ValueError("ref_audio must be an http(s) URL or data: base64 URI")
+            audio, sr = sf.read(buf, dtype="float32", always_2d=False)
+            if isinstance(audio, np.ndarray) and audio.ndim > 1:
+                audio = np.mean(audio, axis=-1)
+            return np.asarray(audio, dtype=np.float32), int(sr)
+
+        loop = asyncio.get_running_loop()
+        wav_np, sr = await loop.run_in_executor(None, _fetch_sync)
+        return wav_np.tolist(), sr

    def _build_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
        """Build TTS parameters from request.
@@ -164,9 +251,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
        else:
            params["instruct"] = [""]

-        # Voice clone parameters (used with Base task)
-        if request.ref_audio is not None:
-            params["ref_audio"] = [request.ref_audio]
+        # Voice clone: ref_audio resolved in create_speech(), not here.
        if request.ref_text is not None:
            params["ref_text"] = [request.ref_text]
        if request.x_vector_only_mode is not None:
@@ -221,11 +306,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                if validation_error:
                    return self.create_error_response(validation_error)

-                # Build TTS parameters and prompt
+                # Must use prompt_token_ids (not text prompt): the AR Talker
+                # operates on codec tokens; text token IDs exceed codec vocab.
+                # model.preprocess replaces all embeddings, so placeholder value
+                # is irrelevant -- but length must match to avoid excess padding.
                tts_params = self._build_tts_params(request)
-                prompt_text = self._build_tts_prompt(request.input)
+
+                if request.ref_audio is not None:
+                    wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
+                    tts_params["ref_audio"] = [[wav_list, sr]]
+
+                ph_len = self._estimate_prompt_len(tts_params)
                prompt = {
-                    "prompt": prompt_text,
+                    "prompt_token_ids": [1] * ph_len,
                    "additional_information": tts_params,
                }
            else:
@@ -282,6 +375,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
            if hasattr(sample_rate, "item"):
                sample_rate = sample_rate.item()

+            # Streaming accumulates chunks as a list; concat first.
+            if isinstance(audio_tensor, list):
+                import torch
+
+                audio_tensor = torch.cat(audio_tensor, dim=-1)
            # Convert tensor to numpy
            if hasattr(audio_tensor, "float"):
                audio_tensor = audio_tensor.float().detach().cpu().numpy()
@@ -504,19 +504,39 @@ class Qwen3TTSConfig(PretrainedConfig):
        self.tts_bos_token_id = tts_bos_token_id
        self.tts_eos_token_id = tts_eos_token_id

-        # TODO: remove these dummy values after
-        self.image_token_id = 0  # dummy image token id
-        self.video_token_id = 0  # dummy video token id
-        self.vision_start_token_id = 0  # dummy vision start token id
+        # Dummy vision token IDs that must never collide with real codec tokens.
+        # mrope scans prompt_token_ids for these; using -1 ensures no false match.
+        self.image_token_id = -1
+        self.video_token_id = -1
+        self.vision_start_token_id = -1
        self.vision_config = PretrainedConfig()  # dummy vision config
        self.vision_config.spatial_merge_size = 1

+    @property
+    def codec_frame_rate_hz(self) -> float | None:
+        pos_per_sec = getattr(self.talker_config, "position_id_per_seconds", None)
+        if pos_per_sec is None:
+            return None
+        try:
+            fps = float(pos_per_sec)
+        except (TypeError, ValueError):
+            return None
+        return fps if fps > 0 else None
+
    def get_text_config(self, **kwargs):
        # vLLM expects text config to expose hidden_size/num_attention_heads.
        # For Qwen3 TTS, the talker config is the text model config.
        config = self.talker_config
-        # if hasattr(config, "rope_parameters"):
-        #     delattr(config, "rope_parameters")
+        # Code2Wav is a pure convolutional waveform decoder; it does NOT use
+        # rotary position embeddings.  When hf_overrides sets architectures
+        # to [Qwen3TTSCode2Wav], strip rope_parameters so that the model
+        # runner sees uses_mrope == False and skips mrope position computation
+        # on codec tokens.  Each stage loads its own config instance, so this
+        # in-place mutation does not affect the Talker stage.
+        archs = getattr(self, "architectures", []) or []
+        if any("Code2Wav" in str(a) for a in archs):
+            if hasattr(config, "rope_parameters"):
+                delattr(config, "rope_parameters")
        return config


@@ -1,102 +0,0 @@
-# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from transformers.feature_extraction_utils import BatchFeature
-from transformers.processing_utils import ProcessingKwargs, ProcessorMixin
-
-
-class Qwen3TTSProcessorKwargs(ProcessingKwargs, total=False):
-    _defaults = {
-        "text_kwargs": {
-            "padding": False,
-            "padding_side": "left",
-        }
-    }
-
-
-class Qwen3TTSProcessor(ProcessorMixin):
-    r"""
-    Constructs a Qwen3TTS processor.
-
-    Args:
-        tokenizer ([`Qwen2TokenizerFast`], *optional*):
-            The text tokenizer.
-        chat_template (`Optional[str]`, *optional*):
-            The Jinja template to use for formatting the conversation.
-            If not provided, the default chat template is used.
-    """
-
-    attributes = ["tokenizer"]
-    tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
-
-    def __init__(self, tokenizer=None, chat_template=None):
-        super().__init__(tokenizer, chat_template=chat_template)
-
-    def __call__(self, text=None, **kwargs) -> BatchFeature:
-        """
-        Main method to prepare for the model one or several sequences(s) and audio(s).
-        This method forwards the `text` and `kwargs` arguments to
-        Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None`
-        to encode the text.
-
-        Args:
-            text (`str`, `List[str]`, `List[List[str]]`):
-                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
-                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
-                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
-        """
-
-        if text is None:
-            raise ValueError("You need to specify either a `text` input to process.")
-
-        output_kwargs = self._merge_kwargs(
-            Qwen3TTSProcessorKwargs,
-            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
-            **kwargs,
-        )
-        if not isinstance(text, list):
-            text = [text]
-
-        texts_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
-
-        return BatchFeature(
-            data={**texts_inputs},
-            tensor_type=kwargs.get("return_tensors"),
-        )
-
-    def batch_decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
-        refer to the docstring of this method for more information.
-        """
-        return self.tokenizer.batch_decode(*args, **kwargs)
-
-    def decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
-        the docstring of this method for more information.
-        """
-        return self.tokenizer.decode(*args, **kwargs)
-
-    def apply_chat_template(self, conversations, chat_template=None, **kwargs):
-        if isinstance(conversations[0], dict):
-            conversations = [conversations]
-        return super().apply_chat_template(conversations, chat_template, **kwargs)
-
-    @property
-    def model_input_names(self):
-        tokenizer_input_names = self.tokenizer.model_input_names
-        return list(dict.fromkeys(tokenizer_input_names))
-
-
-__all__ = ["Qwen3TTSProcessor"]
@@ -0,0 +1,239 @@
+from __future__ import annotations
+
+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
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+from .qwen3_tts_tokenizer import Qwen3TTSTokenizer
+
+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."""
+
+    input_modalities = "audio"
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.model_path = vllm_config.model_config.model
+
+        self.have_multimodal_outputs = True
+        self.has_preprocess = False
+        self.has_postprocess = False
+        # Generation-only stage (no logits / sampling).
+        self.requires_raw_input_tokens = True
+
+        self._speech_tokenizer: Qwen3TTSTokenizer | None = None
+        self._num_quantizers: int | None = None
+        self._decode_upsample_rate: int | None = None
+        self._output_sample_rate: int | None = None
+        self._logged_codec_stats = False
+
+    @staticmethod
+    def _module_device(module: nn.Module) -> torch.device:
+        try:
+            return next(module.parameters()).device
+        except StopIteration:
+            for _, buf in module.named_buffers(recurse=True):
+                return buf.device
+            return torch.device("cpu")
+
+    def _ensure_speech_tokenizer_loaded(self) -> Qwen3TTSTokenizer:
+        if self._speech_tokenizer is not None:
+            return self._speech_tokenizer
+
+        # 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(
+                f"{self.model_path}/speech_tokenizer/preprocessor_config.json not found. "
+                "Please make sure the checkpoint contains the required HF preprocessing files."
+            )
+
+        tok = Qwen3TTSTokenizer.from_pretrained(
+            speech_tokenizer_dir,
+            torch_dtype=torch.bfloat16,
+            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:
+            raise ValueError("speech_tokenizer decoder_config.num_quantizers not found")
+        num_q = int(num_q)
+        if num_q <= 0:
+            raise ValueError(f"Invalid speech_tokenizer num_quantizers={num_q}")
+
+        try:
+            upsample = int(tok.get_decode_upsample_rate())
+        except Exception as e:
+            raise ValueError(f"Failed to get decode upsample rate: {e}") from e
+        if upsample <= 0:
+            raise ValueError(f"Invalid decode upsample rate: {upsample}")
+
+        try:
+            out_sr = int(tok.get_output_sample_rate())
+        except Exception as e:
+            raise ValueError(f"Failed to get output sample rate: {e}") from e
+
+        self._speech_tokenizer = tok
+        self._num_quantizers = num_q
+        self._decode_upsample_rate = upsample
+        self._output_sample_rate = out_sr
+        return tok
+
+    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.
+        if input_ids.numel() == 0:
+            return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32)
+        return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
+
+    def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None:
+        return None
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: Any = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Decode codec codes into audio waveform.
+
+        input_ids layout: [codec_context_frames, *flat_codes]
+        where flat_codes is codebook-major [q*F].
+        """
+        tok = self._ensure_speech_tokenizer_loaded()
+        assert self._num_quantizers is not None
+        assert self._output_sample_rate is not None
+
+        sr_val = self._output_sample_rate
+        empty_ret = (
+            torch.zeros((0,), dtype=torch.float32),
+            torch.tensor(sr_val, dtype=torch.int32),
+        )
+
+        if input_ids is None:
+            return empty_ret
+
+        q = int(self._num_quantizers)
+        ids = input_ids.reshape(-1).to(dtype=torch.long)
+        n_tokens = ids.numel()
+
+        if n_tokens == 0:
+            return empty_ret
+
+        # input_ids[0] = codec_context_frames (prepended by stage_input_processor).
+        ctx_frames = int(ids[0].item())
+        ids = ids[1:]
+        n_tokens = ids.numel()
+
+        if n_tokens == 0:
+            return empty_ret
+
+        # Warmup / dummy_run: not divisible by num_quantizers.
+        if n_tokens % q != 0:
+            logger.warning(
+                "Code2Wav input_ids length %d not divisible by num_quantizers %d, "
+                "likely a warmup run; returning empty audio.",
+                n_tokens,
+                q,
+            )
+            return empty_ret
+
+        total_frames = n_tokens // q
+
+        # Reshape codebook-major flat [q*F] -> [q, F] -> [F, q] for SpeechTokenizer.
+        codes_fq = ids.reshape(q, total_frames).transpose(0, 1).contiguous()
+
+        if not self._logged_codec_stats and total_frames > 1:
+            self._logged_codec_stats = True
+            try:
+                uniq = int(torch.unique(codes_fq).numel())
+                cmin = int(codes_fq.min().item())
+                cmax = int(codes_fq.max().item())
+                head = codes_fq[: min(2, total_frames), : min(8, q)].cpu().tolist()
+                logger.info(
+                    "Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] head=%s",
+                    total_frames,
+                    q,
+                    uniq,
+                    cmin,
+                    cmax,
+                    head,
+                )
+            except Exception:
+                pass
+
+        wavs, sr = tok.decode({"audio_codes": codes_fq})
+        if not wavs:
+            raise ValueError("SpeechTokenizer code2wav produced empty waveform list.")
+        audio_np = wavs[0].astype(np.float32, copy=False)
+
+        # Trim left-context waveform samples (streaming sliding window).
+        if ctx_frames > 0:
+            upsample = self._decode_upsample_rate
+            if upsample is None or upsample <= 0:
+                raise ValueError(f"Invalid decode upsample rate: {upsample}")
+            cut = ctx_frames * upsample
+            if cut < audio_np.shape[0]:
+                audio_np = audio_np[cut:]
+            else:
+                logger.warning(
+                    "Context trim %d >= decoded length %d; returning empty audio.",
+                    cut,
+                    audio_np.shape[0],
+                )
+                return empty_ret
+
+        audio_tensor = torch.from_numpy(audio_np).to(dtype=torch.float32).reshape(-1)
+        sr_tensor = torch.tensor(int(sr), dtype=torch.int32)
+        return audio_tensor, sr_tensor
+
+    def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
+        if isinstance(model_outputs, OmniOutput):
+            return model_outputs
+
+        if not (isinstance(model_outputs, tuple) and len(model_outputs) == 2):
+            raise TypeError(f"Qwen3TTSCode2Wav expected (audio_tensor, sr) outputs, got {type(model_outputs)}")
+
+        audio_tensor, sr = model_outputs
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={
+                "model_outputs": audio_tensor,
+                "sr": sr,
+            },
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # SpeechTokenizer weights live under `speech_tokenizer/` and are loaded
+        # lazily from that directory. Ignore main checkpoint weights.
+        return set()
@@ -0,0 +1,495 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.config.vllm import set_current_vllm_config
+from vllm.forward_context import set_forward_context
+from vllm.model_executor.layers.layernorm import RMSNorm
+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 .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig
+
+
+class _LocalPredictorKVCache:
+    """Minimal local KV cache + attention metadata for running
+    code_predictor inside one worker (independent of engine KV)."""
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        max_seq_len: int,
+        max_batch_size: int,
+        device: torch.device,
+    ) -> 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)
+            ],
+        )
+
+        # 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.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,
+        )
+
+        # 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(
+        self,
+        *,
+        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.
+
+        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}")
+
+        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), {}
+
+        # 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 = int(query_lens_i32[i].item())
+            sl = int(seq_lens_i32[i].item())
+            start = sl - ql
+            pos_list.append(torch.arange(start, sl, dtype=torch.int64))
+        positions_cpu = torch.cat(pos_list, dim=0)
+
+        # 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 = int(query_lens_i32[i].item())
+            sl = int(seq_lens_i32[i].item())
+            start = sl - ql
+            for p in range(start, sl):
+                block_idx = p // self.block_size
+                offset = p % self.block_size
+                block_id = int(self._block_table[i, block_idx].item())
+                slot_mapping[cursor] = block_id * self.block_size + offset
+                cursor += 1
+
+        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)
+
+        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
+
+
+class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module):
+    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)
+            ]
+        )
+        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.
+        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)]
+        )
+
+    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]
+        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)
+        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"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        loaded_params: set[str] = set()
+        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:
+                    continue
+                mapped = name.replace(weight_name, param_name)
+                if mapped.endswith(".bias") and mapped not in params_dict:
+                    continue
+                if is_pp_missing_parameter(mapped, self):
+                    continue
+                if mapped.endswith("scale"):
+                    mapped = maybe_remap_kv_scale_name(mapped, params_dict)
+                    if mapped is None:
+                        continue
+                param = params_dict.get(mapped)
+                if param is None:
+                    continue
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                if weight_loader == default_weight_loader:
+                    weight_loader(param, loaded_weight)
+                else:
+                    weight_loader(param, loaded_weight, shard_id)
+                loaded_params.add(mapped)
+                break
+            else:
+                mapped = maybe_remap_kv_scale_name(name, params_dict)
+                if mapped is None:
+                    continue
+                if name.endswith(".bias") and mapped not in params_dict:
+                    continue
+                if is_pp_missing_parameter(mapped, self):
+                    continue
+                param = params_dict.get(mapped)
+                if param is None:
+                    continue
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                loaded_params.add(mapped)
+        return loaded_params
+
+
+class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
+    """vLLM-native code_predictor used by the AR talker (residual codebooks)."""
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        config: Qwen3TTSTalkerCodePredictorConfig,
+        talker_config: Qwen3TTSTalkerConfig,
+        prefix: str = "code_predictor",
+    ) -> None:
+        super().__init__()
+        self._vllm_config = vllm_config
+        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)]
+        )
+
+        if config.hidden_size != talker_config.hidden_size:
+            self.small_to_mtp_projection = nn.Linear(talker_config.hidden_size, config.hidden_size, bias=True)
+        else:
+            self.small_to_mtp_projection = nn.Identity()
+
+        self._kv_cache: _LocalPredictorKVCache | 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]] = []
+            other_weights: list[tuple[str, torch.Tensor]] = []
+            for name, w in weights:
+                if name.startswith("model."):
+                    model_weights.append((name[len("model.") :], w))
+                else:
+                    other_weights.append((name, w))
+
+            loaded_model = self.model.load_weights(model_weights)
+            loaded |= {f"model.{n}" for n in loaded_model}
+
+            params = dict(self.named_parameters(remove_duplicate=False))
+            for name, w in other_weights:
+                if name not in params:
+                    continue
+                default_weight_loader(params[name], w)
+                loaded.add(name)
+            return loaded
+
+    def _maybe_init_kv_cache(self, device: torch.device) -> None:
+        if self._kv_cache is not None:
+            return
+        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.
+        hs = inputs_embeds.to(dtype=torch.bfloat16).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
+        )
+
+        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)
+
+        # 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
+        )
+
+        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
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        layer0_code: torch.Tensor,
+        layer0_embed: torch.Tensor,
+        last_talker_hidden: torch.Tensor,
+        do_sample: bool = True,
+        temperature: float = 1.0,
+        top_k: int = 50,
+        top_p: float = 1.0,
+    ) -> torch.Tensor:
+        """Full autoregressive prediction of residual codebooks 1..Q-1.
+
+        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.
+
+        Returns:
+            audio_codes: [B, Q] all codebook tokens (layer0 + residuals).
+        """
+        bsz = int(layer0_code.shape[0])
+        num_groups = int(self.config.num_code_groups)
+        max_steps = num_groups - 1
+
+        # Reset KV cache for a fresh sequence.
+        self.reset_cache()
+
+        # 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]
+
+        all_codes = [layer0_code.reshape(bsz, 1)]
+        past_seq_len = 2
+
+        for step in range(1, num_groups):
+            # Sample or argmax from logits.
+            if do_sample and temperature > 0:
+                scaled = logits / 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)
+                next_ids = torch.multinomial(probs, num_samples=1)  # [B, 1]
+            else:
+                next_ids = logits.argmax(dim=-1, keepdim=True)  # [B, 1]
+            all_codes.append(next_ids)
+
+            # 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
+
+        return torch.cat(all_codes, dim=1)  # [B, Q]
@@ -80,16 +80,21 @@ class Qwen3TTSTokenizer:
        """
        inst = cls()

-        AutoConfig.register("qwen3_tts_tokenizer_25hz", Qwen3TTSTokenizerV1Config)
-        AutoModel.register(Qwen3TTSTokenizerV1Config, Qwen3TTSTokenizerV1Model)
+        load_feature_extractor = bool(kwargs.pop("load_feature_extractor", True))

        AutoConfig.register("qwen3_tts_tokenizer_12hz", Qwen3TTSTokenizerV2Config)
        AutoModel.register(Qwen3TTSTokenizerV2Config, Qwen3TTSTokenizerV2Model)

-        inst.feature_extractor = AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path)
+        AutoConfig.register("qwen3_tts_tokenizer_25hz", Qwen3TTSTokenizerV1Config)
+        AutoModel.register(Qwen3TTSTokenizerV1Config, Qwen3TTSTokenizerV1Model)
+
        inst.model = AutoModel.from_pretrained(pretrained_model_name_or_path, **kwargs)
        inst.config = inst.model.config

+        inst.feature_extractor = (
+            AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path) if load_feature_extractor else None
+        )
+
        inst.device = getattr(inst.model, "device", None)
        if inst.device is None:
            # fallback: infer from first parameter device
@@ -212,12 +217,7 @@ class Qwen3TTSTokenizer:
        audios: AudioInput,
        sr: int | None = None,
        return_dict: bool = True,
-    ) -> (
-        Qwen3TTSTokenizerV1EncoderOutput
-        | Qwen3TTSTokenizerV2EncoderOutput
-        | tuple[list[torch.Tensor], list[torch.Tensor] | None, list[torch.Tensor] | None]
-        | tuple[list[torch.Tensor]]
-    ):
+    ) -> Qwen3TTSTokenizerV1EncoderOutput | Qwen3TTSTokenizerV2EncoderOutput | tuple:
        """
        Batch-encode audio into discrete codes (and optional conditioning, depending on 25Hz/12Hz).

@@ -50,8 +50,18 @@ _OMNI_MODELS = {
    ),
    "Qwen3TTSForConditionalGeneration": (
        "qwen3_tts",
+        "qwen3_tts_talker",
+        "Qwen3TTSTalkerForConditionalGeneration",
+    ),
+    "Qwen3TTSTalkerForConditionalGeneration": (
        "qwen3_tts",
-        "Qwen3TTSModelForGeneration",
+        "qwen3_tts_talker",
+        "Qwen3TTSTalkerForConditionalGeneration",
+    ),
+    "Qwen3TTSCode2Wav": (
+        "qwen3_tts",
+        "qwen3_tts_code2wav",
+        "Qwen3TTSCode2Wav",
    ),
 }

@@ -1,22 +1,101 @@
+async_chunk: true
 stage_args:
  - stage_id: 0
-    stage_type: llm  # Use llm stage type to launch OmniLLM
+    stage_type: llm
    runtime:
      devices: "0"
      max_batch_size: 1
    engine_args:
      model_stage: qwen3_tts
-      model_arch: Qwen3TTSForConditionalGeneration
+      model_arch: Qwen3TTSTalkerForConditionalGeneration
+      # Force stage-specific registered architecture.
+      hf_overrides:
+        architectures: [Qwen3TTSTalkerForConditionalGeneration]
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      enforce_eager: false
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: latent
+      gpu_memory_utilization: 0.3
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 512
+      max_model_len: 4096
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
+    # Use named connector to apply runtime.connectors.extra.
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: false
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3TTSCode2Wav
+      # Force stage-specific registered architecture.
+      hf_overrides:
+        architectures: [Qwen3TTSCode2Wav]
      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      enforce_eager: true
      trust_remote_code: true
      async_scheduling: false
      enable_prefix_caching: false
-      engine_output_type: audio  # Final output: audio waveform
-      gpu_memory_utilization: 0.1
+      engine_output_type: audio
+      gpu_memory_utilization: 0.2
      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 1000000
-
+      # Must be divisible by num_code_groups and cover (left_context + chunk).
+      max_num_batched_tokens: 8192
+      # async_chunk appends windows per step; max_model_len must cover accumulated stream.
+      max_model_len: 32768
+    engine_input_source: [0]
    final_output: true
    final_output_type: audio
+    # Distributed connector configuration
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.0
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+        # Frame-aligned codec streaming transport.
+        codec_streaming: true
+        # Connector polling / timeout (unit: loop count, sleep interval in seconds).
+        connector_get_sleep_s: 0.01
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        # Align with Omni: small chunks with sufficient context overlap.
+        codec_chunk_frames: 25
+        codec_left_context_frames: 25
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
@@ -0,0 +1,76 @@
+"""Stage input processor for Qwen3-TTS: Talker -> Code2Wav."""
+
+from typing import Any
+
+import torch
+
+
+def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None:
+    audio_codes = pooling_output.get("audio_codes")
+    if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
+        return None
+    if audio_codes.ndim == 2:
+        frame = audio_codes[-1]
+        if frame.numel() == 0 or not bool(frame.any().item()):
+            return None
+        return frame.to(torch.long).reshape(-1)
+    if audio_codes.ndim == 1:
+        return audio_codes.to(torch.long).reshape(-1)
+    raise ValueError(f"Invalid audio_codes shape for Qwen3-TTS async_chunk: {tuple(audio_codes.shape)}")
+
+
+def talker2code2wav_async_chunk(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any],
+    request: Any,
+) -> dict[str, Any] | None:
+    if not isinstance(pooling_output, dict):
+        return None
+
+    request_id = request.external_req_id
+
+    connector = getattr(transfer_manager, "connector", None)
+    raw_cfg = getattr(connector, "config", {}) or {}
+    cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+    chunk_size = int(cfg.get("codec_chunk_frames", 25))
+    left_context_size = int(cfg.get("codec_left_context_frames", 25))
+    if chunk_size <= 0 or left_context_size < 0:
+        raise ValueError(
+            f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, "
+            f"codec_left_context_frames={left_context_size}"
+        )
+
+    finished = bool(request.is_finished())
+
+    frame = _extract_last_frame(pooling_output)
+    if frame is not None:
+        codec_codes = frame.cpu().tolist()
+        transfer_manager.code_prompt_token_ids[request_id].append(codec_codes)
+
+    length = len(transfer_manager.code_prompt_token_ids[request_id])
+    chunk_length = length % chunk_size
+
+    if chunk_length != 0 and not finished:
+        return None
+
+    context_length = chunk_length if chunk_length != 0 else chunk_size
+
+    if length <= 0:
+        return {
+            "code_predictor_codes": [],
+            "finished": torch.tensor(bool(finished), dtype=torch.bool),
+        }
+
+    end_index = min(length, left_context_size + context_length)
+    ctx_frames = max(0, int(end_index - context_length))
+    window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
+
+    # Pack context + chunk into codebook-major flat codes for adapter.
+    code_predictor_codes = torch.tensor(window_frames).transpose(0, 1).reshape(-1).tolist()
+
+    # Build final prompt_token_ids with ctx_frames header for Qwen3-TTS Code2Wav.
+    # The model expects input_ids layout: [ctx_frames, *flat_codes].
+    return {
+        "code_predictor_codes": [int(ctx_frames)] + code_predictor_codes,
+        "finished": torch.tensor(bool(finished), dtype=torch.bool),
+    }
@@ -1,22 +1,92 @@
+async_chunk: true
 stage_args:
  - stage_id: 0
-    stage_type: llm  # Use llm stage type to launch OmniLLM
+    stage_type: llm
    runtime:
      devices: "0"
      max_batch_size: 1
    engine_args:
      model_stage: qwen3_tts
-      model_arch: Qwen3TTSForConditionalGeneration
+      model_arch: Qwen3TTSTalkerForConditionalGeneration
+      hf_overrides:
+        architectures: [Qwen3TTSTalkerForConditionalGeneration]
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: latent
+      gpu_memory_utilization: 0.3
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 512
+      max_model_len: 4096
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: false
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3TTSCode2Wav
+      hf_overrides:
+        architectures: [Qwen3TTSCode2Wav]
      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      enforce_eager: true
      trust_remote_code: true
      async_scheduling: false
      enable_prefix_caching: false
-      engine_output_type: audio  # Final output: audio waveform
-      gpu_memory_utilization: 0.1
+      engine_output_type: audio
+      gpu_memory_utilization: 0.2
      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 1000000
-
+      max_num_batched_tokens: 8192
+      max_model_len: 32768
+    engine_input_source: [0]
    final_output: true
    final_output_type: audio
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.0
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+        codec_streaming: true
+        connector_get_sleep_s: 0.01
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        codec_chunk_frames: 25
+        codec_left_context_frames: 25
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
@@ -238,9 +238,12 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
            pooler_output.append(mm_payload)
        else:
            raise RuntimeError("Unsupported diffusion output type")
+        # [Omni] Copy req_id mappings to avoid async scheduling mutation.
+        req_ids_output_copy = self.input_batch.req_ids.copy()
+        req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
        output = OmniModelRunnerOutput(
-            req_ids=self.input_batch.req_ids,
-            req_id_to_index=self.input_batch.req_id_to_index,
+            req_ids=req_ids_output_copy,
+            req_id_to_index=req_id_to_index_output_copy,
            sampled_token_ids=[],
            logprobs=None,
            prompt_logprobs_dict={},
@@ -380,9 +380,12 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                pooler_output.append(mm_payload)
        else:
            raise RuntimeError("Unsupported diffusion output type")
+        # [Omni] Copy req_id mappings to avoid async scheduling mutation.
+        req_ids_output_copy = self.input_batch.req_ids.copy()
+        req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
        output = OmniModelRunnerOutput(
-            req_ids=self.input_batch.req_ids,
-            req_id_to_index=self.input_batch.req_id_to_index,
+            req_ids=req_ids_output_copy,
+            req_id_to_index=req_id_to_index_output_copy,
            sampled_token_ids=[],
            logprobs=None,
            prompt_logprobs_dict={},
@@ -70,16 +70,24 @@ class OmniGPUModelRunner(GPUModelRunner):
    @instrument(span_name="Loading (GPU)")
    def load_model(self, *args, **kwargs) -> None:
        super().load_model(*args, **kwargs)
+
        # TODO move this model specific logic to a separate class
-        if hasattr(self.model, "talker_mtp") and self.model.talker is not None:
-            self.talker_mtp = self.model.talker_mtp
+        # TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.
+        talker_mtp = getattr(self.model, "talker_mtp", None)
+        if talker_mtp is not None:
+            self.talker_mtp = talker_mtp  # type: ignore[assignment]
            cudagraph_mode = self.compilation_config.cudagraph_mode
            assert cudagraph_mode is not None
-            if cudagraph_mode.has_full_cudagraphs():
-                self.talker_mtp = CUDAGraphWrapper(
-                    self.model.talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
-                )
-            hidden_size = self.model_config.hf_config.talker_config.text_config.hidden_size
+            # Only wrap talker_mtp in CUDAGraphWrapper for Omni models that
+            # have a separate .talker sub-module.  TTS models' code predictor
+            # has internal AR loops / torch.multinomial — not graph-safe.
+            has_separate_talker = getattr(self.model, "talker", None) is not None
+            if cudagraph_mode.has_full_cudagraphs() and has_separate_talker:
+                self.talker_mtp = CUDAGraphWrapper(talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL)
+            # TTS exposes mtp_hidden_size; Omni uses hf_text_config.hidden_size.
+            hidden_size = int(
+                getattr(self.model, "mtp_hidden_size", 0) or getattr(self.model_config.hf_text_config, "hidden_size")
+            )
            max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size)
            self.talker_mtp_input_ids = self._make_buffer(max_batch_size, dtype=torch.int32)
            self.talker_mtp_inputs_embeds = self._make_buffer(
@@ -644,6 +652,11 @@ class OmniGPUModelRunner(GPUModelRunner):
                input_ids = None
                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
                model_kwargs = self._init_model_kwargs()
+            elif getattr(getattr(self, "model", None), "has_preprocess", False):
+                # Capture CUDA graph with inputs_embeds path so replay reads
+                # from the same buffer that _preprocess writes into.
+                input_ids = self.input_ids.gpu[:num_tokens_padded]
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
            else:
                input_ids = self.input_ids.gpu[:num_tokens_padded]
                inputs_embeds = None
@@ -986,6 +999,11 @@ class OmniGPUModelRunner(GPUModelRunner):
            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
            model_kwargs = self._init_model_kwargs()
            input_ids = self.input_ids.gpu[:num_input_tokens]
+        elif getattr(self.model, "has_preprocess", False):
+            # Use pre-allocated buffer for CUDA graph compatibility.
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = self._init_model_kwargs()
        else:
            # For text-only models, we use token ids as input.
            # While it is possible to use embeddings as input just like the
@@ -1047,9 +1065,16 @@ class OmniGPUModelRunner(GPUModelRunner):
                span_len = int(e) - int(s)

                # call the custom process function
+                embed_slice = inputs_embeds[s:e] if inputs_embeds is not None else None
                req_input_ids, req_embeds, update_dict = self.model.preprocess(
-                    input_ids=input_ids[s:e], input_embeds=inputs_embeds[s:e], **req_infos
+                    input_ids=input_ids[s:e], input_embeds=embed_slice, **req_infos
                )
+                if inputs_embeds is None:
+                    inputs_embeds = torch.empty(
+                        (input_ids.shape[0], req_embeds.shape[-1]),
+                        device=req_embeds.device,
+                        dtype=req_embeds.dtype,
+                    )

                if hasattr(self.model, "talker_mtp") and span_len == 1:
                    last_talker_hidden, text_step = update_dict.pop("mtp_inputs")
@@ -1093,6 +1118,9 @@ class OmniGPUModelRunner(GPUModelRunner):
            max_num_scheduled_tokens=1,
            use_cascade_attn=False,
        )
+        # Force eager for unwrapped code predictors (AR loops / multinomial).
+        if not isinstance(self.talker_mtp, CUDAGraphWrapper):
+            _cudagraph_mode = CUDAGraphMode.NONE
        num_tokens_padded = batch_desc.num_tokens
        req_input_ids = self.talker_mtp_input_ids.gpu[:num_tokens_padded]
        req_embeds = self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded]
@@ -1104,11 +1132,12 @@ class OmniGPUModelRunner(GPUModelRunner):
            req_embeds, code_predictor_codes = self.talker_mtp(req_input_ids, req_embeds, last_talker_hidden, text_step)
        # update the inputs_embeds and code_predictor_codes
        code_predictor_codes_cpu = code_predictor_codes.detach().to("cpu").contiguous()
+        out_key = getattr(self.model, "talker_mtp_output_key", "code_predictor_codes")
        for idx, req_id in enumerate(decode_req_ids):
            req_index = self.input_batch.req_ids.index(req_id)
            start_offset = int(self.query_start_loc.cpu[req_index])
            inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
-            update_dict = {"code_predictor_codes": code_predictor_codes_cpu[idx : idx + 1]}
+            update_dict = {out_key: code_predictor_codes_cpu[idx : idx + 1]}
            self._merge_additional_information_update(req_id, update_dict)

    def _model_forward(
@@ -1142,7 +1171,9 @@ class OmniGPUModelRunner(GPUModelRunner):
        self._omni_last_model_output = model_output
        return model_output

-    def _merge_additional_information_update(self, req_id: str, upd: dict) -> None:
+    def _merge_additional_information_update(self, req_id: str, upd: dict | None) -> None:
+        if not isinstance(upd, dict):
+            return
        req_state = self.requests.get(req_id)
        if req_state is None:
            return