ASR Engine¶

The engine subsystem (oasr/engine/) is the top-level inference façade. It takes raw audio in and produces transcripts out, hiding the complexity of feature extraction, dynamic batching, paged KV-cache management, encoder forward passes, and CTC decoding behind a small public API.

OASR's engine design is vLLM-inspired: a single step loop drives both offline (single-pass batched) and streaming (chunk-by-chunk with paged KV cache) work in one pool, with length-aware bucketing and CPU/GPU overlap.

1. Purpose and Responsibilities¶

ASREngine is responsible for:

Loading the model and its config from a WeNet-format checkpoint directory.
Owning shared GPU resources — the paged KV block pool, CNN cache manager, CTC state manager — through ModelRunner.
Routing requests through the scheduler into either the offline executor or the streaming step path.
Driving the step loop: schedule → ingest fbank → forward → decode → finalise.
Exposing a high-level API:
transcribe(audio | List[audio]) for one-shot streaming use.
transcribe_offline(audio | List[audio]) for batch-only offline use.
transcribe_outputs(...) when the strings are not enough — it returns the whole RequestOutput, which is where words, confidence, timestamps, nbest_texts and finish_reason live.
add_request / feed_chunk / step / run for explicit control.

The engine no longer has a separate OfflineEngine subclass — pass streaming=False to add_request / transcribe, or use the transcribe_offline convenience helper, to take the offline path.

2. High-Level Architecture¶

                       ASREngine
   ┌──────────────────────────────────────────────────────────┐
   │                                                          │
   │   ┌───────────────┐   ┌─────────────┐                    │
   │   │ InputProcessor│   │  Scheduler  │                    │
   │   │ (load/fbank/  │   │ (admission, │                    │
   │   │  chunk split) │   │  batching)  │                    │
   │   └───────┬───────┘   └──────┬──────┘                    │
   │           │                  │                           │
   │           ▼                  ▼                           │
   │       Request*          SchedulerOutput                  │
   │           │                  │                           │
   │           │   ┌──────────────┴────────────┐              │
   │           │   ▼                           ▼              │
   │           │ OfflineExecutor       streaming step path    │
   │           │ (micro-batches,       (per-step batched      │
   │           │  CPU/GPU overlap)      paged forward)        │
   │           │   │                           │              │
   │           │   ▼                           ▼              │
   │           │  ┌──────────────────────────────┐            │
   │           └─▶│         ModelRunner          │            │
   │              │  (offline forward,           │            │
   │              │   streaming forward+cache)   │            │
   │              └──────────────┬───────────────┘            │
   │                             ▼                            │
   │                  ┌─────────────────────┐                 │
   │                  │   OutputProcessor   │                 │
   │                  │ (CTC decode +       │                 │
   │                  │  detokenization)    │                 │
   │                  └─────────────────────┘                 │
   └──────────────────────────────────────────────────────────┘

3. Internal Structure¶

File	Class	Role
`engine.py`	`ASREngine`	Top-level façade. Owns one of each subsystem and the step loop. Use `transcribe(...)` for streaming and `transcribe_offline(...)` for batch.
`config.py`	`EngineConfig`	Unified dataclass aggregating model / cache / feature / decoding / detokenization settings. Auto-detects SentencePiece model and `units.txt`.
`request.py`	`Request`, `RequestOutput`, `RequestState`	Single-request representation, output container, lifecycle enum (`WAITING → RUNNING → FINISHED`).
`scheduler.py`	`Scheduler`, `SchedulerOutput`	Dynamic-batching admission control, length bucketing, and offline micro-batch partition (`split_offline_batch`: count/preferred, padded-frame, or sequence-packed). See `scheduler.md`.
`input_processor.py`	`InputProcessor`	Audio loading, batched GPU fbank/MFCC, streaming chunk-split, CMVN-free (CMVN is in the model).
`model_runner.py`	`ModelRunner`	Wraps `ConformerModel`. Owns the cache managers; runs `forward_offline`, `forward_streaming_step`, and the batched paged path.
`executor/offline.py`	`OfflineExecutor`	Runs each scheduler-partitioned micro-batch (fbank → forward → decode → finalise) back-to-back; sequence-packed forward when `enable_sequence_packing` is set.
`executor/streaming.py`	`StreamingExecutor`	Chunk-by-chunk streaming with paged KV cache; partial outputs per tick, final on drain.
`output_processor.py`	`OutputProcessor`	CTC decode (GPU prefix beam / GPU WFST) and SentencePiece-or-units detokenization.

4. Core Algorithms and Workflows¶

4.1 Offline transcription (single pass)¶

add_request(audio, streaming=False):
    waveform = load + scale(audio_scale)
    num_frames = exact Kaldi snip_edges count
    enqueue Request in scheduler._offline_waiting

step():
    sched = scheduler.schedule()
    if sched.offline_batch:
        outputs = OfflineExecutor.run(sched.offline_batch)

OfflineExecutor.run then:

Ask the scheduler to partition the batch into micro-batches (Scheduler.split_offline_batch: length-bucketed / padded-frame / packed).
Run batched GPU fbank over each micro-batch.
Run micro-batches back-to-back on the default stream (fbank → forward → decode → finalise); with enable_sequence_packing each micro-batch is one gapless varlen-attention packed row.
Restore the original input order before returning.

