Skip to content

Cache Manager

The cache subsystem (oasr/cache/) manages all per-stream GPU state required for chunk-by-chunk Conformer inference: paged attention KV cache, fixed-size CNN left-context cache, and per-stream CTC decoder state. It is the foundation that lets the engine run many concurrent streaming requests on a single shared GPU memory pool without tearing down or re-allocating buffers between chunks.

1. Purpose and Responsibilities

The cache manager exists to:

  1. Pre-allocate all GPU memory needed for streaming KV-cache up front, avoiding per-chunk cudaMalloc calls.
  2. Page the attention KV cache so that a fixed pool of physical blocks can be shared across many concurrent streams (vLLM-style).
  3. Track per-stream state (logical block lists, frame counters, block_table / cache_seqlens tensors for paged attention) and keep it in sync with what the model writes into the pool.
  4. Evict old context once a stream exceeds its left-context window.
  5. Pool lightweight CTC decoder states so allocating a new stream does not require a fresh GPU buffer.
  6. Expose a single per-stream handle (StreamContext) so engine code can interact with all three caches without juggling stream IDs.

2. High-Level Architecture

                         ┌────────────────────────────┐
                         │      CacheConfig            │
                         │  (one shared dataclass)     │
                         └──────────────┬─────────────┘
                                        │
                ┌───────────────────────┼───────────────────────┐
                │                       │                       │
                ▼                       ▼                       ▼
     ┌────────────────────┐  ┌────────────────────┐  ┌────────────────────┐
     │     BlockPool      │  │  CnnCacheManager   │  │ CtcStateCacheMgr   │
     │ (shared K/V pool)  │  │ (per-stream tensor)│  │ (1 decoder + pool) │
     └─────────┬──────────┘  └─────────┬──────────┘  └─────────┬──────────┘
               │                       │                       │
               ▼                       │                       │
     ┌────────────────────┐            │                       │
     │AttentionCacheMgr   │            │                       │
     │ (per-stream block  │            │                       │
     │  table + seqlens)  │            │                       │
     └─────────┬──────────┘            │                       │
               │                       │                       │
               └───────────────────────┴────┬──────────────────┘
                                            ▼
                                  ┌──────────────────┐
                                  │  StreamContext   │
                                  │  (one per req)   │
                                  └──────────────────┘

BlockPool is the only object holding the large preallocated K and V tensors. AttentionCacheManager consumes block IDs from it; CnnCacheManager allocates a small per-stream tensor of its own; CtcStateCacheManager delegates to a single shared GpuStreamingDecoder engine and pools per-stream StreamState objects. StreamContext is a thin facade that bundles a stream_id with all three managers.

3. Internal Structure and Key Classes

File Class / Type Responsibility
types.py CacheConfig Master configuration dataclass (layers, heads, dims, chunk size, block size, pool capacity, dtype, device) and derived helpers (max_logical_blocks, cnn_cache_frames, kv_last_dim).
block_pool.py BlockPool Two large GPU tensors (_k_pool, _v_pool) of shape (num_layers, max_num_blocks, block_size_frames, n_kv_head, head_dim) plus a thread-safe free list of physical block IDs.
attention_cache.py AttentionCacheManager, _StreamKVState Per-stream logical-to-physical block mapping; supports both dense (commit + get_stacked_cache) and paged (prepare_chunk + commit_chunk_paged + PagedKVCache) modes; handles eviction.
cnn_cache.py CnnCacheManager Per-stream tensor of shape (num_layers, 1, kernel_size - 1, hidden_dim) overwritten in place each chunk (no paging needed — fixed size per stream).
ctc_state.py CtcStateCacheManager Single shared GpuStreamingDecoder plus a pool of StreamState objects; exposes a StreamHandle per call site so callers see the same API as a stand-alone decoder.
stream.py StreamContext Per-request facade that ties the three managers together using a single stream_id.

CacheConfig highlights

