[Perf][CI] Use Whisper on GPU when memory permits

[akshatvishu]

Purpose

Previously, Whisper transcription used a GPU only when more than one GPU was available (n > 1). This meant that single-GPU runners, including , for example, H100 and MI300X machines, used the CPU even when the GPU had enough free memory.

The helper now checks GPUs from the highest index to the lowest. It uses the first GPU with at least 16 GiB of free VRAM. This keeps the existing preference for the last GPU and avoids checking devices that are not needed.

The 16 GiB threshold is only used to choose a device. It does not reserve memory. It leaves enough room for Whisper small and the large-v3 fallback.

Test Plan

pytest -sv \
  'tests/e2e/online_serving/test_qwen3_tts_base_expansion.py::test_voice_clone_streaming_001[async_chunk]' \
  --run-level full_model

ENV:

ROCM Version : 7.2.53211-c2d9476115
vLLM Version : 0.26.0
vLLM-Omni Version : 0.1.dev2403+g67c54777b.rocm (git sha: 67c5477, date: ocm)

GPU: MI300 x 1

Test Result

On current main 67c5477 :

================= 1 passed, 15 warnings in 1942.26s (0:32:22) ==================

On this branch d9d9645:

================== 1 passed, 15 warnings in 254.84s (0:04:14) ==================

Logs:

• qwen3_tts_test_pref.log
• qwen3_tts_test_main.log

cc: @linyueqian