Where the waveform comes from matters to step 2. A caller that fills ASREngine.new_audio_buffer(n) — page-locked host memory, which is what the Rust front-end does after the codec — lets collate DMA each row straight into the padded device batch. Anything else is packed into a pinned staging buffer first, because B separate copies out of pageable memory would each synchronise the stream before staging themselves. The pack is a second full copy of the audio, and skipping it measures 1.13–1.19× on this path; max_pinned_audio_seconds bounds how much the engine will page-lock per request, and max_pinned_audio_total_seconds (default 3600 s, ~230 MB at 16 kHz) how much across requests at once: that memory is process-global, and PyTorch's host cache keeps every block it frees. Past either cap a request gets ordinary heap memory, and so does one whose page-lock fails; the next request is offered pinned memory again. An idle engine also empties the host cache when it holds more than twice the total cap above its baseline. That happens after bursts of different lengths, since cached blocks are never reused across power-of-two size classes. It is never emptied unconditionally: page-locking afresh costs the next request 0.5-8 ms.

4.2 Streaming transcription (chunk-by-chunk)¶

add_request(audio, streaming=True) does no fbank work:

prepare_streaming(req):
    waveform = load + scale
    samples_per_chunk = ceil(stride * frame_shift_samples)
    audio_chunks = deque of CPU float32 chunks
    samples_enqueued = total
    num_frames = exact Kaldi snip_edges count   # for length bucketing
    audio_final = True                          # all audio attached up-front

Per step:

1. schedule()
   → newly_admitted streaming → ModelRunner.allocate_stream
2. For every running stream with pending audio:
       extract_streaming_batch  → batched fbank on _feat_stream
       (one GPU kernel call across all streams)
       record event; main stream waits before reading feature_buffer
3. For every stream with a full encoder window in feature_buffer:
       group by offset
       _forward_batched_paged(group)   if all paged & full window
         else
       _forward_single(req)            for partial / final / mismatched
4. For every stream whose audio is exhausted and feature_buffer is drained:
       output = OutputProcessor.finalize_streaming
       free_stream + scheduler.finish_request

4.3 Batched paged forward (`_forward_batched_paged`)¶

Pre-condition: every request in group has identical offset, a full window of features, and is using paged attention.

1. att_mgr.prepare_chunks_batched([sid for sid in group])
       # one BlockPool.allocate(B), B writes into per-stream block_table

2. block_tables, cache_seqlens = stack each stream's view
   batched_bt = cat(block_tables, dim=0)        # (B, blocks_per_seq)
   batched_cs = cat(cache_seqlens, dim=0)       # (B,)

3. batched_caches = [PagedKVCache(...) per layer with shared bt/cs]

4. xs   = torch.stack(feature_chunk_per_req)    # (B, window, F)
   cnn  = stack per-stream cnn_cache or placeholder

5. log_probs, new_cnn = model.forward_chunk_paged(xs, offset, ...)

6. For each stream in group:
       commit_chunk_paged(actual_frames, new_cnn[:, b:b+1])
       offset += actual_frames; feature_cursor += stride
       results[req.request_id] = log_probs[b:b+1]

This is the single biggest streaming throughput lever — the per-layer matmuls were launch-bound at B=1. Batching the lockstep cohort amortises ~num_layers × (linear + attention + conv + ffn) kernel launches across all in-flight streams.

4.4 Chunk-by-chunk streaming with `feed_chunk`¶

For real-time serving where audio is produced incrementally:

rid = engine.add_streaming_request()
while not eof:
    chunk = mic.read(...)
    engine.feed_chunk(rid, chunk, is_last=eof)
    outs = engine.step()       # may return partial RequestOutput

feed_chunk:

Looks up the request via scheduler.find_request(rid) (O(1)).
Appends to req.audio_chunks and updates samples_enqueued.
Sets req.audio_final = is_last so the engine knows when to flush.

The engine tolerates feed-before-admission (chunks queue up) and feed-after-admission (consumed by the next step).

Chunk size is the client's choice. Each step, a stream takes queued chunks until it has one encoder stride of new audio (stride × frame_shift samples, 640 ms at the defaults) or its queue is empty, because the encoder advances every stream by one stride per step. It used to take one chunk per step, so a client sending 20 ms frames needed 32 steps per encoder window and fell behind real time for good once a step took longer than 20 ms.

4.5 Step loop (annotated)¶

def step(self) -> List[RequestOutput]:
    sched = self._scheduler.schedule()
    outputs = []

    # 1. allocate cache for newly admitted streaming requests
    for req in sched.newly_admitted:
        if req.streaming:
            self._model_runner.allocate_stream(req)

    # 2. offline batch (if any) — micro-batches run back-to-back
    if sched.offline_batch:
        outputs.extend(self._executor.run(sched.offline_batch))

    running = sched.running_streams
    if running:
        # 3. batched fbank across every active stream with pending audio,
        #    on the dedicated _feat_stream so it overlaps the previous
        #    step's encoder forward.  extract_streaming_batch orders the
        #    producer->consumer hand-off itself (it appends into
        #    feature_buffer on *this* stream); the wait below is a belt for
        #    our own read, not the protection.  See `.artifacts/known_issues.md`
        #    "Lessons that outlived their bugs".
        needs_feat = [r for r in running if r.has_pending_audio]
        if needs_feat:
            self._input_processor.extract_streaming_batch(
                needs_feat, cuda_stream=self._feat_stream,
            )
            torch.cuda.current_stream().wait_stream(self._feat_stream)

        # 4. for every stream with a full encoder window, run forward.
        #    Streams with shared (offset) take the batched paged path;
        #    others fall back to per-stream forward.
        ready = [r for r in running if r.has_ready_encoder_chunk(window)]
        if ready:
            log_probs_map = self._model_runner.forward_streaming_step(ready)
            for req in ready:
                lp = log_probs_map.get(req.request_id)
                if lp is not None:
                    outputs.append(
                        self._output_processor.decode_streaming_chunk(req, lp)
                    )

        # 5. finalise streams whose audio is exhausted and feature
        #    buffer drained (may happen the same step as the last forward)
        for req in list(running):
            if (not req.has_pending_audio) \
                    and (not req.has_ready_encoder_chunk(window)):
                final = self._output_processor.finalize_streaming(req)
                req.output = final
                outputs.append(final)
                self._model_runner.free_stream(req)
                self._scheduler.finish_request(req.request_id)

    return outputs

