[Feat][OmniVoice]Support Varlen Attn, Request-Batch and Step-Execution

[sphinxkkkbc]

---

Purpose

Support request batching and step execution for OmniVoice by implementing model-specific helpers following the existing execute_stepwise() lifecycle.

Bit-exact parity with outputs generated before this PR is not guaranteed because the attention backend and packed execution layout have changed. WER results are listed below. The OmniVoice decoder is unchanged from main, and each request is decoded separately using exact-length codec tokens.

An earlier version of this PR implemented padded Decoder batching, but it was later reverted. With the packed varlen Generator, generated codec tokens are now kept at exact per-request lengths, so batching Decoder execution would require padding them again to the longest request. The current implementation therefore decodes completed requests individually using exact-length codec tokens. I have not benchmarked padded Decoder batching against per-request exact-length decoding, so it is unclear which approach is faster.

1. Add packed variable-length attention to the OmniVoice generator

The previous Generator layout stored conditional and unconditional inputs as separate rows. Since their sequence lengths differ, unconditional inputs had to be padded to the conditional sequence length. In a request batch, shorter
requests were also padded to the longest request, introducing substantial attention and Transformer overhead.

This PR changes Generator inputs to a packed, token-major 2D layout:

[cond0, uncond0, cond1, uncond1, ...]

with tensor shapes:

input_ids:     [total_seq_len, num_codebooks]
audio_mask:    [total_seq_len]
hidden_states: [total_seq_len, hidden_size]
logits:        [num_codebooks, total_seq_len, audio_vocab_size]

For B requests, CFG produces 2B independent sequences. Their cumulative offsets are represented by cu_seqs, allowing PyTorch varlen attention to keep requests and cond/uncond branches mutually invisible while retaining full bidirectional attention within each sequence.

Position IDs are reset independently for every packed sequence:

[0..cond0_len-1, 0..uncond0_len-1, 0..cond1_len-1, 0..uncond1_len-1, ...]

_position_ids_from_cu_seqs() derives these sequence-local positions directly from cu_seqs.

Per-request random generators are retained so one request does not consume or change another request's Gumbel sampling stream.

2. Make packed varlen attention compatible with CUDA Graphs

OmniVoice already supported CUDA Graphs, but the previous graph inputs and keys were based on padded [batch, codebook, sequence] tensors.

CUDA Graph keys are now: (request_batch_size, token_bucket)

Captured static inputs use:

input_ids:  [token_bucket, num_codebooks]
audio_mask: [token_bucket]
cu_seqs:    [2 * request_batch_size + 2]

The 2B+2 cu_seqs layout represents:

• 2B real cond/uncond sequences;
• the initial zero offset;
• one additional tail sequence used for Graph bucket padding.

In eager mode, the tail has zero length: [..., real_total, real_total]

In Graph mode, the last offset is changed to the selected bucket: [..., real_total, token_bucket]

This keeps the metadata shape static while isolating padded tail tokens from all real requests.

max_q and max_k use the fixed packed token bucket during graph capture, so the real maximum sequence length does not need to be part of the graph key.

Common request batch sizes derived from default config are pre-captured. Other batch sizes and oversized token shapes are captured lazily.

3. Support request-batch and step-execution lifecycles

Request-batch mode packs the input_ids and audio_mask values produced by _prepare_request_input() along the token dimension. cond_lens and target_lens are used to construct cu_seqs once before Transformer execution and to locate each request's target logits and generated codec tokens.

Step-execution stores each request's persistent latents as: [request_seq_len, num_codebooks]

InputBatch.make_batch() therefore concatenates active requests directly along dimension 0 without model-specific repadding. denoise_step() constructs the packed metadata, runs one varlen Generator step, and returns the same 2D layout so the Runner can scatter rows back to individual request states.

The standard lifecycle methods are used:

• prepare_encode() initializes request-local state and sampling metadata;
• denoise_step() executes one packed Generator step;
• step_scheduler() persists the request slice and advances its step;
• post_decode() decodes the completed request.

Generated codec tokens are maintained at exact per-request lengths.

Test Plan

Unit Tests