[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.

[linyueqian]

Thanks for turning this around so fast, and for keeping the change this small. The capacity check is the right shape, the platform API usage is correct, and I like that the parametrized test pins the boundary at 16 GiB - 1 rather than just testing the happy path.

I verified the API surface before commenting: get_free_memory(device) is implemented on all five real platforms (cuda, rocm, npu, xpu, musa), and the base-class NotImplementedError is unreachable here because UnspecifiedOmniPlatform.get_device_count() returns 0 and your n > 0 guard catches that case. So there is no platform where this raises.

One thing worth tightening before merge, on the multi-GPU path, plus a matching test case. Details inline. Neither blocks the single-GPU win you measured, which is clearly worth having.

[linyueqian]

This gate now applies to multi-GPU hosts too, which the old n > 1 branch never did, and that adds a failure mode the previous code could not have.

Concretely: the Entrypoint Test with H100 lane runs tests/entrypoints/ on mirror_hardwares: h100_2, and that includes test_qwen3_omni_realtime_websocket.py, which transcribes. Its deploy config puts stage 1 (gpu_memory_utilization: 0.6) and stage 2 (gpu_memory_utilization: 0.1) both on devices: "1" — exactly the device this code targets at n - 1. That is roughly 0.7 of the card reserved, leaving on the order of 24 GiB nominal free on an 80 GiB H100.

So today it clears the 16 GiB bar, and nothing breaks on merge. But the margin is not large, the CI overlay in tests/helpers/stage_config.py sets utilization independently of this file, and if a lane's utilization rises or a two-GPU job lands on a smaller card, transcription silently drops to CPU and that lane eats the same ~6x slowdown this PR exists to remove. Under n > 1 that outcome was structurally impossible.

Two ways to close it, either is fine by me:

1. Keep the multi-GPU branch unconditional, since a spare device exists by definition, and apply the threshold only when n == 1:

   if n > 1:
       device_index = n - 1
   elif n == 1 and current_omni_platform.get_free_memory(current_omni_platform.get_torch_device(0)) >= _MIN_FREE_VRAM:
       device_index = 0

2. Better, scan all visible devices and take the one with the most free memory, using 16 GiB as a floor. That is robust to which device happens to be busy, and it would also help the multi-GPU case rather than merely preserving it.

Unrelated but worth a line in the PR description: each transcription runs in its own spawned process, so N concurrent calls each independently observe "16 GiB free" and each load their own model. Whisper small is 1 to 2 GiB so 16 GiB absorbs that easily, but the OOM that motivated 4e4458e was specifically a concurrency failure, so it is worth stating explicitly that the parallel case was considered.

Minor: consider hoisting 16 * 1024**3 to a module-level constant so the value is greppable and the line fits more comfortably.

[akshatvishu]

Thanks for the review!

• I used the max-free-device approach. The helper now checks all visible devices and selects the one with the most free memory, provided it has at least 16 GiB available. If no device meets the threshold, Whisper falls back to CPU. When devices have equal free memory, the existing last-device behavior is preserved.

• Added coverage for multi-GPU hosts where all devices have less than 16 GiB free, cases where a non-last device has the most free memory, and equal-memory ties.

• Added _MIN_FREE_VRAM and both the selection logic and tests now use it.

• Updated the PR description

[linyueqian]

Good boundary case here. The one gap is the scenario that this PR actually changes semantics for: multi-GPU where the target device is below the threshold. Something like

(2, 8 * 1024**3, "cpu"),

would pin that behavior. Right now (3, 16 GiB, "cuda:2") only covers multi-GPU on the passing side, so if someone later restores the unconditional n > 1 branch (which is one of the options I suggested on media.py), no test would notice either way. Whichever direction you settle on for the multi-GPU path, a case asserting it would be good to have.

[linyueqian]

Re-reviewed at 8bdb09b. The scan is implemented as suggested, _MIN_FREE_VRAM is hoisted, and the new parametrized cases cover both the passing and failing sides of the multi-GPU path. I re-ran the selection logic against all six cases and every one matches its expectation.

One correction to my own round-1 comment first. The concurrency caveat I raised is not a live risk, so the description does not need to hedge on it. convert_audio_file_to_text submits into a ProcessPoolExecutor(max_workers=1), and every concurrent path in tests/helpers/runtime.py calls assert_*_response from the main thread inside the as_completed loop rather than from the worker thread, so transcriptions are serialized end to end. There is no pytest-xdist in .buildkite/ either. N validators never hold N Whisper models at once.

What I did miss in round 1 is the cost of the probe itself, which I steered you toward. Details inline, with numbers measured on an 8-GPU H200. It is a small change and it keeps every one of your new test cases passing unmodified. Nothing here undoes the single-GPU win, which is clearly worth having.

@yenuo26 could you sanity check the placement heuristic when you get a chance, since this sits on the validation path of every GPU lane. I verified the platform API is correct on all five backends and measured the probe cost on real hardware, but a second opinion on whether most-free or first-fit is the right default would help.

[linyueqian]

[important] get_free_memory is torch.cuda.mem_get_info(device), and that call goes through c10::cuda::CUDAGuard, so probing a device instantiates a CUDA primary context on it. Scanning all of them therefore leaves a context on every visible GPU, including the ones the model server is running on. The old n > 1 path only ever touched a single device.

Measured on an idle 8-GPU H200 (torch 2.11.0+cu130), with one card deliberately busy at 3.7 GiB free:

	VRAM delta	wall
this PR, probe all 8	+528 to +531 MiB on every GPU, about 4.1 GiB total, of which 528 MiB lands on the busy card	3.11 s
reverse first-fit	+531 MiB on one idle GPU, +0 MiB on the busy card	2.47 s

Lanes in this repo top out at gpu: 4, so the realistic cost is nearer 2 GiB than 4 GiB, and it will not OOM an H100 at 0.6 to 0.7 utilization. It is still stray VRAM and startup time landing on the server's own card, in a change whose whole purpose is to save time.

Scanning from the top and stopping at the first device that clears the floor avoids it:

for candidate_index in range(n - 1, -1, -1):
    candidate_device = current_omni_platform.get_torch_device(candidate_index)
    if current_omni_platform.get_free_memory(candidate_device) >= _MIN_FREE_VRAM:
        device_index = candidate_index
        break

I ran this against your six parametrized cases and it returns the same device for all of them, so the test table stays exactly as written. Most-free is not worth paying for here: anything clearing 16 GiB holds whisper small at roughly 2 GiB, or the large-v3 escalation at roughly 10 GiB, with room to spare. It also restores the pre-PR multi-GPU behavior exactly, while keeping the single-GPU improvement you measured.

[akshatvishu]

Done. Switched it to reverse first-fit and added checks for probe order and early exit. All six existing cases still pass unchanged. I also updated the PR description and removed the concurrency note.

[linyueqian]

[suggestion] The comment this replaces recorded why device 0 was off limits, namely the single-GPU OOM behind #3822. The new one explains what the loop does but not what it is guarding against, which leaves 16 GiB looking like an arbitrary number to the next reader. One clause pointing back at that OOM would keep the history attached to the constant.

[akshatvishu]

done!

[linyueqian]

Re-reviewed the delta 8bdb09be..74fcc56d. Both round-2 findings are resolved and I verified them rather than taking the replies at face value. This looks good to merge from my side.

The scan is now reverse first-fit with an early exit, which is what the probe-cost finding was about: on the six parametrized cases it touches one device in four of them and two in the other two, instead of instantiating a CUDA primary context on every visible GPU. I transcribed the loop out of the head and ran it against all six cases, and both the selected device and the probe order match what the test asserts in every one.

test_transcribe_selects_device_by_available_memory now records the probe order and asserts it, which is a better test than I asked for. I asked for the behaviour to be pinned; recording the order pins the mechanism, so a future change that reintroduces a full scan fails loudly even if it happens to select the same device.

The comment fix landed too: the 16 GiB floor now carries its reason, naming the single-GPU server-plus-Whisper OOM from #3822, so the constant is no longer a bare number to the next reader.

One loose end, not a blocker. I asked @yenuo26 in round 2 to sanity check most-free against first-fit as a placement default and that never got picked up. I am comfortable proceeding without it: the platform API is verified on all five backends, the probe cost was measured on real hardware, and first-fit restores the pre-PR multi-GPU behaviour exactly while keeping the single-GPU win. If anyone disagrees with the default later it is a small, well-tested function to revisit.

Also worth noting the per-PR Buildkite lane has not run on this head, since the change is confined to tests/helpers/ and the PR carries no ready label. The GitHub Actions checks and readthedocs are green. Whoever merges may want the label first, given this sits on the validation path of every GPU lane.

[linyueqian]

Re-approving at eb878320. My earlier approval was recorded against 74fcc56d, and the branch has since been rebased and restructured, so that approval no longer covers what is here. I re-checked rather than assuming a rebase is a no-op.

Everything the approval rested on survived. The reverse first-fit scan with the early break is intact, _MIN_FREE_VRAM still carries the #3822 reason, and test_transcribe_selects_device_by_available_memory still asserts the probe order, so a regression to a full scan would still fail rather than pass quietly.

The footprint shrank from 63 insertions to 51, which looked at first like the rebase had dropped a hunk. It had not: the selection logic was extracted into a named _select_whisper_device() helper, which is tighter to diff and lets the test target the function directly instead of driving it through _whisper_transcribe_in_current_process. I confirmed the helper is actually wired in and not orphaned, defined at media.py:709 and called at :746.

One genuinely new behaviour that my earlier approval did not cover, and which I think is an improvement worth naming. The previous head had no model cache: whisper.load_model ran inside _whisper_transcribe_in_current_process, so every transcription reloaded the model and re-picked a device. This head adds _WHISPER_MODELS and moves selection into the load path, with a comment saying the device is picked on first load and the model stays there for the worker's lifetime. Since convert_audio_file_to_text submits into a ProcessPoolExecutor(max_workers=1), that worker process reuses the cached model, so this removes a repeated load rather than merely moving it. It also ties the chosen device to the device the model is actually resident on, which the previous arrangement did not guarantee.

The caveat from my last approval is unchanged and has now had a chance to clear, so it is worth restating more firmly. This PR carries the ready label, but no Buildkite lane has reported on eb878320: the only green checks are DCO, the 3.11 and 3.12 builds, docs and pre-commit, all of which are lint and packaging rather than behaviour. That is the second PR I have looked at today where ready is applied and the per-PR lane still does not run, so this may be a lane trigger problem rather than anything about this change. Given this function sits on the validation path of every GPU lane, I would want that understood before merge rather than after.

Verified before approving: mergeable=true against current main. state=blocked reflects the outstanding required checks and the fact that my approval sat on a superseded commit, which this review resolves.

[akshatvishu]

Rebased onto current main and included the follow-up changes from #6936. Local tests passed and pre-commit passed. Waiting for AMD and CUDA CI validation on d9e0aac

[andyluo7]

I verified current head d9e0aac8 against #6936. The two continuation patches are patch-identical after the rebase: c6e99caf matches 35e287c3, and d9e0aac8 matches f4361867. The relevant final file contents also match; the remaining test-ready.yml delta is unrelated newer-main CI work.

The existing ready label has not produced an exact-head CUDA Buildkite build, and there is no exact-head AMD build. @yenuo26, could you please retrigger CUDA ready, add merge-test and amd-test, and run AMD with DEBUG_TEST_YAML=merge,ready or a separate merge suite? The evidence gate is complete Base and CustomVoice pytest/quality summaries on AMD plus the corresponding CUDA ready/merge TTS jobs. I will verify the exact SHA, rendered suite, target jobs, and terminal logs before treating the timeout fix as proven.

[andyluo7]

@linyueqian, you originally applied the ready label and previously noted that it was not producing a Buildkite lane. After the rebase, current head d9e0aac8 still has no CUDA Buildkite build. Could you please toggle/reapply ready or otherwise retrigger the ready lane on this exact head? The rendered diff correctly retains the Qwen3-TTS CustomVoice job.

[andyluo7]

Current-head validation addendum for d9e0aac8:

• tests/helpers/tests/test_media.py: 16 passed.
• Buildkite uploader/skip suites: 47 passed.
• Changed-file pre-commit (including mypy and Buildkite schema), compileall, and git diff --check: passed.
• A local merge-tree against current main (5d20f6b9) is clean.
• Diff-aware rendering retains CUDA ready TTS · Qwen3-TTS CustomVoice Test and CUDA merge TTS · Qwen3-TTS CustomVoice Test plus TTS · Qwen3-TTS Base Test.
• Runtime placement remains platform-generic: it asks the existing platform API for device count, device objects, and free memory, then uses the returned device string. No ROCm/CUDA or model-name branch was added.

This closes the local/static validation. The remaining gate is still exact-head AMD Base/CustomVoice completion plus the corresponding CUDA ready/merge jobs.

[andyluo7]

@hsliuustc0106, thank you for routing #6935 and #6938. Could you also retrigger this canonical PR at exact head d9e0aac8? Its existing ready label predates the current-main rebase and still has not produced a CUDA Buildkite build. Please toggle/reapply ready, add amd-test and merge-test, and include AMD merge-suite coverage (DEBUG_TEST_YAML=merge,ready or a separate merge run). #6936 is closed in favor of this original credited PR.

[andyluo7]

@linyueqian, could you please re-review current head d9e0aac8 when convenient? Your last approval covered the pre-rebase implementation. The approved placement logic and probe-order test are patch-equivalent; the only later additions are the typed fake adjustment required by current-main mypy and the Buildkite source_file_dependencies entries that retain the affected CUDA TTS jobs for this helper-only diff. Local unit/uploader/schema checks and a current-main merge-tree are clean. Hardware validation remains separately pending and will not be inferred from approval.

[andyluo7]

@akshatvishu current main is now 3b6ce984 and includes changes to both CUDA ready/merge pipeline files. I simulated the merge locally: it is conflict-free, and the tests/helpers/media.py source dependencies added by this PR remain present.

One comment should be corrected during the rebase: in .buildkite/cuda/test-merge.yml, the MiniCPMO note still says Whisper “only uses an accelerator when a spare GPU exists.” That describes the old n > 1 policy; this PR now uses reverse first-fit on any device with at least 16 GiB free, including a sufficiently free single GPU. The hardware allocation itself does not need to change here, but the explanation should reflect the new capacity-based behavior.

After rebasing, the new exact head will still need the requested CUDA ready/merge and AMD merge validation before merge.

[akshatvishu]

One comment should be corrected during the rebase: in .buildkite/cuda/test-merge.yml, the MiniCPMO note still says Whisper “only uses an accelerator when a spare GPU exists.” That describes the old n > 1 policy; this PR now uses reverse first-fit on any device with at least 16 GiB free, including a sufficiently free single GPU. The hardware allocation itself does not need to change here, but the explanation should reflect the new capacity-based behavior.

Thanks for catching it @andyluo7 ; will update the comment soon!

Update: Added it in 2406a3a

[andyluo7]

Current-head follow-up on 2406a3a: I verified that the MiniCPM-o allocation comment now matches the 16 GiB reverse-first-fit behavior, the commit is signed off, and all GitHub static checks are green.

GitHub currently displays APPROVED, but the latest approval was submitted against superseded head eb878320; there is no review submitted against 2406a3a. There is also no Buildkite context on the current SHA. The ready label predates this head.

@linyueqian, could you please re-review exact head 2406a3a?

@yenuo26 @hsliuustc0106, this account cannot change repository labels. Could you please:

1. Remove and reapply ready for the exact-head ready build. As demonstrated on [CI] Isolate layerwise offload memory measurements #6938, that event should launch both CUDA ready and the current AMD L2 fallback.
2. Add merge-test for CUDA merge.
3. For the AMD Base/CustomVoice gate, run the AMD build with DEBUG_TEST_YAML=merge (or merge,ready). Until [CI][ROCm] Match CUDA/NPU L2/L3 label routing #6966 is merged into the source head, the existing AMD bootstrap still defaults PR builds to L2 even when merge-test is present.

The remaining runtime gate is:

• CUDA ready: Qwen3-TTS CustomVoice Test;
• CUDA merge: Qwen3-TTS CustomVoice Test and Qwen3-TTS Base Test;
• AMD merge: CustomVoice and Base, including terminal pytest and quality-validation summaries rather than only successful serving requests.

This is the original credited PR, so I am not modifying the author branch.

[linyueqian]

Re-reviewed at the current head 2406a3a3, as asked. The approval stands; nothing in the delta since it changes what I signed off on.

The branch was force-pushed after my approval, so eb878320 is no longer an ancestor and GitHub is carrying the approval forward across a history it cannot line up. I compared the PR's own content at both points rather than the raw commit range, diffing each against its own merge base with main so the rebased upstream commits drop out. The production surface is identical: tests/helpers/media.py is the same change on both sides. Three things landed on top, none behavioral. tests/helpers/media.py was added to the pipeline file triggers so Whisper helper changes route to the TTS lanes. The MiniCPM-o allocation comment was rewritten to state the 16 GiB free-VRAM threshold and why h100_2 guarantees a second card with headroom, which now matches the code it sits above. And the test fakes gained typing, init_kwargs annotated and the mp_context assertion split into an isinstance(mp_context, BaseContext) check plus the get_start_method() call.

Separately, the Buildkite problem from my earlier note has a cause and a fix. ready was on this PR with zero lanes ever reporting, which looked like a trigger bug. It is: pushing to a branch that already carries the label fires nothing, so a label set before the last push leaves per-PR CI permanently unrun while the PR shows BLOCKED as though a check had failed. Cycling the label is what re-fires it, and I have just done that here, so the four lanes should report shortly. Worth knowing generally, since it silently affects any PR whose head moves after labelling.

[linyueqian]

Re-approving at 2406a3a3 so the record is pinned to the head that was actually reviewed. My previous approval was recorded against eb878320, which the force-push removed from the branch, and a carried-forward approval on a history GitHub cannot line up is not the state this should merge in.

The delta since that approval is CI and test hygiene only, as set out in my earlier comment: tests/helpers/media.py added to the pipeline file triggers so Whisper helper changes route to the TTS lanes, the MiniCPM-o allocation comment rewritten to state the 16 GiB free-VRAM threshold and the h100_2 rationale, and the test fakes typed with a BaseContext import, an init_kwargs annotation, and the mp_context assertion split into an isinstance check plus the get_start_method() call. The production change is byte-identical.

Lanes at this head, now that cycling the label has actually fired them: general, Intel, and NPU green. AMD build 11410 fails only mi300_1: Diffusion / Model Test, on the four SenseNova paged-decode tests, each raising ImportError: vllm.vllm_flash_attn requires the CUDA flash attention extensions (_vllm_fa2_C or _vllm_fa3_C). On ROCm, use upstream flash_attn. That is the ROCm image lacking those extensions, it reproduces on unrelated PRs, and this branch touches no diffusion code. Every other AMD step passed.

This also closes the open question from my earlier round, where ready was on the PR with no lane having ever reported. That was not a per-PR oversight: a label already present when the head moves triggers nothing, so per-PR CI stays unrun while the PR displays as blocked. Cycling the label is the fix, and the lanes above are the result.