5. Data Flow¶

The engine is waveform-only — file decoding happens at the entry point (the serving front-end via oasr-asr, or the bench/test harness), never in the engine. audio_scale is applied on the GPU after padding in collate (offline).

Request.audio = waveform (Tensor / ndarray, at the model sample rate, or None)
        │
   (offline)                        (streaming)
        │                                │
        ▼                                ▼
prepare_offline                    prepare_streaming
  → canonicalises Request.audio      → Request.audio_chunks
    in place (1-D f32 CPU)           → Request.num_frames (exact estimate)
  → Request.num_frames
        │                                │
  collate: pad + scale          append_streaming_chunk: scale (CPU)
  on the GPU (after padding)         per chunk
        │                                │
        └──────────────── Scheduler ─────┘
                              │
            ┌─────────────────┴──────────────────┐
            ▼                                    ▼
     OfflineExecutor                    streaming step:
       collate/cpu                    extract_streaming_batch  → _feat_stream
       _run_stage:                          (writes feature_buffer)
         forward_offline                  forward_streaming_step
         decode_offline                   _forward_batched_paged | _forward_single
                                          decode_streaming_chunk → CTC
            │                                    │
            ▼                                    ▼
                       RequestOutput
                       (text, tokens, scores, finished)

6. Configuration Options¶

EngineConfig aggregates every knob. Key groups:

Model and runtime¶

Field	Default	Description
`ckpt_dir`	`""`	A directory in any supported checkpoint format, or a Hugging Face Hub repo id. See checkpoints.md.
`checkpoint_name`	converter default	Filename inside `ckpt_dir`.
`device`	`"cuda"`	Target device.
`dtype`	`torch.float16`	Model + cache precision.
`quantization`	`None`	Online quantization of the model's GEMM layers: a scheme name (`"int8_w8a8"`) or `{"scheme": ..., "ignore": "glob,...", "module_types": "linear,pointwise_conv", "kv_cache": ...}` (`QuantizationConfig`; `module_types` picks the layer kinds, and a convolution kind needs `oasr quantize`). `None` follows the checkpoint — one written by `oasr quantize` is served quantized with no flag; naming one of its schemes is a no-op, and another scheme, `ignore` or `module_types` raises (they select layers for online quantization, so they also need a `scheme`). `kv_cache="fp8_e4m3"` stores an AR decoder's KV cache in fp8 with per-(token, head) scales, independently of `scheme` (cache_manager.md §10.6). See quantization.md.
`audio_scale`	`32768.0`	Multiplied into the float waveform. Per framework — adopted from `FeatureSpec.audio_scale` unless set explicitly. See features.md.
`service_mode`	`"streaming"`	Pins the engine to `"streaming"` or `"offline"` for its whole lifecycle; mismatched requests are rejected at admission. `"offline"` builds no streaming backend and no paged pool.
`use_cuda_graphs`	`True`	Capture the steady-state streaming encoder forward (`GraphedEncoderForward`, one graph per shape key). Also the master gate for the offline capture below.
`streaming_graph_cache_growth`	`1.5`	Growth ratio of the streaming `cache_t1` graph-key ladder above 512 frames. `1.0` restores the legacy flat 64-frame rounding, which makes that axis unbounded under `num_left_chunks=-1`. See §9.7.
`streaming_graph_batch_ladder`	`None`	Batch widths streaming encoder graphs are captured and pre-warmed at. `None` = powers of two up to `max_batch_size` with the padding lane on, every width `1..max_batch_size` with it off.
`streaming_graph_pad_batch`	`None` (auto)	Pad a cohort up to the next ladder width through the padding lane (§9.8). Auto = on when the streaming cache has more than one rung (off for a trained fixed window).
`streaming_graph_max_shapes`	`512`	Ceiling on `widths x cache rungs` pre-warmed at construction (~25 ms and a few MiB each), and on live captures. Past it the ladder is truncated to the low widths plus the cap, and the rest capture lazily.
`use_offline_cuda_graphs`	`True`	Capture the offline encoder forward by `(B_bucket, T_bucket)` (`GraphedOfflineForward`). See §9.7.
`offline_graph_batch_buckets`	`None`	Batch widths to capture at. `None` takes `preferred_batch_size` when set, else powers of two up to `max_batch_size`. Set it below `max_batch_size` to exclude the widest batches, where capture is a small net loss.
`offline_graph_frame_granularity`	`64`	Time-axis rounding, in feature frames. Overridden to `1` — exact, so no padding — for a fixed-window frontend.
`offline_graph_max_frames`	`4096`	Refuse to capture past this padded width; longer requests run eager.
`offline_graph_max_captures`	`64`	Live captures. Past this, new shapes run eager and are counted.
`long_form`	`False`	Fan a request longer than a fixed-window frontend's window out into consecutive windows, decode them through the normal batched path in parallel, and stitch one output. The caller sees one request id, so HTTP/gRPC need no change. `long_form_overlap_seconds` plus a word-level overlap merge recover most of the boundary accuracy: the words a window's head repeats from the window before are dropped from the text, the word timings, the tokens and their timestamps alike, and speech segments are cut at the middle of the shared audio. Only where windows share audio: a VAD seam, cut in silence, keeps every word on both sides.
`recycle_streaming_history`	`False`	At the streaming cache ceiling, recycle the oldest KV block instead of finalising the request with `finish_reason="length"`.
`overlap_admit`	`False`	Offline only (`oasr-server --overlap-admit`): run each request's admission prep (waveform canonicalisation, frame-count stamp) on a daemon thread, so it overlaps `step()`. The caller gets its id back before the prep runs, so a prep failure comes back as an output with `error_stage="offline_prep"`, and `abort_request` also withdraws requests still being prepared. GIL-bound: it helps only with a deep admission backlog, and can cost a little at low concurrency.
`offline_decode_overlap`	`True`	Collect a micro-batch's queued decode read-back on the next tick, behind that tick's forward (§9.4b). Needs `offline_collate_prefetch`; families without `decode_offline_async` are unaffected.
`gc_freeze`	`True`	Freeze the Python heap after construction so full collections skip it (§9.4c). Process-global and refcounted.

