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

backend/charts/batch-pipeline-worker/templates/deployment.yaml

@@ -124,6 +124,8 @@ spec:
             # Pump processes sharing this pod's GPU; "1" is the single-process kill switch.
             - {name: BATCH_PIPELINE_WORKER_PROCESSES, value: {{ $processes | quote }}}
             - {name: BATCH_PIPELINE_WORKER_IO_CONCURRENCY, value: {{ .Values.worker.ioConcurrency | default 5 | quote }}}
+            # Prepare the budgeted flight prefix before the GPU lock; "false" is the kill switch.
+            - {name: BATCH_PIPELINE_PREPARE_FLIGHT_ENABLED, value: {{ .Values.worker.prepareFlightEnabled | quote }}}
             - {name: BATCH_PIPELINE_WORKER_ENABLED, value: "true"}
             - {name: BATCH_PIPELINE_ENABLED, value: "false"}
             - {name: BATCH_PIPELINE_PERCENT, value: "0"}
@@ -152,6 +154,10 @@ spec:
             - {name: BATCH_PIPELINE_FLIGHT_BATCHING_ENABLED, value: {{ .Values.worker.flightBatchingEnabled | quote }}}
             - {name: BATCH_PIPELINE_FLIGHT_BATCH_SECONDS, value: {{ .Values.worker.flightBatchSeconds | quote }}}
             - {name: BATCH_PIPELINE_SHARED_EMBEDDING_TRUNK_ENABLED, value: {{ .Values.worker.sharedEmbeddingTrunkEnabled | quote }}}
+            # Full-unit diarization reuses the VAD pass's raw segmentation scores; "false" is the kill switch.
+            - {name: BATCH_PIPELINE_VAD_SEGMENTATION_REUSE_ENABLED, value: {{ .Values.worker.vadSegmentationReuseEnabled | quote }}}
+            # CUDA-graph TDT decoding, asserted per call; "false" is the kill switch (asserted eager).
+            - {name: BATCH_PIPELINE_ASR_CUDA_GRAPHS_ENABLED, value: {{ .Values.worker.asrCudaGraphsEnabled | quote }}}
             - name: ENCRYPTION_SECRET
               valueFrom:
                 secretKeyRef:

backend/charts/batch-pipeline-worker/values.yaml

@@ -78,6 +78,10 @@ worker:
   # Shared artifact IO pool, including first decode before the GPU lock.
   # 1 disables preparation and restores serial artifact IO.
   ioConcurrency: 5
+  # A waiting lane prepares every leading unit of its flight that fits the
+  # shared 1800-audio-second decode budget, not only the first, so the GPU
+  # section does not wait on artifact fetches. false restores first-only.
+  prepareFlightEnabled: true
   shadowPercent: 0
   shadowPaidRecoveryPerHour: 200
   shadowRecoveryApiUrl: ""
@@ -93,6 +97,17 @@ worker:
   # Kill switch: one speaker-embedding trunk forward per 10s chunk of a unit,
   # pooled per local speaker. Ship on; false restores pyannote's per-pair forward.
   sharedEmbeddingTrunkEnabled: true
+  # Serial chain: a unit diarized at full length takes the raw segmentation
+  # scores its own VAD pass computed (same model, chunking and forward batch,
+  # so identical scores) instead of a second segmentation forward. Compacted
+  # speech-only waveforms still segment themselves. false restores the
+  # second forward.
+  vadSegmentationReuseEnabled: true
+  # CUDA graph TDT decoding (same label-looping algorithm and words as eager;
+  # fewer kernel launches). The chosen mode is asserted around every call.
+  # L4 paired run 2026-11-02: ASR stage -26%, identical words. false restores
+  # asserted eager decoding.
+  asrCudaGraphsEnabled: true
   emptyAsrRecoveryEnabled: false
   # Activation follows the CPU role's separately approved secret-access change.
   independentLanguageRecoveryEnabled: false

backend/config/batch_pipeline_providers.py

@@ -35,6 +35,15 @@ class ProviderConfig:
     # per local speaker, instead of one per (chunk, speaker) pair. Never packs
     # across units. Kill switch; the flight path always shares the trunk.
     shared_embedding_trunk: bool = True
