[dghosef]

Since #1072, Qwen3-Omni doesn't start on an 80 GB GPU (H100, A100-80GB) with either qwen3omni_2gpu.yaml or qwen3omni_colocated.yaml. Capture sizing for the 32768-patch vision bucket runs out of memory with ~2.5 GB free. With that one bucket removed, it starts but captures 0 of the Thinker's 68 graphs, and A2T drops from ~21x real time to 1.8x.

A torch memory snapshot taken before and after sizing shows what the encoder captures hold on the Thinker GPU: 2,048 MiB of FlashInfer float workspace (the 512 MB default × 2 slots × audio and vision), 883 MiB of 8 MiB int workspaces (one per captured wrapper, 110 of them), and 538 MiB of output buffers. This PR:

• declares the encoder workspaces explicitly: 256 MB for audio, 128 MB for vision. FA3 (Hopper) plans don't use the float workspace at all. FA2's captured plans always split KV at the bucket's segment ceiling, which needs ~230 MB for AuT at bs=8 and ~95 MB for the 8192-patch vision bucket. 64 MB would fail capture sizing off Hopper.
• captures Omni's audio encoder at batch 8 only. Smaller batches pad with empty rows, so every batch still replays. The override is on Omni's AudioEncoderSubmodule, so Qwen3-ASR keeps [1, 2, 4, 8].
• stops the vision buckets at 8192 patches (1920x1080). Larger images run the block loop eagerly.
• cuts thinker_kv in qwen3omni_colocated.yaml from 256 to 176 pages. Without it, colocated clears its startup floor by only 46 MiB and captures almost nothing.

Testing

All on H200, torch 2.11.0+cu130. "80 GB" means the GPU is capped to emulate an H100. Pre-#1072 main (c713c20) reproduces the #858 numbers measured on H100 (4.9 GB free colocated, 64/68 Thinker graphs on 2gpu), so the emulation matches. A2T is input-audio seconds / wall seconds at c=1 over 20 LibriSpeech clips (94 s of audio); higher is faster.

80 GB	main	pre-#1072	this PR
2gpu: Thinker graphs	doesn't start	64/68	54/68 (+ AuT 12/12, vision 4/6)
2gpu: A2T / TTFT	–	21.2x / 0.083 s	22.3x / 0.077 s
2gpu, FA2 forced on both encoders	–	–	starts, 53/68, 22.7x
colocated: graphs	doesn't start	Talker 24/24, Thinker 46/68	Talker 24/24, Thinker 23/68, AuT 10/12
colocated: A2T / TTFT	–	20.4x / 0.080 s	21.0x / 0.079 s
real H100 80GB, 2gpu: graphs	doesn't start (2,688 MiB free at the vision bucket)	–	54/68 (+ AuT 12/12, vision 4/6)
real H100 80GB, colocated: graphs	doesn't start (1,856 MiB free)	–	Talker 24/24, Thinker 26/68, AuT 10/12; vision and Code2Wav eager
real H100 80GB, both configs: A2T (20 clips), 448² and 2560x1440 images, t2s/s2s speech	–	–	0 failed, WER 0.5–0.7%

Uncapped H200 (141 GB), main vs this PR's encoder changes: every graph is captured on both, A2T is unchanged (21.3 vs 21.4x colocated, 22.3 vs 22.3x 2gpu), and the PR leaves ~3.3 GB more free. Dropping the 16384/32768 vision buckets costs ≤3% at 2560x1440 and 2–7% at 3840x2160 (~20–50 ms on a 0.8 s request).

Correctness:

• A2T transcripts: 19/20 identical to the eager-encoder run on every arm. Clip 4 differs everywhere, pre-Qwen3-Omni: encoders on ragged attention, no flash_attn #1072 included; the model answers it instead of transcribing it.
• Image descriptions: correct from 448² to 2560x1440 on every arm, FA2 included.
• Tests: the 29 test/modular files covering Omni, AuT, Qwen3-ASR, ragged attention and capture give 1,040 passed, 18 skipped on both this PR and main. The real-weight Omni prompt parity and Qwen3-ASR 0.6B/1.7B tests pass on both.

The real-H100 rows ran on torch 2.12.1+cu132 with torchaudio 2.11.0+cpu, PR head de5272d against main 7cc4b04.

Known, not changed here:

• On 2gpu at 80 GB, a 5-minute (308 s) audio clip OOMs in the AuT frontend at request time (HTTP 500, the server stays up); clips up to 4 min work. Merged into current main it works (a little more headroom after capture).
• Colocated's 176 Thinker pages cap one request at about 22K tokens: a 25K-token prompt gets a clean 400 ("needs 190 pages ... the pool has 175"), on H100 and H200. Main's 256 pages accepted it.
• On 2gpu at 80 GB, a 3840x2160 image (32,400 patches, eager) can OOM at request time on the Thinker GPU. On main this config can't start at all.
• A fresh uv pip install -e .[all] resolves torch 2.12.1 with torchaudio 2.11.1, which fails to load. Every Omni request then returns 500. All runs above pin torch 2.11.0 / torchaudio 2.11.0.

🤖 Generated with Claude Code