[Benchmark] Benchmark Running Samples for Qwen3 Omni and Dataset Preparation (#212)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>

benchmarks/README.md

@@ -0,0 +1,42 @@
+# Benchmarks Overview and Architecture
+
+This document explains the benchmark architecture across all benchmark assets in this repo. It describes what we measure, and where to find or plug in new scenarios. Per-task details remain in subfolder READMEs (e.g., `benchmarks/<model>/README.md`).
+
+## Scope and goals
+- Establish repeatable latency/throughput measurements for multimodal LLM pipelines.
+- Provide both HF Transformers (offline) and vLLM-Omni (multi-stage/pipeline) baselines.
+- Make it easy to plug in new datasets and models with minimal changes to the runner scripts.
+
+## Dataset and inputs
+- Default example: SeedTTS top-100 prompts (`benchmarks/build_dataset/top100.txt`) via `benchmarks/build_dataset/`.
+- Extensible: drop in new prompt files or modality-aligned payloads; keep the expected format for the consuming scripts (e.g., one prompt per line).
+- If you add a new dataset, document it under `benchmarks/<model>/README.md` and point scripts to your data path.
+
+## Directory layout
+- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
+- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
+- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
+
+## Reference workflows
+- **HF Transformers (offline, single process)**  
+  Script (example): `benchmarks/<model>/transformers/eval_qwen3_moe_omni_transformers.sh`  
+  Outputs: `benchmark_results/perf_stats.json`, `benchmark_results/results.json`, `benchmark_results/audio/` (if audio is produced).
+
+- **vLLM-Omni end-to-end pipeline**  
+  Script (example): `benchmarks/<model>/vllm_omni/eval_qwen3_moe_omni.sh`  
+  Outputs: `vllm_omni/logs/*.stats.jsonl` (per-stage/overall latency & TPS), `vllm_omni/logs/stage*.log`, `vllm_omni/outputs/` (text/audio artifacts).
+
+- **Adding a new task/model**  
+  1) Create `benchmarks/<model>/transformers/` and/or `benchmarks/<model>/vllm_omni/` with scripts referencing your model and dataset.  
+  2) Add a task README describing dataset, configs, and expected outputs.  
+  3) Keep the output/log structure similar for easy comparison (perf_stats/results/audio or text outputs; stats.jsonl/logs for pipeline).
+
+## Metrics to watch
+- **Throughput**: `overall_tps`, `*_tps_avg` per stage.
+- **Latency distribution**: look for long tails in `*.stats.jsonl`.
+- **Quality/completeness**: missing outputs or errors in stage logs indicate pipeline failures or misconfigurations.
+
+## Troubleshooting
+- Verify GPU/driver/FlashAttention2 requirements for your chosen model/config.
+- Ensure network access for dataset/model downloads (Google Drive, Hugging Face, etc.).
+- If outputs are missing or slow, inspect per-stage logs and `*.stats.jsonl` for errors, stragglers, or contention.

benchmarks/build_dataset/download_process_data_seedtts.md

@@ -0,0 +1,82 @@
+# Benchmark Dataset Preparation Guide
+
+This guide describes how to download and prepare the SeedTTS test dataset for benchmarking Qwen-Omni models.
+
+## Prerequisites
+
+- Python 3.8+
+- `gdown` for downloading from Google Drive
+- Access to the benchmark scripts
+
+## Steps
+
+### 1. Navigate to the Dataset Directory
+
+```bash
+cd benchmarks/build_dataset
+```
+
+### 2. Install Dependencies
+
+```bash
+pip install gdown
+```
+
+### 3. Download the SeedTTS Test Dataset
+
+Download the dataset from Google Drive:
+
+```bash
+gdown --id 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
+```
+
+### 4. Extract the Dataset
+
+```bash
+tar -xf seedtts_testset.tar
+```
+
+### 5. Prepare the Metadata File
+
+Copy the English metadata file to the working directory:
+
+```bash
+cp seedtts_testset/en/meta.lst meta.lst
+```
+
+### 6. Extract Prompts
+
+Extract the first N prompts from the metadata file:
+
+```bash
+# Extract top 100 prompts (adjust -n for different amounts)
+python extract_tts_prompts.py -i meta.lst -o top100.txt -n 100
+```
+
+**Options:**
+- `-i, --input`: Input metadata file (default: `meta.lst`)
+- `-o, --output`: Output prompts file (default: `prompts.txt`)
+- `-n, --num_lines`: Number of prompts to extract (required)
+
+### 7. Clean Up (Optional)
+
+Remove temporary files to save disk space:
+
+```bash
+rm -rf seedtts_testset
+rm seedtts_testset.tar
+rm meta.lst
+```
+
+## Quick Start (All-in-One)
+
+```bash
+# Full setup and benchmark
+cd benchmarks/build_dataset
+pip install gdown
+gdown --id 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
+tar -xf seedtts_testset.tar
+cp seedtts_testset/en/meta.lst meta.lst
+python extract_tts_prompts.py -i meta.lst -o top100.txt -n 100
+rm -rf seedtts_testset seedtts_testset.tar meta.lst
+```

