Add Qwen3-TTS-12Hz-0.6B-CustomVoice support

[qaqjx]

What does this PR do?

Adds initial serving support for Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice.

• Registers the qwen3_tts model and adds its serving YAML and dependencies.
• Implements checkpoint configuration parsing, weight loading, request preprocessing, and PCM audio postprocessing.
• Declares the Talker and Codec partitions and their graph walks.
• Supports continuous batching and streaming codec output.
• Adds CUDA Graph support for batched decode and piecewise capture of inner loops.
• Adds configurable FlashInfer attention backends so Qwen3-TTS can use the compatible FA2 kernels.
• Adds unit tests and CUDA real-checkpoint integration tests covering graph transitions, batching, CUDA Graph execution, tokenizer/prefill behavior, and codec output.

Partition and graph-walk layout

Qwen3-TTS is split into two independently scheduled partitions:

• Talker uses the KV-cache engine and owns talker_prefill and talker_decode. Prefill consumes the text, speaker, and language inputs, then transitions to the autoregressive decode loop. Each decode iteration predicts one complete 16-group codec frame and feeds its embedding back into the next iteration.
• Codec is a stateless consumer partition with the codec_chunk walk. Codec tokens are streamed from Talker and buffered with left context. Once enough new frames are available, the conductor schedules codec_chunk, emits the resulting PCM audio to the client, and re-arms the partition for the next chunk.

This separation allows waveform decoding and audio streaming to overlap with subsequent Talker decode iterations instead of waiting for the full token sequence.

How was it tested?

• .venv/bin/ruff check .
• FLASHINFER_WORKSPACE_BASE=/tmp/mstar-flashinfer .venv/bin/pytest -q test/modular/test_qwen3_tts_model.py
  • 27 passed

Real-weight CUDA integration tests are included in test/integration/test_qwen3_tts_real_weights.py for GPU environments with the Qwen3-TTS checkpoint available.

Checklist

• ruff check . passes
• Added or updated tests / docs where relevant

[NSagan271]

Overall looks good. The model needs to be added to the docs/cli, and as a note, and I made #199 for the issues with the cuda graphable sampler that you had to worker around.

I also got a few issues when running it on a GPU:

(1) I tried running it on an H100 and I ran out of memory when capturing the code2wav cuda graphs until I removed batch size 16 from the codec cuda graph. I'd recommend either reducing the maximum batch size to 8, or perhaps the codec decoder can be made more lightweight (I see it is now a wrapper around the decoder from the qwen_tts package; it can, in the future maybe, be ported over and optimized, like the Qwen3-Omni codec).

(2) When testing with

python -m benchmark.runner \
    --url localhost:8000 \
    --model orpheus \
    --profiling-type closed_loop \
    --request-type text_to_speech \
    --num-requests 20 \
    --inference-system ours \
    --num-warmup 0 \
    --max-concurrency 4 \
    --dataset seed_tts \
    --output-dir .bench_outs

the requests were giving correct output audio, but all were running to max tokens, with the remainder of the audio silent. I think either the model is not properly outputting EOS or there is a bug in check_stop.

(3) I got

/bin/sh: 1: sox: not found                                                                                                                                                                                           
2026-07-31 17:53:38,435 WARNING [worker_0] sox: SoX could not be found!                                                                                                                                              
                                                                                                                                                                                                                     
    If you do not have SoX, proceed here:                                                                                                                                                                            
     - - - http://sox.sourceforge.net/ - - -                                                                                                                                                                         
                                                                                                                                                                                                                     
    If you do (or think that you should) have SoX, double-check your
    path variables.

on startup, even though sox appears to exist in my Python environment.

[qaqjx]

Overall looks good. The model needs to be added to the docs/cli, and as a note, and I made #199 for the issues with the cuda graphable sampler that you had to worker around.

I also got a few issues when running it on a GPU:

(1) I tried running it on an H100 and I ran out of memory when capturing the code2wav cuda graphs until I removed batch size 16 from the codec cuda graph. I'd recommend either reducing the maximum batch size to 8, or perhaps the codec decoder can be made more lightweight (I see it is now a wrapper around the decoder from the qwen_tts package; it can, in the future maybe, be ported over and optimized, like the Qwen3-Omni codec).

(2) When testing with

