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 : ℕ) (tokenShape : Spec.Shape) : Type

A key/value cache with an explicit sequence length and token shape.

• keys : TorchLean.Tensor α (seqLen :: tokenShape)

  Cached key tensors, indexed by time.

• values : TorchLean.Tensor α (seqLen :: tokenShape)

  Cached value tensors, indexed by time.

def NN.Examples.BugZoo.KVCache.singletonToken {α : Type} [TorchLean.Storage α] {tokenShape : Spec.Shape} (x : TorchLean.Tensor α tokenShape) : TorchLean.Tensor α (1 :: tokenShape)

View one token tensor as a length-one sequence.

@[simp]

theorem NN.Examples.BugZoo.KVCache.singletonToken_get {α : Type} [TorchLean.Storage α] {tokenShape : Spec.Shape} (x : TorchLean.Tensor α tokenShape) (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 : ℕ} {tokenShape : Spec.Shape} (past : TorchLean.Tensor α (seqLen :: tokenShape)) (newToken : TorchLean.Tensor α tokenShape) : TorchLean.Tensor α ((seqLen + 1) :: tokenShape)

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

@[simp]

theorem NN.Examples.BugZoo.KVCache.appendToken_last {α : Type} [TorchLean.Storage α] {seqLen : ℕ} {tokenShape : Spec.Shape} (past : TorchLean.Tensor α (seqLen :: tokenShape)) (newToken : TorchLean.Tensor α tokenShape) : (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 : ℕ} {tokenShape : Spec.Shape} (cache : Cache α seqLen tokenShape) (newKey newValue : TorchLean.Tensor α tokenShape) : Cache α (seqLen + 1) tokenShape

Append both key and value tensors to the KV cache.

theorem NN.Examples.BugZoo.KVCache.appendKV_last_key {α : Type} [TorchLean.Storage α] {seqLen : ℕ} {tokenShape : Spec.Shape} (cache : Cache α seqLen tokenShape) (newKey newValue : TorchLean.Tensor α tokenShape) : (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 : ℕ} {tokenShape : Spec.Shape} (cache : Cache α seqLen tokenShape) (newKey newValue : TorchLean.Tensor α tokenShape) : (appendKV cache newKey newValue).values[seqLen] = newValue

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