[Refactor] Migrate Legacy Stage Configs to the Pipeline Registry (#5031)

Signed-off-by: zwhzzz0821 <2831474076@qq.com>
Co-authored-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/model_executor/models/hunyuan_video/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/hunyuan_video/pipeline.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HunyuanVideo-1.5 pipeline topology."""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+HUNYUAN_VIDEO_15_PIPELINE = PipelineConfig(
+    model_type="hunyuan_video_15",
+    model_arch="HunyuanVideo15Pipeline",
+    diffusers_class_name="HunyuanVideo15Pipeline",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="dit",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="video",
+            model_arch="HunyuanVideo15Pipeline",
+        ),
+    ),
+)

vllm_omni/model_executor/models/step_audio2/__init__.py

@@ -1,3 +0,0 @@
-from .step_audio2 import StepAudio2ForConditionalGeneration
-
-__all__ = ["StepAudio2ForConditionalGeneration"]

vllm_omni/model_executor/models/step_audio2/pipeline.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Step-Audio2 pipeline topologies (frozen)."""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.step_audio2"
+_HF_ARCHITECTURES = (
+    "StepAudio2ForCausalLM",
+    "StepAudio2ForConditionalGeneration",
+)
+
+
+STEP_AUDIO2_PIPELINE = PipelineConfig(
+    model_type="step_audio_2",
+    model_arch="StepAudio2ForConditionalGeneration",
+    hf_architectures=_HF_ARCHITECTURES,
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="step_audio2_thinker",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="latent",
+            model_arch="StepAudio2ThinkerForConditionalGeneration",
+            async_chunk_process_next_stage_input_func=f"{_PROC}.thinker2token2wav_async_chunk",
+            sampling_constraints={"detokenize": True},
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="token2wav",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            model_arch="StepAudio2Token2WavModel",
+            sync_process_input_func=f"{_PROC}.thinker2token2wav",
+            sampling_constraints={"detokenize": False},
+        ),
+    ),
+)
+
+
+STEP_AUDIO2_ASR_PIPELINE = PipelineConfig(
+    model_type="step_audio_2_asr",
+    model_arch="StepAudio2ForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="step_audio2_thinker",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="text",
+            model_arch="StepAudio2ThinkerForConditionalGeneration",
+            sampling_constraints={"detokenize": True},
+        ),
+    ),
+)