Autoregressive decode¶

Only read by incremental decode families — see decoding.md.

Field	Default	Description
`decode_method`	model default	Selects among `model.capabilities`; validated at construction.
`decode_options`	`{}`	Per-family knobs, resolved against the strategy's `options_cls`. An unknown key raises.
`decode_steps_per_tick`	32	Batched decoder steps allowed per engine tick.
`max_tick_ms`	25	Wall-clock cap per tick. This is the knob that bounds latency, not the step count. `0` disables.
`max_decode_slots`	`max_batch_size`	In-flight AR requests before admission pauses.
`decode_admit_window_ms`	0	Coalesce near-simultaneous arrivals into one decode group. Costs up to one window of first-token latency for an isolated request, so it is off by default.
`decode_kv_budget_gib`	derived	Byte ceiling for AR decoder KV. `0` disables the budget; `None` derives it — see §6.1.

Streaming chunking¶

Field	Default	Description
`chunk_size`	16	Encoder output frames per chunk. Must match training.
`num_left_chunks`	-1	Past chunks kept in attention cache (-1 = unlimited).

Batching¶

Field	Default	Description
`max_batch_size`	32	Encoder forward `B`. In streaming mode caps the running pool; in offline mode is the GPU forward width. Offline admission per `step()` is one length-bucketed batch of up to `max_batch_size`, run as a single forward.
`preferred_batch_size`	`None`	When set, scheduler snaps streaming admission and offline micro-batches to one of these sizes; engine pre-warms the encoder CUDA-Graph cache at each value; defaults `feature_graph_batch_buckets`. `max_wait_time` is the escape valve. See scheduler.md §4.6.
`length_bucket_ratio`	0.0	Length-aware policies (`window`, `bucket`, `sjf`): floor on `min_len/max_len` in an offline batch.
`max_offline_pad_ratio`	4.0	Length-aware policies: hard cap on padded/useful compute.
`max_wait_time`	0.2	Starvation bound (seconds).
`schedule_policy`	`"window"`	`fcfs` / `window` / `bucket` / `sjf`. See scheduler.md §4.1.1.
`length_window_factor`	4	`window` policy: batches are chosen from the oldest `max_batch_size × factor` requests.
`streaming_cohort_admit`	`True`	Length-sort waiting streams before admission. Streams that have waited `max_wait_time` go first, oldest first; without that bound the sort is last-in, first-out under load, because a waiting stream's length is the audio it has buffered. See scheduler.md §4.2.

Paged KV cache¶

Field	Default	Description
`max_num_blocks`	`None` (derived)	Total physical blocks in the shared pool. `None` derives it from free VRAM, capped at what the block table can address — see §6.1. An explicit value caps each stream at `max_num_blocks // max_batch_size` blocks with unlimited history; the old default of 2048 cut every stream at ~41 s at `max_batch_size=32`. Inert in `service_mode="offline"`, which builds no pool.
`gpu_memory_utilization`	0.90	Share of the device the engine may occupy in total (weights + caches + activations) when it derives a capacity. Read only when something is left to derive.
`block_size_frames`	16	Frames per block (= chunk_size by default).
`max_blocks_per_seq`	512	Block-table width. With unlimited history this, times `max_batch_size`, is also the ceiling a derived pool is capped at — blocks past it cannot be addressed.
`use_paged_cache`	`True`	False falls back to dense `forward_chunk`.

6.1 VRAM-aware capacity sizing¶

Two capacities are memory rather than compute, and both can be left to the engine (oasr/engine/memory.py):

Left unset	Derived as
`max_num_blocks=None` (streaming)	`available / bytes_per_block`, capped at `max_batch_size × blocks_per_seq`
`decode_kv_budget_gib=None` (AR families)	`available`, clamped up to one row so a tight card still admits work; a row's bytes follow `quantization.kv_cache` (fp8 + scales ≈ 0.52x of fp16)

where

available = total × gpu_memory_utilization − resident − activation_reserve

