core/nn/kv_cache library

Per-block KV cache for autoregressive attention.

A MHACache holds the running K and V tensors for every head of one MultiHeadAttention layer. During autoregressive sampling the forward pass appends the K/V produced from each new token to these buffers via a row-wise concat, letting the next step reuse all past projections instead of recomputing them.

An EncoderCache is a list of MHACaches (one per block) — the state carried across steps by a full TransformerEncoder / GPT.generate loop.

The caches are strictly inference-time state: their tensors are created without requiresGrad. Building a fresh cache (MHACache.empty / EncoderCache.forGpt) makes every entry null; the first append sets it, subsequent appends concat the new rows onto the existing buffer.

Classes

EncoderCache

MHACache