[New Model] Step-Audio2 (#464)

Signed-off-by: linyueqian <linyueqian@outlook.com>
Signed-off-by: wuli666 <djjpro975@gmail.com>
Co-authored-by: linyueqian <linyueqian@outlook.com>
Co-authored-by: wuli666 <djjpro975@gmail.com>

.buildkite/test-merge.yml

@@ -164,6 +164,28 @@ steps:
   - group: ":card_index_dividers: E2E Test"
     depends_on: upload-merge-pipeline
     steps:
+      - label: "TTS · Step-Audio2 E2E Test"
+        timeout_in_minutes: 20
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              pytest -s -v tests/e2e/offline_inference/test_step_audio2.py -m 'advanced_model' --run-level 'advanced_model'
+            "
+        agents:
+          queue: "gpu_4_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
       - label: "TTS · Qwen3-TTS CustomVoice Test"
         source_file_dependencies:
           - tests/e2e/online_serving/test_qwen3_tts_customvoice.py

examples/offline_inference/step_audio2/README.md

@@ -0,0 +1,375 @@
+# Step-Audio2 Offline Inference Examples
+
+This directory contains examples for running offline inference with Step-Audio2 using vLLM-Omni.
+
+## Model Overview
+
+Step-Audio2 is a two-stage audio model:
+
+- **Stage 0 (Thinker)**: Audio understanding → Text + Audio tokens
+  - Input: Audio (16kHz)
+  - Output: Text transcription + Audio tokens for synthesis
+
+- **Stage 1 (Token2Wav)**: Audio synthesis
+  - Input: Audio tokens + Speaker prompt wav
+  - Output: Synthesized audio waveform (24kHz)
+
+## Hardware Requirements
+
+| Mode | GPU Configuration | VRAM Required |
+|------|-------------------|---------------|
+| ASR (S2T) | 1x GPU | ~20-25GB |
+| TTS/S2ST (single GPU) | 1x GPU | ~40-50GB |
+| TTS/S2ST (multi GPU) | 2x GPU | GPU0: ~28GB, GPU1: ~22GB |
+
+**Tested on:**
+- 1x NVIDIA H100 80GB (single-card S2ST)
+- 2x NVIDIA A10 40GB (multi-card S2ST)
+
+**Notes:**
+- Single GPU mode requires high VRAM due to both stages sharing memory
+- Multi GPU mode separates Stage 0 (Thinker) and Stage 1 (Token2Wav) across GPUs
+- VRAM usage can be adjusted via `gpu_memory_utilization` in stage config
+
+## Performance Benchmark
+
+### vLLM-Omni vs Official Step-Audio2
+
+Single request latency comparison between vLLM-Omni and official Step-Audio2 implementation.
+
+| Task | Tokens | vllm-omni | Step-Audio2 | Speedup |
+|------|--------|-----------|-------------|---------|
+| S2ST | ~85 | 5.45s | 7.36s | **1.35x** |
+| S2ST | ~160 | 6.67s | 13.92s | **2.09x** |
+| S2ST | ~315 | 9.42s | 31.50s | **3.34x** |
+| TTS | ~1024 | 16.65s | ~87s | **~5.2x** |
+
+**Key observations:**
+- Speedup increases with sequence length due to vLLM's efficient KV cache management
+- TTS (pure generation) shows the largest speedup (~5x)
+- S2ST benefits from optimized multi-stage pipeline
+
+**Benchmark environment:**
+- GPU: NVIDIA H100 80GB (single card)
+- Model: Step-Audio2-mini
+- Warmup: 1 run, Measured: 3 runs (averaged)
+
+### Async Chunk Streaming Performance
+
+Comparison between sequential (non-async) and async chunk modes via `/v1/audio/speech` TTS endpoint.
+
+| Mode | Mean TTFP | Mean E2E | Mean RTF | Audio Throughput |
+|------|-----------|----------|----------|-----------------|
+| Sequential | 4316ms | 4316ms | 0.938 | 1.07x realtime |
+| **Async Chunk** | **1437ms** | 4362ms | 0.949 | 1.06x realtime |
+| **Improvement** | **-67% (3x faster)** | ~same | ~same | ~same |
+
+**Key observations:**
+- Async chunk reduces **time-to-first-audio (TTFP) by 67%** (4.3s → 1.4s)
+- E2E latency remains comparable — async chunk overlaps Thinker decode with Token2Wav synthesis
+- RTF < 1 in both modes (real-time capable)
+- Sequential mode: TTFP ≈ E2E (must wait for all audio tokens before synthesis starts)
+- Async chunk mode: Token2Wav starts after first 28 tokens (chunk_size=25 + lookahead=3)
+
+**Benchmark environment:**
+- GPU: 4x NVIDIA RTX 3090 24GB (TP=2 for Thinker, 1 GPU for Token2Wav)
+- Model: Step-Audio2-mini
+- Endpoint: `/v1/audio/speech` (10 prompts, concurrency=1)
+- Measured via `bench_tts_serve.py`
+
+## Installation
+
+Make sure you have installed vLLM-Omni and all required dependencies:
+
+```bash
+# Install vLLM-Omni
+pip install vllm-omni
+
+# Install Step-Audio2 (REQUIRED for Token2Wav stage)
+pip install step-audio2
+```
+
+## Model Setup
+
+### Option 1: Auto-download from HuggingFace (Recommended)
+
+The script will **automatically download** the model on first run:
+
+```bash
+# Just run without specifying --model, it will auto-download stepfun-ai/Step-Audio2-mini
+python end2end.py --query-type audio_to_text
+
+# Or explicitly specify the HuggingFace model
+python end2end.py --query-type audio_to_text --model stepfun-ai/Step-Audio2-mini
+```
+
+Models will be cached in `~/.cache/huggingface/hub/` for future use.
+
+**Available models**:
+- `stepfun-ai/Step-Audio2-mini` (smaller, faster)
+- `stepfun-ai/Step-Audio2-7B` (larger, better quality)
+
+### Option 2: Manual Download (for offline use)
+
+Download and use locally:
+
+```bash
+# Download from HuggingFace
+huggingface-cli download stepfun-ai/Step-Audio2-mini --local-dir ./models/Step-Audio2-mini
+
+# Then use the local path
+python end2end.py --query-type audio_to_text --model ./models/Step-Audio2-mini
+```
+
+Ensure the model directory contains:
+```
+Step-Audio2-mini/
+├── config.json
+├── model.safetensors (or pytorch_model.bin)
+├── tokenizer.json
+├── tokenizer_config.json
+└── token2wav/                           # Token2Wav models (REQUIRED)
+    ├── speech_tokenizer_v2_25hz.onnx   # Audio tokenizer
+    ├── campplus.onnx                    # Speaker encoder
+    ├── flow.yaml                        # Flow model config
+    ├── flow.pt                          # Flow model weights
+    └── hift.pt                          # HiFT vocoder weights
+```
+
+## Usage Examples
+
+### 1. Audio to Text (ASR - Speech Recognition)
+
+Transcribe audio to text:
+
+```bash
+# Quick start - Using default model and test audio
+python end2end.py --query-type audio_to_text
+
+# Using your own audio file (model will auto-download)
+python end2end.py --query-type audio_to_text \
+    --audio-path /path/to/input.wav
+
+# With specific model
+python end2end.py --query-type audio_to_text \
+    --audio-path /path/to/input.wav \
+    --model stepfun-ai/Step-Audio2-7B
+
+# With custom question
+python end2end.py --query-type audio_to_text \
+    --audio-path input.wav \
+    --question "What is the speaker saying?"
+```
+
+**Output**: Text transcription saved to `output_step_audio2/00000_text.txt`
+
+### 2. Text to Audio (TTS - Speech Synthesis)
+
+Convert text to speech:
+
+```bash
+# Basic TTS (model auto-downloads)
+python end2end.py --query-type text_to_audio \
+    --text "Hello, this is a test of Step Audio 2 synthesis."
+
+# With specific model
+python end2end.py --query-type text_to_audio \
+    --text "Hello, this is a test." \
+    --model stepfun-ai/Step-Audio2-7B
+```
+
+**Note**: Speaker voice is controlled by the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable or the default prompt wav bundled with the model.
+
+**Output**:
+- Text: `output_step_audio2/00000_text.txt`
+- Audio: `output_step_audio2/00000_output.wav` (24kHz)
+
+### 3. Audio to Audio (Voice Conversion)
+
+Process input audio and generate output audio:
+
+```bash
+# Basic voice conversion (model auto-downloads)
+python end2end.py --query-type audio_to_audio \
+    --audio-path /path/to/source_audio.wav
+
+# With specific model
+python end2end.py --query-type audio_to_audio \
+    --audio-path source.wav \
+    --model stepfun-ai/Step-Audio2-7B
+```
+
+This mode:
+1. Understands the content in `--audio-path` (source)
+2. Generates audio output with the default voice
+
+**Note**: To use a custom speaker voice, set the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable.
+
+### Advanced Options
+
+```bash
+# Use custom stage configuration
+python end2end.py --query-type audio_to_text \
+    --stage-configs-path /path/to/custom_config.yaml
+
+# Multiple prompts (for batch testing)
+python end2end.py --query-type audio_to_text \
+    --audio-path input.wav \
+    --num-prompts 5
+
+# Custom output directory
+python end2end.py --query-type text_to_audio \
+    --text "Test synthesis" \
+    --output-dir ./my_outputs
+
+# Enable detailed logging
+python end2end.py --query-type audio_to_text \
+    --audio-path input.wav \
+    --enable-stats
+
+# Adjust generation parameters
+python end2end.py --query-type audio_to_text \
+    --audio-path input.wav \
+    --max-tokens 2048
+
+# Use custom speaker voice via environment variable
+STEP_AUDIO2_DEFAULT_PROMPT_WAV=/path/to/speaker.wav python end2end.py \
+    --query-type text_to_audio \
+    --text "Hello world"
+
+# Use Ray backend for distributed processing
+python end2end.py --query-type text_to_audio \
+    --text "Hello world" \
+    --worker-backend ray \
+    --ray-address "auto"
+```
+
+## Configuration
+
+### Stage Configuration
+
+The default configuration (`step_audio_2.yaml`) uses:
+
+- **Stage 0 (Thinker)**: GPU 0, 80% memory
+- **Stage 1 (Token2Wav)**: GPU 1, 30% memory
+
+For **single GPU** setup, edit the config to use `devices: "0"` for both stages.
+
+### Sampling Parameters
+
+- **Thinker (Stage 0)**:
+  - Temperature: 0.7 (balanced creativity)
+  - Top-p: 0.9
+  - Max tokens: 1024 (configurable)
+
+- **Token2Wav (Stage 1)**:
+  - Temperature: 0.0 (deterministic)
+  - Operates in generation mode (not sampling)
+
+## Common Issues
+
+### 1. ImportError: No module named 's3tokenizer'
+
+**Solution**: Install Step-Audio2 package:
+```bash
+pip install step-audio2
+```
+
+### 2. FileNotFoundError: prompt_wav file not found
+
+**Solution**: Set the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable to a valid audio file:
+```bash
+export STEP_AUDIO2_DEFAULT_PROMPT_WAV=/path/to/speaker.wav
+python end2end.py --query-type text_to_audio --text "Hello"
+```
+Or ensure the default prompt wav (`default_female.wav`) exists in your model directory.
+
+### 3. FileNotFoundError: token2wav models not found
+
+**Solution**: Ensure your model directory has the complete `token2wav/` subdirectory with all ONNX and PyTorch models.
+
+### 4. CUDA Out of Memory
+
+**Solutions**:
+- Use single GPU mode (set both stages to `devices: "0"`)
+- Reduce `gpu_memory_utilization` in config
+- Reduce `max_num_batched_tokens`
+- Process fewer prompts at once
+
+### 5. Model not found in registry
+
+**Solution**: Ensure you're using vLLM-Omni's entry point with `--omni` flag or install vllm-omni properly:
+```bash
+pip install vllm-omni
+```
+
+## Output Files
+
+The script generates files in the output directory (default: `output_step_audio2/`):
+
+```
+output_step_audio2/
+├── 00000_text.txt        # Text output from Thinker stage
+├── 00000_output.wav      # Audio output from Token2Wav stage (24kHz)
+├── 00001_text.txt        # (if multiple prompts)
+└── 00001_output.wav
+```
+
+## Performance Tips
+
+1. **First run is slow**: Stage initialization takes 20-60 seconds
+2. **Single GPU**: Set both stages to `devices: "0"` in config
+3. **Multiple prompts**: Use `--num-prompts N` for batch testing
+4. **Ray backend**: For multi-node or advanced scheduling
+5. **Logging**: Use `--enable-stats` to debug performance issues
+
+## Speaker Voice Configuration
+
+The Token2Wav stage requires a speaker prompt wav for voice conditioning. It is automatically resolved in this order:
+
+1. `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable (if set)
+2. `{model_dir}/assets/default_female.wav`
+3. `{model_dir}/default_female.wav`
+
+If none are found, set the environment variable explicitly:
+```bash
+export STEP_AUDIO2_DEFAULT_PROMPT_WAV=/path/to/speaker.wav
+```
+
+**Guidelines for custom speaker prompt:**
+
+- **Duration**: 3-10 seconds recommended
+- **Quality**: Clean audio, minimal background noise
+- **Format**: WAV, MP3, FLAC (will be resampled internally)
+- **Content**: Clear speech, representative of target voice
+
+## Example Workflow
+
+Complete example from audio to final output:
+
+```bash
+# 1. ASR: Transcribe audio
+python end2end.py --query-type audio_to_text \
+    --audio-path interview.wav \
+    --model ./models/Step-Audio2-7B \
+    --output-dir ./outputs
+
+# 2. Check the transcription
+cat ./outputs/00000_text.txt
+
+# 3. TTS: Synthesize new speech (with custom voice)
+STEP_AUDIO2_DEFAULT_PROMPT_WAV=./speaker_samples/female_voice.wav \
+python end2end.py --query-type text_to_audio \
+    --text "The quick brown fox jumps over the lazy dog" \
+    --model ./models/Step-Audio2-7B \
+    --output-dir ./outputs
+
+# 4. Listen to the result
+# Audio saved to: ./outputs/00000_output.wav
+```
+
+## References
+
+- [Step-Audio2 Paper](https://arxiv.org/abs/2507.16632)
+- [vLLM-Omni Documentation](https://vllm-omni.readthedocs.io)
+- [Model on HuggingFace](https://huggingface.co/stepfun-ai/Step-Audio-2-mini)

examples/offline_inference/step_audio2/end2end.py

@@ -0,0 +1,527 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+This example shows how to use vLLM-Omni for running offline inference
+with Step-Audio2 model for audio understanding and generation.
+
+Step-Audio2 is a two-stage model:
+  - Stage 0 (Thinker): Audio understanding → Text + Audio tokens
+  - Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz)
+
+Usage examples:
+  # Audio to text (ASR)
+  python end2end.py --query-type audio_to_text --audio-path input.wav --model-path "$MODEL_PATH"
+
+  # Text to audio (TTS with default voice)
+  python end2end.py --query-type text_to_audio --model-path "$MODEL_PATH"
+
+  # Audio to audio (Voice conversion with default voice)
+  python end2end.py --query-type audio_to_audio --audio-path input.wav --model-path "$MODEL_PATH"
+"""
+
+import os
+import sys
+from pathlib import Path
+from typing import NamedTuple
+
+import librosa  # noqa: TID251
+import numpy as np
+import soundfile as sf
+import torch
+from vllm.assets.audio import AudioAsset
+from vllm.sampling_params import SamplingParams
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.entrypoints.omni import Omni
+
+SEED = 42
+
+
+class QueryResult(NamedTuple):
+    inputs: dict
+    limit_mm_per_prompt: dict[str, int]
+
+
+def get_audio_to_text_query(
+    audio_path: str | None = None,
+    question: str | None = None,
+    sampling_rate: int = 16000,
+) -> QueryResult:
+    """
+    Audio Speech Recognition (ASR) - Convert audio to text.
+
+    Args:
+        audio_path: Path to audio file
+        question: Question about the audio (optional)
+        sampling_rate: Target sampling rate (16kHz for Step-Audio2)
+
+    Returns:
+        QueryResult with prompt and audio data
+    """
+    if question is None:
+        question = "请记录下你所听到的语音内容。"
+
+    prompt = (
+        f"<|im_start|>system\n{question}<|im_end|>\n<|im_start|>user\n<audio_patch><|im_end|>\n<|im_start|>assistant\n"
+    )
+
+    # Load audio
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        # Use default test audio
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+        # Resample to 16kHz if needed
+        if audio_data[1] != sampling_rate:
+            audio_signal = librosa.resample(audio_data[0], orig_sr=audio_data[1], target_sr=sampling_rate)
+            audio_data = (audio_signal.astype(np.float32), sampling_rate)
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "audio": audio_data,
+            },
+        },
+        limit_mm_per_prompt={"audio": 1},
+    )
+
+
+def get_text_to_audio_query(
+    text: str | None = None,
+    sampling_rate: int = 16000,
+) -> QueryResult:
+    """
+    Text-to-Speech (TTS) - Convert text to audio.
+
+    Args:
+        text: Text to synthesize
+        sampling_rate: Target sampling rate (16kHz for Step-Audio2)
+
+    Returns:
+        QueryResult with text prompt and speaker wav path
+
+    Note:
+        For audio generation, the prompt must end with "<tts_start>"
+        to signal the model to generate audio tokens instead of just text.
+    """
+    if text is None:
+        text = "Hello, this is a test of Step Audio 2 text to speech synthesis."
+
+    prompt = (
+        "<|im_start|>system\n"
+        "请逐字朗读用户提供的文本，不要改写或补充。<|im_end|>\n"
+        "<|im_start|>user\n"
+        f"{text}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+        "<tts_start>"
+    )
+
+    inputs = {
+        "prompt": prompt,
+    }
+
+    return QueryResult(
+        inputs=inputs,
+        limit_mm_per_prompt={},
+    )
+
+
+def get_audio_to_audio_query(
+    audio_path: str | None = None,
+    question: str | None = None,
+    sampling_rate: int = 16000,
+) -> QueryResult:
+    """
+    Audio-to-Audio - Voice conversion/cloning.
+
+    Args:
+        audio_path: Path to source audio file
+        question: Question/instruction about the audio
+        sampling_rate: Target sampling rate (16kHz for Step-Audio2)
+
+    Returns:
+        QueryResult with audio input
+
+    Note:
+        This mode processes input audio and generates output audio.
+        Speaker voice is controlled by STEP_AUDIO2_DEFAULT_PROMPT_WAV env var.
+    """
+    if question is None:
+        question = "请仔细聆听这段语音，然后复述其内容。"
+
+    prompt = (
+        "<|im_start|>system\n"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>user\n"
+        "<audio_patch><|im_end|>\n"
+        "<|im_start|>assistant\n"
+        "<tts_start>"
+    )
+
+    # Load source audio
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+        if audio_data[1] != sampling_rate:
+            audio_signal = librosa.resample(audio_data[0], orig_sr=audio_data[1], target_sr=sampling_rate)
+            audio_data = (audio_signal.astype(np.float32), sampling_rate)
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "audio": audio_data,
+            },
+        },
+        limit_mm_per_prompt={"audio": 1},
+    )
+
+
+query_map = {
+    "audio_to_text": get_audio_to_text_query,
+    "text_to_audio": get_text_to_audio_query,
+    "audio_to_audio": get_audio_to_audio_query,
+}
+
+
+def main(args):
+    # Use local model path or HuggingFace model name
+    if not args.model:
+        sys.exit(1)
+
+    model_name = args.model
+
+    # Check if it's a local path
+    if os.path.isdir(model_name):
+        # Verify essential files exist
+        config_path = os.path.join(model_name, "config.json")
+        if not os.path.exists(config_path):
+            sys.exit(1)
+
+    # Get query configuration
+    query_func = query_map[args.query_type]
+
+    # Build query based on type
+    if args.query_type == "audio_to_text":
+        query_result = query_func(
+            audio_path=args.audio_path,
+            question=args.question,
+            sampling_rate=args.sampling_rate,
+        )
+    elif args.query_type == "text_to_audio":
+        query_result = query_func(
+            text=args.text,
+            sampling_rate=args.sampling_rate,
+        )
+    elif args.query_type == "audio_to_audio":
+        query_result = query_func(
+            audio_path=args.audio_path,
+            question=args.question,
+            sampling_rate=args.sampling_rate,
+        )
+    else:
+        raise ValueError(f"Unknown query type: {args.query_type}")
+
+    # Initialize vLLM-Omni with Step-Audio2
+    # Resolve stage config path
+    if args.stage_configs_path:
+        stage_config_path = args.stage_configs_path
+    else:
+        configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/model_executor/stage_configs"
+        if args.query_type == "audio_to_text":
+            # ASR only needs Thinker (Stage 0), no Token2Wav
+            stage_config_path = str(configs_dir / "step_audio_2_asr.yaml")
+            if not os.path.exists(stage_config_path):
+                stage_config_path = str(configs_dir / "step_audio_2.yaml")
+        else:
+            # TTS/S2ST need both stages
+            stage_config_path = str(configs_dir / "step_audio_2.yaml")
+
+    omni_llm = Omni(
+        model=model_name,
+        log_stats=args.enable_stats,
+        log_file=("step_audio2_pipeline.log" if args.enable_stats else None),
+        init_sleep_seconds=args.init_sleep_seconds,
+        batch_timeout=args.batch_timeout,
+        init_timeout=args.init_timeout,
+        shm_threshold_bytes=args.shm_threshold_bytes,
+        stage_configs_path=stage_config_path,
+        worker_backend=args.worker_backend,
+        ray_address=args.ray_address,
+    )
+
+    # Configure sampling parameters for each stage
+    # Stage 0 (Thinker): Generate text and audio tokens
+    # Note: For audio generation tasks, we need to allow the model to generate audio tokens
+    # Audio tokens are in range [151696, 158257], so we should not stop too early
+    # Qwen EOS token is typically 151645 (<|im_end|>), but we want to allow audio generation
+    # So we don't set stop_token_ids for audio tasks, but we do set it for text-only tasks
+    stop_token_ids = None
+    if args.query_type == "audio_to_text":
+        # For ASR, stop at EOS token
+        stop_token_ids = [151645]  # Qwen EOS token <|im_end|>
+
+    # Adjust sampling parameters based on query type
+    # For audio generation tasks, we need higher repetition_penalty to avoid text loops
+    # and allow the model to transition to audio token generation
+    repetition_penalty = 1.1 if args.query_type in ["text_to_audio", "audio_to_audio"] else 1.05
+
+    thinker_sampling_params = SamplingParams(
+        temperature=0.7,
+        top_p=0.9,
+        top_k=-1,
+        max_tokens=args.max_tokens,
+        seed=SEED,
+        detokenize=True,
+        repetition_penalty=repetition_penalty,
+        stop_token_ids=stop_token_ids,
+    )
+
+    # Stage 1 (Token2Wav): Convert audio tokens to waveform
+    # Note: This stage doesn't actually sample, it's deterministic generation
+    token2wav_sampling_params = SamplingParams(
+        temperature=0.0,
+        top_p=1.0,
+        top_k=-1,
+        max_tokens=1,
+        seed=SEED,
+        detokenize=False,
+    )
+
+    if args.query_type == "audio_to_text":
+        # ASR: single stage, only Thinker
+        sampling_params_list = [thinker_sampling_params]
+    else:
+        # TTS/S2ST: two stages
+        sampling_params_list = [thinker_sampling_params, token2wav_sampling_params]
+
+    # Prepare prompts
+    if args.num_prompts > 1:
+        prompts = [query_result.inputs for _ in range(args.num_prompts)]
+    else:
+        prompts = [query_result.inputs]
+
+    # Run inference
+    omni_outputs = omni_llm.generate(prompts, sampling_params_list)
+
+    # Create output directory
+    output_dir = args.output_dir
+    os.makedirs(output_dir, exist_ok=True)
+
+    def _get_request_index(req_id: str, fallback: int) -> int:
+        """Parse request index from vLLM request_id.
+
+        vLLM may return IDs like "0_<uuid>". Use the leading integer if present,
+        otherwise fall back to the loop index.
+        """
+        try:
+            return int(req_id)
+        except (TypeError, ValueError):
+            if isinstance(req_id, str):
+                head = req_id.split("_", 1)[0]
+                if head.isdigit():
+                    return int(head)
+        return fallback
+
+    # Process outputs from each stage
+    for stage_idx, stage_outputs in enumerate(omni_outputs):
+        # request_output may be a single object or a list
+        raw_output = stage_outputs.request_output
+        outputs_list = raw_output if isinstance(raw_output, list) else [raw_output]
+
+        if stage_outputs.final_output_type == "text":
+            # Stage 0 (Thinker) text output
+            for i, output in enumerate(outputs_list):
+                request_id = _get_request_index(output.request_id, i)
+                text_output = output.outputs[0].text
+
+                # Save to file
+                out_txt = os.path.join(output_dir, f"{request_id:05d}_text.txt")
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.write(f"Prompt:\n{prompts[request_id]['prompt']}\n\n")
+                    f.write(f"Output:\n{text_output}\n")
+                print(f"Text output saved to: {out_txt}")
+                print(f"Text: {text_output[:200]}")
+
+        elif stage_outputs.final_output_type == "audio":
+            # Stage 1 (Token2Wav) audio output
+            for i, output in enumerate(outputs_list):
+                request_id = _get_request_index(output.request_id, i)
+
+                # Get audio tensor (24kHz)
+                mm_out = {}
+                if hasattr(output, "multimodal_output") and output.multimodal_output:
+                    mm_out = output.multimodal_output
+                elif output.outputs and hasattr(output.outputs[0], "multimodal_output"):
+                    mm_out = output.outputs[0].multimodal_output or {}
+                audio_tensor = None
+                for key in ("audio", "wav", "waveform", "audio_pcm", "pcm", "model_outputs"):
+                    if key in mm_out and mm_out[key] is not None:
+                        audio_tensor = mm_out[key]
+                        break
+
+                if audio_tensor is not None:
+                    # Save audio file
+                    output_wav = os.path.join(output_dir, f"{request_id:05d}_output.wav")
+                    if isinstance(audio_tensor, list):
+                        chunks = []
+                        for chunk in audio_tensor:
+                            if chunk is None:
+                                continue
+                            if isinstance(chunk, torch.Tensor):
+                                chunk_np = chunk.detach().cpu().numpy()
+                            else:
+                                chunk_np = np.asarray(chunk)
+                            if chunk_np.size > 0:
+                                chunks.append(chunk_np)
+                        if not chunks:
+                            continue
+                        audio_numpy = np.concatenate(chunks, axis=-1)
+                    elif isinstance(audio_tensor, torch.Tensor):
+                        audio_numpy = audio_tensor.detach().cpu().numpy()
+                    else:
+                        audio_numpy = np.asarray(audio_tensor)
+                    sf.write(output_wav, audio_numpy, samplerate=24000)
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Demo on using vLLM-Omni for offline inference with Step-Audio2")
+
+    # Model and config
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        required=True,
+        help=(
+            "Model path (required). Can be:\n"
+            "  - Local path: /path/to/step-audio-2\n"
+            "  - HuggingFace ID: stepfun-ai/Step-Audio2-mini"
+        ),
+    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to stage config YAML file (default: use built-in step_audio_2.yaml)",
+    )
+
+    # Query configuration
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="audio_to_text",
+        choices=query_map.keys(),
+        help="Query type: audio_to_text (ASR), text_to_audio (TTS), audio_to_audio (voice conversion)",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to input audio file (for audio_to_text and audio_to_audio)",
+    )
+    parser.add_argument(
+        "--text",
+        "-t",
+        type=str,
+        default=None,
+        help="Text to synthesize (for text_to_audio mode)",
+    )
+    parser.add_argument(
+        "--question",
+        type=str,
+        default=None,
+        help="Question/instruction about the audio",
+    )
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=16000,
+        help="Audio sampling rate for input (default: 16000 Hz)",
+    )
+
+    # Generation parameters
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=1024,
+        help="Maximum tokens to generate in Thinker stage (default: 1024)",
+    )
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of times to repeat the prompt (for testing)",
+    )
+
+    # System configuration
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable detailed statistics logging",
+    )
+    parser.add_argument(
+        "--init-sleep-seconds",
+        type=int,
+        default=20,
+        help="Sleep seconds after starting each stage (default: 20)",
+    )
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Batch timeout in seconds (default: 5)",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=300,
+        help="Initialization timeout in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Shared memory threshold in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--worker-backend",
+        type=str,
+        default="multi_process",
+        choices=["multi_process", "ray"],
+        help="Worker backend (default: multi_process)",
+    )
+    parser.add_argument(
+        "--ray-address",
+        type=str,
+        default=None,
+        help="Ray cluster address (for ray backend)",
+    )
+
+    # Output
+    parser.add_argument(
+        "--output-dir",
+        "-o",
+        type=str,
+        default="output_step_audio2",
+        help="Output directory for generated files (default: output_step_audio2)",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/online_serving/step_audio2/README.md

@@ -0,0 +1,193 @@
+# Step-Audio2 Online Serving
+
+This directory contains examples for running Step-Audio2 with vLLM-Omni's online serving API.
+
+## Installation
+
+Please refer to [README.md](../../../README.md)
+
+## Launch the Server
+
+```bash
+# Async chunk mode (recommended — lower first-packet latency for TTS)
+vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml \
+    --trust-remote-code --enforce-eager
+```
+
+Sequential mode:
+```bash
+vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/step_audio_2.yaml \
+    --trust-remote-code --enforce-eager
+```
+
+With local model:
+```bash
+vllm serve /path/to/Step-Audio-2-mini --omni --port 8092 \
+    --trust-remote-code --enforce-eager
+```
+
+## Send Requests
+
+### TTS via `/v1/audio/speech` (Recommended)
+
+```bash
+cd examples/online_serving/step_audio2
+
+# Python client
+python openai_speech_client.py --text "你好世界"
+
+# With custom system prompt
+python openai_speech_client.py --text "Hello, how are you?" \
+    --instructions "You are a friendly assistant."
+
+# Save to specific file
+python openai_speech_client.py --text "你好世界" -o output.wav
+```
+
+Or via curl:
+
+```bash
+curl -X POST http://localhost:8092/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{"model":"stepfun-ai/Step-Audio2-mini","input":"你好世界","voice":"default"}' \
+    --output output.wav
+```
+
+This endpoint bypasses the chat template and directly triggers TTS mode. It supports async chunk streaming for low first-packet latency.
+
+**Note**: Speaker voice is controlled by `STEP_AUDIO2_DEFAULT_PROMPT_WAV` env var on the server side.
+
+### Chat Completions (ASR / S2ST)
+
+```bash
+# Audio to Text (ASR)
+python openai_chat_completion_client.py --query-type audio_to_text
+
+# Audio to Audio (S2ST)
+python openai_chat_completion_client.py --query-type audio_to_audio --audio-path /path/to/input.wav
+```
+
+| Argument | Description |
+|----------|-------------|
+| `--query-type`, `-q` | Query type: `audio_to_text`, `text_to_audio`, `audio_to_audio` |
+| `--audio-path`, `-a` | Path to input audio file (local or URL) |
+| `--text`, `-t` | Text to synthesize (for TTS mode) |
+| `--prompt`, `-p` | Custom prompt/question |
+| `--output-dir`, `-o` | Output directory for audio files (default: `output_online`) |
+| `--api-base` | API base URL (default: `http://localhost:8092/v1`) |
+
+### Curl (Chat Completions)
+
+```bash
+# Audio to Text
+bash run_curl.sh audio_to_text
+
+# Text to Audio
+bash run_curl.sh text_to_audio
+
+# Audio to Audio
+bash run_curl.sh audio_to_audio
+```
+
+## Query Types
+
+### 1. Audio to Text (ASR)
+
+Transcribe audio to text.
+
+```bash
+python openai_chat_completion_client.py \
+    --query-type audio_to_text \
+    --audio-path /path/to/speech.wav \
+    --prompt "Transcribe this audio."
+```
+
+### 2. Text to Audio (TTS)
+
+Convert text to speech:
+
+```bash
+# Via speech endpoint (recommended, returns WAV directly)
+python openai_speech_client.py --text "Hello, welcome to Step-Audio2."
+
+# Via chat completions
+python openai_chat_completion_client.py \
+    --query-type text_to_audio \
+    --text "Hello, welcome to Step-Audio2."
+```
+
+### 3. Audio to Audio (S2ST)
+
+Process input audio and generate text transcription + audio output.
+
+```bash
+python openai_chat_completion_client.py \
+    --query-type audio_to_audio \
+    --audio-path /path/to/source.wav
+```
+
+## Output
+
+- **Text output**: Printed to console
+- **Audio output**: Saved to `output_online/audio_0.wav` (24kHz WAV)
+
+## API Format
+
+Step-Audio2 uses the OpenAI-compatible chat completions API:
+
+```json
+{
+  "model": "stepfun-ai/Step-Audio2-mini",
+  "messages": [
+    {
+      "role": "system",
+      "content": [{"type": "text", "text": "Transcribe the audio."}]
+    },
+    {
+      "role": "user",
+      "content": [
+        {"type": "audio_url", "audio_url": {"url": "..."}},
+        {"type": "text", "text": "Please transcribe."}
+      ]
+    }
+  ],
+  "sampling_params_list": [
+    {"temperature": 0.7, "max_tokens": 1024},
+    {"temperature": 0.0, "max_tokens": 1}
+  ]
+}
+```
+
+## Performance
+
+### Async Chunk vs Sequential
+
+Benchmark via `/v1/audio/speech` (4x RTX 3090, 10 prompts, concurrency=1):
+
+| Mode | Mean TTFP | Mean E2E | Mean RTF |
+|------|-----------|----------|----------|
+| Sequential | 4316ms | 4316ms | 0.938 |
+| **Async Chunk** | **1437ms** | 4362ms | 0.949 |
+
+Async chunk reduces TTFP by **67%** by streaming audio token chunks from Thinker to Token2Wav as they are generated. RTF < 1 in both modes (real-time capable).
+
+## Troubleshooting
+
+### Server not responding
+- Check if the server is running: `curl http://localhost:8092/v1/models`
+- Verify the port number matches
+
+### FileNotFoundError: prompt_wav file not found
+- Ensure `default_female.wav` exists at `{model_dir}/assets/default_female.wav`
+- Or set `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable when launching the server
+
+### Audio not generated
+- For TTS, use the `/v1/audio/speech` endpoint (recommended) or `openai_speech_client.py`
+- For chat completions TTS, ensure the prompt ends with `<tts_start>`
+- Check server logs for errors
+
+### Out of memory
+- Reduce `gpu_memory_utilization` in stage configs
+- Use a smaller batch size

examples/online_serving/step_audio2/openai_chat_completion_client.py

@@ -0,0 +1,342 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+OpenAI-compatible client for Step-Audio2 online serving.
+
+Step-Audio2 is a two-stage model:
+  - Stage 0 (Thinker): Audio understanding -> Text + Audio tokens
+  - Stage 1 (Token2Wav): Audio tokens -> Waveform (24kHz)
+
+Usage examples:
+  # Start server first:
+  vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092
+
+  # Audio to text (ASR)
+  python openai_chat_completion_client.py --query-type audio_to_text
+
+  # Text to audio (TTS)
+  python openai_chat_completion_client.py --query-type text_to_audio --text "Hello world"
+
+  # Audio to audio (Voice conversion)
+  python openai_chat_completion_client.py --query-type audio_to_audio --audio-path input.wav
+"""
+
+import base64
+import logging
+import os
+
+import requests
+from openai import OpenAI
+from vllm.assets.audio import AudioAsset
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+logger = logging.getLogger(__name__)
+
+SEED = 42
+
+
+def encode_base64_content_from_file(file_path: str) -> str:
+    """Encode a local file to base64 format."""
+    with open(file_path, "rb") as f:
+        return base64.b64encode(f.read()).decode("utf-8")
+
+
+def encode_base64_content_from_url(url: str) -> str:
+    """Encode content from URL to base64 format."""
+    with requests.get(url) as response:
+        response.raise_for_status()
+        return base64.b64encode(response.content).decode("utf-8")
+
+
+def get_audio_url(audio_path: str | None = None) -> str:
+    """Convert audio path to URL format for the API.
+
+    Args:
+        audio_path: Local file path or URL. If None, uses default test audio.
+
+    Returns:
+        Audio URL (either original URL or base64 data URL)
+    """
+    if not audio_path:
+        return AudioAsset("mary_had_lamb").url
+
+    if audio_path.startswith(("http://", "https://")):
+        return audio_path
+
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+
+    # Detect MIME type
+    audio_path_lower = audio_path.lower()
+    if audio_path_lower.endswith(".wav"):
+        mime_type = "audio/wav"
+    elif audio_path_lower.endswith((".mp3", ".mpeg")):
+        mime_type = "audio/mpeg"
+    elif audio_path_lower.endswith(".ogg"):
+        mime_type = "audio/ogg"
+    elif audio_path_lower.endswith(".flac"):
+        mime_type = "audio/flac"
+    else:
+        mime_type = "audio/wav"
+
+    audio_base64 = encode_base64_content_from_file(audio_path)
+    return f"data:{mime_type};base64,{audio_base64}"
+
+
+def get_audio_to_text_query(
+    audio_path: str | None = None,
+    custom_prompt: str | None = None,
+) -> dict:
+    """Build query for Audio-to-Text (ASR) mode."""
+    question = custom_prompt or "Please transcribe the audio content."
+    audio_url = get_audio_url(audio_path)
+
+    return {
+        "role": "user",
+        "content": [
+            {"type": "audio_url", "audio_url": {"url": audio_url}},
+            {"type": "text", "text": question},
+        ],
+    }
+
+
+def get_text_to_audio_query(
+    text: str | None = None,
+    custom_prompt: str | None = None,
+) -> dict:
+    """Build query for Text-to-Audio (TTS) mode.
+
+    Returns a list of messages: user message + assistant message with <tts_start>.
+    The assistant message triggers TTS mode and must be used with
+    continue_final_message=True so the chat template does not append <|im_end|>.
+    """
+    text_to_speak = text or "Hello, this is a test of Step Audio 2 text to speech."
+
+    return [
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": text_to_speak}],
+        },
+        {
+            "role": "assistant",
+            "content": "<tts_start>",
+        },
+    ]
+
+
+def get_audio_to_audio_query(
+    audio_path: str | None = None,
+    custom_prompt: str | None = None,
+) -> dict:
+    """Build query for Audio-to-Audio (S2ST) mode.
+
+    Returns a list of messages: user message (with audio) + assistant message
+    with <tts_start>. Must be used with continue_final_message=True.
+    """
+    question = custom_prompt or "Please listen to this audio and repeat its content."
+    audio_url = get_audio_url(audio_path)
+
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "audio_url", "audio_url": {"url": audio_url}},
+                {"type": "text", "text": question},
+            ],
+        },
+        {
+            "role": "assistant",
+            "content": "<tts_start>",
+        },
+    ]
+
+
+def get_system_prompt(query_type: str) -> dict:
+    """Get system prompt based on query type."""
+    if query_type == "audio_to_text":
+        system_text = "You are a speech recognition assistant. Transcribe the audio accurately."
+    elif query_type == "text_to_audio":
+        system_text = "You are a text-to-speech assistant. Read the text aloud exactly as provided."
+    else:  # audio_to_audio
+        system_text = "You are an audio processing assistant. Listen and repeat the audio content."
+
+    return {
+        "role": "system",
+        "content": [{"type": "text", "text": system_text}],
+    }
+
+
+QUERY_MAP = {
+    "audio_to_text": get_audio_to_text_query,
+    "text_to_audio": get_text_to_audio_query,
+    "audio_to_audio": get_audio_to_audio_query,
+}
+
+
+def run_step_audio2(args) -> None:
+    """Run Step-Audio2 inference via OpenAI-compatible API."""
+
+    # Initialize OpenAI client
+    client = OpenAI(
+        api_key="EMPTY",
+        base_url=args.api_base,
+    )
+
+    model_name = args.model
+
+    # Sampling parameters for each stage
+    thinker_sampling_params = {
+        "temperature": 0.7,
+        "top_p": 0.9,
+        "top_k": -1,
+        "max_tokens": args.max_tokens,
+        "seed": SEED,
+        "detokenize": True,
+        "repetition_penalty": 1.1 if args.query_type != "audio_to_text" else 1.05,
+    }
+
+    # For ASR, stop at EOS token
+    if args.query_type == "audio_to_text":
+        thinker_sampling_params["stop_token_ids"] = [151645]
+
+    token2wav_sampling_params = {
+        "temperature": 0.0,
+        "top_p": 1.0,
+        "top_k": -1,
+        "max_tokens": 1,
+        "seed": SEED,
+        "detokenize": False,
+    }
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        token2wav_sampling_params,
+    ]
+
+    # Build query based on type
+    query_func = QUERY_MAP[args.query_type]
+    if args.query_type == "audio_to_text":
+        user_message = query_func(audio_path=args.audio_path, custom_prompt=args.prompt)
+    elif args.query_type == "text_to_audio":
+        user_message = query_func(text=args.text, custom_prompt=args.prompt)
+    else:  # audio_to_audio
+        user_message = query_func(audio_path=args.audio_path, custom_prompt=args.prompt)
+
+    system_message = get_system_prompt(args.query_type)
+
+    logger.info(f"Query type: {args.query_type}")
+    logger.info(f"Model: {model_name}")
+    logger.info("Sending request to server...")
+
+    # Build messages list.
+    # TTS and S2ST return a list of messages (user + assistant with <tts_start>);
+    # ASR returns a single user message dict.
+    if isinstance(user_message, list):
+        messages = [system_message] + user_message
+    else:
+        messages = [system_message, user_message]
+
+    # For TTS/S2ST, use continue_final_message so <tts_start> in the
+    # assistant turn is not followed by <|im_end|>.
+    extra_body: dict = {"sampling_params_list": sampling_params_list}
+    if args.query_type in ("text_to_audio", "audio_to_audio"):
+        extra_body["continue_final_message"] = True
+        extra_body["add_generation_prompt"] = False
+
+    # Send request
+    chat_completion = client.chat.completions.create(
+        messages=messages,
+        model=model_name,
+        extra_body=extra_body,
+    )
+
+    # Process response
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    audio_count = 0
+    for choice in chat_completion.choices:
+        if choice.message.audio:
+            # Audio response
+            audio_data = base64.b64decode(choice.message.audio.data)
+            audio_file = os.path.join(args.output_dir, f"audio_{audio_count}.wav")
+            with open(audio_file, "wb") as f:
+                f.write(audio_data)
+            logger.info(f"Audio saved to: {audio_file}")
+            audio_count += 1
+
+        if choice.message.content:
+            # Text response
+            logger.info(f"Text output: {choice.message.content}")
+
+    if audio_count == 0 and args.query_type != "audio_to_text":
+        logger.warning("No audio output received. Check server logs for details.")
+
+    logger.info("Done!")
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="OpenAI-compatible client for Step-Audio2 online serving")
+
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="audio_to_text",
+        choices=QUERY_MAP.keys(),
+        help="Query type: audio_to_text (ASR), text_to_audio (TTS), audio_to_audio",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to input audio file (local path or URL)",
+    )
+    parser.add_argument(
+        "--text",
+        "-t",
+        type=str,
+        default=None,
+        help="Text to synthesize (for text_to_audio mode)",
+    )
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        type=str,
+        default=None,
+        help="Custom prompt/question",
+    )
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        default="stepfun-ai/Step-Audio2-mini",
+        help="Model name",
+    )
+    parser.add_argument(
+        "--api-base",
+        type=str,
+        default="http://localhost:8092/v1",
+        help="API base URL",
+    )
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=1024,
+        help="Maximum tokens for Thinker stage",
+    )
+    parser.add_argument(
+        "--output-dir",
+        "-o",
+        type=str,
+        default="output_online",
+        help="Output directory for audio files",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_step_audio2(args)

examples/online_serving/step_audio2/openai_speech_client.py

@@ -0,0 +1,130 @@
+"""OpenAI-compatible client for Step-Audio2 TTS via /v1/audio/speech endpoint.
+
+Examples:
+    # Basic TTS
+    python openai_speech_client.py --text "你好世界"
+
+    # With custom system prompt
+    python openai_speech_client.py --text "Hello, how are you?" \
+        --instructions "You are a friendly assistant."
+
+    # Save to specific file
+    python openai_speech_client.py --text "你好世界" -o output.wav
+"""
+
+import argparse
+
+import httpx
+
+DEFAULT_API_BASE = "http://localhost:8092"
+DEFAULT_API_KEY = "EMPTY"
+
+
+def run_tts_generation(args) -> None:
+    """Run TTS generation via /v1/audio/speech API."""
+    payload = {
+        "model": args.model,
+        "input": args.text,
+        "voice": args.voice,
+        "response_format": args.response_format,
+    }
+
+    if args.instructions:
+        payload["instructions"] = args.instructions
+
+    print(f"Model: {args.model}")
+    print(f"Text: {args.text}")
+    print(f"Voice: {args.voice}")
+    print("Generating audio...")
+
+    api_url = f"{args.api_base}/v1/audio/speech"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {args.api_key}",
+    }
+
+    with httpx.Client(timeout=300.0) as client:
+        response = client.post(api_url, json=payload, headers=headers)
+
+    if response.status_code != 200:
+        print(f"Error: {response.status_code}")
+        print(response.text)
+        return
+
+    try:
+        text = response.content.decode("utf-8")
+        if text.startswith('{"error"'):
+            print(f"Error: {text}")
+            return
+    except UnicodeDecodeError:
+        pass
+
+    output_path = args.output or "tts_output.wav"
+    with open(output_path, "wb") as f:
+        f.write(response.content)
+    print(f"Audio saved to: {output_path}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Step-Audio2 TTS client via /v1/audio/speech",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+    parser.add_argument(
+        "--api-base",
+        type=str,
+        default=DEFAULT_API_BASE,
+        help=f"API base URL (default: {DEFAULT_API_BASE})",
+    )
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=DEFAULT_API_KEY,
+        help="API key (default: EMPTY)",
+    )
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        default="stepfun-ai/Step-Audio2-mini",
+        help="Model name/path",
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        required=True,
+        help="Text to synthesize",
+    )
+    parser.add_argument(
+        "--voice",
+        type=str,
+        default="default",
+        help="Voice name (default: default)",
+    )
+    parser.add_argument(
+        "--instructions",
+        type=str,
+        default=None,
+        help="System prompt for the Thinker stage",
+    )
+    parser.add_argument(
+        "--response-format",
+        type=str,
+        default="wav",
+        choices=["wav", "pcm"],
+        help="Audio output format (default: wav)",
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=str,
+        default=None,
+        help="Output audio file path (default: tts_output.wav)",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_tts_generation(args)

examples/online_serving/step_audio2/run_curl.sh

@@ -0,0 +1,149 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Step-Audio2 curl client for online serving
+# Usage: bash run_curl.sh [audio_to_text|text_to_audio|audio_to_audio]
+
+QUERY_TYPE="${1:-audio_to_text}"
+API_BASE="${API_BASE:-http://localhost:8092}"
+
+# Validate query type
+if [[ ! "$QUERY_TYPE" =~ ^(audio_to_text|text_to_audio|audio_to_audio)$ ]]; then
+    echo "Error: Invalid query type '$QUERY_TYPE'"
+    echo "Usage: $0 [audio_to_text|text_to_audio|audio_to_audio]"
+    echo "  audio_to_text: Speech recognition (ASR)"
+    echo "  text_to_audio: Text-to-speech (TTS)"
+    echo "  audio_to_audio: Voice conversion"
+    exit 1
+fi
+
+SEED=42
+
+# Default test audio URL
+MARY_HAD_LAMB_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
+
+# Sampling parameters for Thinker stage
+thinker_sampling_params='{
+  "temperature": 0.7,
+  "top_p": 0.9,
+  "top_k": -1,
+  "max_tokens": 1024,
+  "seed": 42,
+  "detokenize": true,
+  "repetition_penalty": 1.05
+}'
+
+# Sampling parameters for Token2Wav stage
+token2wav_sampling_params='{
+  "temperature": 0.0,
+  "top_p": 1.0,
+  "top_k": -1,
+  "max_tokens": 1,
+  "seed": 42,
+  "detokenize": false
+}'
+
+# Build request based on query type
+case "$QUERY_TYPE" in
+  audio_to_text)
+    system_content='[{"type": "text", "text": "You are a speech recognition assistant. Transcribe the audio accurately."}]'
+    user_content='[
+      {"type": "audio_url", "audio_url": {"url": "'"$MARY_HAD_LAMB_URL"'"}},
+      {"type": "text", "text": "Please transcribe this audio."}
+    ]'
+    # Add stop token for ASR
+    thinker_sampling_params='{
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "top_k": -1,
+      "max_tokens": 1024,
+      "seed": 42,
+      "detokenize": true,
+      "repetition_penalty": 1.05,
+      "stop_token_ids": [151645]
+    }'
+    ;;
+  text_to_audio)
+    system_content='[{"type": "text", "text": "You are a text-to-speech assistant. Read the text aloud exactly as provided."}]'
+    user_content='[
+      {"type": "text", "text": "Hello, this is a test of Step Audio 2 text to speech synthesis.<tts_start>"}
+    ]'
+    thinker_sampling_params='{
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "top_k": -1,
+      "max_tokens": 1024,
+      "seed": 42,
+      "detokenize": true,
+      "repetition_penalty": 1.1
+    }'
+    ;;
+  audio_to_audio)
+    system_content='[{"type": "text", "text": "You are an audio processing assistant. Listen and repeat the audio content."}]'
+    user_content='[
+      {"type": "audio_url", "audio_url": {"url": "'"$MARY_HAD_LAMB_URL"'"}},
+      {"type": "text", "text": "Please listen to this audio and repeat its content.<tts_start>"}
+    ]'
+    thinker_sampling_params='{
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "top_k": -1,
+      "max_tokens": 1024,
+      "seed": 42,
+      "detokenize": true,
+      "repetition_penalty": 1.1
+    }'
+    ;;
+esac
+
+sampling_params_list='[
+  '"$thinker_sampling_params"',
+  '"$token2wav_sampling_params"'
+]'
+
+echo "Query type: $QUERY_TYPE"
+echo "API base: $API_BASE"
+echo "Sending request..."
+echo ""
+
+output=$(curl -sS -X POST "${API_BASE}/v1/chat/completions" \
+    -H "Content-Type: application/json" \
+    -d @- <<EOF
+{
+  "model": "stepfun-ai/Step-Audio2-mini",
+  "sampling_params_list": $sampling_params_list,
+  "messages": [
+    {
+      "role": "system",
+      "content": $system_content
+    },
+    {
+      "role": "user",
+      "content": $user_content
+    }
+  ]
+}
+EOF
+)
+
+# Extract and display text content
+text_content=$(echo "$output" | jq -r '.choices[0].message.content // empty')
+if [[ -n "$text_content" ]]; then
+    echo "Text output: $text_content"
+fi
+
+# Check for audio content
+audio_data=$(echo "$output" | jq -r '.choices[0].message.audio.data // empty')
+if [[ -n "$audio_data" ]]; then
+    echo "Audio output received (base64 encoded)"
+    echo "To save audio, use the Python client or decode the base64 data"
+fi
+
+# Check for errors
+error=$(echo "$output" | jq -r '.error // empty')
+if [[ -n "$error" ]]; then
+    echo "Error: $error"
+fi
+
+echo ""
+echo "Done!"

tests/e2e/offline_inference/stage_configs/step_audio2_ci.yaml

@@ -0,0 +1,37 @@
+# Stage config for Step-Audio2 CI tests
+# Single Thinker stage (audio/text understanding → text). ASR/text-output
+# tests don't need Token2Wav; a two-stage config would deadlock on stage-1 skip.
+
+# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
+stage_args:
+  - stage_id: 0
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: StepAudio2ThinkerForConditionalGeneration
+      tensor_parallel_size: 1
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      max_model_len: 4096
+      max_num_batched_tokens: 4096
+      max_num_seqs: 1
+      gpu_memory_utilization: 0.8
+      skip_mm_profiling: true
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: text  # Output text + audio tokens
+    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: 256
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.05

tests/e2e/offline_inference/test_step_audio2.py

@@ -0,0 +1,303 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E tests for Step-Audio2 model with audio input and audio output.
+
+These tests verify:
+1. Processor loads correctly and handles audio input
+2. Field configuration for vLLM batching (flat_from_sizes)
+3. Tensor shapes are correct (3D not 4D for profiling)
+4. Basic E2E inference works with audio input
+"""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.runtime import OmniRunner
+
+models = ["stepfun-ai/Step-Audio-2-mini"]
+
+stage_configs = [str(Path(__file__).parent / "stage_configs" / "step_audio2_ci.yaml")]
+
+test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+
+
+def create_dummy_audio(sample_rate: int = 16000, duration_sec: float = 1.0) -> tuple[np.ndarray, int]:
+    """Create a dummy audio signal (sine wave) for testing."""
+    t = np.linspace(0, duration_sec, int(sample_rate * duration_sec), dtype=np.float32)
+    audio = 0.5 * np.sin(2 * np.pi * 440 * t)
+    return audio, sample_rate
+
+
+class StepAudio2OmniRunner(OmniRunner):
+    """Test runner specialized for Step-Audio2 model."""
+
+    def get_omni_inputs(
+        self,
+        prompts: list[str] | str,
+        system_prompt: str | None = None,
+        audios=None,
+        **kwargs,
+    ) -> list[dict]:
+        """Construct Omni input format for Step-Audio2."""
+        if system_prompt is None:
+            system_prompt = "你是一个语音对话助手，能够理解音频输入并生成语音回复。"
+
+        audio_placeholder = "<audio_patch>"
+
+        if isinstance(prompts, str):
+            prompts = [prompts]
+
+        num_prompts = len(prompts)
+
+        if audios is None:
+            audios_list = [None] * num_prompts
+        elif isinstance(audios, list):
+            audios_list = audios if len(audios) == num_prompts else [audios[0]] * num_prompts
+        else:
+            audios_list = [audios] * num_prompts
+
+        omni_inputs = []
+        for i, prompt_text in enumerate(prompts):
+            user_content = ""
+            multi_modal_data = {}
+
+            audio = audios_list[i]
+            if audio is not None:
+                user_content += audio_placeholder
+                multi_modal_data["audio"] = audio
+
+            user_content += prompt_text
+
+        full_prompt = (
+            f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
+            f"<|im_start|>user\n{user_content}<|im_end|>\n"
+            f"<|im_start|>assistant\n"
+        )
+
+        input_dict = {"prompt": full_prompt}
+        if multi_modal_data:
+            input_dict["multi_modal_data"] = multi_modal_data
+
+        omni_inputs.append(input_dict)
+
+        return omni_inputs
+
+
+@pytest.fixture(scope="session")
+def step_audio2_runner():
+    return StepAudio2OmniRunner
+
+
+@pytest.mark.core_model
+def test_processor_loads_correctly():
+    """Test that Step-Audio2 processor can be imported and initialized."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
+        StepAudio2MultiModalProcessor,
+        StepAudio2ProcessingInfo,
+        StepAudio2Processor,
+    )
+
+    assert StepAudio2MultiModalProcessor is not None
+    assert StepAudio2ProcessingInfo is not None
+    assert StepAudio2Processor is not None
+
+
+@pytest.mark.core_model
+def test_audio_preprocessing():
+    """Test that audio preprocessing produces correct tensor shapes."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
+        log_mel_spectrogram,
+        padding_mels,
+    )
+
+    audio, sample_rate = create_dummy_audio(sample_rate=16000, duration_sec=1.0)
+
+    mel = log_mel_spectrogram(audio)
+
+    assert mel.ndim == 2, f"Expected 2D tensor, got {mel.ndim}D"
+    assert mel.shape[0] == 128, f"Expected 128 mel bins, got {mel.shape[0]}"
+
+    mels = [mel, mel]
+    padded_mels, mel_lens = padding_mels(mels)
+
+    assert padded_mels.ndim == 3, f"Expected 3D tensor, got {padded_mels.ndim}D"
+    assert padded_mels.shape[0] == 2, f"Expected batch size 2, got {padded_mels.shape[0]}"
+    assert padded_mels.shape[1] == 128, f"Expected 128 mel bins, got {padded_mels.shape[1]}"
+
+
+@pytest.mark.core_model
+def test_feature_length_calculation():
+    """Test that audio feature length calculation is correct."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
+        calculate_audio_feature_length,
+    )
+
+    assert calculate_audio_feature_length(1000) == 125
+
+    assert calculate_audio_feature_length(100) > 0
+
+    assert calculate_audio_feature_length(1) >= 1
+
+
+@pytest.mark.core_model
+def test_token_config_constants():
+    """Test that token configuration constants are set correctly."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import (
+        DEFAULT_TOKEN_CONFIG,
+        STEP_AUDIO2_AUDIO_END,
+        STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID,
+        STEP_AUDIO2_AUDIO_START,
+        STEP_AUDIO2_AUDIO_VOCAB_SIZE,
+        STEP_AUDIO2_TEXT_MAX,
+    )
+
+    assert STEP_AUDIO2_TEXT_MAX < STEP_AUDIO2_AUDIO_START, "Text tokens should come before audio tokens"
+    assert STEP_AUDIO2_AUDIO_END == STEP_AUDIO2_AUDIO_START + STEP_AUDIO2_AUDIO_VOCAB_SIZE - 1
+    assert STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID > STEP_AUDIO2_TEXT_MAX
+    assert STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID < STEP_AUDIO2_AUDIO_START
+
+    assert DEFAULT_TOKEN_CONFIG.text_max == STEP_AUDIO2_TEXT_MAX
+    assert DEFAULT_TOKEN_CONFIG.audio_start == STEP_AUDIO2_AUDIO_START
+
+
+@pytest.mark.core_model
+def test_mm_field_config_structure():
+    """Test that multimodal field config is properly structured for vLLM batching."""
+    from vllm.multimodal.inputs import MultiModalFieldConfig
+
+    audio_lens = torch.tensor([100, 200, 150], dtype=torch.int32)
+
+    field_config = MultiModalFieldConfig.flat_from_sizes("audio", audio_lens)
+
+    assert field_config is not None
+
+
+@pytest.mark.core_model
+def test_audio_encoder_output_shape():
+    """Test that audio encoder produces correct output shapes."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
+        Adaptor,
+        AudioEncoder,
+    )
+
+    encoder = AudioEncoder(n_mels=128, n_ctx=1500, n_state=512, n_head=8, n_layer=6)
+
+    batch_size = 2
+    n_mels = 128
+    time_steps = 400
+    x = torch.randn(batch_size, n_mels, time_steps)
+    x_len = torch.tensor([time_steps, time_steps // 2], dtype=torch.int32)
+
+    encoded, encoded_lens = encoder(x, x_len)
+
+    assert encoded.ndim == 3, f"Expected 3D tensor, got {encoded.ndim}D"
+    assert encoded.shape[0] == batch_size
+    assert encoded.shape[2] == 512
+
+    adapter = Adaptor(n_state=512, n_hidden=4096, kernel_size=3, stride=2)
+    adapted = adapter(encoded)
+
+    assert adapted.ndim == 3
+    assert adapted.shape[0] == batch_size
+    assert adapted.shape[2] == 4096
+
+
+@pytest.mark.core_model
+def test_token_separation():
+    """Test that token separation works correctly."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
+        StepAudio2ThinkerForConditionalGeneration,
+    )
+
+    token_ids = [100, 200, 151700, 300, 151800, 400]
+
+    text_tokens, audio_tokens = StepAudio2ThinkerForConditionalGeneration.separate_tokens(token_ids)
+
+    assert text_tokens == [100, 200, 300, 400]
+    assert audio_tokens == [4, 104]
+
+
+@pytest.mark.core_model
+def test_has_audio_output():
+    """Test detection of audio tokens in output."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
+        StepAudio2ThinkerForConditionalGeneration,
+    )
+
+    text_only = [100, 200, 300]
+    assert not StepAudio2ThinkerForConditionalGeneration.has_audio_output(text_only)
+
+    with_audio = [100, 200, 151700, 300]
+    assert StepAudio2ThinkerForConditionalGeneration.has_audio_output(with_audio)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.parametrize("test_config", test_params)
+def test_audio_to_text_and_audio(step_audio2_runner: type[StepAudio2OmniRunner], test_config: tuple[str, str]) -> None:
+    """Test processing audio input and generating text + audio output."""
+    model, stage_config_path = test_config
+
+    with step_audio2_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
+        audio = create_dummy_audio(sample_rate=16000, duration_sec=2.0)
+
+        outputs = runner.generate_multimodal(
+            prompts="请复述这段音频的内容。",
+            audios=audio,
+        )
+
+        assert len(outputs) > 0
+
+        text_output = None
+        for stage_output in outputs:
+            if stage_output.final_output_type == "text":
+                text_output = stage_output
+                break
+
+        assert text_output is not None
+        assert text_output.request_output is not None
+        text_content = text_output.request_output.outputs[0].text
+        assert text_content is not None
+        assert len(text_content.strip()) > 0
+
+        audio_output = None
+        for stage_output in outputs:
+            if stage_output.final_output_type == "audio":
+                audio_output = stage_output
+                break
+
+        if audio_output is not None:
+            assert audio_output.request_output is not None
+            audio_tensor = audio_output.request_output.outputs[0].multimodal_output.get("audio")
+            if audio_tensor is not None:
+                assert audio_tensor.numel() > 0
+
+
+@pytest.mark.advanced_model
+@pytest.mark.parametrize("test_config", test_params)
+def test_text_only_input(step_audio2_runner: type[StepAudio2OmniRunner], test_config: tuple[str, str]) -> None:
+    """Test processing text-only input (no audio)."""
+    model, stage_config_path = test_config
+
+    with step_audio2_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
+        outputs = runner.generate_multimodal(
+            prompts="你好，请用中文回答。",
+            audios=None,
+        )
+
+        assert len(outputs) > 0
+
+        text_output = None
+        for stage_output in outputs:
+            if stage_output.final_output_type == "text":
+                text_output = stage_output
+                break
+
+        assert text_output is not None
+        assert text_output.request_output is not None
+        text_content = text_output.request_output.outputs[0].text
+        assert text_content is not None
+        assert len(text_content.strip()) > 0

tests/e2e/online_serving/stage_configs/step_audio2_ci.yaml

@@ -0,0 +1,81 @@
+# Stage config for Step-Audio2 CI tests
+# Stage 0: Thinker (audio understanding → text + audio tokens)
+# Stage 1: Token2Wav (audio tokens → waveform)
+
+# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
+stage_args:
+  - stage_id: 0
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: StepAudio2ThinkerForConditionalGeneration
+      tensor_parallel_size: 1
+      hf_overrides:
+          architectures: ["StepAudio2ThinkerForConditionalGeneration"]
+
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      max_model_len: 1024
+      max_num_batched_tokens: 1024
+      max_num_seqs: 1
+      gpu_memory_utilization: 0.7
+      skip_mm_profiling: true
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: text  # Output text + audio tokens
+    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: 256
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+
+      model_stage: token2wav
+      model_arch: StepAudio2Token2WavModel
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      hf_overrides:
+          architectures: ["StepAudio2Token2WavModel"]
+      gpu_memory_utilization: 0.2
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      max_num_batched_tokens: 1024
+      engine_output_type: audio  # Output waveform
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav
+    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
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1

tests/e2e/online_serving/test_step_audio2.py

@@ -0,0 +1,368 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+E2E tests for Step-Audio2 online serving with concurrent requests.
+
+Tests:
+1. Single request via API
+2. Concurrent throughput (1/2/4/8 concurrent requests)
+3. Latency percentiles (p50/p95/p99)
+"""
+
+import asyncio
+import os
+import socket
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import aiohttp
+import numpy as np
+import pytest
+from vllm.utils.network_utils import get_open_port
+
+MODEL = "stepfun-ai/Step-Audio-2-mini"
+STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "step_audio2_ci.yaml")
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+
+class OmniServer:
+    """Omniserver for vLLM-Omni tests."""
+
+    def __init__(
+        self,
+        model: str,
+        serve_args: list[str],
+        *,
+        env_dict: dict[str, str] | None = None,
+    ) -> None:
+        self.model = model
+        self.serve_args = serve_args
+        self.env_dict = env_dict
+        self.proc: subprocess.Popen | None = None
+        self.host = "127.0.0.1"
+        self.port = get_open_port()
+
+    def _start_server(self) -> None:
+        """Start the vLLM-Omni server subprocess."""
+        env = os.environ.copy()
+        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+        if self.env_dict is not None:
+            env.update(self.env_dict)
+
+        cmd = [
+            sys.executable,
+            "-m",
+            "vllm_omni.entrypoints.cli.main",
+            "serve",
+            self.model,
+            "--omni",
+            "--host",
+            self.host,
+            "--port",
+            str(self.port),
+        ] + self.serve_args
+
+        print(f"Launching OmniServer with: {' '.join(cmd)}")
+        self.proc = subprocess.Popen(
+            cmd,
+            env=env,
+            cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+        )
+
+        max_wait = 600  # 10 minutes
+        start_time = time.time()
+        while time.time() - start_time < max_wait:
+            try:
+                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+                    sock.settimeout(1)
+                    result = sock.connect_ex((self.host, self.port))
+                    if result == 0:
+                        print(f"Server ready on {self.host}:{self.port}")
+                        return
+            except Exception:
+                pass
+            time.sleep(2)
+
+        raise RuntimeError(f"Server failed to start within {max_wait} seconds")
+
+    def __enter__(self):
+        self._start_server()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.proc:
+            self.proc.terminate()
+            try:
+                self.proc.wait(timeout=30)
+            except subprocess.TimeoutExpired:
+                self.proc.kill()
+                self.proc.wait()
+
+
+def create_dummy_audio_base64(duration_sec: float = 5.0, sample_rate: int = 16000) -> str:
+    """Create dummy audio as base64."""
+    import base64
+    import io
+    import wave
+
+    num_samples = int(duration_sec * sample_rate)
+    t = np.linspace(0, duration_sec, num_samples, dtype=np.float32)
+    audio = (0.5 * np.sin(2 * np.pi * 440 * t) * 32767).astype(np.int16)
+
+    buffer = io.BytesIO()
+    with wave.open(buffer, "wb") as wav_file:
+        wav_file.setnchannels(1)
+        wav_file.setsampwidth(2)
+        wav_file.setframerate(sample_rate)
+        wav_file.writeframes(audio.tobytes())
+
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+def create_s2st_request(audio_base64: str) -> dict:
+    """Create S2ST request payload."""
+    # Use audio_url format with base64 data URL
+    audio_url = f"data:audio/wav;base64,{audio_base64}"
+    return {
+        "model": MODEL,
+        "messages": [
+            {"role": "system", "content": "请仔细聆听这段语音，然后复述其内容。"},
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "audio_url",
+                        "audio_url": {"url": audio_url},
+                    },
+                    {
+                        "type": "text",
+                        "text": "<tts_start>",
+                    },
+                ],
+            },
+        ],
+        "max_tokens": 512,
+        "temperature": 0.7,
+        "top_p": 0.9,
+    }
+
+
+async def send_request(session: aiohttp.ClientSession, url: str, payload: dict) -> tuple[float, int, str]:
+    """Send single request and return (latency, num_tokens, response_text)."""
+    start = time.perf_counter()
+    async with session.post(url, json=payload) as response:
+        result = await response.json()
+    latency = time.perf_counter() - start
+
+    if "choices" in result and len(result["choices"]) > 0:
+        text = result["choices"][0].get("message", {}).get("content", "")
+        tokens = result.get("usage", {}).get("completion_tokens", len(text.split()))
+    else:
+        text = str(result)
+        tokens = 0
+
+    return latency, tokens, text
+
+
+async def benchmark_concurrent(
+    url: str,
+    payload: dict,
+    num_concurrent: int,
+    num_requests: int,
+    warmup: int = 2,
+) -> dict:
+    """Run concurrent benchmark with true sustained concurrency.
+
+    Uses semaphore to maintain exactly num_concurrent requests in flight at all times,
+    rather than batch-and-wait approach. Measures real wall clock time.
+    """
+    semaphore = asyncio.Semaphore(num_concurrent)
+
+    async def bounded_request(session: aiohttp.ClientSession):
+        async with semaphore:
+            return await send_request(session, url, payload)
+
+    async with aiohttp.ClientSession() as session:
+        if warmup > 0:
+            print(f"    Warming up ({warmup} requests)...", end=" ", flush=True)
+            for _ in range(warmup):
+                try:
+                    await send_request(session, url, payload)
+                except Exception as e:
+                    print(f"Warmup failed: {e}")
+            print("done")
+
+        print(f"    Running {num_requests} requests with concurrency={num_concurrent}...", end=" ", flush=True)
+        wall_start = time.perf_counter()
+
+        tasks = [bounded_request(session) for _ in range(num_requests)]
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        wall_time = time.perf_counter() - wall_start
+        print(f"done ({wall_time:.2f}s)")
+
+    all_latencies = []
+    all_tokens = []
+    errors = 0
+
+    for result in results:
+        if isinstance(result, Exception):
+            print(f"    Request failed: {result}")
+            errors += 1
+            continue
+        latency, tokens, _ = result
+        all_latencies.append(latency)
+        all_tokens.append(tokens)
+
+    if not all_latencies:
+        return {"error": "All requests failed"}
+
+    latencies = np.array(all_latencies)
+    total_tokens = sum(all_tokens)
+    successful_requests = len(all_latencies)
+
+    return {
+        "num_concurrent": num_concurrent,
+        "num_requests": num_requests,
+        "successful_requests": successful_requests,
+        "errors": errors,
+        "wall_time": wall_time,
+        "p50_latency": float(np.percentile(latencies, 50)),
+        "p95_latency": float(np.percentile(latencies, 95)),
+        "p99_latency": float(np.percentile(latencies, 99)),
+        "mean_latency": float(np.mean(latencies)),
+        "min_latency": float(np.min(latencies)),
+        "max_latency": float(np.max(latencies)),
+        "throughput_req_per_sec": successful_requests / wall_time,
+        "throughput_tokens_per_sec": total_tokens / wall_time,
+    }
+
+
+class TestStepAudio2OnlineServing:
+    """Online serving tests for Step-Audio2."""
+
+    @pytest.fixture(scope="class")
+    def server(self):
+        """Start vLLM-Omni server."""
+        with OmniServer(MODEL, ["--stage-configs-path", STAGE_CONFIG]) as omni_server:
+            yield f"http://{omni_server.host}:{omni_server.port}/v1/chat/completions"
+
+    @pytest.mark.asyncio
+    async def test_single_request(self, server):
+        """Test single request works."""
+        audio_base64 = create_dummy_audio_base64(duration_sec=2.0)
+        payload = create_s2st_request(audio_base64)
+
+        async with aiohttp.ClientSession() as session:
+            latency, tokens, text = await send_request(session, server, payload)
+
+        assert latency > 0
+        assert tokens >= 0
+        print(f"Single request: latency={latency:.2f}s, tokens={tokens}")
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("num_concurrent", [1, 2, 4])
+    async def test_concurrent_throughput(self, server, num_concurrent):
+        """Test concurrent throughput at different concurrency levels."""
+        audio_base64 = create_dummy_audio_base64(duration_sec=5.0)
+        payload = create_s2st_request(audio_base64)
+
+        num_requests = num_concurrent * 3
+
+        stats = await benchmark_concurrent(server, payload, num_concurrent, num_requests)
+
+        assert "error" not in stats
+        print(f"\nConcurrency={num_concurrent}:")
+        print(f"  p50={stats['p50_latency']:.2f}s, p95={stats['p95_latency']:.2f}s")
+        print(f"  throughput={stats['throughput_req_per_sec']:.2f} req/s")
+
+
+async def run_throughput_benchmark(
+    server_url: str,
+    audio_duration: float = 5.0,
+    concurrency_levels: list[int] = [1, 2, 4, 8],
+    requests_per_level: int = 10,
+    warmup: int = 2,
+):
+    """Run full throughput benchmark."""
+    audio_base64 = create_dummy_audio_base64(duration_sec=audio_duration)
+    payload = create_s2st_request(audio_base64)
+
+    print(f"{'=' * 70}")
+    print("Step-Audio2 Online Serving Throughput Benchmark")
+    print(f"{'=' * 70}")
+    print(f"Server: {server_url}")
+    print(f"Audio duration: {audio_duration}s")
+    print(f"Requests per level: {requests_per_level}")
+    print(f"Warmup requests: {warmup}")
+    print(f"Concurrency levels: {concurrency_levels}")
+    print(f"{'=' * 70}")
+
+    results = []
+    for num_concurrent in concurrency_levels:
+        print(f"\n[Concurrency={num_concurrent}]")
+        stats = await benchmark_concurrent(server_url, payload, num_concurrent, requests_per_level, warmup)
+        results.append(stats)
+
+        if "error" not in stats:
+            print(
+                f"    Result: p50={stats['p50_latency']:.2f}s, "
+                f"p95={stats['p95_latency']:.2f}s, "
+                f"throughput={stats['throughput_req_per_sec']:.2f} req/s, "
+                f"{stats['throughput_tokens_per_sec']:.1f} tok/s"
+            )
+
+    print(f"\n{'=' * 80}")
+    print("SUMMARY")
+    print(f"{'=' * 80}")
+    print(
+        f"{'Conc':<6} {'Reqs':<6} {'Err':<5} {'Wall(s)':<9} {'p50':<8} {'p95':<8} {'p99':<8} {'Req/s':<8} {'Tok/s':<8}"
+    )
+    print("-" * 80)
+
+    for stats in results:
+        if "error" in stats:
+            print(f"{stats.get('num_concurrent', '?'):<6} ERROR: {stats['error']}")
+            continue
+        print(
+            f"{stats['num_concurrent']:<6} "
+            f"{stats['successful_requests']:<6} "
+            f"{stats['errors']:<5} "
+            f"{stats['wall_time']:<9.2f} "
+            f"{stats['p50_latency']:<8.2f} "
+            f"{stats['p95_latency']:<8.2f} "
+            f"{stats['p99_latency']:<8.2f} "
+            f"{stats['throughput_req_per_sec']:<8.2f} "
+            f"{stats['throughput_tokens_per_sec']:<8.1f}"
+        )
+
+    print("=" * 80)
+    return results
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Step-Audio2 Online Serving Benchmark")
+    parser.add_argument(
+        "--server-url", type=str, default="http://localhost:8000/v1/chat/completions", help="vLLM server URL"
+    )
+    parser.add_argument("--audio-duration", type=float, default=5.0, help="Input audio duration in seconds")
+    parser.add_argument("--concurrency", type=str, default="1,2,4,8", help="Comma-separated concurrency levels to test")
+    parser.add_argument("--requests-per-level", type=int, default=10, help="Number of requests per concurrency level")
+    parser.add_argument("--warmup", type=int, default=2, help="Number of warmup requests before each benchmark")
+
+    args = parser.parse_args()
+    concurrency_levels = [int(x) for x in args.concurrency.split(",")]
+
+    asyncio.run(
+        run_throughput_benchmark(
+            args.server_url,
+            args.audio_duration,
+            concurrency_levels,
+            args.requests_per_level,
+            args.warmup,
+        )
+    )

tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+# Token2Wav imports the optional `step-audio2` extras (flashcosyvoice / s3tokenizer
+# / hyperpyyaml) at module load. Skip instead of erroring at collection time on CI
+# images that do not install them (e.g. the CPU core_model runner); a genuine
+# import error (e.g. vLLM API drift) still surfaces. See
+# examples/offline_inference/step_audio2/README.md.
+try:
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_token2wav import (
+        StepAudio2Token2WavForConditionalGeneration,
+    )
+except ModuleNotFoundError as exc:
+    pytest.skip(
+        f"Step-Audio2 Token2Wav extras not installed ({exc.name})",
+        allow_module_level=True,
+    )
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummyModelConfig:
+    def __init__(self) -> None:
+        self.hf_config = SimpleNamespace(
+            token2wav_path="/tmp/fake-token2wav",
+            token2wav_float16=False,
+        )
+        self.model = "/tmp/fake-model"
+
+    def get_hidden_size(self) -> int:
+        return 8
+
+
+def _build_model() -> StepAudio2Token2WavForConditionalGeneration:
+    vllm_config = SimpleNamespace(
+        model_config=_DummyModelConfig(),
+        device_config=SimpleNamespace(device="cpu"),
+    )
+    return StepAudio2Token2WavForConditionalGeneration(vllm_config=vllm_config)
+
+
+def test_step_audio2_token2wav_sync_path_not_misdetected_by_empty_runtime_info(monkeypatch):
+    model = _build_model()
+
+    monkeypatch.setattr("os.path.exists", lambda _: True)
+    monkeypatch.setattr(
+        model.token2wav,
+        "forward",
+        lambda *args, **kwargs: torch.zeros(1, 16, dtype=torch.float32),
+    )
+
+    out = model.forward(
+        input_ids=torch.tensor([1, 2, 3]),
+        positions=torch.tensor([0, 1, 2]),
+        runtime_additional_information=[{}],
+    )
+
+    audio = out.multimodal_outputs["model_outputs"]
+    # sync path returns the full waveform wrapped in a single-element list,
+    # same contract as the async-chunk path (see tests below)
+    assert isinstance(audio, list)
+    assert len(audio) == 1
+    assert isinstance(audio[0], torch.Tensor)
+    assert model._stream_states == []
+
+
+def test_step_audio2_token2wav_async_chunk_batch_guard():
+    model = _build_model()
+
+    with pytest.raises(RuntimeError, match="only supports batch=1"):
+        model._forward_async_chunk(
+            input_ids=torch.tensor([1, 2, 3]),
+            runtime_additional_information=[
+                {"left_context_size": 0},
+                {"left_context_size": 1},
+            ],
+        )
+
+
+def test_step_audio2_token2wav_async_chunk_last_chunk_resets_state(monkeypatch):
+    model = _build_model()
+
+    monkeypatch.setattr("os.path.exists", lambda _: True)
+    calls = {"setup": 0, "stream": 0, "reset": 0}
+
+    def _setup_stream_for(prompt_wav, state):
+        calls["setup"] += 1
+        state.setup_done = True
+        state.stream_cache = {}
+        state.hift_cache_dict = {
+            "mel": torch.zeros(1, 1, 0),
+            "source": torch.zeros(1, 1, 0),
+            "speech": torch.zeros(1, 0),
+        }
+
+    def _stream_chunk_for(audio_tokens, prompt_wav, last_chunk, state):
+        calls["stream"] += 1
+        return torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32)
+
+    def _reset_stream_for(state):
+        calls["reset"] += 1
+        state.setup_done = False
+        state.finished = True
+        state.stream_cache = None
+        state.hift_cache_dict = {}
+
+    monkeypatch.setattr(model.token2wav, "setup_stream_for", _setup_stream_for)
+    monkeypatch.setattr(model.token2wav, "stream_chunk_for", _stream_chunk_for)
+    monkeypatch.setattr(model.token2wav, "reset_stream_for", _reset_stream_for)
+
+    out = model.forward(
+        input_ids=torch.tensor([10, 11, 12]),
+        positions=torch.tensor([0, 1, 2]),
+        runtime_additional_information=[{"left_context_size": 1}],
+    )
+
+    audio_list = out.multimodal_outputs["model_outputs"]
+    assert isinstance(audio_list, list)
+    assert len(audio_list) == 1
+    assert torch.allclose(audio_list[0], torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32))
+    assert calls == {"setup": 1, "stream": 1, "reset": 1}
+
+
+def test_step_audio2_token2wav_async_chunk_empty_eof_returns_zero_chunk():
+    model = _build_model()
+
+    out = model._forward_async_chunk(
+        input_ids=torch.tensor([], dtype=torch.int64),
+        runtime_additional_information=[{"left_context_size": 1}],
+    )
+
+    audio_list = out.multimodal_outputs["model_outputs"]
+    assert isinstance(audio_list, list)
+    assert len(audio_list) == 1
+    assert torch.equal(audio_list[0], torch.zeros(1, dtype=torch.float32))
+    assert model._stream_states
+    assert model._stream_states[0].finished is True

tests/model_executor/stage_input_processors/test_step_audio2_async_chunk.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import (
+    DEFAULT_STREAM_CONFIG,
+    DEFAULT_TOKEN_CONFIG,
+)
+from vllm_omni.model_executor.stage_input_processors.step_audio2 import (
+    thinker2token2wav_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _req(external_req_id: str, *, prompt_token_ids: list[int], all_token_ids: list[int], finished: bool):
+    return SimpleNamespace(
+        external_req_id=external_req_id,
+        prompt_token_ids=prompt_token_ids,
+        all_token_ids=all_token_ids,
+        is_finished=lambda: finished,
+    )
+
+
+def test_step_audio2_async_chunk_uses_decode_only_tokens_not_prompt_history():
+    audio_start = DEFAULT_TOKEN_CONFIG.audio_start
+    audio_eos = DEFAULT_TOKEN_CONFIG.audio_eos
+    transfer_manager = SimpleNamespace(code_prompt_token_ids=defaultdict(list))
+
+    prompt = [11, 12, audio_start + 5]  # historical prompt audio token
+    generated = [99, audio_start + 1, audio_start + 2, audio_start + audio_eos]
+    request = _req(
+        "rid-decode-only",
+        prompt_token_ids=prompt,
+        all_token_ids=prompt + generated,
+        finished=True,
+    )
+
+    payload = thinker2token2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+    )
+
+    assert payload is not None
+    assert payload["code_predictor_codes"] == [1, 2]
+    assert payload["left_context_size"] == 1
+    assert payload["finished"].item() is True
+
+
+def test_step_audio2_async_chunk_returns_none_when_not_enough_tokens():
+    audio_start = DEFAULT_TOKEN_CONFIG.audio_start
+    required = DEFAULT_STREAM_CONFIG.chunk_size + DEFAULT_STREAM_CONFIG.pre_lookahead_len
+    transfer_manager = SimpleNamespace(code_prompt_token_ids=defaultdict(list))
+
+    generated = [audio_start + i for i in range(required - 1)]
+    request = _req(
+        "rid-not-ready",
+        prompt_token_ids=[1, 2, 3],
+        all_token_ids=[1, 2, 3] + generated,
+        finished=False,
+    )
+
+    payload = thinker2token2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+    )
+
+    assert payload is None
+
+
+def test_step_audio2_async_chunk_emits_non_last_chunk_and_advances_consumed_by_chunk_size():
+    audio_start = DEFAULT_TOKEN_CONFIG.audio_start
+    chunk_size = DEFAULT_STREAM_CONFIG.chunk_size
+    required = chunk_size + DEFAULT_STREAM_CONFIG.pre_lookahead_len
+    transfer_manager = SimpleNamespace(code_prompt_token_ids=defaultdict(list))
+
+    generated = [audio_start + i for i in range(required + 10)]
+    request = _req(
+        "rid-ready",
+        prompt_token_ids=[7, 8],
+        all_token_ids=[7, 8] + generated,
+        finished=False,
+    )
+
+    payload = thinker2token2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+    )
+
+    assert payload is not None
+    assert payload["left_context_size"] == 0
+    assert payload["finished"].item() is False
+    assert payload["code_predictor_codes"] == list(range(required))
+    assert transfer_manager.code_prompt_token_ids["rid-ready"] == list(range(chunk_size))
+
+
+def test_step_audio2_async_chunk_emits_eof_when_finished_with_no_remaining_audio():
+    transfer_manager = SimpleNamespace(code_prompt_token_ids=defaultdict(list))
+    transfer_manager.code_prompt_token_ids["rid-eof"] = [1, 2, 3]
+
+    request = _req(
+        "rid-eof",
+        prompt_token_ids=[10, 11],
+        all_token_ids=[
+            10,
+            11,
+            DEFAULT_TOKEN_CONFIG.audio_start + 1,
+            DEFAULT_TOKEN_CONFIG.audio_start + 2,
+            DEFAULT_TOKEN_CONFIG.audio_start + 3,
+        ],
+        finished=True,
+    )
+
+    payload = thinker2token2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+    )
+
+    assert payload == {
+        "code_predictor_codes": [],
+        "left_context_size": 1,
+        "finished": torch.tensor(True, dtype=torch.bool),
+    }

vllm_omni/entrypoints/openai/serving_speech.py

@@ -80,6 +80,7 @@ _MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"}
 _HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"}
 _HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"}
 _GLM_TTS_MODEL_STAGES = {"glm_tts"}
+_STEP_AUDIO2_TTS_MODEL_STAGES = {"step_audio2_thinker"}
 _TTS_MODEL_STAGES: set[str] = (
     _VOXTRAL_TTS_MODEL_STAGES
     | _QWEN3_TTS_MODEL_STAGES
@@ -94,6 +95,7 @@ _TTS_MODEL_STAGES: set[str] = (
     | _MOSS_TTS_MODEL_STAGES
     | _MOSS_TTS_FULL_MODEL_STAGES
     | _GLM_TTS_MODEL_STAGES
+    | _STEP_AUDIO2_TTS_MODEL_STAGES
 )
 _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
     "fish_tts",
@@ -655,6 +657,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return "higgs_audio_v3"
         if model_stage in _GLM_TTS_MODEL_STAGES:
             return "glm_tts"
+        if model_stage in _STEP_AUDIO2_TTS_MODEL_STAGES:
+            return "step_audio2"
         return None
 
     def _get_custom_voice_dir(self) -> str | None:
@@ -1356,6 +1360,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
     def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
         """Validate TTS request parameters. Returns error message or None."""
+        if self._tts_model_type == "step_audio2":
+            # StepAudio2 only requires non-empty input text
+            if not request.input or not request.input.strip():
+                return "Input text cannot be empty"
+            return None
         if self._tts_model_type == "voxtral_tts":
             return self._validate_voxtral_tts_request(request)
         if self._tts_model_type == "fish_tts":
@@ -2858,6 +2867,33 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 "multi_modal_data": {"audio": [(audio.audio_array, audio.sampling_rate)]},
             }
 
+    # ---- Step-Audio2 helpers ----
+
+    def _build_step_audio2_prompt(
+        self,
+        request: OpenAICreateSpeechRequest,
+    ) -> dict[str, Any]:
+        """Build prompt for Step-Audio2 TTS.
+
+        Constructs the chat prompt with ``<tts_start>`` as the last token
+        of the assistant turn (without ``<|im_end|>``), so the thinker
+        continues generating audio tokens.
+
+        Prompt format::
+            <|im_start|>system\\n{system_prompt}<|im_end|>\\n
+            <|im_start|>user\\n{input_text}<|im_end|>\\n
+            <|im_start|>assistant\\n<tts_start>
+        """
+        system_prompt = getattr(request, "instructions", None) or "You are a voice assistant. Read the text aloud."
+        text = request.input
+
+        raw_prompt = (
+            f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
+            f"<|im_start|>user\n{text}<|im_end|>\n"
+            f"<|im_start|>assistant\n<tts_start>"
+        )
+        return {"prompt": raw_prompt}
+
     # ---- Fish Speech helpers ----
 
     def _build_fish_speech_prompt(
@@ -3202,7 +3238,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             if err:
                 raise ValueError(err)
 
-        if self._is_fish_speech:
+        if self._tts_model_type == "step_audio2":
+            if not request.input or not request.input.strip():
+                raise ValueError("Input text cannot be empty")
+            prompt = self._build_step_audio2_prompt(request)
+            tts_params = {}
+        elif self._is_fish_speech:
             validation_error = self._validate_fish_tts_request(request)
             if validation_error:
                 raise ValueError(validation_error)
@@ -3413,7 +3454,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             tts_params = {}
             prompt = {"prompt": request.input}
 
-        if self._is_fish_speech:
+        if self._tts_model_type == "step_audio2":
+            model_type = "step_audio2"
+        elif self._is_fish_speech:
             model_type = "fish_speech"
         elif self._tts_model_type == "covo_audio":
             model_type = "covo_audio"

vllm_omni/model_executor/models/registry.py

@@ -53,6 +53,32 @@ _OMNI_MODELS = {
         "qwen3_omni_code2wav",
         "Qwen3OmniMoeCode2Wav",
     ),
+    # Step-Audio2 models
+    "StepAudio2ForCausalLM": (
+        "step_audio2",
+        "step_audio2",
+        "StepAudio2ForConditionalGeneration",
+    ),
+    "StepAudio2ForConditionalGeneration": (
+        "step_audio2",
+        "step_audio2",
+        "StepAudio2ForConditionalGeneration",
+    ),
+    "StepAudio2ThinkerForConditionalGeneration": (
+        "step_audio2",
+        "step_audio2_thinker",
+        "StepAudio2ThinkerForConditionalGeneration",
+    ),
+    "StepAudio2Token2WavModel": (
+        "step_audio2",
+        "step_audio2_token2wav",
+        "StepAudio2Token2WavForConditionalGeneration",
+    ),
+    "StepAudio2Token2WavForConditionalGeneration": (
+        "step_audio2",
+        "step_audio2_token2wav",
+        "StepAudio2Token2WavForConditionalGeneration",
+    ),
     "CosyVoice3Model": (
         "cosyvoice3",
         "cosyvoice3",

vllm_omni/model_executor/models/step_audio2/__init__.py

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

vllm_omni/model_executor/models/step_audio2/step_audio2.py

@@ -0,0 +1,195 @@
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.sequence import IntermediateTensors
+
+from .step_audio2_thinker import (
+    StepAudio2DummyInputsBuilder,
+    StepAudio2MultiModalProcessor,
+    StepAudio2ProcessingInfo,
+)
+
+logger = init_logger(__name__)
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    StepAudio2MultiModalProcessor,
+    info=StepAudio2ProcessingInfo,
+    dummy_inputs=StepAudio2DummyInputsBuilder,
+)
+class StepAudio2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
+    """
+    Step-Audio2 Main Controller
+
+    Manages two-stage inference pipeline:
+    - Stage 1 (Thinker): Audio understanding and token generation
+    - Stage 2 (Token2Wav): Audio token to waveform synthesis
+
+    Usage:
+        # Stage 1: Thinker
+        model = StepAudio2ForConditionalGeneration(
+            vllm_config=config,
+            model_stage="thinker"
+        )
+
+        # Stage 2: Token2Wav
+        model = StepAudio2ForConditionalGeneration(
+            vllm_config=config,
+            model_stage="token2wav"
+        )
+    """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.have_multimodal_outputs = True
+
+        config = vllm_config.model_config.hf_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+
+        self.config = config
+        self.multimodal_config = multimodal_config
+        self.vllm_config = vllm_config
+
+        raw_model_stage = vllm_config.model_config.model_stage
+        # Normalise: "step_audio2_thinker" → "thinker"
+        self.model_stage = "thinker" if raw_model_stage in ("thinker", "step_audio2_thinker") else raw_model_stage
+
+        if self.model_stage == "thinker":
+            self.thinker = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                prefix=maybe_prefix(prefix, "thinker"),
+                hf_config=config,
+                architectures=["StepAudio2ThinkerForConditionalGeneration"],
+            )
+            self.model = self.thinker
+            self.token2wav = None
+
+            logger.info("Initialized Step-Audio2 Thinker (Stage 0)")
+
+        elif self.model_stage == "token2wav":
+            self.thinker = None
+            self.token2wav = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                prefix=maybe_prefix(prefix, "token2wav"),
+                hf_config=config,
+                architectures=["StepAudio2Token2WavForConditionalGeneration"],
+            )
+            self.model = self.token2wav
+
+            logger.info("Initialized Step-Audio2 Token2Wav (Stage 1)")
+
+        else:
+            raise ValueError(
+                f"Invalid model_stage: {self.model_stage}. Must be 'thinker'/'step_audio2_thinker' or 'token2wav'"
+            )
+
+        self.make_empty_intermediate_tensors = (
+            self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None
+        )
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        """Get placeholder string for a modality
+
+        Returns:
+            For audio: "<audio_patch>" (matches processor's audio_token)
+        """
+        if modality == "audio":
+            return "<audio_patch>"
+        return None
+
+    def get_language_model(self) -> nn.Module:
+        """Get the underlying language model."""
+        if self.model_stage == "thinker":
+            return self.thinker.get_language_model()
+        # Token2Wav is not a real language model. Return the wrapper itself so
+        # vLLM interface helpers resolve embed_input_ids() on the adapter layer.
+        return self.token2wav
+
+    def get_multimodal_embeddings(self, **kwargs):
+        """Get multimodal embeddings - only used in Thinker stage."""
+        if self.model_stage == "thinker":
+            return self.thinker.get_multimodal_embeddings(**kwargs)
+        return None
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """Compatibility helper used by older call sites."""
+        if self.model_stage == "thinker":
+            return self.model.get_input_embeddings(input_ids)
+        return self.model.embed_input_ids(input_ids)
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        is_multimodal=None,
+    ) -> torch.Tensor:
+        """Explicit vLLM embedding hook for both stages."""
+        if self.model_stage == "token2wav":
+            return self.model.embed_input_ids(input_ids)
+        return self.model.get_input_embeddings(input_ids, multimodal_embeddings)
+
+    def embed_multimodal(self, **kwargs):
+        """Delegate multimodal embedding to thinker stage only."""
+        if self.model_stage == "thinker":
+            return self.model.embed_multimodal(**kwargs)
+        return None
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs,
+    ):
+        """
+        Forward pass through the model
+
+        For Thinker:
+            Returns hidden states/logits
+        For Token2Wav:
+            Returns waveform
+        """
+        return self.model.forward(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        """Compute logits from hidden states"""
+        return self.model.compute_logits(hidden_states)
+
+    def load_weights(self, weights):
+        """Load weights"""
+        return self.model.load_weights(weights)
+
+    def move_submodules_to_devices(
+        self,
+        *,
+        thinker_device: str | torch.device | None = None,
+        token2wav_device: str | torch.device | None = None,
+    ) -> None:
+        """
+        Optionally move thinker/token2wav to different devices
+
+        Example:
+            model.move_submodules_to_devices(
+                thinker_device='cuda:0',
+                token2wav_device='cuda:1',
+            )
+        """
+        if thinker_device is not None and self.thinker is not None:
+            self.thinker.to(torch.device(thinker_device))
+            logger.info(f"Moved Thinker to {thinker_device}")
+
+        if token2wav_device is not None and self.token2wav is not None:
+            self.token2wav.to(torch.device(token2wav_device))
+            logger.info(f"Moved Token2Wav to {token2wav_device}")

vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py

@@ -0,0 +1,181 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Step-Audio2 configuration constants - Single Source of Truth."""
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class StepAudio2TokenConfig:
+    """
+    Step-Audio2 token ID ranges and vocabulary configuration.
+
+    Token ID Layout (based on vocab_size=158720):
+    - Text tokens: 0 - 151688 (standard Qwen tokenizer)
+    - Special tokens: 151689 - 151695 (reserved)
+    - Audio tokens: 151696 - 158257 (vocab size 6562)
+
+    These values are architecture constants for Step-Audio2 and do not
+    change between model variants (mini/7B).
+    """
+
+    # Text token range
+    text_max: int = 151688
+    """Maximum text token ID (inclusive)"""
+
+    # Audio token range (absolute IDs in full vocabulary)
+    audio_start: int = 151696
+    """First audio token ID in absolute vocabulary"""
+
+    audio_vocab_size: int = 6562
+    """Total number of audio tokens"""
+
+    audio_eos: int = 6561
+    """Audio EOS token ID (relative to audio_start, used for padding)"""
+
+    # Special tokens
+    audio_patch_token_id: int = 151690
+    """<audio_patch> placeholder token ID"""
+
+    @property
+    def audio_end(self) -> int:
+        """Last audio token ID in absolute vocabulary"""
+        return self.audio_start + self.audio_vocab_size - 1
+
+
+@dataclass(frozen=True)
+class StepAudio2EncoderConfig:
+    """
+    Step-Audio2 audio encoder architecture configuration.
+
+    These values define the encoder structure and should match your model weights.
+    """
+
+    # Mel spectrogram parameters
+    n_mels: int = 128
+    """Number of mel frequency bins"""
+
+    # Encoder architecture
+    n_audio_ctx: int = 1500
+    """Audio context length (max sequence length)"""
+
+    n_audio_state: int = 512
+    """Audio encoder hidden state dimension"""
+
+    n_audio_head: int = 8
+    """Number of attention heads in encoder"""
+
+    n_audio_layer: int = 6
+    """Number of encoder layers"""
+
+    # Adapter parameters
+    kernel_size: int = 3
+    """Convolution kernel size for adapter"""
+
+    adapter_stride: int = 2
+    """Stride for adapter convolution (downsampling factor)"""
+
+
+@dataclass(frozen=True)
+class StepAudio2Token2WavConfig:
+    """Token2Wav synthesis configuration."""
+
+    n_timesteps: int = 10
+    """Number of diffusion timesteps for flow model inference"""
+
+    sample_rate: int = 24000
+    """Output audio sample rate"""
+
+
+@dataclass(frozen=True)
+class StepAudio2StreamConfig:
+    """Streaming (chunked) inference configuration for Token2Wav."""
+
+    chunk_size: int = 25
+    """Number of audio tokens consumed per chunk."""
+
+    pre_lookahead_len: int = 3
+    """Conformer encoder lookahead tokens beyond the chunk boundary."""
+
+    up_rate: int = 2
+    """Token-to-mel upsampling rate inside the flow model."""
+
+    mel_cache_len: int = 8
+    """HiFT overlap mel frames (160 ms at 50 Hz mel rate)."""
+
+    n_timesteps: int = 10
+    """Flow model ODE solver steps."""
+
+    estimator_cache_keep: int = 100
+    """Attention-cache pruning window for the flow estimator."""
+
+
+@dataclass(frozen=True)
+class StepAudio2ModelConfig:
+    """
+    Step-Audio2 complete model configuration - Single Source of Truth.
+
+    All architecture constants are defined here. The stage YAML and model config.json
+    can override these values if needed, but this serves as the default.
+    """
+
+    # LLM configuration (Qwen2)
+    hidden_size: int = 4096
+    """LLM hidden dimension (Qwen2-7B: 4096, Qwen2-1.5B: 1536, Qwen2-0.5B: 896)"""
+
+    # Token configuration
+    text_max: int = 151688
+    audio_start: int = 151696
+    audio_vocab_size: int = 6562
+    audio_eos: int = 6561
+    audio_patch_token_id: int = 151690
+
+    # Encoder configuration
+    n_mels: int = 128
+    n_audio_ctx: int = 1500
+    n_audio_state: int = 512
+    n_audio_head: int = 8
+    n_audio_layer: int = 6
+    kernel_size: int = 3
+    adapter_stride: int = 2
+
+
+# Default configuration instances
+DEFAULT_TOKEN_CONFIG = StepAudio2TokenConfig()
+DEFAULT_ENCODER_CONFIG = StepAudio2EncoderConfig()
+DEFAULT_TOKEN2WAV_CONFIG = StepAudio2Token2WavConfig()
+DEFAULT_STREAM_CONFIG = StepAudio2StreamConfig()
+DEFAULT_MODEL_CONFIG = StepAudio2ModelConfig()
+
+STREAM_SOURCE_CACHE_LEN = DEFAULT_STREAM_CONFIG.mel_cache_len * 480  # 3840 samples
+
+# Export constants for backward compatibility
+STEP_AUDIO2_TEXT_MAX = DEFAULT_TOKEN_CONFIG.text_max
+STEP_AUDIO2_AUDIO_START = DEFAULT_TOKEN_CONFIG.audio_start
+STEP_AUDIO2_AUDIO_VOCAB_SIZE = DEFAULT_TOKEN_CONFIG.audio_vocab_size
+STEP_AUDIO2_AUDIO_EOS = DEFAULT_TOKEN_CONFIG.audio_eos
+STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID = DEFAULT_TOKEN_CONFIG.audio_patch_token_id
+STEP_AUDIO2_AUDIO_END = DEFAULT_TOKEN_CONFIG.audio_end
+
+# Default prompt wav path for Token2Wav (can be overridden by env var)
+STEP_AUDIO2_DEFAULT_PROMPT_WAV = "default_female.wav"
+
+
+__all__ = [
+    "StepAudio2TokenConfig",
+    "StepAudio2EncoderConfig",
+    "StepAudio2Token2WavConfig",
+    "StepAudio2StreamConfig",
+    "StepAudio2ModelConfig",
+    "DEFAULT_TOKEN_CONFIG",
+    "DEFAULT_ENCODER_CONFIG",
+    "DEFAULT_TOKEN2WAV_CONFIG",
+    "DEFAULT_STREAM_CONFIG",
+    "DEFAULT_MODEL_CONFIG",
+    "STREAM_SOURCE_CACHE_LEN",
+    "STEP_AUDIO2_TEXT_MAX",
+    "STEP_AUDIO2_AUDIO_START",
+    "STEP_AUDIO2_AUDIO_VOCAB_SIZE",
+    "STEP_AUDIO2_AUDIO_EOS",
+    "STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID",
+    "STEP_AUDIO2_AUDIO_END",
+]

vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py

@@ -0,0 +1,898 @@
+"""
+Step-Audio2 Thinker - Stage 1 LLM for Audio Understanding
+
+This file contains:
+- Audio encoder components (AudioEncoder, Adaptor)
+- Audio preprocessing utilities (mel-spectrogram extraction)
+- Multi-modal processing (StepAudio2MultiModalProcessor)
+- Main thinker model (StepAudio2ThinkerForConditionalGeneration)
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, TypedDict
+
+import librosa  # noqa: TID251
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
+from transformers import AutoTokenizer
+from transformers.processing_utils import ProcessorMixin
+from vllm.config import VllmConfig
+from vllm.inputs import MultiModalDataDict
+from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.utils import (
+    _merge_multimodal_embeddings,
+    flatten_bn,
+    init_vllm_registered_model,
+    maybe_prefix,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import (
+    MultiModalFieldConfig,
+    MultiModalKwargsItems,
+    NestedTensors,
+)
+from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    PromptReplacement,
+    PromptUpdate,
+)
+from vllm.sequence import IntermediateTensors
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import (
+    DEFAULT_MODEL_CONFIG,
+    DEFAULT_TOKEN_CONFIG,
+    STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID,
+)
+
+if TYPE_CHECKING:
+    pass
+
+
+# ============================================================================
+# Audio Preprocessing Utilities
+# ============================================================================
+
+
+def _mel_filters(n_mels: int) -> torch.Tensor:
+    """Generate mel filter banks"""
+    assert n_mels in {80, 128}, f"Unsupported n_mels: {n_mels}"
+    return torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=n_mels))
+
+
+def _normalize_audio(audio: Any) -> torch.Tensor:
+    """Normalize audio to torch.Tensor"""
+    if isinstance(audio, tuple):
+        audio = audio[0]
+    if isinstance(audio, np.ndarray):
+        audio = torch.from_numpy(audio)
+    elif not torch.is_tensor(audio):
+        audio = torch.tensor(audio, dtype=torch.float32)
+    return audio.float()
+
+
+def log_mel_spectrogram(audio: Any, n_mels: int = 128, padding: int = 479) -> torch.Tensor:
+    """
+    Convert audio to log mel-spectrogram
+
+    Args:
+        audio: Audio waveform (any format)
+        n_mels: Number of mel frequency bins (default: 128)
+        padding: Right padding for STFT (default: 479)
+
+    Returns:
+        Log mel-spectrogram tensor of shape (n_mels, T)
+    """
+    audio = _normalize_audio(audio)
+    if padding > 0:
+        audio = F.pad(audio, (0, padding))
+    window = torch.hann_window(400, device=audio.device)
+    stft = torch.stft(audio, 400, 160, window=window, return_complex=True)
+    magnitudes = stft[..., :-1].abs() ** 2
+    filters = _mel_filters(n_mels).to(audio.device)
+    mel_spec = filters @ magnitudes
+
+    log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+    log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
+    log_spec = (log_spec + 4.0) / 4.0
+    return log_spec
+
+
+def padding_mels(data: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Pad mel spectrograms to same length
+
+    Args:
+        data: List of mel spectrograms, each of shape (n_mels, T)
+
+    Returns:
+        padded_feats: Padded tensor of shape (batch, n_mels, max_T)
+        feats_lengths: Lengths of each sequence
+    """
+    feats_lengths = torch.tensor([s.size(1) - 2 for s in data], dtype=torch.int32)
+    feats = [s.t() for s in data]
+    padded_feats = pad_sequence(feats, batch_first=True, padding_value=0)
+    return padded_feats.transpose(1, 2), feats_lengths
+
+
+# ============================================================================
+# Encoder Utilities
+# ============================================================================
+
+
+def make_non_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
+    """Create non-padding mask from lengths"""
+    batch_size = lengths.size(0)
+    max_len = max_len if max_len > 0 else lengths.max().item()
+    seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device)
+    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
+    seq_length_expand = lengths.unsqueeze(-1)
+    mask = seq_range_expand >= seq_length_expand
+    return ~mask
+
+
+def mask_to_bias(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Convert mask to attention bias"""
+    assert mask.dtype == torch.bool
+    assert dtype in [torch.float32, torch.bfloat16, torch.float16]
+    mask = mask.to(dtype)
+    mask = (1.0 - mask) * -1.0e10
+    return mask
+
+
+# ============================================================================
+# Audio Encoder Components
+# ============================================================================
+
+
+class MultiHeadAttention(nn.Module):
+    """Multi-head attention for audio encoder"""
+
+    def __init__(self, n_state: int, n_head: int):
+        super().__init__()
+        self.n_head = n_head
+        self.query = nn.Linear(n_state, n_state)
+        self.key = nn.Linear(n_state, n_state, bias=False)
+        self.value = nn.Linear(n_state, n_state)
+        self.out = nn.Linear(n_state, n_state)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor | None = None,
+    ):
+        q = self.query(x)
+        k = self.key(x)
+        v = self.value(x)
+
+        wv, qk = self.qkv_attention(q, k, v, mask)
+        return self.out(wv), qk
+
+    def qkv_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None):
+        B, T, D = q.shape
+        head_dim = D // self.n_head
+
+        # Reshape: (B, T, D) -> (B, n_head, T, head_dim)
+        q = q.view(B, T, self.n_head, head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_head, head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_head, head_dim).transpose(1, 2)
+
+        # Use F.scaled_dot_product_attention
+        # mask is additive bias (0 or -1e10), pass directly as attn_mask
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+
+        # Reshape back: (B, n_head, T, head_dim) -> (B, T, D)
+        out = out.transpose(1, 2).reshape(B, T, D)
+        return out, None
+
+
+class ResidualAttentionBlock(nn.Module):
+    """Residual attention block for audio encoder"""
+
+    def __init__(self, n_state: int, n_head: int):
+        super().__init__()
+
+        self.attn = MultiHeadAttention(n_state, n_head)
+        self.attn_ln = nn.LayerNorm(n_state)
+
+        n_mlp = n_state * 4
+        self.mlp = nn.Sequential(nn.Linear(n_state, n_mlp), nn.GELU(), nn.Linear(n_mlp, n_state))
+        self.mlp_ln = nn.LayerNorm(n_state)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor | None = None,
+    ):
+        x = x + self.attn(self.attn_ln(x), mask=mask)[0]
+        x = x + self.mlp(self.mlp_ln(x))
+        return x
+
+
+class AudioEncoder(nn.Module):
+    """
+    Step-Audio2 Audio Encoder
+
+    Lightweight audio encoder (6 layers, 512 hidden)
+    Optimized for 25s audio chunks at 25Hz
+    """
+
+    def __init__(self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int):
+        super().__init__()
+        self.conv1 = nn.Conv1d(n_mels, n_state, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1)
+        self.positional_embedding = nn.Embedding(n_ctx, n_state)
+
+        self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList(
+            [ResidualAttentionBlock(n_state, n_head) for _ in range(n_layer)]
+        )
+        self.avg_pooler = nn.AvgPool1d(2, stride=2)
+        self.after_norm = nn.LayerNorm(n_state)
+
+    def forward(self, x: torch.Tensor, x_len: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            x: mel spectrogram, shape (batch_size, n_mels, T)
+            x_len: length of each audio, shape (batch_size,)
+
+        Returns:
+            x: encoded features, shape (batch_size, T//4, n_state)
+            x_len: updated lengths, shape (batch_size,)
+        """
+        T = x.size(-1)
+        x = F.gelu(self.conv1(x))
+        x = F.gelu(self.conv2(x))
+        x = x.permute(0, 2, 1)  # (B, T // 2, n_state)
+
+        # Create attention mask
+        mask = make_non_pad_mask(x_len, T).unsqueeze(1)  # (B, 1, T)
+        mask = mask_to_bias(mask[:, :, (T + 1) % 2 :: 2], x.dtype)  # (B, 1, T // 2)
+
+        # Add positional embedding
+        x = (x + self.positional_embedding.weight[: x.shape[1], :]).to(x.dtype).contiguous()
+
+        # Apply attention blocks
+        for block in self.blocks:
+            x = block(x, mask.unsqueeze(1))
+
+        # Pool and normalize
+        x = x.permute(0, 2, 1)
+        x = self.avg_pooler(x)
+        x = x.permute(0, 2, 1)
+        x_len = (x_len + 1) // 2 // 2
+        x = self.after_norm(x)
+
+        return x, x_len
+
+
+class Adaptor(nn.Module):
+    """
+    Adaptor to project audio features to LLM dimension
+
+    Maps from n_state (512) to n_hidden (LLM dimension, e.g., 4096)
+    with optional downsampling via convolution
+    """
+
+    def __init__(
+        self, n_state: int = 512, n_hidden: int = 4096, kernel_size: int = 3, stride: int = 2, adapter_state: int = 2048
+    ):
+        super().__init__()
+        self.stride = stride
+
+        if self.stride != -1:
+            self.conv = nn.Conv1d(n_state, n_state, kernel_size, stride, padding=1)
+
+        self.linear1 = nn.Linear(n_state, adapter_state)
+        self.relu = nn.ReLU()
+        self.linear2 = nn.Linear(adapter_state, n_hidden)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: audio features, shape (batch_size, T, n_state)
+
+        Returns:
+            x: projected features, shape (batch_size, T//stride, n_hidden)
+        """
+        if self.stride != -1:
+            x = x.permute(0, 2, 1)  # (B, n_state, T)
+            x = F.gelu(self.conv(x))
+            x = x.permute(0, 2, 1)  # (B, T//stride, n_state)
+
+        x = self.linear1(x)
+        x = self.relu(x)
+        x = self.linear2(x)
+
+        return x
+
+
+# ============================================================================
+# Feature Length Calculation Utility
+# ============================================================================
+
+
+def calculate_audio_feature_length(mel_length: int) -> int:
+    """
+    Calculate output feature length after encoder + adapter processing.
+
+    Processing chain:
+    1. Mel spectrogram input: length = mel_length
+    2. Encoder (AudioEncoder):
+       - conv1: kernel=3, padding=1, stride=1 → no change
+       - conv2: kernel=3, padding=1, stride=2 → T // 2
+       - avg_pool: kernel=2, stride=2 → T // 4
+       - Output length: (mel_length + 1) // 2 // 2 ≈ mel_length // 4
+    3. Adapter (Adaptor):
+       - conv: kernel=3, padding=1, stride=2 → T // 2
+       - Output length: (encoder_output - 1) // 2 + 1
+
+    Args:
+        mel_length: Length of mel spectrogram (time steps)
+
+    Returns:
+        Final feature length after encoder + adapter
+
+    Example:
+        >>> calculate_audio_feature_length(1000)
+        125
+    """
+    encoder_output_len = (mel_length + 1) // 4
+
+    adapter_output_len = (encoder_output_len - 1) // 2 + 1
+
+    return max(1, adapter_output_len)
+
+
+class StepAudio2Processor(ProcessorMixin):
+    """Processor for Step-Audio2 that handles text tokenization and audio mel-spectrograms."""
+
+    attributes = ["tokenizer"]
+    attribute_class = {"tokenizer": "PreTrainedTokenizerBase"}
+
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+        self.audio_token = "<audio_patch>"
+
+    def __call__(self, text=None, audio=None, **kwargs):
+        """Process text and audio inputs."""
+        if text is None:
+            text = ""
+
+        allowed = {
+            "padding",
+            "truncation",
+            "max_length",
+            "return_tensors",
+            "add_special_tokens",
+            "pad_to_multiple_of",
+            "stride",
+            "return_attention_mask",
+            "return_token_type_ids",
+            "return_overflowing_tokens",
+            "return_offsets_mapping",
+            "return_special_tokens_mask",
+            "verbose",
+            "is_split_into_words",
+        }
+        tok_kwargs = {k: v for k, v in kwargs.items() if k in allowed}
+        if "return_tensors" not in tok_kwargs:
+            tok_kwargs["return_tensors"] = "pt"
+
+        encoded = self.tokenizer(text, **tok_kwargs)
+
+        if audio is None:
+            encoded["audio_mels"] = torch.empty((0, 128, 0))
+            encoded["audio_lens"] = torch.tensor([], dtype=torch.int32)
+            return encoded
+
+        audio_list = list(audio) if isinstance(audio, (list, tuple)) else [audio]
+
+        mels = [log_mel_spectrogram(a) for a in audio_list]
+        audio_mels, audio_lens = padding_mels(mels)
+
+        encoded["audio_mels"] = audio_mels
+        encoded["audio_lens"] = audio_lens
+        return encoded
+
+
+class StepAudio2AudioInputs(TypedDict):
+    """Audio inputs for Step-Audio2"""
+
+    audio_mels: torch.Tensor
+    """Shape: (batch, num_mel_bins, time_steps)"""
+
+    audio_lens: list[int]
+    """Shape: (batch,) - length of each audio in time steps"""
+
+
+class StepAudio2ProcessingInfo(BaseProcessingInfo):
+    """Processing info for Step-Audio2"""
+
+    def get_tokenizer(self, **kwargs):
+        return AutoTokenizer.from_pretrained(
+            self.ctx.model_config.model,
+            trust_remote_code=True,
+            fix_mistral_regex=True,
+            **kwargs,
+        )
+
+    def get_hf_processor(self, **kwargs: object):
+        return StepAudio2Processor(self.get_tokenizer())
+
+    def build_data_parser(self) -> MultiModalDataParser:
+        return MultiModalDataParser(target_sr=16000)
+
+    # Backward/branch compatibility:
+    # some code paths still call get_data_parser() on ProcessingInfo.
+    def get_data_parser(self) -> MultiModalDataParser:
+        return MultiModalDataParser(target_sr=16000)
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"audio": None}
+
+    def get_mm_max_tokens_per_item(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> Mapping[str, int]:
+        max_audio_tokens = 250
+        return {"audio": max_audio_tokens}
+
+
+class StepAudio2DummyInputsBuilder(BaseDummyInputsBuilder[StepAudio2ProcessingInfo]):
+    """Dummy inputs builder for Step-Audio2"""
+
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        num_audios = mm_counts.get("audio", 0)
+        return "<audio_patch>" * num_audios
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, object] | None = None,
+    ) -> MultiModalDataDict:
+        audio_len = 16000 * 25
+        num_audios = mm_counts.get("audio", 0)
+        return {"audio": self._get_dummy_audios(length=audio_len, num_audios=num_audios)}
+
+
+class StepAudio2MultiModalProcessor(BaseMultiModalProcessor[StepAudio2ProcessingInfo]):
+    """Multi-modal processor for Step-Audio2"""
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        """Map processor outputs to vLLM multimodal fields.
+
+        audio_mels is a padded 3D tensor per audio item with matching audio_lens.
+        Use flat_from_sizes to keep one mm item per audio, so downstream mm_kwargs
+        alignment does not collapse to a single item.
+        """
+        audio_lens = hf_inputs.get("audio_lens", torch.empty(0))
+        if isinstance(audio_lens, torch.Tensor):
+            lens_tensor = audio_lens.flatten()
+        elif audio_lens is None:
+            lens_tensor = torch.empty(0, dtype=torch.int32)
+        else:
+            lens_tensor = torch.tensor(
+                list(audio_lens) if hasattr(audio_lens, "__iter__") else [int(audio_lens)], dtype=torch.int32
+            )
+
+        return dict(
+            audio_mels=MultiModalFieldConfig.flat_from_sizes("audio", lens_tensor),
+            audio_lens=MultiModalFieldConfig.batched("audio"),
+        )
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+
+        audio_token = getattr(processor, "audio_token", "<audio_patch>")
+
+        audio_items = mm_items.get("audio", [])
+        num_audio_items = len(audio_items) if audio_items else 0
+
+        out_mm_data = out_mm_kwargs.get_data()
+        audio_lens = out_mm_data.get("audio_lens")
+
+        feature_lens = []
+        if audio_lens is not None:
+            if isinstance(audio_lens, torch.Tensor):
+                audio_lens_list = audio_lens.tolist()
+            else:
+                audio_lens_list = list(audio_lens) if hasattr(audio_lens, "__iter__") else [audio_lens]
+
+            for length in audio_lens_list:
+                if length > 0:
+                    feature_len = calculate_audio_feature_length(int(length))
+                    feature_lens.append(feature_len)
+
+        # CRITICAL: Align feature_lens with mm_items['audio'] count
+        if num_audio_items > 0:
+            if len(feature_lens) < num_audio_items:
+                default_feature_len = 250
+                pad_count = num_audio_items - len(feature_lens)
+                feature_lens.extend([default_feature_len] * pad_count)
+            elif len(feature_lens) > num_audio_items:
+                feature_lens = feature_lens[:num_audio_items]
+        elif not feature_lens:
+            feature_lens = [250]
+
+        def get_replacement_audio(item_idx: int):
+            """Generate replacement tokens for audio placeholder.
+
+            Following Qwen2.5-Omni pattern: returns list[int] of token IDs.
+            """
+            if item_idx >= len(feature_lens):
+                num_features = 1
+            else:
+                num_features = feature_lens[item_idx]
+
+            return [STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID] * num_features
+
+        return [
+            PromptReplacement(
+                modality="audio",
+                target=audio_token,
+                replacement=get_replacement_audio,
+            )
+        ]
+
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object] | None = None,
+        tok_kwargs: Mapping[str, object] | None = None,
+        **_: object,
+    ):
+        """Call HF processor and post-process outputs"""
+        mm_data = dict(mm_data)
+        audios = mm_data.pop("audios", [])
+
+        mm_kwargs = mm_kwargs or {}
+        tok_kwargs = tok_kwargs or {}
+
+        # CRITICAL: Ensure audios is ALWAYS a list to prevent string iteration
+        if audios:
+            if isinstance(audios, str):
+                mm_data["audio"] = [audios]
+            elif isinstance(audios, (list, tuple)):
+                mm_data["audio"] = audios
+            else:
+                mm_data["audio"] = [audios]
+
+        hf_inputs = super()._call_hf_processor(
+            prompt=prompt,
+            mm_data=mm_data,
+            mm_kwargs=mm_kwargs,
+            tok_kwargs=tok_kwargs,
+        )
+
+        if "audio_mels" not in hf_inputs:
+            hf_inputs["audio_mels"] = torch.empty((0, 128, 0))
+        if "audio_lens" not in hf_inputs:
+            hf_inputs["audio_lens"] = torch.tensor([], dtype=torch.int32)
+
+        return hf_inputs
+
+
+# ============================================================================
+# Main Thinker Model
+# ============================================================================
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    StepAudio2MultiModalProcessor,
+    info=StepAudio2ProcessingInfo,
+    dummy_inputs=StepAudio2DummyInputsBuilder,
+)
+class StepAudio2ThinkerForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
+    """
+    Step-Audio2 Thinker - Stage 1 LLM
+
+    Architecture:
+        AudioEncoder (6 layers, 512 hidden) →
+        Adaptor (512 → LLM dim) →
+        Qwen2 LLM (vocab=64012)
+    """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.have_multimodal_outputs = True
+
+        config = vllm_config.model_config.hf_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+
+        self.config = config
+        self.multimodal_config = multimodal_config
+
+        audio_enc_cfg = getattr(config, "audio_encoder_config", {})
+
+        def _get_config_value(cfg, attr_name: str, default_value):
+            """Get value from config object or dict"""
+            if isinstance(cfg, dict):
+                return cfg.get(attr_name, default_value)
+            else:
+                return getattr(cfg, attr_name, default_value)
+
+        n_mels = _get_config_value(audio_enc_cfg, "n_mels", DEFAULT_MODEL_CONFIG.n_mels)
+        n_audio_ctx = _get_config_value(audio_enc_cfg, "n_audio_ctx", DEFAULT_MODEL_CONFIG.n_audio_ctx)
+        n_audio_state = _get_config_value(audio_enc_cfg, "n_audio_state", DEFAULT_MODEL_CONFIG.n_audio_state)
+        n_audio_head = _get_config_value(audio_enc_cfg, "n_audio_head", DEFAULT_MODEL_CONFIG.n_audio_head)
+        n_audio_layer = _get_config_value(audio_enc_cfg, "n_audio_layer", DEFAULT_MODEL_CONFIG.n_audio_layer)
+        kernel_size = _get_config_value(audio_enc_cfg, "kernel_size", DEFAULT_MODEL_CONFIG.kernel_size)
+        adapter_stride = _get_config_value(audio_enc_cfg, "adapter_stride", DEFAULT_MODEL_CONFIG.adapter_stride)
+
+        n_hidden = None
+
+        if audio_enc_cfg:
+            n_hidden = _get_config_value(audio_enc_cfg, "llm_dim", None)
+
+        if n_hidden is None:
+            text_config = getattr(config, "text_config", None)
+            if text_config is not None:
+                n_hidden = _get_config_value(text_config, "hidden_size", None)
+
+        if n_hidden is None:
+            n_hidden = _get_config_value(config, "hidden_size", None)
+
+        if n_hidden is None:
+            n_hidden = DEFAULT_MODEL_CONFIG.hidden_size
+
+        self.encoder = AudioEncoder(
+            n_mels=n_mels,
+            n_ctx=n_audio_ctx,
+            n_state=n_audio_state,
+            n_head=n_audio_head,
+            n_layer=n_audio_layer,
+        )
+
+        self.adapter = Adaptor(
+            n_state=n_audio_state,
+            n_hidden=n_hidden,
+            kernel_size=kernel_size,
+            stride=adapter_stride,
+        )
+
+        from transformers import PretrainedConfig
+
+        text_config = getattr(config, "text_config", None)
+        lm_config = text_config if isinstance(text_config, PretrainedConfig) else config
+        architectures = getattr(lm_config, "architectures", None) or getattr(config, "architectures", None)
+        if not architectures:
+            model_type = getattr(lm_config, "model_type", None) or getattr(config, "model_type", None)
+            model_type_to_arch = {
+                "qwen2": "Qwen2ForCausalLM",
+                "qwen2_5": "Qwen2_5ForCausalLM",
+            }
+            arch = model_type_to_arch.get(model_type)
+            architectures = [arch] if arch else ["Qwen2ForCausalLM"]
+        lm_vllm_config = vllm_config.with_hf_config(lm_config, architectures=architectures)
+        self.language_model = init_vllm_registered_model(
+            vllm_config=lm_vllm_config, hf_config=lm_config, prefix=maybe_prefix(prefix, "language_model")
+        )
+
+        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        if modality.startswith("audio"):
+            return "<audio_patch>"
+        raise ValueError("Only audio modality is supported")
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def _parse_and_validate_audio_input(self, **kwargs: object) -> StepAudio2AudioInputs | None:
+        """Parse audio inputs from kwargs"""
+        audio_mels = kwargs.get("audio_mels", None)
+        audio_lens = kwargs.get("audio_lens", None)
+
+        if audio_mels is None:
+            return None
+
+        if isinstance(audio_mels, torch.Tensor):
+            if audio_mels.dim() >= 4:
+                audio_mels = audio_mels.flatten(0, 1)
+            elif audio_mels.dim() == 3:
+                # Already (B, n_mels, T)
+                pass
+            elif audio_mels.dim() == 2:
+                audio_mels = audio_mels.unsqueeze(0)
+            else:
+                return None
+        else:
+            audio_mels = flatten_bn(audio_mels, concat=True)
+
+        if isinstance(audio_lens, torch.Tensor):
+            if audio_lens.dim() >= 2:
+                audio_lens = audio_lens.flatten(0, 1)
+            elif audio_lens.dim() == 0:
+                audio_lens = audio_lens.unsqueeze(0)
+            # dim == 1 is already fine
+        else:
+            audio_lens = flatten_bn(audio_lens, concat=True)
+
+        expected_n_mels = 128
+        if audio_mels.ndim == 3 and audio_mels.size(1) != expected_n_mels:
+            return None
+
+        return StepAudio2AudioInputs(
+            audio_mels=audio_mels.to(self.dtype).to(self.device),
+            audio_lens=audio_lens,
+        )
+
+    def _process_audio_input(self, audio_input: StepAudio2AudioInputs) -> tuple[torch.Tensor, ...]:
+        """Process audio mels through encoder and adapter"""
+        audio_mels = audio_input["audio_mels"]
+        audio_lens = audio_input["audio_lens"]
+        if not isinstance(audio_lens, torch.Tensor):
+            audio_lens = torch.tensor(audio_lens, device=self.device)
+        else:
+            audio_lens = audio_lens.to(self.device)
+
+        audio_features, audio_lens = self.encoder(audio_mels, audio_lens)
+
+        audio_features = self.adapter(audio_features)
+
+        audio_feature_lens = (audio_lens - 1) // 2 + 1
+
+        audio_feature_list = [audio_features[i, : audio_feature_lens[i]] for i in range(audio_features.size(0))]
+
+        return audio_feature_list
+
+    def get_multimodal_embeddings(self, **kwargs) -> NestedTensors | None:
+        """Get multimodal embeddings"""
+        audio_input = self._parse_and_validate_audio_input(**kwargs)
+        if audio_input is None:
+            return None
+        else:
+            audio_embeddings = self._process_audio_input(audio_input)
+            return audio_embeddings
+
+    def embed_multimodal(self, **kwargs: object) -> NestedTensors:
+        """vLLM multimodal encoder entrypoint."""
+        audio_embeddings = self.get_multimodal_embeddings(**kwargs)
+        if audio_embeddings is None:
+            return []
+        return audio_embeddings
+
+    def get_input_embeddings(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: NestedTensors | None = None,
+    ) -> torch.Tensor:
+        """Get input embeddings with multimodal fusion"""
+        # Under newer vLLM versions the registered LM wrapper exposes
+        # get_input_embeddings(), while the inner `.model` may be a bare
+        # Qwen2Model without that compatibility method.
+        inputs_embeds = self.language_model.embed_input_ids(input_ids)
+        if multimodal_embeddings is not None:
+            inputs_embeds = _merge_multimodal_embeddings(
+                input_ids, inputs_embeds, multimodal_embeddings, STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID
+            )
+        return inputs_embeds
+
+    def get_language_model(self) -> nn.Module:
+        return self.language_model
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ):
+        """Forward pass - returns OmniOutput for multi-stage pipeline"""
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+        elif inputs_embeds is None:
+            audio_embeddings = self.get_multimodal_embeddings(**kwargs)
+            inputs_embeds = self.get_input_embeddings(input_ids, audio_embeddings)
+            input_ids = None
+
+        hidden_states = self.language_model(
+            input_ids,
+            positions,
+            intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+        return OmniOutput(
+            text_hidden_states=hidden_states,
+            multimodal_outputs={},
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        return self.language_model.compute_logits(hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights with name mapping from HuggingFace to vLLM structure"""
+        from vllm.model_executor.models.utils import AutoWeightsLoader
+
+        # Map HuggingFace weight names to our model structure
+        # HF: model.*, lm_head.* → Our: language_model.model.*, language_model.lm_head.*
+        # HF: encoder.*, adapter.* → Our: encoder.*, adapter.* (no change)
+        def weight_mapper(name: str) -> str:
+            if name.startswith("lm_head."):
+                return name.replace("lm_head.", "language_model.lm_head.")
+            elif name.startswith("model."):
+                return name.replace("model.", "language_model.model.")
+            else:
+                # encoder.* and adapter.* remain unchanged
+                return name
+
+        # Apply mapping to weights
+        mapped_weights = [(weight_mapper(name), tensor) for name, tensor in weights]
+
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(mapped_weights)
+
+    @staticmethod
+    def separate_tokens(
+        token_ids: list[int],
+        text_max: int = DEFAULT_TOKEN_CONFIG.text_max,
+        audio_start: int = DEFAULT_TOKEN_CONFIG.audio_start,
+    ) -> tuple[list[int], list[int]]:
+        """Separate generated tokens into text and audio tokens"""
+        text_tokens = [tid for tid in token_ids if tid < text_max]
+        audio_tokens = [tid - audio_start for tid in token_ids if tid >= audio_start]
+        return text_tokens, audio_tokens
+
+    @staticmethod
+    def has_audio_output(token_ids: list[int], audio_start: int = DEFAULT_TOKEN_CONFIG.audio_start) -> bool:
+        """Check if generated tokens contain audio tokens"""
+        return any(tid >= audio_start for tid in token_ids)
+
+
+class StepAudio2OutputProcessor:
+    """Helper class to process Step-Audio2 outputs"""
+
+    @staticmethod
+    def process_output(output_ids: torch.Tensor, tokenizer, remove_audio_padding: bool = True) -> dict:
+        """Process model output and separate text/audio tokens"""
+        if isinstance(output_ids, torch.Tensor):
+            output_ids = output_ids.squeeze().tolist()
+
+        # Separate tokens
+        text_tokens, audio_tokens = StepAudio2ThinkerForConditionalGeneration.separate_tokens(output_ids)
+
+        # Remove audio padding if requested
+        if remove_audio_padding and audio_tokens:
+            audio_tokens = [t for t in audio_tokens if t < DEFAULT_TOKEN_CONFIG.audio_eos]
+
+        # Decode text
+        text = tokenizer.decode(text_tokens, skip_special_tokens=False)
+
+        return {
+            "text": text,
+            "text_tokens": text_tokens,
+            "audio_tokens": audio_tokens,
+            "has_audio": len(audio_tokens) > 0,
+            "all_tokens": output_ids,
+        }

vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py

@@ -0,0 +1,759 @@
+import io
+import os
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from typing import Any
+
+import librosa  # noqa: TID251
+import numpy as np
+import onnxruntime
+import s3tokenizer
+import soundfile as sf
+import torch
+import torch.nn as nn
+import torchaudio
+import torchaudio.compliance.kaldi as kaldi
+from flashcosyvoice.modules.hifigan import HiFTGenerator
+from flashcosyvoice.utils.audio import mel_spectrogram
+from hyperpyyaml import load_hyperpyyaml
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.models.interfaces import SupportsPP
+from vllm.model_executor.models.utils import WeightsMapper
+from vllm.sequence import IntermediateTensors
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.sampler import Sampler
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import (
+    DEFAULT_STREAM_CONFIG,
+    DEFAULT_TOKEN2WAV_CONFIG,
+    STEP_AUDIO2_DEFAULT_PROMPT_WAV,
+    STREAM_SOURCE_CACHE_LEN,
+)
+
+logger = init_logger(__name__)
+
+
+def fade_in_out(
+    fade_in_mel: torch.Tensor,
+    fade_out_mel: torch.Tensor,
+    window: torch.Tensor,
+) -> torch.Tensor:
+    """Cross-fade two overlapping waveform segments using a Hamming window.
+
+    The window is split in half: the first half ramps *up* (fade-in) and
+    the second half ramps *down* (fade-out).  The overlap region of
+    ``fade_in_mel`` is blended with the tail of ``fade_out_mel``.
+    """
+    mel_overlap_len = int(window.shape[0] / 2)
+    fade_in_mel = fade_in_mel.clone()
+    fade_in_mel[..., :mel_overlap_len] = (
+        fade_in_mel[..., :mel_overlap_len] * window[:mel_overlap_len]
+        + fade_out_mel[..., -mel_overlap_len:] * window[mel_overlap_len:]
+    )
+    return fade_in_mel
+
+
+@dataclass
+class _StreamState:
+    """Per-request streaming state for concurrent Token2Wav sessions."""
+
+    setup_done: bool = False
+    stream_cache: dict | None = None
+    hift_cache_dict: dict = field(default_factory=dict)
+    finished: bool = False
+
+
+class StepAudio2Token2WavCore(nn.Module):
+    """Core Token2Wav model - handles token to waveform conversion"""
+
+    def __init__(
+        self,
+        model_path: str,
+        float16: bool = False,
+        device: str = "cuda",
+        n_timesteps: int = DEFAULT_TOKEN2WAV_CONFIG.n_timesteps,
+    ):
+        super().__init__()
+        self.model_path = model_path
+        self.float16 = float16
+        self.device = torch.device(device)
+        self.n_timesteps = n_timesteps
+
+        self._models_loaded = False
+        self._audio_tokenizer = None
+        self._spk_model = None
+        self._flow = None
+        self._hift = None
+
+        self.cache = {}
+
+        # Streaming state (constants from centralised config)
+        self.mel_cache_len = DEFAULT_STREAM_CONFIG.mel_cache_len
+        self.source_cache_len = STREAM_SOURCE_CACHE_LEN
+        self.speech_window: torch.Tensor | None = None  # created lazily on device
+        self.stream_cache: dict | None = None  # flow model causal cache
+        self.hift_cache_dict: dict = {}  # HiFT vocoder overlap cache
+
+    def _ensure_models_loaded(self):
+        """Lazy load models on first use"""
+        if self._models_loaded:
+            return
+
+        from vllm.logger import init_logger
+
+        logger = init_logger(__name__)
+        logger.info(f"Loading Token2Wav models from: {self.model_path}")
+
+        self._audio_tokenizer = (
+            s3tokenizer.load_model(f"{self.model_path}/speech_tokenizer_v2_25hz.onnx").to(self.device).eval()
+        )
+
+        option = onnxruntime.SessionOptions()
+        option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+        option.intra_op_num_threads = 1
+        self._spk_model = onnxruntime.InferenceSession(
+            f"{self.model_path}/campplus.onnx", sess_options=option, providers=["CPUExecutionProvider"]
+        )
+
+        with open(f"{self.model_path}/flow.yaml") as f:
+            configs = load_hyperpyyaml(f)
+        self._flow = configs["flow"]
+        if self.float16:
+            self._flow.half()
+        self._flow.load_state_dict(
+            torch.load(f"{self.model_path}/flow.pt", map_location="cpu", weights_only=True), strict=True
+        )
+        self._flow.to(self.device).eval()
+
+        self._hift = HiFTGenerator()
+        hift_state_dict = {
+            k.replace("generator.", ""): v
+            for k, v in torch.load(f"{self.model_path}/hift.pt", map_location="cpu", weights_only=True).items()
+        }
+        self._hift.load_state_dict(hift_state_dict, strict=True)
+        self._hift.to(self.device).eval()
+
+        self._models_loaded = True
+        logger.info("Token2Wav models loaded successfully")
+
+    @property
+    def audio_tokenizer(self):
+        """Lazy-loaded audio tokenizer"""
+        self._ensure_models_loaded()
+        return self._audio_tokenizer
+
+    @property
+    def spk_model(self):
+        """Lazy-loaded speaker model"""
+        self._ensure_models_loaded()
+        return self._spk_model
+
+    @property
+    def flow(self):
+        """Lazy-loaded flow model"""
+        self._ensure_models_loaded()
+        return self._flow
+
+    @property
+    def hift(self):
+        """Lazy-loaded HiFT vocoder"""
+        self._ensure_models_loaded()
+        return self._hift
+
+    def _prepare_prompt(self, prompt_wav: str):
+        """Prepare prompt audio for conditioning"""
+        # Prefer soundfile/librosa path to avoid torchaudio->torchcodec runtime coupling.
+        try:
+            audio_np_raw, sample_rate = sf.read(prompt_wav, dtype="float32", always_2d=False)
+        except Exception:
+            audio_np_raw, sample_rate = librosa.load(prompt_wav, sr=None, mono=True)
+        if isinstance(audio_np_raw, np.ndarray) and audio_np_raw.ndim > 1:
+            audio_np_raw = np.mean(audio_np_raw, axis=-1)
+        sample_rate = int(sample_rate)
+
+        # 16 kHz branch: s3tokenizer + speaker embedding
+        audio_16k = audio_np_raw.astype(np.float32)
+        if sample_rate != 16000:
+            audio_16k = librosa.resample(y=audio_16k, orig_sr=sample_rate, target_sr=16000)
+        audio = torch.from_numpy(audio_16k)
+        mels = s3tokenizer.log_mel_spectrogram(audio)
+        mels, mels_lens = s3tokenizer.padding([mels])
+        prompt_speech_tokens, prompt_speech_tokens_lens = self.audio_tokenizer.quantize(
+            mels.to(self.device), mels_lens.to(self.device)
+        )
+
+        spk_feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000)
+        spk_feat = spk_feat - spk_feat.mean(dim=0, keepdim=True)
+        spk_emb = torch.tensor(
+            self.spk_model.run(None, {self.spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()})[0],
+            device=self.device,
+        )
+
+        # 24 kHz branch: mel spectrogram for flow model conditioning
+        # Must resample from the ORIGINAL audio, not the 16 kHz version.
+        audio_24k = audio_np_raw.astype(np.float32)
+        if sample_rate != 24000:
+            audio_24k = librosa.resample(y=audio_24k, orig_sr=sample_rate, target_sr=24000)
+        audio_24k_t = torch.from_numpy(audio_24k).unsqueeze(0)  # [1, T]
+        prompt_mel = mel_spectrogram(audio_24k_t).transpose(1, 2).squeeze(0)  # [T, num_mels]
+        prompt_mels = prompt_mel.unsqueeze(0).to(self.device)
+        prompt_mels_lens = torch.tensor([prompt_mels.shape[1]], dtype=torch.int32, device=self.device)
+        prompt_mels = torch.nn.functional.pad(
+            prompt_mels,
+            (0, 0, 0, prompt_speech_tokens.shape[1] * self.flow.up_rate - prompt_mels.shape[1]),
+            mode="replicate",
+        )
+
+        return prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens
+
+    def forward(
+        self, generated_speech_tokens: list, prompt_wav: str, return_bytes: bool = True
+    ) -> torch.Tensor | bytes:
+        """
+        Convert audio tokens to waveform
+
+        Args:
+            generated_speech_tokens: List of audio token IDs (0-6560)
+            prompt_wav: Path to prompt wav file for speaker conditioning
+            return_bytes: If True, return WAV bytes; else return tensor
+
+        Returns:
+            Audio waveform as bytes or tensor
+        """
+        if prompt_wav not in self.cache:
+            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
+        prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens = self.cache[prompt_wav]
+
+        generated_speech_tokens = torch.tensor([generated_speech_tokens], dtype=torch.int32, device=self.device)
+        generated_speech_tokens_lens = torch.tensor(
+            [generated_speech_tokens.shape[1]], dtype=torch.int32, device=self.device
+        )
+
+        with torch.amp.autocast(str(self.device.type), dtype=torch.float16 if self.float16 else torch.float32):
+            mel = self.flow.inference(
+                generated_speech_tokens,
+                generated_speech_tokens_lens,
+                prompt_speech_tokens,
+                prompt_speech_tokens_lens,
+                prompt_mels,
+                prompt_mels_lens,
+                spk_emb,
+                self.n_timesteps,
+            )
+
+        wav, _ = self.hift(speech_feat=mel)
+
+        if return_bytes:
+            output = io.BytesIO()
+            torchaudio.save(output, wav.cpu(), sample_rate=DEFAULT_TOKEN2WAV_CONFIG.sample_rate, format="wav")
+            return output.getvalue()
+        else:
+            return wav
+
+    # ------------------------------------------------------------------
+    # Streaming (chunked) inference
+    # ------------------------------------------------------------------
+
+    def setup_stream(self, prompt_wav: str) -> None:
+        """Initialize flow + HiFT caches for chunked streaming inference.
+
+        Must be called once before the first ``stream_chunk()`` call for a
+        given utterance.
+        """
+        if prompt_wav not in self.cache:
+            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
+        (
+            prompt_speech_tokens,
+            _,
+            spk_emb,
+            prompt_mels,
+            _,
+        ) = self.cache[prompt_wav]
+
+        # Lazily create the Hamming window on the correct device
+        if self.speech_window is None or self.speech_window.device != self.device:
+            self.speech_window = torch.from_numpy(np.hamming(2 * self.source_cache_len)).to(
+                device=self.device, dtype=torch.float32
+            )
+
+        # Flow model: setup causal cache with prompt tokens + lookahead
+        pre_lookahead = DEFAULT_STREAM_CONFIG.pre_lookahead_len
+        self.stream_cache = self.flow.setup_cache(
+            torch.cat([prompt_speech_tokens, prompt_speech_tokens[:, :pre_lookahead]], dim=1),
+            prompt_mels,
+            spk_emb,
+            n_timesteps=self.n_timesteps,
+        )
+
+        # HiFT vocoder: empty overlap buffers
+        self.hift_cache_dict = {
+            "mel": torch.zeros(1, prompt_mels.shape[2], 0, device=self.device),
+            "source": torch.zeros(1, 1, 0, device=self.device),
+            "speech": torch.zeros(1, 0, device=self.device),
+        }
+
+    def stream_chunk(
+        self,
+        audio_tokens: list[int],
+        prompt_wav: str,
+        last_chunk: bool = False,
+    ) -> torch.Tensor:
+        """Process one chunk of audio tokens and return a waveform segment.
+
+        Args:
+            audio_tokens: Token IDs for this chunk (includes lookahead for
+                non-last chunks).
+            prompt_wav: Path to the prompt wav (for cached speaker embeddings).
+            last_chunk: Whether this is the final chunk of the utterance.
+
+        Returns:
+            1-D float tensor of audio samples (24 kHz).
+        """
+        if self.stream_cache is None:
+            raise ValueError("stream_cache not initialised – call setup_stream() first")
+
+        if prompt_wav not in self.cache:
+            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
+        _, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav]
+
+        token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device)
+
+        with torch.amp.autocast(
+            str(self.device.type),
+            dtype=torch.float16 if self.float16 else torch.float32,
+        ):
+            chunk_mel, self.stream_cache = self.flow.inference_chunk(
+                token=token_tensor,
+                spk=spk_emb,
+                cache=self.stream_cache,
+                last_chunk=last_chunk,
+                n_timesteps=self.n_timesteps,
+            )
+
+        # Trim estimator attention cache to avoid unbounded growth
+        keep = DEFAULT_STREAM_CONFIG.estimator_cache_keep
+        est_att = self.stream_cache.get("estimator_att_cache")
+        if est_att is not None and est_att.shape[4] > (prompt_mels.shape[1] + keep):
+            self.stream_cache["estimator_att_cache"] = torch.cat(
+                [est_att[:, :, :, :, : prompt_mels.shape[1]], est_att[:, :, :, :, -keep:]],
+                dim=4,
+            )
+
+        # ---- HiFT vocoder with overlap-add ----
+        hift_cache_mel = self.hift_cache_dict["mel"]
+        hift_cache_source = self.hift_cache_dict["source"]
+        hift_cache_speech = self.hift_cache_dict["speech"]
+
+        mel = torch.cat([hift_cache_mel, chunk_mel], dim=2)
+        speech, source = self.hift(mel, hift_cache_source)
+
+        # Cross-fade overlap region for smooth concatenation
+        if hift_cache_speech.shape[-1] > 0:
+            speech = fade_in_out(speech, hift_cache_speech, self.speech_window)
+
+        # Update vocoder cache for the next chunk
+        self.hift_cache_dict = {
+            "mel": mel[..., -self.mel_cache_len :].clone().detach(),
+            "source": source[:, :, -self.source_cache_len :].clone().detach(),
+            "speech": speech[:, -self.source_cache_len :].clone().detach(),
+        }
+
+        # For non-last chunks, trim the tail that will be re-synthesised
+        if not last_chunk:
+            speech = speech[:, : -self.source_cache_len]
+
+        return speech.squeeze(0)  # [samples]
+
+    def reset_stream(self) -> None:
+        """Clear all streaming state after an utterance is complete."""
+        self.stream_cache = None
+        self.hift_cache_dict = {}
+
+    # ------------------------------------------------------------------
+    # Per-request streaming (does NOT mutate self – uses external state)
+    # ------------------------------------------------------------------
+
+    def setup_stream_for(self, prompt_wav: str, state: _StreamState) -> None:
+        """Initialise flow + HiFT caches into *state* (no self mutation)."""
+        if prompt_wav not in self.cache:
+            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
+        prompt_speech_tokens, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav]
+
+        if self.speech_window is None or self.speech_window.device != self.device:
+            self.speech_window = torch.from_numpy(np.hamming(2 * self.source_cache_len)).to(
+                device=self.device, dtype=torch.float32
+            )
+
+        pre_lookahead = DEFAULT_STREAM_CONFIG.pre_lookahead_len
+        state.stream_cache = self.flow.setup_cache(
+            torch.cat(
+                [prompt_speech_tokens, prompt_speech_tokens[:, :pre_lookahead]],
+                dim=1,
+            ),
+            prompt_mels,
+            spk_emb,
+            n_timesteps=self.n_timesteps,
+        )
+        state.hift_cache_dict = {
+            "mel": torch.zeros(1, prompt_mels.shape[2], 0, device=self.device),
+            "source": torch.zeros(1, 1, 0, device=self.device),
+            "speech": torch.zeros(1, 0, device=self.device),
+        }
+        state.setup_done = True
+
+    def stream_chunk_for(
+        self,
+        audio_tokens: list[int],
+        prompt_wav: str,
+        last_chunk: bool,
+        state: _StreamState,
+    ) -> torch.Tensor:
+        """Process one chunk using *state* (no self mutation except speech_window)."""
+        if state.stream_cache is None:
+            raise ValueError("stream_cache not initialised – call setup_stream_for() first")
+
+        if prompt_wav not in self.cache:
+            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
+        _, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav]
+
+        token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device)
+
+        with torch.amp.autocast(
+            str(self.device.type),
+            dtype=torch.float16 if self.float16 else torch.float32,
+        ):
+            chunk_mel, state.stream_cache = self.flow.inference_chunk(
+                token=token_tensor,
+                spk=spk_emb,
+                cache=state.stream_cache,
+                last_chunk=last_chunk,
+                n_timesteps=self.n_timesteps,
+            )
+
+        # Trim estimator attention cache to avoid unbounded growth
+        keep = DEFAULT_STREAM_CONFIG.estimator_cache_keep
+        est_att = state.stream_cache.get("estimator_att_cache")
+        if est_att is not None and est_att.shape[4] > (prompt_mels.shape[1] + keep):
+            state.stream_cache["estimator_att_cache"] = torch.cat(
+                [
+                    est_att[:, :, :, :, : prompt_mels.shape[1]],
+                    est_att[:, :, :, :, -keep:],
+                ],
+                dim=4,
+            )
+
+        # ---- HiFT vocoder with overlap-add ----
+        hift_cache_mel = state.hift_cache_dict["mel"]
+        hift_cache_source = state.hift_cache_dict["source"]
+        hift_cache_speech = state.hift_cache_dict["speech"]
+
+        mel = torch.cat([hift_cache_mel, chunk_mel], dim=2)
+        speech, source = self.hift(mel, hift_cache_source)
+
+        if hift_cache_speech.shape[-1] > 0:
+            speech = fade_in_out(speech, hift_cache_speech, self.speech_window)
+
+        state.hift_cache_dict = {
+            "mel": mel[..., -self.mel_cache_len :].clone().detach(),
+            "source": source[:, :, -self.source_cache_len :].clone().detach(),
+            "speech": speech[:, -self.source_cache_len :].clone().detach(),
+        }
+
+        if not last_chunk:
+            speech = speech[:, : -self.source_cache_len]
+
+        return speech.squeeze(0)
+
+    @staticmethod
+    def reset_stream_for(state: _StreamState) -> None:
+        """Clear streaming caches for the given state."""
+        state.stream_cache = None
+        state.hift_cache_dict = {}
+        state.setup_done = False
+        state.finished = True
+
+
+class StepAudio2Token2WavForConditionalGeneration(nn.Module, SupportsPP):
+    """vLLM-compatible wrapper for Step-Audio2 Token2Wav"""
+
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "model.": "token2wav.",
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.config = vllm_config.model_config.hf_config
+
+        model_path = getattr(self.config, "token2wav_path", None)
+        model_name_or_path = vllm_config.model_config.model
+        if model_path is None:
+            # Resolve HF repo names to local cache path
+            if not os.path.isdir(model_name_or_path):
+                from huggingface_hub import snapshot_download
+
+                model_name_or_path = snapshot_download(model_name_or_path)
+            model_path = f"{model_name_or_path}/token2wav"
+
+        float16 = getattr(self.config, "token2wav_float16", False)
+
+        device = "cuda"
+        if hasattr(vllm_config, "device_config") and vllm_config.device_config:
+            device = str(vllm_config.device_config.device)
+
+        self.token2wav = StepAudio2Token2WavCore(model_path=model_path, float16=float16, device=device)
+
+        self.have_multimodal_outputs = True
+        # Required for the runner to pass async_chunk info via
+        # runtime_additional_information on every forward step.
+        self.enable_update_additional_information = True
+
+        # Per-request streaming states (ordered list matching batch order).
+        self._stream_states: list[_StreamState] = []
+
+        # Resolve default prompt wav path from model directory.
+        # Priority: env var > {model_dir}/assets/default_female.wav > bare filename (cwd)
+        self._default_prompt_wav = self._resolve_prompt_wav(model_name_or_path)
+
+        self.make_empty_intermediate_tensors = lambda: None
+
+    @staticmethod
+    def _resolve_prompt_wav(model_root: str) -> str:
+        """Find the default prompt wav, searching model assets directory."""
+        env_val = os.environ.get("STEP_AUDIO2_DEFAULT_PROMPT_WAV")
+        if env_val and os.path.exists(env_val):
+            return env_val
+        # Search common locations: model dir, then bundled assets in this package
+        pkg_dir = os.path.dirname(__file__)
+        candidates = [
+            os.path.join(model_root, "assets", STEP_AUDIO2_DEFAULT_PROMPT_WAV),
+            os.path.join(model_root, STEP_AUDIO2_DEFAULT_PROMPT_WAV),
+            os.path.join(pkg_dir, "assets", STEP_AUDIO2_DEFAULT_PROMPT_WAV),
+            os.path.join(pkg_dir, STEP_AUDIO2_DEFAULT_PROMPT_WAV),
+            STEP_AUDIO2_DEFAULT_PROMPT_WAV,
+        ]
+        for path in candidates:
+            if os.path.exists(path):
+                return path
+        # Return the env var (even if not found) or the bare filename as fallback
+        return env_val or STEP_AUDIO2_DEFAULT_PROMPT_WAV
+
+    def get_language_model(self) -> torch.nn.Module:
+        # Return the wrapper so vLLM interface helpers resolve
+        # embed_input_ids() on this adapter instead of the bare core model.
+        return self
+
+    @property
+    def sampler(self):
+        return Sampler()
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        is_multimodal=None,
+    ) -> torch.Tensor:
+        hidden_size = self.vllm_config.model_config.get_hidden_size()
+        return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, hidden_size)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        additional: dict | None = None,
+        runtime_additional_information: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor | bytes:
+        """
+        Forward pass for Token2Wav
+
+        Args:
+            input_ids: Audio token IDs (from stage input processor)
+            positions: Position IDs (unused)
+            intermediate_tensors: Intermediate tensors (unused)
+            inputs_embeds: Input embeddings (unused)
+            additional: Additional information dict (deprecated)
+            runtime_additional_information: Per-request dicts from the
+                async_chunk processor.  Each dict contains ``audio_tokens``
+                (list[int]) and ``last_chunk`` (bool).
+            **kwargs: Other kwargs
+
+        Returns:
+            OmniOutput with waveform tensor
+        """
+        # ----- async_chunk streaming path -----
+        # Only enter when at least one entry carries ``left_context_size``
+        # (the field the framework forwards from the async_chunk payload).
+        # The runner always passes runtime_additional_information (possibly
+        # [{}]), so a bare truthiness check would incorrectly enter this
+        # branch in synchronous mode.
+        if runtime_additional_information and any(
+            "left_context_size" in info for info in runtime_additional_information
+        ):
+            return self._forward_async_chunk(input_ids, runtime_additional_information, **kwargs)
+
+        # ----- profiling / warmup guard -----
+        # Profile/warmup passes use all-zero token IDs. Real audio tokens
+        # are non-zero, so check sum == 0 rather than a length threshold
+        # (short utterances may have fewer than 50 tokens).
+        if isinstance(input_ids, torch.Tensor):
+            is_profile_mode = input_ids.numel() == 0 or input_ids.sum().item() == 0
+        else:
+            is_profile_mode = not input_ids if not hasattr(input_ids, "__len__") else len(input_ids) == 0
+
+        if is_profile_mode:
+            device = torch.device("cpu")
+            if isinstance(inputs_embeds, torch.Tensor) and inputs_embeds.numel() > 0:
+                device = inputs_embeds.device
+            elif isinstance(input_ids, torch.Tensor) and input_ids.numel() > 0:
+                device = input_ids.device
+
+            dummy_audio = torch.zeros(1000, device=device, dtype=torch.float32)
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={"model_outputs": [dummy_audio]},
+            )
+
+        # ----- synchronous (full-sequence) path -----
+        if isinstance(input_ids, torch.Tensor):
+            if input_ids.dim() == 1:
+                input_ids_list = input_ids.cpu().tolist()
+            else:
+                input_ids_list = input_ids[0].cpu().tolist()
+        else:
+            input_ids_list = input_ids if isinstance(input_ids, list) else [input_ids]
+
+        prompt_wav = self._default_prompt_wav
+
+        if not os.path.exists(prompt_wav):
+            raise FileNotFoundError(f"Token2Wav: prompt_wav file not found: {prompt_wav}")
+
+        waveform_tensor = self.token2wav(input_ids_list, prompt_wav=prompt_wav, return_bytes=False)
+
+        if waveform_tensor.dim() == 2 and waveform_tensor.shape[0] == 1:
+            waveform_tensor = waveform_tensor.squeeze(0)
+
+        waveform_tensor_cpu = waveform_tensor.detach().cpu().contiguous()
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={"model_outputs": [waveform_tensor_cpu]},
+        )
+
+    # ------------------------------------------------------------------
+    # Async-chunk streaming forward
+    # ------------------------------------------------------------------
+
+    def _forward_async_chunk(
+        self,
+        input_ids: torch.Tensor,
+        runtime_additional_information: list[dict[str, Any]],
+        **kwargs: Any,
+    ) -> OmniOutput:
+        """Handle one async_chunk forward step (batch=1).
+
+        The end-of-stream flag is encoded as ``left_context_size == 1``
+        (the only field the framework transparently forwards from the
+        async_chunk payload to ``request.additional_information``).
+        """
+        # --- batch=1 guard ---
+        batch_size = sum(1 for info in runtime_additional_information if "left_context_size" in info)
+        if batch_size != 1:
+            raise RuntimeError(
+                f"Token2Wav async_chunk only supports batch=1, got {batch_size}. "
+                "Batch>1 requires framework support for batch_req_ids."
+            )
+
+        info = next(info for info in runtime_additional_information if "left_context_size" in info)
+        last_chunk = info.get("left_context_size", 0) == 1
+
+        # --- Manage single stream state ---
+        # If the previous request was preempted (setup_done but not finished),
+        # its stale caches must be discarded before we start a new session.
+        if self._stream_states and not self._stream_states[0].finished:
+            state = self._stream_states[0]
+            if state.setup_done and not last_chunk:
+                # Heuristic: a non-last chunk arriving into an already-active
+                # state that we did NOT set up this call means a new request
+                # reused the slot.  For a genuinely continuing request the
+                # state is fine; we cannot distinguish without request IDs,
+                # so we only forcibly reset on an explicit new-session signal
+                # (setup_done=True but stream_cache already consumed all
+                # previous data -- detected below when stream_chunk fails).
+                pass
+        if not self._stream_states or self._stream_states[0].finished:
+            self._stream_states = [_StreamState()]
+        state = self._stream_states[0]
+
+        # --- Extract audio tokens ---
+        audio_tokens = input_ids.flatten().cpu().tolist()
+
+        # Empty chunk (e.g. EOF with no audio tokens)
+        if not audio_tokens:
+            if state.setup_done:
+                self.token2wav.reset_stream_for(state)
+            else:
+                state.finished = True
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={"model_outputs": [torch.zeros(1, dtype=torch.float32)]},
+            )
+
+        prompt_wav = self._default_prompt_wav
+        if not os.path.exists(prompt_wav):
+            raise FileNotFoundError(f"Token2Wav: prompt_wav file not found: {prompt_wav}")
+
+        try:
+            if not state.setup_done:
+                logger.info("Token2Wav: initialising streaming caches")
+                self.token2wav.setup_stream_for(prompt_wav, state)
+
+            waveform = self.token2wav.stream_chunk_for(
+                audio_tokens,
+                prompt_wav=prompt_wav,
+                last_chunk=last_chunk,
+                state=state,
+            )
+        except Exception:
+            logger.exception("Token2Wav: stream error, resetting state")
+            self.token2wav.reset_stream_for(state)
+            raise
+
+        if last_chunk:
+            self.token2wav.reset_stream_for(state)
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={"model_outputs": [waveform.detach().cpu().contiguous()]},
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        return hidden_states
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> SamplerOutput | None:
+        return None
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loaded_params = set()
+        for name, _ in self.token2wav.named_parameters():
+            loaded_params.add(f"token2wav.{name}")
+        for name, _ in self.token2wav.named_buffers():
+            loaded_params.add(f"token2wav.{name}")
+        return loaded_params
+
+    def load_weights_without_buffers(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loaded_params = set()
+        for name, _ in self.token2wav.named_parameters():
+            loaded_params.add(f"token2wav.{name}")
+        return loaded_params

vllm_omni/model_executor/stage_configs/step_audio_2.yaml

@@ -0,0 +1,122 @@
+# 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

@@ -0,0 +1,35 @@
+# 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

@@ -0,0 +1,78 @@
+# 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_input_processors/step_audio2.py

@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stage input processor for Step-Audio2: Thinker → Token2Wav transition."""
+
+from typing import Any
+
+import torch
+from vllm.inputs import TextPrompt
+from vllm.logger import init_logger
+
+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 (
+    DEFAULT_STREAM_CONFIG,
+    DEFAULT_TOKEN_CONFIG,
+)
+
+logger = init_logger(__name__)
+
+
+def _ensure_list(x):
+    """Convert ConstantList / tensor / iterable to Python list."""
+    if hasattr(x, "_x"):
+        return list(x._x)
+    elif isinstance(x, torch.Tensor):
+        return x.tolist()
+    elif not isinstance(x, list):
+        return list(x) if hasattr(x, "__iter__") else [x]
+    return list(x)
+
+
+# =========================
+# Async chunk processor
+# =========================
+
+
+def thinker2token2wav_async_chunk(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any],
+    request: OmniEngineCoreRequest,
+    is_finished: bool = False,
+) -> dict[str, Any] | None:
+    """
+    Async chunk processor: stream audio tokens from Thinker to Token2Wav.
+
+    Unlike Qwen3-Omni which passes hidden states via pooling_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.
+
+    The flow model's conformer encoder requires ``pre_lookahead_len`` extra
+    tokens beyond the chunk boundary.  Each chunk therefore contains
+    ``chunk_size + pre_lookahead_len`` tokens, but the consumed pointer
+    only advances by ``chunk_size`` so the lookahead tokens are re-used as
+    the start of the next chunk's context (cached inside the encoder).
+
+    Payload adapts to the framework's ``_poll_single_request`` non-AR path:
+    - ``code_predictor_codes`` → ``request.prompt_token_ids`` → ``input_ids``
+    - ``left_context_size``   → ``request.additional_information``
+      → ``runtime_additional_information`` (0 = not last, 1 = last chunk)
+    - ``finished``            → marks request as finished in adapter
+
+    Args:
+        transfer_manager: OmniChunkTransferAdapter instance (has
+            ``code_prompt_token_ids`` defaultdict for state tracking).
+        pooling_output: Pooler 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.
+    """
+    audio_start = DEFAULT_TOKEN_CONFIG.audio_start
+    audio_eos = DEFAULT_TOKEN_CONFIG.audio_eos
+    finished = bool(is_finished or request.is_finished())
+
+    # Only look at decode (generated) tokens — the prompt may contain
+    # historical audio tokens from prior conversation turns.
+    all_token_ids = _ensure_list(request.all_token_ids)
+    prompt_len = len(_ensure_list(request.prompt_token_ids))
+    generated_ids = all_token_ids[prompt_len:]
+
+    # Extract audio tokens and convert to 0-based IDs for Token2Wav
+    audio_tokens = [tid - audio_start for tid in generated_ids if tid >= audio_start]
+    # Remove padding / EOS tokens
+    audio_tokens = [t for t in audio_tokens if t < audio_eos]
+
+    # Flow model streaming parameters (from centralised config)
+    chunk_size = DEFAULT_STREAM_CONFIG.chunk_size
+    pre_lookahead_len = DEFAULT_STREAM_CONFIG.pre_lookahead_len
+
+    # consumed = number of tokens whose mel output has been produced.
+    # We track this via transfer_manager.code_prompt_token_ids[request_id].
+    request_id = request.external_req_id
+    consumed = len(transfer_manager.code_prompt_token_ids[request_id])
+    available = len(audio_tokens) - consumed
+
+    if finished:
+        # 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),
+            }
+        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),
+        }
+    else:
+        # Non-last chunk: need chunk_size + pre_lookahead_len tokens
+        required = chunk_size + pre_lookahead_len
+        if available < required:
+            return None
+        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),
+        }
+
+
+def thinker2token2wav(
+    source_outputs: list[Any],
+    prompt: OmniTokensPrompt | TextPrompt | None = None,
+    requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    """
+    Process thinker outputs to create token2wav inputs.
+
+    Workflow:
+    1. Extract generated token IDs from thinker output
+    2. Separate audio tokens from text tokens
+    3. Package audio tokens for token2wav stage
+
+    Step-Audio2 token ranges:
+    - Text tokens: 0 - 151688
+    - Audio tokens: 151696 - 158257 (vocab size 6562)
+
+    Args:
+        source_outputs: Resolved list of upstream (thinker) engine outputs
+        prompt: Original prompt data (unused for Token2Wav in default-voice mode)
+        requires_multimodal_data: Whether multimodal data is required
+
+    Returns:
+        List of OmniTokensPrompt for token2wav stage
+
+    """
+    if not source_outputs:
+        raise ValueError("source_outputs cannot be empty")
+
+    thinker_outputs = source_outputs
+    token2wav_inputs = []
+
+    # Token configuration uses the fixed Step-Audio2 model constants.
+    audio_start = DEFAULT_TOKEN_CONFIG.audio_start
+    audio_eos = DEFAULT_TOKEN_CONFIG.audio_eos  # Relative to audio start
+
+    # Process each thinker output
+    for i, thinker_output in enumerate(thinker_outputs):
+        output = thinker_output.outputs[0]
+
+        # Only look at decode (generated) tokens — the prompt may contain
+        # historical audio tokens from prior conversation turns.
+        gen_ids = output.token_ids
+
+        # Convert to CPU list if needed (for multi-process/multi-GPU setup)
+        if isinstance(gen_ids, torch.Tensor):
+            gen_ids = gen_ids.cpu().tolist()
+
+        # Separate audio tokens from text tokens
+        # Audio tokens are >= audio_start
+        audio_tokens = [
+            tid - audio_start  # Convert to 0-based for Token2Wav
+            for tid in gen_ids
+            if tid >= audio_start
+        ]
+
+        # Remove padding tokens (anything >= EOS)
+        audio_tokens = [t for t in audio_tokens if t < audio_eos]
+
+        if not audio_tokens:
+            # No audio tokens generated, skip Token2Wav for this request
+            logger.info(f"Request {i}: No audio tokens generated, skipping Token2Wav stage")
+            continue
+
+        logger.debug("Creating Token2Wav input with %d audio tokens", len(audio_tokens))
+
+        token2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=audio_tokens,  # Pass original tokens (no encoding)
+                multi_modal_data=None,
+                additional_information=None,  # Not used due to framework limitation
+            )
+        )
+
+    return token2wav_inputs