+    # Serial chain: a full-unit diarization takes the raw segmentation scores
+    # its own VAD pass computed instead of a second forward. Kill switch.
+    vad_segmentation_reuse: bool = True
+    # A waiting lane prepares every leading unit of its flight that fits the
+    # shared decode budget, not only the first. Kill switch.
+    prepare_flight: bool = True
+    # CUDA graph TDT label-looping decoding: same greedy algorithm and words as
+    # eager; the chosen mode is asserted around every transcription. Kill switch.
+    asr_cuda_graphs: bool = True
 
     def __post_init__(self):
         if self.pool_kind not in {'dedicated_batch', 'shared_live'}:
@@ -96,4 +105,19 @@ def provider_config() -> ProviderConfig:
             'BATCH_PIPELINE_SHARED_EMBEDDING_TRUNK_ENABLED',
             os.getenv('BATCH_PIPELINE_SHARED_EMBEDDING_TRUNK_ENABLED', 'true').strip().lower() != 'false',
         ),
+        # Ship-on: only an explicit 'false' restores the second segmentation forward.
+        vad_segmentation_reuse=rollout_default_on(
+            'BATCH_PIPELINE_VAD_SEGMENTATION_REUSE_ENABLED',
+            os.getenv('BATCH_PIPELINE_VAD_SEGMENTATION_REUSE_ENABLED', 'true').strip().lower() != 'false',
+        ),
+        # Ship-on: only an explicit 'false' restores first-decode-only preparation.
+        prepare_flight=rollout_default_on(
+            'BATCH_PIPELINE_PREPARE_FLIGHT_ENABLED',
+            os.getenv('BATCH_PIPELINE_PREPARE_FLIGHT_ENABLED', 'true').strip().lower() != 'false',
+        ),
+        # Ship-on: only an explicit 'false' restores asserted eager decoding.
+        asr_cuda_graphs=rollout_default_on(
+            'BATCH_PIPELINE_ASR_CUDA_GRAPHS_ENABLED',
+            os.getenv('BATCH_PIPELINE_ASR_CUDA_GRAPHS_ENABLED', 'true').strip().lower() != 'false',
+        ),
     )

backend/config/flag_budget.json

@@ -20,7 +20,7 @@
   },
   "seed_review_by": "2026-11-15",
   "budget": 169,
