[RFC]: Shared CUDA Graph runner for model-local audio submodules

Motivation.

CUDA Graph is now on the critical path for vllm-omni audio serving, especially for Code2Wav / tokenizer / codec / acoustic / DiT-style submodules. The current outer vLLM model-forward graph path is relatively aligned with upstream vLLM through CompilationConfig.cudagraph_mode, batch descriptors, and the CUDA Graph dispatcher. However, many audio submodules still need model-local inner CUDA Graph wrappers because they run below or outside the normal vLLM model-forward boundary.

Today those inner wrappers are fragmented across models. Each one owns some combination of:

• enable/disable policy
• capture size and batch bucket selection
• warmup and capture lifecycle
• static input/output buffers
• replay and eager fallback logic
• output trimming semantics
• outer stream-capture guards
• stats/logging/debug controls
• model-specific config knobs

Recent issues show the same classes of problems recurring:

• #4466 / #4525: Qwen3-Omni Code2Wav graph replay used nominal output length instead of eager-equivalent output length.
• #4316: MOSS-TTS batching and CUDA Graph design are coupled because the wrapper is effectively B=1.
• #4562: HiggsAudioV3 hits graph-capture constraints from data-dependent output shapes.
• #2287: prior RFC identifies duplicated graph management and proposes moving graph ownership toward the runner layer.

This RFC proposes a smaller long-term maintainability step: keep model-specific audio semantics local, but centralize the common CUDA Graph mechanics.

Proposed Change.

Introduce a shared internal CUDA Graph runner for model-local audio submodules, plus thin model-specific adapters.

The shared runner owns:

• graph enablement policy using enforce_eager and compilation_config.cudagraph_mode
• bucket lookup for one or more static axes, for example (batch, frames)
• warmup/capture/replay lifecycle
• static buffer ownership
• torch.cuda.is_current_stream_capturing() fallback behavior
• eager fallback when no graph matches
• hit/fallback stats and logging
• common test helpers

Each model adapter owns:

• input tensor schema
• bucket key construction
• how runtime inputs are copied/padded into static buffers
• how to call the underlying module during capture
• how outputs are sliced/cloned after replay
• output length contract
• model-specific fallback constraints

The output length contract should be explicit. Suggested enum-style categories:

• nominal: output length is input_len * scale
• captured_minus_padding: output length is derived from captured output length minus padded frames
• per_row_lengths: batched output needs per-request trimming
• custom: adapter supplies a model-specific trim function

The goal is not to replace vLLM upstream CUDA Graph dispatch. For normal model-forward / AR decode paths, we should continue using vLLM's cudagraph_mode machinery. This runner is only for inner audio components that cannot naturally use the upstream wrapper boundary yet.

Suggested Migration Plan

1. Add the shared runner and adapter interface without changing behavior.
2. Migrate Qwen3-TTS / Qwen3-Omni Code2Wav first, because #4525 already exposes the output-length contract problem.
3. Add conformance tests that every migrated adapter must pass:
   • exact bucket
   • padded bucket
   • fallback
   • batch > 1 when supported
   • variable-length batch when supported
   • output aliasing after replay
   • active outer stream-capture fallback
4. Migrate one second family, preferably MOSS-TTS or MiMo Audio, to validate batch-axis and per-row-length support.
5. Gradually replace model-specific env knobs with common config where practical.

Alternatives Considered

1. Keep model-specific wrappers only.

   • Lowest short-term risk.
   • Continues duplicating graph behavior and correctness bugs.

2. Move all graph logic into the runner immediately, as in the spirit of #2287.

   • Most architecturally clean.
   • Higher migration risk because audio submodule boundaries and output semantics differ significantly by model.

3. Rely only on upstream vLLM CUDA Graphs.

   • Best alignment with upstream.
   • Not sufficient today for Code2Wav/tokenizer/codec/DiT paths that execute outside the normal model-forward graph boundary.

4. Shared runner plus model adapters.

   • Proposed option.
   • Keeps model semantics local while centralizing the mechanics that keep causing repeated bugs.

Open Questions

1. Should this RFC supersede part of #2287, or should it be treated as an incremental implementation path for #2287?
2. Should cudagraph_mode=NONE disable inner audio CUDA Graphs as well as outer model-forward graphs?
3. Which model family should be the second migration target after Qwen3-TTS / Qwen3-Omni?
4. Should per-row output lengths be required in the adapter API from day one?
5. What common metrics should all inner graph runners expose?

Feedback Period.

One week.

CC List.

@gcanlin, @linyueqian, @yangyonggit, @jajmangold

Any Other Things.

Related upstream references:

• vLLM CUDA Graph design: https://docs.vllm.ai/en/latest/design/cuda_graphs/
• vLLM CUDAGraphWrapper: https://docs.vllm.ai/en/stable/api/vllm/compilation/cuda_graph/
• vLLM CudagraphDispatcher: https://docs.vllm.ai/en/stable/api/vllm/v1/cudagraph_dispatcher/

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.