pytest -v \
  tests/entrypoints/openai_api/test_serving_speech.py \
  tests/model_executor/models/omnivoice/test_cuda_graph_generator.py \
  tests/model_executor/models/omnivoice/test_pipeline_batching.py \
  tests/model_executor/models/omnivoice/test_mask_dtype.py \
  tests/model_executor/models/omnivoice/test_fused_projection_load.py

The CUDA Graph tests continue the validation coverage from the earlier implementation for packed Generator. Additional tests were added for the varlen layout to verify dynamic cu_seqs updates across replays of the same graph key, and eager/Graph parity.

E2E Tests

pytest -v \
  tests/e2e/online_serving/test_omnivoice_expansion.py \
  tests/e2e/online_serving/test_omnivoice_parity.py

The E2E tests verify:

• request-batch execution correctness test;
• step execution generates valid audio without an error response, batch step-execution correctness test;
• B=1 step execution produces the same seeded WAV bytes as B=1 request-mode execution.

Perf and Acc

• request-rate=inf: batch vs per-req
• request-rate=1: batch vs step-execution

Serve Command:

vllm serve k2-fsa/OmniVoice \
  --omni \
  --port 8091 \
  --trust-remote-code \
  --max-num-seqs 8 \
  --request-batch-max-wait-ms 50 \
  --step-execution

Remove --step-execution for benchmarking request-batch.

Benchmark Command:

vllm bench serve --omni \
  --host 127.0.0.1 \
  --port 8091 \
  --model k2-fsa/OmniVoice \
  --backend openai-audio-speech \
  --endpoint /v1/audio/speech \
  --dataset-name seed-tts-text \
  --dataset-path /path/to/seedtts_testset \
  --seed-tts-locale en \
  --num-prompts 32 \
  --num-warmups 2 \
  --extra-body '{
    "language": "English",
    "seed": 42,
    "extra_params": {
      "num_inference_steps": 32,
      "guidance_scale": 2.0
    }
  }' \
  --max-concurrency 4 \
  --seed-tts-wer-eval \
  --request-rate inf \
  --percentile-metrics e2el,audio_rtf,audio_ttfp,audio_duration \
  --save-result \
  --result-dir ./omnivoice_results

Environment:

• vLLM Version: 0.28.0
• vLLM-Omni Commit: c4197ffe163c23017f179bf2f766c67766aca249

---

Test Results

Unit Test: 1 failed, 302 passed. The single failure is unrelated to this PR and is not addressed here.

FAILED tests/entrypoints/openai_api/test_serving_speech.py::TestTTSAsyncOffloading::test_prepare_speech_generation_awaits_voxtral_async - ValueError: Invalid speaker 'test'. Supported: bar

E2E : 14 passed

Performance

1. Batch vs Per-Request

C=4

Metrics	C=4 Batch	C=4 per-req	∆
Benchmark duration (s)	15.96	24.04	-33.61%
Request throughput (req/s)	2.00	1.33	+50.38%
Mean E2EL (ms)	1993.46	2852.22	-30.11%
Median E2EL (ms)	2257.81	2920.93	-22.70%
P99 E2EL (ms)	2290.35	3400.14	-32.64%
Audio throughput (audio duration/s)	7.35	4.88	+50.61%
Total Token throughput (tok/s)	26.12	17.35	+50.55%
Mean AUDIO_RTF	0.58	0.84	-30.95%
Median AUDIO_RTF	0.56	0.78	-28.21%
P99 AUDIO_RTF	0.95	1.57	-39.49%
Mean AUDIO_TTFP (ms)	1991.95	2851.77	-30.15%
Median AUDIO_TTFP (ms)	2256.20	2920.65	-22.75%
P99 AUDIO_TTFP (ms)	2288.25	3399.69	-32.69%

C=8