python -m benchmark.runner \
    --url localhost:8000 \
    --model orpheus \
    --profiling-type closed_loop \
    --request-type text_to_speech \
    --num-requests 20 \
    --inference-system ours \
    --num-warmup 0 \
    --max-concurrency 4 \
    --dataset seed_tts \
    --output-dir .bench_outs

the requests were giving correct output audio, but all were running to max tokens, with the remainder of the audio silent. I think either the model is not properly outputting EOS or there is a bug in check_stop.

(3) I got

/bin/sh: 1: sox: not found                                                                                                                                                                                           
2026-07-31 17:53:38,435 WARNING [worker_0] sox: SoX could not be found!                                                                                                                                              
                                                                                                                                                                                                                     
    If you do not have SoX, proceed here:                                                                                                                                                                            
     - - - http://sox.sourceforge.net/ - - -                                                                                                                                                                         
                                                                                                                                                                                                                     
    If you do (or think that you should) have SoX, double-check your
    path variables.

on startup, even though sox appears to exist in my Python environment.

Thanks for pointing this out. I also added a modular test that verifies the CLI and benchmark registrations.

(1) Fixed. CodecSubmodule.MAX_BATCH_SIZE is now 8 and its capture buckets are [1, 2, 4, 8]. The scheduler therefore rejects larger Codec micro-batches instead of attempting the unsafe batch-16 graph. The real server startup log confirms that only Codec batches 8, 4, 2, and 1 are captured. I also added unit coverage for the cap and for rejecting an oversized batch. The final H20 process used 55,158 MiB after startup and inference. I did not have an H100 available for an additional local rerun, but the implementation now enforces the safe ceiling requested in this review.

(2) Fixed. We confirmed through a CUDA-Graph-versus-eager A/B test that the full Talker CUDA Graph was the cause of the missing EOS, rather than the Talker weights or check_stop alone. The captured graph contains the stateful recurrent chain -- main-token sampling, the mutable CodePredictor scratch cache, the 16-code embedding sum, and the embedding fed into the next Talker frame. Under repeated graph replay, this path produced stale/incorrect recurrent state after the text condition was exhausted, so the Talker stopped emitting a natural codec EOS and the loop continued until max_output_tokens.

The safe fix is therefore to keep the outer Talker recurrent walk eager while retaining the independently captured 15-step CodePredictor loop and the Codec graphs. This preserves the useful inner-loop optimization without capturing request-owned recurrent state and sampling together. Before the fix, a 128-frame request always returned 10.24 seconds of audio and typically became all-zero after about 3.7–4.0 seconds. With the final path, six real-server requests all stopped naturally at 47–56 frames, returned 3.76–4.48 seconds of non-silent audio, and never reached the 256-frame safety limit. The final steady-state result on H20 was 1.13 seconds mean latency and 0.278 mean RTF over five measured requests.

(3) Fixed. The 12 Hz speech-tokenizer decoder does not use SoX; the warning came from importing the broad qwen_tts package, which eagerly imports the unrelated 25 Hz tokenizer and pysox. The worker now loads only the two exact upstream 12 Hz decoder modules through their installed package paths, without executing the broad package __init__ files. A cold-subprocess regression test verifies that the decoder
classes load, sox is absent from sys.modules, and neither stdout nor stderr contains the SoX warning.

[NSagan271]

@qaqjx you're right that the full Talker CUDA Graph was the cause of the missing EOS, and your safe fix does work, but I think I know the (fixable) root cause.

The runner re-runs preprocess on every replay, but with the capture slot's dummy request ids:

engine_inputs = ModelInputsFromEngine(request_ids=dummy_rids, ...)   # cuda_graph_runner.py:1403
real_inputs = submodule.preprocess(..., engine_inputs=engine_inputs) # :1416

Only alloc_manager.request_states is remapped dummy→real. The submodule's own request_states is not, and per_request_states is None on this path (Note: I think this is something that we should change upstream in the cuda graph runner... I will think about it more, but it's in any case not a blocker for your PR).