resident is read from the driver (torch.cuda.mem_get_info) after the model is on the device, so it covers the weights, the CUDA context and anything another process holds — there is no separate "weights" term to get wrong. activation_reserve is 1.5 × a measured probe forward at the widest shape the engine will run (one chunk window at max_batch_size in streaming mode; the frontend's fixed window, else 30 s, in offline mode), floored at 256 MiB. The unspent 1 − utilization is the headroom for what the probe cannot see: CUDA graph capture pools, an AR family's prefill transient, allocator fragmentation.

The derivation is logged in full (paged KV pool derived from VRAM: … | total=… resident=… cap=… activation_reserve=… → available=…) so the numbers are auditable. Nothing is measured — and no probe forward runs — unless a capacity was actually left unset. Both are unset by default.

bytes_per_block and blocks_per_seq are in the geometry the pool is built with (EngineConfig.streaming_cache_geometry). For most encoders that is block_size_frames per page and, with unlimited history, max_blocks_per_seq pages per stream. An encoder with a trained attention window (Nemotron) instead pages by chunk_size and retains ceil(window / chunk_size) + 1 pages, so its derived pool is exactly that per stream: 160 pages (0.23 GiB) at 32 streams, where sizing from the raw config gave 16384 (24 GiB).

On a non-CUDA device there is nothing to measure, so max_num_blocks=None sizes the pool to its ceiling, max_batch_size × blocks_per_seq. On CUDA, when not even the minimum viable pool fits, construction fails with the arithmetic and the levers (max_batch_size, num_left_chunks, gpu_memory_utilization, or an explicit max_num_blocks) rather than deriving a pool that OOMs at allocation or degrades every transcript. decode_kv_budget_gib=0 turns the byte budget off outright; None derives it.

Worth knowing: on a large card the pool is usually capped by max_blocks_per_seq × max_batch_size, not by VRAM (a 32 GiB card had 27 GiB available and took 3 GiB). The engine logs when that happens — the remaining lever is then the per-stream ceiling, not the pool size.

Feature extraction¶

Field	Default	Description
`feature_config`	`FeatureConfig(dither=0.0)`	80-dim log-mel FBANK at 16 kHz; deterministic. Extracted on the GPU via the batched fbank/mfcc kernels.

Decoding and detokenization¶

Field	Default	Description
`decoder_type`	`"ctc_cuda"`	`ctc_cuda` (GPU CTC beam) / `ctc_wfst` (in-tree GPU WFST).
`ctc_decoder_config`	`GpuDecoderConfig()`	GPU CTC config (beam, blank ID, thresholds).
`wfst_decoder_config`	`DecoderConfig(search_type="wfst")`	GPU WFST decoder config.
`fst_path`	`None`	Required for `ctc_wfst`: a prebuilt `.img`, or a k2 `HLG.pt` exported at load.
`sentencepiece_model`	auto-detected	`.model` in `ckpt_dir`.
`unit_table`	auto-detected	`units.txt` / `words.txt` fallback.

Derived properties¶

subsampling_rate = 4              # Conv2dSubsampling
right_context    = 6
decoding_window  = (chunk_size - 1) * 4 + right_context + 1   # input frames per chunk
stride           = subsampling_rate * chunk_size              # frame advance
required_cache_size = chunk_size * num_left_chunks            # dense mode

That formula describes a centred subsampling front-end, which needs a receptive field beyond its chunk. An encoder whose front-end it does not describe overrides BaseEncoder.streaming_geometry(chunk_size) and returns its own (decoding_window, stride) — for a cached causal subsampling those are equal, because a chunk consumes exactly chunk_size * subsampling_rate frames with no lookahead. That hook is also where an encoder refuses a chunk size it cannot serve; the backend calls it once, at construction, so an unserviceable chunk_size fails there rather than drifting silently at request time.

build_cache_config(model_config) derives a CacheConfig from the loaded encoder dimensions — including, when the encoder declares a trained fixed_attention_window, the retained cache itself (see docs/cache_manager.md §10.2: how much history the model may attend to is part of its mask, not an operator preference, so num_left_chunks is ignored there and a warning says so).

7. Usage Examples¶

7.1 Offline batch transcription¶

import torchaudio
from oasr.engine import ASREngine, EngineConfig

engine = ASREngine(EngineConfig(ckpt_dir="/path/to/ckpt", service_mode="offline"))

# The engine is waveform-only — decode files yourself (here in the harness;
# in serving, oasr-asr does this at the entry point).
def wav(p): return torchaudio.load(p)[0].squeeze(0).float()

text  = engine.transcribe_offline(wav("audio.wav"))
texts = engine.transcribe_offline([wav(p) for p in ("a.wav", "b.wav", "c.wav")])

# Everything a decode produces, not just the string:
from oasr.engine import DecodingOptions

(out,) = engine.transcribe_outputs(
    [wav("audio.wav")], streaming=False, decoding=DecodingOptions(word_timestamps=True)
)
for w in out.words or ():
    print(f"{w.start:6.2f}-{w.end:6.2f}  {w.confidence:.2f}  {w.word}")

7.2 Streaming, attached audio¶

import torchaudio
from oasr.engine import ASREngine, EngineConfig

engine = ASREngine(EngineConfig(ckpt_dir="/path/to/ckpt"))

def wav(p): return torchaudio.load(p)[0].squeeze(0).float()

# All audio handed in up-front; engine splits it into chunks and runs to
# completion (real-time clients instead use add_streaming_request + feed_chunk).
texts = engine.transcribe([wav(p) for p in ("a.wav", "b.wav", "c.wav")], streaming=True)

7.3 Streaming, chunk-by-chunk feed (real-time serving)¶

rid = engine.add_streaming_request()
for chunk in mic_chunks(samples_per_chunk=4000):
    engine.feed_chunk(rid, chunk, is_last=False)
    outputs = engine.step()           # may return partial transcripts
    for o in outputs:
        if o.request_id == rid and not o.finished:
            print("partial:", o.text)

engine.feed_chunk(rid, last_chunk, is_last=True)
final = engine.run()                  # drain until finalised

7.4 Mixed offline + streaming on one engine¶

engine = ASREngine(EngineConfig(ckpt_dir=..., max_batch_size=8))
for path in offline_paths:
    engine.add_request(path, streaming=False)
for path in streaming_paths:
    engine.add_request(path, streaming=True)

results = engine.run()                # one engine handles both pools

7.5 Selecting GPU CTC¶

engine = ASREngine(EngineConfig(
    ckpt_dir=...,
    decoder_type="ctc_cuda",
    ctc_decoder_config=GpuDecoderConfig(beam_size=10, blank_threshold=0.95),
))

8. Error Handling and Edge Cases¶

Situation	Behaviour
`feed_chunk` for unknown / finalised id	`KeyError` from `Scheduler.find_request → None`.
Block pool exhaustion	`RuntimeError` from `BlockPool.allocate`. Currently fatal — size `max_num_blocks` for worst case, or set it to `None` and let the engine derive it (§6.1). With eviction on (`num_left_chunks >= 0`) the invariant is checked at construction instead.
Audio shorter than one window with `audio_final=True`	`_forward_single` flushes whatever frames remain (special `is_final_window` path).
`chunk.size(1) < context` (less than `right_context+1` input frames) and not final	Skipped — the engine waits for more audio.
First chunk of paged stream	`prepare_chunks_batched` lazily allocates `block_table` / `cache_seqlens` before writing the first physical block.
First-chunk CNN cache is `(0,0,0,0)` placeholder	Batched forward passes the placeholder; the model handles per-layer initialisation.
`ckpt_dir` missing required files	Fails fast inside `load_wenet_checkpoint` (out of engine scope).
Unsupported `decoder_type`	`OutputProcessor` raises during decode.
Aborting a running stream	`abort_request` removes from the scheduler and frees the cache via `ModelRunner.free_stream`.
Audio sample-rate mismatch	`InputProcessor.load_audio` resamples to `sample_rate`.
Force-flush firing inside an offline batch	All bucket guards skipped — batch may be highly padded. Acceptable to bound starvation.
More waiting streams than free slots	Admission is bounded by `max_wait_time`: overdue streams are admitted first, oldest first, and only the rest are length-sorted.

Thread safety¶

ASREngine is thread-safe. Every public entry — add_request, add_streaming_request, feed_chunk, abort_request, step, run, num_running / num_waiting, and transcribe (which composes them) — acquires a process-wide re-entrant threading.RLock (run() re-enters to call step()). Uncontended cost is negligible next to GPU work.

It guards two pieces of mutable state:

the scheduler queues — _streaming_waiting, _offline_waiting, _running, _index in oasr/engine/scheduler.py;
per-request audio — request.audio_chunks and request.audio_final, mutated by feed_chunk and concurrently read and popped by step().

The lock is coarse on purpose. step() holds it for the full step, so a concurrent feed_chunk from another thread waits up to one step. That is acceptable because the GIL serializes Python anyway and CUDA releases it during the forward — the lock makes data-structure consistency explicit rather than buying parallelism. A finer split (scheduler-only lock, per-request audio locks, dropping the lock during forward/decode) is possible but has never been the bottleneck.

Do not add threads inside one engine to gain throughput — CPython's GIL means it cannot work, and the engine is GPU-bound and already batching. Horizontal scale is one process per GPU; see serving.md.

Under serving, the PyO3 dispatcher is the only Python caller and runs single-threaded; HTTP and gRPC handlers stay on tokio and never touch the GIL.

Concurrency tests live in tests/engine/test_concurrent.py behind the concurrent marker and cover a 16-thread stress on add_streaming_request, feeders racing aborters on a shared request-id pool, and concurrent readers of the queue-depth properties:

pytest tests/engine/test_concurrent.py -m concurrent -v

Failure isolation¶

An exception escaping step() becomes an INTERNAL error for every in-flight request, so both executors contain it instead.

Executor	Policy	Why
`OfflineExecutor`	Re-runs a failed micro-batch one request at a time, so only the culprit is rejected — except on OOM, where the batch fails whole	Retrying under memory pressure is how one over-large request cascades
`StreamingExecutor`	Fails the whole ready cohort	The batched forward has no per-stream boundary, and retrying per stream would double-commit the chunk for streams that already advanced
AR decode (`IncrementalArStrategy.advance`)	Fails the decode group whose step raised, returns its decoder KV, and keeps advancing the others; rows that finished before the failure keep their finals	A group's rows share one batched forward, and its partly-written step cannot be replayed. A dense merge that fails part-way fails the group it was merging into, whose cache it had begun to release

Terminal outputs carry finish_reason="error" plus RequestOutput.error_stage, which becomes the stage label on the oasr_requests_failed_total metric (offline_prep, offline_forward, offline_oom, prefill_oom, decode_step, decode_oom, streaming_forward, streaming_features) — see serving.md. offline_prep is specific to overlap_admit. Without it, a waveform that cannot be prepared raises to the admitting caller. With it, the caller already has the request id, so the failure is answered as an output.

9. Performance Considerations¶

Streaming throughput is dominated by _forward_batched_paged. Pick max_batch_size to fit the GPU. Larger batches amortise launch overhead but spread per-step latency.
A CUDA stream overlaps streaming fbank and forward. _feat_stream runs the streaming fbank kernel; the default stream waits on a recorded event before reading feature_buffer. The offline executor runs its micro-batches sequentially (GPU fbank → forward → decode), so it needs no extra stream — and measurably does not want one: at every batch width the offline step thread, not the GPU, is the critical path, so a prefetch only moves work it cannot remove.
Per-step host→device metadata goes through oasr.utils.staging.to_device. Slot ids, block ids and per-row offsets are built host-side every step, and CUDA synchronises the stream before staging a pageable copy — so the obvious torch.tensor(ids, device="cuda") costs whatever is queued, not the 40 bytes it moves. One sat immediately after the streaming encoder forward, which meant the host waited out that forward on every step and could never run a step ahead of the device. to_device builds the tensor in pinned memory and copies it non_blocking=True; readers on the same stream need no further ordering, readers on another stream need the usual wait_stream.
Offline batches are length-sorted within a window (schedule_policy "window", the default; scheduler §4.1.1). With the forward graph-captured, a batch's padding is GPU time: FIFO batches of 32 LJSpeech utterances computed ~1.48× the useful frames. Choosing each batch as the least-padded length-run that contains the oldest request, from the oldest 4 × max_batch_size, measured 1.27× offline Conformer throughput (1.32× on the transducer, whose greedy loop also runs to the batch's longest row) with the accuracy gate unchanged. max_batch_frames, max_offline_pad_ratio and length_bucket_ratio bound its runs as they bound the other policies' batches. 4a. A Python scalar written through an index into a CUDA tensor blocks the host. t[i, j] = 0 and t[idx_t] = 0 both lower to a pageable 4-byte host-to-device copy, which waits for everything queued on the stream — measured 41.6 ms of host blocked behind queued GEMMs, against 0.04 ms for the same write with a device scalar, fill_ on a slice, or a staged index. The cache manager's admission and eviction paths did this once per admitted stream and once per step for every capped-history stream; see AttentionCacheManager._zero. 4b. The offline decode read-back is collected a tick late (offline_decode_overlap, with offline_collate_prefetch). A blocking read-back ends a tick with the GPU's queue empty, so its host tail — token extraction, detokenization, finalisation — ran in front of an idle GPU (~1.1 ms of a ~9.3 ms Conformer tick at B = 32). A family that implements DecodeStrategy.decode_offline_async (GPU CTC) queues its read-back into pinned memory behind an event; the executor collects it on the next tick, after that tick's forward and decode are queued. 1.09–1.17× on Conformer offline (B = 64 … 1), transcripts unchanged. Two details carry it: the CTC offline launcher no longer reads the selected-frame count back to size its tile loop (the fused kernels mask per row, so it loops to the input length), and a batch with nothing staged behind it is collected in its own tick — otherwise a lone request's result would wait a tick and an issue-only tick would read as idle to the serving dispatcher. 4c. The Python heap is frozen after construction (gc_freeze). A full garbage collection over the constructed engine's ~440k tracked objects stops the step loop for 162–170 ms, a tick in which every stream stalls; frozen, collections walk only what was allocated since. Refcounted across engines in one process (oasr.engine.gc_freeze); a dropped engine still returns its device memory through reference counting.
NVTX profiling. The step loop is annotated with nvtx_push("engine.step") → schedule / allocate_stream / offline_batch / extract_fbank / forward_streaming / decode_streaming / finalize_streams. Capture with nsys profile/ncu to get per-stage timing without touching code.
The offline forward is CUDA-graph captured too (use_offline_cuda_graphs, oasr/engine/offline_graph.py). It was not, for a long time, and the gap was large: at B=1 one offline forward issued ~437 kernel launches for 0.99 ms of GPU work and took ~9.7 ms, so eager wall time stayed flat from B=1 to B=32 — the signature of a host-issue-bound path. Capture is worth 7.1x / 3.1x / 1.3x at B=1 / 8 / 32 on the u2++ Conformer and 6.7x / 3.4x / 1.5x on Zipformer-large, transcripts unchanged. At B=64 it is a ~2% loss: the host already keeps up there, so the copy into the static input buffer and the clone out are pure additions. Exclude the widest widths with offline_graph_batch_buckets if you serve them.

