[Qwen3TTS][Feat] Code2Wav batched decoding

[JuanPZuluaga]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

In this simple PR, we add support for Code2Wav batched decoding. Note that in the stage configs for the offline benchmarking: stage-0 BS >= stage-1 BS. Otherwise, stage-1 wouldn't reach the number of requests that are specified by the batch size (probably in async scheduling wouldn't be a problem?). Other mods:

• added batched decoding in benchmark (can be removed if we don't want it, I can also refactor it a bit to support larger batches).

Test Plan

Test Result

1. I added a benchmark prompts txt.
2. batch size for stage 0 is always 4.
3. run evaluation with bs=1 and bs=4 for stage-1
4. compared results.

Baseline with 12 samples:

# First batch
INFO 02-21 22:29:02 [stats.py:500] [Overall Summary]
INFO 02-21 22:29:02 [stats.py:500] +-----------------------------+-----------+
INFO 02-21 22:29:02 [stats.py:500] | Field                       |     Value |
INFO 02-21 22:29:02 [stats.py:500] +-----------------------------+-----------+
INFO 02-21 22:29:02 [stats.py:500] | e2e_requests                |         4 |
INFO 02-21 22:29:02 [stats.py:500] | e2e_wall_time_ms            | 5,036.018 |
INFO 02-21 22:29:02 [stats.py:500] | e2e_total_tokens            |       321 |
INFO 02-21 22:29:02 [stats.py:500] | e2e_avg_time_per_request_ms | 1,259.005 |
INFO 02-21 22:29:02 [stats.py:500] | e2e_avg_tokens_per_s        |    63.741 |
INFO 02-21 22:29:02 [stats.py:500] | e2e_stage_0_wall_time_ms    | 4,406.654 |
INFO 02-21 22:29:02 [stats.py:500] | e2e_stage_1_wall_time_ms    |   630.117 |
INFO 02-21 22:29:02 [stats.py:500] +-----------------------------+-----------+

# Second batch
INFO 02-21 22:29:08 [stats.py:500] [Overall Summary]
INFO 02-21 22:29:08 [stats.py:500] +-----------------------------+-----------+
INFO 02-21 22:29:08 [stats.py:500] | Field                       |     Value |
INFO 02-21 22:29:08 [stats.py:500] +-----------------------------+-----------+
INFO 02-21 22:29:08 [stats.py:500] | e2e_requests                |         4 |
INFO 02-21 22:29:08 [stats.py:500] | e2e_wall_time_ms            | 5,340.400 |
INFO 02-21 22:29:08 [stats.py:500] | e2e_total_tokens            |       424 |
INFO 02-21 22:29:08 [stats.py:500] | e2e_avg_time_per_request_ms | 1,335.100 |
INFO 02-21 22:29:08 [stats.py:500] | e2e_avg_tokens_per_s        |    79.395 |
INFO 02-21 22:29:08 [stats.py:500] | e2e_stage_0_wall_time_ms    | 5,067.353 |
INFO 02-21 22:29:08 [stats.py:500] | e2e_stage_1_wall_time_ms    |   274.454 |
INFO 02-21 22:29:08 [stats.py:500] +-----------------------------+-----------+

# Third batch
INFO 02-21 22:29:12 [stats.py:500] [Overall Summary]
INFO 02-21 22:29:12 [stats.py:500] +-----------------------------+-----------+
INFO 02-21 22:29:12 [stats.py:500] | Field                       |     Value |
INFO 02-21 22:29:12 [stats.py:500] +-----------------------------+-----------+
INFO 02-21 22:29:12 [stats.py:500] | e2e_requests                |         4 |
INFO 02-21 22:29:12 [stats.py:500] | e2e_wall_time_ms            | 4,708.225 |
INFO 02-21 22:29:12 [stats.py:500] | e2e_total_tokens            |       408 |
INFO 02-21 22:29:12 [stats.py:500] | e2e_avg_time_per_request_ms | 1,177.056 |
INFO 02-21 22:29:12 [stats.py:500] | e2e_avg_tokens_per_s        |    86.657 |
INFO 02-21 22:29:12 [stats.py:500] | e2e_stage_0_wall_time_ms    | 4,435.627 |
INFO 02-21 22:29:12 [stats.py:500] | e2e_stage_1_wall_time_ms    |   273.613 |
INFO 02-21 22:29:12 [stats.py:500] +-----------------------------+-----------+

Stage 1 batching with the same 12 samples:

# First batch
INFO 02-21 22:31:51 [stats.py:500] [Overall Summary]
INFO 02-21 22:31:51 [stats.py:500] +-----------------------------+-----------+
INFO 02-21 22:31:51 [stats.py:500] | Field                       |     Value |
INFO 02-21 22:31:51 [stats.py:500] +-----------------------------+-----------+
INFO 02-21 22:31:51 [stats.py:500] | e2e_requests                |         4 |
INFO 02-21 22:31:51 [stats.py:500] | e2e_wall_time_ms            | 5,345.297 |
INFO 02-21 22:31:51 [stats.py:500] | e2e_total_tokens            |       333 |
INFO 02-21 22:31:51 [stats.py:500] | e2e_avg_time_per_request_ms | 1,336.324 |
INFO 02-21 22:31:51 [stats.py:500] | e2e_avg_tokens_per_s        |    62.298 |
INFO 02-21 22:31:51 [stats.py:500] | e2e_stage_0_wall_time_ms    | 4,767.987 |
INFO 02-21 22:31:51 [stats.py:500] | e2e_stage_1_wall_time_ms    |   578.176 |
INFO 02-21 22:31:51 [stats.py:500] +-----------------------------+-----------+

# Second batch
INFO 02-21 22:31:56 [stats.py:500] [Overall Summary]
INFO 02-21 22:31:56 [stats.py:500] +-----------------------------+-----------+
INFO 02-21 22:31:56 [stats.py:500] | Field                       |     Value |
INFO 02-21 22:31:56 [stats.py:500] +-----------------------------+-----------+
INFO 02-21 22:31:56 [stats.py:500] | e2e_requests                |         4 |
INFO 02-21 22:31:56 [stats.py:500] | e2e_wall_time_ms            | 5,120.655 |
INFO 02-21 22:31:56 [stats.py:500] | e2e_total_tokens            |       425 |
INFO 02-21 22:31:56 [stats.py:500] | e2e_avg_time_per_request_ms | 1,280.164 |
INFO 02-21 22:31:56 [stats.py:500] | e2e_avg_tokens_per_s        |    82.997 |
INFO 02-21 22:31:56 [stats.py:500] | e2e_stage_0_wall_time_ms    | 4,899.338 |
INFO 02-21 22:31:56 [stats.py:500] | e2e_stage_1_wall_time_ms    |   222.618 |
INFO 02-21 22:31:56 [stats.py:500] +-----------------------------+-----------+

# Third batch
INFO 02-21 22:32:01 [stats.py:500] [Overall Summary]
INFO 02-21 22:32:01 [stats.py:500] +-----------------------------+-----------+
INFO 02-21 22:32:01 [stats.py:500] | Field                       |     Value |
INFO 02-21 22:32:01 [stats.py:500] +-----------------------------+-----------+
INFO 02-21 22:32:01 [stats.py:500] | e2e_requests                |         4 |
INFO 02-21 22:32:01 [stats.py:500] | e2e_wall_time_ms            | 5,365.096 |
INFO 02-21 22:32:01 [stats.py:500] | e2e_total_tokens            |       432 |
INFO 02-21 22:32:01 [stats.py:500] | e2e_avg_time_per_request_ms | 1,341.274 |
INFO 02-21 22:32:01 [stats.py:500] | e2e_avg_tokens_per_s        |    80.520 |
INFO 02-21 22:32:01 [stats.py:500] | e2e_stage_0_wall_time_ms    | 5,139.006 |
INFO 02-21 22:32:01 [stats.py:500] | e2e_stage_1_wall_time_ms    |   227.250 |
INFO 02-21 22:32:01 [stats.py:500] +-----------------------------+-----------+

One can see the speed up in: e2e_stage_1_wall_time_ms between the batched and the baseline.

The benchmark results clearly show the improvement:

Baseline stage-1 times: 630ms, 274ms, 273ms
Batched stage-1 times: 578ms, 222ms, 227ms

The command to run the batched decoding:

python3 examples/offline_inference/qwen3_tts/end2end.py --query-type CustomVoice --txt-prompts examples/offline_inference/qwen3_tts/benchmark_prompts.txt --batch-size 4 --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml --output-dir benchmark_output_final --log-stats

(if one want to run the baseline, only needs to change the stage_configs).

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please providing the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please pasting the results comparison before and after, or e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b877279efd

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Split inputs by per-request lengths before batched decode

This split logic uses forward_context.ubatch_slices, but in GPU generation those slices come from maybe_create_ubatch_slices(...) and can represent microbatch partitions rather than request boundaries when ubatching is enabled. In that mode, one request can be split into multiple pseudo-requests here, so model_outputs/sr lengths no longer match input_batch.num_reqs, which then fails in gpu_generation_model_runner.sample_tokens (length check for dict list outputs) or misroutes audio to requests.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Keep Base cloning metadata when overriding prompts from file

When --txt-prompts is set, this branch rebuilds additional_information with only generic text/speaker fields for every query_type. If users run --query-type Base, the talker preprocessing path requires Base-specific clone inputs like ref_audio (and often ref_text), so this new path raises runtime validation errors instead of generating audio.

Useful? React with 👍 / 👎.

[lishunyang12]

A few comments on the batched decode path and edge cases.

[lishunyang12]

The Codex bot flagged this too — ubatch_slices may not map 1:1 to request boundaries under microbatching. Worth verifying this doesn't split a single request across slices.

[JuanPZuluaga]

Fixed in the last commit. now, we check for "token_slice"

[lishunyang12]

Looks good, the token_slice guard makes sense.

[lishunyang12]

Does SpeechTokenizer.decode() officially support a list of dicts? The old code passed a single dict — if the upstream API hasn't changed, this could silently produce wrong results.

[JuanPZuluaga]

Indeed, it officially supports it, see here

[lishunyang12]

Got it, thanks for the link.

[lishunyang12]

If tok.decode(valid_codes) returns fewer waveforms than expected, wavs[j] will IndexError. An assert len(wavs) == len(valid_codes) would help.

[JuanPZuluaga]

fixed

[lishunyang12]

Thanks.

[lishunyang12]

Any reason not to always pass a list here? The conditional unwrapping means omni.generate() gets different types depending on batch size.

[JuanPZuluaga]

fixed

[lishunyang12]

Since --batch-size must match a CUDA graph capture size, a runtime power-of-two check would save users from a cryptic CUDA graph error.

[JuanPZuluaga]

fixed

[hsliuustc0106]

@vllm-omni-reviewer

[github-actions]

🤖 VLLM-Omni PR Review

Code Review: [Qwen3TTS][Feat] Code2Wav batched decoding

1. Overview

This PR adds support for batched decoding in the Code2Wav stage of Qwen3TTS, enabling multiple audio generation requests to be processed in a single forward pass through the SpeechTokenizer. The changes include:

• Core batching logic in qwen3_tts_code2wav.py
• Updated end2end.py example with batch processing support
• New stage config file for batched processing
• Benchmark prompts file for testing
• Documentation updates in README

Overall Assessment: Positive. The implementation is well-structured and the benchmark results demonstrate meaningful performance improvements in stage-1 wall time.

---

2. Code Quality

Strengths

• Clean separation of concerns with the _split_request_ids helper method
• Proper handling of edge cases (empty inputs, invalid lengths, context trimming)
• Good use of existing vLLM infrastructure (ubatch_slices) for batch handling

Concerns

Import reordering noise: Multiple files have import statements reordered (moving from vllm import ... after other imports). While not incorrect, this adds significant noise to the diff. These should ideally be in a separate commit or PR.

Silent warning removal: In the original code, when context trim >= decoded length, a warning was logged. In the batched version, this silently continues without any indication:

# Original (lines 179-186):
if cut < audio_np.shape[0]:
    audio_np = audio_np[cut:]
else:
    logger.warning(
        "Context trim %d >= decoded length %d; returning empty audio.",
        cut,
        audio_np.shape[0],
    )
    return empty_ret

# New (lines 218-220):
if cut < audio_np.shape[0]:
    audio_np = audio_np[cut:]
else:
    continue  # context trim >= decoded length

Consider adding a debug-level log for this case.

Type annotation inconsistency: The forward method signature shows tuple[torch.Tensor, torch.Tensor] as return type hint but actually returns OmniOutput:

def forward(
    self,
    input_ids: torch.Tensor | None = None,
    ...
) -> OmniOutput:  # Good - signature is correct

---

3. Architecture & Design

Good Patterns

• The _split_request_ids method cleanly abstracts batch splitting logic
• Leveraging get_forward_context().ubatch_slices for batch boundary detection is elegant
• The parsing loop correctly handles mixed valid/invalid requests in a batch

Design Questions

Batch size constraint: The README mentions that batch size must match CUDA graph capture sizes. This is an important constraint that should be validated programmatically:

batch_size = args.batch_size
# Consider adding validation:
# valid_sizes = [1, 2, 4, 8, 16, 32, 64, 128, 256]
# if batch_size not in valid_sizes:
#     logger.warning("batch-size should be a power of 2 for CUDA graph compatibility")

Single-item batch workaround: The code has a workaround for single-item batches:

batch_input = batch[0] if len(batch) == 1 else batch

This suggests the API might benefit from consistently handling list inputs regardless of length.

---

4. Security & Safety

Resource Management

• No memory leaks detected; tensors are properly managed
• The batched approach should be more memory-efficient than sequential processing

Input Validation

• Good validation for n % q != 0 (warmup/dummy runs)
• Empty input handling is robust

Potential Issue

The _split_request_ids method falls back to returning [ids] if forward context is unavailable or ubatch_slices is None. This could mask issues if batching is expected but the context isn't properly set up:

def _split_request_ids(self, ids: torch.Tensor) -> list[torch.Tensor]:
    if is_forward_context_available():
        slices = get_forward_context().ubatch_slices
        if slices is not None and len(slices) > 1:
            ...
    return [ids]  # Fallback - could mask issues

Consider adding a debug log when falling back to single-batch mode.

---

5. Testing & Documentation

Documentation

• README update is clear and explains the batching feature well
• The constraint about stage-0 BS >= stage-1 BS is important and well-documented

Missing Test Coverage

The PR lacks unit tests for the new batching functionality. Consider adding tests for:

1. Single request (baseline)
2. Multiple requests with valid codes
3. Mixed batch with some invalid requests
4. Empty batch handling
5. Edge case: context trim >= decoded length

Test Results

The benchmark results clearly show the improvement:

• Baseline stage-1 times: 630ms, 274ms, 273ms
• Batched stage-1 times: 578ms, 222ms, 227ms

The improvement is visible, though the first batch shows less improvement (possibly due to warmup).

---

6. Specific Suggestions

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py

Line 117-128: Add docstring explaining ubatch_slices:

def _split_request_ids(self, ids: torch.Tensor) -> list[torch.Tensor]:
    """Split concatenated input_ids into per-request segments using forward context.
    
    Uses ubatch_slices from forward context which contains either:
    - int: number of tokens in the request
    - slice object with token_slice attribute
    
    Returns list of per-request id tensors, or [ids] if not in batched context.
    """

Line 218-220: Add debug logging for silent edge case:

if cut < audio_np.shape[0]:
    audio_np = audio_np[cut:]
else:
    logger.debug(
        "Context trim %d >= decoded length %d for request %d; returning empty audio.",
        cut, audio_np.shape[0], idx
    )
    continue

examples/offline_inference/qwen3_tts/end2end.py

Line 263-281: Add comment explaining template approach:

# Load prompts from text file if provided.
# Use the default query as a template so task-specific fields
# (e.g. ref_audio for Base) are preserved; only override text.
if args.txt_prompts:

Line 295: Consider adding batch size validation:

batch_size = args.batch_size
if batch_size & (batch_size - 1) != 0:  # Not a power of 2
    logger.warning(
        "batch-size %d is not a power of 2; CUDA graph capture may fail. "
        "Use sizes like 1, 2, 4, 8, 16...", batch_size
    )

Import Changes

Consider reverting the import reordering changes in:

• examples/offline_inference/qwen3_omni/end2end.py
• tests/conftest.py
• tests/engine/test_async_omni_engine_abort.py
• tests/entrypoints/openai_api/test_image_server.py
• tests/entrypoints/test_omni_llm.py
• vllm_omni/benchmarks/patch/patch.py
• vllm_omni/entrypoints/omni.py
• vllm_omni/entrypoints/omni_stage.py
• vllm_omni/entrypoints/openai/api_server.py

Or move them to a separate "style: fix import ordering" commit.

---

7. Approval Status

LGTM with suggestions

The core implementation is solid and the feature works as intended. The benchmark results demonstrate real performance improvements. However, I recommend addressing the following before merging:

Required:

1. Add debug logging for the silent edge case (context trim >= decoded length)

Recommended:
2. Separate import reordering into a different commit
3. Add batch size validation with warning for non-power-of-2 values
4. Consider adding unit tests for the batching functionality

Optional:
5. Add docstring to _split_request_ids method
6. Add debug log when falling back to single-batch mode in _split_request_ids

The PR is functionally correct and ready for merge after addressing the required item. The recommended items would improve maintainability but are not blocking.

---

This review was generated automatically by the VLLM-Omni PR Reviewer Bot
using glm-5.

[JuanPZuluaga]

Thanks for the comments! @lishunyang12. I think i've addressed most of them.

Do you think we could improve a bit the e2e script? I could:

• add some logic to run a larger benchmarking with the provided .txt
• should we keep the benchmark text samples? Or do we indicate that one could pass a txt file?

Additionally,

• would it be useful to run a larger benchmark and add the results here? I could run on a NVIDIA GeForce RTX 5090 Blackwell sm120.

[lishunyang12]

I think the e2e script is fine as-is for this PR — the txt file option is already there which is nice. Adding a larger benchmark result would be a plus but not blocking. Up to you if you want to include it.

[lishunyang12]

All previous comments addressed — LGTM.

[JuanPZuluaga]

I think the e2e script is fine as-is for this PR — the txt file option is already there which is nice. Adding a larger benchmark result would be a plus but not blocking. Up to you if you want to include it.

Thanks @lishunyang12. Let's merge this and i'll run a benchmark together with: #1438

[lishunyang12]

Solve precommit and DCO please.

[hsliuustc0106]

@vllm-omni-reviewer

[github-actions]

🤖 VLLM-Omni PR Review

Code Review: [Qwen3TTS][Feat] Code2Wav batched decoding

1. Overview

This PR adds support for batched decoding in the Code2Wav stage of Qwen3TTS. The changes include:

• Core model changes: Refactored qwen3_tts_code2wav.py to handle multiple requests in a single forward pass
• Example script updates: Added --batch-size and --txt-prompts arguments to end2end.py
• Configuration: Added qwen3_tts_batch.yaml with batched stage configs
• Documentation: Updated README with batched decoding instructions

Overall Assessment: Positive. The implementation is well-structured and the benchmark results demonstrate meaningful speedup in stage-1 processing. However, there are a few areas that need attention.

---

2. Code Quality

Strengths

• Good validation for power-of-two batch sizes (CUDA graph alignment)
• Comprehensive handling of edge cases (empty inputs, warmup runs, invalid token counts)
• Clear logging for debugging

Issues & Suggestions

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py:120-134

The _split_request_ids method has a complex condition that could be clearer:

if slices is not None and len(slices) > 1 and not any(hasattr(s, "token_slice") for s in slices):

Consider adding a docstring explaining when each branch is taken, or extracting the condition into a helper method like _is_batched_context().

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py:165-168

The assertion after batched decode could fail silently in production:

wavs, _ = tok.decode(valid_codes)
assert len(wavs) == len(valid_codes), f"Code2Wav returned {len(wavs)} waveforms for {len(valid_codes)} requests"

Suggestion: Consider raising a more descriptive RuntimeError instead of assert, as assertions can be disabled with -O flag.

examples/offline_inference/qwen3_tts/end2end.py:270-286

The file reading lacks error handling:

if args.txt_prompts:
    with open(args.txt_prompts) as f:
        lines = [line.strip() for line in f if line.strip()]

Suggestion: Add explicit error handling:

if args.txt_prompts:
    try:
        with open(args.txt_prompts) as f:
            lines = [line.strip() for line in f if line.strip()]
    except FileNotFoundError:
        raise FileNotFoundError(f"Prompts file not found: {args.txt_prompts}")
    if not lines:
        raise ValueError(f"No valid prompts found in {args.txt_prompts}")

---

3. Architecture & Design

Strengths

• Clean separation between single-request and batched code paths
• Good use of OmniOutput abstraction for multimodal outputs
• Configuration-driven approach for enabling batching

Design Considerations

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py:144-226

The forward method has grown significantly. Consider extracting helper methods:

1. _parse_request_codes(ids, q) - Parse single request's codes
2. _build_empty_output(num_requests) - Create empty output structure
3. _trim_audio(audio_np, ctx_frames, upsample) - Handle context trimming

This would improve readability and testability.

Batch size coordination between stages

The README correctly notes that both stages need max_batch_size >= batch_size. However, this constraint isn't validated at runtime. Consider adding a validation check in the Omni class initialization.

---

4. Security & Safety

Input Validation

examples/offline_inference/qwen3_tts/end2end.py:248-255

Good validation for batch size being a power of two. Consider also adding:

• Upper bound check (e.g., batch_size <= 64) to prevent memory issues
• Validation that batch_size matches the stage config's max_batch_size

Resource Management

The batched processing in end2end.py correctly iterates through batches, but there's no explicit cleanup between batches. If memory pressure is a concern, consider adding:

for batch_start in range(0, len(inputs), batch_size):
    batch = inputs[batch_start : batch_start + batch_size]
    omni_generator = omni.generate(batch, sampling_params_list=None)
    # ... process ...
    torch.cuda.empty_cache()  # Optional: explicit cleanup between batches

---

5. Testing & Documentation

Documentation

examples/offline_inference/qwen3_tts/README.md:90-104

The documentation is clear and helpful. Minor suggestion: add an example of the expected performance improvement with concrete numbers from the benchmark.

Test Coverage

Missing unit tests: The PR lacks unit tests for:

1. _split_request_ids method with various slice configurations
2. Batched forward pass with mixed valid/invalid requests
3. Edge cases (empty batch, single-item batch, all-invalid batch)

Suggested test cases:

def test_split_request_ids_batched():
    # Test with multiple requests
    ...

def test_split_request_ids_single():
    # Test with single request (no batching)
    ...

def test_forward_empty_batch():
    # Test with all empty/invalid requests
    ...

def test_forward_mixed_batch():
    # Test with some valid, some invalid requests
    ...

Import Order Changes

The changes in tests/entrypoints/openai_api/test_image_server.py, tests/entrypoints/test_omni_llm.py, and vllm_omni/benchmarks/patch/patch.py appear to be lint fixes. These should ideally be in a separate PR to keep changes focused.

---

6. Specific Suggestions

qwen3_tts_code2wav.py:120-134

def _split_request_ids(self, ids: torch.Tensor) -> list[torch.Tensor]:
    """
    Split concatenated input_ids into per-request segments using forward context.
    ...
    """
    if not is_forward_context_available():
        return [ids]
    
    slices = get_forward_context().ubatch_slices
    # Check if we're in a batched context with simple integer slices
    if slices is None or len(slices) <= 1:
        return [ids]
    
    # Skip if slices contain token_slice objects (different batching mode)
    if any(hasattr(s, "token_slice") for s in slices):
        return [ids]
    
    boundaries = [0]
    for s in slices:
        boundaries.append(boundaries[-1] + s)
    return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)]