preprocess (qwen3_tts/submodules.py:321) calls _get_batch_suppress_mask(engine_inputs.request_ids), which reads self.request_state(rid)["generated_frames"]. A dummy id auto-creates an empty state with generated_frames = 0, causing **codec EOS to be masked to -inffor every request on every frame**. Then, the Talker can never stop and runs totalker_max_tokens`.

Proposed fix: keep preprocess a pure function of its inputs. prepare_inputs is always called by the engine with real requests (the runner never calls it), and the runner passes the real inputs list through to preprocess at replay (cuda_graph_runner.py:1379-1416), so carry the flag on ARNodeInputs:

# prepare_inputs, decode branch
suppress_eos = int(state.get("generated_frames", 0)) < self.config.generation.min_new_tokens
return ARNodeInputs(..., tensor_inputs={"suppress_eos": torch.tensor([suppress_eos], device=...)})

# preprocess
mask[:, eos] = torch.cat([item.tensor_inputs["suppress_eos"] for item in inputs])

Two things to keep in mind: single_request_inputs must declare the same tensor_inputs key (e.g., suppress_eos), since padding slots are capture_template.clone() and ARNodeInputs.clone() only copies keys that exist. Otherwise, it works at full batch and KeyErrors on the first padded one. And preprocess currently builds the mask for both walks, so _build_prefill needs to set it too (or the mask becomes decode-only).

I verified this on GPU (starting from HEAD~1): with that change, the whole-walk capture from the previous commit emits codec EOS normally. I'd restore get_cuda_graph_configs / can_use_cuda_graphs rather than drop the full-decode graph.

I also did another pass and found some minor issues; I'll make comments on those; otherwise overall looks good.

[NSagan271]

Some minor comments; main thing is #184 (comment) on the root cause behind the missing EOS.

[qaqjx]

@qaqjx you're right that the full Talker CUDA Graph was the cause of the missing EOS, and your safe fix does work, but I think I know the (fixable) root cause.

The runner re-runs preprocess on every replay, but with the capture slot's dummy request ids:

engine_inputs = ModelInputsFromEngine(request_ids=dummy_rids, ...)   # cuda_graph_runner.py:1403
real_inputs = submodule.preprocess(..., engine_inputs=engine_inputs) # :1416

Only alloc_manager.request_states is remapped dummy→real. The submodule's own request_states is not, and per_request_states is None on this path (Note: I think this is something that we should change upstream in the cuda graph runner... I will think about it more, but it's in any case not a blocker for your PR).

preprocess (qwen3_tts/submodules.py:321) calls _get_batch_suppress_mask(engine_inputs.request_ids), which reads self.request_state(rid)["generated_frames"]. A dummy id auto-creates an empty state with generated_frames = 0, causing **codec EOS to be masked to -inffor every request on every frame**. Then, the Talker can never stop and runs totalker_max_tokens`.

Proposed fix: keep preprocess a pure function of its inputs. prepare_inputs is always called by the engine with real requests (the runner never calls it), and the runner passes the real inputs list through to preprocess at replay (cuda_graph_runner.py:1379-1416), so carry the flag on ARNodeInputs:

# prepare_inputs, decode branch
suppress_eos = int(state.get("generated_frames", 0)) < self.config.generation.min_new_tokens
return ARNodeInputs(..., tensor_inputs={"suppress_eos": torch.tensor([suppress_eos], device=...)})

# preprocess
mask[:, eos] = torch.cat([item.tensor_inputs["suppress_eos"] for item in inputs])

Two things to keep in mind: single_request_inputs must declare the same tensor_inputs key (e.g., suppress_eos), since padding slots are capture_template.clone() and ARNodeInputs.clone() only copies keys that exist. Otherwise, it works at full batch and KeyErrors on the first padded one. And preprocess currently builds the mask for both walks, so _build_prefill needs to set it too (or the mask becomes decode-only).

I verified this on GPU (starting from HEAD~1): with that change, the whole-walk capture from the previous commit emits codec EOS normally. I'd restore get_cuda_graph_configs / can_use_cuda_graphs rather than drop the full-decode graph.

I also did another pass and found some minor issues; I'll make comments on those; otherwise overall looks good.

Thanks for digging into this. Your analysis is correct: the use of dummy request IDs was the underlying cause. I followed your suggested approach by computing suppress_eos in prepare_inputs for both prefill and decode, carrying it through ARNodeInputs, and building the mask only from those input tensors in preprocess; I also added the key to single_request_inputs and restored the full Talker decode CUDA Graph.

[NSagan271]

LGTM, I also tested out the latest changes and it works

[NSagan271]

@qaqjx I'm going to go ahead and merge this PR (so I can rebase #201 off of it and test the fix to #199)