ClikaRT::nn::KVCache

class

Header: ClikaRT/nn/kv_cache.h

The serving-side KV cache: per-layer storage rows (attention K/V, conv windows, recurrent slabs) over a slot-indexed batch.

Built once per serving configuration with make(config); the attention ops append each step's K/V through it and attend over what it holds. Rows lay out CONTINUOUS (per-slot buffers) or PAGED (a shared block pool) per the config's mode; a KVQuantSpec on a row quantizes on append and dequantizes on attend, invisibly to the caller. Move-only handle; the cache owns its device storage for its whole lifetime.

auto cache = ClikaRT::nn::KVCache::make(config);
// per step: the attention ops read/write through the cache's tensors

Nested types

Name	Description
StepIndices	The device-side index tensors one varlen decode/prefill step consumes, derived from the host-side step description (see prepare_step).

Static member functions

make()

static KVCache make(const KVCacheConfig& config, Span<const KVLayerSpec> layer_specs)

Build a cache with one layer per layer_specs row (size must equal config.num_layers). Every composition goes through this one entry: a uniform full-attention stack is N identical AttentionKV rows; a windowed/full alternation varies window per row; recurrent stacks mix in the state kinds. Each row's Layout::Auto resolves from config.mode; a row may pin Continuous/Paged explicitly and mix freely with the rest. When any row lays out Paged, config.paged must carry the pool geometry; paged rows share one block table, with one pool per distinct (kv_heads, head_dim) row geometry. Container bookkeeping (seq_lens) reads a paged row when one exists, else the first token-indexed layer, so a stack whose layer 0 is a state row works.

Raises ClikaRT::Error on a spec the cache cannot build: a spec/num_layers size mismatch, an empty state geometry, max_tokens_per_seq <= 0 with a non-windowed continuous row present, a paged row asking for an asymmetric v_head_dim (it stays on continuous rows), a negative kv_heads, a per-row k_quant/v_quant override, or a per-row k_scale/v_scale calibration override that breaks its contract (see the field).

Declared in ClikaRT/nn/kv_cache.h, line 423

Member functions

KVCache()

KVCache(KVCache&& other) noexcept

Declared in ClikaRT/nn/kv_cache.h, line 426

operator=()

KVCache& operator=(KVCache&& other) noexcept

Move-only handle; the moved-from object is empty.

Declared in ClikaRT/nn/kv_cache.h, line 428

~KVCache()

~KVCache()

Releases the cache's device storage.

Declared in ClikaRT/nn/kv_cache.h, line 430

keys()

Tensor keys(int layer) const

This layer's K / V storage view; bind as past_key/out_present_key (and value) into group_query_attention_varlen. Raises ClikaRT::Error (InvalidArgument) when layer is outside [0, num_layers); the message names the index and the layer count.

Declared in ClikaRT/nn/kv_cache.h, line 437

values()

Tensor values(int layer) const

The value-side sibling of keys(layer): this layer's V storage view, same binding use.

Throws

• ClikaRT::Error: as keys does.

Declared in ClikaRT/nn/kv_cache.h, line 441

conv_state()

Tensor conv_state(int layer, int index = 0) const

A State layer's i-th conv window slab [max_seqs, dim_i, width_i], a shared handle; bind it (or a slot-row view of it) as causal_conv_update's state and the op updates it in place. Raises on an out-of-range layer, a non-State layer, or an out-of-range index.

Declared in ClikaRT/nn/kv_cache.h, line 449

recurrent_state()

Tensor recurrent_state(int layer) const

A State layer's recurrent slab [max_seqs, heads, k_dim, v_dim]. Raises on an out-of-range layer, a non-State layer, or one that declared no recurrent state.

Declared in ClikaRT/nn/kv_cache.h, line 455

seq_lens()

Tensor seq_lens() const

Per-slot effective token counts, [max_seqs] Int32 on the cache device.

Declared in ClikaRT/nn/kv_cache.h, line 459

mode()

KVCacheMode mode() const

The serving mode this cache was built with, the value every Layout::Auto row resolved to. The predicate that selects the decode addressing path (cheap, allocates nothing).

Declared in ClikaRT/nn/kv_cache.h, line 464

num_layers()

int num_layers() const

The number of layers this cache was built with: the valid layer range of keys / values / conv_state / recurrent_state is [0, num_layers()). Never fails; allocates nothing.

Declared in ClikaRT/nn/kv_cache.h, line 468

paged_block_table()

Tensor paged_block_table() const

The paged block table, [max_seqs, max_blocks_per_seq] Int32, or an undefined tensor for a cache with no paged rows. Use mode() for the serving-mode test; this accessor is for gathering a sequence's blocks.

Declared in ClikaRT/nn/kv_cache.h, line 473

reserve()

Tensor reserve(Tensor slot_ids, Span<const std::int32_t> n_new_per_seq)

Paged only: reserve n_new_per_seq[i] new tokens for slot slot_ids[i], ONCE per decode step (acquire blocks + advance seq_len). Returns the per-sequence pre-advance write offset, [B] Int32 on the cache device. Raises ClikaRT::Error on a continuous cache or pool exhaustion.

Declared in ClikaRT/nn/kv_cache.h, line 480

prepare_step()

StepIndices prepare_step(
    Span<const std::int32_t> cu_seqlens_q,
    Span<const std::int32_t> past_len,
    Span<const std::int32_t> slot_ids
)