vllm_omni/model_executor/models/wan2_2/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/wan2_2/pipeline.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Wan2.2 TI2V pipeline topology."""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+WAN2_2_TI2V_PIPELINE = PipelineConfig(
+    model_type="wan2_2_ti2v",
+    model_arch="WanPipeline",
+    diffusers_class_name="WanPipeline",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="dit",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="video",
+            model_arch="WanPipeline",
+        ),
+    ),
+)

vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml

@@ -1,38 +0,0 @@
-# Stage config for running HunyuanVideo-1.5 DiT with ModelOpt FP8 auto-detect.
-# Single GPU. Bump `tensor_parallel_size` and `devices` for multi-GPU TP.
-#
-# Use with a ModelOpt FP8 checkpoint (e.g. produced by
-# examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py).
-
-stage_args:
-  - stage_id: 0
-    stage_type: diffusion
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: dit
-      model_class_name: HunyuanVideo15Pipeline
-      max_num_seqs: 1
-      enforce_eager: true
-      trust_remote_code: true
-      distributed_executor_backend: "mp"
-      # HunyuanVideo-1.5 peak VRAM is dominated by the single-shot VAE decode of
-      # the full video (the DiT contributes little), so tiling the decode is the
-      # main memory lever — far more than FP8 (e.g. 480x832x33: ~117 -> ~47 GiB).
-      # Default it on here; override with --vae-use-tiling/false if you have headroom.
-      vae_use_tiling: true
-      parallel_config:
-        tensor_parallel_size: 1
-
-    final_output: true
-    final_output_type: video
-    is_comprehension: false
-    default_sampling_params:
-      seed: 42
-
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 1

vllm_omni/model_executor/stage_configs/mimo_audio.yaml

@@ -1,61 +0,0 @@
-
-# MiMo Audio deploy: 2-stage thinker+talker → code2wav.
-#
-# Default mode is async-chunk streaming on a single GPU (both stages on
-# device 0) through SharedMemoryConnector. For the legacy 2-GPU sync
-# pipeline (stage 1 on a second card), pass ``--no-async-chunk
-# --stage-overrides '{"1": {"devices": "1", "max_model_len": 18192, "max_num_batched_tokens": 18192}}'``
-# to the serve CLI.
-async_chunk: true
-dtype: bfloat16
-
-connectors:
-  connector_of_shared_memory:
-    name: SharedMemoryConnector
-    extra:
-      shm_threshold_bytes: 365536
-      codec_streaming: true
-      connector_get_sleep_s: 0.001
-      connector_get_max_wait_first_chunk: 3000
-      connector_get_max_wait: 300
-      codec_chunk_frames: 30
-      codec_left_context_frames: 40
-
-stages:
-  - stage_id: 0
-    max_num_seqs: 1
-    gpu_memory_utilization: 0.5
-    enforce_eager: true
-    trust_remote_code: true
-    enable_prefix_caching: false
-    max_num_batched_tokens: 8192
-    max_model_len: 8192
-    devices: "0"
-    output_connectors:
-      to_stage_1: connector_of_shared_memory
-    default_sampling_params:
-      temperature: 0.6
-      top_p: 0.95
-      top_k: 50
-      max_tokens: 18192
-      seed: 42
-      repetition_penalty: 1.1
-
-  - stage_id: 1
-    max_num_seqs: 1
-    gpu_memory_utilization: 0.4
-    enforce_eager: true
-    trust_remote_code: true
-    enable_prefix_caching: false
-    async_scheduling: false
-    max_num_batched_tokens: 8192
-    max_model_len: 8192
-    devices: "0"
-    input_connectors:
-      from_stage_0: connector_of_shared_memory
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 18192
-      seed: 42

vllm_omni/model_executor/stage_configs/step_audio_2.yaml

@@ -1,122 +0,0 @@
-# Stage config for Step-Audio2 with vLLM-Omni
-#
-# Two-stage pipeline:
-#   Stage 0 (Thinker): Audio understanding → Text + Audio tokens
-#   Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz)
-#
-# Hardware requirement: 1-2 GPUs (can run on single GPU with devices: "0" for both)
-
-# Model configuration reference (for documentation)
-# Actual values are set in hf_overrides below
-#
-# Token configuration:
-#   text_max: 151688              # Maximum text token ID (inclusive)
-#   audio_start: 151696           # First audio token ID in absolute vocabulary
-#   audio_vocab_size: 6562        # Total number of audio tokens
-#   audio_eos: 6561               # Audio EOS token ID (relative to audio_start)
-#   audio_patch_token_id: 151690  # <audio_patch> placeholder token ID
-#
-# Encoder configuration:
-#   n_mels: 128                   # Number of mel frequency bins
-#   n_audio_ctx: 1500             # Audio context length
-#   n_audio_state: 512            # Audio encoder hidden state dimension
-#   n_audio_head: 8               # Number of attention heads
-#   n_audio_layer: 6              # Number of encoder layers
-#   kernel_size: 3                # Convolution kernel size for adapter
-#   adapter_stride: 2             # Stride for adapter convolution
-
-stage_args:
-  # Stage 0: Thinker (Audio Understanding + Token Generation)
-  - stage_id: 0
-    runtime:
-      process: true              # Run in separate process
-      devices: "0,1,2,3"          # Use 3 GPUs for tensor parallelism
-      max_batch_size: 1         # Single sample at a time
-    engine_args:
-      model_stage: step_audio2_thinker
-      model_arch: StepAudio2ThinkerForConditionalGeneration
-      tensor_parallel_size: 4   # Split model across 3 GPUs
-        # NOTE: All default values are defined in step_audio2_constants.py (DEFAULT_MODEL_CONFIG)
-        # Only override here if your model differs from the defaults
-        #
-        # Default values (for reference, from constants.py):
-        #   hidden_size: 4096              # Qwen2-7B (change to 1536 for 1.5B, 896 for 0.5B)
-        #   text_max: 151688
-        #   audio_start: 151696
-        #   audio_vocab_size: 6562
-        #   audio_eos: 6561
-        #   audio_patch_token_id: 151690
-        #   n_mels: 128
-        #   n_audio_ctx: 1500
-        #   n_audio_state: 512
-        #   n_audio_head: 8
-        #   n_audio_layer: 6
-        #   kernel_size: 3
-        #   adapter_stride: 2
-        #
-        # Uncomment and modify only if needed:
-        # hidden_size: 4096
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.8
-      enforce_eager: true        # Required for now
-      trust_remote_code: true
-      enable_prefix_caching: false
-      max_num_batched_tokens: 8192
-      engine_output_type: text   # Outputs text + audio tokens
-    is_comprehension: true       # This stage does comprehension
-    final_output: true          # Can be final output (for ASR mode)
-    final_output_type: text
-    default_sampling_params:
-      temperature: 0.7
-      top_p: 0.9
-      top_k: -1
-      max_tokens: 1024
-      seed: 42
-      detokenize: true
-      repetition_penalty: 1.05
-
-  # Stage 1: Token2Wav (Audio Synthesis)
-  - stage_id: 1
-    runtime:
-      process: true
-      devices: "3"              # Use GPU 3 (separate from Stage 0's 0-2)
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2wav
-      model_arch: StepAudio2Token2WavModel
-      hf_overrides:
-        architectures: ["StepAudio2Token2WavForConditionalGeneration"]
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-
-      gpu_memory_utilization: 0.3  # Token2Wav needs less memory
-      enforce_eager: true
-      trust_remote_code: true
-      async_scheduling: false
-      enable_prefix_caching: false
-      max_num_batched_tokens: 4096
-      engine_output_type: audio    # Outputs waveform
-    engine_input_source: [0]       # Takes input from Stage 0
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav
-    final_output: true            # Final output stage
-    final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 1
-      seed: 42
-      detokenize: false
-
-# Runtime configuration
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1             # Wait for full input before triggering
-    max_inflight: 1             # Process one at a time
-
-  edges:
-    - from: 0                   # Thinker → Token2Wav
-      to: 1
-      window_size: -1           # Wait for complete output from Thinker

vllm_omni/model_executor/stage_configs/step_audio_2_asr.yaml

@@ -1,35 +0,0 @@
-# Stage config for Step-Audio2 ASR (speech-to-text only)
-#
-# Single-stage pipeline: only Thinker, no Token2Wav.
-# Use this for audio_to_text tasks where no audio output is needed.
-
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    runtime:
-      process: true
-      devices: "0,1"
-      max_batch_size: 1
-    engine_args:
-      model_stage: step_audio2_thinker
-      model_arch: StepAudio2ThinkerForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.8
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      max_num_batched_tokens: 2048
-      tensor_parallel_size: 2
-      engine_output_type: text
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
-    default_sampling_params:
-      temperature: 0.7
-      top_p: 0.9
-      top_k: -1
-      max_tokens: 1024
-      seed: 42
-      detokenize: true
-      repetition_penalty: 1.05

vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml

@@ -1,78 +0,0 @@
-# Stage config for Step-Audio2 with async_chunk support
-#
-# Two-stage pipeline with streaming between stages:
-#   Stage 0 (Thinker): Audio understanding → Text + Audio tokens (streamed in chunks)
-#   Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz)
-#
-# Key difference from step_audio_2.yaml:
-#   - async_chunk: true enables streaming from Thinker to Token2Wav
-#   - Thinker uses custom_process_next_stage_input_func to send audio token chunks
-#   - Token2Wav receives chunks incrementally instead of waiting for full output
-#
-# Hardware requirement: 1-2 GPUs (can run on single GPU with devices: "0" for both)
-
-async_chunk: true
-stage_args:
-  # Stage 0: Thinker (Audio Understanding + Token Generation)
-  - stage_id: 0
-    stage_type: llm
-    runtime:
-      process: true
-      devices: "0,1"
-      max_batch_size: 1
-    engine_args:
-      model_stage: step_audio2_thinker
-      model_arch: StepAudio2ThinkerForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.8
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      max_num_batched_tokens: 1024
-      tensor_parallel_size: 2
-      engine_output_type: latent
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav_async_chunk
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
-    default_sampling_params:
-      temperature: 0.7
-      top_p: 0.9
-      top_k: -1
-      max_tokens: 1024
-      seed: 42
-      detokenize: true
-      repetition_penalty: 1.05
-
-  # Stage 1: Token2Wav (Audio Synthesis)
-  - stage_id: 1
-    stage_type: llm
-    runtime:
-      process: true
-      devices: "3"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2wav
-      model_arch: StepAudio2Token2WavModel
-      worker_type: generation
-      hf_overrides:
-        architectures: ["StepAudio2Token2WavForConditionalGeneration"]
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      gpu_memory_utilization: 0.3
-      enforce_eager: true
-      trust_remote_code: true
-      async_scheduling: false
-      enable_prefix_caching: false
-      max_num_batched_tokens: 2048
-      engine_output_type: audio
-    engine_input_source: [0]
-    final_output: true
-    final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 1
-      seed: 42
-      detokenize: false

vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml

@@ -1,41 +0,0 @@
-# Stage config for running Wan2.2 TI2V-5B DiT with ModelOpt FP8 auto-detect.
-# Single GPU (TI2V-5B fits 80GB BF16; FP8 drops by ~half).
-# For the A14B MoE variants, bump `tensor_parallel_size` and `devices`.
-#
-# Use with a ModelOpt FP8 checkpoint (e.g. produced by
-# examples/quantization/quantize_wan2_2_modelopt_fp8.py).
-
-stage_args:
-  - stage_id: 0
-    stage_type: diffusion
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: dit
-      # Wan2.2 TI2V-5B is served by the base Wan22Pipeline, which auto-detects
-      # expand_timesteps (TI2V-5B mode) from model_index.json. (The standalone
-      # Wan22TI2VPipeline was consolidated into Wan22Pipeline upstream.)
-      model_class_name: Wan22Pipeline
-      max_num_seqs: 1
-      enforce_eager: true
-      trust_remote_code: true
-      distributed_executor_backend: "mp"
-      # Tile the VAE decode by default. Wan2.2's 480p peak is DiT-dominated (so
-      # FP8 is the main lever here), but tiling adds headroom for higher
-      # resolutions / longer clips at negligible cost. Override with --vae-use-tiling/false.
-      vae_use_tiling: true
-      parallel_config:
-        tensor_parallel_size: 1
-
-    final_output: true
-    final_output_type: video
-    is_comprehension: false
-    default_sampling_params:
-      seed: 42
-
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 1

vllm_omni/model_executor/stage_input_processors/mimo_audio.py

@@ -17,8 +17,8 @@ logger = init_logger(__name__)
 # sequence fed to stage-1 must not exceed this, otherwise gpu_input_batch
 # add_request will fail with a broadcast error when copying prompt_token_ids
 # into token_ids_cpu. Keep in sync with the stage-1 ``max_model_len`` in
-# ``vllm_omni/model_executor/stage_configs/mimo_audio.yaml`` and the offline
-# example ``examples/offline_inference/mimo_audio/end2end.py``.
+# ``vllm_omni/deploy/mimo_audio.yaml`` and the offline example
+# ``examples/offline_inference/mimo_audio/end2end.py``.
 MAX_CODE2WAV_TOKENS = 18192
 
 # Minimum safe values for codec streaming parameters.

vllm_omni/model_executor/stage_input_processors/step_audio2.py

@@ -7,6 +7,7 @@ import torch
 from vllm.inputs import TextPrompt
 from vllm.logger import init_logger
 
+from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import (
@@ -35,14 +36,14 @@ def _ensure_list(x):
 
 def thinker2token2wav_async_chunk(
     transfer_manager: Any,
-    pooling_output: dict[str, Any],
+    multimodal_output: dict[str, Any] | None,
     request: OmniEngineCoreRequest,
     is_finished: bool = False,
-) -> dict[str, Any] | None:
+) -> OmniPayloadStruct | None:
     """
     Async chunk processor: stream audio tokens from Thinker to Token2Wav.
 
