perf(batch): prepared flights, VAD segmentation reuse and CUDA-graph TDT on the GPU worker (+61% L4 bench)

Merged
GI
2d ago
+798-64

What changed and why

Productionizes the GPU-efficiency gains from the batch-worker hill climb (avo-lite avo/gpueff), as three ship-on kill switches on the dedicated batch GPU worker. Batch at 100% and moving Japanese off Soniox both need the extra GPU headroom.

LeverEnv kill switch (false = off)Helm valueSource
Prepare the waiting lane's whole flight within the shared 1800-audio-second decode budget; return finished preparations when the GPU section endsBATCH_PIPELINE_PREPARE_FLIGHT_ENABLEDworker.prepareFlightEnabledticks t5 + t6
Reuse the VAD pass's raw segmentation scores for a full-unit diarization (no second segmentation forward)BATCH_PIPELINE_VAD_SEGMENTATION_REUSE_ENABLEDworker.vadSegmentationReuseEnabledtick t8
CUDA-graph TDT decoding, mode asserted around every call (mismatch quarantines)BATCH_PIPELINE_ASR_CUDA_GRAPHS_ENABLEDworker.asrCudaGraphsEnabledlane/gpe-cgraph, turned on by the L4 run below

The switches are read once at startup through ProviderConfig, like the shared embedding trunk (#24352). They are registered in config/feature-flags.yaml as ops_kill and added to backend/config/flag_budget.json (count 222) and config/deployment-setting-classification.json.

Not ported:

  • The T4 FP32 ASR fallback (t7): prod batch runs only on L4.
  • The ASR length split (t13): each extra ASR microbatch makes its own dispatch-time spend admission inside the GPU section. Prod's admission stage averages 1.49 s per call; the bench simulates 0.05 s. More calls would cost more GPU-held time than they save.

Product invariants affected

none

How it was verified

Dev bench (gpueff-avo/eval/run.sh): the real worker (load_worker) runs on 248 public-audio units (7055 audio-s) with the prod chart env, 2 lanes and flights of 8. Storage IO is simulated with fixed sleeps. The gate (gate.py) checks no lost speech, no result-class changes, word edit rate ≤1%, speaker agreement ≥0.97 and no words on silent units.

L4 (g2-standard-8, engine-pool, real L4, simulated IO):

runaudio-h/GPU-hGPU section sASR sdiarize sCUDA peak GiB
main ×270.35 / 70.2499.913.019.92.10
t5+t6+t8 ×2113.65 / 112.8561.713.515.92.10
this PR ×2112.91 / 112.8961.710.015.92.25
cgraph tree, eager ×2 / graphs ×257.13, 57.02 / 58.77, 58.63123 / 119.613.6 / 10.1—2.10 / 2.25

Every run's gate against the L4 main run passed with word edit rate 0.0, speaker agreement 1.0, 0 lost speech and 0 class changes. That is +61% audio-h/GPU-h on L4 against main. CUDA graphs cut ASR by 26% with identical words. On the full stack the bench objective stays flat, because the L4 bench is then bound by its simulated artifact IO, so that saving shows up as less GPU-busy time.

Caveats (simulated IO): the prepared-flight gain comes from overlapping fetches with the GPU section. Prod fetch_decode averages 0.27 s per unit (2 h aggregate) against ~0.43 s simulated, and prod is not saturated. Expect the prod gain from that lever to be smaller than the bench shows, and largest under batch load. VAD reuse and CUDA graphs remove GPU work outright, so they don't depend on IO.

Safety of the early release (t6): the pump cancels its lost-lease watcher on GPU entry and decides refunds (skipped) before dispatch, so releasing finished preparations after the GPU section can't change lease, refund or publication outcomes. Only done futures are closed, and close is idempotent. The decoded-memory bound is unchanged: 1800 s of preparation plus a 7200 s group, and buffers are freed sooner. If any prepared lease is lost while waiting, the whole prepared prefix is refunded with no attempt charge. That costs some throughput, never correctness.

Tests

  • backend/test.sh through omi-heavy, BACKEND_UNIT_TEST_FILE_LIST = all 69 tests/unit/test_batch_pipeline*.py files plus test_batch_worker_vad_recovery.py, test_batch_unit_policy.py and test_flag_budget.py (72 files): all passed on the t5/t6/t8 commit. On the final tree, 71 passed. test_batch_pipeline_admission.py::test_slow_dependency_cannot_grant_after_deadline[seal] failed one wall-clock < 1 s assertion on a loaded host. The file is untouched here, and it passed 50/50 on rerun.
  • New and updated tests:
    • Prepared prefix within the budget; a mismatched prefix is never reused; lost lease → whole prefix refunded with attempt 0; release when the GPU section ends; a running fetch stays reserved; kill switch (test_batch_pipeline_worker_bundle.py).
    • VAD-score reuse identity, consume-once, weak drop, shape fallback, kill switch, wiring only without flight batching, and exact equality with serial pyannote (test_batch_pipeline_flight.py).
    • Graph or eager mode persisting through NeMo's timestamp rebuild, plus violation quarantine (test_batch_pipeline_runtime.py).
    • Helm kill switches (test_batch_pipeline_chart.py).
  • The pyannote-backed equality tests skip locally (no torch). Run on the dev L4 pod against pyannote 4.0.3 (pytest --noconftest -k pyannote), 9 passed, including test_vad_segmentation_reuse_matches_serial_pyannote_exactly.
  • make preflight: see the CI run.

UI evidence

UI-Evidence: none -- backend GPU worker only

Failure class (fixes)

Failure-Class: none

🤖 Generated with Claude Code