@dataclass
class CacheConfig:
    num_layers: int = 12
    n_kv_head: int = 4
    head_dim: int = 64
    hidden_dim: int = 256
    kernel_size: int = 15           # depthwise conv kernel
    chunk_size: int = 16            # encoder output frames per chunk
    num_left_chunks: int = -1       # -1 = unlimited history
    block_size_frames: int = 16     # frames per physical block
    max_num_blocks: int = 1024      # total pool size
    max_blocks_per_seq: int = 512   # block_table capacity per stream
    device: torch.device = ...
    dtype:  torch.dtype  = torch.bfloat16

Derived properties of interest:

  • kv_last_dim → head_dim * 2 — K and V are packed along the last dim in the dense API.
  • cnn_cache_frames → kernel_size - 1 — left-context frames stored by the causal depthwise conv per layer.
  • max_logical_blocks → ceil(chunk_size * num_left_chunks / block_size_frames) or None when history is unlimited.
  • blocks_per_stream — how many logical blocks one stream may hold before it is at capacity: max_logical_blocks when eviction is on, else min(max_blocks_per_seq, max_num_blocks // max_batch_size).
  • max_stream_frames → blocks_per_stream * block_size_frames — the encoder-frame ceiling per stream.

With eviction enabled (num_left_chunks ≥ 0) memory is bounded by construction: size the pool so max_num_blocks ≥ max_batch_size * max_logical_blocks, with headroom for short-lived oversubscription during cohort transitions. Or don't size it at all: EngineConfig.max_num_blocks=None derives the block count from free VRAM at engine construction and, with eviction on, reduces to checking that the retained history fits (see engine.md §6.1).

With unlimited history (num_left_chunks = -1, the default) eviction is off, so every stream has a finite ceiling — max_stream_frames, logged at INFO when the config is built (at the shipped defaults: 64 blocks ≈ 1024 encoder frames ≈ 41 s of audio at 4× subsampling). A stream that reaches it is finalized cleanly with the transcript decoded so far and finish_reason="length": AttentionCacheManager.at_capacity(stream_id) is consulted by the streaming backend before it dispatches a chunk, so the allocator can no longer raise BlockPool exhausted from inside the encoder forward (which the serving dispatcher used to fan out to every in-flight request). Raise max_num_blocks to lift the ceiling, or set num_left_chunks to bound history by eviction instead.

CacheConfig.__post_init__ rejects two geometries outright:

  • chunk_size > block_size_frames — one block is allocated per chunk, so a wider chunk would spill into the next logical block, which PagedKVCache's two-block write path reads as the following chunk's page (silent KV corruption).
  • max_blocks_per_seq < max_logical_blocks — the block table cannot address the history the config asks for.

4. Core Algorithms and Workflows

4.1 Paged KV-cache lifecycle (per stream)

Paged mode is the production path used by ASREngine. For each new chunk:

prepare_chunk(stream_id):
    if stream is new:
        allocate per-stream block_table (zeros, int32, max_blocks_per_seq)
        allocate per-stream cache_seqlens (zeros, int32, shape (1,))
    block_id = pool.allocate(1)
    state.logical_blocks.append(block_id)
    state.block_table[0, len(logical_blocks) - 1] = block_id

# Model writes K/V directly into pool[block_id] via _paged_write_kv.

commit_chunk_paged(stream_id, chunk_frames):
    state.num_committed_frames += chunk_frames
    state.cache_seqlens[0] = state.num_committed_frames
    while len(state.logical_blocks) > max_logical_blocks:
        evicted = state.logical_blocks.pop(0)
        pool.free([evicted])
        # shift block_table left by one slot
        # adjust num_committed_frames + cache_seqlens

Key invariants:

  • The model never sees a physical block ID — it walks the per-stream block_table indexed by logical block.
  • cache_seqlens[0] is the source of truth for paged attention; the manager keeps it consistent across commit and eviction.
  • Eviction is left-shift, preserving FIFO order. A stream's history is always the most recent num_left_chunks chunks.

4.2 Dense KV-cache mode (legacy / forward_chunk)

Dense mode is kept for the non-paged code path and for tests. commit() allocates one block per chunk, splits packed (K|V) into K and V slabs, writes them into the pool, then get_stacked_cache() gathers all blocks back into a contiguous (num_layers, n_kv_head, T, head_dim*2) tensor — the shape ConformerEncoder.forward_chunk expects. Eviction is identical to paged mode.

4.3 Batched paged forward path

AttentionCacheManager.prepare_chunks_batched(stream_ids) performs a single BlockPool.allocate(B) call, then writes the resulting block IDs into each stream's own block_table[0, idx]. The engine (ModelRunner._forward_batched_paged) then concatenates each stream's (block_table, cache_seqlens) views (via get_paged_state_views) into batched (B, max_blocks_per_seq) and (B,) tensors and constructs one PagedKVCache per layer that all B streams share.

This is the most performance-critical path: it amortises ~num_layers × several kernel launches across all in-flight streams.

4.4 CNN cache update

CNN cache is fixed-size per stream and overwritten in place every chunk:

allocate_stream(sid):
    caches[sid] = zeros(num_layers, 1, cnn_cache_frames, hidden_dim)

update(sid, new_cnn_cache):
    caches[sid].copy_(new_cnn_cache)   # shape-checked; in-place

4.5 CTC state pooling

allocate_stream(sid, batch, vocab_size):
    if pool not empty:
        state = pool.pop()
        decoder.reset_state(state, batch, vocab_size)
    else:
        state = decoder.create_state(batch, vocab_size)
    states[sid] = state

free_stream(sid):
    pool.append(states.pop(sid))   # state's GPU buffer is NOT freed

StreamHandle(decoder, state) lets the caller treat the (decoder, state) pair as a single object whose decode_chunk / finalize_stream / step / config mirror the stand-alone GpuStreamingDecoder API.

5. Data Flow and Component Interaction

                ┌──────────── Caller (e.g. ASREngine.step) ────────────┐
                │                                                       │
                │    1. attention.prepare_chunk(sid)                    │
                │    2. att_caches  = attention.get_paged_caches(sid)   │
                │    3. cnn_cache   = cnn.get_cache(sid)                │
                │    4. logits, new_cnn = model.forward_chunk_paged(    │
                │            xs, offset, att_caches, cnn_cache, ...)    │
                │    5. attention.commit_chunk_paged(sid, frames)       │
                │    6. cnn.update(sid, new_cnn)                        │
                │    7. ctc.get_decoder(sid).decode_chunk(logits)       │
                │                                                       │
                └───────────────────────────────────────────────────────┘
                           ▲           ▲                ▲           ▲
                           │           │                │           │
                           │ allocate  │ get_kv_view    │ blocks    │ pool/restore
                           ▼           │                ▼           │
                     ┌──────────────────────────┐  ┌─────────────────┐
                     │       BlockPool          │  │ GpuStreaming    │
                     │  K/V tensors + freelist  │  │ Decoder + state │
                     └──────────────────────────┘  └─────────────────┘

StreamContext collapses steps 1–7 into a single prepare_chunk() / commit_chunk_paged() / get_decoder() interface so the engine's hot loop reads cleanly.

6. Configuration Options

EngineConfig.build_cache_config(model_config) derives CacheConfig from the loaded model. The fields you typically tune:

Field Default When to change
chunk_size 16 Must match training. Longer chunk = higher accuracy, higher per-chunk latency.
num_left_chunks -1 -1 is unlimited; set to a positive value (e.g. 4) to bound memory in long-running streams.
block_size_frames 16 Equal to chunk_size is simplest (one block per chunk). Smaller values reduce eviction granularity but increase block-table size.
max_num_blocks None in EngineConfig Pool capacity. EngineConfig derives it from free VRAM by default, capped at max_batch_size × max_blocks_per_seq. An explicit value must cover all concurrent streams' worst-case logical block counts plus overhead; with unlimited history each stream is capped at max_num_blocks // max_batch_size blocks.
max_blocks_per_seq 512 Width of the per-stream block_table. Must be ≥ max_logical_blocks.
dtype float16 bfloat16 works on Ampere+. The K/V pool is the dominant memory consumer.

Pool memory in bytes:

2 (K and V) × num_layers × max_num_blocks × block_size_frames
            × n_kv_head × head_dim × bytes_per_element(dtype)

For a default Conformer (num_layers=12, n_kv_head=4, head_dim=64, block_size=16) with FP16 this is 192 KiB per block: 384 MiB for 2048 blocks, and 3 GiB for the 16384 a derived pool reaches at max_batch_size=32 (the block-table ceiling, 32 × 512). oasr.engine.memory.bytes_per_kv_block is the same arithmetic in code, and tests/engine/test_memory.py pins it against a real BlockPool allocation — the formula and the allocator must not drift.

7. Usage Examples

7.1 Full streaming loop (paged mode)

from oasr import GpuDecoderConfig
from oasr.cache import (
    CacheConfig, BlockPool,
    AttentionCacheManager, CnnCacheManager, CtcStateCacheManager,
    StreamContext,
)

cfg = CacheConfig(num_layers=12, n_kv_head=4, head_dim=64, hidden_dim=256,
                  kernel_size=15, chunk_size=16, num_left_chunks=4,
                  block_size_frames=16, max_num_blocks=2048)

pool    = BlockPool(cfg)
att_mgr = AttentionCacheManager(pool, cfg)
cnn_mgr = CnnCacheManager(cfg)
ctc_mgr = CtcStateCacheManager(GpuDecoderConfig(beam_size=10))

sid = 42
att_mgr.allocate_stream(sid)
cnn_mgr.allocate_stream(sid)
ctc_mgr.allocate_stream(sid, batch=1, vocab_size=5000)
ctx = StreamContext(sid, att_mgr, cnn_mgr, ctc_mgr)

for chunk in audio_chunks:                    # produced by feature extractor
    ctx.prepare_chunk()
    att_caches = ctx.get_att_caches()         # List[PagedKVCache]
    cnn_cache  = ctx.get_cnn_cache()
    # cache_t1 is the history the stream *holds*, not its offset: with
    # num_left_chunks=4 the offset grows forever while the cache does not.
    log_probs, new_cnn = model.forward_chunk_paged(
        chunk, offset, att_caches, cnn_cache,
        cache_t1=att_mgr.cached_frames([sid]))
    ctx.commit_chunk_paged(log_probs.size(1), new_cnn)
    ctx.get_decoder().decode_chunk(log_probs)
    offset += log_probs.size(1)

result = ctx.get_decoder().finalize_stream()
ctx.free()

7.2 Dense mode (legacy)

ctx = StreamContext(sid, att_mgr, cnn_mgr, ctc_mgr)

for chunk in audio_chunks:
    att_cache = ctx.get_att_cache()
    cnn_cache = ctx.get_cnn_cache()
    log_probs, new_att, new_cnn = model.forward_chunk(
        chunk, offset, required_cache_size, att_cache, cnn_cache)
    new_kv_chunk = new_att[:, :, -chunk_size:, :]
    ctx.commit_chunk(new_kv_chunk, new_cnn)
    ctx.get_decoder().decode_chunk(log_probs)

7.3 Batched paged forward (multi-stream, single encoder call)

# All B streams share an offset and have a full window ready.
att_mgr.prepare_chunks_batched([req.stream_id for req in group])

block_tables, cache_seqlens = [], []
for req in group:
    bt, cs = req.stream_context.get_paged_state_views()
    block_tables.append(bt); cache_seqlens.append(cs)

batched_bt = torch.cat(block_tables, dim=0)
batched_cs = torch.cat(cache_seqlens, dim=0)

batched_caches = [PagedKVCache(
    k_cache=pool.get_kv_view(l)[0],
    v_cache=pool.get_kv_view(l)[1],
    block_table=batched_bt, cache_seqlens=batched_cs,
    block_size=cfg.block_size_frames, host_seqlen=offset)
    for l in range(cfg.num_layers)]

log_probs, new_cnn = model.forward_chunk_paged(
    xs_stacked, offset, batched_caches, cnn_stacked,
    cache_t1=att_mgr.cached_frames([req.stream_id for req in group]))

8. Error Handling and Edge Cases

Condition Where Behaviour
Pool exhausted (pool.allocate(n) with len(free) < n) BlockPool.allocate Raises RuntimeError with the requested vs. free count. Caller (engine) must currently treat this as fatal — there is no built-in preemption.
Duplicate allocate_stream(sid) All three managers Raises ValueError.
Operating on a freed / unknown stream id All three managers KeyError from _get_state / caches[sid] / states[sid].
prepare_chunk not called before get_paged_caches / get_paged_state_views AttentionCacheManager Raises RuntimeError — block_table is None until first prepare.
Shape mismatch in commit / update AttentionCacheManager.commit, CnnCacheManager.update Raises ValueError with both expected and actual shapes.
chunk_frames > block_size_frames in dense commit AttentionCacheManager.commit Raises ValueError. Dense mode requires one block per chunk.
First-chunk (0, 0, 0, 0) placeholder cnn_cache in batched forward ModelRunner._forward_batched_paged Skips the per-stream stack and passes the placeholder; the model handles it layer-wise.
Eviction inside batched forward commit_chunk_paged Per-stream eviction shifts each stream's block_table independently — no cross-stream interaction.

BlockPool.allocate / free are guarded by a threading.Lock. The other managers are not thread-safe internally. They are expected to be called only from the engine's main step loop. OfflineExecutor runs CPU-side fbank on a producer thread but never touches the cache managers; only the consumer thread that performs the GPU forward calls into them.

9. Performance Considerations

  1. Pool over-provisioning. Allocate enough max_num_blocks to cover peak concurrent demand. Falling back into a RuntimeError mid-stream is fatal; running out is the most common production failure mode. EngineConfig.max_num_blocks=None removes the hand-computation: it fills what free VRAM allows, capped at what the block table can address.
  2. Avoid cudaMalloc in the hot loop. Per-stream allocations (block_table, cache_seqlens) are lazy — they happen on first prepare_chunk. CTC states are pooled. The CNN cache tensor is allocated once at allocate_stream. After that, the streaming step does no allocator activity except the freelist popleft inside BlockPool.allocate.
  3. Use prepare_chunks_batched when admitting or stepping a B-stream cohort. One pool lock + one CPU-side popleft loop is substantially cheaper than B independent prepares, even though the per-stream block-table writes are unavoidable.
  4. Cohort admission (EngineConfig.streaming_cohort_admit=True) length-sorts the waiting streams, bounded by max_wait_time (see scheduler.md §4.2). It no longer gates on offsets: _forward_batched_paged batches streams at any offset.
  5. gather_kv_blocks (dense mode) does an index_select per layer and is O(num_layers × num_blocks). Paged mode avoids this entirely — prefer it for long-context streams.
  6. block_size_frames == chunk_size is the simplest invariant: each chunk is exactly one block, eviction matches chunk granularity, and the dense commit shape check trivially holds.
  7. No Python scalar is written through an index into a persistent tensor. block_table[slot, i] = block_id and block_table[slots_t, kept_t] = 0 both lower to a pageable 4-byte host-to-device copy, which waits for everything queued on the stream (41.6 ms measured behind queued GEMMs). Admission writes a row through one staged copy (to_device) and fill_ on a slice, eviction blanks the vacated column from a device scalar (_zero). On Nemotron — a prefilled window, so five block writes per admission and an eviction every step — these were most of the 16 synchronisations per tick.
  8. The padding lane. BlockPool(config, scratch_blocks=1) allocates one physical block past max_num_blocks that is never on the free list (and that free rejects), and AttentionCacheManager(pool, config, pad_lanes=1) adds one slot row past max_batch_size whose block table points only at it, at the cache length every stream starts with. The paged streaming backend pads a cohort up to a captured width with that row (SlotStateCache gets the extra slot too), so padding rows write K/V and conv state only where no stream will ever read. StreamSlotPool still hands out [0, max_batch_size), so no stream can be bound to the lane. See engine.md §9.8.

10. Two storage disciplines, and how to add a cache

Every streaming cache here — and every one an encoder has needed so far — is one of exactly two things. Knowing which one you are adding is the whole design decision:

extent addressing eviction examples
Fixed, slot-addressed (oasr/cache/state.py) known when the engine is built; never grows one persistent tensor per declaration, with the stream's slot_id as one of its axes none — overwritten in place each chunk convolutional left-context, a subsampling stage's tail, per-layer recurrent state
Growing, block-addressed (attention_cache.py, decoder_kv.py) grows with the stream shared BlockPool + per-stream block table LRU shift / recycle attention K/V, AR decoder K/V

StreamSlotPool is the shared addressing primitive for both: the paged manager's block_table and cache_seqlens rows are slot-indexed too.

The decoder K/V of the autoregressive families is the second growing kind, on its own pool with its own policy: it grows one token per step and never evicts, so DecoderKVCacheManager allocates a page as a row crosses a boundary and frees the whole row when the request finishes. Two consequences follow from having no eviction. Admission must reserve each row's declared ceiling (can_admit) — a pool that can run out mid-generation fails a request that is already half-answered, where refusing at admission is one rejected batch with a reason. And a row's pages are its own, so the repeated-index select that expands a beam grid would alias two slots onto one page; PagedDecoderKv.select refuses, and beam search runs on dense per-group buffers instead. This pool is what kv_storage="auto" (the default) selects for a decoder that captures its steps — a graph needs addresses that outlive a batch, and a per-group capacity buffer moves with every prefill. See decoding.md.

An AED's cross-attention K/V lives in the same pool as a fixed-extent region: a second slot per row created with prefill_len equal to the encoder window and no growth (decoder_state.build_cross_kv), so admission reserves exactly the pages it maps, and select / merge treat it like the self-attention's rows. It is block-addressed although it never grows because it is read by the same paged attention path as the self-attention, and because the alternative — a dense tensor per decode group — is the one thing that kept a Whisper step out of a CUDA graph.

10.1 Adding a fixed-extent cache — declare it

CnnCacheManager is the single-spec instance of SlotStateCache. To add more, an encoder declares them and reads them back by name; nothing in the cache layer or the engine changes:

# on the encoder
@property
def streaming_state_specs(self):
    return (
        StreamStateSpec("subsample.0", (kernel - 1, freq_bins, channels), slot_axis=0),
        ...
    )

# in its chunk forward
def forward_chunk_paged(self, xs, offset, att_caches, cnn_cache, att_mask, cache_t1, states):
    view = states["subsample.0"]
    cached = view.gather()            # (B, kernel-1, F, C)
    view.scatter(new_tail)            # in place, no per-chunk allocation

Two things carry consequences:

  • slot_axis is per spec. It is why the conv cache keeps its historical (layers, slots, frames, dim) layout — and therefore its buffer address, which GraphedEncoderForward captures by reference. It is also exactly the per-kind batch dimension a Zipformer-style recurrent state needs (§10.4).
  • Allocation zeroes the slot, and that is the initial value, not hygiene. A zero left-context is precisely the padding an offline pass applies, so a stream's first chunk computes what an offline pass would.
  • persistent=True marks a state that describes the request, not the audio. Nemotron's per-stream language-prompt id is the one: the paged runtime writes it at admission from BaseAsrModel.stream_conditioning (SlotStateCache.fill_slot, a fill_ with no host-to-device copy), the chunk forward gathers it like any other state (so a captured graph reads each replay's rows), and reset_stream keeps it. A VAD turn reset rewinds what the encoder heard; zeroing the prompt with it would decode the next turn under slot 0, which is a language too. A runtime that does not write it is refused for such a model at construction (carries_encoder_conditioning).