end2end.py:270-286

Add validation:

if args.txt_prompts:
    if not os.path.exists(args.txt_prompts):
        raise FileNotFoundError(f"Prompts file not found: {args.txt_prompts}")
    with open(args.txt_prompts) as f:
        lines = [line.strip() for line in f if line.strip()]
    if not lines:
        raise ValueError(f"No valid prompts found in {args.txt_prompts}")

qwen3_tts_code2wav.py:165-168

Replace assertion with proper error:

wavs, _ = tok.decode(valid_codes)
if len(wavs) != len(valid_codes):
    raise RuntimeError(
        f"SpeechTokenizer returned {len(wavs)} waveforms for "
        f"{len(valid_codes)} requests"
    )

---

7. Approval Status

LGTM with suggestions

The PR is fundamentally sound and achieves its stated goal of enabling batched decoding for Code2Wav. The benchmark results demonstrate the expected performance improvement.

Before merging, please address:

1. Required: Add error handling for file not found in end2end.py
2. Required: Replace assert with proper exception in forward method
3. Recommended: Add unit tests for the new batched functionality
4. Optional: Consider extracting helper methods in the forward function for better readability
5. Optional: Move import order fixes to a separate PR

The import reordering changes in test files and patch.py are unrelated to the main feature and should ideally be separated, but this is a minor concern.