Metrics	C=8 Batch	C=8 per-req	∆
Benchmark duration (s)	14.75	24.44	-39.65%
Request throughput (req/s)	2.17	1.31	+65.65%
Mean E2EL (ms)	3579.96	5416.61	-33.91%
Median E2EL (ms)	3575.87	6025.53	-40.65%
P99 E2EL (ms)	4092.70	6495.25	-36.99%
Audio throughput (audio duration/s)	7.96	4.80	+65.83%
Total Token throughput (tok/s)	28.28	17.06	+65.77%
Mean AUDIO_RTF	1.08	1.60	-32.50%
Median AUDIO_RTF	1.02	1.54	-33.77%
P99 AUDIO_RTF	1.98	3.42	-42.11%
Mean AUDIO_TTFP (ms)	3577.48	5416.21	-33.95%
Median AUDIO_TTFP (ms)	3573.22	6025.31	-40.70%
P99 AUDIO_TTFP (ms)	4091.90	6494.96	-37.00%

2. Batch vs Step-Execution (C=4, request-rate=1)

Metrics	Step-Execution	Request-batch	∆
Benchmark duration (s)	32.73	32.70	-0.09%
Request throughput (req/s)	0.98	0.98	0.00%
Mean E2EL (ms)	981.70	943.41	-3.90%
Median E2EL (ms)	934.53	785.74	-15.92%
P99 E2EL (ms)	1764.82	1974.45	+11.88%
Audio throughput (audio duration/s)	3.59	3.59	0.00%
Total Token throughput (tok/s)	12.74	12.75	+0.08%
Mean AUDIO_RTF	0.28	0.27	-3.57%
Median AUDIO_RTF	0.24	0.21	-12.50%
P99 AUDIO_RTF	0.84	0.98	+16.67%
Mean AUDIO_TTFP (ms)	981.21	942.79	-3.92%
Median AUDIO_TTFP (ms)	934.14	785.35	-15.93%
P99 AUDIO_TTFP (ms)	1764.36	1974.01	+11.88%

Accuracy

All request passed for each benchmark with 32 requests. All benchmark result with Median WER=0.0000

Mode	Mean WER
Baseline (Original)	0.0074
Request-batch + varlen (C=4, inf)	0.0074
Request-batch + varlen (C=8, inf)	0.0074
Step-execution (C=4, rate=1)	0.0048
Request-batch (C=4, rate=1)	0.0096

---

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.

(anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/model_integration.md.

Module owners: @tzhouam @gcanlin

@sphinxkkkbc, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[sphinxkkkbc]

@linyueqian could you help review this PR? Thanks!

[linyueqian]

Thanks for the writeup, the packed varlen layout and the 2B+2 cu_seqs trick with the padded tail sequence are a clean way to keep the graph metadata static. Numbers on the request-batch path look convincing.

Three blocking items and a few smaller ones, all inline.

Blocking

1. pipeline_omnivoice.py:584: declaring supports_request_batch = True routes through execute_model_batch, where allow_single_output=False. Returning a bare DiffusionOutput on a per-request validation error now raises RuntimeError for every batch size including B=1, where main returned a clean user error.
2. pipeline_omnivoice.py:376: in step mode the same error return is discarded by the runner, state.latents stays None, and _prepare_latents raises outside the per-request try, failing every concurrent request in the batch.
3. omni_base.py:182: defaulting diffusion_batch_size to max_num_seqs changes startup behavior for all 51 diffusion pipelines. 43 of them lack supports_request_batch and will now fail to start under --max-num-seqs > 1 where they previously ran serially. This is unrelated to the PR title and has no test.

Test coverage

test_cuda_graph_generator.py is the only core_model test, and its _SyntheticGenerator never calls _varlen_attn, so it covers the padding, slicing and replay bookkeeping but not the attention change itself. The e2e tests are slow plus tts plus L4, so they only run in the weekly sweep, which does pick up the new test_omnivoice_parity.py automatically through the marker-driven collection in .buildkite/cuda/test-weekly.yml:93. test_omnivoice_parity.py also sets OMNIVOICE_CUDA_GRAPH=0 on both sides, so the CUDA-graph plus varlen combination, which is the riskiest part of this change, has no automated coverage at all. Could you add a core_model test that drives a small real OmniVoiceGenerator through _varlen_attn in both eager and graph mode, including a replay with a nonzero padded tail?

Reviewed at 7a096c39668eb3295c016f792a16d4b3ea21f577. Static review only, no GPU run on my side.

[linyueqian]

[High] This changes the diffusion batch default for every pipeline, not just OmniVoice, and the PR body does not mention it.

diffusion_batch_size is force-assigned onto od_config.max_num_seqs at stage_init_utils.py:1444 and stage_engine_startup.py:1561, overwriting whatever the stage config resolved. With the old default of 1, --max-num-seqs N was silently ignored by diffusion stages. Now it lands, which is what makes the --max-num-seqs 8 e2e test in this PR work, and it does match docs/user_guide/diffusion/execution_modes.md:232.

The side effect: only 8 of 51 diffusion pipelines set supports_request_batch = True, and diffusion_engine.py:229 raises at startup when max_num_seqs > 1 without it. Serve commands that run today (serially) will fail to start after this. It also still clobbers a per-stage engine_args.max_num_seqs with the global CLI value.

Please split this into its own PR with a unit test in tests/engine/test_async_omni_engine_stage_init.py, or at minimum call it out in the body and in the execution-modes doc.

[sphinxkkkbc]

This change was removed in d5f0bfb. #6405 is already opened, so we can support that PR.

Also, #6525 resolved the issue where max_num_seqs was being overridden to 1. Now we can configure this in deploy.yaml. I've set the default max_num_seqs to 8 in 3354dbf.

[linyueqian]

[High] Returning a bare DiffusionOutput here now raises RuntimeError instead of surfacing the user error.

Declaring supports_request_batch = True routes OmniVoice through execute_model_batch, which calls _normalize_pipeline_outputs(..., allow_single_output=False) (diffusion_model_runner.py:650). That helper raises

RuntimeError: OmniVoicePipeline.forward returned a single DiffusionOutput;
request-batch forward must return list[DiffusionOutput].

at diffusion_model_runner.py:84, for every batch size including B=1. On main OmniVoice went through execute_model with allow_single_output=True, so the same prompt produced a clean per-request error.

Trigger: offline Omni.generate with two clips in multi_modal_data["audio"] (line 258 above), or a dict prompt with no text key (line 276).

Fix direction: allocate outputs = [None] * len(req.requests) up front, place the error DiffusionOutput in that request's slot, and keep the remaining requests running.

[sphinxkkkbc]

Resolved in d5f0bfb, added

outputs = [None] * len(req.requests)
prepared_indices: list[int] = []

to handle this and preserve the original request indices to prevent slot mismatch

Also added regression test in tests/model_executor/models/omnivoice/test_pipeline_batching.py for changes in pipeline_omnivoice.py

[linyueqian]

[High] In step mode this error return is discarded, and the malformed request takes down the whole step batch.

diffusion_model_runner.py:714 calls prepare_encode(state) for its side effects only and drops the return value. When _prepare_request_input returns an error, state.latents is never assigned, so InputBatch.make_batch hits _prepare_latents and raises ValueError("All requests must have 'latents' initialized.") (input_batch.py:347). That happens at execute_stepwise line 776, outside the per-request try at line 810, so one bad request fails every concurrent request in the batch.

The annotation on line 371 also matches neither the interface (interface.py:60 declares -> StepRequestState) nor the success path, which returns None. Suggest recording the error on the state and letting the runner emit it per request, then fixing the annotation.

[sphinxkkkbc]

This was initially resolved in d5f0bfb, but after merging with the latest main, I found that #5810 already fixed it a few days ago. I've reverted it in f78c290 and now directly using the handling from main.

Annotation fixed.

[linyueqian]

[Medium] The pre-warm set does not cover the concurrency this PR targets, so the hot path falls into exact-size lazy capture.

The graph key is now (request_batch_size, packed_token_total), but this tuple is hard-coded to (1, 2, 3, 4) and the largest bucket is 1024 (configs/omnivoice.py:88).

Worked example with the parity-test prompt: RuleDurationEstimator gives target_len about 80, so the packed total per request is about 200 (cond_len about 120 plus target_len 80).

• B=4 gives about 800, bucket 1024, fine.
• B=8 gives about 1600, past every bucket. _find_bucket returns None and line 629 sets bucket = seq_len exactly.

So every distinct combination of request lengths becomes its own lazy capture of a 28-layer graph on the hot path, thrashing the 16-entry cache. B in {5..8} is never pre-warmed at all, and tests/e2e/online_serving/test_omnivoice_expansion.py:34 runs this model with --max-num-seqs 8.

Fix direction: derive the capture batch sizes from max_num_seqs, and round oversized totals up to a capped ladder instead of capturing at the exact packed length.

[sphinxkkkbc]

Resolved in d5f0bfb6ae982bd21f73d3e4756aad458202f637, now _OmniVoiceCUDAGraphForward will read od_config and derive buckets and batches to be captured. I configured the buckets in a ladder pattern (shown below), which properly activates the LRU cache. With MAX_LAZY_GRAPH = 16, lazy-captured buckets are rounded up to multiples of 128. In benchmarking (using the command described), most scenarios do not trigger lazy capture, except for three specific (batch, seq_len) combinations: (7, 1574), (8, 1511), and (7, 1511). I believe the current capture strategy is reasonable.

This PR measures ~7 GiB peak memory after pre-capture, versus ~5 GiB on the main branch.

Long-text generation issues persist as noted in #6333. #6409 proposes a chunk-based solution with a per-request limit of ~15 seconds for ~375 tokens in a single forward pass, so capturing extremely long sequences is unnecessary.

B=1: [128, 192, 256, 320, 384, 448, 512]
B=2: [256, 320, 384, 448, 512, 640, 768]
B=3: [384, 448, 512, 640, 768, 896, 1024]
B=4: [512, 640, 768, 896, 1024]
B=5: [640, 768, 896, 1024]
B=6: [768, 896, 1024]
B=7: [896, 1024]
B=8: [1024]

Regression test added in tests/model_executor/models/omnivoice/test_cuda_graph_generator.py for LRU.

[linyueqian]

[Medium] This reverses a deliberate decision on main, and the class docstring above still describes the old behavior.

main used a per-instance torch.cuda.graph_pool_handle() with an explicit comment:

Lazy-init per-instance pool handle: isolates OmniVoice CUDA Graph memory from other vllm modules (unlike get_global_graph_pool which shares a single pool across all captured graphs and can cause memory aliasing when two graphs replay concurrently).

That comment is gone, but the class docstring at lines 486 to 490 still says graphs "share a single per-instance pool handle, which isolates OmniVoice CUDA memory from other vllm modules". Please either restore the per-instance pool, or state why the aliasing concern no longer applies and update the docstring to match.

[sphinxkkkbc]

Reverted in d5f0bfb. This change was introduced while debugging an issue during testing, but there were two potential variables involved, and this change turned out to be one of them. After double-checking, I found no reason to keep it. I'll comment below with more details.

[linyueqian]

[Medium] Lines 488 to 551 duplicate omnivoice_generator.py:956-1016 nearly verbatim.

The CFG fuse, the log_probs[..., mask_id] = -inf mask, the layer penalty, the Gumbel position noise, the top-k select and the cond/uncond mirror write are copied line for line, comments included. The only differences are guidance_scales[i] vs guidance_scale, generators[i] vs request_generator, and where sample_tokens comes from. Any future correctness fix has to land in both places.

Extracting one _unmask_one_request(...) helper owned by the generator would also remove denoise_step's reach into generator._prepare_embeddings, _transformer_forward, _get_logits and _cuda_graph_fwd, which are all private.

[sphinxkkkbc]

Resolved in d5f0bfb

[linyueqian]

[Minor] Two small things in this method.

1. This line has no input_ids.is_cuda guard, unlike the request-mode path at omnivoice_generator.py:929. Worth keeping the two consistent.
2. Line 450 reads state.extra.get("guidance", self.guidance_scale), but prepare_encode sets state.guidance (line 420), not state.extra["guidance"]. The .get always falls through, so the per-request guidance plumbing is dead. Either write state.extra["guidance"] in prepare_encode or read state.guidance here.

[sphinxkkkbc]

Both resolved in d5f0bfb, regression test for Q2 added in tests/model_executor/models/omnivoice/test_pipeline_batching.py

[linyueqian]

[Minor] bucket is not None cannot be false: line 629 assigns bucket = seq_len when _find_bucket returns None, so the check is dead and only bucket != seq_len matters.

Also, _lazy_graphs is documented as LRU but behaves as FIFO: the hit path on line 639 never calls move_to_end, so a frequently used key is still evicted by popitem(last=False).

[sphinxkkkbc]

Both issues have been fixed.

Detailed comments on the _lazy_graphs FIFO issue have been inlined in the bucket issue. This was an existing unresolved issue on main.

[linyueqian]

[Minor] This block is out of scope for OmniVoice and unvalidated.

OmniVoice ignores both fields: prepare_encode reads self.num_step and self.guidance_scale from the model config, never state.sampling. So this changes behavior only for other diffusion TTS pipelines, with no test in this PR.

Two details if it stays:

• guidance_scale_provided = True is redundant. OmniDiffusionRequest.__post_init__ already sets it whenever guidance_scale is not None (vllm_omni/diffusion/request.py:66).
• int(...) and float(...) run on unvalidated client input. extra_params: {"num_inference_steps": "abc"} raises ValueError and surfaces as a 500 rather than a 400.

[sphinxkkkbc]

These two changes are primarily required for scheduler correctness:

1. StepScheduler._get_total_steps() reads sampling_params.num_inference_steps. If the value remains only in extra_args and is not promoted to the standard sampling field, the scheduler receives None and fails when evaluating int(None).

2. Request-batch compatibility is determined from the standard sampling fields, including guidance_scale. If the request value remains only in extra_args, requests with different guidance scales may receive the same compatibility key and be incorrectly admitted to the same request batch.

Regression coverage was added in tests/entrypoints/openai_api/test_serving_speech.py

[linyueqian]

[Minor] _varlen_attn is a private torch API imported at module scope, so it becomes a hard load-time requirement for OmniVoice on every platform.

What is the minimum torch version, and is torch.nn.attention.varlen present on the npu, rocm and xpu images? A short note in the module docstring plus the minimum version in the model doc would help, since there is no other use of this API in the repo.

[sphinxkkkbc]

Resolved in 9432df3. Now using the in-tree class diffusion.attention.layer.Attention with an explicit attn_mask fallback when varlen_attn is not supported.

As for the version question, the torch version I used is the one installed with vLLM(2.13), and the minimum version that supports varlen_attn is torch 2.10.

[sphinxkkkbc]

Inline comments are replied. Perf number still hold.

Test coverage updates:

1. test_omnivoice_parity.py now verifies both graph and eager modes.
2. test_cuda_graph_generator.py now uses the small real OmniVoiceGenerator in all cases rather than the _SyntheticGenerator. The results and commands have also been updated with descriptions.

[sphinxkkkbc]

@linyueqian PTAL, thanks!

[linyueqian]

One practical follow-up to the approval, because it is approved but still cannot merge.

Auto-merge is armed, so this will land by itself once the required checks are satisfied. The blocker is that buildkite/vllm-omni is a required check and it is still showing build 14666, which is red on the mooncake test_concurrent_put_get_threaded_both_sides MD5 mismatch. That failure is not yours, as covered in the review, but GitHub does not care whose it is.

I just tried clearing it without troubling you: cycling the ready label re-fired the NPU lane, but the general lane stayed pinned to the old failed build. The general pipeline triggers on push, so a label cycle cannot replace a stale general result. Only a new commit will.

So when you have a moment, merging current main into the branch will give you a fresh general run, and the flake is intermittent so it should come back clean. That is also worth doing on its own terms, since main has moved since 70a2b19 and you have already been bitten twice by tests on main that call APIs this branch changes. I would rather you spend the push on a real rebase than on an empty commit.

Nothing else is outstanding from my side.

[sphinxkkkbc]

Rebased cleanly with no conflicts, ran the unit tests locally with no new errors raised. The ready label needs to be cycled.