-  "count": 219,
+  "count": 222,
   "names": {
     "BATCH_VENDOR_ROUTING_ENABLED": {"owner": "vroute-b / two-lane lead", "review_by": "2026-11-29"},
     "BATCH_VENDOR_ROUTING_PERCENT": {"owner": "vroute-b / two-lane lead", "review_by": "2026-11-29"},
@@ -60,6 +60,9 @@
     "BATCH_PIPELINE_FLIGHT_BATCHING_ENABLED": {"seeded": "2026-10-25"},
     "BATCH_PIPELINE_FLIGHT_BATCH_SECONDS": {"seeded": "2026-10-25"},
     "BATCH_PIPELINE_SHARED_EMBEDDING_TRUNK_ENABLED": {"owner": "batch record lane / gpu efficiency sprint", "review_by": "2026-12-02"},
+    "BATCH_PIPELINE_VAD_SEGMENTATION_REUSE_ENABLED": {"owner": "planscost / gpu efficiency sprint", "review_by": "2026-12-02"},
+    "BATCH_PIPELINE_PREPARE_FLIGHT_ENABLED": {"owner": "planscost / gpu efficiency sprint", "review_by": "2026-12-02"},
+    "BATCH_PIPELINE_ASR_CUDA_GRAPHS_ENABLED": {"owner": "planscost / gpu efficiency sprint", "review_by": "2026-12-02"},
     "BATCH_PIPELINE_FRESH_GPU_PRIORITY_ENABLED": {"seeded": "2026-10-25"},
     "BATCH_PIPELINE_LANGUAGE_ELIGIBILITY_ENABLED": {"owner": "dazheng (batch record lane / rollout lead)", "review_by": "2026-11-15"},
     "BATCH_PIPELINE_LANGUAGE_INHERITANCE_ENABLED": {"owner": "dazheng (batch record lane / rollout lead)", "review_by": "2026-11-15"},

backend/diarizer/batch_pipeline.md

@@ -76,6 +76,20 @@ compares both paths on CPU with random weights (exact VAD spans and turns;
 segmentation within 1e-5, embeddings within 1e-4) whenever torch and pyannote
 are installed.
 
+**Serial VAD segmentation reuse** (`BATCH_PIPELINE_VAD_SEGMENTATION_REUSE_ENABLED`,
+chart `vadSegmentationReuseEnabled`, **default on**; `false` is the kill
+switch). With flight packing off, a unit diarized at full length (speech-only
+off, dense units kept over 80%, no evidence) used to run the segmentation model
+twice: once in VAD, once in `get_segmentations`. `VadSegmentationReuse` keeps
+the raw multilabel scores the VAD inference passes to its
+`pre_aggregation_hook` and hands them to the diarizer for that same waveform
+array (identity, held weakly). It installs only when `unsupported_reason`
+proves the two inferences share the model, chunking and forward batch, so the
+scores are the ones the second forward would produce; the result does not
+change. Compacted waveforms, other units and shape mismatches run pyannote's
+own segmentation. Scores of VAD-silent units are dropped at once; the rest live
+no longer than the decoded waveform (about 11% of its bytes).
+
 **Enablement gate.** CPU equality does not cover the L4: packing moves a unit's
 last chunk out of its own batch-of-one forward into a shared batch, so cuDNN may
 pick a different kernel, and a ~1e-5 embedding move or a powerset frame near 0.5

backend/diarizer/batch_pipeline.py

@@ -1,4 +1,4 @@
 """Conversation-sized, synchronous GPU chain. Invoke only on the GPU worker lane.
 
 Storage, lineage and model adapters are injected; importing this module does no
 IO and loads no models. One instance owns one GPU, with a lock across the entire
@@ -61,7 +61,7 @@
 from utils.batch_pipeline.errors import FatalGpuError
 from utils.batch_pipeline.language_metrics import record_route
 from utils.batch_pipeline.stage_timings import StageTimings
-from utils.batch_pipeline.worker_io import DecodeBudget, FirstDecode, PendingResults
+from utils.batch_pipeline.worker_io import DecodeBudget, FirstDecode, PendingResults, PreparedFlight
 from utils.batch_pipeline import worker_metrics
 from utils.observability.fallback import record_fallback
 
@@ -297,9 +297,14 @@
         self._futures: dict[int, Future] = {}
         self._submitted = 0
         self._consumed = 0
-        if prepared is not None and self._items and prepared.item is self._items[0]:
-            self._futures[0] = prepared.future
-            self._submitted = 1
+        if prepared is not None:
+            # Leading decodes prepared outside the GPU section, matched by
+            # position and object identity only.
+            for index, (item, future) in enumerate(prepared.decodes()):
+                if index >= len(self._items) or self._items[index] is not item:
+                    break
+                self._futures[index] = future
+                self._submitted = index + 1
 
     def _fill(self):
         while self._submitted < len(self._items):
@@ -476,10 +481,24 @@
 
     def prepare_first(self, items):
         # IO=1 is the existing serial rollback. No models or privacy shortcuts.
+        # Prepares the longest admitted prefix that fits the shared decode
+        # budget, so the GPU section does not wait on artifact fetches; the
+        # kill switch restores first-item-only preparation.
         if self._io is None or not items:
             return None
-        future = self._decode_budget.submit(self._io, self._decode_context, items[0])
-        return FirstDecode(self._decode_budget, items[0], future) if future is not None else None
+        handles = []
+        try:
+            for item in items if self.config.prepare_flight else items[:1]:
+                future = self._decode_budget.submit(self._io, self._decode_context, item)
+                if future is None:
+                    break
+                handles.append(FirstDecode(self._decode_budget, item, future))
+        except BaseException:
+            for handle in handles:
+                handle.cancel()
+                handle.close()
+            raise
+        return PreparedFlight(handles) if handles else None
 
     def _fatal(self, exc):
         self._unavailable = True
@@ -637,6 +656,7 @@
             if self._unavailable:
                 raise FatalGpuError('GPU chain unavailable; restart required')
             pending, group, seconds = [], [], 0
+            handle = prepared
             # Bound decoded host memory to two audio-hours (~460 MB float32).
             try:
                 for item in items:
@@ -652,12 +672,18 @@
                 for receipt in pending:
                     receipt.close()
                 raise
+            release = getattr(handle, 'release_settled', None)
+            if release is not None:
+                # Every group has dropped its waveforms: return the prepared
+                # reservations now so the waiting lane can prepare its whole
+                # flight. The owner's final close still drains the rest.
+                release()
             return PendingResults(
                 lambda: [result for receipt in pending for result in receipt.settle()],
                 [future for receipt in pending for future in receipt.futures],
             )
 
     def _process(self, items, *, authorize_dispatch=None, prepared=None) -> PendingResults:
         states: dict[str, dict[str, Any]] = {}
         windows, failed, results = [], set(), []
         failure_records = []

backend/diarizer/batch_pipeline_benchmark.py

@@ -21,7 +21,7 @@
     torch = None
     nemo_asr = None
 
-from config.batch_pipeline_providers import ProviderConfig
+from config.batch_pipeline_providers import ProviderConfig, provider_config
 from diarizer.batch_pipeline import pack_windows, speech_windows
 from diarizer.batch_pipeline_models import ParakeetWords, decode_pcm16_wav, pin_parakeet_decoding
 from utils.batch_pipeline.contracts import Priority, WorkItem
@@ -61,8 +61,9 @@ def main():
     raw_spans = json.loads(Path(args.vad_json).read_text()) if args.vad_json else [(0, duration)]
     spans = [(s['start'], s['end']) if isinstance(s, dict) else tuple(s) for s in raw_spans]
     model = nemo_asr.models.ASRModel.from_pretrained(args.model).cuda().eval()
-    pin_parakeet_decoding(model)
-    provider = ParakeetWords(model, enforce_eager=True)
+    cuda_graphs = provider_config().asr_cuda_graphs
+    pin_parakeet_decoding(model, cuda_graphs=cuda_graphs)
+    provider = ParakeetWords(model, cuda_graphs=cuda_graphs)
     rows = []
     report = {
         'audio_sha256': hashlib.sha256(body).hexdigest(),
@@ -72,7 +73,7 @@ def main():
         'device': torch.cuda.get_device_name(),
         'torch': torch.__version__,
         'model': args.model,
-        'cuda_graph_policy': 'persistently_disabled_and_asserted',
+        'cuda_graph_policy': 'enabled_and_asserted' if cuda_graphs else 'persistently_disabled_and_asserted',
         'rows': rows,
         'scope': 'ASR only; whole conversation input, internal windows; no persistence',
     }

backend/diarizer/batch_pipeline_flight.py

@@ -24,15 +24,19 @@
 its two model-feature hooks read the precomputed values.
 
 With packing off, ``install_shared_trunk`` gives the serial chain the same
-per-chunk trunk forward, one unit at a time.
+per-chunk trunk forward, one unit at a time, and ``VadSegmentationReuse``
+hands a full-unit diarization the raw segmentation scores its own VAD computed.
 """
 
 from collections.abc import Mapping
+from contextlib import contextmanager
 from dataclasses import dataclass
 from importlib import metadata
 import logging
 import math
+import threading
 import warnings
+import weakref
 from typing import Any, Callable, Sequence
 
 import numpy as np
@@ -66,6 +70,7 @@
 # trunk_embedding_rows mirrors SpeakerDiarization.get_embeddings of these
 # releases; any other pyannote.audio keeps its stock per-pair forward.
 SHARED_TRUNK_PYANNOTE_VERSIONS = ('4.0.',)
+VAD_SEGMENTATION_KEY = 'omi_vad_segmentation'
 
 
 def is_out_of_memory(exc) -> bool:
@@ -285,6 +290,113 @@ def get_embeddings(file, binary_segmentations, exclude_overlap=False, hook=None)
     return True
 
 
+class VadSegmentationReuse:
+    """Serial chain: diarize a full unit on the raw scores its own VAD pass computed.
+
+    The serial chain runs the shared community-1 segmentation model twice on a
+    full-unit diarization: once inside VAD, once inside ``get_segmentations``.
+    Both Inference objects slide the same model over the same waveform with the
+    same chunking and forward batch (``unsupported_reason`` proves it), and the
+    VAD keeps the raw multilabel scores until its ``pre_aggregation_hook``. So
+    those scores are exactly what the diarizer's skip-aggregation inference
+    returns; this hands them over instead of running the forwards again.
+
+    Scores are keyed by the identity of the live waveform array the VAD saw
+    (held weakly, dropped when the array is freed), so a speech-only compacted
+    waveform, a different unit or another flight never matches and runs
+    pyannote's own segmentation. Nothing is shared across units.
+    ``BATCH_PIPELINE_VAD_SEGMENTATION_REUSE_ENABLED=false`` restores the second
+    forward.
+    """
+
+    def __init__(self, diarizer, vad):
+        self.inference = diarizer._segmentation
+        self.counts = {'reused': 0, 'shape_mismatch': 0}
+        self._local = threading.local()
+        # Reentrant: a weakref callback may run on this thread under the lock.
+        self._lock = threading.RLock()
+        self._scores: dict[int, tuple[Any, Any]] = {}
+        vad_inference = vad._segmentation
+        aggregate_hook = vad_inference.pre_aggregation_hook
+        original = diarizer.get_segmentations
+
+        def pre_aggregation_hook(scores):
+            self._capture(scores)
+            return aggregate_hook(scores)
+
+        def get_segmentations(file, hook=None):
+            cached = file.get(VAD_SEGMENTATION_KEY) if isinstance(file, Mapping) else None
+            return cached if cached is not None else original(file, hook=hook)
+
+        vad_inference.pre_aggregation_hook = pre_aggregation_hook
+        # Same bypass as the shared-trunk hooks; pyannote's own call otherwise.
+        object.__setattr__(diarizer, 'get_segmentations', get_segmentations)
+
+    @classmethod
+    def install(cls, diarizer, vad):
+        reason = 'runtime_unavailable' if SlidingWindowFeature is None else unsupported_reason(diarizer, vad)
+        if reason is not None:
+            logger.warning('event=batch_vad_segmentation_reuse outcome=disabled reason=%s', reason)
+            return None
+        logger.info('event=batch_vad_segmentation_reuse outcome=enabled')
+        return cls(diarizer, vad)
+
+    @contextmanager
+    def capturing(self, samples):
+        """Record the raw scores of the VAD call made on ``samples``."""
+        self._local.samples = samples
+        try:
+            yield
+        finally:
+            self._local.samples = None
+
+    def _capture(self, scores):
+        samples = getattr(self._local, 'samples', None)
+        if samples is None:
+            return
+        key = id(samples)
+
+        def dropped(ref):
+            with self._lock:
+                if self._scores.get(key, (None,))[0] is ref:
+                    del self._scores[key]
+
+        try:
+            ref = weakref.ref(samples, dropped)
+        except TypeError:
+            return  # Not weakly referenceable: diarization segments it itself.
+        with self._lock:
+            self._scores[key] = (ref, scores)
+
+    def forget(self, samples):
+        """Drop the scores of a unit that will not be diarized (VAD silence)."""
+        with self._lock:
+            self._scores.pop(id(samples), None)
+
+    def _expected_chunks(self, num_samples):
+        """Inference.slide's chunk count: complete chunks plus a padded last one."""
+        window_size = self.inference.model.audio.get_num_samples(self.inference.duration)
+        step_size = round(self.inference.step * SAMPLE_RATE)
+        complete = (num_samples - window_size) // step_size + 1 if num_samples >= window_size else 0
+        has_last_chunk = (num_samples < window_size) or (num_samples - window_size) % step_size > 0
+        return complete + int(has_last_chunk)
+
+    def take(self, samples):
+        """The unit's VAD scores as get_segmentations returns them, or None."""
+        with self._lock:
+            entry = self._scores.pop(id(samples), None)
+        if entry is None or entry[0]() is not samples:
+            return None
+        scores = entry[1]
+        if scores.ndim != 3 or scores.shape[0] != self._expected_chunks(len(samples)):
+            self.counts['shape_mismatch'] += 1
+            logger.warning('event=batch_vad_segmentation_reuse outcome=fallback reason=shape_mismatch')
+            return None
+        self.counts['reused'] += 1
+        frames = SlidingWindow(start=0.0, duration=self.inference.duration, step=self.inference.step)
+        return SlidingWindowFeature(scores, frames)
+
+
 class FlightPyannote:
     """Packed VAD/diarization features for one process-owned pyannote pair."""