benchmarks/build_dataset/extract_tts_prompts.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+"""
+Extract prompts from meta.lst and save them to a txt file.
+
+Each line in meta.lst has the format:
+ID|prompt_text|audio_path|target_text
+
+This script extracts the prompt_text (second field) from the first N lines.
+"""
+
+import argparse
+from pathlib import Path
+
+
+def extract_prompts(input_file: str, output_file: str, num_lines: int) -> None:
+    """
+    Extract prompts from meta.lst and save to output file.
+
+    Args:
+        input_file: Path to the meta.lst file
+        output_file: Path to the output txt file
+        num_lines: Number of lines to process
+    """
+    prompts = []
+
+    with open(input_file, encoding="utf-8") as f:
+        for i, line in enumerate(f):
+            if i >= num_lines:
+                break
+
+            line = line.strip()
+            if not line:  # Skip empty lines
+                continue
+
+            parts = line.split("|")
+            if len(parts) >= 2:
+                prompt = parts[1]  # The prompt is the second field
+                prompts.append(prompt)
+
+    # Write prompts to output file
+    with open(output_file, "w", encoding="utf-8") as f:
+        for prompt in prompts:
+            f.write(prompt + "\n")
+
+    # Print result stats
+    print(f"Extracted {len(prompts)} prompts from first {num_lines} lines")
+    print(f"Saved to: {output_file}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Extract prompts from meta.lst file")
+    parser.add_argument(
+        "-i", "--input", type=str, default="meta.lst", help="Input meta.lst file path (default: meta.lst)"
+    )
+    parser.add_argument(
+        "-o", "--output", type=str, default="prompts.txt", help="Output txt file path (default: prompts.txt)"
+    )
+    parser.add_argument(
+        "-n", "--num_lines", type=int, required=True, help="Number of lines to extract from the beginning"
+    )
+
+    args = parser.parse_args()
+
+    # Check if input file exists
+    if not Path(args.input).exists():
+        print(f"Error: Input file '{args.input}' not found")
+        return
+
+    extract_prompts(args.input, args.output, args.num_lines)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/qwen3-omni/README.md

@@ -0,0 +1,86 @@
+# Benchmarks Guide
+
+This README explains how to (1) prepare benchmark datasets and (2) run the provided Qwen3-Omni benchmarks.
+
+## 1) Prepare the dataset (SeedTTS top100)
+
+```bash
+cd benchmarks/build_dataset
+pip install gdown
+
+# Download SeedTTS test set from Google Drive
+gdown --id 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
+
+# Extract
+tar -xf seedtts_testset.tar
+
+# Copy metadata and extract top-100 prompts
+cp seedtts_testset/en/meta.lst meta.lst
+python extract_prompts.py -i meta.lst -o top100.txt -n 100
+
+# (Optional) clean up to save space
+rm -rf seedtts_testset seedtts_testset.tar meta.lst
+```
+
+Artifacts:
+- `benchmarks/build_dataset/top100.txt` — 100 text prompts (one per line).
+
+## 2) Run benchmarks
+
+All commands assume repo root (`vllm-omni`).
+
+### A. Transformers benchmark (offline, HF Transformers)
+
+```
+bash benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh
+```
+
+What it does:
+- Runs `qwen3_omni_moe_transformers.py` over `top100.txt` with `--num_prompts 100`.
+- Outputs to `benchmarks/qwen3-omni/transformers/benchmark_results/`:
+  - `perf_stats.json` — aggregated & per-prompt TPS/latency (thinker/talker/code2wav/overall).
+  - `results.json` — per-prompt outputs and audio paths.
+  - `audio/` — ~100 generated `.wav` files.
+
+Key checks:
+- `overall_tps` and `*_tps_avg` should be non-zero and reasonably stable.
+- Investigate any 0/NaN or unusually low TPS / long-tail latency.
+
+### B. vLLM Omni end-to-end benchmark (pipeline)
+
+```
+bash benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh
+```
+
+What it does:
+- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--enable-stats`.
+- Uses `benchmarks/build_dataset/top100.txt` and writes to:
+  - Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
+    - `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
+    - `omni_llm_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
+    - `omni_llm_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
+  - Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
+
+Key checks:
+- Overall stats: end-to-end latency/TPS should be reasonable.
+- Orchestrator stats: per-stage latency should be stable; investigate long tails.
+- Stage logs: ensure no errors and no unusually slow stages.
+
+
+## Performance snapshot
+
+The chart below summarizes our measured Qwen3-Omni MoE end-to-end benchmark, comparing vLLM-Omni against HF Transformers. It shows the overall throughput advantage for vLLM-Omni. These are actual experiment results—please refer to this performance when evaluating or reproducing the benchmark.
+
+![vLLM-Omni vs HF](./vllm-omni-vs-hf.png)
+
+## Directory layout
+- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
+- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
+- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
+- `benchmarks/<model>/vllm-omni-vs-hf.png` — current performance snapshot (overall throughput comparison).
+- `benchmarks/<model>/transformers/` — HF Transformers benchmarks (offline reference).
+
+## Troubleshooting
+- Make sure GPU/driver/FlashAttention2 requirements are met for the chosen model.
+- If downloads fail, confirm network access to Google Drive (`gdown`) and Hugging Face.
+- If audio files are missing, check for errors in stage logs or model generation.***

benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+# Qwen3-Omni Transformers Benchmark Evaluation Script
+# This script must be run from the vllm-omni root directory
+
+# Get the directory where this script is located
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+# Navigate to vllm-omni root directory (4 levels up from script location)
+VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
+cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
+
+echo "Working directory: $(pwd)"
+# Verify we're in the correct directory and run benchmark
+if [[ ! -f "benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py" ]]; then
+    echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
+else
+    cd benchmarks/qwen3-omni/transformers
+
+    python qwen3_omni_moe_transformers.py --prompts_file ../../build_dataset/top100.txt --num_prompts 100
+
+    echo "Logs and outputs are saved to $(pwd)/benchmark_results:"
+    echo "  - perf_stats.json    Aggregated/per-prompt TPS and latency (thinker/talker/code2wav/overall)"
+    echo "  - results.json       Per-prompt outputs and audio paths"
+    echo "  - audio/             Generated wav files, there should be 100 wav file generated"
+    echo "Key checks: overall_tps and *_tps_avg should be non-zero and stable; investigate 0/NaN or unusually low TPS/long-tail latency."
+fi

benchmarks/qwen3-omni/transformers/qwen3_omni_moe_model.py

@@ -0,0 +1,265 @@
+import time
+
+import torch
+from transformers import Qwen3OmniMoeForConditionalGeneration
+
+
+class Qwen3OmniMoeForConditionalGenerationWithLogging(Qwen3OmniMoeForConditionalGeneration):
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor | None = None,
+        speaker: str = "Ethan",
+        use_audio_in_video: bool = False,
+        return_audio: bool | None = None,
+        thinker_max_new_tokens: int = 1024,
+        thinker_eos_token_id: int = 151645,
+        talker_max_new_tokens: int = 4096,
+        talker_do_sample: bool = True,
+        talker_top_k: int = 50,
+        talker_top_p: float = 1.0,
+        talker_temperature: float = 0.9,
+        talker_repetition_penalty: float = 1.05,
+        **kwargs,
+    ):
+        total_t0 = time.time()
+        perf_stats = {
+            "thinker_tokens": 0,
+            "thinker_time_s": 0.0,
+            "thinker_tps": 0.0,
+            "talker_tokens": 0,
+            "talker_time_s": 0.0,
+            "talker_tps": 0.0,
+            "code2wav_tokens": 0,
+            "code2wav_time_s": 0.0,
+            "code2wav_tps": 0.0,
+            "total_tokens": 0,
+            "total_time_s": 0.0,
+            "total_tps": 0.0,
+        }
+        if return_audio and not self.has_talker:
+            raise ValueError(
+                "Cannot use talker when talker module not initialized. "
+                "Use `enable_talker` method or set enable_talker in config "
+                "to enable talker."
+            )
+        if return_audio is None:
+            return_audio = self.has_talker
+
+        shared_kwargs = {"use_audio_in_video": use_audio_in_video}
+        thinker_kwargs = {
+            "max_new_tokens": thinker_max_new_tokens,
+            "eos_token_id": thinker_eos_token_id,
+        }
+
+        talker_kwargs = {}
+        token2wav_kwargs = {}
+        if return_audio:
+            speaker_id = self.config.talker_config.speaker_id.get(speaker.lower())
+            if speaker_id is None:
+                raise NotImplementedError(f"Speaker {speaker} not implemented")
+            if input_ids.shape[0] != 1:
+                raise NotImplementedError("Qwen3-Omni currently does not support batched inference with audio output")
+            talker_suppressed_tokens = [
+                i
+                for i in range(
+                    self.config.talker_config.text_config.vocab_size - 1024,
+                    self.config.talker_config.text_config.vocab_size,
+                )
+                if i != self.config.talker_config.codec_eos_token_id
+            ]  # Suppress additional special tokens, should not be predicted
+            talker_kwargs = {
+                "max_new_tokens": talker_max_new_tokens,
+                "do_sample": talker_do_sample,
+                "top_k": talker_top_k,
+                "top_p": talker_top_p,
+                "temperature": talker_temperature,
+                "eos_token_id": self.config.talker_config.codec_eos_token_id,
+                "repetition_penalty": talker_repetition_penalty,
+                "suppress_tokens": talker_suppressed_tokens,
+                "output_hidden_states": True,
+                "return_dict_in_generate": True,
+            }
+            token2wav_kwargs = {}
+
+        for key, value in kwargs.items():
+            if key.startswith("thinker_"):
+                thinker_kwargs[key[len("thinker_") :]] = value
+            elif key.startswith("talker_"):
+                talker_kwargs[key[len("talker_") :]] = value
+            elif key.startswith("token2wav_"):
+                token2wav_kwargs[key[len("token2wav_") :]] = value
+            # Process special input values
+            elif key == "feature_attention_mask":
+                thinker_kwargs[key] = value
+                talker_kwargs["audio_feature_lengths"] = torch.sum(value, dim=1)
+            elif key in ("input_features", "attention_mask"):
+                thinker_kwargs[key] = value
+            # Put other key to shared kwargs
+            else:
+                shared_kwargs[key] = value
+
+        # Merge kwargs
+        for key, value in shared_kwargs.items():
+            if key not in thinker_kwargs:
+                thinker_kwargs[key] = value
+            if key not in talker_kwargs and key in ["image_grid_thw", "video_grid_thw", "video_second_per_grid"]:
+                talker_kwargs[key] = value
+            if key not in token2wav_kwargs:
+                token2wav_kwargs[key] = value
+
+        # 1. Generate from thinker module
+        generate_audio = return_audio and self.has_talker
+        if generate_audio:
+            thinker_kwargs["output_hidden_states"] = True
+            thinker_kwargs["return_dict_in_generate"] = True
+
+        t0 = time.time()
+        thinker_result = self.thinker.generate(input_ids=input_ids, **thinker_kwargs)
+        t1 = time.time()
+        perf_stats["thinker_time_s"] = max(0.0, t1 - t0)
+        try:
+            prompt_len = int(input_ids.shape[1]) if input_ids is not None else 0
+            total_len = int(thinker_result.sequences.shape[-1])
+            thinker_out_len = max(0, total_len - prompt_len)
+        except Exception:
+            thinker_out_len = 0
+        perf_stats["thinker_tokens"] = thinker_out_len
+        perf_stats["thinker_tps"] = (
+            (thinker_out_len / perf_stats["thinker_time_s"]) if perf_stats["thinker_time_s"] > 0 else 0.0
+        )
+
+        if not generate_audio:
+            perf_stats["total_tokens"] = perf_stats["thinker_tokens"]
+            perf_stats["total_time_s"] = time.time() - total_t0
+            perf_stats["total_tps"] = (
+                (perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
+            )
+            # attach stats to self
+            setattr(self, "_perf_stats_last", perf_stats)
+            if not hasattr(self, "_perf_stats_history"):
+                setattr(self, "_perf_stats_history", [])
+            self._perf_stats_history.append(perf_stats)
+            return thinker_result, None
+
+        # 2. Prepare talker input
+        thinker_embed = torch.cat([hidden_states[0] for hidden_states in thinker_result.hidden_states], dim=1).to(
+            self.talker.device
+        )  # [1 t d]
+        thinker_hidden = torch.cat(
+            [
+                hidden_states[self.config.talker_config.accept_hidden_layer]
+                for hidden_states in thinker_result.hidden_states
+            ],
+            dim=1,
+        ).to(self.talker.device)  # [1 t d]
+
+        im_start_indexes = torch.cat(
+            (
+                torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(),
+                torch.tensor([thinker_result.sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype),
+            ),
+            dim=-1,
+        ).to(self.talker.device)  # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here.
+        multimodal_mask = (
+            (thinker_result.sequences == self.config.thinker_config.audio_token_id) |
+            (thinker_result.sequences == self.config.thinker_config.image_token_id) |
+            (thinker_result.sequences == self.config.thinker_config.video_token_id)
+        ).to(self.talker.device)  # [1 t] # fmt: skip
+
+        talker_special_tokens = torch.tensor(
+            [[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
+            device=self.thinker.device,
+            dtype=input_ids.dtype,
+        )
+        tts_bos_embed, tts_eos_embed, tts_pad_embed = (
+            self.talker.text_projection(self.thinker.get_input_embeddings()(talker_special_tokens))
+            .to(self.talker.device)
+            .chunk(3, dim=1)
+        )  # 3 * [1 1 d]
+
+        talker_input_embeds = []  # [1 t d]
+        talker_input_ids = []
+        # For every chatml parts
+        for i in range(len(im_start_indexes) - 1):
+            im_start_index = im_start_indexes[i]
+            segment_end_index = im_start_indexes[i + 1]
+            role_token = input_ids[0][im_start_index + 1]
+            # Talker should ignore thinker system prompt
+            if role_token == self.config.system_token_id:
+                continue
+            # Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs
+            elif role_token == self.config.user_token_id:
+                talker_user_part = self._get_talker_user_parts(
+                    im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
+                )
+                talker_input_embeds.append(talker_user_part)
+                talker_input_ids.append(thinker_result.sequences[:, im_start_index:segment_end_index])
+            # Take assistant output (for now)
+            elif role_token == self.config.assistant_token_id and i == len(im_start_indexes) - 2:
+                talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts(
+                    im_start_index,
+                    segment_end_index,
+                    speaker_id,
+                    thinker_embed,
+                    tts_pad_embed,
+                    tts_bos_embed,
+                    tts_eos_embed,
+                )
+                talker_input_embeds.append(talker_assistant_embeds)
+                talker_input_ids.append(talker_assistant_ids)
+            # History assistant output (ignore for now)
+            elif role_token == self.config.assistant_token_id and i != len(im_start_indexes) - 2:
+                continue
+            else:
+                raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)")
+        talker_input_embed = torch.cat([embed.to(self.talker.device) for embed in talker_input_embeds], dim=1)
+        talker_input_id = torch.cat([embed.to(self.talker.device) for embed in talker_input_ids], dim=1)
+        t2 = time.time()
+        talker_result = self.talker.generate(
+            inputs_embeds=talker_input_embed,
+            trailing_text_hidden=trailing_text_hidden,
+            tts_pad_embed=tts_pad_embed,
+            talker_input_ids=talker_input_id,  # Not use input_ids to prevent repetition penalty out of bound
+            **talker_kwargs,
+        )
+        t3 = time.time()
+        perf_stats["talker_time_s"] = max(0.0, t3 - t2)
+        talker_codes = (
+            torch.stack([hid[-1] for hid in talker_result.hidden_states if hid[-1] is not None], dim=1)
+            .transpose(1, 2)
+            .to(self.code2wav.device)
+        )
+        try:
+            # codes shape: (B, num_quantizers, T). We log T as token length.
+            perf_stats["talker_tokens"] = int(talker_codes.shape[-1])
+        except Exception:
+            perf_stats["talker_tokens"] = 0
+        perf_stats["talker_tps"] = (
+            (perf_stats["talker_tokens"] / perf_stats["talker_time_s"]) if perf_stats["talker_time_s"] > 0 else 0.0
+        )
+        t4 = time.time()
+        talker_wavs = self.code2wav.chunked_decode(talker_codes, chunk_size=300, left_context_size=25).float()
+        t5 = time.time()
+        perf_stats["code2wav_time_s"] = max(0.0, t5 - t4)
+        perf_stats["code2wav_tokens"] = perf_stats["talker_tokens"]  # same T, not times 16
+        perf_stats["code2wav_tps"] = (
+            (perf_stats["code2wav_tokens"] / perf_stats["code2wav_time_s"])
+            if perf_stats["code2wav_time_s"] > 0
+            else 0.0
+        )
+        perf_stats["total_tokens"] = perf_stats["thinker_tokens"] + perf_stats["talker_tokens"]
+        perf_stats["total_time_s"] = time.time() - total_t0
+        perf_stats["total_tps"] = (
+            (perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
+        )
+        setattr(self, "_perf_stats_last", perf_stats)
+        if not hasattr(self, "_perf_stats_history"):
+            setattr(self, "_perf_stats_history", [])
+        self._perf_stats_history.append(perf_stats)
+        return thinker_result, talker_wavs.float()
+
+
+__all__ = [
+    "Qwen3OmniMoeForConditionalGenerationWithLogging",
+]

benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py

@@ -0,0 +1,275 @@
+import argparse
+import json
+import os
+
+import soundfile as sf
+from qwen3_omni_moe_model import Qwen3OmniMoeForConditionalGenerationWithLogging
+from qwen_omni_utils import process_mm_info
+from tqdm import tqdm
+from transformers import Qwen3OmniMoeProcessor
+
+MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+# MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
+
+
+def load_prompts(prompts_file: str) -> list[str]:
+    """Load prompts from a text file, one prompt per line."""
+    prompts = []
+    with open(prompts_file, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                prompts.append(line)
+    return prompts
+
+
+def run_benchmark(
+    model,
+    processor,
+    prompts: list[str],
+    output_dir: str = "benchmark_results",
+    speaker: str = "Ethan",
+    use_audio_in_video: bool = True,
+):
+    """
+    Run benchmark on a list of prompts and collect performance stats.
+
+    Args:
+        model: The Qwen3OmniMoe model
+        processor: The Qwen3OmniMoe processor
+        prompts: List of text prompts to process
+        output_dir: Directory to save results
+        speaker: Speaker voice for audio output
+        use_audio_in_video: Whether to use audio in video
+
+    Returns:
+        tuple: (aggregated_stats, results, audio_outputs)
+            - aggregated_stats: dict with aggregated performance statistics
+            - results: list of dicts with per-prompt results
+            - audio_outputs: list of audio tensors/arrays (or None if no audio)
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    audio_dir = os.path.join(output_dir, "audio")
+    os.makedirs(audio_dir, exist_ok=True)
+
+    all_stats = []
+    results = []
+    audio_outputs = []
+
+    for idx, prompt in enumerate(tqdm(prompts, desc="Processing prompts")):
+        conversation = [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": prompt}],
+            },
+        ]
+
+        # Preparation for inference
+        text = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
+        audios, images, videos = process_mm_info(conversation, use_audio_in_video=use_audio_in_video)
+        inputs = processor(
+            text=text,
+            audio=audios,
+            images=images,
+            videos=videos,
+            return_tensors="pt",
+            padding=True,
+            use_audio_in_video=use_audio_in_video,
+        )
+        inputs = inputs.to(model.device).to(model.dtype)
+
+        # Inference: Generation of the output text and audio
+        text_ids, audio = model.generate(
+            **inputs, speaker=speaker, thinker_return_dict_in_generate=True, use_audio_in_video=use_audio_in_video
+        )
+
+        # Decode output text
+        output_text = processor.batch_decode(
+            text_ids.sequences[:, inputs["input_ids"].shape[1] :],
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False,
+        )[0]
+
+        # Collect performance stats
+        perf_stats = None
+        if hasattr(model, "_perf_stats_last"):
+            perf_stats = model._perf_stats_last.copy()
+            perf_stats["prompt_idx"] = idx
+            perf_stats["prompt"] = prompt
+            all_stats.append(perf_stats)
+
+        # Save audio and collect audio output
+        audio_path = None
+        audio_data = None
+        if audio is not None:
+            audio_data = audio.reshape(-1).detach().cpu().numpy()
+            audio_path = os.path.join(audio_dir, f"output_{idx:04d}.wav")
+            sf.write(
+                audio_path,
+                audio_data,
+                samplerate=24000,
+            )
+            audio_outputs.append(audio_data)
+        else:
+            audio_outputs.append(None)
+
+        # Save result
+        result = {
+            "idx": idx,
+            "prompt": prompt,
+            "output": output_text,
+            "audio_path": audio_path,
+            "perf_stats": perf_stats,
+        }
+        results.append(result)
+
+    # Aggregate statistics
+    aggregated_stats = aggregate_stats(all_stats)
+
+    # Save all results
+    results_path = os.path.join(output_dir, "results.json")
+    with open(results_path, "w", encoding="utf-8") as f:
+        json.dump(results, f, ensure_ascii=False, indent=2)
+
+    # Save aggregated stats
+    stats_path = os.path.join(output_dir, "perf_stats.json")
+    with open(stats_path, "w", encoding="utf-8") as f:
+        json.dump({"aggregated": aggregated_stats, "per_prompt": all_stats}, f, ensure_ascii=False, indent=2)
+
+    # Count saved audio files
+    num_audio_saved = sum(1 for a in audio_outputs if a is not None)
+    print(f"\nSaved {num_audio_saved} audio files to {audio_dir}/")
+
+    return aggregated_stats, results, audio_outputs
+
+
+def aggregate_stats(all_stats: list[dict]) -> dict:
+    """Aggregate performance statistics from multiple runs."""
+    if not all_stats:
+        return {}
+
+    keys = [
+        "thinker_tokens",
+        "thinker_time_s",
+        "thinker_tps",
+        "talker_tokens",
+        "talker_time_s",
+        "talker_tps",
+        "code2wav_tokens",
+        "code2wav_time_s",
+        "code2wav_tps",
+        "total_tokens",
+        "total_time_s",
+        "total_tps",
+    ]
+
+    aggregated = {
+        "num_samples": len(all_stats),
+    }
+
+    for key in keys:
+        values = [s.get(key, 0) for s in all_stats if key in s]
+        if values:
+            aggregated[f"{key}_sum"] = sum(values)
+            aggregated[f"{key}_avg"] = sum(values) / len(values)
+            aggregated[f"{key}_min"] = min(values)
+            aggregated[f"{key}_max"] = max(values)
+
+    # Calculate overall throughput
+    total_tokens = aggregated.get("total_tokens_sum", 0)
+    total_time = aggregated.get("total_time_s_sum", 0)
+    if total_time > 0:
+        aggregated["overall_tps"] = total_tokens / total_time
+
+    return aggregated
+
+
+def print_stats(stats: dict):
+    """Print performance statistics in a formatted way."""
+    print("\n" + "=" * 60)
+    print("Performance Statistics Summary")
+    print("=" * 60)
+
+    print(f"\nNumber of samples: {stats.get('num_samples', 0)}")
+
+    print("\n--- Thinker ---")
+    print(f"  Total tokens:  {stats.get('thinker_tokens_sum', 0):.0f}")
+    print(f"  Total time:    {stats.get('thinker_time_s_sum', 0):.2f}s")
+    print(f"  Avg TPS:       {stats.get('thinker_tps_avg', 0):.2f}")
+    print(f"  Min TPS:       {stats.get('thinker_tps_min', 0):.2f}")
+    print(f"  Max TPS:       {stats.get('thinker_tps_max', 0):.2f}")
+
+    print("\n--- Talker ---")
+    print(f"  Total tokens:  {stats.get('talker_tokens_sum', 0):.0f}")
+    print(f"  Total time:    {stats.get('talker_time_s_sum', 0):.2f}s")
+    print(f"  Avg TPS:       {stats.get('talker_tps_avg', 0):.2f}")
+    print(f"  Min TPS:       {stats.get('talker_tps_min', 0):.2f}")
+    print(f"  Max TPS:       {stats.get('talker_tps_max', 0):.2f}")
+
+    print("\n--- Code2Wav ---")
+    print(f"  Total tokens:  {stats.get('code2wav_tokens_sum', 0):.0f}")
+    print(f"  Total time:    {stats.get('code2wav_time_s_sum', 0):.2f}s")
+    print(f"  Avg TPS:       {stats.get('code2wav_tps_avg', 0):.2f}")
+    print(f"  Min TPS:       {stats.get('code2wav_tps_min', 0):.2f}")
+    print(f"  Max TPS:       {stats.get('code2wav_tps_max', 0):.2f}")
+
+    print("\n--- Overall ---")
+    print(f"  Total tokens:  {stats.get('total_tokens_sum', 0):.0f}")
+    print(f"  Total time:    {stats.get('total_time_s_sum', 0):.2f}s")
+    print(f"  Overall TPS:   {stats.get('overall_tps', 0):.2f}")
+    print(f"  Avg TPS:       {stats.get('total_tps_avg', 0):.2f}")
+    print(f"  Min TPS:       {stats.get('total_tps_min', 0):.2f}")
+    print(f"  Max TPS:       {stats.get('total_tps_max', 0):.2f}")
+
+    print("=" * 60 + "\n")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Qwen3-Omni Benchmark Script")
+    parser.add_argument(
+        "--prompts_file",
+        type=str,
+        default="benchmark/build_dataset/top100.txt",
+        help="Path to the prompts file (one prompt per line)",
+    )
+    parser.add_argument(
+        "--output_dir", type=str, default="benchmark_results", help="Directory to save benchmark results"
+    )
+    parser.add_argument("--model_path", type=str, default=MODEL_PATH, help="Path to the model")
+    parser.add_argument("--speaker", type=str, default="Ethan", help="Speaker voice for audio output")
+    parser.add_argument("--num_prompts", type=int, default=None, help="Number of prompts to process (default: all)")
+    args = parser.parse_args()
+
+    # Load model and processor
+    print(f"Loading model from {args.model_path}...")
+    model = Qwen3OmniMoeForConditionalGenerationWithLogging.from_pretrained(
+        args.model_path,
+        dtype="auto",
+        device_map="auto",
+        attn_implementation="flash_attention_2",
+    )
+    processor = Qwen3OmniMoeProcessor.from_pretrained(args.model_path)
+
+    # Benchmark mode
+    print(f"Loading prompts from {args.prompts_file}...")
+    prompts = load_prompts(args.prompts_file)
+
+    if args.num_prompts:
+        prompts = prompts[: args.num_prompts]
+
+    print(f"Running benchmark on {len(prompts)} prompts...")
+
+    aggregated_stats, results, audio_outputs = run_benchmark(
+        model=model,
+        processor=processor,
+        prompts=prompts,
+        output_dir=args.output_dir,
+        speaker=args.speaker,
+    )
+
+    print_stats(aggregated_stats)
+    print(f"\nResults saved to {args.output_dir}/")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+# Qwen3-Omni Benchmark Evaluation Script
+# This script must be run from the vllm-omni root directory
+
+# Get the directory where this script is located
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+# Navigate to vllm-omni root directory (4 levels up from script location)
+VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
+cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
+
+echo "Working directory: $(pwd)"
+
+# Verify we're in the correct directory and run benchmark
+if [[ ! -d "benchmarks/qwen3-omni/vllm_omni" ]]; then
+    echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
+else
+    log_dir=benchmarks/qwen3-omni/vllm_omni/logs
+    outputs_dir=benchmarks/qwen3-omni/vllm_omni/outputs
+    end2end_script_path=examples/offline_inference/qwen3_omni/end2end.py
+    build_dataset_path=benchmarks/build_dataset/top100.txt
+
+    python $end2end_script_path --output-wav $outputs_dir \
+                      --query-type text \
+                      --txt-prompts $build_dataset_path \
+                      --enable-stats \
+                      --log-dir $log_dir
+    echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
+    echo "  - omni_llm_pipeline_text                       run dir/base name"
+    echo "  - omni_llm_pipeline_text.orchestrator.stats.jsonl  orchestrator-stage latency stats"
+    echo "  - omni_llm_pipeline_text.overall.stats.jsonl       overall latency/TPS stats"
+    echo "  - omni_llm_pipeline_text.stage0.log                per-stage detailed logs"
+    echo "  - omni_llm_pipeline_text.stage1.log"
+    echo "  - omni_llm_pipeline_text.stage2.log"
+    echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
+    echo "  - outputs/             Generated txt and wav files, there should be 100 text and wav files generated respectively"
+fi

examples/offline_inference/qwen3_omni/end2end.py

@@ -170,9 +170,16 @@ def main(args):
     else:
         query_result = query_func()
 
+    if not args.enable_stats:
+        log_file = None
+    else:
+        log_file = os.path.join(args.log_dir, f"omni_llm_pipeline_{args.query_type}")
+
     omni_llm = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
+        log_file=log_file,
+        log_stats=args.enable_stats,
     )
 
     thinker_sampling_params = SamplingParams(
@@ -364,6 +371,12 @@ def parse_args():
         default=16000,
         help="Sampling rate for audio loading (default: 16000).",
     )
+    parser.add_argument(
+        "--log-dir",
+        type=str,
+        default="logs",
+        help="Log directory (default: logs).",
+    )
     parser.add_argument(
         "--modalities",
         type=str,