diff --git a/mlx_lm/models/cache.py b/mlx_lm/models/cache.py
index b84c9d650..29312ff86 100644
--- a/mlx_lm/models/cache.py
+++ b/mlx_lm/models/cache.py
@@ -3,7 +3,7 @@
 import copy
 from collections import deque
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Sequence, Tuple
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -1761,3 +1761,1023 @@ def stats_by_type(self):
                 "n_bytes": self._n_bytes_by_type[cache_type],
             }
         return result
+
+
+# ---------------------------------------------------------------------------
+# SnapKV-D: post-prefill KV eviction with position-preserving decode
+#
+# Long-context decode reads the whole KV cache every step. After the prompt is
+# prefilled, most middle prompt rows contribute little to future attention, so
+# keeping attention sinks + a recent window + the top observation-window-scored
+# middle rows within a budget (SnapKV, arXiv:2404.14469) and evicting the rest
+# cuts the per-token KV read proportionally. The retained rows are a sparse
+# subset of the prompt, so RoPE position and physical row count must diverge:
+# PositionPreservingKVCache tracks the true sequence position in ``offset`` for
+# future rotations while storing only the retained rows, and records each row's
+# true position so a prefix trim (prompt-cache reuse) stays exact.
+# ---------------------------------------------------------------------------
+
+
+def snapkv_keep_indices(
+    seq_len: int,
+    budget: int,
+    scores: Sequence[float],
+    *,
+    sink_tokens: int = 4,
+    recent_tokens: Optional[int] = None,
+    min_tokens: int = 128,
+) -> Tuple[int, ...]:
+    """Return the sorted prompt positions retained by the SnapKV-D policy.
+
+    Keeps ``sink_tokens`` leading rows, a recent window, and the highest-scoring
+    remaining rows up to ``budget``. Returns all positions unchanged when the
+    prompt is at or below ``min_tokens`` or the budget covers it.
+    """
+    if seq_len < 0:
+        raise ValueError("seq_len must be non-negative")
+    if budget <= 0:
+        raise ValueError("budget must be positive")
+    if len(scores) < seq_len:
+        raise ValueError(f"scores length {len(scores)} < seq_len {seq_len}")
+    if seq_len == 0 or seq_len <= min_tokens or budget >= seq_len:
+        return tuple(range(seq_len))
+
+    budget = min(budget, seq_len)
+    sink_count = min(max(0, sink_tokens), max(1, budget // 8), budget)
+    retained = set(range(sink_count))
+
+    remaining = budget - len(retained)
+    if remaining > 0:
+        recent_count = (
+            recent_tokens if recent_tokens is not None else max(1, budget // 8)
+        )
+        recent_count = min(recent_count, remaining)
+        retained.update(range(seq_len - recent_count, seq_len))
+
+    remaining = budget - len(retained)
+    if remaining > 0:
+        ranked = sorted(
+            (i for i in range(seq_len) if i not in retained),
+            key=lambda i: (float(scores[i]), i),
+            reverse=True,
+        )
+        retained.update(ranked[:remaining])
+
+    return tuple(sorted(retained))
+
+
+class PositionPreservingKVCache(_BaseCache):
+    """KV cache with a true-position ``offset`` and compact retained storage.
+
+    Intended for post-prefill eviction followed by single-token decode. The
+    physical K/V rows may be a sparse subset of the logical prompt; the
+    ``positions`` metadata records each row's true sequence position so a prefix
+    trim (prompt-cache reuse) can shorten the logical ``offset`` without
+    pretending sparse rows are contiguous. Speculative rollback trims a
+    generated suffix and is tracked separately.
+    """
+
+    step = 256
+
+    def __init__(
+        self,
+        keys=None,
+        values=None,
+        *,
+        offset: int = 0,
+        protected_stored: Optional[int] = None,
+        positions: Optional[Sequence[int]] = None,
+    ):
+        self.keys = keys
+        self.values = values
+        self.offset = int(offset)
+        self._stored = 0 if keys is None else int(keys.shape[2])
+        self._protected_stored = (
+            self._stored if protected_stored is None else int(protected_stored)
+        )
+        if self._protected_stored < 0 or self._protected_stored > self._stored:
+            raise ValueError("protected_stored must be between 0 and stored rows")
+        self._positions = self._coerce_positions(positions)
+        self._speculating = False
+        self._speculative_appends = deque()
+
+    def _coerce_positions(self, positions):
+        if positions is None:
+            if self._stored == 0:
+                return ()
+            if self.offset == self._stored:
+                return tuple(range(self._stored))
+            return None
+        out = tuple(int(p) for p in positions)
+        if len(out) != self._stored:
+            raise ValueError("positions length must match stored rows")
+        if any(p < 0 for p in out):
+            raise ValueError("positions must be non-negative")
+        if tuple(sorted(out)) != out:
+            raise ValueError("positions must be sorted")
+        return out
+
+    @staticmethod
+    def _encode_positions(positions):
+        if positions is None:
+            return "-"
+        return ",".join(map(str, positions))
+
+    @staticmethod
+    def _decode_positions(raw):
+        if raw == "-":
+            return None
+        if raw == "":
+            return ()
+        return tuple(int(p) for p in raw.split(","))
+
+    def update_and_fetch(self, keys, values):
+        prev = self._stored
+        n_new = int(keys.shape[2])
+        start_pos = self.offset
+        required = prev + n_new
+        if self.keys is None or required > self.keys.shape[2]:
+            bsz, n_kv_heads, _, k_head_dim = keys.shape
+            v_head_dim = values.shape[3]
+            grow = ((required - prev + self.step - 1) // self.step) * self.step
+            new_k = mx.zeros((bsz, n_kv_heads, grow, k_head_dim), keys.dtype)
+            new_v = mx.zeros((bsz, n_kv_heads, grow, v_head_dim), values.dtype)
+            if self.keys is not None:
+                self.keys = mx.concatenate([self.keys[..., :prev, :], new_k], axis=2)
+                self.values = mx.concatenate(
+                    [self.values[..., :prev, :], new_v], axis=2
+                )
+            else:
+                self.keys, self.values = new_k, new_v
+
+        self.keys[..., prev:required, :] = keys
+        self.values[..., prev:required, :] = values
+        self._stored = required
+        self.offset += n_new
+        if self._positions is not None:
+            self._positions = self._positions + tuple(
+                range(start_pos, start_pos + n_new)
+            )
+        if self._speculating and n_new:
+            self._speculative_appends.append(n_new)
+        return (
+            self.keys[..., : self._stored, :],
+            self.values[..., : self._stored, :],
+        )
+
+    def size(self):
+        return self._stored
+
+    @property
+    def state(self):
+        if self.keys is None:
+            return None
+        return (
+            self.keys[..., : self._stored, :],
+            self.values[..., : self._stored, :],
+        )
+
+    @state.setter
+    def state(self, v):
+        if v is None:
+            self.keys = None
+            self.values = None
+            self._stored = 0
+            self.offset = 0
+            self._protected_stored = 0
+            self._positions = ()
+            self._speculating = False
+            self._speculative_appends = deque()
+            return
+        self.keys, self.values = v
+        self._stored = int(self.keys.shape[2])
+        self.offset = self._stored
+        self._protected_stored = self._stored
+        self._positions = tuple(range(self._stored))
+        self._speculating = False
+        self._speculative_appends = deque()
+
+    @property
+    def meta_state(self):
+        return tuple(
+            map(
+                str,
+                (
+                    self.offset,
+                    self._stored,
+                    self._protected_stored,
+                    self._encode_positions(self._positions),
+                ),
+            )
+        )
+
+    @meta_state.setter
+    def meta_state(self, v):
+        vals = tuple(v)
+        ints = tuple(map(int, vals[:3]))
+        if len(ints) == 2:
+            self.offset, self._stored = ints
+            self._protected_stored = self._stored
+        else:
+            self.offset, self._stored, self._protected_stored = ints
+        if len(vals) >= 4:
+            self._positions = self._decode_positions(str(vals[3]))
+            if self._positions is not None and len(self._positions) != self._stored:
+                raise ValueError("stored row count does not match position metadata")
+        elif self.offset == self._stored:
+            self._positions = tuple(range(self._stored))
+        else:
+            self._positions = None
+        self._speculating = False
+        self._speculative_appends = deque()
+
+    def is_trimmable(self):
+        return self._speculating or self._positions is not None
+
+    @property
+    def protected_stored(self):
+        return self._protected_stored
+
+    @property
+    def positions(self):
+        return self._positions
+
+    def start_speculation(self):
+        self._speculating = True
+        self._speculative_appends.clear()
+
+    def stop_speculation(self):
+        self._speculating = False
+        self._speculative_appends.clear()
+
+    def _trim_speculative_suffix(self, n):
+        if not self._speculative_appends:
+            raise RuntimeError("No speculative append is available to trim")
+        latest = self._speculative_appends.pop()
+        if n > latest:
+            self._speculative_appends.append(latest)
+            raise RuntimeError(
+                f"Cannot trim {n} tokens from PositionPreservingKVCache: "
+                f"latest speculative append has {latest} tokens."
+            )
+        removable = self._stored - self._protected_stored
+        if n > removable:
+            self._speculative_appends.append(latest)
+            raise RuntimeError(
+                f"Cannot trim {n} tokens from PositionPreservingKVCache: "
+                f"only {removable} appended rows are removable."
+            )
+        self._stored -= n
+        self.offset -= n
+        if self._positions is not None and n:
+            self._positions = self._positions[:-n]
+        if n < latest:
+            self._speculative_appends.append(latest - n)
+        if self.offset < 0:
+            raise ValueError("trim would make true offset negative")
+        return n
+
+    def _trim_logical_prefix(self, n):
+        if self._positions is None:
+            raise RuntimeError(
+                "PositionPreservingKVCache has no position metadata for prefix trim"
+            )
+        n = min(int(n), self.offset)
+        if n <= 0:
+            return 0
+        new_offset = self.offset - n
+        old_positions = self._positions
+        keep_physical = [i for i, pos in enumerate(old_positions) if pos < new_offset]
+        new_positions = tuple(old_positions[i] for i in keep_physical)
+        if len(keep_physical) != self._stored:
+            if keep_physical:
+                self.keys = _take_positions(
+                    self.keys[..., : self._stored, :], keep_physical
+                )
+                self.values = _take_positions(
+                    self.values[..., : self._stored, :], keep_physical
+                )
+            else:
+                self.keys = self.keys[..., :0, :]
+                self.values = self.values[..., :0, :]
+        self._protected_stored = sum(
+            1 for pos in old_positions[: self._protected_stored] if pos < new_offset
+        )
+        self._stored = len(new_positions)
+        self._positions = new_positions
+        self.offset = new_offset
+        return n
+
+    def trim(self, n):
+        if n <= 0:
+            return 0
+        n = int(n)
+        if self._speculating:
+            return self._trim_speculative_suffix(n)
+        return self._trim_logical_prefix(n)
+
+    def make_mask(self, n_tokens, window_size=None, return_array: bool = False):
+        if window_size is not None:
+            raise ValueError(
+                "PositionPreservingKVCache does not support sliding-window masks"
+            )
+        if n_tokens == 1 and not return_array:
+            return None
+        prefix = mx.ones((n_tokens, self._stored), dtype=mx.bool_)
+        causal = mx.tril(mx.ones((n_tokens, n_tokens), dtype=mx.bool_))
+        return mx.concatenate([prefix, causal], axis=1)
+
+    def empty(self):
+        return self.keys is None
+
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return (
+            self.keys[..., : self._stored, :].nbytes
+            + self.values[..., : self._stored, :].nbytes
+        )
+
+
+def _take_positions(x, keep: Sequence[int]):
+    return mx.take(x, mx.array(keep, dtype=mx.int32), axis=2)
+
+
+@dataclass(frozen=True)
+class SnapKVEvictionResult:
+    cache: list
+    evicted: bool
+    true_offset: int
+    retained_tokens: int
+    original_tokens: int
+    kv_layers: int
+    compact_cache_nbytes: int
+
+
+def evict_prompt_cache(
+    prompt_cache: List[Any],
+    keep_indices: Sequence[int],
+    *,
+    true_offset: int,
+) -> SnapKVEvictionResult:
+    """Replace plain ``KVCache`` layers with position-preserving compact caches.
+
+    ``keep_indices`` are the sorted prompt positions to retain (see
+    ``snapkv_keep_indices``); every other layer type is left untouched.
+    """
+    keep = tuple(int(i) for i in keep_indices)
+    if any(i < 0 for i in keep):
+        raise ValueError("keep_indices must be non-negative")
+    if tuple(sorted(keep)) != keep:
+        raise ValueError("keep_indices must be sorted")
+
+    out = list(prompt_cache)
+    original_tokens = int(true_offset)
+    kv_layers = 0
+    for idx, cache in enumerate(prompt_cache):
+        if type(cache) is not KVCache or cache.keys is None:
+            continue
+        if keep and keep[-1] >= cache.offset:
+            raise ValueError("keep index exceeds KV cache offset")
+        keys, values = cache.state
+        new_keys = _take_positions(keys, keep) if keep else keys[..., :0, :]
+        new_values = _take_positions(values, keep) if keep else values[..., :0, :]
+        out[idx] = PositionPreservingKVCache(
+            new_keys,
+            new_values,
+            offset=true_offset,
+            positions=keep,
+        )
+        original_tokens = max(original_tokens, int(cache.offset))
+        kv_layers += 1
+
+    retained = len(keep)
+    compact_cache_nbytes = sum(int(getattr(cache, "nbytes", 0) or 0) for cache in out)
+    return SnapKVEvictionResult(
+        cache=out,
+        evicted=kv_layers > 0 and retained < original_tokens,
+        true_offset=int(true_offset),
+        retained_tokens=retained,
+        original_tokens=original_tokens,
+        kv_layers=kv_layers,
+        compact_cache_nbytes=compact_cache_nbytes,
+    )
+
+
+class SnapKVAttentionCapture:
+    """Capture windowed attention scores by wrapping ``mx.fast`` SDPA.
+
+    Used as a context manager around a prefill. The model output still uses the
+    original fused kernel; this hook separately scores only the final
+    ``window`` observation-window query rows, in small chunks reduced
+    immediately to a per-key vector, so it never retains a prompt-sized dense
+    attention matrix. ``snap_scores`` then returns a per-prompt-position score.
+
+    The patch is process-global for the duration of the ``with`` block, so a
+    prefill scored this way should not run concurrently with unscored prefills.
+    """
+
+    def __init__(self, window: int = 24, score_chunk_size: int = 8):
+        if window <= 0:
+            raise ValueError("window must be positive")
+        if score_chunk_size <= 0:
+            raise ValueError("score_chunk_size must be positive")
+        self.window = int(window)
+        self.score_chunk_size = int(score_chunk_size)
+        self._orig = None
+        self.snap = None
+
+    def __enter__(self):
+        self._orig = mx.fast.scaled_dot_product_attention
+        mx.fast.scaled_dot_product_attention = self._capture
+        return self
+
+    def __exit__(self, *_exc):
+        mx.fast.scaled_dot_product_attention = self._orig
+
+    def _align_add(self, current, vec):
+        if current is None:
+            return vec
+        if current.shape[0] < vec.shape[0]:
+            current = mx.concatenate(
+                [current, mx.zeros((vec.shape[0] - current.shape[0],), current.dtype)]
+            )
+        elif vec.shape[0] < current.shape[0]:
+            vec = mx.concatenate(
+                [vec, mx.zeros((current.shape[0] - vec.shape[0],), vec.dtype)]
+            )
+        return current + vec
+
+    @staticmethod
+    def _causal_chunk_mask(q_start, q_end, q_len, key_len):
+        key_positions = mx.arange(key_len)
+        query_positions = key_len - q_len + mx.arange(q_start, q_end)
+        return key_positions[None, :] <= query_positions[:, None]
+
+    @staticmethod
+    def _slice_mask(mask, q_start, q_end, q_len):
+        if mask is None or isinstance(mask, str):
+            return mask
+        if len(mask.shape) >= 2 and mask.shape[-2] == q_len:
+            prefix = (slice(None),) * (len(mask.shape) - 2)
+            return mask[prefix + (slice(q_start, q_end), slice(None))]
+        return mask
+
+    @staticmethod
+    def _slice_query_heads(mask, start, end):
+        if mask is None or isinstance(mask, str):
+            return mask
+        if len(mask.shape) >= 4 and mask.shape[-3] >= end:
+            prefix = (slice(None),) * (len(mask.shape) - 3)
+            return mask[prefix + (slice(start, end), slice(None), slice(None))]
+        return mask
+
+    @staticmethod
+    def _apply_mask(scores, mask):
+        if mask is None or isinstance(mask, str):
+            return scores
+        return mx.where(mask, scores, -1e9) if mask.dtype == mx.bool_ else scores + mask
+
+    def _score_chunk(self, q_chunk, k, scale, chunk_mask):
+        _bsz, q_heads, _chunk, _dim = q_chunk.shape
+        kv_heads = k.shape[1]
+        if kv_heads == q_heads:
+            scores = (q_chunk @ k.swapaxes(-1, -2)) * scale
+            return self._apply_mask(scores, chunk_mask)
+        if q_heads % kv_heads != 0:
+            raise ValueError("query heads must be a multiple of KV heads")
+        repeats = q_heads // kv_heads
+        groups = []
+        for kv_head in range(kv_heads):
+            q_start = kv_head * repeats
+            q_end = q_start + repeats
+            q_group = q_chunk[:, q_start:q_end, :, :]
+            k_group = k[:, kv_head : kv_head + 1, :, :]
+            scores = (q_group @ k_group.swapaxes(-1, -2)) * scale
+            scores = self._apply_mask(
+                scores, self._slice_query_heads(chunk_mask, q_start, q_end)
+            )
+            groups.append(scores)
+        return mx.concatenate(groups, axis=1)
+
+    def _capture(self, q, k, v, *, scale, mask=None, **kwargs):
+        out = self._orig(q, k, v, scale=scale, mask=mask, **kwargs)
+        _bsz, q_heads, q_len, _dim = q.shape
+        key_len = k.shape[2]
+        snap_window = min(self.window, q_len)
+        q_start = q_len - snap_window
+        for chunk_start in range(q_start, q_len, self.score_chunk_size):
+            chunk_end = min(q_len, chunk_start + self.score_chunk_size)
+            q_chunk = q[:, :, chunk_start:chunk_end, :]
+            chunk_mask = self._slice_mask(mask, chunk_start, chunk_end, q_len)
+            if isinstance(chunk_mask, str) and chunk_mask == "causal":
+                chunk_mask = self._causal_chunk_mask(
+                    chunk_start, chunk_end, q_len, key_len
+                )
+            scores = self._score_chunk(q_chunk, k, scale, chunk_mask)
+            weights = mx.softmax(scores.astype(mx.float32), axis=-1)
+            snap = weights.sum(axis=(0, 1, 2))
+            mx.eval(snap)
+            self.snap = self._align_add(self.snap, snap)
+            mx.eval(self.snap)
+        return out
+
+    def snap_scores(self, seq_len: int) -> List[float]:
+        if self.snap is None:
+            return [0.0] * seq_len
+        mx.eval(self.snap)
+        vals = self.snap.tolist()
+        if len(vals) < seq_len:
+            vals.extend([0.0] * (seq_len - len(vals)))
+        return [float(v) for v in vals[:seq_len]]
+
+
+def compact_prompt_cache(
+    model,
+    prompt,
+    *,
+    budget: int,
+    window: int = 24,
+    sink_tokens: int = 4,
+    recent_tokens: Optional[int] = None,
+    min_tokens: int = 128,
+    score_chunk_size: int = 8,
+) -> SnapKVEvictionResult:
+    """Prefill ``prompt`` and return a SnapKV-D-compacted prompt cache.
+
+    Convenience wrapper: prefills ``prompt`` (a 1-D token sequence) under a
+    ``SnapKVAttentionCapture``, scores it, and evicts every full-attention
+    ``KVCache`` layer down to ``budget`` retained rows. The returned
+    ``SnapKVEvictionResult.cache`` is ready to decode from at the true prompt
+    offset. No-op (all rows kept) for prompts at or below ``min_tokens``.
+    """
+    prompt = list(int(t) for t in prompt)
+    cache = make_prompt_cache(model)
+    with SnapKVAttentionCapture(
+        window=window, score_chunk_size=score_chunk_size
+    ) as capture:
+        logits = model(mx.array([prompt]), cache=cache)
+        mx.eval(logits, [c.state for c in cache])
+    seq_len = len(prompt)
+    scores = capture.snap_scores(seq_len)
+    keep = snapkv_keep_indices(
+        seq_len,
+        budget,
+        scores,
+        sink_tokens=sink_tokens,
+        recent_tokens=recent_tokens,
+        min_tokens=min_tokens,
+    )
+    return evict_prompt_cache(cache, keep, true_offset=seq_len)
+
+
+# DuoAttention: head-partitioned KV eviction (retrieval vs streaming heads)
+#
+# DuoAttention (arXiv:2410.10819) observes that only a subset of attention
+# heads ("retrieval heads") need the full long context; the remaining
+# ("streaming heads") attend well with just attention sinks plus a recent
+# window. Keeping the full KV only for retrieval heads and a sink+recent slice
+# for streaming heads cuts the decode-time KV read further than a uniform
+# SnapKV-D budget, at no quality cost on the streaming heads.
+#
+# This is the SnapKV-D second stage: after the caller decides which prompt
+# positions each KV head keeps, ``HeadPartitionedKVCache`` stores the union of
+# retained positions once and exposes a per-head prefix mask so each head
+# attends only to its own retained rows. ``offset`` still tracks the true
+# sequence position for future RoPE. The head-CLASSIFICATION policy (which heads
+# are retrieval vs streaming) is intentionally out of scope here: this module
+# only provides the cache representation and the eviction op that applies a
+# caller-supplied ``head_keep_indices``.
+# ---------------------------------------------------------------------------
+
+
+class HeadPartitionedKVCache(_BaseCache):
+    """Sparse KV cache with independent retained positions per KV head.
+
+    This is the representation DuoAttention eviction needs. Storage uses the
+    union of all retained prefix positions, but ``make_mask`` exposes a per-head
+    prefix mask so retrieval heads can attend long-range retained rows while
+    streaming heads attend only their own recency rows. ``offset`` remains the
+    true logical sequence position for future RoPE.
+    """
+
+    def __init__(
+        self,
+        keys=None,
+        values=None,
+        *,
+        offset: int = 0,
+        positions: Optional[Sequence[int]] = None,
+        head_positions: Optional[Sequence[Sequence[int]]] = None,
+        query_heads: Optional[int] = None,
+        protected_stored: Optional[int] = None,
+    ):
+        self.keys = keys
+        self.values = values
+        self.offset = int(offset)
+        self._stored = 0 if keys is None else int(keys.shape[2])
+        self._positions = self._coerce_positions(positions)
+        self._head_positions = self._coerce_head_positions(head_positions)
+        self.query_heads = None if query_heads is None else int(query_heads)
+        self._head_position_mask = self._build_head_position_mask()
+        self._protected_stored = (
+            self._stored if protected_stored is None else int(protected_stored)
+        )
+        if self._protected_stored < 0 or self._protected_stored > self._stored:
+            raise ValueError("protected_stored must be between 0 and stored rows")
+        self._speculating = False
+        self._speculative_appends = deque()
+
+    def _coerce_positions(self, positions: Optional[Sequence[int]]):
+        if positions is None:
+            if self._stored == 0:
+                return ()
+            if self.offset == self._stored:
+                return tuple(range(self._stored))
+            raise ValueError("positions are required for sparse head-partitioned KV")
+        out = tuple(int(p) for p in positions)
+        if len(out) != self._stored:
+            raise ValueError("positions length must match stored rows")
+        if any(p < 0 for p in out):
+            raise ValueError("positions must be non-negative")
+        if tuple(sorted(out)) != out:
+            raise ValueError("positions must be sorted")
+        return out
+
+    def _coerce_head_positions(self, head_positions):
+        if self.keys is None:
+            n_heads = 0
+        else:
+            n_heads = int(self.keys.shape[1])
+        if head_positions is None:
+            return tuple(self._positions for _ in range(n_heads))
+        out = []
+        valid = set(self._positions)
+        for row in head_positions:
+            vals = tuple(int(p) for p in row)
+            if tuple(sorted(vals)) != vals:
+                raise ValueError("head positions must be sorted")
+            if any(p not in valid for p in vals):
+                raise ValueError("head positions must be a subset of positions")
+            out.append(vals)
+        if n_heads and len(out) != n_heads:
+            raise ValueError("head_positions length must match KV heads")
+        return tuple(out)
+
+    def _build_head_position_mask(self):
+        rows = []
+        for row in self._head_positions:
+            keep = set(row)
+            rows.append([pos in keep for pos in self._positions])
+        return tuple(tuple(row) for row in rows)
+
+    @staticmethod
+    def _encode_positions(positions):
+        return ",".join(map(str, positions))
+
+    @staticmethod
+    def _decode_positions(raw: str):
+        if raw == "":
+            return ()
+        return tuple(int(p) for p in raw.split(","))
+
+    @staticmethod
+    def _encode_head_positions(head_positions):
+        return "|".join(",".join(map(str, positions)) for positions in head_positions)
+
+    @classmethod
+    def _decode_head_positions(cls, raw: str):
+        if raw == "":
+            return ()
+        return tuple(cls._decode_positions(part) for part in raw.split("|"))
+
+    @property
+    def positions(self):
+        return self._positions
+
+    @property
+    def head_positions(self):
+        return self._head_positions
+
+    @property
+    def protected_stored(self):
+        return self._protected_stored
+
+    def size(self):
+        return self._stored
+
+    @property
+    def state(self):
+        if self.keys is None:
+            return None
+        return self.keys[..., : self._stored, :], self.values[..., : self._stored, :]
+
+    @state.setter
+    def state(self, v):
+        if v is None:
+            self.keys = None
+            self.values = None
+            self._stored = 0
+            self.offset = 0
+            self._positions = ()
+            self._head_positions = ()
+            self.query_heads = None
+            self._head_position_mask = ()
+            self._protected_stored = 0
+            self._speculating = False
+            self._speculative_appends = deque()
+            return
+        self.keys, self.values = v
+        self._stored = int(self.keys.shape[2])
+        self.offset = self._stored
+        self._positions = tuple(range(self._stored))
+        self._head_positions = tuple(
+            self._positions for _ in range(int(self.keys.shape[1]))
+        )
+        self.query_heads = None
+        self._head_position_mask = self._build_head_position_mask()
+        self._protected_stored = self._stored
+        self._speculating = False
+        self._speculative_appends = deque()
+
+    @property
+    def meta_state(self):
+        return tuple(
+            map(
+                str,
+                (
+                    self.offset,
+                    self._stored,
+                    self._protected_stored,
+                    self._encode_positions(self._positions),
+                    self._encode_head_positions(self._head_positions),
+                    "" if self.query_heads is None else self.query_heads,
+                ),
+            )
+        )
+
+    @meta_state.setter
+    def meta_state(self, v):
+        vals = tuple(v)
+        if len(vals) < 5:
+            raise ValueError("HeadPartitionedKVCache meta_state is incomplete")
+        self.offset, self._stored, self._protected_stored = map(int, vals[:3])
+        self._positions = self._decode_positions(str(vals[3]))
+        self._head_positions = self._decode_head_positions(str(vals[4]))
+        self.query_heads = None if len(vals) < 6 or str(vals[5]) == "" else int(vals[5])
+        if len(self._positions) != self._stored:
+            raise ValueError("stored row count does not match position metadata")
+        if self.keys is not None and len(self._head_positions) != int(
+            self.keys.shape[1]
+        ):
+            raise ValueError("head position count does not match KV heads")
+        self._head_position_mask = self._build_head_position_mask()
+        self._speculating = False
+        self._speculative_appends = deque()
+
+    def update_and_fetch(self, keys, values):
+        n_new = int(keys.shape[2])
+        start_pos = self.offset
+        if self.keys is None:
+            self.keys = keys
+            self.values = values
+            self._stored = n_new
+            self.offset = n_new
+            self._positions = tuple(range(n_new))
+            self._head_positions = tuple(
+                self._positions for _ in range(int(keys.shape[1]))
+            )
+            self.query_heads = None
+            self._head_position_mask = self._build_head_position_mask()
+            self._protected_stored = self._stored
+            return self.state
+
+        if int(keys.shape[1]) != int(self.keys.shape[1]):
+            raise ValueError("new keys must have the same KV head count")
+        self.keys = mx.concatenate([self.keys[..., : self._stored, :], keys], axis=2)
+        self.values = mx.concatenate(
+            [self.values[..., : self._stored, :], values], axis=2
+        )
+        new_positions = tuple(range(start_pos, start_pos + n_new))
+        self._stored += n_new
+        self.offset += n_new
+        self._positions = self._positions + new_positions
+        self._head_positions = tuple(
+            row + new_positions for row in self._head_positions
+        )
+        self._head_position_mask = self._build_head_position_mask()
+        if self._speculating and n_new:
+            self._speculative_appends.append(n_new)
+        return self.state
+
+    def make_mask(
+        self,
+        n_tokens: int,
+        window_size=None,
+        return_array: bool = False,
+        *,
+        query_heads: Optional[int] = None,
+    ):
+        if window_size is not None:
+            raise ValueError(
+                "HeadPartitionedKVCache does not support sliding-window masks"
+            )
+        if n_tokens <= 0:
+            raise ValueError("n_tokens must be positive")
+        prefix = mx.array(self._head_position_mask, dtype=mx.bool_)
+        if query_heads is not None:
+            effective_query_heads = int(query_heads)
+        elif self.query_heads is not None:
+            effective_query_heads = int(self.query_heads)
+        else:
+            effective_query_heads = None
+        if effective_query_heads is not None:
+            kv_heads = prefix.shape[0]
+            if effective_query_heads % kv_heads != 0:
+                raise ValueError("query_heads must be a multiple of KV heads")
+            prefix = mx.repeat(prefix, effective_query_heads // kv_heads, axis=0)
+        prefix = mx.broadcast_to(
+            prefix[:, None, :],
+            (prefix.shape[0], n_tokens, self._stored),
+        )
+        causal = mx.tril(mx.ones((n_tokens, n_tokens), dtype=mx.bool_))
+        causal = mx.broadcast_to(
+            causal[None, :, :],
+            (prefix.shape[0], n_tokens, n_tokens),
+        )
+        return mx.concatenate([prefix, causal], axis=2)[None, ...]
+
+    def is_trimmable(self):
+        return self._speculating or self._positions is not None
+
+    def start_speculation(self):
+        self._speculating = True
+        self._speculative_appends.clear()
+
+    def stop_speculation(self):
+        self._speculating = False
+        self._speculative_appends.clear()
+
+    def trim(self, n):
+        if n <= 0:
+            return 0
+        n = int(n)
+        old_positions = self._positions
+        if self._speculating:
+            if not self._speculative_appends:
+                raise RuntimeError("No speculative append is available to trim")
+            latest = self._speculative_appends.pop()
+            if n > latest:
+                self._speculative_appends.append(latest)
+                raise RuntimeError("trim exceeds latest speculative append")
+            removable = self._stored - self._protected_stored
+            if n > removable:
+                self._speculative_appends.append(latest)
+                raise RuntimeError("trim exceeds removable speculative rows")
+            if n < latest:
+                self._speculative_appends.append(latest - n)
+        n = min(n, self.offset)
+        new_offset = self.offset - n
+        keep_physical = [i for i, pos in enumerate(self._positions) if pos < new_offset]
+        if len(keep_physical) != self._stored:
+            if keep_physical:
+                self.keys = _take_positions(
+                    self.keys[..., : self._stored, :], keep_physical
+                )
+                self.values = _take_positions(
+                    self.values[..., : self._stored, :], keep_physical
+                )
+            else:
+                self.keys = self.keys[..., :0, :]
+                self.values = self.values[..., :0, :]
+        self._positions = tuple(old_positions[i] for i in keep_physical)
+        self._head_positions = tuple(
+            tuple(pos for pos in row if pos < new_offset)
+            for row in self._head_positions
+        )
+        self._stored = len(self._positions)
+        self._protected_stored = sum(
+            1 for pos in old_positions[: self._protected_stored] if pos < new_offset
+        )
+        self.offset = new_offset
+        self._head_position_mask = self._build_head_position_mask()
+        return n
+
+    def empty(self):
+        return self.keys is None
+
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return (
+            self.keys[..., : self._stored, :].nbytes
+            + self.values[..., : self._stored, :].nbytes
+        )
+
+
+def _union_positions(
+    head_keep_indices: Sequence[Sequence[int]],
+) -> Tuple[int, ...]:
+    out = sorted({int(pos) for row in head_keep_indices for pos in row})
+    return tuple(out)
+
+
+def _validate_head_keep_indices(head_keep_indices, n_heads: int, offset: int):
+    if len(head_keep_indices) != n_heads:
+        raise ValueError("head_keep_indices length must match KV heads")
+    out = []
+    for row in head_keep_indices:
+        vals = tuple(int(pos) for pos in row)
+        if tuple(sorted(vals)) != vals:
+            raise ValueError("head keep indices must be sorted")
+        if any(pos < 0 or pos >= offset for pos in vals):
+            raise ValueError("head keep index exceeds KV cache offset")
+        out.append(vals)
+    return tuple(out)
+
+
+def _take_head_partitioned(x, union_keep, head_keep_indices):
+    bsz, n_heads, _tokens, dim = x.shape
+    zero = mx.zeros((bsz, 1, 1, dim), dtype=x.dtype)
+    head_chunks = []
+    for head, keep in enumerate(head_keep_indices):
+        keep_set = set(keep)
+        rows = []
+        for pos in union_keep:
+            rows.append(
+                x[:, head : head + 1, pos : pos + 1, :] if pos in keep_set else zero
+            )
+        if rows:
+            head_chunks.append(mx.concatenate(rows, axis=2))
+        else:
+            head_chunks.append(x[:, head : head + 1, :0, :])
+    return mx.concatenate(head_chunks, axis=1) if head_chunks else x[:, :0, :0, :]
+
+
+@dataclass(frozen=True)
+class HeadPartitionedEvictionResult:
+    cache: list
+    evicted: bool
+    true_offset: int
+    union_retained_tokens: int
+    per_head_retained_tokens: Tuple[int, ...]
+    original_tokens: int
+    kv_layers: int
+    compact_cache_nbytes: int
+
+
+def evict_prompt_cache_by_head(
+    prompt_cache: List[Any],
+    head_keep_indices: Sequence[Sequence[int]],
+    *,
+    true_offset: int,
+    query_heads: Optional[int] = None,
+) -> HeadPartitionedEvictionResult:
+    """Replace plain ``KVCache`` layers with a DuoAttention-ready cache.
+
+    ``head_keep_indices`` gives, per KV head, the sorted prompt positions that
+    head retains (retrieval heads typically keep everything; streaming heads
+    keep only sinks + a recent window). Every other layer type is left
+    untouched. The head-classification decision is the caller's responsibility.
+    """
+    out = list(prompt_cache)
+    original_tokens = int(true_offset)
+    kv_layers = 0
+    union_keep = _union_positions(head_keep_indices)
+    per_head_counts = tuple(len(row) for row in head_keep_indices)
+    for idx, cache in enumerate(prompt_cache):
+        if type(cache) is not KVCache or cache.keys is None:
+            continue
+        keys, values = cache.state
+        n_heads = int(keys.shape[1])
+        head_keep = _validate_head_keep_indices(
+            head_keep_indices, n_heads, int(cache.offset)
+        )
+        union_keep = _union_positions(head_keep)
+        new_keys = _take_head_partitioned(keys, union_keep, head_keep)
+        new_values = _take_head_partitioned(values, union_keep, head_keep)
+        out[idx] = HeadPartitionedKVCache(
+            new_keys,
+            new_values,
+            offset=true_offset,
+            positions=union_keep,
+            head_positions=head_keep,
+            query_heads=query_heads,
+        )
+        original_tokens = max(original_tokens, int(cache.offset))
+        kv_layers += 1
+
+    compact_cache_nbytes = sum(int(getattr(cache, "nbytes", 0) or 0) for cache in out)
+    return HeadPartitionedEvictionResult(
+        cache=out,
+        evicted=kv_layers > 0
+        and any(count < original_tokens for count in per_head_counts),
+        true_offset=int(true_offset),
+        union_retained_tokens=len(union_keep),
+        per_head_retained_tokens=per_head_counts,
+        original_tokens=original_tokens,
+        kv_layers=kv_layers,
+        compact_cache_nbytes=compact_cache_nbytes,
+    )
diff --git a/mlx_lm/server.py b/mlx_lm/server.py
index 14d904cd8..6c26caeba 100644
--- a/mlx_lm/server.py
+++ b/mlx_lm/server.py
@@ -40,7 +40,13 @@
     make_text_state_machine,
     stream_generate,
 )
-from .models.cache import LRUPromptCache, make_prompt_cache
+from .models.cache import (
+    LRUPromptCache,
+    SnapKVAttentionCapture,
+    evict_prompt_cache,
+    make_prompt_cache,
+    snapkv_keep_indices,
+)
 from .sample_utils import make_logits_processors, make_sampler
 from .utils import _parse_size, load, sharded_load
 
@@ -354,6 +360,11 @@ def _load(self, model_path, adapter_path=None, draft_model_path=None):
         is_batchable = is_batchable and all(
             hasattr(c, "merge") for c in make_prompt_cache(model)
         )
+        # SnapKV-D scores a prefill with a process-global attention hook, which
+        # cannot separate per-request scores in a shared batch, so it runs on
+        # the single-request path.
+        if self.cli_args.kv_eviction == "snapkv":
+            is_batchable = False
 
         # Update the member variables
         self.model_key = (model_path, adapter_path, draft_model_path)
@@ -868,6 +879,38 @@ def get_next_request(timeout=None):
                         # generation
                         batch_results.pop(uid, None)
 
+    def _maybe_snapkv_compact(self, model, cache, rest, made_fresh):
+        """Post-prefill SnapKV-D eviction for a fresh full prefill.
+
+        Returns ``(cache, decode_prompt)``. When ``--kv-eviction snapkv`` is on
+        and the fresh prompt is longer than the floor, this prefills
+        ``rest[:-1]`` under a scoring hook, compacts every full-attention layer
+        to ``--kv-budget`` retained rows, and returns the last prompt token to
+        decode from. Otherwise it returns ``(cache, rest)`` unchanged.
+        """
+        args = self.cli_args
+        if (
+            getattr(args, "kv_eviction", "none") != "snapkv"
+            or not made_fresh
+            or len(rest) <= args.kv_min_tokens
+        ):
+            return cache, rest
+
+        prefill_ids = rest[:-1]
+        with SnapKVAttentionCapture(window=args.kv_window) as capture:
+            model(mx.array([prefill_ids]), cache=cache)
+            mx.eval([c.state for c in cache])
+        scores = capture.snap_scores(len(prefill_ids))
+        keep = snapkv_keep_indices(
+            len(prefill_ids),
+            args.kv_budget,
+            scores,
+            sink_tokens=args.kv_sink_tokens,
+            min_tokens=args.kv_min_tokens,
+        )
+        result = evict_prompt_cache(cache, keep, true_offset=len(prefill_ids))
+        return result.cache, rest[-1:]
+
     def _serve_single(self, request):
         rqueue, request, args = request
 
@@ -915,11 +958,18 @@ def progress(tokens_processed, tokens_total):
             )
             ctx.prompt_cache_count = len(prompt) - len(rest)
             cache_key = prompt[:]
+            made_fresh = cache is None
             if cache is None:
                 cache = make_prompt_cache(self.model_provider.model)
                 if self.model_provider.draft_model is not None:
                     cache += make_prompt_cache(self.model_provider.draft_model)
 
+            # Post-prefill SnapKV-D eviction (opt-in): on a full miss, prefill
+            # the prompt under a scoring hook, compact the cache to the budget,
+            # and decode from the last token. Only for a fresh full prefill so
+            # the scored/retained positions are exact.
+            cache, rest = self._maybe_snapkv_compact(model, cache, rest, made_fresh)
+
             # Process the prompt and generate tokens
             stop_state = stop_matcher.make_state()
             for gen in stream_generate(
@@ -1848,6 +1898,42 @@ def main():
         type=_parse_size,
         help="Maximum size in bytes of the KV caches",
     )
+    parser.add_argument(
+        "--kv-eviction",
+        type=str,
+        default="none",
+        choices=["none", "snapkv"],
+        help=(
+            "Post-prefill KV cache eviction policy. 'snapkv' compacts each "
+            "full-attention layer to --kv-budget retained rows (sinks + recent "
+            "+ top attention-scored) after a fresh prefill, cutting long-context "
+            "decode KV reads. Default 'none'. Runs on the single-request path."
+        ),
+    )
+    parser.add_argument(
+        "--kv-budget",
+        type=int,
+        default=512,
+        help="Retained KV rows per layer when --kv-eviction is snapkv",
+    )
+    parser.add_argument(
+        "--kv-window",
+        type=int,
+        default=24,
+        help="SnapKV observation window (final prompt query rows used to score)",
+    )
+    parser.add_argument(
+        "--kv-sink-tokens",
+        type=int,
+        default=4,
+        help="Leading attention-sink rows always retained by SnapKV eviction",
+    )
+    parser.add_argument(
+        "--kv-min-tokens",
+        type=int,
+        default=128,
+        help="Prompts at or below this length are never SnapKV-compacted",
+    )
     parser.add_argument(
         "--pipeline",
         action="store_true",
diff --git a/tests/test_server_kv_eviction.py b/tests/test_server_kv_eviction.py
new file mode 100644
index 000000000..1d6159f39
--- /dev/null
+++ b/tests/test_server_kv_eviction.py
@@ -0,0 +1,97 @@
+# Copyright © 2024 Apple Inc.
+
+import unittest
+from argparse import Namespace
+
+import mlx.core as mx
+from test_snapkv_cache import TinyAttnModel
+
+from mlx_lm.models.cache import PositionPreservingKVCache
+from mlx_lm.server import ResponseGenerator
+
+
+def _args(**overrides):
+    base = dict(
+        kv_eviction="none",
+        kv_budget=48,
+        kv_window=24,
+        kv_sink_tokens=4,
+        kv_min_tokens=128,
+    )
+    base.update(overrides)
+    return Namespace(**base)
+
+
+class _StubGenerator:
+    """Exercise ResponseGenerator._maybe_snapkv_compact without a live server."""
+
+    def __init__(self, args):
+        self._args = args
+
+    @property
+    def cli_args(self):
+        return self._args
+
+    _maybe_snapkv_compact = ResponseGenerator._maybe_snapkv_compact
+
+
+class TestSnapKVServerCompaction(unittest.TestCase):
+    def setUp(self):
+        mx.random.seed(0)
+        self.model = TinyAttnModel(vocab=64, dim=32, n_layers=2, n_heads=4)
+        mx.eval(self.model.parameters())
+        self.prompt = [int(t) for t in mx.random.randint(0, 64, shape=(260,)).tolist()]
+
+    def _fresh_cache(self):
+        return self.model.make_cache()
+
+    def test_eviction_off_is_passthrough(self):
+        gen = _StubGenerator(_args(kv_eviction="none"))
+        cache = self._fresh_cache()
+        out_cache, decode = gen._maybe_snapkv_compact(
+            self.model, cache, self.prompt, made_fresh=True
+        )
+        self.assertIs(out_cache, cache)
+        self.assertEqual(decode, self.prompt)
+
+    def test_snapkv_compacts_fresh_long_prompt(self):
+        gen = _StubGenerator(_args(kv_eviction="snapkv", kv_budget=48))
+        cache = self._fresh_cache()
+        out_cache, decode = gen._maybe_snapkv_compact(
+            self.model, cache, self.prompt, made_fresh=True
+        )
+        # Decodes from the last prompt token.
+        self.assertEqual(decode, self.prompt[-1:])
+        # Full-attention layers are compacted, offset preserved at prompt-1.
+        for c in out_cache:
+            self.assertIsInstance(c, PositionPreservingKVCache)
+            self.assertEqual(c.offset, len(self.prompt) - 1)
+            self.assertEqual(c.size(), 48)
+        # And it actually decodes.
+        logits = self.model(mx.array([decode]), cache=out_cache)
+        mx.eval(logits)
+        self.assertEqual(logits.shape, (1, 1, 64))
+
+    def test_short_prompt_not_compacted(self):
+        gen = _StubGenerator(_args(kv_eviction="snapkv", kv_min_tokens=128))
+        short = self.prompt[:100]
+        cache = self._fresh_cache()
+        out_cache, decode = gen._maybe_snapkv_compact(
+            self.model, cache, short, made_fresh=True
+        )
+        self.assertIs(out_cache, cache)
+        self.assertEqual(decode, short)
+
+    def test_partial_hit_not_compacted(self):
+        # made_fresh=False means a partial prefix-cache hit; skip compaction.
+        gen = _StubGenerator(_args(kv_eviction="snapkv"))
+        cache = self._fresh_cache()
+        out_cache, decode = gen._maybe_snapkv_compact(
+            self.model, cache, self.prompt, made_fresh=False
+        )
+        self.assertIs(out_cache, cache)
+        self.assertEqual(decode, self.prompt)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_snapkv_cache.py b/tests/test_snapkv_cache.py
new file mode 100644
index 000000000..9cdbd1eca
--- /dev/null
+++ b/tests/test_snapkv_cache.py
@@ -0,0 +1,225 @@
+# Copyright © 2024 Apple Inc.
+
+import unittest
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mlx_lm.models.cache import (
+    KVCache,
+    PositionPreservingKVCache,
+    SnapKVAttentionCapture,
+    compact_prompt_cache,
+    evict_prompt_cache,
+    snapkv_keep_indices,
+    trim_prompt_cache,
+)
+
+
+class TinyAttnModel(nn.Module):
+    """A tiny causal-attention LM that routes through mx.fast SDPA (so the
+    SnapKV capture hook fires) and caches with KVCache."""
+
+    def __init__(self, vocab=64, dim=32, n_layers=2, n_heads=4):
+        super().__init__()
+        self.embed = nn.Embedding(vocab, dim)
+        self.n_heads = n_heads
+        self.head_dim = dim // n_heads
+        self.scale = self.head_dim**-0.5
+        self.qkv = [nn.Linear(dim, 3 * dim, bias=False) for _ in range(n_layers)]
+        self.o = [nn.Linear(dim, dim, bias=False) for _ in range(n_layers)]
+        self.out = nn.Linear(dim, vocab, bias=False)
+        self.n_layers = n_layers
+
+    def make_cache(self):
+        return [KVCache() for _ in range(self.n_layers)]
+
+    def __call__(self, inputs, cache=None):
+        B, T = inputs.shape
+        x = self.embed(inputs)
+        if cache is None:
+            cache = [None] * self.n_layers
+        for qkv, o, c in zip(self.qkv, self.o, cache):
+            H, d = self.n_heads, self.head_dim
+            q, k, v = mx.split(qkv(x), 3, axis=-1)
+            q = q.reshape(B, T, H, d).transpose(0, 2, 1, 3)
+            k = k.reshape(B, T, H, d).transpose(0, 2, 1, 3)
+            v = v.reshape(B, T, H, d).transpose(0, 2, 1, 3)
+            if c is not None:
+                k, v = c.update_and_fetch(k, v)
+            key_len = k.shape[2]
+            if T > 1:
+                qpos = key_len - T + mx.arange(T)
+                mask = (mx.arange(key_len)[None, :] <= qpos[:, None]).astype(x.dtype)
+                mask = mx.where(mask > 0, 0.0, -1e9)
+            else:
+                mask = None
+            att = mx.fast.scaled_dot_product_attention(
+                q, k, v, scale=self.scale, mask=mask
+            )
+            att = att.transpose(0, 2, 1, 3).reshape(B, T, H * d)
+            x = x + o(att)
+        return self.out(x)
+
+
+class TestSnapKVKeepIndices(unittest.TestCase):
+    def test_keeps_sinks_recent_and_top_scored(self):
+        seq_len = 200
+        scores = [0.0] * seq_len
+        scores[100] = 9.0  # a strong middle token
+        scores[150] = 8.0
+        keep = snapkv_keep_indices(
+            seq_len, budget=64, scores=scores, sink_tokens=4, min_tokens=128
+        )
+        self.assertEqual(len(keep), 64)
+        self.assertEqual(keep, tuple(sorted(keep)))
+        for s in range(4):  # sinks (budget//8 == 8 here, so all 4 fit)
+            self.assertIn(s, keep)
+        self.assertIn(seq_len - 1, keep)  # recent window
+        self.assertIn(100, keep)  # top-scored middle tokens survive
+        self.assertIn(150, keep)
+
+    def test_noop_paths(self):
+        # budget covers the prompt, or prompt <= min_tokens -> keep everything.
+        self.assertEqual(
+            snapkv_keep_indices(50, 999, [0.0] * 50, min_tokens=128),
+            tuple(range(50)),
+        )
+        self.assertEqual(
+            snapkv_keep_indices(80, 16, [0.0] * 80, min_tokens=128),
+            tuple(range(80)),
+        )
+
+    def test_validation(self):
+        with self.assertRaises(ValueError):
+            snapkv_keep_indices(10, 0, [0.0] * 10)
+        with self.assertRaises(ValueError):
+            snapkv_keep_indices(10, 4, [0.0] * 3)
+        with self.assertRaises(ValueError):
+            snapkv_keep_indices(-1, 4, [])
+
+
+class TestPositionPreservingKVCache(unittest.TestCase):
+    def _fill(self, cache, n):
+        x = mx.random.uniform(shape=(1, 2, n, 4))
+        return cache.update_and_fetch(x, x)
+
+    def test_offset_and_stored_diverge_after_eviction(self):
+        # Build from a sparse set of positions: stored rows < true offset.
+        keys = mx.random.uniform(shape=(1, 2, 3, 4))
+        c = PositionPreservingKVCache(keys, keys, offset=100, positions=(0, 1, 99))
+        self.assertEqual(c.size(), 3)
+        self.assertEqual(c.offset, 100)
+        self.assertEqual(c.positions, (0, 1, 99))
+        self.assertTrue(c.is_trimmable())
+
+    def test_state_meta_state_roundtrip(self):
+        keys = mx.random.uniform(shape=(1, 2, 3, 4))
+        c = PositionPreservingKVCache(keys, keys, offset=100, positions=(0, 1, 99))
+        restored = PositionPreservingKVCache.from_state(c.state, c.meta_state)
+        self.assertEqual(restored.offset, 100)
+        self.assertEqual(restored.size(), 3)
+        self.assertEqual(restored.positions, (0, 1, 99))
+
+    def test_logical_prefix_trim_drops_out_of_range_positions(self):
+        keys = mx.random.uniform(shape=(1, 2, 3, 4))
+        c = PositionPreservingKVCache(keys, keys, offset=100, positions=(0, 40, 99))
+        trimmed = c.trim(70)  # new offset 30 -> keep only position 0
+        self.assertEqual(trimmed, 70)
+        self.assertEqual(c.offset, 30)
+        self.assertEqual(c.positions, (0,))
+        self.assertEqual(c.size(), 1)
+
+    def test_speculative_suffix_trim(self):
+        keys = mx.random.uniform(shape=(1, 2, 5, 4))
+        c = PositionPreservingKVCache(keys, keys, offset=5, positions=tuple(range(5)))
+        c.start_speculation()
+        self._fill(c, 3)  # append 3 speculative rows
+        self.assertEqual(c.offset, 8)
+        c.trim(2)  # roll back 2 of the appended rows
+        self.assertEqual(c.offset, 6)
+        self.assertEqual(c.size(), 6)
+        c.stop_speculation()
+
+    def test_growth_and_nbytes(self):
+        c = PositionPreservingKVCache()
+        self._fill(c, 300)  # crosses the step=256 growth boundary
+        self.assertEqual(c.size(), 300)
+        self.assertEqual(c.offset, 300)
+        self.assertGreater(c.nbytes, 0)
+
+
+class TestEvictPromptCache(unittest.TestCase):
+    def test_evicts_kv_layers_preserving_offset(self):
+        cache = [KVCache(), KVCache()]
+        x = mx.random.uniform(shape=(1, 2, 300, 4))
+        for c in cache:
+            c.update_and_fetch(x, x)
+        scores = [0.0] * 300
+        scores[123] = 5.0
+        keep = snapkv_keep_indices(300, 32, scores, sink_tokens=4, min_tokens=128)
+        result = evict_prompt_cache(cache, keep, true_offset=300)
+        self.assertTrue(result.evicted)
+        self.assertEqual(result.retained_tokens, len(keep))
+        self.assertEqual(result.kv_layers, 2)
+        for c in result.cache:
+            self.assertIsInstance(c, PositionPreservingKVCache)
+            self.assertEqual(c.offset, 300)  # true position preserved
+            self.assertEqual(c.size(), len(keep))
+        self.assertLess(result.compact_cache_nbytes, x.nbytes * 2)
+
+    def test_non_kv_layers_untouched(self):
+        from mlx_lm.models.cache import ArraysCache
+
+        cache = [ArraysCache(size=1)]
+        cache[0][0] = mx.zeros((1, 4))
+        result = evict_prompt_cache(cache, (0,), true_offset=1)
+        self.assertIs(result.cache[0], cache[0])
+        self.assertEqual(result.kv_layers, 0)
+
+
+class TestCompactPromptCacheEndToEnd(unittest.TestCase):
+    def test_compaction_scores_and_decodes(self):
+        mx.random.seed(0)
+        model = TinyAttnModel(vocab=64, dim=32, n_layers=2, n_heads=4)
+        mx.eval(model.parameters())
+        prompt = [int(t) for t in mx.random.randint(0, 64, shape=(260,)).tolist()]
+
+        result = compact_prompt_cache(model, prompt, budget=48, min_tokens=128)
+        self.assertTrue(result.evicted)
+        self.assertEqual(result.retained_tokens, 48)
+        for c in result.cache:
+            self.assertEqual(c.offset, len(prompt))  # RoPE offset preserved
+            self.assertEqual(c.size(), 48)
+
+        # The compacted cache decodes: append the next token, offset advances.
+        logits = model(mx.array([[prompt[-1]]]), cache=result.cache)
+        mx.eval(logits)
+        self.assertEqual(logits.shape, (1, 1, 64))
+        for c in result.cache:
+            self.assertEqual(c.offset, len(prompt) + 1)
+
+    def test_capture_produces_nontrivial_scores(self):
+        mx.random.seed(1)
+        model = TinyAttnModel(vocab=64, dim=32, n_layers=2, n_heads=4)
+        mx.eval(model.parameters())
+        prompt = [int(t) for t in mx.random.randint(0, 64, shape=(200,)).tolist()]
+        cache = model.make_cache()
+        with SnapKVAttentionCapture(window=24) as cap:
+            model(mx.array([prompt]), cache=cache)
+        scores = cap.snap_scores(len(prompt))
+        self.assertEqual(len(scores), len(prompt))
+        self.assertGreater(max(scores), 0.0)  # the hook actually fired
+
+    def test_short_prompt_is_noop(self):
+        mx.random.seed(2)
+        model = TinyAttnModel()
+        mx.eval(model.parameters())
+        prompt = [int(t) for t in mx.random.randint(0, 64, shape=(64,)).tolist()]
+        result = compact_prompt_cache(model, prompt, budget=16, min_tokens=128)
+        self.assertFalse(result.evicted)
+        self.assertEqual(result.retained_tokens, 64)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_snapkv_duoattention.py b/tests/test_snapkv_duoattention.py
new file mode 100644
index 000000000..3b409add2
--- /dev/null
+++ b/tests/test_snapkv_duoattention.py
@@ -0,0 +1,210 @@
+# Copyright © 2024 Apple Inc.
+
+import unittest
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mlx_lm.models.cache import (
+    HeadPartitionedKVCache,
+    KVCache,
+    evict_prompt_cache_by_head,
+)
+
+
+class TinyAttnModel(nn.Module):
+    """A tiny causal-attention LM that routes through mx.fast SDPA and caches
+    with KVCache (mirrors tests/test_snapkv_cache.py)."""
+
+    def __init__(self, vocab=64, dim=32, n_layers=2, n_heads=4):
+        super().__init__()
+        self.embed = nn.Embedding(vocab, dim)
+        self.n_heads = n_heads
+        self.head_dim = dim // n_heads
+        self.scale = self.head_dim**-0.5
+        self.qkv = [nn.Linear(dim, 3 * dim, bias=False) for _ in range(n_layers)]
+        self.o = [nn.Linear(dim, dim, bias=False) for _ in range(n_layers)]
+        self.out = nn.Linear(dim, vocab, bias=False)
+        self.n_layers = n_layers
+
+    def make_cache(self):
+        return [KVCache() for _ in range(self.n_layers)]
+
+    def __call__(self, inputs, cache=None):
+        B, T = inputs.shape
+        x = self.embed(inputs)
+        if cache is None:
+            cache = [None] * self.n_layers
+        for qkv, o, c in zip(self.qkv, self.o, cache):
+            H, d = self.n_heads, self.head_dim
+            q, k, v = mx.split(qkv(x), 3, axis=-1)
+            q = q.reshape(B, T, H, d).transpose(0, 2, 1, 3)
+            k = k.reshape(B, T, H, d).transpose(0, 2, 1, 3)
+            v = v.reshape(B, T, H, d).transpose(0, 2, 1, 3)
+            if c is not None:
+                k, v = c.update_and_fetch(k, v)
+            key_len = k.shape[2]
+            if T > 1:
+                qpos = key_len - T + mx.arange(T)
+                mask = (mx.arange(key_len)[None, :] <= qpos[:, None]).astype(x.dtype)
+                mask = mx.where(mask > 0, 0.0, -1e9)
+            else:
+                mask = None
+            att = mx.fast.scaled_dot_product_attention(
+                q, k, v, scale=self.scale, mask=mask
+            )
+            att = att.transpose(0, 2, 1, 3).reshape(B, T, H * d)
+            x = x + o(att)
+        return self.out(x)
+
+
+class TestHeadPartitionedKVCache(unittest.TestCase):
+    def _build(self):
+        # Two KV heads, four stored rows (union of head-0/head-1 keeps).
+        keys = mx.random.uniform(shape=(1, 2, 4, 4))
+        # union positions = {0, 1, 98, 99}; head 0 keeps all, head 1 keeps
+        # only the sink (0) and recent (99).
+        return HeadPartitionedKVCache(
+            keys,
+            keys,
+            offset=100,
+            positions=(0, 1, 98, 99),
+            head_positions=((0, 1, 98, 99), (0, 99)),
+            query_heads=4,
+        )
+
+    def test_construction_and_accessors(self):
+        c = self._build()
+        self.assertEqual(c.size(), 4)
+        self.assertEqual(c.offset, 100)
+        self.assertEqual(c.positions, (0, 1, 98, 99))
+        self.assertEqual(c.head_positions, ((0, 1, 98, 99), (0, 99)))
+        self.assertTrue(c.is_trimmable())
+        self.assertGreater(c.nbytes, 0)
+        self.assertFalse(c.empty())
+
+    def test_state_meta_state_roundtrip(self):
+        c = self._build()
+        restored = HeadPartitionedKVCache.from_state(c.state, c.meta_state)
+        self.assertEqual(restored.offset, 100)
+        self.assertEqual(restored.size(), 4)
+        self.assertEqual(restored.positions, (0, 1, 98, 99))
+        self.assertEqual(restored.head_positions, ((0, 1, 98, 99), (0, 99)))
+        self.assertEqual(restored.query_heads, 4)
+
+    def test_make_mask_per_head_semantics(self):
+        c = self._build()
+        # No query_heads expansion: one row per KV head.
+        mask = c.make_mask(1, query_heads=2)
+        self.assertEqual(mask.shape, (1, 2, 1, 4 + 1))
+        prefix = mask[0, :, 0, :4]  # (kv_heads, stored)
+        # head 0 attends to all four retained rows; head 1 only to {0, 99}.
+        self.assertEqual(prefix[0].tolist(), [True, True, True, True])
+        self.assertEqual(prefix[1].tolist(), [True, False, False, True])
+        # The trailing block is causal over the new tokens.
+        self.assertEqual(mask[0, 0, 0, 4].tolist(), True)
+
+    def test_make_mask_query_head_expansion(self):
+        c = self._build()
+        mask = c.make_mask(1, query_heads=4)  # 4 query heads over 2 KV heads
+        self.assertEqual(mask.shape, (1, 4, 1, 4 + 1))
+        prefix = mask[0, :, 0, :4]
+        # query heads 0,1 map to KV head 0; 2,3 map to KV head 1.
+        self.assertEqual(prefix[0].tolist(), [True, True, True, True])
+        self.assertEqual(prefix[1].tolist(), [True, True, True, True])
+        self.assertEqual(prefix[2].tolist(), [True, False, False, True])
+        self.assertEqual(prefix[3].tolist(), [True, False, False, True])
+
+    def test_trim_drops_out_of_range_positions(self):
+        c = self._build()
+        trimmed = c.trim(50)  # new offset 50 -> keep only positions {0, 1}
+        self.assertEqual(trimmed, 50)
+        self.assertEqual(c.offset, 50)
+        self.assertEqual(c.positions, (0, 1))
+        self.assertEqual(c.head_positions, ((0, 1), (0,)))
+        self.assertEqual(c.size(), 2)
+
+    def test_update_and_fetch_appends_and_advances(self):
+        c = self._build()
+        new = mx.random.uniform(shape=(1, 2, 1, 4))
+        keys, values = c.update_and_fetch(new, new)
+        self.assertEqual(c.offset, 101)
+        self.assertEqual(c.size(), 5)
+        self.assertEqual(keys.shape[2], 5)
+        # The appended row (position 100) is visible to every head.
+        self.assertEqual(c.head_positions[0][-1], 100)
+        self.assertEqual(c.head_positions[1][-1], 100)
+
+
+class TestEvictPromptCacheByHead(unittest.TestCase):
+    def test_asymmetric_per_head_eviction(self):
+        cache = [KVCache(), KVCache()]
+        x = mx.random.uniform(shape=(1, 2, 300, 4))
+        for c in cache:
+            c.update_and_fetch(x, x)
+        original = sum(c.nbytes for c in cache)  # dense keys+values, all rows
+        # head 0 = retrieval: keeps a long window (sinks + first 200 + recent).
+        # head 1 = streaming: keeps only sinks (0..3) + recent (296..299).
+        # The union drops the 200..295 middle band, so storage shrinks.
+        head0 = tuple(range(200)) + tuple(range(296, 300))
+        head1 = tuple(range(4)) + tuple(range(296, 300))
+        head_keep = (head0, head1)
+        result = evict_prompt_cache_by_head(
+            cache, head_keep, true_offset=300, query_heads=4
+        )
+        self.assertTrue(result.evicted)
+        self.assertEqual(result.kv_layers, 2)
+        self.assertEqual(result.per_head_retained_tokens, (204, 8))
+        # union = head 0's retained rows (head 1's are a subset of them).
+        self.assertEqual(result.union_retained_tokens, 204)
+        for c in result.cache:
+            self.assertIsInstance(c, HeadPartitionedKVCache)
+            self.assertEqual(c.offset, 300)
+            self.assertEqual(c.head_positions, head_keep)
+        self.assertLess(result.compact_cache_nbytes, original)
+
+    def test_non_kv_layers_untouched(self):
+        from mlx_lm.models.cache import ArraysCache
+
+        cache = [ArraysCache(size=1)]
+        cache[0][0] = mx.zeros((1, 4))
+        result = evict_prompt_cache_by_head(cache, ((0,), (0,)), true_offset=1)
+        self.assertIs(result.cache[0], cache[0])
+        self.assertEqual(result.kv_layers, 0)
+
+
+class TestHeadPartitionedEndToEnd(unittest.TestCase):
+    def test_prefill_evict_by_head_then_decode(self):
+        mx.random.seed(0)
+        model = TinyAttnModel(vocab=64, dim=32, n_layers=2, n_heads=4)
+        mx.eval(model.parameters())
+        prompt = [int(t) for t in mx.random.randint(0, 64, shape=(200,)).tolist()]
+
+        cache = model.make_cache()
+        logits = model(mx.array([prompt]), cache=cache)
+        mx.eval(logits, [c.state for c in cache])
+
+        seq_len = len(prompt)
+        # The model has 4 KV heads: heads 0,1 are retrieval (keep the full
+        # context), heads 2,3 are streaming (sinks + recent window only).
+        retrieval = tuple(range(seq_len))
+        streaming = tuple(range(4)) + tuple(range(seq_len - 8, seq_len))
+        head_keep = (retrieval, retrieval, streaming, streaming)
+        result = evict_prompt_cache_by_head(
+            cache, head_keep, true_offset=seq_len, query_heads=model.n_heads
+        )
+        self.assertTrue(result.evicted)
+        for c in result.cache:
+            self.assertIsInstance(c, HeadPartitionedKVCache)
+            self.assertEqual(c.offset, seq_len)
+
+        # Decode one token from the head-partitioned cache.
+        out = model(mx.array([[prompt[-1]]]), cache=result.cache)
+        mx.eval(out)
+        self.assertEqual(out.shape, (1, 1, 64))
+        for c in result.cache:
+            self.assertEqual(c.offset, seq_len + 1)
+
+
+if __name__ == "__main__":
+    unittest.main()