Three things about the shape space are worth knowing before tuning it:

It is already small. 200 mixed-length utterances at max_batch_size=32 produce 3 captured shapes, not hundreds, because the scheduler length-sorts and bucket-fills before the forward ever sees a batch.
preferred_batch_size is not the lever it looks like. It changes the captured shape count not at all (the B axis was already 2 wide), and a single preferred value costs a max_wait_time stall on every queue drain — measured 3.4x slower end-to-end, recovered entirely by lowering max_wait_time. Use it for admission shaping, not for capture control.
Sequence packing enlarges the shape space rather than collapsing it. The intuition says packing pins B at 1 and bounds T by max_packed_frames — one graph for everything. Measured, the same 200 utterances that need 3 keys unpacked need 6 packed, because a packed key is (S, T_total, max_seg, bias_size) and the last two are a sum and a sum-of-squares of segment lengths. It is capturable — build_packed_layout_device rebuilds the layout with no host round-trip and no data-dependent shape — and fully captured it beats the non-packed path by 13–15 % at B=32 and 7–13 % at B=64, while losing ~10 % at B=8. Wiring it needs a packed-specific 4-axis key and bucket padding of the packed row; that padding moves real segments by ~1.9e-1 in bf16 past a kernel-selection threshold, so it would be bit-exact to itself rather than to B=1. Until then packed batches run eager and are counted under fallback_failed.

