Prefix Cache Architecture

1. The Radix Cache

RadixTreeCache is the outer lookup structure. It is not the model's attention cache itself. Its job is only to answer: “Have we seen this token prefix before?”

2. The Live Request Cache

Every request gets a fresh runtime cache stack via context.model.newCache(parameters: ...). On a hit, AFM restores data into this live cache and then continues generation from there.

3. The Per-Layer Cache Classes

Inside the live cache stack, each layer may use a different implementation such as KVCacheSimple, RotatingKVCache, or ArraysCache. These are not separate top-level systems; they are layer-specific cache shapes.

Architectural Levels

Level	What It Does	Examples
1. Prefix Index	Looks up shared prompt-token prefixes.	RadixTreeCache
2. Live Request Cache	Fresh runtime cache stack for the current generation.	generationCache, batch prefill cache
3. Per-Layer Cache Objects	Hold the actual model state for each cached layer.	KVCacheSimple, RotatingKVCache, MambaCache, CacheList

Cache Families

Family	Used For	Tensor Layout / Indexing
KVCacheSimple	Standard attention KV state.	[batch, kv_heads, seq_len, head_dim]; sequence grows on axis 2.
RotatingKVCache	Sliding-window / ring-buffer attention.	[batch, kv_heads, cache_slots, head_dim]; logical time is tracked by offset and ring cursor idx on axis 2.
ArraysCache / MambaCache	Recurrent or state-space layers.	No single universal 4D layout. Stores an array-list of tensors; shape is model-specific. MambaCache is a 2-slot ArraysCache.
CacheList	Composite wrapper for multiple sub-caches in one logical layer.	No single tensor layout. Flattens and delegates indexing/shape semantics to child caches.

Per-Layer Cache Types

KVCacheSimple RotatingKVCache ArraysCache MambaCache CacheList

The prefix cache stores a full per-layer snapshot. Different layer classes require different restore semantics.

Observed Model Mappings

• Qwen3.5 A3B: recurrent layers via MambaCache / ArraysCache.
• Nemotron Nano: mixed recurrent + regular KV layers.
• GPT-OSS 20B: sliding-window attention via RotatingKVCache.

Step A: Find a Prefix

• findPrefix(tokens) returns the longest cached prefix.
• The result includes the matched length and the cached layer payloads.

Step B: Compute a Safe Prefix

• effectiveCachedPrefix = min(prefixLen, inputCount - 16)
• This leaves a suffix so the model reprocesses the prompt boundary.
• For recurrent caches on an exact full replay, AFM bypasses restore entirely.

Primary Decider: The Model Architecture

AFM does not arbitrarily pick cache classes at runtime. The model implementation in MLX defines what each layer needs when it builds caches inside newCache(...) or the model's cache factory.

• Plain transformer attention layers usually use KVCacheSimple.
• Sliding-window attention layers use RotatingKVCache.
• Recurrent or state-space layers use ArraysCache / MambaCache.
• Composite layers may wrap multiple sub-caches in CacheList.

Secondary Decider: Model Parameters

Some models also change cache shape based on configuration such as sliding-window size or whether a layer is an attention layer versus a recurrent/state-space layer.

In practice, the decisive code lives in the vendor model definitions and cache helpers, especially AFMKit's vendor/MLX/mlx-swift-lm/Libraries/MLXLMCommon/KVCache.swift and the specific model file's makeCache / newCache logic.

Example: GPT-OSS 20B

Sliding-window layers choose RotatingKVCache, which is why metadata restore matters there.

Example: Qwen3.5 A3B

Recurrent/state-space layers use array-based caches, which is why exact full replay had to be bypassed.

Example: Standard Attention

Normal transformer attention typically uses KVCacheSimple, the least complicated restore case.

Qwen Exact-Replay Divergence

Recurrent cache layers could produce shortened completions when the cache restore covered the entire prompt. AFM now treats that as unsafe and logs outcome=exact-replay-bypass.

GPT-OSS Rotating-Cache Crash

A rotating cache hit needed metadata restoration, but the original radix store only saved tensor state. A naive metadata restore was also unsafe because ring-buffer cursor values must stay coherent with restored state length. AFM now stores metadata, restores it carefully, and skips unsafe physical truncation for rotating caches.