BugZoo: KV-cache contracts

LLM inference-engine bug reports include cache-shift bugs, RoPE/position mismatches, shape mistakes, and resource/configuration errors. TorchLean does not verify a full serving engine, paged attention allocator, or multi-GPU scheduler. The useful first step is still precise: represent the cache update as a typed tensor operation and prove the append invariant we rely on.

Reference:

• Liu et al., "A First Look at Bugs in LLM Inference Engines", 2025.

This file proves the cache append boundary. A stronger future theorem should connect cached decode to full-sequence attention:

$$ \operatorname{decodeWithCache}(\mathit{prefix},\mathit{newToken}) =\operatorname{fullAttention}(\mathit{prefix}\mathbin{+\!\!+}[\mathit{newToken}]). $$

under the same mask, RoPE/position encoding, and numeric semantics.

structure NN.Examples.BugZoo.KVCache.Cache (α : Type) [TorchLean.Storage α] (seqLen headDim : ℕ) : Type

A key/value cache with an explicit sequence length and head dimension.

• keys : TorchLean.Tensor α [seqLen, headDim]

  Cached key vectors, indexed by time.

• values : TorchLean.Tensor α [seqLen, headDim]

  Cached value vectors, indexed by time.

def NN.Examples.BugZoo.KVCache.singletonToken {α : Type} [TorchLean.Storage α] {headDim : ℕ} (x : TorchLean.Tensor α [headDim]) : TorchLean.Tensor α [1, headDim]

View one token vector as a length-one sequence.

@[simp]

theorem NN.Examples.BugZoo.KVCache.singletonToken_get {α : Type} [TorchLean.Storage α] {headDim : ℕ} (x : TorchLean.Tensor α [headDim]) (i : Fin 1) : (singletonToken x)[i] = x

Reading the only position of a singleton sequence gives the token back.

def NN.Examples.BugZoo.KVCache.appendToken {α : Type} [TorchLean.Storage α] {seqLen headDim : ℕ} (past : TorchLean.Tensor α [seqLen, headDim]) (newToken : TorchLean.Tensor α [headDim]) : TorchLean.Tensor α [seqLen + 1, headDim]

Append one token vector to a sequence cache along the time axis.

@[simp]

theorem NN.Examples.BugZoo.KVCache.appendToken_last {α : Type} [TorchLean.Storage α] {seqLen headDim : ℕ} (past : TorchLean.Tensor α [seqLen, headDim]) (newToken : TorchLean.Tensor α [headDim]) : (appendToken past newToken)[seqLen] = newToken

The appended token lands at index seqLen, that is, at the end.

This is the fact an off-by-one in a KV cache would break: writing the new token over the last cached position instead of after it. The type already forces the length to grow by one, and this lemma pins down where the new entry goes.

def NN.Examples.BugZoo.KVCache.appendKV {α : Type} [TorchLean.Storage α] {seqLen headDim : ℕ} (cache : Cache α seqLen headDim) (newKey newValue : TorchLean.Tensor α [headDim]) : Cache α (seqLen + 1) headDim

Append both key and value vectors to the KV cache.

theorem NN.Examples.BugZoo.KVCache.appendKV_last_key {α : Type} [TorchLean.Storage α] {seqLen headDim : ℕ} (cache : Cache α seqLen headDim) (newKey newValue : TorchLean.Tensor α [headDim]) : (appendKV cache newKey newValue).keys[seqLen] = newKey

The newly appended key is exactly the final key in the updated cache.

theorem NN.Examples.BugZoo.KVCache.appendKV_last_value {α : Type} [TorchLean.Storage α] {seqLen headDim : ℕ} (cache : Cache α seqLen headDim) (newKey newValue : TorchLean.Tensor α [headDim]) : (appendKV cache newKey newValue).values[seqLen] = newValue

The newly appended value is exactly the final value in the updated cache.