Padding is applied on both paths, not just the captured one. An encoder need not be invariant to trailing padding and the shipped ones are not — Zipformer's SimpleDownsample fills its last window by replicating the final frame — so if only the graph padded to a bucket, an utterance would decode differently depending on whether its shape happened to be captured. See GraphedOfflineForward.pad_time.

The streaming graph ladder is pre-warmed exhaustively, and that is what flattens the tail. The streaming graph is keyed (B_active, T_input, cache_t1_bucket) and a miss captures a graph on a live tick, at ~30 ms. Both axes used to be under-covered:
cache_t1 was rounded to a flat 64 frames. With the default num_left_chunks=-1 that axis is unbounded, so a long stream reached a new bucket every 64 encoder frames and captured a graph for as long as it ran — measured on a 120 s stream, 48 captures over 191 ticks, a capture on a quarter of all ticks, and it never settled (24 in each half of the run). streaming_graph_cache_growth makes the rungs geometric above 512 frames, turning that into 11–16 rungs that cover a stream's whole life. The trade is a wider over-read of the paged K/V, measured at ~1 % of streaming throughput, because a replay only grows 1.04 → 1.41 ms from an empty cache to 82 s of history while an eager forward is flat at ~12 ms.
with history eviction (num_left_chunks >= 0, or recycle_streaming_history) the axis was unbounded for a different reason: cache_t1 was the stream's offset, which keeps growing after the cache stops. It is the cached length now (AttentionCacheManager.cached_frames), which eviction holds at max_logical_blocks - 1 blocks, and an evicting stream's ladder is built to exactly that, so it is a handful of rungs the pre-warm covers. Positions are unaffected: they start at offset - cache_seqlens per row.
the batch axis was pre-warmed at {1, max_batch_size} only, so every active width between them captured on first appearance. It now defaults to every width.