10.2 A trained fixed attention window — prefill it

CacheConfig.prefill_kv_window hands a stream its whole retained K/V window, zeroed, at admission. This is for an encoder whose attention span is part of the trained mask (Nemotron's sliding_window) rather than "whatever history fits", and its purpose is uniformity: a Transformer-XL relative-position table's distances are cache_seqlens + i - j, so one table shared across a cohort is only correct when every stream reports the same cached length. A per-row table would mean running the positional projection per row, which at B = 32 costs more than the encoder layer it feeds.

Two things fall out for free, and one obligation:

  • cache_t1 becomes constant, so GraphedEncoderForward captures one graph per batch size instead of one per (B, cache_t1 bucket);
  • the pool requirement is max_batch_size * blocks_per_stream from admission — the invariant CacheConfig.__post_init__ already enforces when eviction is on;
  • the blocks must be zeroed, not merely reserved. oasr.fmha gives a past-the-length column zero softmax weight but still multiplies it into P @ V, so a NaN there propagates where no mask can intercept it.

10.3 What is deliberately not on the declared axis

  • The per-layer conv cache is not re-declared as a StreamStateSpec by each encoder. conv_kernel_size is already a declaration, just a compact one, and it feeds CacheConfig.cnn_cache_frames, CnnCacheManager.buffer and the graph cache's address capture. Rewriting three call sites for uniformity, on the one path with a known graph-capture invariant, buys nothing functional.

10.4 An encoder whose whole cache is fixed-extent — the slot runtime

Zipformer's streaming state is all fixed-extent: per layer a left-context key, nonlin-attention and two value caches (a rolling window of left_context // downsample frames, not a growing one), two conv tails, plus the embed's cached frames and a processed-length counter — 98 tensors for the 16-layer release. It is not paged because it does not grow, and because each stack's window has its own width, which one uniform BlockPool geometry cannot hold. So it is declared instead: ZipformerEncoder.slot_state_specs is one StreamStateSpec per tensor, derived from the encoder's own get_streaming_init_states (on the meta device) so its order, shapes and batch axes cannot drift from the list API, and streaming_kind="slot" hands it to SlotStreamingBackend:

  • one SlotStateCache holds every tensor, the stream slot on the tensor's own batch axis (slot_axis), zero being every tensor's initial value;
  • a tick gathers the active slots into the list, runs one batched chunk forward through the list API, and scatters the new list back — bit-identical to the "stateful" runtime, which cats every stream's lists together each tick;
  • stable addresses make the whole step one CUDA graph per batch width. Every window is full (a short final one is padded with streaming_pad_value), so the width is the only key, and the padding lane (streaming_graph_pad_batch, on whenever steps are graphed) keeps it to a power-of-two ladder — 6 captures instead of 32 at a 32-stream pool, 0.43 GiB of graphs instead of 1.7;
  • capture warms up and records the step, scatter included, against a scratch slot past the streams', so a width captured while streams are live advances none of them;
  • reset (rule 13) zeroes the slot in place and rewinds offset.

Measured on the causal release at a 32-stream pool: 5.95× over the stateful runtime, transcripts unchanged; see docs/models.md for the geometry.

10.5 Other extension points

  • Custom stream identifiers. Stream IDs are arbitrary integers — the engine assigns them monotonically, but external code can use any scheme as long as IDs are unique across active streams.
  • Alternative pool layouts. BlockPool exposes get_kv_view(layer) and get_kv_block_view(layer, block_id) as the only tensor entry points. A subclass could change the underlying storage (e.g. swap K/V into separate per-layer tensors) without touching AttentionCacheManager.
  • Custom decoders. CtcStateCacheManager is hard-wired to GpuStreamingDecoder. To plug in a different decoder you would mirror the create_state / reset_state / decode_chunk / finalize_stream interface and either subclass the manager or write a parallel one.
  • Per-stream metadata. _StreamKVState is private. If you need to carry per-stream data (e.g. user-supplied tags, biasing graph pointers) prefer holding it on the engine Request object — keep the cache layer single-purpose.

10.6 FP8 KV for the AR decoders

EngineConfig(quantization={"kv_cache": "fp8_e4m3"}) (or oasr-server --quantization-option kv_cache=fp8_e4m3) stores the decoder KV — the self-attention cache and an AED's cross-attention K/V — as float8_e4m3fn with one fp32 scale per (token, kv-head). Both storage modes:

data scales
paged (BlockPool with dtype=float8_e4m3fn) (layers, blocks, page, H_kv, D) fp8 (layers, blocks, page, H_kv) fp32, zeroed like the data
dense (DecoderKv(kv_cache="fp8_e4m3")) (B, H_kv, cap, D) fp8 per layer (B, H_kv, cap) fp32 per layer
  • Written by one launch per layer. quantize_kv_store computes each vector's scale from that vector alone (the least power of two putting its amax inside ±448 — no calibration, batch-invariant) and scatters K and V to the row's own positions, paged by flat slot or dense by offset.
  • Read by the attention kernel, never dequantized wholesale. append hands the attention QuantizedKv (data + scales) where it handed a tensor; oasr.layers.Attention passes it to the fp8 variants of the C++ lane (kernels.md), which dequantize per tile in shared memory.
  • A prefill attends its prompt fresh. The prompt's K/V are quantized into the cache as they are written, but an empty quantized cache answers mask_kwargs / append as if trimmed, so the prompt's own causal, left-padded attention runs in the served dtype; every later step reads the fp8 cache. (Whisper's prompt still reads the cross K/V through the fp8 path — they are cached.)
  • Bookkeeping is unchanged. select / merge move scales with their rows; paged pages are freed the same way; a dense cache needs a capacity buffer. The pool's format and the decoder's must agree (build_kv raises otherwise). Step graphs capture the quantize-and-scatter and the fp8 read like any other kernel, and replay is bit-identical to the eager step.
  • Sizing counts it. IncrementalArStrategy._kv_bytes_per_token is 2 · layers · H_kv · (D + 4) bytes instead of 2 · layers · H_kv · D · 2 — 1.88x fewer at D = 64, 1.94x at D = 128 — so the VRAM-derived decode_kv_budget_gib affords that many more rows and the pool built from a byte budget holds that many more blocks. Rows actually admitted are still capped by max_decode_slots (default max_batch_size): raise it to spend the room.
  • Refused where it cannot be served, at engine construction, by name: a decode family with no AR cache, a decoder that does not list the format in kv_cache_dtypes, an fp32 model on CUDA, a head dim the fp8 gather cannot tile, or OASR_ATTN_BACKEND=cute|sdpa. The word-timing alignment pass is one teacher-forced forward that stores nothing and keeps the served dtype.

11. Quick Reference

allocate_stream(sid)         # all three managers
StreamContext(sid, ...)      # bundle them

per chunk:
  ctx.prepare_chunk()        # paged mode only
  ctx.get_att_cache[s]()     # dense | paged
  ctx.get_cnn_cache()
  ... model.forward_chunk[_paged](...) ...
  ctx.commit_chunk[_paged](...)
  ctx.get_decoder().decode_chunk(logits)

at end:
  ctx.get_decoder().finalize_stream()
  ctx.free()                 # releases blocks, CNN cache, pools CTC state