---

This review was generated automatically by the VLLM-Omni PR Reviewer Bot
using glm-5.

[JuanPZuluaga]

@lishunyang12 pre-commit and DCO fixed.

[hsliuustc0106]

lgtm

[linyueqian]

LGTM. Tested locally with batch_size=4 on 12 prompts (3 batches). All 12 wav files generated successfully - valid audio, 24kHz, 3-10s each.

Batch 1: e2e_wall_time=13,228ms, stage_1=1,242ms, 4 reqs, 392 tokens
Batch 2: e2e_wall_time=11,245ms, stage_1=543ms,   4 reqs, 459 tokens
Batch 3: e2e_wall_time=16,137ms, stage_1=422ms,   4 reqs, 508 tokens

output_3_c51959d6-550e-433a-b809-d93041ee6d04.wav
output_3_bf152fed-2769-4a78-877e-85a82d2ff558.wav
output_3_26fb9402-8219-4acb-a5f9-ae5c308b5f8c.wav
output_2_d978694c-3a14-4059-8d21-e81365ffffd5.wav
output_2_4829ba90-78ab-4093-bcb5-0d19ed2471a1.wav
output_2_4af2d240-4ba2-4e65-ac99-a6114d780fda.wav
output_1_87782150-4c8d-44eb-8005-ebc5972e0b2a.wav
output_1_6fc12b9f-c75a-42a6-bf35-d4ab7d0b906f.wav
output_1_3de4114b-e78e-46ee-9fd5-d1badcba623a.wav
output_0_ffeb1bc8-65bc-444e-a2ce-0a75049a35be.wav
output_0_aebb20a7-ef5f-419b-8b75-513ec83d6cb8.wav
output_0_85ea6d24-4315-450e-9c2a-f991efae5795.wav

[Gaohan123]

@JuanPZuluaga Could you please resolve the merged CI failure? Thanks. https://buildkite.com/vllm/vllm-omni/builds/3246/steps/canvas

[hsliuustc0106]

@JuanPZuluaga Could you please resolve the merged CI failure? Thanks. https://buildkite.com/vllm/vllm-omni/builds/3246/steps/canvas

I think the failures are not related to this PR and known to us. they also fail in other PR merges