The rung is also the width the eager fallback runs at (a saturated cache, a failed capture, graphs off): it is chosen before the graph/eager branch, like the offline padding above. The fallback used to run at the exact cached length, which sizes the positional projection and the bias matmul differently, and in fp32 a 2.5-minute stream then differed from its graph replay on 909 of 991 chunks.

Together: p99 31.4 → 3.2 ms on long streams and max 32.6 → 4.0 ms at sustained concurrency 16, with zero live-tick captures in both. The cost is startup — ~25 ms and a few MiB per shape — bounded by streaming_graph_max_shapes. Transcripts are unchanged.

The width axis is a ladder, not every width (streaming_graph_pad_batch, auto). Every width × every cache rung is 576 shapes at max_batch_size=64 — over the 512 budget, so the draining tail ran eager. A cohort is now padded up to the next power-of-two width through a padding lane: one slot row past max_batch_size whose block table points only at a scratch block outside the free list (BlockPool(scratch_blocks=1), AttentionCacheManager(pad_lanes=1)), fed zeros from the lane's own cache length. Every encoder op is row-local, so padding rows compute garbage only for themselves; padding happens before the graph/eager branch, so a shape that falls back decodes exactly as one that replays. At max_batch_size=64: 1.59× streaming throughput, p99 18.2 → 6.1 ms, 512 → 63 shapes, startup 22 → 7.6 s, graph VRAM 3.0 → 0.8 GiB, transcripts byte-identical. Auto pads only when the cache has more than one rung: a trained fixed window (Nemotron) has exactly one, so its exact widths are cheap captures and padding would only add compute.

A stream's last full window is batched. It used to take the per-stream fallback, and for an encoder with a declared geometry the finalize pad makes every stream's last window exactly one window long — one serial B=1 forward per finishing stream (~400 ms of a 664 ms Nemotron backlog). Batched, and consumed whole as the fallback consumed it: Nemotron streaming 1.78×.

Pool sizing. By default (max_num_blocks=None) the pool is derived from free VRAM at construction (§6.1), which is what makes one config portable across card sizes. An explicit pool must cover max_batch_size × max_logical_blocks plus headroom: with unlimited history each stream is capped at its share, max_num_blocks // max_batch_size, and ends there with finish_reason="length". The pool's size does not cost throughput: freeing a stream's blocks costs its own blocks, not the pool's.

10. Extension Points¶

Every extension lands through a registry — subclass a base, register under a name, and the engine picks it up by configuration. Never edit the engine core to add a variant. docs/architecture.md is the authoritative map (seams table + cookbook); the summary:

New model architecture → model package under oasr/models/<arch>/ (BaseAsrModel / BaseEncoder) + a CheckpointConverter, registered via register_model. The engine touches models only through the base-class surface (forward_offline / encode_offline / forward_chunk_paged / cache_spec / capabilities / encoder.streaming_kind). See docs/checkpoints.md for the converter contract and native format.
New decode family → DecodeStrategy + @register_decode_strategy. Declare consumes ("log_probs" / "hidden" / "both") and, for label-synchronous AR families, incremental = True with the bounded begin_offline / advance(StepBudget) / has_pending protocol. Selected per deployment by EngineConfig.decode_method (validated against model.capabilities); per-family knobs go on the strategy's options_cls and per-request knobs ride on DecodingOptions (Request.decoding). Full axis: docs/decoding.md.
New streaming runtime → StreamingEncoderBackend + @register_streaming_backend, keyed by the encoder's streaming_kind. Implement allocate / forward_step / free + window geometry. An encoder whose whole cache is fixed-extent needs no new runtime: declare it (slot_state_specs + streaming_pad_value, streaming_kind="slot") and the slot runtime owns the state and graph-captures the step. Without the declaration, stack_streaming_states / unstack_streaming_states still get batched (one B = N forward) stateful streaming.
New batching / partition policy → @register_batching_policy / @register_partition_policy; select via EngineConfig.schedule_policy and the partition flags.
New tokenizer kind → Tokenizer subclass + register_tokenizer, selected by the converter-emitted TokenizerSpec.kind. See docs/tokenizers.md.
Feature extraction is checkpoint-derived: converters emit a FeatureSpec the engine materializes into FeatureConfig; an explicit EngineConfig.feature_config still overrides (with a loud mismatch warning).
Streaming preemption. Currently absent. To add it, extend the scheduler to evict a low-priority running stream, then have ASREngine call ModelRunner.free_stream and re-queue the request. No public API changes required.

11. Quick Reference¶

# construction
engine = ASREngine(EngineConfig(ckpt_dir=..., **knobs))

# request submission
rid = engine.add_request(audio, streaming=True/False, priority=0)
rid = engine.add_streaming_request()       # then feed_chunk
engine.feed_chunk(rid, chunk, is_last=False)

# per-request decoding options (n-best, AR generation cap, sampling, LLM prompt)
from oasr.engine import DecodingOptions
rid = engine.add_request(audio, streaming=False,
                         decoding=DecodingOptions(n_best=3, max_new_tokens=64))

# stepping
outputs = engine.step()                    # one tick, returns partials+finals
finals  = engine.run()                     # drain until empty

# convenience
texts = engine.transcribe(audio_or_list, streaming=...)

# mid-flight control
engine.abort_request(rid)

# observability
engine.num_running    # int
engine.num_waiting    # int