-    Unlike Qwen3-Omni which passes hidden states via pooling_output,
+    Unlike Qwen3-Omni which passes hidden states via multimodal_output,
     Step Audio2 generates audio tokens as part of the LLM autoregressive
     output (token IDs >= audio_start). This processor extracts those audio
     tokens from the running token stream and sends them in chunks.
@@ -62,14 +63,13 @@ def thinker2token2wav_async_chunk(
     Args:
         transfer_manager: OmniChunkTransferAdapter instance (has
             ``code_prompt_token_ids`` defaultdict for state tracking).
-        pooling_output: Pooler output dict (unused — audio tokens are in
+        multimodal_output: Multimodal output dict (unused — audio tokens are in
             the token ID stream, not in hidden states).
         request: Current engine request with access to all_token_ids.
         is_finished: Whether the upstream request has finished generating.
 
     Returns:
-        dict with framework-compatible fields, or None if the chunk is
-        not yet ready.
+        Structured connector payload, or None if the chunk is not yet ready.
     """
     audio_start = DEFAULT_TOKEN_CONFIG.audio_start
     audio_eos = DEFAULT_TOKEN_CONFIG.audio_eos
@@ -100,18 +100,22 @@ def thinker2token2wav_async_chunk(
         # Last chunk: send all remaining tokens
         if available <= 0:
             # No audio tokens at all (text-only response) — send EOF marker
-            return {
-                "code_predictor_codes": [],
-                "left_context_size": 1,  # 1 = last chunk
-                "finished": torch.tensor(True, dtype=torch.bool),
-            }
+            return OmniPayloadStruct(
+                codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+                meta=MetaStruct(
+                    left_context_size=1,  # 1 = last chunk
+                    finished=torch.tensor(True, dtype=torch.bool),
+                ),
+            )
         remaining_tokens = audio_tokens[consumed:]
         transfer_manager.code_prompt_token_ids[request_id].extend(remaining_tokens)
-        return {
-            "code_predictor_codes": remaining_tokens,
-            "left_context_size": 1,  # 1 = last chunk
-            "finished": torch.tensor(True, dtype=torch.bool),
-        }
+        return OmniPayloadStruct(
+            codes=CodesStruct(audio=torch.tensor(remaining_tokens, dtype=torch.long)),
+            meta=MetaStruct(
+                left_context_size=1,  # 1 = last chunk
+                finished=torch.tensor(True, dtype=torch.bool),
+            ),
+        )
     else:
         # Non-last chunk: need chunk_size + pre_lookahead_len tokens
         required = chunk_size + pre_lookahead_len
@@ -120,11 +124,13 @@ def thinker2token2wav_async_chunk(
         chunk_tokens = audio_tokens[consumed : consumed + required]
         # Only consume chunk_size tokens; the lookahead portion is re-used
         transfer_manager.code_prompt_token_ids[request_id].extend(audio_tokens[consumed : consumed + chunk_size])
-        return {
-            "code_predictor_codes": chunk_tokens,
-            "left_context_size": 0,  # 0 = not last chunk
-            "finished": torch.tensor(False, dtype=torch.bool),
-        }
+        return OmniPayloadStruct(
+            codes=CodesStruct(audio=torch.tensor(chunk_tokens, dtype=torch.long)),
+            meta=MetaStruct(
+                left_context_size=0,  # 0 = not last chunk
+                finished=torch.tensor(False, dtype=torch.bool),
+            ),
+        )
 
 
 def thinker2token2wav(