Convert one step's host-side description (the [B+1] query cumulative cu_seqlens_q, the [B] per-slot committed lengths past_len, and the [B] cache slot ids) into the device-side index tensors a varlen attention forward binds. One call per step, before the layer loop:

• On a PAGED cache this also RESERVES the step's blocks (once, covering every layer; do not call reserve separately) and returns the block table gathered into slot order as kvcache_start ([B, blocks_per_seq]); cu_seqlens_k is [B] post-append lengths in decode (every query length 1) or the [B+1] cumulative in prefill.
• On a CONTINUOUS cache slot_ids is the [B] slot ids on device (the attention op's slot_ids input: the cache row each batch row appends to and attends from), cu_seqlens_k is the [B] per-seq post-append lengths, and kvcache_start is the rank-1 [B] selector the op does not read the values of (it carries the same slot ids; every write offset derives from cu_seqlens_k). Bind all three to the attention call; on a paged cache slot_ids is undefined and binds as absent.

Steady-state decode repeats the same cu_seqlens_q / slot list every step; those uploads are cached by content and re-run only when the batch composition changes. Drive a cache's steps from ONE thread (the uploads are cached without locking). Raises ClikaRT::Error on mismatched span lengths or pool exhaustion.

Declared in ClikaRT/nn/kv_cache.h, line 547

admit()

std::int64_t admit(
    int batch_idx,
    Tensor prompt_ids,
    std::int64_t addressable_len = kAddressableAll
)

Paged only: admit a freshly-recycled slot with its prompt. Hashes the prompt's full blocks, binds the longest already-cached prefix into the slot (sharing physical blocks with the sessions that produced them), and returns the number of leading tokens whose K/V is already cached, always block-aligned and strictly less than the prompt length. Forward only prompt_ids[num_cached:] with past length num_cached. Call once per session, before the first reserve. prompt_ids is [N] Int32 (any device). Raises ClikaRT::Error on a continuous cache or a non-fresh slot.

addressable_len (default kAddressableAll): the count of leading tokens whose K/V depends on the token ids alone. A multimodal prompt passes the position of its FIRST media token; caching then covers only the pure-text prefix before it (blocks at or past the boundary are never shared), at zero cost to text-only sessions. Pass the SAME value to this session's retire.

Declared in ClikaRT/nn/kv_cache.h, line 577

retire()

void retire(
    int batch_idx,
    Tensor full_token_ids,
    std::int64_t addressable_len = kAddressableAll
)

Paged only: retire a FINISHED session gracefully. Pass the session's full committed token history ([N] Int32, prompt + generated, exactly the tokens whose K/V the slot holds): their blocks stay cached for future admits (a follow-up chat turn resubmits the transcript, so its prefix hits), then the slot is freed. Use evict_batch instead on the cancel/error path; it drops everything and caches nothing. addressable_len as on admit: blocks at or past the boundary are dropped, never cached.

Declared in ClikaRT/nn/kv_cache.h, line 590

fill()

void fill(
    int layer,
    int batch_idx,
    Tensor k,
    Tensor v
)

Fill layer layer's K/V planes for slot batch_idx in ONE call, the write verb for Kind::CrossAttentionKV rows (encoder cross-attention: project the encoder states per layer, hand both projections here). k / v are token-major [t_enc, kv_heads, head_dim]; a dtype differing from the cache's kv_dtype converts inside the copy. Sets the filled layer's slot length to t_enc, IDEMPOTENTLY: a refill with the SAME t_enc overwrites in place; a different t_enc on a non-empty slot raises (evict the slot first). Raises on a non-CrossAttentionKV layer, a shape/rank mismatch, or t_enc beyond max_tokens_per_seq.

Declared in ClikaRT/nn/kv_cache.h, line 603

update_position()

void update_position(int batch_idx, std::int64_t new_pos)

Set the effective seq_len of batch_idx to new_pos (bookkeeping only; does not trim storage). Call after each step's append. Raises on error.

Declared in ClikaRT/nn/kv_cache.h, line 608

evict_batch()

void evict_batch(int batch_idx)

Free a slot's stored K/V (its seq_lens() entry returns to 0).

Declared in ClikaRT/nn/kv_cache.h, line 612

abandon_batch()

void abandon_batch(int batch_idx)

Free a slot after a FAILED session, with host-side bookkeeping only: no kernel launch, no device copy, no stream wait, whatever the cache carries. Block references the host bookkeeping knows are released exactly as evict_batch releases them (a shared prefix block other slots reference is untouched); anything newer is dropped, bounded to this slot, and never listed as free again. The slot's storage hygiene runs at its next use, on that caller's stream, before any read. Never fails for device reasons; a bad slot index reports INVALID_ARGUMENT. Prefer evict_batch when the session ended normally (it keeps the pool's bookkeeping exact); use this when the session ended in an error its device may not have survived.

Declared in ClikaRT/nn/kv_cache.h, line 626

reset()

void reset()

Drop all stored K/V across every slot and layer.

Declared in ClikaRT/nn/kv_cache.h, line 630

Static data members

kAddressableAll

static std::int64_t kAddressableAll = -1

The admit/retire addressability default: every token's K/V is a pure function of the token ids, so the whole prompt participates in prefix caching, the text-only posture. A MULTIMODAL session must pass its first media position instead (see admit): the K/V under a media placeholder comes from pixels/audio the ids cannot identify, so id-keyed sharing at or past it would serve another session's media content.

Declared in ClikaRT/nn/kv_cache.h, line 558