diff --git a/gmlx/gen/generation.py b/gmlx/gen/generation.py
index 9c627b2b..334ba1a6 100644
--- a/gmlx/gen/generation.py
+++ b/gmlx/gen/generation.py
@@ -1212,7 +1212,7 @@ def generate_speculative_owned(
         # Owned rounds take the layers serve takes; the stock walks
         # decline. No later converter runs here, so convert now.
         from gmlx.cache.kv_policy import quantize_stack, resolve_and_report
-        from gmlx.spec.engine import mtp_kv_decline
+        from gmlx.spec.kv_quant import mtp_kv_decline
 
         decline = mtp_kv_decline(lm)
         policy = resolve_and_report(
diff --git a/gmlx/serve/patches/__init__.py b/gmlx/serve/patches/__init__.py
index b35e5ee8..b0473d0e 100644
--- a/gmlx/serve/patches/__init__.py
+++ b/gmlx/serve/patches/__init__.py
@@ -214,11 +214,14 @@ def install_server_patches(cfg, *, reload_fn=None) -> None:
         lora_rows.configure("rows", 1)
     lora_rows.install_lora_gen_args()
     lora_rows.install_row_channel()
-    import gmlx.spec.engine as spec_engine
-    spec_engine.install_full_prompt_mtp_prefill()
-    spec_engine.install_owned_spec_engine()
-    spec_engine.install_continuous_batch_admission()
-    spec_engine.install_spec_kv_quant()
+    from gmlx.spec.admission import install_continuous_batch_admission
+    from gmlx.spec.engine import install_owned_spec_engine
+    from gmlx.spec.kv_quant import install_spec_kv_quant
+    from gmlx.spec.mtp_prefill import install_full_prompt_mtp_prefill
+    install_full_prompt_mtp_prefill()
+    install_owned_spec_engine()
+    install_continuous_batch_admission()
+    install_spec_kv_quant()
     from ..batch_sched import install_decode_priority_sched
     install_decode_priority_sched()
     from gmlx.cache.apc_pooling import (
diff --git a/gmlx/spec/admission.py b/gmlx/spec/admission.py
new file mode 100644
index 00000000..755153d4
--- /dev/null
+++ b/gmlx/spec/admission.py
@@ -0,0 +1,429 @@
+"""Continuous-batch admission for the owned speculative engine.
+
+Split out of ``gmlx.spec.engine``.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import mlx.core as mx
+
+from gmlx.envflags import env_bool
+from gmlx.spec.engine import _debug_note
+from gmlx.spec.kv_quant import batch_liftable, lift_single_cache
+
+_log = logging.getLogger(__name__)
+
+
+_CONTINUOUS_BATCH_FLAG = "_kq_gguf_continuous_batch"
+_RELEASED_FLAG = "_kq_gguf_spec_released"
+_RELEASE_PENDING_FLAG = "_kq_gguf_spec_release_pending"
+
+def install_continuous_batch_admission() -> None:
+    """Let new requests prefill and inject during speculative decode.
+
+    Without this, mlx-vlm's ``is_speculative`` gate blocks all prefills while
+    speculative decode is in-flight, and ``extend()`` raises on non-empty
+    speculative batches. This installs five patches:
+
+    1. Disables the ``is_speculative`` admission gate (lets prefills run
+       during decode).
+    2. Overrides ``extend()`` to buffer new batches instead of raising.
+    3. Overrides ``__len__()`` to auto-promote buffered batches when the
+       current batch finishes.
+    4. Overrides ``next()`` to process pending injections - updates outer
+       tracking state, emits first tokens, queues for the generator.
+    5. Releases a finished batch's request state (target KV, captured
+       hidden, shared KV, drafter KV) the moment its last row finishes.
+
+    The generator-side injection (extending caches + drafter mid-flight)
+    happens in ``_owned_decode_rounds_batch`` via ``model._generator_injections``.
+    """
+    from mlx_vlm.generate import ar as _ar
+
+    SpecBatch = _ar.SpeculativeGenerationBatch
+    if getattr(SpecBatch, _CONTINUOUS_BATCH_FLAG, False):
+        return
+
+    # 1. Remove admission gate
+    SpecBatch.is_speculative = False
+
+    _orig_len = SpecBatch.__len__
+
+    # 5. Release request state at finish. BatchGenerator parks the finished
+    # batch in _generation_batch until the next request's prefill completes
+    # (only PromptProcessingBatch.generate's extend replaces it), so every
+    # heavy attr -- the full target KV, the captured full-prompt hidden, the
+    # prefill shared-KV, the rounds generator (whose delegation frame re-pins
+    # all of the above), and the drafter's own head KV -- survives that whole
+    # prefill window. At deep context that stacks two requests' footprints
+    # for many minutes (d200k gemma-4-31b: ~65 GB across an ~18-minute
+    # prefill) and runs the box to the wire ceiling. Drop it all on the
+    # finishing step instead.
+    def _release_heavy_state(self) -> bool:
+        """Drop request state from a finished batch. Returns False when the
+        rounds generator is mid-step on another thread (a client abort racing
+        the engine); ``__len__`` retries on the engine thread."""
+        if getattr(self, _RELEASED_FLAG, False):
+            return True
+        rounds = getattr(self, "_rounds_iter", None)
+        if rounds is not None:
+            try:
+                # Terminal-token finishes already ran the inner loop's own
+                # cleanup; close() is then a no-op resume. Aborted requests
+                # close here, firing the mid-round rollback + retirement.
+                rounds.close()
+            except ValueError:
+                setattr(self, _RELEASE_PENDING_FLAG, True)
+                return False
+            except Exception:
+                _log.warning("spec batch release: rounds close failed", exc_info=True)
+        self._rounds_iter = None
+        self.prompt_cache = []
+        self.hidden = None
+        self.shared_kv_states = None
+        self.prompt_tokens = None
+        self.first_tokens = None
+        if getattr(self, "draft_kind", None) == "mtp":
+            drafter = getattr(self, "draft_model", None)
+            model = getattr(self, "model", None)
+            if drafter is not None and model is not None:
+                try:
+                    drafter.reset(model)  # drops the head's request KV
+                except Exception:
+                    _log.warning(
+                        "spec batch release: drafter reset failed", exc_info=True
+                    )
+        setattr(self, _RELEASED_FLAG, True)
+        setattr(self, _RELEASE_PENDING_FLAG, False)
+        mx.clear_cache()
+        return True
+
+    def _release_if_finished(self) -> None:
+        if _orig_len(self) == 0:
+            _release_heavy_state(self)
+            return
+        _shed_finished_attr_rows(self)
+
+    def _shed_finished_attr_rows(self) -> None:
+        """Per-row release of the batch-held start-time snapshots.
+
+        The live rounds generator sheds a finished or filtered row's KV,
+        drafter state, and its own hidden/shared_kv slices at the next
+        round boundary; the batch object's prefill-time copies (hidden,
+        shared_kv_states, prompt_tokens, first_tokens) stayed resident
+        until the whole batch finished. Slice them by the surviving rows
+        instead. Injected rows carry no snapshot here (their state rides
+        the injection queue into the generator), so the snapshot covers
+        the first first_tokens.shape[0] physical rows only. Slices are
+        lazy and ride the tick's eval; nothing here forces a sync.
+
+        Runs only once the rounds generator holds the state: pre-start,
+        _start_rounds still needs the snapshots row-aligned with the
+        caches (finished rows included; the generator stop_checks them
+        out itself), so a first-token finish must not slice here."""
+        if self._rounds_iter is None:
+            return
+        ft = getattr(self, "first_tokens", None)
+        if ft is None or getattr(self, _RELEASED_FLAG, False):
+            return
+        rows = getattr(self, "_kq_attr_rows", None)
+        if rows is None:
+            try:
+                rows = self._kq_attr_rows = list(range(ft.shape[0]))
+            except Exception:
+                return
+        keep = [p for p in rows
+                if p < len(self._finished) and not self._finished[p]]
+        if len(keep) == len(rows):
+            return
+        if not keep:
+            self.hidden = None
+            self.shared_kv_states = None
+            self.prompt_tokens = None
+            self.first_tokens = None
+            self._kq_attr_rows = []
+            return
+        keep_set = set(keep)
+        pos = [i for i, p in enumerate(rows) if p in keep_set]
+        idx = mx.array(pos, dtype=mx.int32)
+        for name in ("hidden", "prompt_tokens", "first_tokens"):
+            arr = getattr(self, name, None)
+            if arr is not None:
+                setattr(self, name, arr[idx])
+        kv = getattr(self, "shared_kv_states", None)
+        if isinstance(kv, dict) and kv:
+            # New dict, new arrays: the generator may still hold (and
+            # slice) the originals; never mutate a possibly shared dict.
+            self.shared_kv_states = {
+                k: (K[idx], V[idx]) for k, (K, V) in kv.items()}
+        self._kq_attr_rows = keep
+
+    # 2. Buffer extend() instead of raising
+    def _buffered_extend(self, other):
+        active = sum(not d for d in self._finished)
+        if active == 0:
+            pending = getattr(self, "_pending_injections", [])
+            self.__dict__.pop("_kq_attr_rows", None)
+            self.__dict__.update(other.__dict__)
+            self._pending_injections = pending
+            setattr(self, _RELEASED_FLAG, False)
+            setattr(self, _RELEASE_PENDING_FLAG, False)
+            return
+        if not hasattr(self, "_pending_injections"):
+            self._pending_injections = []
+        self._pending_injections.append(other)
+        _debug_note(f"[mtp] extend buffered: +{len(other._all_uids)} rows "
+                    f"(pending={len(self._pending_injections)}, "
+                    f"active={active})")
+
+    SpecBatch.extend = _buffered_extend
+
+    # 3. Auto-promote buffered batches when current is done
+    def _len_with_promotion(self):
+        if getattr(self, _RELEASE_PENDING_FLAG, False) and _orig_len(self) == 0:
+            _release_heavy_state(self)
+        active = _orig_len(self)
+        if active == 0:
+            pending = getattr(self, "_pending_injections", None)
+            if pending:
+                other = pending.pop(0)
+                remaining = pending[:]
+                self.__dict__.pop("_kq_attr_rows", None)
+                self.__dict__.update(other.__dict__)
+                self._pending_injections = remaining
+                setattr(self, _RELEASED_FLAG, False)
+                setattr(self, _RELEASE_PENDING_FLAG, False)
+                return _orig_len(self)
+        return active
+
+    SpecBatch.__len__ = _len_with_promotion
+
+    _orig_filter = SpecBatch.filter
+
+    def _compact_prestart_rows(self, keep) -> None:
+        """Physically drop rows from a batch whose rounds generator has
+        not started: filter the caches through their own filter (lifting
+        host caches first) and slice snapshots plus bookkeeping to the
+        same keep list. Pre-start, the batch object owns all state, so
+        the drop frees the rows' bytes immediately instead of marking
+        them finished and waiting for a generator that has no round
+        boundary yet."""
+        idx = mx.array(keep, dtype=mx.int32)
+        self.prompt_cache = [_lift_host_cache(c) for c in self.prompt_cache]
+        for c in self.prompt_cache:
+            c.filter(idx)
+        for name in ("hidden", "prompt_tokens", "first_tokens"):
+            arr = getattr(self, name, None)
+            if arr is not None:
+                setattr(self, name, arr[idx])
+        kv = getattr(self, "shared_kv_states", None)
+        if isinstance(kv, dict) and kv:
+            self.shared_kv_states = {
+                k: (K[idx], V[idx]) for k, (K, V) in kv.items()}
+        self._all_uids = [self._all_uids[i] for i in keep]
+        self.uids = list(self._all_uids)
+        self.max_tokens = [self.max_tokens[i] for i in keep]
+        self._num_tokens = [self._num_tokens[i] for i in keep]
+        self._finished = [False] * len(keep)
+        self.__dict__.pop("_kq_attr_rows", None)
+
+    def _filter_with_release(self, keep):
+        # Pre-start strict subset (a cancel or a governor retire landing
+        # before the first tick): compact physically. Live or degenerate
+        # cases keep the upstream mark-finished contract; the running
+        # generator sheds the row at its next round boundary and the
+        # snapshot shed below covers the batch-held copies.
+        if (len(keep) < len(self.uids)
+                and keep
+                and self._rounds_iter is None
+                and not getattr(self, _RELEASED_FLAG, False)
+                and getattr(self, "first_tokens", None) is not None
+                and self.uids == self._all_uids
+                and not any(self._finished)
+                and all(batch_liftable(c) for c in self.prompt_cache)):
+            _compact_prestart_rows(self, list(keep))
+            return
+        _orig_filter(self, keep)
+        _release_if_finished(self)
+
+    SpecBatch.filter = _filter_with_release
+
+    # 4. Process pending injections in next() before advancing the generator
+    _orig_next = SpecBatch.next
+
+    def _note_last_tokens(self, responses) -> None:
+        # Last delivered token per uid: the bonus a preempt rebuild restarts
+        # from (its KV is not yet in the cache at a round boundary).
+        stash = getattr(self, "_kq_last_tokens", None)
+        if stash is None:
+            stash = self._kq_last_tokens = {}
+        for r in responses:
+            if r.token is not None:
+                stash[r.uid] = int(r.token)
+
+    def _lift_host_cache(c):
+        """Promote a single-sequence host cache to its batch class so the
+        rebuilt batch generator can extend/filter it (same lift the
+        injection path applies to incoming caches)."""
+        if hasattr(c, "filter") and hasattr(c, "extend"):
+            return c
+        return lift_single_cache(c)
+
+    def _preempt_scalar(self) -> bool:
+        """Preempt a live scalar (B=1) spec generation so queued rows can
+        join: close the generator, deliver the closed round's undelivered
+        tail (the scalar path yields one token per next(), so a close
+        usually lands mid-round; those tokens are verified and their KV
+        stays in the cache), lift the caches to batch classes, and mark
+        the batch armless (hidden=None); _start_rounds then rebuilds it on
+        the batch loop, whose first injection drain admits the waiters.
+        The rebuild resumes from the round's bonus token, whose KV is not
+        in the cache. GMLX_MTP_PREEMPT=0 leaves the old drain-wait
+        behavior.
+
+        The rebuilt row carries no APC retirement context (batch-loop rows
+        start with retire_ctxs None), so the preempted request's prefix is
+        not offered back to the prompt cache when it finishes."""
+        if not env_bool("GMLX_MTP_PREEMPT", True):
+            return False
+        if not getattr(self, "_sent_first", False):
+            return False
+        last = getattr(self, "_kq_last_tokens", {}).get(self._all_uids[0])
+        if last is None:
+            return False
+        # Every cache must be batch-liftable before the generator
+        # closes. A quantized or kvarn B=1 cache lifts to fp16. Anything
+        # else unliftable declines into the drain-wait.
+        if not all(batch_liftable(c) for c in self.prompt_cache):
+            return False
+        it = self._rounds_iter
+        captured = []
+        if it is not None:
+            self._rounds_iter = None
+            self.model._kq_preempt_capture = captured
+            try:
+                it.close()
+            finally:
+                try:
+                    del self.model._kq_preempt_capture
+                except AttributeError:
+                    pass
+        responses = []
+        uid = self._all_uids[0]
+        for tok in captured:
+            if self._finished[0]:
+                break
+            tok = int(tok)
+            self._num_tokens[0] += 1
+            finish = self._finish_reason(0, tok)
+            if finish is not None:
+                self._finished[0] = True
+            responses.append(self.Response(
+                uid=uid, token=tok, token_logprob=0.0, finish_reason=finish))
+            last = tok
+        self._kq_preempt_responses = responses
+        if self._finished[0]:
+            # The captured tail finished the row; nothing to rebuild. The
+            # pending injections promote through __len__ once drained.
+            self._refresh_uids()
+            return False
+        self.prompt_cache = [_lift_host_cache(c) for c in self.prompt_cache]
+        self.first_tokens = mx.array([int(last)], dtype=self.token_dtype)
+        self.hidden = None
+        self.shared_kv_states = None
+        self.prompt_tokens = None
+        self.model._kq_rebuild_emitted = [int(self._num_tokens[0])]
+        _debug_note("[mtp] preempt: scalar generation rebuilt for "
+                    "continuous batching")
+        return True
+
+    def _next_with_injection(self):
+        pending = getattr(self, "_pending_injections", None)
+        # Physical-row uids for the owned rounds loop (it has no batch
+        # object): read once at generator start, injected rows carry theirs.
+        try:
+            self.model._kq_row_uids = list(self._all_uids)
+        except AttributeError:      # attribute-less model stand-ins
+            pass
+        # Mid-flight adoption works only when the batch rounds generator is
+        # running: it drains model._generator_injections at its round
+        # boundaries. The scalar (B=1) generator never does, so a live
+        # scalar host is preempted first: its generator closes at the round
+        # boundary and the batch is rebuilt armless on the batch loop.
+        # `_all_uids` is an mlx-vlm generator internal (stable under the
+        # ==0.6.3 pin); re-verify this batch-vs-scalar signal on a pin lift.
+        preempted = False
+        if pending and len(self._all_uids) == 1:
+            preempted = _preempt_scalar(self)
+        # The preempt capture: verified tokens the closed round had not yet
+        # delivered. They precede everything this call returns.
+        pre_responses = self.__dict__.pop("_kq_preempt_responses", None) or []
+        if pending and (len(self._all_uids) > 1 or preempted):
+            responses = list(pre_responses)
+            gen_inj = getattr(self.model, "_generator_injections", None)
+            if gen_inj is None:
+                self.model._generator_injections = []
+                gen_inj = self.model._generator_injections
+
+            for other in pending:
+                B_new = len(other._all_uids)
+                base_row = len(self._all_uids)
+                self._all_uids.extend(other._all_uids)
+                self._num_tokens.extend([0] * B_new)
+                self._finished.extend([False] * B_new)
+                self.max_tokens.extend(other.max_tokens)
+
+                mx.eval(other.first_tokens)
+                first_list = other.first_tokens.tolist()
+                for row in range(B_new):
+                    abs_row = base_row + row
+                    tok = int(first_list[row])
+                    self._num_tokens[abs_row] = 1
+                    finish = self._finish_reason(abs_row, tok)
+                    if finish is not None:
+                        self._finished[abs_row] = True
+                    responses.append(
+                        self.Response(
+                            uid=other._all_uids[row],
+                            token=tok,
+                            token_logprob=0.0,
+                            finish_reason=finish,
+                        )
+                    )
+
+                gen_inj.append(
+                    {
+                        "uids": list(other._all_uids),
+                        "prompt_cache": other.prompt_cache,
+                        "hidden": other.hidden,
+                        "shared_kv_states": other.shared_kv_states,
+                        "prompt_tokens": other.prompt_tokens,
+                        "first_tokens": other.first_tokens,
+                        "first_tokens_list": first_list,
+                        # The running generator froze max(max_tokens) at
+                        # start; injected rows carry their own budgets.
+                        "max_tokens": list(other.max_tokens),
+                    }
+                )
+
+            pending.clear()
+            self._refresh_uids()
+
+            more = _orig_next(self)
+            responses.extend(more)
+            _note_last_tokens(self, responses)
+            _release_if_finished(self)
+            return responses
+
+        responses = pre_responses + _orig_next(self)
+        _note_last_tokens(self, responses)
+        _release_if_finished(self)
+        return responses
+
+    SpecBatch.next = _next_with_injection
+    setattr(SpecBatch, _CONTINUOUS_BATCH_FLAG, True)
+    _debug_note(
+        "[mtp] continuous batch: admission gate removed, mid-flight injection enabled"
+    )
diff --git a/gmlx/spec/ckpt.py b/gmlx/spec/ckpt.py
new file mode 100644
index 00000000..f54aa984
--- /dev/null
+++ b/gmlx/spec/ckpt.py
@@ -0,0 +1,991 @@
+"""Checkpoint-tier arming, exact anchors, and the plain ckpt decode path.
+
+Split out of ``gmlx.spec.engine``; the engine keeps the L1 view, layout
+probes, and env flags these helpers read.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+
+from gmlx.envflags import env_int
+from gmlx.spec.engine import (
+    _L1View,
+    _SPEC_APC_DISABLED,
+    _SPEC_APC_RETIRE_DISABLED,
+    _SPEC_APC_SIDECAR_DISABLED,
+    _ckpt_active,
+    _ckpt_layout_for,
+    _ckpt_layout_live,
+    _live_kv_quant_config,
+)
+
+_log = logging.getLogger(__name__)
+
+
+def _l1_lookup_and_arm_store(batch, manager, mode, l0_prefix) -> int:
+    """Consult the shared APCManager below L0 and arm the stock post-prefill
+    store (mid-prefill exact checkpoints + post-prefill exact store / block
+    harvest, all owned by stock ``PromptProcessingBatch.generate``) by
+    populating ``_apc_manager`` / ``_apc_mode`` / ``_apc_meta``.
+
+    Returns the restored L1 prefix length (0 on miss, or when L0 already
+    restored -- L0 carries full-prompt hidden and is always preferred).
+
+    ``meta["prefix_len"]`` stays 0 by design: the owned prefill keeps
+    ``_processed_prompt_columns`` in absolute token space (it trims
+    ``_input_ids`` in place, unlike stock warm batches which are constructed
+    with suffix-only rows), and ``_row_real_tokens_processed`` -- which gates
+    the mid-prefill checkpoint store -- is only correct in that space with a
+    zero meta prefix. The one cost is that a block-tier harvest re-walks
+    restored prefix blocks, but ``store_kv_blocks`` dedups by hash chain
+    (acquire+release of existing blocks, no data copies).
+    """
+    view = _L1View(batch.model, manager, mode)
+    ids_list = [int(t) for t in batch._mtp_full_input_ids[0].tolist()]
+    prompt_kwargs = batch._prompt_kwargs or {}
+    extra_hash = view._apc_extra_hash(prompt_kwargs)
+    ckpt = _ckpt_active(batch.model, mode, int(manager.block_size))
+    held_blocks = []
+    l1_prefix = 0
+    if l0_prefix == 0 and len(ids_list) >= 2:
+        warm = None
+        blocks = []
+        prefix_len = 0
+        tier = "exact"
+        pick = view._apc_pick_for((0, ids_list, 0, prompt_kwargs, None, None))
+        # Same trivial-pick floor as the admission wrapper: a sub-block
+        # exact restore saves nothing and its nonzero l1_prefix would skip
+        # the L0 hidden store for this request.
+        if (pick is not None and not pick.get("matched_blocks")
+                and 0 < int(pick.get("prefix_len") or 0)
+                < int(manager.block_size)):
+            pick = None
+        if pick is not None:
+            warm = pick.get("warm_cache")
+            blocks = list(pick.get("matched_blocks") or ())
+            prefix_len = int(pick.get("prefix_len") or 0)
+            extra_hash = int(pick.get("extra_hash", extra_hash))
+            if warm is None and blocks:
+                from mlx_vlm import apc as _apc
+
+                warm = _apc.make_warm_kv_cache(
+                    blocks, min_capacity_tokens=len(ids_list) + 1
+                )
+                tier = "block"
+        if ckpt:
+            # Checkpoint tier: the longest salted sidecar + block chain
+            # wins only when strictly longer than the exact-tier pick.
+            # Media guards mirror the stock exact probe.
+            from gmlx.cache.snapshot import ckpt_lookup
+            min_p = max(prefix_len,
+                        view._apc_safe_prefix_lookup_min(ids_list))
+            cw, cp = ckpt_lookup(
+                manager,
+                ids_list,
+                extra_hash=extra_hash,
+                min_prefix_tokens=min_p,
+                layout=_ckpt_layout_live(batch, int(manager.block_size)),
+            )
+            if (
+                cw is not None
+                and cp > prefix_len
+                and view._apc_suffix_is_text_only(ids_list, cp)
+            ):
+                if blocks:
+                    manager.release(blocks)
+                    blocks = []
+                warm, prefix_len, tier = cw, cp, "ckpt"
+        elif mode == "exact":
+            # Exact-tier anchor: the shared-system-prefix clone in the
+            # gmlx anchor LRU wins only when strictly longer than the
+            # stock exact pick. Media guards mirror the stock probe.
+            from gmlx.cache.snapshot import anchor_exact_lookup
+            min_p = max(prefix_len,
+                        view._apc_safe_prefix_lookup_min(ids_list))
+            aw, ap = anchor_exact_lookup(
+                manager, ids_list, extra_hash=extra_hash,
+                min_prefix_tokens=min_p)
+            if (aw is not None and ap > prefix_len
+                    and view._apc_suffix_is_text_only(ids_list, ap)):
+                if blocks:
+                    manager.release(blocks)
+                    blocks = []
+                warm, prefix_len, tier = aw, ap, "anchor"
+        if warm and 0 < prefix_len < len(ids_list) and tier in (
+                "exact", "anchor"):
+            # Same batch-aware merge admission applies to its picks: raw
+            # exact/anchor clones carry single-row leaves (left_padding
+            # None, scalar offsets) and crash the batch cache classes'
+            # update path (mx.depends on a None) when the suffix forwards.
+            # kv_quant_config re-quantizes the float snapshot to the live
+            # _make_cache layer types under serve kv_bits (stored exact
+            # entries stay float; a float row joining a quantized batch
+            # breaks the update path).
+            from mlx_vlm import apc as _apc
+            warm, _ = _apc.make_warm_batch_exact_cache_multi(
+                [warm], prefix_lens=[prefix_len],
+                kv_quant_config=_live_kv_quant_config(batch.model))
+        if warm and 0 < prefix_len < len(ids_list):
+            batch.prompt_cache = warm
+            # Matched blocks stay acquired until the stock post-prefill
+            # harvest releases them (the warm-cache concatenation is
+            # lazy; the pool must not recycle these blocks before it
+            # materializes).
+            held_blocks = blocks
+            l1_prefix = prefix_len
+            # Observability only: the live request view reads the
+            # restored prefix from here (meta keeps prefix_len 0 so the
+            # stock machinery does not account it twice).
+            batch._kq_apc_restored = (int(prefix_len), str(tier))
+            _log.info(
+                "APC L1 hit: prefix=%d suffix=%d tier=%s",
+                prefix_len,
+                len(ids_list) - prefix_len,
+                tier,
+            )
+            # Drafter-KV sidecar: a plain L1 hit restores target KV but
+            # not hidden, so the drafter would re-seed from suffix-only
+            # hidden at the wrong positions (acceptance erodes at depth).
+            # A sidecar covering exactly the restored prefix hands the
+            # owned round a warm drafter start. Stash rides the first
+            # cache entry, same discipline as the retirement context.
+            if not _SPEC_APC_SIDECAR_DISABLED:
+                from gmlx.cache.snapshot import drafter_sidecar_lookup
+                side = drafter_sidecar_lookup(
+                    manager, ids_list, prefix_len, extra_hash)
+                if side:
+                    batch.prompt_cache[0]._kq_apc_drafter_warm = side
+                    _log.info("APC sidecar hit: prefix=%d", prefix_len)
+        elif blocks:
+            manager.release(blocks)
+    batch._mtp_l1_prefix_len = l1_prefix
+    batch._apc_manager = manager
+    batch._apc_mode = mode
+    guard = int(view._apc_exact_checkpoint_len(ids_list) or 0)
+    meta = {
+        "full_input_ids": ids_list,
+        "prefix_len": 0,
+        "extra_hash": extra_hash,
+        "apc_blocks": held_blocks,
+        "checkpoint_len": guard,
+    }
+    batch._apc_meta = [meta]
+    if ckpt:
+        # The checkpoint tier replaces the stock exact-tier stores: the
+        # post-prefill full-cache clone is suppressed here, and the
+        # mid-prefill checkpoint store is superseded by the cursor riding
+        # the wrapped stock store (_install_ckpt_checkpoint_store; the
+        # stock body is suppressed by the cursor's advance). Column
+        # alignment itself still runs on the stock machinery, which
+        # requires _apc_mode == "exact".
+        _ckpt_arm_schedule(batch, meta, guard,
+                           max(l0_prefix, l1_prefix),
+                           int(manager.block_size))
+        batch._apc_harvest_enabled = False
+        batch._kq_ckpt_armed = True
+        from gmlx.cache.snapshot import ckpt_note_armed
+        ckpt_note_armed(manager)
+    elif mode == "exact":
+        _exact_anchor_arm(batch, meta, guard,
+                          max(l0_prefix, l1_prefix))
+    return l1_prefix
+
+
+def _gcd(a: int, b: int) -> int:
+    while b:
+        a, b = b, a % b
+    return a
+
+
+def _ckpt_unit(batch, block_size: int) -> int:
+    """The natural chunk grid: lcm(prefill_step_size, block_size)."""
+    step = int(getattr(batch, "prefill_step_size", 0) or 0)
+    return block_size if step <= 0 else \
+        step * block_size // _gcd(step, block_size)
+
+
+def _ckpt_cursor_init(batch, guard: int, restored: int,
+                      block_size: int) -> tuple[list, int, int]:
+    """Boundary schedule for the checkpoint cursor: an ordered
+    ``[(position, kind), ...]`` list plus ``(terminal, interval)``.
+
+    Boundaries sit on the natural chunk grid, lcm(prefill_step_size,
+    block_size) -- an off-grid boundary truncates a chunk, and gated-delta
+    state is chunk-shape sensitive (certified: any grid change drifts).
+    Interval points above the restored prefix, then the terminal (the
+    grid point at or below the stock guard column). GMLX_APC_CKPT_INTERVAL
+    tokens, default 4096, snapped up to the grid; 0 = terminal-only.
+    Later stages add store positions by appending boundaries here, never
+    by new store mechanisms.
+    """
+    unit = _ckpt_unit(batch, block_size)
+    terminal = (guard // unit) * unit
+    if terminal <= max(0, restored):
+        return [], 0, 0
+    raw = env_int("GMLX_APC_CKPT_INTERVAL", 4096)
+    interval = 0 if raw <= 0 else max(unit, (raw // unit) * unit)
+    bounds = []
+    if interval:
+        b = ((max(0, restored) // interval) + 1) * interval
+        while b < terminal:
+            bounds.append((b, "boundary"))
+            b += interval
+    bounds.append((terminal, "boundary"))
+    return bounds, terminal, interval
+
+
+def _ckpt_replay_boundary(batch, meta, restored: int,
+                          block_size: int) -> int | None:
+    """N-1 replay boundary, or None when it cannot earn its pause.
+
+    An identical resend can only adopt a record strictly below the
+    query, and the interval/terminal schedule never places one there
+    for prompts under one interval (the depth e2e's bug 1); N-1 is the
+    deepest position that is both adoptable and drift-free (the warm
+    turn forwards exactly one token), and the pause is free on the cold
+    side -- both prefill loops already stop at N-1 to feed the first
+    decode step, so the boundary lands on a natural chunk edge and
+    perturbs no chunk shape. arr layouts gate on a minimum N:
+    recurrent state is prompt-length-independent (>100 MB per record on
+    27B-class), and short-prompt records would churn the LRU out of the
+    deep-conversation records it exists to protect. Rotating layouts
+    need N-1 at or past the window (below it the store's grid gate
+    declines). Kill switch: GMLX_APC_CKPT_REPLAY=0.
+    """
+    if env_int("GMLX_APC_CKPT_REPLAY", 1) == 0:
+        return None
+    n = len(meta.get("full_input_ids") or ())
+    replay = n - 1
+    if replay < 2 or replay <= max(0, restored):
+        return None
+    tags = _ckpt_layout_live(batch, block_size) or ()
+    if "arr" in tags and n < env_int("GMLX_APC_CKPT_REPLAY_MIN", 1024):
+        return None
+    for t in tags:
+        if t.startswith("rot") and replay < int(t.split(":")[1]):
+            return None
+    return replay
+
+
+def _ckpt_turn_boundaries(batch, meta, restored: int,
+                          block_size: int) -> list[int]:
+    """Render-stable turn boundary positions for the schedule.
+
+    p_stable is the deepest prompt position a next-turn re-render keeps;
+    the gen-prompt/think tail past it is re-rendered away, so records
+    stored only above it can never serve turn 2 (how multi-turn adoption
+    silently died). Every layout gets the grid point at or below
+    p_stable (drift-free for chunk-shape-sensitive state); rot-only
+    layouts also pause exactly at p_stable (attention splits exactly;
+    needs the window wrapped). GMLX_APC_CKPT_TURN=0 disables these and
+    with them the p=N drop gate.
+    """
+    if env_int("GMLX_APC_CKPT_TURN", 1) == 0:
+        return []
+    ids = meta.get("full_input_ids") or ()
+    unit = _ckpt_unit(batch, block_size)
+    tags = _ckpt_layout_live(batch, block_size) or ()
+    ws = [int(t.split(":")[1]) for t in tags if t.startswith("rot")]
+    # Cheapest boundary this layout could arm: the grid needs one unit
+    # of stable prefix; rot-only layouts can also pause exactly at
+    # p_stable once the window wraps. Below that no boundary can land,
+    # so skip the render+tokenize prediction entirely.
+    need = unit if ("arr" in tags or not ws) else min(unit, max(ws))
+    if len(ids) - 1 < need:
+        return []
+    from gmlx.cache.retire_key import lookup_render_ctx, prompt_stable_lcp
+    ctx = lookup_render_ctx(ids)
+    p_stable = prompt_stable_lcp(ctx, ids) if ctx else None
+    if not p_stable or p_stable < 2:
+        return []
+    p_stable = min(int(p_stable), len(ids) - 1)
+    meta["ckpt_p_stable"] = p_stable
+    floor = max(0, restored)
+    out = []
+    grid = (p_stable // unit) * unit
+    if grid > floor:
+        out.append(grid)
+    if ws and "arr" not in tags and p_stable != grid \
+            and p_stable > floor and p_stable >= max(ws):
+        out.append(p_stable)
+    return out
+
+
+def _ckpt_sys_boundary(batch, meta, restored: int,
+                       block_size: int) -> int | None:
+    """Anchor stop at the end of the shared system prefix.
+
+    Sibling fan-out requests share the system prompt and tool schemas
+    and diverge at the first user message, generally between grid
+    points, so the interval schedule alone wastes up to one interval of
+    sibling recompute, and strip-on-extend removes the early boundary
+    the siblings need as the chain deepens (the anchor exemption in
+    _record_insert keeps this one). arr layouts snap the stop down to
+    the chunk grid (off-grid chunking drifts GDN state) and keep the
+    replay byte floor (recurrent state is prompt-length-independent, so
+    a tiny anchor costs the same >100 MB clone as a deep one);
+    attention layouts snap to the block grid, which also satisfies the
+    rotating store's below-window grid gate. GMLX_APC_CKPT_SYS=0
+    disables; GMLX_APC_CKPT_SYS_MIN floors the position (a sub-floor
+    shared prefix re-prefills in milliseconds and is not worth a
+    record).
+    """
+    if env_int("GMLX_APC_CKPT_SYS", 1) == 0:
+        return None
+    ids = meta.get("full_input_ids") or ()
+    tags = _ckpt_layout_for(getattr(batch, "model", None), block_size) or ()
+    floor_min = max(block_size, env_int("GMLX_APC_CKPT_SYS_MIN", 256))
+    if "arr" in tags:
+        floor_min = max(floor_min,
+                        env_int("GMLX_APC_CKPT_REPLAY_MIN", 1024))
+    # Below the floor no anchor can land; skip the render+tokenize
+    # prediction entirely (same rule as the turn boundaries).
+    if len(ids) - 1 < floor_min:
+        return None
+    from gmlx.cache.retire_key import lookup_render_ctx, system_prefix_lcp
+    ctx = lookup_render_ctx(ids)
+    lcp = system_prefix_lcp(ctx, ids) if ctx else None
+    if not lcp:
+        return None
+    unit = _ckpt_unit(batch, block_size) if "arr" in tags else block_size
+    pos = (min(int(lcp), len(ids) - 1) // unit) * unit
+    if pos < floor_min or pos <= max(0, restored):
+        return None
+    meta["ckpt_sys_bound"] = pos
+    return pos
+
+
+def _exact_anchor_boundary(batch, meta, guard: int,
+                           restored: int) -> int | None:
+    """Anchor position for an exact-tier (non-ckpt) model: the sibling
+    divergence point, ungridded (exact clones restore at any position).
+    Clamped to the stock guard column, so the prefill pauses at most
+    twice: once for the anchor, once for the stock guard store.
+    GMLX_APC_CKPT_SYS=0 disables (one switch for both tiers);
+    GMLX_APC_CKPT_SYS_MIN floors the position (a sub-floor shared
+    prefix re-prefills in milliseconds and is not worth a clone).
+    """
+    if env_int("GMLX_APC_CKPT_SYS", 1) == 0:
+        return None
+    ids = meta.get("full_input_ids") or ()
+    floor_min = max(2, env_int("GMLX_APC_CKPT_SYS_MIN", 256))
+    if len(ids) - 1 < floor_min:
+        return None
+    from gmlx.cache.retire_key import lookup_render_ctx, system_prefix_lcp
+    ctx = lookup_render_ctx(ids)
+    lcp = system_prefix_lcp(ctx, ids) if ctx else None
+    if not lcp:
+        _log.info("APC anchor declined: no measurable system prefix "
+                  "(render ctx %s)", "present" if ctx else "missing")
+        return None
+    pos = min(int(lcp), len(ids) - 1)
+    if guard > 0:
+        pos = min(pos, guard)
+    if pos < floor_min or pos <= max(0, restored):
+        return None
+    return pos
+
+
+def _exact_anchor_arm(batch, meta, guard: int, restored: int) -> None:
+    """Schedule the anchor pause by mirroring its position into
+    ``checkpoint_len`` (the key the stock column truncation reads).
+    ``_exact_anchor_store`` hands the column back to the stock guard
+    after firing, so the stock store still runs exactly as unarmed."""
+    pos = _exact_anchor_boundary(batch, meta, guard, restored)
+    if pos is None:
+        return
+    meta["anchor_len"] = pos
+    meta["anchor_guard"] = guard
+    if pos != guard:
+        meta["checkpoint_len"] = pos
+    batch._kq_anchor_armed = True
+    _log.info("APC anchor armed: pos=%d guard=%d", pos, guard)
+
+
+def _exact_anchor_store(batch) -> None:
+    """Anchor store for exact-tier models: one whole-prefix clone at the
+    sibling divergence, into the gmlx anchor LRU. Runs from the wrapped
+    stock store immediately before the stock body; after firing it
+    restores ``checkpoint_len`` to the stock guard column without
+    latching ``checkpoint_done``, so the stock guard store (and its
+    latch) fire untouched."""
+    manager = getattr(batch, "_apc_manager", None)
+    meta_list = getattr(batch, "_apc_meta", None) or []
+    if manager is None or not meta_list or meta_list[0] is None:
+        return
+    meta = meta_list[0]
+    pos = int(meta.get("anchor_len") or 0)
+    if pos <= 0 or meta.get("anchor_done"):
+        return
+    if batch._row_real_tokens_processed(0) != pos:
+        return
+    meta["anchor_done"] = True
+    guard = int(meta.get("anchor_guard") or 0)
+    if int(meta.get("checkpoint_len") or 0) == pos and pos != guard:
+        meta["checkpoint_len"] = guard
+    cache = batch._apc_prompt_cache_for_store(0)
+    if cache is None:
+        return
+    from gmlx.cache.snapshot import anchor_exact_store
+    anchor_exact_store(manager, meta["full_input_ids"][:pos], cache,
+                       extra_hash=int(meta.get("extra_hash", 0)))
+
+
+def _sched_insert(bounds: list, pos: int, kind: str, *,
+                  upgrade: bool = False) -> None:
+    """Insert (pos, kind) keeping order. On collision the existing entry
+    keeps its kind: a colliding position is always grid-aligned or an
+    exact turn boundary, where a plain boundary record adopts freely --
+    identical resend included -- while flipping it to replay would gate
+    turn-2 and branch adoption out on recurrent layouts (and satisfy the
+    p=N drop with a record turn 2 cannot use). ``upgrade`` lets an
+    anchor replace a plain boundary at the same position (strictly more
+    retention, same free adoption), never a replay."""
+    import bisect
+
+    pts = [b for b, _ in bounds]
+    i = bisect.bisect_left(pts, pos)
+    if i < len(pts) and pts[i] == pos:
+        if upgrade and bounds[i][1] == "boundary":
+            bounds[i] = (pos, kind)
+        return
+    bounds.insert(i, (pos, kind))
+
+
+def _ckpt_arm_schedule(batch, meta, guard: int, restored: int,
+                       block_size: int) -> None:
+    """Publish the boundary schedule into the request meta. The head
+    mirrors into ``checkpoint_len`` (an int) because the stock
+    checkpoint-column truncation and store reads exactly that key.
+    ``ckpt_stored_boundaries`` collects every boundary whose store
+    landed (record verified in the index) -- the settled variable the
+    post-prefill p=N decision and the sidecar key set both read;
+    ``ckpt_p_stable_bounds`` is the qualifying set for the p=N drop."""
+    bounds, terminal, interval = _ckpt_cursor_init(
+        batch, guard, restored, block_size)
+    turn = _ckpt_turn_boundaries(batch, meta, restored, block_size)
+    for pos in turn:
+        _sched_insert(bounds, pos, "boundary")
+    sysb = _ckpt_sys_boundary(batch, meta, restored, block_size)
+    if sysb is not None:
+        _sched_insert(bounds, sysb, "anchor", upgrade=True)
+    replay = _ckpt_replay_boundary(batch, meta, restored, block_size)
+    if replay is not None:
+        # Colliding with the anchor keeps the anchor (default no-upgrade):
+        # it adopts identical resends freely, replay semantics add nothing.
+        _sched_insert(bounds, replay, "replay")
+    meta["ckpt_boundaries"] = bounds
+    meta["checkpoint_len"] = int(bounds[0][0]) if bounds else 0
+    meta["ckpt_terminal"] = terminal
+    meta["ckpt_interval"] = interval
+    meta["ckpt_last_stored"] = 0
+    meta["ckpt_stored_boundaries"] = []
+    meta["ckpt_p_stable_bounds"] = turn
+
+
+def _ckpt_mid_prefill_store(batch) -> None:
+    """Checkpoint-tier replacement for the stock mid-prefill exact store.
+
+    Fires at the schedule head, pops it, and mirrors the next head into
+    ``checkpoint_len``, latching ``checkpoint_done`` when the schedule
+    empties. The advance is what suppresses the stock store;
+    ``_install_ckpt_checkpoint_store`` wraps the stock method so the
+    cursor always runs immediately before it -- the ordering is
+    structural, not positional. Advances past failed stores;
+    ``ckpt_last_stored`` records only boundaries that landed.
+    """
+    if not getattr(batch, "_kq_ckpt_armed", False):
+        return
+    manager = getattr(batch, "_apc_manager", None)
+    meta_list = getattr(batch, "_apc_meta", None) or []
+    if manager is None or not meta_list or meta_list[0] is None:
+        return
+    meta = meta_list[0]
+    if meta.get("checkpoint_done"):
+        return
+    checkpoint_len = int(meta.get("checkpoint_len") or 0)
+    if checkpoint_len <= 0:
+        return
+    if batch._row_real_tokens_processed(0) != checkpoint_len:
+        return
+    terminal = int(meta.get("ckpt_terminal") or 0)
+    bounds = meta.get("ckpt_boundaries") or []
+    kind = "boundary"
+    if bounds and int(bounds[0][0]) == checkpoint_len:
+        kind = str(bounds.pop(0)[1])
+    # Inline-heavy skeletons (GDN state >100 MB; kvarn state scales with p
+    # across every attention layer) earn disk only at the terminal --
+    # boundaries superseded within the same prefill do not, and a replay
+    # skeleton would buy restart-repair of an identical resend only,
+    # which does not earn it either.
+    layout = _ckpt_layout_live(batch, int(manager.block_size)) or ()
+    heavy = "arr" in layout or any(t.startswith("kvarn") for t in layout)
+    skel = not heavy or (kind != "replay"
+                         and checkpoint_len >= terminal)
+    from gmlx.cache.snapshot import ckpt_store
+
+    if ckpt_store(
+            manager, meta["full_input_ids"][:checkpoint_len],
+            batch.prompt_cache, extra_hash=int(meta.get("extra_hash", 0)),
+            skeleton_disk=skel, kind=kind):
+        meta["ckpt_last_stored"] = checkpoint_len
+        meta.setdefault("ckpt_stored_boundaries", []).append(checkpoint_len)
+    if bounds:
+        meta["checkpoint_len"] = int(bounds[0][0])
+    else:
+        meta["checkpoint_done"] = True
+
+
+_CKPT_STORE_FLAG = "_kq_ckpt_cursor_store"
+
+
+def _install_ckpt_checkpoint_store() -> None:
+    """Wrap the stock mid-prefill checkpoint store so the cursor runs
+    immediately before it on armed batches (both the owned MTP prefill
+    and the stock prompt_step call the stock method, so one wrap covers
+    both paths). The cursor's advance of ``checkpoint_len`` is what
+    suppresses the stock store -- wrapping makes that ordering
+    structural. Exact-tier anchor batches ride the same wrap with their
+    own single-stop hook. Idempotent."""
+    from mlx_vlm.generate.ar import PromptProcessingBatch
+
+    if getattr(
+        PromptProcessingBatch._store_apc_exact_checkpoints, _CKPT_STORE_FLAG, False
+    ):
+        return
+    _orig = PromptProcessingBatch._store_apc_exact_checkpoints
+
+    def _store_with_ckpt_cursor(self):
+        if getattr(self, "_kq_ckpt_armed", False):
+            _ckpt_mid_prefill_store(self)
+        elif getattr(self, "_kq_anchor_armed", False):
+            _exact_anchor_store(self)
+        _orig(self)
+
+    _store_with_ckpt_cursor.__dict__[_CKPT_STORE_FLAG] = True
+    PromptProcessingBatch._store_apc_exact_checkpoints = _store_with_ckpt_cursor
+
+
+def _snap_fields(batch, manager) -> dict:
+    """Decode-time snapshot ring parameters for a retirement stash.
+
+    ``snap_grid`` anchors snapshot positions to the prefill chunk grid
+    (lcm of step and block size), so a restore replays chunk-exact --
+    but only while one grid unit fits inside the snapshot interval; a
+    serve-sized step (2048) would otherwise push the first snapshot far
+    past prompt end + interval, so it falls back to the block size (the
+    off-grid restore is the scoped-benign case). ``snap_align`` is the
+    block alignment a rotating window store requires below the window;
+    ``snap_offgrid_min`` (= W) is where the store gate stops caring --
+    a wrapped window is whole blocks at any p.
+    """
+    import math
+    from gmlx.cache.snapshot import _DECODE_CKPT_DEFAULT
+    bs = int(manager.block_size)
+    tags = _ckpt_layout_live(batch, bs) or ()
+    step = int(getattr(batch, "prefill_step_size", 0) or 0)
+    grid = math.lcm(step, bs) if step > 0 else bs
+    if grid > env_int("GMLX_APC_DECODE_CKPT", _DECODE_CKPT_DEFAULT):
+        grid = bs
+    rot_w = 0
+    for t in tags:
+        if t.startswith("rot"):
+            rot_w = int(t.split(":")[1])
+            break
+    return {
+        "snap_ok": bool(tags),
+        "snap_grid": grid,
+        "snap_align": bs if rot_w else 1,
+        "snap_offgrid_min": rot_w,
+    }
+
+
+def _plain_ckpt_init(batch) -> None:
+    """Checkpoint-tier lookup + arming for a stock (non-speculative)
+    prompt batch.
+
+    The stock path reaches the tier only here: exact-tier stores are
+    suppressed on ckpt models, so admission's own lookup ladder misses
+    and every ckpt-tier request arrives as a cold single-request batch.
+    Lookup and in-place prefix trim mirror the owned MTP prefill
+    (single-row caches throughout; the batched warm-merge machinery
+    never runs). B=1 unbatched batches only; anything else stays stock.
+    """
+    manager = getattr(batch, "_apc_manager", None)
+    mode = getattr(batch, "_apc_mode", None)
+    meta_list = getattr(batch, "_apc_meta", None) or []
+    if (
+        manager is None
+        or mode != "exact"
+        or len(meta_list) != 1
+        or meta_list[0] is None
+        or len(batch.uids) != 1
+        or batch._right_pad_per_row is not None
+        or batch._inputs_embeds is None
+    ):
+        return
+    bs = int(manager.block_size)
+    if not _ckpt_active(batch.model, mode, bs):
+        return
+    meta = meta_list[0]
+    if int(meta.get("prefix_len") or 0):
+        return  # stock warm row: leave it stock
+    ids_list = [int(t) for t in meta["full_input_ids"]]
+    if len(ids_list) < 2:
+        return
+    extra_hash = int(meta.get("extra_hash", 0))
+    view = _L1View(batch.model, manager, mode)
+    restored = 0
+    from gmlx.cache.snapshot import ckpt_lookup
+    warm, cp = ckpt_lookup(
+        manager,
+        ids_list,
+        extra_hash=extra_hash,
+        min_prefix_tokens=view._apc_safe_prefix_lookup_min(ids_list),
+        layout=_ckpt_layout_live(batch, bs),
+    )
+    if (
+        warm is not None
+        and 0 < cp < len(ids_list)
+        and view._apc_suffix_is_text_only(ids_list, cp)
+    ):
+        batch.prompt_cache = warm
+        batch._input_ids = batch._input_ids[:, cp:]
+        batch._inputs_embeds = batch._inputs_embeds[:, cp:]
+        batch._processed_prompt_columns = cp
+        for k in batch._prompt_length_aware_keys:
+            batch._prompt_kwargs[k] = batch._prompt_kwargs[k][:, cp:, ...]
+        restored = cp
+        batch._kq_apc_restored = (int(cp), "ckpt")     # live request view
+        _log.info("APC L1 hit: prefix=%d suffix=%d tier=ckpt",
+                  cp, len(ids_list) - cp)
+    guard = int(meta.get("checkpoint_len") or 0)
+    _ckpt_arm_schedule(batch, meta, guard, restored, bs)
+    batch._apc_harvest_enabled = False
+    batch._kq_ckpt_armed = True
+    from gmlx.cache.snapshot import ckpt_note_armed
+    ckpt_note_armed(manager)
+    if not _SPEC_APC_RETIRE_DISABLED and batch.prompt_cache:
+        from gmlx.cache.retire_key import lookup_render_ctx
+        batch.prompt_cache[0]._kq_apc_retire = {
+            "full_ids": ids_list,
+            "extra_hash": extra_hash,
+            "mode": "ckpt",
+            "checkpoint_len": int(meta.get("checkpoint_len") or 0),
+            "apc_meta": meta,
+            "render_ctx": lookup_render_ctx(ids_list),
+            "manager": manager,
+            "gen": [],
+            **_snap_fields(batch, manager),
+        }
+
+
+def _plain_anchor_init(batch) -> None:
+    """Arm the exact-tier anchor stop on a stock prompt batch (non-ckpt
+    exact models: DeepSeek-V4-class pooling stacks).
+
+    Restores come from the admission pick (_install_exact_anchor_pick),
+    so this only schedules the store. Warm and right-padded rows are
+    included: a restored prefix is usually far short of the divergence
+    (a bare bos match off some unrelated request), and upstream's
+    checkpoint column and row extraction handle both shapes. Refusing
+    them would skip every row that rides a warm batch, which on a busy
+    server is nearly all of them. The restored prefix becomes the
+    boundary floor, so a row already past the divergence arms nothing.
+    """
+    manager = getattr(batch, "_apc_manager", None)
+    mode = getattr(batch, "_apc_mode", None)
+    meta_list = getattr(batch, "_apc_meta", None) or []
+    if (manager is None or mode != "exact" or len(meta_list) != 1
+            or meta_list[0] is None or len(batch.uids) != 1
+            or batch._inputs_embeds is None):
+        return
+    if _ckpt_active(batch.model, mode, int(manager.block_size)):
+        return                          # ckpt tier owns these models
+    meta = meta_list[0]
+    if len(meta.get("full_input_ids") or ()) < 2:
+        return
+    _exact_anchor_arm(batch, meta, int(meta.get("checkpoint_len") or 0),
+                      int(meta.get("prefix_len") or 0))
+    # Retirement stash, independent of the anchor outcome: exact-tier
+    # rows retire their full post-decode row at filter (the per-turn
+    # store the post-prefill exact store cannot cover), warm rows
+    # included -- the decode cache holds the full sequence either way.
+    if not _SPEC_APC_RETIRE_DISABLED and batch.prompt_cache:
+        from gmlx.cache.retire_key import lookup_render_ctx
+        ids_list = [int(t) for t in meta["full_input_ids"]]
+        batch.prompt_cache[0]._kq_apc_retire = {
+            "full_ids": ids_list,
+            "extra_hash": int(meta.get("extra_hash", 0)),
+            "mode": "exact",
+            "manager": manager,
+            "render_ctx": lookup_render_ctx(ids_list),
+            "gen": [],
+        }
+
+
+_ANCHOR_PICK_FLAG = "_kq_exact_anchor_pick"
+
+
+def _install_exact_anchor_pick() -> None:
+    """Consult the anchor LRU inside the stock admission pick.
+
+    The pick is where a warm prefix belongs: admission builds the batch
+    from it (suffix rows, right padding, warm-cache merge) and every
+    downstream path treats an anchor restore exactly like a stock exact
+    one. The anchor wins only when strictly longer than the stock pick,
+    so it never shortens a restore. Idempotent.
+    """
+    from mlx_vlm.generate.ar import BatchGenerator
+    if getattr(BatchGenerator._apc_pick_for, _ANCHOR_PICK_FLAG, False):
+        return
+    _orig = BatchGenerator._apc_pick_for
+
+    def _pick_with_anchor(self, sequence):
+        pick = _orig(self, sequence)
+        try:
+            if _SPEC_APC_DISABLED or getattr(self, "apc_mode", None) != "exact":
+                return pick
+            manager = getattr(self, "apc_manager", None)
+            if manager is None or _ckpt_active(
+                    getattr(self, "model", None), "exact",
+                    int(manager.block_size)):
+                return pick
+            _uid, ids_list, _mt, prompt_kwargs, _lps, _crit = sequence
+            if not ids_list or len(ids_list) < 2:
+                return pick
+            # Floor trivial exact picks: a sub-block restore (a bare-BOS
+            # match off an unrelated request) saves nothing but suffix-
+            # constructs the batch, knocking the spec path's ids out of
+            # render space (anchor + retirement keys). Real warm picks are
+            # thousands of tokens and pass untouched.
+            if (pick is not None and not pick.get("matched_blocks")
+                    and 0 < int(pick.get("prefix_len") or 0)
+                    < int(manager.block_size)):
+                pick = None
+            have = int((pick or {}).get("prefix_len") or 0)
+            extra_hash = self._apc_extra_hash(prompt_kwargs or {})
+            floor = max(have, self._apc_safe_prefix_lookup_min(ids_list))
+            from gmlx.cache.snapshot import anchor_exact_lookup
+            warm, ap = anchor_exact_lookup(
+                manager, ids_list, extra_hash=extra_hash,
+                min_prefix_tokens=floor)
+            if warm is None or ap <= have or ap >= len(ids_list):
+                return pick
+            if not self._apc_suffix_is_text_only(ids_list, ap):
+                return pick
+            if pick and pick.get("matched_blocks"):
+                manager.release(pick["matched_blocks"])
+            _log.info("APC L1 hit: prefix=%d suffix=%d tier=anchor",
+                      ap, len(ids_list) - ap)
+            return {
+                "matched_blocks": [],
+                "warm_cache": warm,
+                "prefix_len": ap,
+                "extra_hash": extra_hash,
+                "full_input_ids": list(ids_list),
+            }
+        except Exception:
+            _log.warning("APC anchor pick failed; continuing",
+                         exc_info=True)
+            return pick
+
+    _pick_with_anchor.__dict__[_ANCHOR_PICK_FLAG] = True
+    BatchGenerator._apc_pick_for = _pick_with_anchor
+
+
+_PLAIN_DECODE_FLAG = "_kq_ckpt_plain_decode"
+
+
+def _retire_rows(gb) -> dict:
+    """uid -> retire-stash registry on a generation batch.
+
+    Stashes arm on the B=1 prompt batch's cache object (the only stable
+    home before the decode batch exists); the first decode-side touch
+    lifts them here so they survive ``extend`` rebuilding the cache
+    objects at continuous-batch injection."""
+    reg = getattr(gb, "_kq_apc_retire_rows", None)
+    if reg is None:
+        reg = {}
+        gb._kq_apc_retire_rows = reg
+    return reg
+
+
+def _lift_cache_stash(gb) -> None:
+    if not getattr(gb, "prompt_cache", None) or len(gb.uids) != 1:
+        return
+    stash = getattr(gb.prompt_cache[0], "_kq_apc_retire", None)
+    if stash is not None:
+        gb.prompt_cache[0]._kq_apc_retire = None
+        _retire_rows(gb)[gb.uids[0]] = stash
+
+
+def _plain_step_tick(gb, out) -> None:
+    """Per-token accounting + snapshot tick for stock-path retire rows.
+
+    Rows are tracked per uid so accounting survives ``extend`` merges.
+    Runs per step, so a deterministic failure disables the hook for that
+    row on first strike instead of emitting a traceback per token;
+    dropping ``gen`` also quiets retirement (its offset check would skip
+    anyway on a broken count). The decode-time snapshot ring stays B=1
+    (its clones ride the live single-row caches); rows in a B>1 batch
+    retire snapshot-free, under their verbatim key or an LCP cap the
+    tier arm can serve without a ring."""
+    try:
+        _lift_cache_stash(gb)
+        reg = getattr(gb, "_kq_apc_retire_rows", None)
+    except Exception:
+        _log.warning("APC plain decode hook failed; continuing",
+                     exc_info=True)
+        return
+    if not reg:
+        return
+    # _step returns (tokens, lps, top_idx, top_lp); slot 0 is the flat
+    # per-row token list.
+    rows = out[0] if isinstance(out, tuple) else out
+    if rows is None:
+        return
+    solo = len(gb.uids) == 1
+    for i, uid in enumerate(gb.uids):
+        stash = reg.get(uid)
+        if stash is None or "gen" not in stash:
+            continue
+        tok = rows[i] if i < len(rows) else None
+        if tok is None:
+            continue                    # no emission for this row this tick
+        try:
+            if isinstance(tok, (list, tuple)):
+                tok = tok[0]
+            stash["gen"].append(int(tok))
+            if solo and stash.get("mode") == "ckpt":
+                from gmlx.cache.snapshot import decode_ckpt_tick
+                decode_ckpt_tick(stash, gb.prompt_cache, stash["gen"])
+        except Exception:
+            stash.pop("gen", None)
+            stash["snap_ok"] = False
+            _log.warning("APC plain decode hook failed; disabled for "
+                         "this request", exc_info=True)
+
+
+def _plain_retire(stash: dict, prompt_cache: list) -> None:
+    """Retire a finished stock-path row off a single-row cache list.
+
+    Offset invariants mirror ``speculative._retire_b1``: the stock step
+    loop forwards each token as it emits it, so a clean finish leaves
+    ``offset == len(seq)`` (an abort between steps leaves the same).
+    ``stash["mode"]`` picks the tier arm: "ckpt" stores blocks +
+    sidecar, "exact" a whole-row snapshot (DeepSeek-V4-class pooling
+    stacks).
+    """
+    try:
+        manager = stash.get("manager")
+        if manager is None:
+            return
+        gen = [int(t) for t in stash.get("gen") or ()]
+        if not gen:
+            return
+        seq = [int(t) for t in stash["full_ids"]] + gen
+        from gmlx.cache.snapshot import _cache_offset_max, retirement_store
+        offset = _cache_offset_max(prompt_cache)
+        if offset == len(seq) - 1:
+            seq = seq[:-1]
+        elif offset != len(seq):
+            _log.info(
+                "APC retire skipped: cache offset %d != tokens %d", offset, len(seq)
+            )
+            return
+        lcp = None
+        if os.environ.get("GMLX_APC_RETIRE_LCP") != "0":
+            from gmlx.cache.retire_key import next_turn_lcp
+            lcp = next_turn_lcp(stash.get("render_ctx"), seq, gen)
+        max_len = lcp if lcp is not None and lcp < len(seq) else None
+        _log.info("APC retire: seq=%d ctx=%s lcp=%s cap=%s",
+                  len(seq), stash.get("render_ctx") is not None,
+                  lcp, max_len)
+        ok = retirement_store(
+            manager, stash.get("mode") or "ckpt", seq, prompt_cache,
+            row=0,
+            extra_hash=int(stash.get("extra_hash", 0)), max_len=max_len,
+            decode_snaps=stash.get("snaps"))
+        if ok:
+            _log.info("APC retire store: tokens=%d", ok)
+    except Exception:
+        _log.warning("APC retire failed; continuing", exc_info=True)
+
+
+def _install_plain_ckpt_decode() -> None:
+    """Stock-path decode hooks for the retirement store (ckpt + exact).
+
+    Token accounting rides ``_step``; retirement fires from ``filter``
+    for every leaving row (finish or client abort). A lone row retires
+    off its live single-row caches; a row leaving a B>1 batch is first
+    extracted via ``row_snapshot`` (padding-trimmed clones with row-true
+    offsets), so retirement survives concurrency instead of firing only
+    when the batch happens to drain to one row. Stashes live in a
+    uid-keyed registry lifted across ``extend`` (the seam that rebuilds
+    cache objects at continuous-batch injection).
+    GMLX_APC_RETIRE_BATCH=0 restores the lone-row-only v1 scope.
+    Idempotent."""
+    from mlx_vlm.generate.ar import GenerationBatch
+
+    if getattr(GenerationBatch._step, _PLAIN_DECODE_FLAG, False):
+        return
+    _orig_step = GenerationBatch._step
+    _orig_filter = GenerationBatch.filter
+    _orig_extend = GenerationBatch.extend
+
+    def _step_with_ckpt(self):
+        out = _orig_step(self)
+        _plain_step_tick(self, out)
+        return out
+
+    def _filter_with_ckpt(self, keep):
+        try:
+            _lift_cache_stash(self)
+            reg = getattr(self, "_kq_apc_retire_rows", None)
+            if reg and self.prompt_cache:
+                keep_set = set(keep)
+                solo = len(self.uids) == 1
+                batched_ok = os.environ.get(
+                    "GMLX_APC_RETIRE_BATCH") != "0"
+                for i, uid in enumerate(self.uids):
+                    if i in keep_set:
+                        continue
+                    stash = reg.pop(uid, None)
+                    if stash is None:
+                        continue
+                    if solo:
+                        _plain_retire(stash, self.prompt_cache)
+                    elif batched_ok:
+                        from gmlx.cache.snapshot import row_snapshot
+                        rows = row_snapshot(self.prompt_cache, i)
+                        if rows is None:
+                            _log.info("APC retire skipped: row %d "
+                                      "extract unavailable", i)
+                        else:
+                            _plain_retire(stash, rows)
+        except Exception:
+            _log.warning("APC plain retire hook failed; continuing", exc_info=True)
+        _orig_filter(self, keep)
+
+    def _extend_with_ckpt(self, other):
+        try:
+            _lift_cache_stash(self)
+            _lift_cache_stash(other)
+            other_reg = getattr(other, "_kq_apc_retire_rows", None)
+            if other_reg:
+                _retire_rows(self).update(other_reg)
+                other._kq_apc_retire_rows = {}
+        except Exception:
+            _log.warning("APC retire stash carry failed; continuing",
+                         exc_info=True)
+        _orig_extend(self, other)
+
+    _step_with_ckpt.__dict__[_PLAIN_DECODE_FLAG] = True
+    _filter_with_ckpt.__dict__[_PLAIN_DECODE_FLAG] = True
+    _extend_with_ckpt.__dict__[_PLAIN_DECODE_FLAG] = True
+    GenerationBatch._step = _step_with_ckpt
+    GenerationBatch.filter = _filter_with_ckpt
+    GenerationBatch.extend = _extend_with_ckpt
diff --git a/gmlx/spec/engine.py b/gmlx/spec/engine.py
index 6ad84310..88729b5b 100644
--- a/gmlx/spec/engine.py
+++ b/gmlx/spec/engine.py
@@ -20,9 +20,7 @@
 
 import mlx.core as mx
 
-import gmlx.lora_rows as lora_rows
-import gmlx.gen.prefill_decay as prefill_decay
-from gmlx.envflags import env_bool, env_int
+from gmlx.envflags import env_int
 
 _log = logging.getLogger(__name__)
 
@@ -390,2122 +388,6 @@ def _live_kv_quant_config(model=None):
                      "float", exc_info=True)
         return None
 
-
-def _l1_lookup_and_arm_store(batch, manager, mode, l0_prefix) -> int:
-    """Consult the shared APCManager below L0 and arm the stock post-prefill
-    store (mid-prefill exact checkpoints + post-prefill exact store / block
-    harvest, all owned by stock ``PromptProcessingBatch.generate``) by
-    populating ``_apc_manager`` / ``_apc_mode`` / ``_apc_meta``.
-
-    Returns the restored L1 prefix length (0 on miss, or when L0 already
-    restored -- L0 carries full-prompt hidden and is always preferred).
-
-    ``meta["prefix_len"]`` stays 0 by design: the owned prefill keeps
-    ``_processed_prompt_columns`` in absolute token space (it trims
-    ``_input_ids`` in place, unlike stock warm batches which are constructed
-    with suffix-only rows), and ``_row_real_tokens_processed`` -- which gates
-    the mid-prefill checkpoint store -- is only correct in that space with a
-    zero meta prefix. The one cost is that a block-tier harvest re-walks
-    restored prefix blocks, but ``store_kv_blocks`` dedups by hash chain
-    (acquire+release of existing blocks, no data copies).
-    """
-    view = _L1View(batch.model, manager, mode)
-    ids_list = [int(t) for t in batch._mtp_full_input_ids[0].tolist()]
-    prompt_kwargs = batch._prompt_kwargs or {}
-    extra_hash = view._apc_extra_hash(prompt_kwargs)
-    ckpt = _ckpt_active(batch.model, mode, int(manager.block_size))
-    held_blocks = []
-    l1_prefix = 0
-    if l0_prefix == 0 and len(ids_list) >= 2:
-        warm = None
-        blocks = []
-        prefix_len = 0
-        tier = "exact"
-        pick = view._apc_pick_for((0, ids_list, 0, prompt_kwargs, None, None))
-        # Same trivial-pick floor as the admission wrapper: a sub-block
-        # exact restore saves nothing and its nonzero l1_prefix would skip
-        # the L0 hidden store for this request.
-        if (pick is not None and not pick.get("matched_blocks")
-                and 0 < int(pick.get("prefix_len") or 0)
-                < int(manager.block_size)):
-            pick = None
-        if pick is not None:
-            warm = pick.get("warm_cache")
-            blocks = list(pick.get("matched_blocks") or ())
-            prefix_len = int(pick.get("prefix_len") or 0)
-            extra_hash = int(pick.get("extra_hash", extra_hash))
-            if warm is None and blocks:
-                from mlx_vlm import apc as _apc
-
-                warm = _apc.make_warm_kv_cache(
-                    blocks, min_capacity_tokens=len(ids_list) + 1
-                )
-                tier = "block"
-        if ckpt:
-            # Checkpoint tier: the longest salted sidecar + block chain
-            # wins only when strictly longer than the exact-tier pick.
-            # Media guards mirror the stock exact probe.
-            from gmlx.cache.snapshot import ckpt_lookup
-            min_p = max(prefix_len,
-                        view._apc_safe_prefix_lookup_min(ids_list))
-            cw, cp = ckpt_lookup(
-                manager,
-                ids_list,
-                extra_hash=extra_hash,
-                min_prefix_tokens=min_p,
-                layout=_ckpt_layout_live(batch, int(manager.block_size)),
-            )
-            if (
-                cw is not None
-                and cp > prefix_len
-                and view._apc_suffix_is_text_only(ids_list, cp)
-            ):
-                if blocks:
-                    manager.release(blocks)
-                    blocks = []
-                warm, prefix_len, tier = cw, cp, "ckpt"
-        elif mode == "exact":
-            # Exact-tier anchor: the shared-system-prefix clone in the
-            # gmlx anchor LRU wins only when strictly longer than the
-            # stock exact pick. Media guards mirror the stock probe.
-            from gmlx.cache.snapshot import anchor_exact_lookup
-            min_p = max(prefix_len,
-                        view._apc_safe_prefix_lookup_min(ids_list))
-            aw, ap = anchor_exact_lookup(
-                manager, ids_list, extra_hash=extra_hash,
-                min_prefix_tokens=min_p)
-            if (aw is not None and ap > prefix_len
-                    and view._apc_suffix_is_text_only(ids_list, ap)):
-                if blocks:
-                    manager.release(blocks)
-                    blocks = []
-                warm, prefix_len, tier = aw, ap, "anchor"
-        if warm and 0 < prefix_len < len(ids_list) and tier in (
-                "exact", "anchor"):
-            # Same batch-aware merge admission applies to its picks: raw
-            # exact/anchor clones carry single-row leaves (left_padding
-            # None, scalar offsets) and crash the batch cache classes'
-            # update path (mx.depends on a None) when the suffix forwards.
-            # kv_quant_config re-quantizes the float snapshot to the live
-            # _make_cache layer types under serve kv_bits (stored exact
-            # entries stay float; a float row joining a quantized batch
-            # breaks the update path).
-            from mlx_vlm import apc as _apc
-            warm, _ = _apc.make_warm_batch_exact_cache_multi(
-                [warm], prefix_lens=[prefix_len],
-                kv_quant_config=_live_kv_quant_config(batch.model))
-        if warm and 0 < prefix_len < len(ids_list):
-            batch.prompt_cache = warm
-            # Matched blocks stay acquired until the stock post-prefill
-            # harvest releases them (the warm-cache concatenation is
-            # lazy; the pool must not recycle these blocks before it
-            # materializes).
-            held_blocks = blocks
-            l1_prefix = prefix_len
-            # Observability only: the live request view reads the
-            # restored prefix from here (meta keeps prefix_len 0 so the
-            # stock machinery does not account it twice).
-            batch._kq_apc_restored = (int(prefix_len), str(tier))
-            _log.info(
-                "APC L1 hit: prefix=%d suffix=%d tier=%s",
-                prefix_len,
-                len(ids_list) - prefix_len,
-                tier,
-            )
-            # Drafter-KV sidecar: a plain L1 hit restores target KV but
-            # not hidden, so the drafter would re-seed from suffix-only
-            # hidden at the wrong positions (acceptance erodes at depth).
-            # A sidecar covering exactly the restored prefix hands the
-            # owned round a warm drafter start. Stash rides the first
-            # cache entry, same discipline as the retirement context.
-            if not _SPEC_APC_SIDECAR_DISABLED:
-                from gmlx.cache.snapshot import drafter_sidecar_lookup
-                side = drafter_sidecar_lookup(
-                    manager, ids_list, prefix_len, extra_hash)
-                if side:
-                    batch.prompt_cache[0]._kq_apc_drafter_warm = side
-                    _log.info("APC sidecar hit: prefix=%d", prefix_len)
-        elif blocks:
-            manager.release(blocks)
-    batch._mtp_l1_prefix_len = l1_prefix
-    batch._apc_manager = manager
-    batch._apc_mode = mode
-    guard = int(view._apc_exact_checkpoint_len(ids_list) or 0)
-    meta = {
-        "full_input_ids": ids_list,
-        "prefix_len": 0,
-        "extra_hash": extra_hash,
-        "apc_blocks": held_blocks,
-        "checkpoint_len": guard,
-    }
-    batch._apc_meta = [meta]
-    if ckpt:
-        # The checkpoint tier replaces the stock exact-tier stores: the
-        # post-prefill full-cache clone is suppressed here, and the
-        # mid-prefill checkpoint store is superseded by the cursor riding
-        # the wrapped stock store (_install_ckpt_checkpoint_store; the
-        # stock body is suppressed by the cursor's advance). Column
-        # alignment itself still runs on the stock machinery, which
-        # requires _apc_mode == "exact".
-        _ckpt_arm_schedule(batch, meta, guard,
-                           max(l0_prefix, l1_prefix),
-                           int(manager.block_size))
-        batch._apc_harvest_enabled = False
-        batch._kq_ckpt_armed = True
-        from gmlx.cache.snapshot import ckpt_note_armed
-        ckpt_note_armed(manager)
-    elif mode == "exact":
-        _exact_anchor_arm(batch, meta, guard,
-                          max(l0_prefix, l1_prefix))
-    return l1_prefix
-
-
-def _gcd(a: int, b: int) -> int:
-    while b:
-        a, b = b, a % b
-    return a
-
-
-def _ckpt_unit(batch, block_size: int) -> int:
-    """The natural chunk grid: lcm(prefill_step_size, block_size)."""
-    step = int(getattr(batch, "prefill_step_size", 0) or 0)
-    return block_size if step <= 0 else \
-        step * block_size // _gcd(step, block_size)
-
-
-def _ckpt_cursor_init(batch, guard: int, restored: int,
-                      block_size: int) -> tuple[list, int, int]:
-    """Boundary schedule for the checkpoint cursor: an ordered
-    ``[(position, kind), ...]`` list plus ``(terminal, interval)``.
-
-    Boundaries sit on the natural chunk grid, lcm(prefill_step_size,
-    block_size) -- an off-grid boundary truncates a chunk, and gated-delta
-    state is chunk-shape sensitive (certified: any grid change drifts).
-    Interval points above the restored prefix, then the terminal (the
-    grid point at or below the stock guard column). GMLX_APC_CKPT_INTERVAL
-    tokens, default 4096, snapped up to the grid; 0 = terminal-only.
-    Later stages add store positions by appending boundaries here, never
-    by new store mechanisms.
-    """
-    unit = _ckpt_unit(batch, block_size)
-    terminal = (guard // unit) * unit
-    if terminal <= max(0, restored):
-        return [], 0, 0
-    raw = env_int("GMLX_APC_CKPT_INTERVAL", 4096)
-    interval = 0 if raw <= 0 else max(unit, (raw // unit) * unit)
-    bounds = []
-    if interval:
-        b = ((max(0, restored) // interval) + 1) * interval
-        while b < terminal:
-            bounds.append((b, "boundary"))
-            b += interval
-    bounds.append((terminal, "boundary"))
-    return bounds, terminal, interval
-
-
-def _ckpt_replay_boundary(batch, meta, restored: int,
-                          block_size: int) -> int | None:
-    """N-1 replay boundary, or None when it cannot earn its pause.
-
-    An identical resend can only adopt a record strictly below the
-    query, and the interval/terminal schedule never places one there
-    for prompts under one interval (the depth e2e's bug 1); N-1 is the
-    deepest position that is both adoptable and drift-free (the warm
-    turn forwards exactly one token), and the pause is free on the cold
-    side -- both prefill loops already stop at N-1 to feed the first
-    decode step, so the boundary lands on a natural chunk edge and
-    perturbs no chunk shape. arr layouts gate on a minimum N:
-    recurrent state is prompt-length-independent (>100 MB per record on
-    27B-class), and short-prompt records would churn the LRU out of the
-    deep-conversation records it exists to protect. Rotating layouts
-    need N-1 at or past the window (below it the store's grid gate
-    declines). Kill switch: GMLX_APC_CKPT_REPLAY=0.
-    """
-    if env_int("GMLX_APC_CKPT_REPLAY", 1) == 0:
-        return None
-    n = len(meta.get("full_input_ids") or ())
-    replay = n - 1
-    if replay < 2 or replay <= max(0, restored):
-        return None
-    tags = _ckpt_layout_live(batch, block_size) or ()
-    if "arr" in tags and n < env_int("GMLX_APC_CKPT_REPLAY_MIN", 1024):
-        return None
-    for t in tags:
-        if t.startswith("rot") and replay < int(t.split(":")[1]):
-            return None
-    return replay
-
-
-def _ckpt_turn_boundaries(batch, meta, restored: int,
-                          block_size: int) -> list[int]:
-    """Render-stable turn boundary positions for the schedule.
-
-    p_stable is the deepest prompt position a next-turn re-render keeps;
-    the gen-prompt/think tail past it is re-rendered away, so records
-    stored only above it can never serve turn 2 (how multi-turn adoption
-    silently died). Every layout gets the grid point at or below
-    p_stable (drift-free for chunk-shape-sensitive state); rot-only
-    layouts also pause exactly at p_stable (attention splits exactly;
-    needs the window wrapped). GMLX_APC_CKPT_TURN=0 disables these and
-    with them the p=N drop gate.
-    """
-    if env_int("GMLX_APC_CKPT_TURN", 1) == 0:
-        return []
-    ids = meta.get("full_input_ids") or ()
-    unit = _ckpt_unit(batch, block_size)
-    tags = _ckpt_layout_live(batch, block_size) or ()
-    ws = [int(t.split(":")[1]) for t in tags if t.startswith("rot")]
-    # Cheapest boundary this layout could arm: the grid needs one unit
-    # of stable prefix; rot-only layouts can also pause exactly at
-    # p_stable once the window wraps. Below that no boundary can land,
-    # so skip the render+tokenize prediction entirely.
-    need = unit if ("arr" in tags or not ws) else min(unit, max(ws))
-    if len(ids) - 1 < need:
-        return []
-    from gmlx.cache.retire_key import lookup_render_ctx, prompt_stable_lcp
-    ctx = lookup_render_ctx(ids)
-    p_stable = prompt_stable_lcp(ctx, ids) if ctx else None
-    if not p_stable or p_stable < 2:
-        return []
-    p_stable = min(int(p_stable), len(ids) - 1)
-    meta["ckpt_p_stable"] = p_stable
-    floor = max(0, restored)
-    out = []
-    grid = (p_stable // unit) * unit
-    if grid > floor:
-        out.append(grid)
-    if ws and "arr" not in tags and p_stable != grid \
-            and p_stable > floor and p_stable >= max(ws):
-        out.append(p_stable)
-    return out
-
-
-def _ckpt_sys_boundary(batch, meta, restored: int,
-                       block_size: int) -> int | None:
-    """Anchor stop at the end of the shared system prefix.
-
-    Sibling fan-out requests share the system prompt and tool schemas
-    and diverge at the first user message, generally between grid
-    points, so the interval schedule alone wastes up to one interval of
-    sibling recompute, and strip-on-extend removes the early boundary
-    the siblings need as the chain deepens (the anchor exemption in
-    _record_insert keeps this one). arr layouts snap the stop down to
-    the chunk grid (off-grid chunking drifts GDN state) and keep the
-    replay byte floor (recurrent state is prompt-length-independent, so
-    a tiny anchor costs the same >100 MB clone as a deep one);
-    attention layouts snap to the block grid, which also satisfies the
-    rotating store's below-window grid gate. GMLX_APC_CKPT_SYS=0
-    disables; GMLX_APC_CKPT_SYS_MIN floors the position (a sub-floor
-    shared prefix re-prefills in milliseconds and is not worth a
-    record).
-    """
-    if env_int("GMLX_APC_CKPT_SYS", 1) == 0:
-        return None
-    ids = meta.get("full_input_ids") or ()
-    tags = _ckpt_layout_for(getattr(batch, "model", None), block_size) or ()
-    floor_min = max(block_size, env_int("GMLX_APC_CKPT_SYS_MIN", 256))
-    if "arr" in tags:
-        floor_min = max(floor_min,
-                        env_int("GMLX_APC_CKPT_REPLAY_MIN", 1024))
-    # Below the floor no anchor can land; skip the render+tokenize
-    # prediction entirely (same rule as the turn boundaries).
-    if len(ids) - 1 < floor_min:
-        return None
-    from gmlx.cache.retire_key import lookup_render_ctx, system_prefix_lcp
-    ctx = lookup_render_ctx(ids)
-    lcp = system_prefix_lcp(ctx, ids) if ctx else None
-    if not lcp:
-        return None
-    unit = _ckpt_unit(batch, block_size) if "arr" in tags else block_size
-    pos = (min(int(lcp), len(ids) - 1) // unit) * unit
-    if pos < floor_min or pos <= max(0, restored):
-        return None
-    meta["ckpt_sys_bound"] = pos
-    return pos
-
-
-def _exact_anchor_boundary(batch, meta, guard: int,
-                           restored: int) -> int | None:
-    """Anchor position for an exact-tier (non-ckpt) model: the sibling
-    divergence point, ungridded (exact clones restore at any position).
-    Clamped to the stock guard column, so the prefill pauses at most
-    twice: once for the anchor, once for the stock guard store.
-    GMLX_APC_CKPT_SYS=0 disables (one switch for both tiers);
-    GMLX_APC_CKPT_SYS_MIN floors the position (a sub-floor shared
-    prefix re-prefills in milliseconds and is not worth a clone).
-    """
-    if env_int("GMLX_APC_CKPT_SYS", 1) == 0:
-        return None
-    ids = meta.get("full_input_ids") or ()
-    floor_min = max(2, env_int("GMLX_APC_CKPT_SYS_MIN", 256))
-    if len(ids) - 1 < floor_min:
-        return None
-    from gmlx.cache.retire_key import lookup_render_ctx, system_prefix_lcp
-    ctx = lookup_render_ctx(ids)
-    lcp = system_prefix_lcp(ctx, ids) if ctx else None
-    if not lcp:
-        _log.info("APC anchor declined: no measurable system prefix "
-                  "(render ctx %s)", "present" if ctx else "missing")
-        return None
-    pos = min(int(lcp), len(ids) - 1)
-    if guard > 0:
-        pos = min(pos, guard)
-    if pos < floor_min or pos <= max(0, restored):
-        return None
-    return pos
-
-
-def _exact_anchor_arm(batch, meta, guard: int, restored: int) -> None:
-    """Schedule the anchor pause by mirroring its position into
-    ``checkpoint_len`` (the key the stock column truncation reads).
-    ``_exact_anchor_store`` hands the column back to the stock guard
-    after firing, so the stock store still runs exactly as unarmed."""
-    pos = _exact_anchor_boundary(batch, meta, guard, restored)
-    if pos is None:
-        return
-    meta["anchor_len"] = pos
-    meta["anchor_guard"] = guard
-    if pos != guard:
-        meta["checkpoint_len"] = pos
-    batch._kq_anchor_armed = True
-    _log.info("APC anchor armed: pos=%d guard=%d", pos, guard)
-
-
-def _exact_anchor_store(batch) -> None:
-    """Anchor store for exact-tier models: one whole-prefix clone at the
-    sibling divergence, into the gmlx anchor LRU. Runs from the wrapped
-    stock store immediately before the stock body; after firing it
-    restores ``checkpoint_len`` to the stock guard column without
-    latching ``checkpoint_done``, so the stock guard store (and its
-    latch) fire untouched."""
-    manager = getattr(batch, "_apc_manager", None)
-    meta_list = getattr(batch, "_apc_meta", None) or []
-    if manager is None or not meta_list or meta_list[0] is None:
-        return
-    meta = meta_list[0]
-    pos = int(meta.get("anchor_len") or 0)
-    if pos <= 0 or meta.get("anchor_done"):
-        return
-    if batch._row_real_tokens_processed(0) != pos:
-        return
-    meta["anchor_done"] = True
-    guard = int(meta.get("anchor_guard") or 0)
-    if int(meta.get("checkpoint_len") or 0) == pos and pos != guard:
-        meta["checkpoint_len"] = guard
-    cache = batch._apc_prompt_cache_for_store(0)
-    if cache is None:
-        return
-    from gmlx.cache.snapshot import anchor_exact_store
-    anchor_exact_store(manager, meta["full_input_ids"][:pos], cache,
-                       extra_hash=int(meta.get("extra_hash", 0)))
-
-
-def _sched_insert(bounds: list, pos: int, kind: str, *,
-                  upgrade: bool = False) -> None:
-    """Insert (pos, kind) keeping order. On collision the existing entry
-    keeps its kind: a colliding position is always grid-aligned or an
-    exact turn boundary, where a plain boundary record adopts freely --
-    identical resend included -- while flipping it to replay would gate
-    turn-2 and branch adoption out on recurrent layouts (and satisfy the
-    p=N drop with a record turn 2 cannot use). ``upgrade`` lets an
-    anchor replace a plain boundary at the same position (strictly more
-    retention, same free adoption), never a replay."""
-    import bisect
-
-    pts = [b for b, _ in bounds]
-    i = bisect.bisect_left(pts, pos)
-    if i < len(pts) and pts[i] == pos:
-        if upgrade and bounds[i][1] == "boundary":
-            bounds[i] = (pos, kind)
-        return
-    bounds.insert(i, (pos, kind))
-
-
-def _ckpt_arm_schedule(batch, meta, guard: int, restored: int,
-                       block_size: int) -> None:
-    """Publish the boundary schedule into the request meta. The head
-    mirrors into ``checkpoint_len`` (an int) because the stock
-    checkpoint-column truncation and store reads exactly that key.
-    ``ckpt_stored_boundaries`` collects every boundary whose store
-    landed (record verified in the index) -- the settled variable the
-    post-prefill p=N decision and the sidecar key set both read;
-    ``ckpt_p_stable_bounds`` is the qualifying set for the p=N drop."""
-    bounds, terminal, interval = _ckpt_cursor_init(
-        batch, guard, restored, block_size)
-    turn = _ckpt_turn_boundaries(batch, meta, restored, block_size)
-    for pos in turn:
-        _sched_insert(bounds, pos, "boundary")
-    sysb = _ckpt_sys_boundary(batch, meta, restored, block_size)
-    if sysb is not None:
-        _sched_insert(bounds, sysb, "anchor", upgrade=True)
-    replay = _ckpt_replay_boundary(batch, meta, restored, block_size)
-    if replay is not None:
-        # Colliding with the anchor keeps the anchor (default no-upgrade):
-        # it adopts identical resends freely, replay semantics add nothing.
-        _sched_insert(bounds, replay, "replay")
-    meta["ckpt_boundaries"] = bounds
-    meta["checkpoint_len"] = int(bounds[0][0]) if bounds else 0
-    meta["ckpt_terminal"] = terminal
-    meta["ckpt_interval"] = interval
-    meta["ckpt_last_stored"] = 0
-    meta["ckpt_stored_boundaries"] = []
-    meta["ckpt_p_stable_bounds"] = turn
-
-
-def _ckpt_mid_prefill_store(batch) -> None:
-    """Checkpoint-tier replacement for the stock mid-prefill exact store.
-
-    Fires at the schedule head, pops it, and mirrors the next head into
-    ``checkpoint_len``, latching ``checkpoint_done`` when the schedule
-    empties. The advance is what suppresses the stock store;
-    ``_install_ckpt_checkpoint_store`` wraps the stock method so the
-    cursor always runs immediately before it -- the ordering is
-    structural, not positional. Advances past failed stores;
-    ``ckpt_last_stored`` records only boundaries that landed.
-    """
-    if not getattr(batch, "_kq_ckpt_armed", False):
-        return
-    manager = getattr(batch, "_apc_manager", None)
-    meta_list = getattr(batch, "_apc_meta", None) or []
-    if manager is None or not meta_list or meta_list[0] is None:
-        return
-    meta = meta_list[0]
-    if meta.get("checkpoint_done"):
-        return
-    checkpoint_len = int(meta.get("checkpoint_len") or 0)
-    if checkpoint_len <= 0:
-        return
-    if batch._row_real_tokens_processed(0) != checkpoint_len:
-        return
-    terminal = int(meta.get("ckpt_terminal") or 0)
-    bounds = meta.get("ckpt_boundaries") or []
-    kind = "boundary"
-    if bounds and int(bounds[0][0]) == checkpoint_len:
-        kind = str(bounds.pop(0)[1])
-    # Inline-heavy skeletons (GDN state >100 MB; kvarn state scales with p
-    # across every attention layer) earn disk only at the terminal --
-    # boundaries superseded within the same prefill do not, and a replay
-    # skeleton would buy restart-repair of an identical resend only,
-    # which does not earn it either.
-    layout = _ckpt_layout_live(batch, int(manager.block_size)) or ()
-    heavy = "arr" in layout or any(t.startswith("kvarn") for t in layout)
-    skel = not heavy or (kind != "replay"
-                         and checkpoint_len >= terminal)
-    from gmlx.cache.snapshot import ckpt_store
-
-    if ckpt_store(
-            manager, meta["full_input_ids"][:checkpoint_len],
-            batch.prompt_cache, extra_hash=int(meta.get("extra_hash", 0)),
-            skeleton_disk=skel, kind=kind):
-        meta["ckpt_last_stored"] = checkpoint_len
-        meta.setdefault("ckpt_stored_boundaries", []).append(checkpoint_len)
-    if bounds:
-        meta["checkpoint_len"] = int(bounds[0][0])
-    else:
-        meta["checkpoint_done"] = True
-
-
-_CKPT_STORE_FLAG = "_kq_ckpt_cursor_store"
-
-
-def _install_ckpt_checkpoint_store() -> None:
-    """Wrap the stock mid-prefill checkpoint store so the cursor runs
-    immediately before it on armed batches (both the owned MTP prefill
-    and the stock prompt_step call the stock method, so one wrap covers
-    both paths). The cursor's advance of ``checkpoint_len`` is what
-    suppresses the stock store -- wrapping makes that ordering
-    structural. Exact-tier anchor batches ride the same wrap with their
-    own single-stop hook. Idempotent."""
-    from mlx_vlm.generate.ar import PromptProcessingBatch
-
-    if getattr(
-        PromptProcessingBatch._store_apc_exact_checkpoints, _CKPT_STORE_FLAG, False
-    ):
-        return
-    _orig = PromptProcessingBatch._store_apc_exact_checkpoints
-
-    def _store_with_ckpt_cursor(self):
-        if getattr(self, "_kq_ckpt_armed", False):
-            _ckpt_mid_prefill_store(self)
-        elif getattr(self, "_kq_anchor_armed", False):
-            _exact_anchor_store(self)
-        _orig(self)
-
-    _store_with_ckpt_cursor.__dict__[_CKPT_STORE_FLAG] = True
-    PromptProcessingBatch._store_apc_exact_checkpoints = _store_with_ckpt_cursor
-
-
-def _snap_fields(batch, manager) -> dict:
-    """Decode-time snapshot ring parameters for a retirement stash.
-
-    ``snap_grid`` anchors snapshot positions to the prefill chunk grid
-    (lcm of step and block size), so a restore replays chunk-exact --
-    but only while one grid unit fits inside the snapshot interval; a
-    serve-sized step (2048) would otherwise push the first snapshot far
-    past prompt end + interval, so it falls back to the block size (the
-    off-grid restore is the scoped-benign case). ``snap_align`` is the
-    block alignment a rotating window store requires below the window;
-    ``snap_offgrid_min`` (= W) is where the store gate stops caring --
-    a wrapped window is whole blocks at any p.
-    """
-    import math
-    from gmlx.cache.snapshot import _DECODE_CKPT_DEFAULT
-    bs = int(manager.block_size)
-    tags = _ckpt_layout_live(batch, bs) or ()
-    step = int(getattr(batch, "prefill_step_size", 0) or 0)
-    grid = math.lcm(step, bs) if step > 0 else bs
-    if grid > env_int("GMLX_APC_DECODE_CKPT", _DECODE_CKPT_DEFAULT):
-        grid = bs
-    rot_w = 0
-    for t in tags:
-        if t.startswith("rot"):
-            rot_w = int(t.split(":")[1])
-            break
-    return {
-        "snap_ok": bool(tags),
-        "snap_grid": grid,
-        "snap_align": bs if rot_w else 1,
-        "snap_offgrid_min": rot_w,
-    }
-
-
-def _plain_ckpt_init(batch) -> None:
-    """Checkpoint-tier lookup + arming for a stock (non-speculative)
-    prompt batch.
-
-    The stock path reaches the tier only here: exact-tier stores are
-    suppressed on ckpt models, so admission's own lookup ladder misses
-    and every ckpt-tier request arrives as a cold single-request batch.
-    Lookup and in-place prefix trim mirror the owned MTP prefill
-    (single-row caches throughout; the batched warm-merge machinery
-    never runs). B=1 unbatched batches only; anything else stays stock.
-    """
-    manager = getattr(batch, "_apc_manager", None)
-    mode = getattr(batch, "_apc_mode", None)
-    meta_list = getattr(batch, "_apc_meta", None) or []
-    if (
-        manager is None
-        or mode != "exact"
-        or len(meta_list) != 1
-        or meta_list[0] is None
-        or len(batch.uids) != 1
-        or batch._right_pad_per_row is not None
-        or batch._inputs_embeds is None
-    ):
-        return
-    bs = int(manager.block_size)
-    if not _ckpt_active(batch.model, mode, bs):
-        return
-    meta = meta_list[0]
-    if int(meta.get("prefix_len") or 0):
-        return  # stock warm row: leave it stock
-    ids_list = [int(t) for t in meta["full_input_ids"]]
-    if len(ids_list) < 2:
-        return
-    extra_hash = int(meta.get("extra_hash", 0))
-    view = _L1View(batch.model, manager, mode)
-    restored = 0
-    from gmlx.cache.snapshot import ckpt_lookup
-    warm, cp = ckpt_lookup(
-        manager,
-        ids_list,
-        extra_hash=extra_hash,
-        min_prefix_tokens=view._apc_safe_prefix_lookup_min(ids_list),
-        layout=_ckpt_layout_live(batch, bs),
-    )
-    if (
-        warm is not None
-        and 0 < cp < len(ids_list)
-        and view._apc_suffix_is_text_only(ids_list, cp)
-    ):
-        batch.prompt_cache = warm
-        batch._input_ids = batch._input_ids[:, cp:]
-        batch._inputs_embeds = batch._inputs_embeds[:, cp:]
-        batch._processed_prompt_columns = cp
-        for k in batch._prompt_length_aware_keys:
-            batch._prompt_kwargs[k] = batch._prompt_kwargs[k][:, cp:, ...]
-        restored = cp
-        batch._kq_apc_restored = (int(cp), "ckpt")     # live request view
-        _log.info("APC L1 hit: prefix=%d suffix=%d tier=ckpt",
-                  cp, len(ids_list) - cp)
-    guard = int(meta.get("checkpoint_len") or 0)
-    _ckpt_arm_schedule(batch, meta, guard, restored, bs)
-    batch._apc_harvest_enabled = False
-    batch._kq_ckpt_armed = True
-    from gmlx.cache.snapshot import ckpt_note_armed
-    ckpt_note_armed(manager)
-    if not _SPEC_APC_RETIRE_DISABLED and batch.prompt_cache:
-        from gmlx.cache.retire_key import lookup_render_ctx
-        batch.prompt_cache[0]._kq_apc_retire = {
-            "full_ids": ids_list,
-            "extra_hash": extra_hash,
-            "mode": "ckpt",
-            "checkpoint_len": int(meta.get("checkpoint_len") or 0),
-            "apc_meta": meta,
-            "render_ctx": lookup_render_ctx(ids_list),
-            "manager": manager,
-            "gen": [],
-            **_snap_fields(batch, manager),
-        }
-
-
-def _plain_anchor_init(batch) -> None:
-    """Arm the exact-tier anchor stop on a stock prompt batch (non-ckpt
-    exact models: DeepSeek-V4-class pooling stacks).
-
-    Restores come from the admission pick (_install_exact_anchor_pick),
-    so this only schedules the store. Warm and right-padded rows are
-    included: a restored prefix is usually far short of the divergence
-    (a bare bos match off some unrelated request), and upstream's
-    checkpoint column and row extraction handle both shapes. Refusing
-    them would skip every row that rides a warm batch, which on a busy
-    server is nearly all of them. The restored prefix becomes the
-    boundary floor, so a row already past the divergence arms nothing.
-    """
-    manager = getattr(batch, "_apc_manager", None)
-    mode = getattr(batch, "_apc_mode", None)
-    meta_list = getattr(batch, "_apc_meta", None) or []
-    if (manager is None or mode != "exact" or len(meta_list) != 1
-            or meta_list[0] is None or len(batch.uids) != 1
-            or batch._inputs_embeds is None):
-        return
-    if _ckpt_active(batch.model, mode, int(manager.block_size)):
-        return                          # ckpt tier owns these models
-    meta = meta_list[0]
-    if len(meta.get("full_input_ids") or ()) < 2:
-        return
-    _exact_anchor_arm(batch, meta, int(meta.get("checkpoint_len") or 0),
-                      int(meta.get("prefix_len") or 0))
-    # Retirement stash, independent of the anchor outcome: exact-tier
-    # rows retire their full post-decode row at filter (the per-turn
-    # store the post-prefill exact store cannot cover), warm rows
-    # included -- the decode cache holds the full sequence either way.
-    if not _SPEC_APC_RETIRE_DISABLED and batch.prompt_cache:
-        from gmlx.cache.retire_key import lookup_render_ctx
-        ids_list = [int(t) for t in meta["full_input_ids"]]
-        batch.prompt_cache[0]._kq_apc_retire = {
-            "full_ids": ids_list,
-            "extra_hash": int(meta.get("extra_hash", 0)),
-            "mode": "exact",
-            "manager": manager,
-            "render_ctx": lookup_render_ctx(ids_list),
-            "gen": [],
-        }
-
-
-_ANCHOR_PICK_FLAG = "_kq_exact_anchor_pick"
-
-
-def _install_exact_anchor_pick() -> None:
-    """Consult the anchor LRU inside the stock admission pick.
-
-    The pick is where a warm prefix belongs: admission builds the batch
-    from it (suffix rows, right padding, warm-cache merge) and every
-    downstream path treats an anchor restore exactly like a stock exact
-    one. The anchor wins only when strictly longer than the stock pick,
-    so it never shortens a restore. Idempotent.
-    """
-    from mlx_vlm.generate.ar import BatchGenerator
-    if getattr(BatchGenerator._apc_pick_for, _ANCHOR_PICK_FLAG, False):
-        return
-    _orig = BatchGenerator._apc_pick_for
-
-    def _pick_with_anchor(self, sequence):
-        pick = _orig(self, sequence)
-        try:
-            if _SPEC_APC_DISABLED or getattr(self, "apc_mode", None) != "exact":
-                return pick
-            manager = getattr(self, "apc_manager", None)
-            if manager is None or _ckpt_active(
-                    getattr(self, "model", None), "exact",
-                    int(manager.block_size)):
-                return pick
-            _uid, ids_list, _mt, prompt_kwargs, _lps, _crit = sequence
-            if not ids_list or len(ids_list) < 2:
-                return pick
-            # Floor trivial exact picks: a sub-block restore (a bare-BOS
-            # match off an unrelated request) saves nothing but suffix-
-            # constructs the batch, knocking the spec path's ids out of
-            # render space (anchor + retirement keys). Real warm picks are
-            # thousands of tokens and pass untouched.
-            if (pick is not None and not pick.get("matched_blocks")
-                    and 0 < int(pick.get("prefix_len") or 0)
-                    < int(manager.block_size)):
-                pick = None
-            have = int((pick or {}).get("prefix_len") or 0)
-            extra_hash = self._apc_extra_hash(prompt_kwargs or {})
-            floor = max(have, self._apc_safe_prefix_lookup_min(ids_list))
-            from gmlx.cache.snapshot import anchor_exact_lookup
-            warm, ap = anchor_exact_lookup(
-                manager, ids_list, extra_hash=extra_hash,
-                min_prefix_tokens=floor)
-            if warm is None or ap <= have or ap >= len(ids_list):
-                return pick
-            if not self._apc_suffix_is_text_only(ids_list, ap):
-                return pick
-            if pick and pick.get("matched_blocks"):
-                manager.release(pick["matched_blocks"])
-            _log.info("APC L1 hit: prefix=%d suffix=%d tier=anchor",
-                      ap, len(ids_list) - ap)
-            return {
-                "matched_blocks": [],
-                "warm_cache": warm,
-                "prefix_len": ap,
-                "extra_hash": extra_hash,
-                "full_input_ids": list(ids_list),
-            }
-        except Exception:
-            _log.warning("APC anchor pick failed; continuing",
-                         exc_info=True)
-            return pick
-
-    _pick_with_anchor.__dict__[_ANCHOR_PICK_FLAG] = True
-    BatchGenerator._apc_pick_for = _pick_with_anchor
-
-
-_PLAIN_DECODE_FLAG = "_kq_ckpt_plain_decode"
-
-
-def _retire_rows(gb) -> dict:
-    """uid -> retire-stash registry on a generation batch.
-
-    Stashes arm on the B=1 prompt batch's cache object (the only stable
-    home before the decode batch exists); the first decode-side touch
-    lifts them here so they survive ``extend`` rebuilding the cache
-    objects at continuous-batch injection."""
-    reg = getattr(gb, "_kq_apc_retire_rows", None)
-    if reg is None:
-        reg = {}
-        gb._kq_apc_retire_rows = reg
-    return reg
-
-
-def _lift_cache_stash(gb) -> None:
-    if not getattr(gb, "prompt_cache", None) or len(gb.uids) != 1:
-        return
-    stash = getattr(gb.prompt_cache[0], "_kq_apc_retire", None)
-    if stash is not None:
-        gb.prompt_cache[0]._kq_apc_retire = None
-        _retire_rows(gb)[gb.uids[0]] = stash
-
-
-def _plain_step_tick(gb, out) -> None:
-    """Per-token accounting + snapshot tick for stock-path retire rows.
-
-    Rows are tracked per uid so accounting survives ``extend`` merges.
-    Runs per step, so a deterministic failure disables the hook for that
-    row on first strike instead of emitting a traceback per token;
-    dropping ``gen`` also quiets retirement (its offset check would skip
-    anyway on a broken count). The decode-time snapshot ring stays B=1
-    (its clones ride the live single-row caches); rows in a B>1 batch
-    retire snapshot-free, under their verbatim key or an LCP cap the
-    tier arm can serve without a ring."""
-    try:
-        _lift_cache_stash(gb)
-        reg = getattr(gb, "_kq_apc_retire_rows", None)
-    except Exception:
-        _log.warning("APC plain decode hook failed; continuing",
-                     exc_info=True)
-        return
-    if not reg:
-        return
-    # _step returns (tokens, lps, top_idx, top_lp); slot 0 is the flat
-    # per-row token list.
-    rows = out[0] if isinstance(out, tuple) else out
-    if rows is None:
-        return
-    solo = len(gb.uids) == 1
-    for i, uid in enumerate(gb.uids):
-        stash = reg.get(uid)
-        if stash is None or "gen" not in stash:
-            continue
-        tok = rows[i] if i < len(rows) else None
-        if tok is None:
-            continue                    # no emission for this row this tick
-        try:
-            if isinstance(tok, (list, tuple)):
-                tok = tok[0]
-            stash["gen"].append(int(tok))
-            if solo and stash.get("mode") == "ckpt":
-                from gmlx.cache.snapshot import decode_ckpt_tick
-                decode_ckpt_tick(stash, gb.prompt_cache, stash["gen"])
-        except Exception:
-            stash.pop("gen", None)
-            stash["snap_ok"] = False
-            _log.warning("APC plain decode hook failed; disabled for "
-                         "this request", exc_info=True)
-
-
-def _plain_retire(stash: dict, prompt_cache: list) -> None:
-    """Retire a finished stock-path row off a single-row cache list.
-
-    Offset invariants mirror ``speculative._retire_b1``: the stock step
-    loop forwards each token as it emits it, so a clean finish leaves
-    ``offset == len(seq)`` (an abort between steps leaves the same).
-    ``stash["mode"]`` picks the tier arm: "ckpt" stores blocks +
-    sidecar, "exact" a whole-row snapshot (DeepSeek-V4-class pooling
-    stacks).
-    """
-    try:
-        manager = stash.get("manager")
-        if manager is None:
-            return
-        gen = [int(t) for t in stash.get("gen") or ()]
-        if not gen:
-            return
-        seq = [int(t) for t in stash["full_ids"]] + gen
-        from gmlx.cache.snapshot import _cache_offset_max, retirement_store
-        offset = _cache_offset_max(prompt_cache)
-        if offset == len(seq) - 1:
-            seq = seq[:-1]
-        elif offset != len(seq):
-            _log.info(
-                "APC retire skipped: cache offset %d != tokens %d", offset, len(seq)
-            )
-            return
-        lcp = None
-        if os.environ.get("GMLX_APC_RETIRE_LCP") != "0":
-            from gmlx.cache.retire_key import next_turn_lcp
-            lcp = next_turn_lcp(stash.get("render_ctx"), seq, gen)
-        max_len = lcp if lcp is not None and lcp < len(seq) else None
-        _log.info("APC retire: seq=%d ctx=%s lcp=%s cap=%s",
-                  len(seq), stash.get("render_ctx") is not None,
-                  lcp, max_len)
-        ok = retirement_store(
-            manager, stash.get("mode") or "ckpt", seq, prompt_cache,
-            row=0,
-            extra_hash=int(stash.get("extra_hash", 0)), max_len=max_len,
-            decode_snaps=stash.get("snaps"))
-        if ok:
-            _log.info("APC retire store: tokens=%d", ok)
-    except Exception:
-        _log.warning("APC retire failed; continuing", exc_info=True)
-
-
-def _install_plain_ckpt_decode() -> None:
-    """Stock-path decode hooks for the retirement store (ckpt + exact).
-
-    Token accounting rides ``_step``; retirement fires from ``filter``
-    for every leaving row (finish or client abort). A lone row retires
-    off its live single-row caches; a row leaving a B>1 batch is first
-    extracted via ``row_snapshot`` (padding-trimmed clones with row-true
-    offsets), so retirement survives concurrency instead of firing only
-    when the batch happens to drain to one row. Stashes live in a
-    uid-keyed registry lifted across ``extend`` (the seam that rebuilds
-    cache objects at continuous-batch injection).
-    GMLX_APC_RETIRE_BATCH=0 restores the lone-row-only v1 scope.
-    Idempotent."""
-    from mlx_vlm.generate.ar import GenerationBatch
-
-    if getattr(GenerationBatch._step, _PLAIN_DECODE_FLAG, False):
-        return
-    _orig_step = GenerationBatch._step
-    _orig_filter = GenerationBatch.filter
-    _orig_extend = GenerationBatch.extend
-
-    def _step_with_ckpt(self):
-        out = _orig_step(self)
-        _plain_step_tick(self, out)
-        return out
-
-    def _filter_with_ckpt(self, keep):
-        try:
-            _lift_cache_stash(self)
-            reg = getattr(self, "_kq_apc_retire_rows", None)
-            if reg and self.prompt_cache:
-                keep_set = set(keep)
-                solo = len(self.uids) == 1
-                batched_ok = os.environ.get(
-                    "GMLX_APC_RETIRE_BATCH") != "0"
-                for i, uid in enumerate(self.uids):
-                    if i in keep_set:
-                        continue
-                    stash = reg.pop(uid, None)
-                    if stash is None:
-                        continue
-                    if solo:
-                        _plain_retire(stash, self.prompt_cache)
-                    elif batched_ok:
-                        from gmlx.cache.snapshot import row_snapshot
-                        rows = row_snapshot(self.prompt_cache, i)
-                        if rows is None:
-                            _log.info("APC retire skipped: row %d "
-                                      "extract unavailable", i)
-                        else:
-                            _plain_retire(stash, rows)
-        except Exception:
-            _log.warning("APC plain retire hook failed; continuing", exc_info=True)
-        _orig_filter(self, keep)
-
-    def _extend_with_ckpt(self, other):
-        try:
-            _lift_cache_stash(self)
-            _lift_cache_stash(other)
-            other_reg = getattr(other, "_kq_apc_retire_rows", None)
-            if other_reg:
-                _retire_rows(self).update(other_reg)
-                other._kq_apc_retire_rows = {}
-        except Exception:
-            _log.warning("APC retire stash carry failed; continuing",
-                         exc_info=True)
-        _orig_extend(self, other)
-
-    _step_with_ckpt.__dict__[_PLAIN_DECODE_FLAG] = True
-    _filter_with_ckpt.__dict__[_PLAIN_DECODE_FLAG] = True
-    _extend_with_ckpt.__dict__[_PLAIN_DECODE_FLAG] = True
-    GenerationBatch._step = _step_with_ckpt
-    GenerationBatch.filter = _filter_with_ckpt
-    GenerationBatch.extend = _extend_with_ckpt
-
-
-def _mtp_prefill_init(batch) -> None:
-    """One-time APC lookup + prefix trim for an MTP prompt batch.
-
-    Runs on the first ``prompt_step`` call, or directly from ``generate()``
-    when the prompt is short enough that chunked prefill never fires.
-    Lookup ladder: L0 (SpecPrefixCache: whole-prompt KV + full-prompt
-    hidden, the only tier the drafter can teacher-force from without a cold
-    start) then L1 (shared APCManager: exact / block / disk KV, no hidden).
-    Also arms the stock post-prefill store whenever a manager is reachable,
-    regardless of which tier (if any) hit.
-    """
-    if hasattr(batch, "_mtp_full_input_ids"):
-        return
-    batch._mtp_full_input_ids = batch._input_ids
-    batch._mtp_chunk_hiddens = []
-    batch._mtp_l1_prefix_len = 0
-
-    if batch._inputs_embeds is None:
-        _log.info("KQDBG mtp_prefill_init: inputs_embeds None, ladder skipped")
-        return
-
-    # Gated to B=1 because PromptProcessingBatch prefills one request at a
-    # time today. The restored single-row cache (with its offset) later
-    # merges into the live B>1 decode batch via BatchKVCache.extend during
-    # continuous-batch injection -- so APC absolutely works in a B>1
-    # serving context; the gate is about prefill granularity, not decode
-    # batch size. If mlx-vlm ever coalesces prefills into a multi-row
-    # PromptProcessingBatch, this guard silently disables APC for those
-    # rows. The warning below makes that visible.
-    b = int(batch._input_ids.shape[0])
-    if b > 1:
-        if not _SPEC_APC_DISABLED:
-            _log.warning(
-                "APC skipped: prefill batch B=%d > 1 "
-                "(owned-path APC requires single-request prefill)",
-                b,
-            )
-        return
-
-    # Serve wraps make_cache so mlx-lm-origin entries carry the mlx-vlm
-    # runtime's class identities; embedded and test users reach this init
-    # without that wrapper, and the L1 exact tiers dispatch on the vlm
-    # classes (an mlx-lm ArraysCache misses every adapter rule). Rebind
-    # here so both paths see the same identities. No-op when the entries
-    # are already vlm-origin.
-    from gmlx.cache.compat import rebind_to_runtime_origin
-    rebind_to_runtime_origin(batch.prompt_cache)
-
-    # Upstream admission already restored a prefix and built this batch
-    # suffix-only: the owned ladder's keys (L0 and L1 both) are full-prompt
-    # token ids, so every lookup and store here would run in the wrong
-    # space -- a suffix-keyed L0 entry cross-hits a later turn's suffix and
-    # its restore clobbers the upstream warm cache. Leave these batches to
-    # the stock machinery, which owns their meta and store schedule.
-    up_meta = getattr(batch, "_apc_meta", None) or []
-    if up_meta and isinstance(up_meta[0], dict) \
-            and int(up_meta[0].get("prefix_len") or 0) > 0:
-        batch._mtp_upstream_warm = True
-        return
-
-    restored = 0
-    spec_cache = _get_spec_prefix_cache(batch.model)
-    if spec_cache is not None:
-        hit = spec_cache.lookup(batch._input_ids)
-        if hit is not None:
-            restored, entry = hit
-            spec_cache.restore(entry, batch.prompt_cache)
-            batch._mtp_chunk_hiddens = [entry.hidden]
-            _log.info(
-                "APC hit: prefix=%d suffix=%d",
-                restored,
-                int(batch._input_ids.shape[1]) - restored,
-            )
-
-    manager, mode = _resolve_l1(batch.model)
-    if manager is not None:
-        try:
-            l1_prefix = _l1_lookup_and_arm_store(batch, manager, mode, restored)
-            restored = max(restored, l1_prefix)
-        except Exception:
-            _log.warning("APC L1 failed; continuing cold", exc_info=True)
-
-    # Stash the retirement context so the owned B=1 round can store this
-    # request's full context (prompt + generated) into the shared APC when it
-    # finishes. Keyed on the original full ids (pre-trim) -- the serve-layer
-    # prompt_tokens is suffix-only on a warm turn, so it can't be the key.
-    # The stash lives on the request's first cache entry, not on the model:
-    # the server closes a finished rounds generator lazily (sometimes after
-    # the next request's prefill), so a model-level stash races and retires
-    # under the wrong key. Must run after the L1 block above -- an exact-tier
-    # hit replaces batch.prompt_cache wholesale. B=1 only (this init is gated
-    # to B=1); B>1 retirement is handled per-row at the batch decode's
-    # finish seam.
-    if manager is not None and not _SPEC_APC_RETIRE_DISABLED and batch.prompt_cache:
-        meta = (batch._apc_meta or [{}])[0] or {}
-        full_ids = [int(t) for t in batch._mtp_full_input_ids[0].tolist()]
-        from gmlx.cache.retire_key import lookup_render_ctx
-        batch.prompt_cache[0]._kq_apc_retire = {
-            "full_ids": full_ids,
-            "extra_hash": int(meta.get("extra_hash", 0)),
-            "mode": (
-                "ckpt"
-                if _ckpt_active(batch.model, mode, int(manager.block_size))
-                else mode
-            ),
-            "checkpoint_len": int(meta.get("checkpoint_len", 0) or 0),
-            # Live reference: the sidecar keys on ckpt_last_stored, not
-            # the cursor value frozen above.
-            "apc_meta": meta,
-            # Render context for the next-turn LCP key (None off the server
-            # path or on a media prompt; retirement then keys as before).
-            "render_ctx": lookup_render_ctx(full_ids),
-            **_snap_fields(batch, manager),
-        }
-
-    if restored > 0:
-        batch._input_ids = batch._input_ids[:, restored:]
-        batch._inputs_embeds = batch._inputs_embeds[:, restored:]
-        batch._processed_prompt_columns = restored
-        for k in batch._prompt_length_aware_keys:
-            batch._prompt_kwargs[k] = batch._prompt_kwargs[k][:, restored:, ...]
-        batch._mtp_apc_prefix_len = restored
-
-
-def _mtp_seed_stream_init(batch) -> None:
-    """Arm per-chunk drafter seeding for this request, if eligible.
-
-    Cold full prefill only (v1): any restored prefix (L0/L1/upstream) or warm
-    drafter sidecar keeps the deferred one-shot seed -- correctness identical,
-    seeding then still runs after the first token. Eligibility here plus the
-    per-chunk B re-check in prompt_step; a mid-request stop keeps the partial
-    seed KV (adopted at its true offset) and defers only the remainder.
-
-    The seed KV is request-scoped (built via drafter.make_cache, ridden on
-    batch state and handed over via a prompt_cache[0] stash exactly like the
-    drafter warm sidecar), never the drafter's own _cache: another request's
-    live decode round owns that object.
-    """
-    if hasattr(batch, "_mtp_seed_ctx"):
-        return
-    batch._mtp_seed_ctx = None
-    drafter = getattr(batch, "draft_model", None)
-    if (
-        _SEED_STREAM_DISABLED
-        or drafter is None
-        or not callable(getattr(drafter, "seed_chunk", None))
-        or getattr(drafter, "hidden_capture_limit", None) is not None
-        or int(batch._input_ids.shape[0]) != 1
-        or getattr(batch, "_mtp_upstream_warm", False)
-        or getattr(batch, "_mtp_chunk_hiddens", None)
-        or int(getattr(batch, "_mtp_l1_prefix_len", 0) or 0) != 0
-        or int(getattr(batch, "_processed_prompt_columns", 0) or 0) != 0
-        or not batch.prompt_cache
-        or getattr(batch.prompt_cache[0], "_kq_apc_drafter_warm", None)
-            is not None
-    ):
-        return
-    lp = getattr(batch.prompt_cache[0], "left_padding", None)
-    if isinstance(lp, mx.array) and lp.size and int(lp.max().item()) > 0:
-        return
-    try:
-        drafter.bind(batch.model)
-        seed_kv = drafter.make_cache()
-    except Exception:
-        _log.warning("seed streaming unavailable for this drafter; "
-                     "deferred seed", exc_info=True)
-        return
-    ctx = {
-        "kv": seed_kv,
-        "len": 0,
-        "active": True,
-        # Retain chunk hiddens alongside streaming whenever an L0 store can
-        # arm: the store needs full-prompt hidden. APC off => no retention
-        # while streaming (the capture-memory win lands in that config).
-        "retain": _get_spec_prefix_cache(batch.model) is not None,
-        "retained_from": 0,
-    }
-    batch._mtp_seed_ctx = ctx
-    batch.prompt_cache[0]._kq_seed_stream = ctx
-
-
-def _zero_pad_rows(arr, rows: int):
-    pad = mx.zeros((rows - arr.shape[0],) + tuple(arr.shape[1:]), dtype=arr.dtype)
-    return mx.concatenate([arr, pad], axis=0)
-
-
-def _widen_prompt_rope_state(batch, prompt_kwargs: dict) -> dict:
-    """Continuous-batch admission can grow the spec prompt batch (and decode
-    forwards run at other widths) between chunks; the target caches text
-    mrope deltas at the old width and only slices down, never widens, so the
-    next chunk forward dies on offsets(B) + rope_deltas(B_old) broadcast.
-    Text rows have delta 0, so zero-pad both delta sources to the live width
-    (decode-loop twin of this guard: speculative.py injection path)."""
-    b = batch._input_ids.shape[0]
-    rd = prompt_kwargs.get("rope_deltas")
-    if rd is not None and rd.shape[0] < b:
-        prompt_kwargs = dict(prompt_kwargs)
-        prompt_kwargs["rope_deltas"] = _zero_pad_rows(rd, b)
-    lm = getattr(batch.model, "language_model", batch.model)
-    rd = getattr(lm, "_rope_deltas", None)
-    if rd is not None and rd.shape[0] < b:
-        lm._rope_deltas = _zero_pad_rows(rd, b)
-    return prompt_kwargs
-
-
-def install_full_prompt_mtp_prefill() -> None:
-    """Retain full-prompt hidden through the BatchGenerator MTP prefill so the
-    native head teacher-forces the whole prompt into its KV (llama parity).
-
-    mlx-vlm's ``PromptProcessingBatch`` chunks prefill: intermediate chunks
-    (``prompt_step``) discard the model output (only KV-cache side-effects
-    survive), then ``generate()`` runs the final chunk with
-    ``return_hidden=True``.  The MTP drafter thus only sees hidden for that
-    last chunk -- often 1 token -- and acceptance erodes at depth.
-
-    This patch makes ``prompt_step`` also request ``return_hidden=True`` on
-    MTP batches, accumulating per-chunk hidden in ``_mtp_chunk_hiddens``.
-    ``generate()`` then concatenates them with the final chunk's hidden so
-    ``speculative_hidden_state`` returns full-prompt hidden to the drafter.
-
-    Also installs the owned-path APC surface: the L0 SpecPrefixCache
-    (whole-prompt KV + hidden, in-memory) plus the L1 shared APCManager
-    (exact / block / disk tiers -- the same manager the stock
-    non-speculative path uses, reached via ``model._kq_apc_manager``, which
-    ``_install_apc_manager_stash`` captures at BatchGenerator construction).
-    Kill switch for both tiers: ``GMLX_SPEC_APC=0``.
-
-    Idempotent.  Only MTP batches (``self.draft_kind == "mtp"``) are affected;
-    eagle3 / dflash keep the stock path.
-    """
-    from mlx_vlm.generate.ar import PromptProcessingBatch
-
-    # L1 plumbing is idempotent on its own flags, so it installs (or
-    # repairs) even when the prefill override is already in place.
-    _bind_l1_view()
-    # The L1 disk tier serializes through mlx-vlm's DiskBlockStore, which
-    # has no arm for QSAKVCache and refuses the whole exact snapshot.
-    # Installed here as well as in serve patches so embedded/test users of
-    # the spec engine get disk APC.
-    from gmlx.cache.apc_qsa import install_qsa_apc_support
-    install_qsa_apc_support()
-    _install_apc_manager_stash()
-    _install_ckpt_checkpoint_store()
-    _install_plain_ckpt_decode()
-    _install_exact_anchor_pick()
-
-    if getattr(PromptProcessingBatch, _FULL_PREFILL_FLAG, False):
-        return
-
-    _orig_prompt_step = PromptProcessingBatch.prompt_step
-    _orig_generate = PromptProcessingBatch.generate
-    _orig_init = PromptProcessingBatch.__init__
-
-    def _resolve_mtp_prefill_step() -> int:
-        # Honor the serve path's PREFILL_STEP_SIZE env override
-        # (mlx_vlm.server.generation.get_prefill_step_size) so MTP prefill
-        # can be chunked smaller to cap peak memory.
-        from mlx_vlm.generate.ar import DEFAULT_PREFILL_STEP_SIZE
-
-        return int(os.environ.get("PREFILL_STEP_SIZE", DEFAULT_PREFILL_STEP_SIZE))
-
-    def _mtp_init(self, *args, **kwargs) -> None:
-        _orig_init(self, *args, **kwargs)
-        # Re-enable chunked prefill.  Stock mlx-vlm nulls prefill_step_size
-        # for speculative models because intermediate chunks discard hidden;
-        # our prompt_step captures it, so the gate no longer applies.
-        # Restoring at construction (not first prompt_step) matters: the
-        # scheduler consults needs_processing() first, and with a None step
-        # an APC-less deep prompt would one-shot the whole prefill.
-        if (
-            getattr(self, "draft_kind", None) == "mtp"
-            and self.prefill_step_size is None
-        ):
-            self.prefill_step_size = _resolve_mtp_prefill_step()
-        # Stock (non-speculative) batches get the checkpoint tier here:
-        # lookup, prefix trim, cursor arming, retirement stash.
-        if getattr(self, "draft_kind", None) is None and not _SPEC_APC_DISABLED:
-            try:
-                # Ckpt-active hybrids under kvarn convert their stock B=1
-                # single-stream caches in place before arming, so the
-                # layout signature, lookup, and every store see the same
-                # kvarn classes. Must run here: the outer batch rebuild
-                # (kvarn_serve) would install batch classes the tier is
-                # blind to; after conversion its shared decline predicate
-                # trips instead. kwargs is load-bearing -- stock init
-                # consumes and drops the scheme/bits constructor params.
-                manager, mode = _resolve_l1(self.model)
-                if manager is not None:
-                    from gmlx.cache.kvarn_serve import ensure_ppb_kvarn
-
-                    ensure_ppb_kvarn(
-                        self, kwargs,
-                        ckpt_active=_ckpt_active(
-                            self.model, mode, int(manager.block_size)))
-            except Exception:
-                _log.warning(
-                    "kvarn ckpt cache conversion failed; continuing stock",
-                    exc_info=True,
-                )
-            try:
-                _plain_ckpt_init(self)
-                _plain_anchor_init(self)
-            except Exception:
-                _log.warning(
-                    "APC plain ckpt init failed; continuing stock", exc_info=True
-                )
-
-    def _mtp_prompt_step(self) -> int:
-        if self.draft_kind != "mtp":
-            return _orig_prompt_step(self)
-        # cb_phase flips fine prefill caps by wrapping the stock
-        # prompt_step, but this body replaces it for MTP batches, so the
-        # flip must happen here too: a multi-thousand-token chunk under
-        # the coarse decode caps keeps every layer's transients live in
-        # one command buffer and OOMs the GPU on deep prompts.
-        if os.environ.get("GMLX_CB_PHASE", "1") != "0":
-            from gmlx.serve.cb_phase import flip
-            flip("prefill")
-
-        if not hasattr(self, "_mtp_full_input_ids"):
-            if self.prefill_step_size is None:
-                self.prefill_step_size = _resolve_mtp_prefill_step()
-            # APC lookup (L0 then L1) + prefix trim + store arming.
-            _mtp_prefill_init(self)
-            _mtp_seed_stream_init(self)
-
-        if not self.needs_processing():
-            return 0
-
-        # Depth-decayed step: shrink only when this chunk's score transient
-        # would exceed the cap (see prefill_decay; keeps MoE weight
-        # amortization at shallow depth instead of a global small step).
-        step = prefill_decay.decayed_for_batch(self) or self._inputs_embeds.shape[1]
-        n = min(step, self._inputs_embeds.shape[1] - 1)
-
-        if not hasattr(self, "_mtp_padding_widened"):
-            self._mtp_padding_widened = True
-            for c in self.prompt_cache:
-                lp = getattr(c, "left_padding", None)
-                if isinstance(lp, mx.array) and lp.ndim > 0 and lp.size > 1:
-                    max_lp = int(lp.max().item())
-                    if max_lp >= n:
-                        n = min(max_lp + 1, self._inputs_embeds.shape[1] - 1)
-                    break
-
-        checkpoint_col = self._next_apc_checkpoint_column()
-        if checkpoint_col is not None:
-            n = min(n, checkpoint_col - self._processed_prompt_columns)
-        # Media requests ride this body too: keep image blocks whole (a
-        # boundary inside a block moves to its edge, see media_spans).
-        from gmlx.gen.media_spans import span_aware_prompt_n
-        n = span_aware_prompt_n(self, n)
-        # A final chunk under ~3 simdgroup tiles routes the projections
-        # through the skinny-M kernels, whose accumulation order seeds fp
-        # noise that stacked recurrent (GDN) layers amplify into
-        # first-token divergence. Absorb such a tail into this chunk so
-        # every chunk stays in the wide-GEMM regime. Checkpoint columns
-        # stay exact.
-        min_tail = env_int("GMLX_PREFILL_MIN_TAIL", 48)
-        if checkpoint_col is None and min_tail > 0:
-            rem1 = self._inputs_embeds.shape[1] - 1
-            tail = rem1 - n
-            if 0 < tail < min_tail:
-                n = rem1        # absorb: overshoot bounded by min_tail-1
-        if n <= 0:
-            return 0
-        prompt_kwargs = self._prompt_kwargs_for_step(n)
-        prompt_kwargs = _widen_prompt_rope_state(self, prompt_kwargs)
-        with lora_rows.published(getattr(self, "uids", [])):
-            out = self.model(
-                self._input_ids[:, :n],
-                cache=self.prompt_cache,
-                inputs_embeds=self._inputs_embeds[:, :n],
-                n_to_process=n,
-                return_hidden=True,
-                **prompt_kwargs,
-            )
-        chunk_hidden = out.hidden_states[-1]
-        # Seed streaming: teacher-force this chunk into the request-scoped
-        # head KV at the head's running offset. The shifted span for
-        # columns [c0, c0+n) is prompt[c0+1 : c0+n+1], always in range
-        # because generate() keeps at least one residual column (the n-1
-        # cap above). A failure or a widened batch stops streaming but
-        # keeps the partial KV: the owned round adopts it at its true
-        # offset and seeds only the remainder.
-        seed_ctx = getattr(self, "_mtp_seed_ctx", None)
-        streamed = False
-        if seed_ctx is not None and seed_ctx["active"]:
-            if int(self._input_ids.shape[0]) != 1:
-                seed_ctx["active"] = False
-            else:
-                c0 = int(self._processed_prompt_columns)
-                try:
-                    self.draft_model.seed_chunk(
-                        self._mtp_full_input_ids[:, c0 + 1:c0 + n + 1],
-                        chunk_hidden, seed_ctx["kv"])
-                    seed_ctx["len"] += n
-                    streamed = True
-                except Exception:
-                    _log.warning("seed streaming failed at column %d; "
-                                 "deferred seed for the remainder", c0,
-                                 exc_info=True)
-                    seed_ctx["active"] = False
-        # Teacher-forcing drafters (native MTP heads) seed their KV from the
-        # whole prompt hidden, so every chunk is retained except when the
-        # chunk just streamed and no L0 store is armed (nothing downstream
-        # reads it). Shared-KV drafters (gemma-4 assistant) read only the
-        # last position: keeping just the newest chunk caps capture memory
-        # at O(chunk) instead of O(prompt), GBs at deep context.
-        if callable(getattr(self.draft_model, "prefill_from_target_hidden", None)):
-            if streamed and not seed_ctx["retain"]:
-                pass
-            else:
-                if (seed_ctx is not None and not seed_ctx["retain"]
-                        and not self._mtp_chunk_hiddens):
-                    # Streaming stopped mid-request with no retention so
-                    # far: the retained span starts here, not at column 0.
-                    seed_ctx["retained_from"] = int(
-                        self._processed_prompt_columns)
-                self._mtp_chunk_hiddens.append(chunk_hidden)
-                # Window-limited heads can't use context beyond the trailing
-                # hidden_capture_limit positions; an uncapped capture pins the
-                # whole prompt's hidden (GBs at deep context). The drafter's
-                # teacher-force self-aligns to the trailing h_len positions.
-                limit = getattr(self.draft_model, "hidden_capture_limit", None)
-                if limit:
-                    total = sum(int(h.shape[1]) for h in self._mtp_chunk_hiddens)
-                    if total > limit:
-                        merged = (self._mtp_chunk_hiddens[0]
-                                  if len(self._mtp_chunk_hiddens) == 1
-                                  else mx.concatenate(self._mtp_chunk_hiddens, axis=1))
-                        self._mtp_chunk_hiddens = [merged[:, -limit:]]
-        else:
-            self._mtp_chunk_hiddens = [chunk_hidden]
-        mx.eval([c.state for c in self.prompt_cache] + [chunk_hidden]
-                + ([c.state for c in seed_ctx["kv"]] if streamed else []))
-        self._processed_prompt_columns += n
-        # The ckpt cursor rides the wrapped stock store (see
-        # _install_ckpt_checkpoint_store).
-        self._store_apc_exact_checkpoints()
-        self._inputs_embeds = self._inputs_embeds[:, n:]
-        self._input_ids = self._input_ids[:, n:]
-        for k in self._prompt_length_aware_keys:
-            self._prompt_kwargs[k] = self._prompt_kwargs[k][:, n:, ...]
-        mx.clear_cache()
-        return n
-
-    def _mtp_generate(
-        self, sampler, stop_criteria, compute_logprobs=True, top_logprobs_k=0
-    ):
-        if self.draft_kind == "mtp":
-            # Short prompts never enter prompt_step (chunked prefill is not
-            # needed), so the APC lookup/store arming runs here instead.
-            _mtp_prefill_init(self)
-        result = _orig_generate(
-            self,
-            sampler,
-            stop_criteria,
-            compute_logprobs=compute_logprobs,
-            top_logprobs_k=top_logprobs_k,
-        )
-        from mlx_vlm.generate.ar import SpeculativeGenerationBatch
-
-        if self.draft_kind != "mtp" or not isinstance(
-            result, SpeculativeGenerationBatch
-        ):
-            # Stock-path ckpt batches store the full prompt here, the
-            # moment the MTP path stores it at rounds entry: prefill just
-            # finished, the first token is out, its KV not yet appended.
-            if (
-                getattr(self, "_kq_ckpt_armed", False)
-                and getattr(self, "draft_kind", None) is None
-            ):
-                try:
-                    cache = getattr(result, "prompt_cache", None) or []
-                    stash = getattr(cache[0], "_kq_apc_retire", None) if cache else None
-                    if stash is not None and stash.get("mode") == "ckpt":
-                        from gmlx.cache.snapshot import (
-                            ckpt_full_store_redundant,
-                            ckpt_store,
-                        )
-                        m = stash.get("apc_meta")
-                        if ckpt_full_store_redundant(m):
-                            _log.info("APC ckpt post-prefill store "
-                                      "skipped: render-stable boundary "
-                                      "landed")
-                        elif ckpt_store(
-                                stash["manager"], stash["full_ids"], cache,
-                                extra_hash=int(stash.get("extra_hash", 0))):
-                            if m is not None:
-                                m.setdefault(
-                                    "ckpt_stored_boundaries", []
-                                ).append(len(stash["full_ids"]))
-                except Exception:
-                    _log.warning(
-                        "APC plain post-prefill store failed; continuing", exc_info=True
-                    )
-            return result
-        chunk_hiddens = getattr(self, "_mtp_chunk_hiddens", None)
-        full_ids = getattr(self, "_mtp_full_input_ids", None)
-        l1_prefix = int(getattr(self, "_mtp_l1_prefix_len", 0) or 0)
-        if not chunk_hiddens:
-            # No captured chunks: the whole (remaining) prompt went through
-            # the final generate forward, so stock prompt_tokens/hidden are
-            # already an aligned pair (suffix-only on an L1 hit) and
-            # result.hidden needs no rebuild; with seed streaming and no
-            # retention, result.hidden is already the residual unstreamed
-            # tail (retention accompanies an armed L0 store, so none can
-            # fire here). The L0 store below must still run for the
-            # single-shot case: arch prefill profiles can raise the step
-            # past typical prompt lengths (qwen4exp defaults to 8192), so
-            # sub-step prompts land here and still need their warm-start
-            # entry.
-            full_hidden = result.hidden
-        else:
-            parts = chunk_hiddens + [result.hidden]
-            full_hidden = mx.concatenate(parts, axis=1)
-        seed_ctx = getattr(self, "_mtp_seed_ctx", None)
-        seed_len = int(seed_ctx["len"]) if seed_ctx else 0
-        if chunk_hiddens:
-            if seed_len > 0:
-                # Columns [0, seed_len) are already teacher-forced into
-                # the streamed head KV; hand the owned round only the
-                # residual hidden so its seed call covers exactly the
-                # unstreamed tail at the adopted offset. full_hidden (the
-                # retained span) still feeds the L0 store below, which
-                # needs the whole prompt.
-                rfrom = int(seed_ctx.get("retained_from") or 0)
-                result.hidden = full_hidden[:, seed_len - rfrom:]
-            else:
-                result.hidden = full_hidden
-        if chunk_hiddens and full_ids is not None:
-            # On an L1 hit the captured hidden covers only the forwarded
-            # suffix, so hand the drafter the matching suffix tokens: the
-            # teacher-forcing (token, hidden) pair must stay positionally
-            # aligned. The missing prefix can only affect draft acceptance,
-            # never correctness -- verify catches every draft.
-            result.prompt_tokens = (
-                full_ids[:, l1_prefix:] if l1_prefix > 0 else full_ids
-            )
-
-        # APC L0 store: cache this request's target KV + hidden so a
-        # future request sharing this token prefix skips re-prefill.
-        # Uses result.prompt_cache (SpecBatch owns the cache now),
-        # not self.prompt_cache (empty after _orig_generate).
-        #
-        # B=1 only -- same prefill-granularity gate as the lookup.
-        # The stored single-row snapshot is valid for injection into
-        # a B>1 batch: SpecPrefixCache.restore writes into a fresh
-        # single-row prompt_cache, and BatchKVCache.extend merges
-        # it at the correct per-row offset.
-        #
-        # Skipped on an L1 hit: hidden covers only the suffix, and L0
-        # entries pair full-prompt keys with full-prompt hidden.
-        b = int(full_hidden.shape[0]) if full_ids is not None else 0
-        # With streaming, full_hidden covers the whole prompt only when
-        # retention ran from column 0: after a mid-request streaming stop
-        # the retained span starts past column 0, and with no retention at
-        # all full_hidden is just the residual tail. Neither must ever be
-        # stored as a full-prompt entry.
-        full_covers_prompt = seed_len == 0 or (
-            bool(chunk_hiddens)
-            and int(seed_ctx.get("retained_from") or 0) == 0)
-        spec_cache = (
-            _get_spec_prefix_cache(self.model)
-            if b == 1 and l1_prefix == 0 and full_covers_prompt
-            and not getattr(self, "_mtp_upstream_warm", False) else None
-        )
-        if spec_cache is not None and full_ids is not None:
-            # Window-limited heads only use the trailing capture window;
-            # chunked prefill already trimmed, single-shot must match (an
-            # uncapped entry pins the whole prompt's hidden for nothing).
-            limit = getattr(self.draft_model, "hidden_capture_limit", None)
-            store_hidden = (full_hidden if not limit
-                            else full_hidden[:, -int(limit):])
-            spec_cache.store(full_ids, result.prompt_cache, store_hidden)
-            _log.info(
-                "APC store: tokens=%d layers=%d",
-                int(full_ids.shape[1]),
-                len(result.prompt_cache),
-            )
-        else:
-            _log.debug(
-                "APC store skipped: b=%d l1_prefix=%d upstream_warm=%s "
-                "full_ids=%s",
-                b, l1_prefix,
-                getattr(self, "_mtp_upstream_warm", False),
-                "set" if full_ids is not None else "None",
-            )
-
-        return result
-
-    PromptProcessingBatch.__init__ = _mtp_init
-    PromptProcessingBatch.prompt_step = _mtp_prompt_step
-    PromptProcessingBatch.generate = _mtp_generate
-    setattr(PromptProcessingBatch, _FULL_PREFILL_FLAG, True)
-    if _SPEC_APC_DISABLED:
-        apc_status = "off"
-    elif _L1_BOUND[0]:
-        apc_status = "on: L0+L1"
-    else:
-        apc_status = "on: L0 only"
-    _debug_note(
-        f"[mtp] serve prefill: full-prompt hidden capture installed (APC {apc_status})"
-    )
-
-
-_CONTINUOUS_BATCH_FLAG = "_kq_gguf_continuous_batch"
-_RELEASED_FLAG = "_kq_gguf_spec_released"
-_RELEASE_PENDING_FLAG = "_kq_gguf_spec_release_pending"
-
-
-def dequantize_lift_cache(c):
-    """Dequantize a B=1 QuantizedKVCache into a one-row BatchKVCache.
-
-    QuantizedKVCache has no merge, and the B>1 MTP arm runs fp16 KV
-    anyway, so the lift performs the same conversion the batch-build
-    swap does, at preemption or injection time, as a one-off O(depth)
-    copy. Without this every kv-bits MTP preemption declined into a
-    drain-wait and queued rows stalled behind the live generation."""
-    from mlx_vlm.models.cache import BatchKVCache
-
-    lifted = BatchKVCache([0])
-    L = c.offset
-    if L:
-        keys = mx.dequantize(
-            *(mx.contiguous(t[..., :L, :]) for t in c.keys),
-            group_size=c.group_size, bits=c.bits)
-        values = mx.dequantize(
-            *(mx.contiguous(t[..., :L, :]) for t in c.values),
-            group_size=c.group_size, bits=c.bits)
-        lifted.update_and_fetch(keys, values)
-    stamp = getattr(c, "_gmlx_cascade", None)
-    if stamp is not None:
-        lifted._gmlx_cascade = stamp
-    return lifted
-
-
-def kvarn_lift_cache(c):
-    """Recover a B=1 KVarNKVCache into a one-row fp16 BatchKVCache, the
-    lift for an mlx-kquant without per-row ends (the batched arm then
-    runs fp16 KV).
-
-    The kvarn twin of dequantize_lift_cache: materialize() returns
-    rotated-domain K/V, which stock SDPA would attend with an un-rotated
-    query -- no crash, just wrong logits on every preempted row.
-    _raw_single is the original-domain accessor."""
-    from mlx_vlm.models.cache import BatchKVCache
-
-    lifted = BatchKVCache([0])
-    if c.offset:
-        keys, values = c._raw_single()
-        lifted.update_and_fetch(keys, values)
-    stamp = getattr(c, "_gmlx_cascade", None)
-    if stamp is not None:
-        lifted._gmlx_cascade = stamp
-    return lifted
-
-
-def mtp_kv_decline(lm, *, owned_round: bool = True) -> str | None:
-    """Why this MTP verify walk cannot run on packed KV, or None.
-
-    The owned rounds roll back with trim, and affine packing is
-    per-token along head_dim, so a trim is an offset move: they take the
-    same layers serve takes. The two stock walks slice keys as raw
-    arrays and cannot read a packed tuple back. Shared by serve, run and
-    chat so the three cannot drift.
-    """
-    if not owned_round:
-        return "GMLX_OWNED_ROUND=0 stock rounds have no KV quantization hook"
-    from gmlx.models.qwen35.gdn import stock_gdn_fallback
-
-    mt = None
-    for h in _spec_target_holders(lm):
-        cfg = getattr(h, "config", None)
-        mt = (getattr(h, "model_type", None)
-              or (cfg.get("model_type") if isinstance(cfg, dict)
-                  else getattr(cfg, "model_type", None)))
-        if mt:
-            break
-    if stock_gdn_fallback(mt):
-        return ("the GMLX_QWEN_OWNED=0 stock fallback cannot verify on a "
-                "quantized KV cache")
-    return None
-
-
-def lift_single_cache(c):
-    """Promote a single-sequence cache to its batch class. An affine B=1
-    cache recovers to a one-row fp16 BatchKVCache (the batched MTP arm
-    runs fp16 KV under uniform). A kvarn B=1 cache becomes a one-row
-    BatchKVarNKVCache, buffers and horizon adopted bit-exactly, when the
-    installed mlx-kquant takes per-row ends, and recovers to fp16 rows
-    otherwise. Everything else lifts through its class's merge. The
-    cascade stamp rides along. One lift for the preempted host row and
-    the injected rows, so the two cannot drift."""
-    from gmlx.cache.compat import cache_types
-
-    if isinstance(c, cache_types("QuantizedKVCache")):
-        return dequantize_lift_cache(c)
-    if getattr(c, "kv_quant_scheme", None) == "kvarn" and batch_liftable(c):
-        from gmlx.cache.kvarn_sdpa import kvarn_row_ends_ok
-
-        if not kvarn_row_ends_ok():
-            return kvarn_lift_cache(c)
-        from gmlx.cache.kvarn_cache import BatchKVarNKVCache
-
-        lifted = BatchKVarNKVCache.merge([c])
-        stamp = getattr(c, "_gmlx_cascade", None)
-        if stamp is not None:
-            lifted._gmlx_cascade = stamp
-        return lifted
-    lifted = type(c).merge([c])
-    stamp = getattr(c, "_gmlx_cascade", None)
-    if stamp is not None:
-        lifted._gmlx_cascade = stamp
-    return lifted
-
-
-def batch_liftable(c) -> bool:
-    """Whether _lift_host_cache can promote this cache to a batch class.
-
-    The preemption and pre-start-compaction gates must agree with it: a
-    scheme it can lift but they refuse declines into the drain-wait, and
-    one they admit but it cannot lift raises mid-rebuild.
-    """
-    if hasattr(c, "filter") and hasattr(c, "extend"):
-        return True
-    if getattr(c, "kv_quant_scheme", None) == "kvarn":
-        from gmlx.cache.kvarn_cache import KVarNRotatingKVCache
-
-        # The rotating subclass counts evicted tokens in offset that its
-        # buffers no longer hold: kvarn_lift_cache would misplace rows.
-        return not isinstance(c, KVarNRotatingKVCache)
-    from gmlx.cache.compat import cache_types
-
-    return (isinstance(c, cache_types("QuantizedKVCache"))
-            or hasattr(type(c), "merge"))
-
-
-def install_continuous_batch_admission() -> None:
-    """Let new requests prefill and inject during speculative decode.
-
-    Without this, mlx-vlm's ``is_speculative`` gate blocks all prefills while
-    speculative decode is in-flight, and ``extend()`` raises on non-empty
-    speculative batches. This installs five patches:
-
-    1. Disables the ``is_speculative`` admission gate (lets prefills run
-       during decode).
-    2. Overrides ``extend()`` to buffer new batches instead of raising.
-    3. Overrides ``__len__()`` to auto-promote buffered batches when the
-       current batch finishes.
-    4. Overrides ``next()`` to process pending injections - updates outer
-       tracking state, emits first tokens, queues for the generator.
-    5. Releases a finished batch's request state (target KV, captured
-       hidden, shared KV, drafter KV) the moment its last row finishes.
-
-    The generator-side injection (extending caches + drafter mid-flight)
-    happens in ``_owned_decode_rounds_batch`` via ``model._generator_injections``.
-    """
-    from mlx_vlm.generate import ar as _ar
-
-    SpecBatch = _ar.SpeculativeGenerationBatch
-    if getattr(SpecBatch, _CONTINUOUS_BATCH_FLAG, False):
-        return
-
-    # 1. Remove admission gate
-    SpecBatch.is_speculative = False
-
-    _orig_len = SpecBatch.__len__
-
-    # 5. Release request state at finish. BatchGenerator parks the finished
-    # batch in _generation_batch until the next request's prefill completes
-    # (only PromptProcessingBatch.generate's extend replaces it), so every
-    # heavy attr -- the full target KV, the captured full-prompt hidden, the
-    # prefill shared-KV, the rounds generator (whose delegation frame re-pins
-    # all of the above), and the drafter's own head KV -- survives that whole
-    # prefill window. At deep context that stacks two requests' footprints
-    # for many minutes (d200k gemma-4-31b: ~65 GB across an ~18-minute
-    # prefill) and runs the box to the wire ceiling. Drop it all on the
-    # finishing step instead.
-    def _release_heavy_state(self) -> bool:
-        """Drop request state from a finished batch. Returns False when the
-        rounds generator is mid-step on another thread (a client abort racing
-        the engine); ``__len__`` retries on the engine thread."""
-        if getattr(self, _RELEASED_FLAG, False):
-            return True
-        rounds = getattr(self, "_rounds_iter", None)
-        if rounds is not None:
-            try:
-                # Terminal-token finishes already ran the inner loop's own
-                # cleanup; close() is then a no-op resume. Aborted requests
-                # close here, firing the mid-round rollback + retirement.
-                rounds.close()
-            except ValueError:
-                setattr(self, _RELEASE_PENDING_FLAG, True)
-                return False
-            except Exception:
-                _log.warning("spec batch release: rounds close failed", exc_info=True)
-        self._rounds_iter = None
-        self.prompt_cache = []
-        self.hidden = None
-        self.shared_kv_states = None
-        self.prompt_tokens = None
-        self.first_tokens = None
-        if getattr(self, "draft_kind", None) == "mtp":
-            drafter = getattr(self, "draft_model", None)
-            model = getattr(self, "model", None)
-            if drafter is not None and model is not None:
-                try:
-                    drafter.reset(model)  # drops the head's request KV
-                except Exception:
-                    _log.warning(
-                        "spec batch release: drafter reset failed", exc_info=True
-                    )
-        setattr(self, _RELEASED_FLAG, True)
-        setattr(self, _RELEASE_PENDING_FLAG, False)
-        mx.clear_cache()
-        return True
-
-    def _release_if_finished(self) -> None:
-        if _orig_len(self) == 0:
-            _release_heavy_state(self)
-            return
-        _shed_finished_attr_rows(self)
-
-    def _shed_finished_attr_rows(self) -> None:
-        """Per-row release of the batch-held start-time snapshots.
-
-        The live rounds generator sheds a finished or filtered row's KV,
-        drafter state, and its own hidden/shared_kv slices at the next
-        round boundary; the batch object's prefill-time copies (hidden,
-        shared_kv_states, prompt_tokens, first_tokens) stayed resident
-        until the whole batch finished. Slice them by the surviving rows
-        instead. Injected rows carry no snapshot here (their state rides
-        the injection queue into the generator), so the snapshot covers
-        the first first_tokens.shape[0] physical rows only. Slices are
-        lazy and ride the tick's eval; nothing here forces a sync.
-
-        Runs only once the rounds generator holds the state: pre-start,
-        _start_rounds still needs the snapshots row-aligned with the
-        caches (finished rows included; the generator stop_checks them
-        out itself), so a first-token finish must not slice here."""
-        if self._rounds_iter is None:
-            return
-        ft = getattr(self, "first_tokens", None)
-        if ft is None or getattr(self, _RELEASED_FLAG, False):
-            return
-        rows = getattr(self, "_kq_attr_rows", None)
-        if rows is None:
-            try:
-                rows = self._kq_attr_rows = list(range(ft.shape[0]))
-            except Exception:
-                return
-        keep = [p for p in rows
-                if p < len(self._finished) and not self._finished[p]]
-        if len(keep) == len(rows):
-            return
-        if not keep:
-            self.hidden = None
-            self.shared_kv_states = None
-            self.prompt_tokens = None
-            self.first_tokens = None
-            self._kq_attr_rows = []
-            return
-        keep_set = set(keep)
-        pos = [i for i, p in enumerate(rows) if p in keep_set]
-        idx = mx.array(pos, dtype=mx.int32)
-        for name in ("hidden", "prompt_tokens", "first_tokens"):
-            arr = getattr(self, name, None)
-            if arr is not None:
-                setattr(self, name, arr[idx])
-        kv = getattr(self, "shared_kv_states", None)
-        if isinstance(kv, dict) and kv:
-            # New dict, new arrays: the generator may still hold (and
-            # slice) the originals; never mutate a possibly shared dict.
-            self.shared_kv_states = {
-                k: (K[idx], V[idx]) for k, (K, V) in kv.items()}
-        self._kq_attr_rows = keep
-
-    # 2. Buffer extend() instead of raising
-    def _buffered_extend(self, other):
-        active = sum(not d for d in self._finished)
-        if active == 0:
-            pending = getattr(self, "_pending_injections", [])
-            self.__dict__.pop("_kq_attr_rows", None)
-            self.__dict__.update(other.__dict__)
-            self._pending_injections = pending
-            setattr(self, _RELEASED_FLAG, False)
-            setattr(self, _RELEASE_PENDING_FLAG, False)
-            return
-        if not hasattr(self, "_pending_injections"):
-            self._pending_injections = []
-        self._pending_injections.append(other)
-        _debug_note(f"[mtp] extend buffered: +{len(other._all_uids)} rows "
-                    f"(pending={len(self._pending_injections)}, "
-                    f"active={active})")
-
-    SpecBatch.extend = _buffered_extend
-
-    # 3. Auto-promote buffered batches when current is done
-    def _len_with_promotion(self):
-        if getattr(self, _RELEASE_PENDING_FLAG, False) and _orig_len(self) == 0:
-            _release_heavy_state(self)
-        active = _orig_len(self)
-        if active == 0:
-            pending = getattr(self, "_pending_injections", None)
-            if pending:
-                other = pending.pop(0)
-                remaining = pending[:]
-                self.__dict__.pop("_kq_attr_rows", None)
-                self.__dict__.update(other.__dict__)
-                self._pending_injections = remaining
-                setattr(self, _RELEASED_FLAG, False)
-                setattr(self, _RELEASE_PENDING_FLAG, False)
-                return _orig_len(self)
-        return active
-
-    SpecBatch.__len__ = _len_with_promotion
-
-    _orig_filter = SpecBatch.filter
-
-    def _compact_prestart_rows(self, keep) -> None:
-        """Physically drop rows from a batch whose rounds generator has
-        not started: filter the caches through their own filter (lifting
-        host caches first) and slice snapshots plus bookkeeping to the
-        same keep list. Pre-start, the batch object owns all state, so
-        the drop frees the rows' bytes immediately instead of marking
-        them finished and waiting for a generator that has no round
-        boundary yet."""
-        idx = mx.array(keep, dtype=mx.int32)
-        self.prompt_cache = [_lift_host_cache(c) for c in self.prompt_cache]
-        for c in self.prompt_cache:
-            c.filter(idx)
-        for name in ("hidden", "prompt_tokens", "first_tokens"):
-            arr = getattr(self, name, None)
-            if arr is not None:
-                setattr(self, name, arr[idx])
-        kv = getattr(self, "shared_kv_states", None)
-        if isinstance(kv, dict) and kv:
-            self.shared_kv_states = {
-                k: (K[idx], V[idx]) for k, (K, V) in kv.items()}
-        self._all_uids = [self._all_uids[i] for i in keep]
-        self.uids = list(self._all_uids)
-        self.max_tokens = [self.max_tokens[i] for i in keep]
-        self._num_tokens = [self._num_tokens[i] for i in keep]
-        self._finished = [False] * len(keep)
-        self.__dict__.pop("_kq_attr_rows", None)
-
-    def _filter_with_release(self, keep):
-        # Pre-start strict subset (a cancel or a governor retire landing
-        # before the first tick): compact physically. Live or degenerate
-        # cases keep the upstream mark-finished contract; the running
-        # generator sheds the row at its next round boundary and the
-        # snapshot shed below covers the batch-held copies.
-        if (len(keep) < len(self.uids)
-                and keep
-                and self._rounds_iter is None
-                and not getattr(self, _RELEASED_FLAG, False)
-                and getattr(self, "first_tokens", None) is not None
-                and self.uids == self._all_uids
-                and not any(self._finished)
-                and all(batch_liftable(c) for c in self.prompt_cache)):
-            _compact_prestart_rows(self, list(keep))
-            return
-        _orig_filter(self, keep)
-        _release_if_finished(self)
-
-    SpecBatch.filter = _filter_with_release
-
-    # 4. Process pending injections in next() before advancing the generator
-    _orig_next = SpecBatch.next
-
-    def _note_last_tokens(self, responses) -> None:
-        # Last delivered token per uid: the bonus a preempt rebuild restarts
-        # from (its KV is not yet in the cache at a round boundary).
-        stash = getattr(self, "_kq_last_tokens", None)
-        if stash is None:
-            stash = self._kq_last_tokens = {}
-        for r in responses:
-            if r.token is not None:
-                stash[r.uid] = int(r.token)
-
-    def _lift_host_cache(c):
-        """Promote a single-sequence host cache to its batch class so the
-        rebuilt batch generator can extend/filter it (same lift the
-        injection path applies to incoming caches)."""
-        if hasattr(c, "filter") and hasattr(c, "extend"):
-            return c
-        return lift_single_cache(c)
-
-    def _preempt_scalar(self) -> bool:
-        """Preempt a live scalar (B=1) spec generation so queued rows can
-        join: close the generator, deliver the closed round's undelivered
-        tail (the scalar path yields one token per next(), so a close
-        usually lands mid-round; those tokens are verified and their KV
-        stays in the cache), lift the caches to batch classes, and mark
-        the batch armless (hidden=None); _start_rounds then rebuilds it on
-        the batch loop, whose first injection drain admits the waiters.
-        The rebuild resumes from the round's bonus token, whose KV is not
-        in the cache. GMLX_MTP_PREEMPT=0 leaves the old drain-wait
-        behavior.
-
-        The rebuilt row carries no APC retirement context (batch-loop rows
-        start with retire_ctxs None), so the preempted request's prefix is
-        not offered back to the prompt cache when it finishes."""
-        if not env_bool("GMLX_MTP_PREEMPT", True):
-            return False
-        if not getattr(self, "_sent_first", False):
-            return False
-        last = getattr(self, "_kq_last_tokens", {}).get(self._all_uids[0])
-        if last is None:
-            return False
-        # Every cache must be batch-liftable before the generator
-        # closes. A quantized or kvarn B=1 cache lifts to fp16. Anything
-        # else unliftable declines into the drain-wait.
-        if not all(batch_liftable(c) for c in self.prompt_cache):
-            return False
-        it = self._rounds_iter
-        captured = []
-        if it is not None:
-            self._rounds_iter = None
-            self.model._kq_preempt_capture = captured
-            try:
-                it.close()
-            finally:
-                try:
-                    del self.model._kq_preempt_capture
-                except AttributeError:
-                    pass
-        responses = []
-        uid = self._all_uids[0]
-        for tok in captured:
-            if self._finished[0]:
-                break
-            tok = int(tok)
-            self._num_tokens[0] += 1
-            finish = self._finish_reason(0, tok)
-            if finish is not None:
-                self._finished[0] = True
-            responses.append(self.Response(
-                uid=uid, token=tok, token_logprob=0.0, finish_reason=finish))
-            last = tok
-        self._kq_preempt_responses = responses
-        if self._finished[0]:
-            # The captured tail finished the row; nothing to rebuild. The
-            # pending injections promote through __len__ once drained.
-            self._refresh_uids()
-            return False
-        self.prompt_cache = [_lift_host_cache(c) for c in self.prompt_cache]
-        self.first_tokens = mx.array([int(last)], dtype=self.token_dtype)
-        self.hidden = None
-        self.shared_kv_states = None
-        self.prompt_tokens = None
-        self.model._kq_rebuild_emitted = [int(self._num_tokens[0])]
-        _debug_note("[mtp] preempt: scalar generation rebuilt for "
-                    "continuous batching")
-        return True
-
-    def _next_with_injection(self):
-        pending = getattr(self, "_pending_injections", None)
-        # Physical-row uids for the owned rounds loop (it has no batch
-        # object): read once at generator start, injected rows carry theirs.
-        try:
-            self.model._kq_row_uids = list(self._all_uids)
-        except AttributeError:      # attribute-less model stand-ins
-            pass
-        # Mid-flight adoption works only when the batch rounds generator is
-        # running: it drains model._generator_injections at its round
-        # boundaries. The scalar (B=1) generator never does, so a live
-        # scalar host is preempted first: its generator closes at the round
-        # boundary and the batch is rebuilt armless on the batch loop.
-        # `_all_uids` is an mlx-vlm generator internal (stable under the
-        # ==0.6.3 pin); re-verify this batch-vs-scalar signal on a pin lift.
-        preempted = False
-        if pending and len(self._all_uids) == 1:
-            preempted = _preempt_scalar(self)
-        # The preempt capture: verified tokens the closed round had not yet
-        # delivered. They precede everything this call returns.
-        pre_responses = self.__dict__.pop("_kq_preempt_responses", None) or []
-        if pending and (len(self._all_uids) > 1 or preempted):
-            responses = list(pre_responses)
-            gen_inj = getattr(self.model, "_generator_injections", None)
-            if gen_inj is None:
-                self.model._generator_injections = []
-                gen_inj = self.model._generator_injections
-
-            for other in pending:
-                B_new = len(other._all_uids)
-                base_row = len(self._all_uids)
-                self._all_uids.extend(other._all_uids)
-                self._num_tokens.extend([0] * B_new)
-                self._finished.extend([False] * B_new)
-                self.max_tokens.extend(other.max_tokens)
-
-                mx.eval(other.first_tokens)
-                first_list = other.first_tokens.tolist()
-                for row in range(B_new):
-                    abs_row = base_row + row
-                    tok = int(first_list[row])
-                    self._num_tokens[abs_row] = 1
-                    finish = self._finish_reason(abs_row, tok)
-                    if finish is not None:
-                        self._finished[abs_row] = True
-                    responses.append(
-                        self.Response(
-                            uid=other._all_uids[row],
-                            token=tok,
-                            token_logprob=0.0,
-                            finish_reason=finish,
-                        )
-                    )
-
-                gen_inj.append(
-                    {
-                        "uids": list(other._all_uids),
-                        "prompt_cache": other.prompt_cache,
-                        "hidden": other.hidden,
-                        "shared_kv_states": other.shared_kv_states,
-                        "prompt_tokens": other.prompt_tokens,
-                        "first_tokens": other.first_tokens,
-                        "first_tokens_list": first_list,
-                        # The running generator froze max(max_tokens) at
-                        # start; injected rows carry their own budgets.
-                        "max_tokens": list(other.max_tokens),
-                    }
-                )
-
-            pending.clear()
-            self._refresh_uids()
-
-            more = _orig_next(self)
-            responses.extend(more)
-            _note_last_tokens(self, responses)
-            _release_if_finished(self)
-            return responses
-
-        responses = pre_responses + _orig_next(self)
-        _note_last_tokens(self, responses)
-        _release_if_finished(self)
-        return responses
-
-    SpecBatch.next = _next_with_injection
-    setattr(SpecBatch, _CONTINUOUS_BATCH_FLAG, True)
-    _debug_note(
-        "[mtp] continuous batch: admission gate removed, mid-flight injection enabled"
-    )
-
-
 def install_owned_spec_engine() -> None:
     """Route serve-path MTP through owned engine: B=1 scalar, B>1 batch.
 
@@ -2633,318 +515,3 @@ def _owned_server_rounds(
 
     _gen.run_speculative_server_rounds = _owned_server_rounds
     _debug_note("[mtp] serve round: owned engine installed (B=1 + B>1)")
-
-
-_SPEC_KV_QUANT_FLAG = "_kq_gguf_spec_kv_quant"
-_SPEC_KV_QUANT_WIDTHS = (2, 3, 4, 6, 8)  # mx.quantize affine widths
-
-
-def _spec_kv_quant_params():
-    """resolve_kv_quant_policy kwargs for the KV quantization serve's env
-    asks the trimmable B=1 single-stream cache to honor, else None.
-    Fractional widths and unknown schemes have no such cache; kvarn
-    engages on the scheme alone (widths default like the CLI's)."""
-    if os.environ.get("GMLX_SPEC_KV_QUANT", "1") == "0":
-        return None
-    scheme = os.environ.get("KV_QUANT_SCHEME", "uniform")
-    raw = os.environ.get("KV_BITS", "")
-    if scheme == "kvarn":
-        from gmlx.cache.kvarn_cache import kvarn_widths, parse_tail_tokens
-
-        try:
-            bits = int(raw) if raw else None
-            tail = parse_tail_tokens(os.environ.get("KV_TAIL_TOKENS"))
-        except ValueError:
-            _log.warning(
-                "KV_BITS/KV_TAIL_TOKENS malformed under scheme kvarn; "
-                "B=1 MTP target KV stays fp16"
-            )
-            return None
-        k_bits, v_bits = kvarn_widths(bits)
-        return dict(scheme="kvarn", kv_bits=k_bits, value_bits=v_bits,
-                    tail_tokens=tail)
-    if not raw:
-        return None
-    try:
-        bits = float(raw)
-    except ValueError:
-        return None
-    if bits <= 0:
-        return None
-    if (
-        scheme != "uniform"
-        or bits != int(bits)
-        or int(bits) not in _SPEC_KV_QUANT_WIDTHS
-    ):
-        _log.warning(
-            "KV_BITS=%s scheme=%s: no trimmable single-stream cache; "
-            "B=1 MTP target KV stays fp16",
-            raw,
-            scheme,
-        )
-        return None
-    return dict(scheme="uniform", kv_bits=int(bits),
-                kv_group_size=int(os.environ.get("KV_GROUP_SIZE", "64")))
-
-
-def _stamped_spec_params(lm):
-    """resolve_kv_quant_policy kwargs from the KV policy residency stamped
-    on this model at load, None when nothing is stamped. Per-model env
-    windows are closed by request time, so the stamp rules the boot env;
-    a stamp that quantizes nothing (off, dropped, error) yields {} (stay
-    fp16)."""
-    from gmlx.cache.kvarn_serve import stamped_single_policy
-
-    single = stamped_single_policy(lm)
-    if single is None:
-        return None
-    bits = getattr(single, "bits", None)
-    if not bits or getattr(single, "verdict", None) not in ("full", "partial"):
-        return {}
-    if getattr(single, "scheme", None) == "kvarn":
-        from gmlx.cache.kvarn_cache import KVARN_DEFAULT_TAIL
-
-        tail = single.tail_tokens
-        return dict(scheme="kvarn", kv_bits=int(bits),
-                    value_bits=int(single.value_bits or bits),
-                    tail_tokens=(KVARN_DEFAULT_TAIL if tail is None
-                                 else int(tail)))
-    if int(bits) != bits or int(bits) not in _SPEC_KV_QUANT_WIDTHS:
-        return {}
-    return dict(scheme="uniform", kv_bits=int(bits),
-                kv_group_size=int(single.group_size))
-
-
-def _spec_target_holders(lm) -> tuple:
-    """The spec target and the language model it may wrap: serve hands
-    the cache builder an MTPTextTarget, the CLI the bare model. Every
-    probe on the target reads both."""
-    inner = getattr(lm, "language_model", None)
-    return (lm,) if inner is None or inner is lm else (lm, inner)
-
-
-def _mtp_reads_kv_back(lm) -> bool:
-    """True when the target's verify route re-reads K/V from the prompt
-    cache (spec_helpers._mtp_shared_kv_from_prompt_cache): it computes
-    logits from hidden but owns no verify hook, so the walk rebuilds the
-    drafter's shared K/V from cache state -- raw arrays kvarn records
-    cannot supply."""
-    return any(
-        callable(getattr(h, "speculative_logits_from_hidden", None))
-        and not callable(getattr(h, "speculative_verify_hidden", None))
-        and not callable(getattr(h, "speculative_verify_logits", None))
-        for h in _spec_target_holders(lm)
-    )
-
-
-def _harden_spec_target(lm) -> None:
-    """harden_mtp_rollback on every holder of the target's rollback."""
-    from gmlx.gen.generation import harden_mtp_rollback
-
-    for h in _spec_target_holders(lm):
-        harden_mtp_rollback(h)
-
-
-def _kvarn_spec_reason(lm):
-    """The kvarn declines the shared policy cannot see: the target's own
-    verify contract. Both MTP arms (B=1 and batched) check it."""
-    from gmlx.cache.kvarn_cache import kvarn_unsupported
-
-    reason = kvarn_unsupported(lm)
-    if reason is None and _mtp_reads_kv_back(lm):
-        reason = (
-            "the target's verify path reads shared K/V back "
-            "from the cache (kvarn records are not raw K/V)"
-        )
-    return reason
-
-
-def _kvarn_batch_spec_cache(lm, caches, left_padding, params):
-    """Convert a B>1 MTP target stack to kvarn batch rows when the batched
-    arm is engaged, in place; None when it declines (the caller keeps the
-    fp16 batch stack). The B=1 declines apply: an ineligible model, a
-    target that reads K/V back, a sliding-window stack, and the policy's
-    own drop when mlx-kquant lacks per-row ends. The verify block is not
-    knowable here; batch formation clamps it to the kernels' width."""
-    from gmlx.cache.kv_policy import kv_line, note_once
-    from gmlx.cache.kvarn_cache import (KVARN_DEFAULT_TAIL,
-                                        kvarn_mtp_window_decline)
-    from gmlx.cache.kvarn_serve import (kvarn_batch_policy,
-                                        kvarn_convert_batch_stack)
-
-    def decline(reason):
-        if note_once(lm, "spec-kv-kvarn-batched"):
-            _log.warning(
-                "KV_QUANT_SCHEME=kvarn dropped on the batched MTP path: %s; "
-                "the batch runs fp16 KV", reason)
-        return None
-
-    reason = kvarn_mtp_window_decline(caches) or _kvarn_spec_reason(lm)
-    if reason is not None:
-        return decline(reason)
-    k_bits = int(params["kv_bits"])
-    v_bits = int(params.get("value_bits") or k_bits)
-    tail = params.get("tail_tokens")
-    tail = KVARN_DEFAULT_TAIL if tail is None else int(tail)
-    policy = kvarn_batch_policy(lm, caches, k_bits, v_bits, tail,
-                                mode="batched", mtp=True)
-    if policy.verdict not in ("full", "partial"):
-        return decline(policy.reason)
-    n = kvarn_convert_batch_stack(caches, policy, left_padding, k_bits,
-                                  v_bits, tail)
-    if not n:
-        return decline("no plain KV-cache layers in this arch's stack")
-    _harden_spec_target(lm)
-    if note_once(lm, "spec-kv-batched"):
-        _log.info("%s", kv_line("MTP spec path (batched)", policy))
-    return caches
-
-
-def install_spec_kv_quant() -> None:
-    """Honor KV_BITS on the B=1 MTP serve path.
-
-    Stock ``make_speculative_prompt_cache`` returns plain fp16 caches for
-    ``draft_kind == "mtp", batch_size == 1``, discarding the engine's
-    kv_bits: ``BatchQuantizedKVCache`` cannot trim, and MTP rollback must
-    trim the target. The single-stream ``QuantizedKVCache`` can trim --
-    packing is per-token along head_dim, so trim is an offset move -- and
-    the model rollback already goes through ``is_trimmable()``/``trim()``.
-    The shared KV policy picks the layers: growing KV converts at
-    construction (empty, so conversion is free), quantizable pools pack
-    at rest, and windows, recurrent state, and opt-outs stay fp16 at any
-    nesting depth. Scheme kvarn converts the same B=1 caches to
-    ``KVarNKVCache`` instead (rollback rides the stage/horizon regions),
-    declining targets whose verify path reads shared K/V back from cache
-    state. B>1 MTP under kvarn converts the batch stack to
-    ``BatchKVarNKVCache`` rows when the installed mlx-kquant takes
-    per-row ends (each row rolls back by its own rejected count); under
-    uniform, or on an older mlx-kquant, B>1 keeps fp16 batch KV with a
-    one-shot warning (the packed batch cache cannot trim). Scheme and
-    widths come from the policy stamped on the model at load; the boot
-    env is the fallback for unstamped models. Kill switch:
-    GMLX_SPEC_KV_QUANT=0."""
-    from mlx_vlm.generate import ar as _ar
-    from mlx_vlm.server import generation as _gen
-    from mlx_vlm.speculative import utils as _su
-
-    if getattr(_su.make_speculative_prompt_cache, _SPEC_KV_QUANT_FLAG, False):
-        return
-    if os.environ.get("GMLX_SPEC_KV_QUANT", "1") == "0":
-        return
-    boot_params = _spec_kv_quant_params()
-
-    from gmlx.cache.compat import cache_types
-
-    from gmlx.cache.kv_policy import note_once
-
-    _orig = _su.make_speculative_prompt_cache
-
-    def _decline_kvarn(lm, reason: str):
-        if note_once(lm, "spec-kv-kvarn"):
-            _log.warning(
-                "KV_QUANT_SCHEME=kvarn dropped on the B=1 MTP path: %s", reason
-            )
-
-    def _quantizing_spec_cache(lm, *, draft_kind, batch_size, left_padding, make_cache):
-        from gmlx.cache.kvarn_serve import spec_cache_build
-
-        # The stock make_cache closure passes the boot scheme through;
-        # suspend the serve wrap so spec targets never get a batch kvarn
-        # cache (the verify walk needs trim, which it does not support).
-        with spec_cache_build():
-            caches = _orig(
-                lm,
-                draft_kind=draft_kind,
-                batch_size=batch_size,
-                left_padding=left_padding,
-                make_cache=make_cache,
-            )
-        if draft_kind != "mtp":
-            return caches
-        params = _stamped_spec_params(lm)
-        if params is None:
-            params = boot_params
-        if batch_size != 1:
-            if params and params.get("scheme") == "kvarn":
-                out = _kvarn_batch_spec_cache(lm, caches, left_padding, params)
-                if out is not None:
-                    return out
-            # Force fp16 batch KV: the stock rollback misfiles
-            # BatchQuantizedKVCache as an SSM cache and never trims
-            # rejected drafts.
-            # to_batch_cache also quantizes nested subcaches. Walk
-            # into CacheList entries.
-            from mlx_vlm.models.cache import BatchKVCache
-
-            batch_quant = cache_types("BatchQuantizedKVCache")
-
-            def _swap(c):
-                if isinstance(c, batch_quant):
-                    return BatchKVCache(left_padding), 1
-                inner = getattr(c, "caches", None)
-                if inner is None:
-                    return c, 0
-                subs = [_swap(s) for s in inner]
-                n = sum(k for _, k in subs)
-                if n:
-                    c.caches = tuple(s for s, _ in subs)
-                return c, n
-
-            swapped = 0
-            for e, c in enumerate(caches):
-                caches[e], n_sw = _swap(c)
-                swapped += n_sw
-            if swapped and note_once(lm, "spec-kv-batch"):
-                _log.warning(
-                    "KV quantization with MTP at batch size %d: packed "
-                    "batch rollback is unsupported; %d layers run fp16 KV",
-                    batch_size, swapped)
-            return caches
-        if not params:
-            return caches
-        kind = params["scheme"]
-        decline = mtp_kv_decline(lm)
-        if decline is not None:
-            if note_once(lm, "spec-kv-stock"):
-                _log.warning(
-                    "KV quantization dropped on the MTP path: %s", decline)
-            return caches
-        scheme_reason = None
-        if kind == "kvarn":
-            from gmlx.cache.kvarn_cache import kvarn_mtp_window_decline
-
-            scheme_reason = (kvarn_mtp_window_decline(caches)
-                             or _kvarn_spec_reason(lm))
-            if scheme_reason is not None:
-                _decline_kvarn(lm, scheme_reason)
-                return caches
-        # The shared policy owns layer selection: nested KV members,
-        # pools, windows, and opt-outs at any depth.
-        from gmlx.cache.kv_policy import (kv_line, quantize_stack,
-                                          resolve_kv_quant_policy)
-
-        policy = resolve_kv_quant_policy(
-            caches, mode="single", scheme_reason=scheme_reason, **params)
-        if policy.verdict not in ("full", "partial"):
-            if kind == "kvarn":
-                _decline_kvarn(lm, policy.reason)
-            elif note_once(lm, "spec-kv-dropped"):
-                _log.warning("%s", kv_line("MTP spec path", policy))
-            return caches
-
-        armed, n = quantize_stack(caches, policy)
-        if kind == "kvarn":
-            if not n:
-                _decline_kvarn(lm, "no plain KV-cache layers in this arch's stack")
-                return caches
-            _harden_spec_target(lm)
-        if (n or armed) and note_once(lm, "spec-kv"):
-            _log.info("%s", kv_line("MTP spec path", policy))
-        return caches
-
-    _quantizing_spec_cache.__dict__[_SPEC_KV_QUANT_FLAG] = True
-    _quantizing_spec_cache.__dict__["_gmlx_orig"] = _orig
-    _su.make_speculative_prompt_cache = _quantizing_spec_cache
-    _ar.make_speculative_prompt_cache = _quantizing_spec_cache
-    _gen.make_speculative_prompt_cache = _quantizing_spec_cache
-    _debug_note(f"[mtp] spec cache wrap armed (B=1); boot env {boot_params}")
diff --git a/gmlx/spec/kv_quant.py b/gmlx/spec/kv_quant.py
new file mode 100644
index 00000000..11a95faa
--- /dev/null
+++ b/gmlx/spec/kv_quant.py
@@ -0,0 +1,456 @@
+"""Speculative KV quantization: lift helpers and the install block.
+
+Split out of ``gmlx.spec.engine``.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+
+import mlx.core as mx
+
+from gmlx.spec.engine import _debug_note
+
+_log = logging.getLogger(__name__)
+
+
+def dequantize_lift_cache(c):
+    """Dequantize a B=1 QuantizedKVCache into a one-row BatchKVCache.
+
+    QuantizedKVCache has no merge, and the B>1 MTP arm runs fp16 KV
+    anyway, so the lift performs the same conversion the batch-build
+    swap does, at preemption or injection time, as a one-off O(depth)
+    copy. Without this every kv-bits MTP preemption declined into a
+    drain-wait and queued rows stalled behind the live generation."""
+    from mlx_vlm.models.cache import BatchKVCache
+
+    lifted = BatchKVCache([0])
+    L = c.offset
+    if L:
+        keys = mx.dequantize(
+            *(mx.contiguous(t[..., :L, :]) for t in c.keys),
+            group_size=c.group_size, bits=c.bits)
+        values = mx.dequantize(
+            *(mx.contiguous(t[..., :L, :]) for t in c.values),
+            group_size=c.group_size, bits=c.bits)
+        lifted.update_and_fetch(keys, values)
+    stamp = getattr(c, "_gmlx_cascade", None)
+    if stamp is not None:
+        lifted._gmlx_cascade = stamp
+    return lifted
+
+
+def kvarn_lift_cache(c):
+    """Recover a B=1 KVarNKVCache into a one-row fp16 BatchKVCache, the
+    lift for an mlx-kquant without per-row ends (the batched arm then
+    runs fp16 KV).
+
+    The kvarn twin of dequantize_lift_cache: materialize() returns
+    rotated-domain K/V, which stock SDPA would attend with an un-rotated
+    query -- no crash, just wrong logits on every preempted row.
+    _raw_single is the original-domain accessor."""
+    from mlx_vlm.models.cache import BatchKVCache
+
+    lifted = BatchKVCache([0])
+    if c.offset:
+        keys, values = c._raw_single()
+        lifted.update_and_fetch(keys, values)
+    stamp = getattr(c, "_gmlx_cascade", None)
+    if stamp is not None:
+        lifted._gmlx_cascade = stamp
+    return lifted
+
+
+def mtp_kv_decline(lm, *, owned_round: bool = True) -> str | None:
+    """Why this MTP verify walk cannot run on packed KV, or None.
+
+    The owned rounds roll back with trim, and affine packing is
+    per-token along head_dim, so a trim is an offset move: they take the
+    same layers serve takes. The two stock walks slice keys as raw
+    arrays and cannot read a packed tuple back. Shared by serve, run and
+    chat so the three cannot drift.
+    """
+    if not owned_round:
+        return "GMLX_OWNED_ROUND=0 stock rounds have no KV quantization hook"
+    from gmlx.models.qwen35.gdn import stock_gdn_fallback
+
+    mt = None
+    for h in _spec_target_holders(lm):
+        cfg = getattr(h, "config", None)
+        mt = (getattr(h, "model_type", None)
+              or (cfg.get("model_type") if isinstance(cfg, dict)
+                  else getattr(cfg, "model_type", None)))
+        if mt:
+            break
+    if stock_gdn_fallback(mt):
+        return ("the GMLX_QWEN_OWNED=0 stock fallback cannot verify on a "
+                "quantized KV cache")
+    return None
+
+
+def lift_single_cache(c):
+    """Promote a single-sequence cache to its batch class. An affine B=1
+    cache recovers to a one-row fp16 BatchKVCache (the batched MTP arm
+    runs fp16 KV under uniform). A kvarn B=1 cache becomes a one-row
+    BatchKVarNKVCache, buffers and horizon adopted bit-exactly, when the
+    installed mlx-kquant takes per-row ends, and recovers to fp16 rows
+    otherwise. Everything else lifts through its class's merge. The
+    cascade stamp rides along. One lift for the preempted host row and
+    the injected rows, so the two cannot drift."""
+    from gmlx.cache.compat import cache_types
+
+    if isinstance(c, cache_types("QuantizedKVCache")):
+        return dequantize_lift_cache(c)
+    if getattr(c, "kv_quant_scheme", None) == "kvarn" and batch_liftable(c):
+        from gmlx.cache.kvarn_sdpa import kvarn_row_ends_ok
+
+        if not kvarn_row_ends_ok():
+            return kvarn_lift_cache(c)
+        from gmlx.cache.kvarn_cache import BatchKVarNKVCache
+
+        lifted = BatchKVarNKVCache.merge([c])
+        stamp = getattr(c, "_gmlx_cascade", None)
+        if stamp is not None:
+            lifted._gmlx_cascade = stamp
+        return lifted
+    lifted = type(c).merge([c])
+    stamp = getattr(c, "_gmlx_cascade", None)
+    if stamp is not None:
+        lifted._gmlx_cascade = stamp
+    return lifted
+
+
+def batch_liftable(c) -> bool:
+    """Whether _lift_host_cache can promote this cache to a batch class.
+
+    The preemption and pre-start-compaction gates must agree with it: a
+    scheme it can lift but they refuse declines into the drain-wait, and
+    one they admit but it cannot lift raises mid-rebuild.
+    """
+    if hasattr(c, "filter") and hasattr(c, "extend"):
+        return True
+    if getattr(c, "kv_quant_scheme", None) == "kvarn":
+        from gmlx.cache.kvarn_cache import KVarNRotatingKVCache
+
+        # The rotating subclass counts evicted tokens in offset that its
+        # buffers no longer hold: kvarn_lift_cache would misplace rows.
+        return not isinstance(c, KVarNRotatingKVCache)
+    from gmlx.cache.compat import cache_types
+
+    return (isinstance(c, cache_types("QuantizedKVCache"))
+            or hasattr(type(c), "merge"))
+
+_SPEC_KV_QUANT_FLAG = "_kq_gguf_spec_kv_quant"
+_SPEC_KV_QUANT_WIDTHS = (2, 3, 4, 6, 8)  # mx.quantize affine widths
+
+
+def _spec_kv_quant_params():
+    """resolve_kv_quant_policy kwargs for the KV quantization serve's env
+    asks the trimmable B=1 single-stream cache to honor, else None.
+    Fractional widths and unknown schemes have no such cache; kvarn
+    engages on the scheme alone (widths default like the CLI's)."""
+    if os.environ.get("GMLX_SPEC_KV_QUANT", "1") == "0":
+        return None
+    scheme = os.environ.get("KV_QUANT_SCHEME", "uniform")
+    raw = os.environ.get("KV_BITS", "")
+    if scheme == "kvarn":
+        from gmlx.cache.kvarn_cache import kvarn_widths, parse_tail_tokens
+
+        try:
+            bits = int(raw) if raw else None
+            tail = parse_tail_tokens(os.environ.get("KV_TAIL_TOKENS"))
+        except ValueError:
+            _log.warning(
+                "KV_BITS/KV_TAIL_TOKENS malformed under scheme kvarn; "
+                "B=1 MTP target KV stays fp16"
+            )
+            return None
+        k_bits, v_bits = kvarn_widths(bits)
+        return dict(scheme="kvarn", kv_bits=k_bits, value_bits=v_bits,
+                    tail_tokens=tail)
+    if not raw:
+        return None
+    try:
+        bits = float(raw)
+    except ValueError:
+        return None
+    if bits <= 0:
+        return None
+    if (
+        scheme != "uniform"
+        or bits != int(bits)
+        or int(bits) not in _SPEC_KV_QUANT_WIDTHS
+    ):
+        _log.warning(
+            "KV_BITS=%s scheme=%s: no trimmable single-stream cache; "
+            "B=1 MTP target KV stays fp16",
+            raw,
+            scheme,
+        )
+        return None
+    return dict(scheme="uniform", kv_bits=int(bits),
+                kv_group_size=int(os.environ.get("KV_GROUP_SIZE", "64")))
+
+
+def _stamped_spec_params(lm):
+    """resolve_kv_quant_policy kwargs from the KV policy residency stamped
+    on this model at load, None when nothing is stamped. Per-model env
+    windows are closed by request time, so the stamp rules the boot env;
+    a stamp that quantizes nothing (off, dropped, error) yields {} (stay
+    fp16)."""
+    from gmlx.cache.kvarn_serve import stamped_single_policy
+
+    single = stamped_single_policy(lm)
+    if single is None:
+        return None
+    bits = getattr(single, "bits", None)
+    if not bits or getattr(single, "verdict", None) not in ("full", "partial"):
+        return {}
+    if getattr(single, "scheme", None) == "kvarn":
+        from gmlx.cache.kvarn_cache import KVARN_DEFAULT_TAIL
+
+        tail = single.tail_tokens
+        return dict(scheme="kvarn", kv_bits=int(bits),
+                    value_bits=int(single.value_bits or bits),
+                    tail_tokens=(KVARN_DEFAULT_TAIL if tail is None
+                                 else int(tail)))
+    if int(bits) != bits or int(bits) not in _SPEC_KV_QUANT_WIDTHS:
+        return {}
+    return dict(scheme="uniform", kv_bits=int(bits),
+                kv_group_size=int(single.group_size))
+
+
+def _spec_target_holders(lm) -> tuple:
+    """The spec target and the language model it may wrap: serve hands
+    the cache builder an MTPTextTarget, the CLI the bare model. Every
+    probe on the target reads both."""
+    inner = getattr(lm, "language_model", None)
+    return (lm,) if inner is None or inner is lm else (lm, inner)
+
+
+def _mtp_reads_kv_back(lm) -> bool:
+    """True when the target's verify route re-reads K/V from the prompt
+    cache (spec_helpers._mtp_shared_kv_from_prompt_cache): it computes
+    logits from hidden but owns no verify hook, so the walk rebuilds the
+    drafter's shared K/V from cache state -- raw arrays kvarn records
+    cannot supply."""
+    return any(
+        callable(getattr(h, "speculative_logits_from_hidden", None))
+        and not callable(getattr(h, "speculative_verify_hidden", None))
+        and not callable(getattr(h, "speculative_verify_logits", None))
+        for h in _spec_target_holders(lm)
+    )
+
+
+def _harden_spec_target(lm) -> None:
+    """harden_mtp_rollback on every holder of the target's rollback."""
+    from gmlx.gen.generation import harden_mtp_rollback
+
+    for h in _spec_target_holders(lm):
+        harden_mtp_rollback(h)
+
+
+def _kvarn_spec_reason(lm):
+    """The kvarn declines the shared policy cannot see: the target's own
+    verify contract. Both MTP arms (B=1 and batched) check it."""
+    from gmlx.cache.kvarn_cache import kvarn_unsupported
+
+    reason = kvarn_unsupported(lm)
+    if reason is None and _mtp_reads_kv_back(lm):
+        reason = (
+            "the target's verify path reads shared K/V back "
+            "from the cache (kvarn records are not raw K/V)"
+        )
+    return reason
+
+
+def _kvarn_batch_spec_cache(lm, caches, left_padding, params):
+    """Convert a B>1 MTP target stack to kvarn batch rows when the batched
+    arm is engaged, in place; None when it declines (the caller keeps the
+    fp16 batch stack). The B=1 declines apply: an ineligible model, a
+    target that reads K/V back, a sliding-window stack, and the policy's
+    own drop when mlx-kquant lacks per-row ends. The verify block is not
+    knowable here; batch formation clamps it to the kernels' width."""
+    from gmlx.cache.kv_policy import kv_line, note_once
+    from gmlx.cache.kvarn_cache import (KVARN_DEFAULT_TAIL,
+                                        kvarn_mtp_window_decline)
+    from gmlx.cache.kvarn_serve import (kvarn_batch_policy,
+                                        kvarn_convert_batch_stack)
+
+    def decline(reason):
+        if note_once(lm, "spec-kv-kvarn-batched"):
+            _log.warning(
+                "KV_QUANT_SCHEME=kvarn dropped on the batched MTP path: %s; "
+                "the batch runs fp16 KV", reason)
+        return None
+
+    reason = kvarn_mtp_window_decline(caches) or _kvarn_spec_reason(lm)
+    if reason is not None:
+        return decline(reason)
+    k_bits = int(params["kv_bits"])
+    v_bits = int(params.get("value_bits") or k_bits)
+    tail = params.get("tail_tokens")
+    tail = KVARN_DEFAULT_TAIL if tail is None else int(tail)
+    policy = kvarn_batch_policy(lm, caches, k_bits, v_bits, tail,
+                                mode="batched", mtp=True)
+    if policy.verdict not in ("full", "partial"):
+        return decline(policy.reason)
+    n = kvarn_convert_batch_stack(caches, policy, left_padding, k_bits,
+                                  v_bits, tail)
+    if not n:
+        return decline("no plain KV-cache layers in this arch's stack")
+    _harden_spec_target(lm)
+    if note_once(lm, "spec-kv-batched"):
+        _log.info("%s", kv_line("MTP spec path (batched)", policy))
+    return caches
+
+
+def install_spec_kv_quant() -> None:
+    """Honor KV_BITS on the B=1 MTP serve path.
+
+    Stock ``make_speculative_prompt_cache`` returns plain fp16 caches for
+    ``draft_kind == "mtp", batch_size == 1``, discarding the engine's
+    kv_bits: ``BatchQuantizedKVCache`` cannot trim, and MTP rollback must
+    trim the target. The single-stream ``QuantizedKVCache`` can trim --
+    packing is per-token along head_dim, so trim is an offset move -- and
+    the model rollback already goes through ``is_trimmable()``/``trim()``.
+    The shared KV policy picks the layers: growing KV converts at
+    construction (empty, so conversion is free), quantizable pools pack
+    at rest, and windows, recurrent state, and opt-outs stay fp16 at any
+    nesting depth. Scheme kvarn converts the same B=1 caches to
+    ``KVarNKVCache`` instead (rollback rides the stage/horizon regions),
+    declining targets whose verify path reads shared K/V back from cache
+    state. B>1 MTP under kvarn converts the batch stack to
+    ``BatchKVarNKVCache`` rows when the installed mlx-kquant takes
+    per-row ends (each row rolls back by its own rejected count); under
+    uniform, or on an older mlx-kquant, B>1 keeps fp16 batch KV with a
+    one-shot warning (the packed batch cache cannot trim). Scheme and
+    widths come from the policy stamped on the model at load; the boot
+    env is the fallback for unstamped models. Kill switch:
+    GMLX_SPEC_KV_QUANT=0."""
+    from mlx_vlm.generate import ar as _ar
+    from mlx_vlm.server import generation as _gen
+    from mlx_vlm.speculative import utils as _su
+
+    if getattr(_su.make_speculative_prompt_cache, _SPEC_KV_QUANT_FLAG, False):
+        return
+    if os.environ.get("GMLX_SPEC_KV_QUANT", "1") == "0":
+        return
+    boot_params = _spec_kv_quant_params()
+
+    from gmlx.cache.compat import cache_types
+
+    from gmlx.cache.kv_policy import note_once
+
+    _orig = _su.make_speculative_prompt_cache
+
+    def _decline_kvarn(lm, reason: str):
+        if note_once(lm, "spec-kv-kvarn"):
+            _log.warning(
+                "KV_QUANT_SCHEME=kvarn dropped on the B=1 MTP path: %s", reason
+            )
+
+    def _quantizing_spec_cache(lm, *, draft_kind, batch_size, left_padding, make_cache):
+        from gmlx.cache.kvarn_serve import spec_cache_build
+
+        # The stock make_cache closure passes the boot scheme through;
+        # suspend the serve wrap so spec targets never get a batch kvarn
+        # cache (the verify walk needs trim, which it does not support).
+        with spec_cache_build():
+            caches = _orig(
+                lm,
+                draft_kind=draft_kind,
+                batch_size=batch_size,
+                left_padding=left_padding,
+                make_cache=make_cache,
+            )
+        if draft_kind != "mtp":
+            return caches
+        params = _stamped_spec_params(lm)
+        if params is None:
+            params = boot_params
+        if batch_size != 1:
+            if params and params.get("scheme") == "kvarn":
+                out = _kvarn_batch_spec_cache(lm, caches, left_padding, params)
+                if out is not None:
+                    return out
+            # Force fp16 batch KV: the stock rollback misfiles
+            # BatchQuantizedKVCache as an SSM cache and never trims
+            # rejected drafts.
+            # to_batch_cache also quantizes nested subcaches. Walk
+            # into CacheList entries.
+            from mlx_vlm.models.cache import BatchKVCache
+
+            batch_quant = cache_types("BatchQuantizedKVCache")
+
+            def _swap(c):
+                if isinstance(c, batch_quant):
+                    return BatchKVCache(left_padding), 1
+                inner = getattr(c, "caches", None)
+                if inner is None:
+                    return c, 0
+                subs = [_swap(s) for s in inner]
+                n = sum(k for _, k in subs)
+                if n:
+                    c.caches = tuple(s for s, _ in subs)
+                return c, n
+
+            swapped = 0
+            for e, c in enumerate(caches):
+                caches[e], n_sw = _swap(c)
+                swapped += n_sw
+            if swapped and note_once(lm, "spec-kv-batch"):
+                _log.warning(
+                    "KV quantization with MTP at batch size %d: packed "
+                    "batch rollback is unsupported; %d layers run fp16 KV",
+                    batch_size, swapped)
+            return caches
+        if not params:
+            return caches
+        kind = params["scheme"]
+        decline = mtp_kv_decline(lm)
+        if decline is not None:
+            if note_once(lm, "spec-kv-stock"):
+                _log.warning(
+                    "KV quantization dropped on the MTP path: %s", decline)
+            return caches
+        scheme_reason = None
+        if kind == "kvarn":
+            from gmlx.cache.kvarn_cache import kvarn_mtp_window_decline
+
+            scheme_reason = (kvarn_mtp_window_decline(caches)
+                             or _kvarn_spec_reason(lm))
+            if scheme_reason is not None:
+                _decline_kvarn(lm, scheme_reason)
+                return caches
+        # The shared policy owns layer selection: nested KV members,
+        # pools, windows, and opt-outs at any depth.
+        from gmlx.cache.kv_policy import (kv_line, quantize_stack,
+                                          resolve_kv_quant_policy)
+
+        policy = resolve_kv_quant_policy(
+            caches, mode="single", scheme_reason=scheme_reason, **params)
+        if policy.verdict not in ("full", "partial"):
+            if kind == "kvarn":
+                _decline_kvarn(lm, policy.reason)
+            elif note_once(lm, "spec-kv-dropped"):
+                _log.warning("%s", kv_line("MTP spec path", policy))
+            return caches
+
+        armed, n = quantize_stack(caches, policy)
+        if kind == "kvarn":
+            if not n:
+                _decline_kvarn(lm, "no plain KV-cache layers in this arch's stack")
+                return caches
+            _harden_spec_target(lm)
+        if (n or armed) and note_once(lm, "spec-kv"):
+            _log.info("%s", kv_line("MTP spec path", policy))
+        return caches
+
+    _quantizing_spec_cache.__dict__[_SPEC_KV_QUANT_FLAG] = True
+    _quantizing_spec_cache.__dict__["_gmlx_orig"] = _orig
+    _su.make_speculative_prompt_cache = _quantizing_spec_cache
+    _ar.make_speculative_prompt_cache = _quantizing_spec_cache
+    _gen.make_speculative_prompt_cache = _quantizing_spec_cache
+    _debug_note(f"[mtp] spec cache wrap armed (B=1); boot env {boot_params}")
diff --git a/gmlx/spec/mtp_prefill.py b/gmlx/spec/mtp_prefill.py
new file mode 100644
index 00000000..4faceb77
--- /dev/null
+++ b/gmlx/spec/mtp_prefill.py
@@ -0,0 +1,648 @@
+"""Full-prompt MTP prefill and the seed stream.
+
+Split out of ``gmlx.spec.engine``; the install flags
+(``_FULL_PREFILL_FLAG``, ``_SEED_STREAM_DISABLED``) stay on the engine,
+which the stays band also reads.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+
+import mlx.core as mx
+
+import gmlx.lora_rows as lora_rows
+import gmlx.gen.prefill_decay as prefill_decay
+from gmlx.envflags import env_int
+from gmlx.spec.ckpt import (
+    _install_ckpt_checkpoint_store,
+    _install_exact_anchor_pick,
+    _install_plain_ckpt_decode,
+    _l1_lookup_and_arm_store,
+    _plain_anchor_init,
+    _plain_ckpt_init,
+    _snap_fields,
+)
+from gmlx.spec.engine import (
+    _FULL_PREFILL_FLAG,
+    _L1_BOUND,
+    _SEED_STREAM_DISABLED,
+    _SPEC_APC_DISABLED,
+    _SPEC_APC_RETIRE_DISABLED,
+    _bind_l1_view,
+    _ckpt_active,
+    _debug_note,
+    _get_spec_prefix_cache,
+    _install_apc_manager_stash,
+    _resolve_l1,
+)
+
+_log = logging.getLogger(__name__)
+
+
+def _mtp_prefill_init(batch) -> None:
+    """One-time APC lookup + prefix trim for an MTP prompt batch.
+
+    Runs on the first ``prompt_step`` call, or directly from ``generate()``
+    when the prompt is short enough that chunked prefill never fires.
+    Lookup ladder: L0 (SpecPrefixCache: whole-prompt KV + full-prompt
+    hidden, the only tier the drafter can teacher-force from without a cold
+    start) then L1 (shared APCManager: exact / block / disk KV, no hidden).
+    Also arms the stock post-prefill store whenever a manager is reachable,
+    regardless of which tier (if any) hit.
+    """
+    if hasattr(batch, "_mtp_full_input_ids"):
+        return
+    batch._mtp_full_input_ids = batch._input_ids
+    batch._mtp_chunk_hiddens = []
+    batch._mtp_l1_prefix_len = 0
+
+    if batch._inputs_embeds is None:
+        _log.info("KQDBG mtp_prefill_init: inputs_embeds None, ladder skipped")
+        return
+
+    # Gated to B=1 because PromptProcessingBatch prefills one request at a
+    # time today. The restored single-row cache (with its offset) later
+    # merges into the live B>1 decode batch via BatchKVCache.extend during
+    # continuous-batch injection -- so APC absolutely works in a B>1
+    # serving context; the gate is about prefill granularity, not decode
+    # batch size. If mlx-vlm ever coalesces prefills into a multi-row
+    # PromptProcessingBatch, this guard silently disables APC for those
+    # rows. The warning below makes that visible.
+    b = int(batch._input_ids.shape[0])
+    if b > 1:
+        if not _SPEC_APC_DISABLED:
+            _log.warning(
+                "APC skipped: prefill batch B=%d > 1 "
+                "(owned-path APC requires single-request prefill)",
+                b,
+            )
+        return
+
+    # Serve wraps make_cache so mlx-lm-origin entries carry the mlx-vlm
+    # runtime's class identities; embedded and test users reach this init
+    # without that wrapper, and the L1 exact tiers dispatch on the vlm
+    # classes (an mlx-lm ArraysCache misses every adapter rule). Rebind
+    # here so both paths see the same identities. No-op when the entries
+    # are already vlm-origin.
+    from gmlx.cache.compat import rebind_to_runtime_origin
+    rebind_to_runtime_origin(batch.prompt_cache)
+
+    # Upstream admission already restored a prefix and built this batch
+    # suffix-only: the owned ladder's keys (L0 and L1 both) are full-prompt
+    # token ids, so every lookup and store here would run in the wrong
+    # space -- a suffix-keyed L0 entry cross-hits a later turn's suffix and
+    # its restore clobbers the upstream warm cache. Leave these batches to
+    # the stock machinery, which owns their meta and store schedule.
+    up_meta = getattr(batch, "_apc_meta", None) or []
+    if up_meta and isinstance(up_meta[0], dict) \
+            and int(up_meta[0].get("prefix_len") or 0) > 0:
+        batch._mtp_upstream_warm = True
+        return
+
+    restored = 0
+    spec_cache = _get_spec_prefix_cache(batch.model)
+    if spec_cache is not None:
+        hit = spec_cache.lookup(batch._input_ids)
+        if hit is not None:
+            restored, entry = hit
+            spec_cache.restore(entry, batch.prompt_cache)
+            batch._mtp_chunk_hiddens = [entry.hidden]
+            _log.info(
+                "APC hit: prefix=%d suffix=%d",
+                restored,
+                int(batch._input_ids.shape[1]) - restored,
+            )
+
+    manager, mode = _resolve_l1(batch.model)
+    if manager is not None:
+        try:
+            l1_prefix = _l1_lookup_and_arm_store(batch, manager, mode, restored)
+            restored = max(restored, l1_prefix)
+        except Exception:
+            _log.warning("APC L1 failed; continuing cold", exc_info=True)
+
+    # Stash the retirement context so the owned B=1 round can store this
+    # request's full context (prompt + generated) into the shared APC when it
+    # finishes. Keyed on the original full ids (pre-trim) -- the serve-layer
+    # prompt_tokens is suffix-only on a warm turn, so it can't be the key.
+    # The stash lives on the request's first cache entry, not on the model:
+    # the server closes a finished rounds generator lazily (sometimes after
+    # the next request's prefill), so a model-level stash races and retires
+    # under the wrong key. Must run after the L1 block above -- an exact-tier
+    # hit replaces batch.prompt_cache wholesale. B=1 only (this init is gated
+    # to B=1); B>1 retirement is handled per-row at the batch decode's
+    # finish seam.
+    if manager is not None and not _SPEC_APC_RETIRE_DISABLED and batch.prompt_cache:
+        meta = (batch._apc_meta or [{}])[0] or {}
+        full_ids = [int(t) for t in batch._mtp_full_input_ids[0].tolist()]
+        from gmlx.cache.retire_key import lookup_render_ctx
+        batch.prompt_cache[0]._kq_apc_retire = {
+            "full_ids": full_ids,
+            "extra_hash": int(meta.get("extra_hash", 0)),
+            "mode": (
+                "ckpt"
+                if _ckpt_active(batch.model, mode, int(manager.block_size))
+                else mode
+            ),
+            "checkpoint_len": int(meta.get("checkpoint_len", 0) or 0),
+            # Live reference: the sidecar keys on ckpt_last_stored, not
+            # the cursor value frozen above.
+            "apc_meta": meta,
+            # Render context for the next-turn LCP key (None off the server
+            # path or on a media prompt; retirement then keys as before).
+            "render_ctx": lookup_render_ctx(full_ids),
+            **_snap_fields(batch, manager),
+        }
+
+    if restored > 0:
+        batch._input_ids = batch._input_ids[:, restored:]
+        batch._inputs_embeds = batch._inputs_embeds[:, restored:]
+        batch._processed_prompt_columns = restored
+        for k in batch._prompt_length_aware_keys:
+            batch._prompt_kwargs[k] = batch._prompt_kwargs[k][:, restored:, ...]
+        batch._mtp_apc_prefix_len = restored
+
+
+def _mtp_seed_stream_init(batch) -> None:
+    """Arm per-chunk drafter seeding for this request, if eligible.
+
+    Cold full prefill only (v1): any restored prefix (L0/L1/upstream) or warm
+    drafter sidecar keeps the deferred one-shot seed -- correctness identical,
+    seeding then still runs after the first token. Eligibility here plus the
+    per-chunk B re-check in prompt_step; a mid-request stop keeps the partial
+    seed KV (adopted at its true offset) and defers only the remainder.
+
+    The seed KV is request-scoped (built via drafter.make_cache, ridden on
+    batch state and handed over via a prompt_cache[0] stash exactly like the
+    drafter warm sidecar), never the drafter's own _cache: another request's
+    live decode round owns that object.
+    """
+    if hasattr(batch, "_mtp_seed_ctx"):
+        return
+    batch._mtp_seed_ctx = None
+    drafter = getattr(batch, "draft_model", None)
+    if (
+        _SEED_STREAM_DISABLED
+        or drafter is None
+        or not callable(getattr(drafter, "seed_chunk", None))
+        or getattr(drafter, "hidden_capture_limit", None) is not None
+        or int(batch._input_ids.shape[0]) != 1
+        or getattr(batch, "_mtp_upstream_warm", False)
+        or getattr(batch, "_mtp_chunk_hiddens", None)
+        or int(getattr(batch, "_mtp_l1_prefix_len", 0) or 0) != 0
+        or int(getattr(batch, "_processed_prompt_columns", 0) or 0) != 0
+        or not batch.prompt_cache
+        or getattr(batch.prompt_cache[0], "_kq_apc_drafter_warm", None)
+            is not None
+    ):
+        return
+    lp = getattr(batch.prompt_cache[0], "left_padding", None)
+    if isinstance(lp, mx.array) and lp.size and int(lp.max().item()) > 0:
+        return
+    try:
+        drafter.bind(batch.model)
+        seed_kv = drafter.make_cache()
+    except Exception:
+        _log.warning("seed streaming unavailable for this drafter; "
+                     "deferred seed", exc_info=True)
+        return
+    ctx = {
+        "kv": seed_kv,
+        "len": 0,
+        "active": True,
+        # Retain chunk hiddens alongside streaming whenever an L0 store can
+        # arm: the store needs full-prompt hidden. APC off => no retention
+        # while streaming (the capture-memory win lands in that config).
+        "retain": _get_spec_prefix_cache(batch.model) is not None,
+        "retained_from": 0,
+    }
+    batch._mtp_seed_ctx = ctx
+    batch.prompt_cache[0]._kq_seed_stream = ctx
+
+
+def _zero_pad_rows(arr, rows: int):
+    pad = mx.zeros((rows - arr.shape[0],) + tuple(arr.shape[1:]), dtype=arr.dtype)
+    return mx.concatenate([arr, pad], axis=0)
+
+
+def _widen_prompt_rope_state(batch, prompt_kwargs: dict) -> dict:
+    """Continuous-batch admission can grow the spec prompt batch (and decode
+    forwards run at other widths) between chunks; the target caches text
+    mrope deltas at the old width and only slices down, never widens, so the
+    next chunk forward dies on offsets(B) + rope_deltas(B_old) broadcast.
+    Text rows have delta 0, so zero-pad both delta sources to the live width
+    (decode-loop twin of this guard: speculative.py injection path)."""
+    b = batch._input_ids.shape[0]
+    rd = prompt_kwargs.get("rope_deltas")
+    if rd is not None and rd.shape[0] < b:
+        prompt_kwargs = dict(prompt_kwargs)
+        prompt_kwargs["rope_deltas"] = _zero_pad_rows(rd, b)
+    lm = getattr(batch.model, "language_model", batch.model)
+    rd = getattr(lm, "_rope_deltas", None)
+    if rd is not None and rd.shape[0] < b:
+        lm._rope_deltas = _zero_pad_rows(rd, b)
+    return prompt_kwargs
+
+
+def install_full_prompt_mtp_prefill() -> None:
+    """Retain full-prompt hidden through the BatchGenerator MTP prefill so the
+    native head teacher-forces the whole prompt into its KV (llama parity).
+
+    mlx-vlm's ``PromptProcessingBatch`` chunks prefill: intermediate chunks
+    (``prompt_step``) discard the model output (only KV-cache side-effects
+    survive), then ``generate()`` runs the final chunk with
+    ``return_hidden=True``.  The MTP drafter thus only sees hidden for that
+    last chunk -- often 1 token -- and acceptance erodes at depth.
+
+    This patch makes ``prompt_step`` also request ``return_hidden=True`` on
+    MTP batches, accumulating per-chunk hidden in ``_mtp_chunk_hiddens``.
+    ``generate()`` then concatenates them with the final chunk's hidden so
+    ``speculative_hidden_state`` returns full-prompt hidden to the drafter.
+
+    Also installs the owned-path APC surface: the L0 SpecPrefixCache
+    (whole-prompt KV + hidden, in-memory) plus the L1 shared APCManager
+    (exact / block / disk tiers -- the same manager the stock
+    non-speculative path uses, reached via ``model._kq_apc_manager``, which
+    ``_install_apc_manager_stash`` captures at BatchGenerator construction).
+    Kill switch for both tiers: ``GMLX_SPEC_APC=0``.
+
+    Idempotent.  Only MTP batches (``self.draft_kind == "mtp"``) are affected;
+    eagle3 / dflash keep the stock path.
+    """
+    from mlx_vlm.generate.ar import PromptProcessingBatch
+
+    # L1 plumbing is idempotent on its own flags, so it installs (or
+    # repairs) even when the prefill override is already in place.
+    _bind_l1_view()
+    # The L1 disk tier serializes through mlx-vlm's DiskBlockStore, which
+    # has no arm for QSAKVCache and refuses the whole exact snapshot.
+    # Installed here as well as in serve patches so embedded/test users of
+    # the spec engine get disk APC.
+    from gmlx.cache.apc_qsa import install_qsa_apc_support
+    install_qsa_apc_support()
+    _install_apc_manager_stash()
+    _install_ckpt_checkpoint_store()
+    _install_plain_ckpt_decode()
+    _install_exact_anchor_pick()
+
+    if getattr(PromptProcessingBatch, _FULL_PREFILL_FLAG, False):
+        return
+
+    _orig_prompt_step = PromptProcessingBatch.prompt_step
+    _orig_generate = PromptProcessingBatch.generate
+    _orig_init = PromptProcessingBatch.__init__
+
+    def _resolve_mtp_prefill_step() -> int:
+        # Honor the serve path's PREFILL_STEP_SIZE env override
+        # (mlx_vlm.server.generation.get_prefill_step_size) so MTP prefill
+        # can be chunked smaller to cap peak memory.
+        from mlx_vlm.generate.ar import DEFAULT_PREFILL_STEP_SIZE
+
+        return int(os.environ.get("PREFILL_STEP_SIZE", DEFAULT_PREFILL_STEP_SIZE))
+
+    def _mtp_init(self, *args, **kwargs) -> None:
+        _orig_init(self, *args, **kwargs)
+        # Re-enable chunked prefill.  Stock mlx-vlm nulls prefill_step_size
+        # for speculative models because intermediate chunks discard hidden;
+        # our prompt_step captures it, so the gate no longer applies.
+        # Restoring at construction (not first prompt_step) matters: the
+        # scheduler consults needs_processing() first, and with a None step
+        # an APC-less deep prompt would one-shot the whole prefill.
+        if (
+            getattr(self, "draft_kind", None) == "mtp"
+            and self.prefill_step_size is None
+        ):
+            self.prefill_step_size = _resolve_mtp_prefill_step()
+        # Stock (non-speculative) batches get the checkpoint tier here:
+        # lookup, prefix trim, cursor arming, retirement stash.
+        if getattr(self, "draft_kind", None) is None and not _SPEC_APC_DISABLED:
+            try:
+                # Ckpt-active hybrids under kvarn convert their stock B=1
+                # single-stream caches in place before arming, so the
+                # layout signature, lookup, and every store see the same
+                # kvarn classes. Must run here: the outer batch rebuild
+                # (kvarn_serve) would install batch classes the tier is
+                # blind to; after conversion its shared decline predicate
+                # trips instead. kwargs is load-bearing -- stock init
+                # consumes and drops the scheme/bits constructor params.
+                manager, mode = _resolve_l1(self.model)
+                if manager is not None:
+                    from gmlx.cache.kvarn_serve import ensure_ppb_kvarn
+
+                    ensure_ppb_kvarn(
+                        self, kwargs,
+                        ckpt_active=_ckpt_active(
+                            self.model, mode, int(manager.block_size)))
+            except Exception:
+                _log.warning(
+                    "kvarn ckpt cache conversion failed; continuing stock",
+                    exc_info=True,
+                )
+            try:
+                _plain_ckpt_init(self)
+                _plain_anchor_init(self)
+            except Exception:
+                _log.warning(
+                    "APC plain ckpt init failed; continuing stock", exc_info=True
+                )
+
+    def _mtp_prompt_step(self) -> int:
+        if self.draft_kind != "mtp":
+            return _orig_prompt_step(self)
+        # cb_phase flips fine prefill caps by wrapping the stock
+        # prompt_step, but this body replaces it for MTP batches, so the
+        # flip must happen here too: a multi-thousand-token chunk under
+        # the coarse decode caps keeps every layer's transients live in
+        # one command buffer and OOMs the GPU on deep prompts.
+        if os.environ.get("GMLX_CB_PHASE", "1") != "0":
+            from gmlx.serve.cb_phase import flip
+            flip("prefill")
+
+        if not hasattr(self, "_mtp_full_input_ids"):
+            if self.prefill_step_size is None:
+                self.prefill_step_size = _resolve_mtp_prefill_step()
+            # APC lookup (L0 then L1) + prefix trim + store arming.
+            _mtp_prefill_init(self)
+            _mtp_seed_stream_init(self)
+
+        if not self.needs_processing():
+            return 0
+
+        # Depth-decayed step: shrink only when this chunk's score transient
+        # would exceed the cap (see prefill_decay; keeps MoE weight
+        # amortization at shallow depth instead of a global small step).
+        step = prefill_decay.decayed_for_batch(self) or self._inputs_embeds.shape[1]
+        n = min(step, self._inputs_embeds.shape[1] - 1)
+
+        if not hasattr(self, "_mtp_padding_widened"):
+            self._mtp_padding_widened = True
+            for c in self.prompt_cache:
+                lp = getattr(c, "left_padding", None)
+                if isinstance(lp, mx.array) and lp.ndim > 0 and lp.size > 1:
+                    max_lp = int(lp.max().item())
+                    if max_lp >= n:
+                        n = min(max_lp + 1, self._inputs_embeds.shape[1] - 1)
+                    break
+
+        checkpoint_col = self._next_apc_checkpoint_column()
+        if checkpoint_col is not None:
+            n = min(n, checkpoint_col - self._processed_prompt_columns)
+        # Media requests ride this body too: keep image blocks whole (a
+        # boundary inside a block moves to its edge, see media_spans).
+        from gmlx.gen.media_spans import span_aware_prompt_n
+        n = span_aware_prompt_n(self, n)
+        # A final chunk under ~3 simdgroup tiles routes the projections
+        # through the skinny-M kernels, whose accumulation order seeds fp
+        # noise that stacked recurrent (GDN) layers amplify into
+        # first-token divergence. Absorb such a tail into this chunk so
+        # every chunk stays in the wide-GEMM regime. Checkpoint columns
+        # stay exact.
+        min_tail = env_int("GMLX_PREFILL_MIN_TAIL", 48)
+        if checkpoint_col is None and min_tail > 0:
+            rem1 = self._inputs_embeds.shape[1] - 1
+            tail = rem1 - n
+            if 0 < tail < min_tail:
+                n = rem1        # absorb: overshoot bounded by min_tail-1
+        if n <= 0:
+            return 0
+        prompt_kwargs = self._prompt_kwargs_for_step(n)
+        prompt_kwargs = _widen_prompt_rope_state(self, prompt_kwargs)
+        with lora_rows.published(getattr(self, "uids", [])):
+            out = self.model(
+                self._input_ids[:, :n],
+                cache=self.prompt_cache,
+                inputs_embeds=self._inputs_embeds[:, :n],
+                n_to_process=n,
+                return_hidden=True,
+                **prompt_kwargs,
+            )
+        chunk_hidden = out.hidden_states[-1]
+        # Seed streaming: teacher-force this chunk into the request-scoped
+        # head KV at the head's running offset. The shifted span for
+        # columns [c0, c0+n) is prompt[c0+1 : c0+n+1], always in range
+        # because generate() keeps at least one residual column (the n-1
+        # cap above). A failure or a widened batch stops streaming but
+        # keeps the partial KV: the owned round adopts it at its true
+        # offset and seeds only the remainder.
+        seed_ctx = getattr(self, "_mtp_seed_ctx", None)
+        streamed = False
+        if seed_ctx is not None and seed_ctx["active"]:
+            if int(self._input_ids.shape[0]) != 1:
+                seed_ctx["active"] = False
+            else:
+                c0 = int(self._processed_prompt_columns)
+                try:
+                    self.draft_model.seed_chunk(
+                        self._mtp_full_input_ids[:, c0 + 1:c0 + n + 1],
+                        chunk_hidden, seed_ctx["kv"])
+                    seed_ctx["len"] += n
+                    streamed = True
+                except Exception:
+                    _log.warning("seed streaming failed at column %d; "
+                                 "deferred seed for the remainder", c0,
+                                 exc_info=True)
+                    seed_ctx["active"] = False
+        # Teacher-forcing drafters (native MTP heads) seed their KV from the
+        # whole prompt hidden, so every chunk is retained except when the
+        # chunk just streamed and no L0 store is armed (nothing downstream
+        # reads it). Shared-KV drafters (gemma-4 assistant) read only the
+        # last position: keeping just the newest chunk caps capture memory
+        # at O(chunk) instead of O(prompt), GBs at deep context.
+        if callable(getattr(self.draft_model, "prefill_from_target_hidden", None)):
+            if streamed and not seed_ctx["retain"]:
+                pass
+            else:
+                if (seed_ctx is not None and not seed_ctx["retain"]
+                        and not self._mtp_chunk_hiddens):
+                    # Streaming stopped mid-request with no retention so
+                    # far: the retained span starts here, not at column 0.
+                    seed_ctx["retained_from"] = int(
+                        self._processed_prompt_columns)
+                self._mtp_chunk_hiddens.append(chunk_hidden)
+                # Window-limited heads can't use context beyond the trailing
+                # hidden_capture_limit positions; an uncapped capture pins the
+                # whole prompt's hidden (GBs at deep context). The drafter's
+                # teacher-force self-aligns to the trailing h_len positions.
+                limit = getattr(self.draft_model, "hidden_capture_limit", None)
+                if limit:
+                    total = sum(int(h.shape[1]) for h in self._mtp_chunk_hiddens)
+                    if total > limit:
+                        merged = (self._mtp_chunk_hiddens[0]
+                                  if len(self._mtp_chunk_hiddens) == 1
+                                  else mx.concatenate(self._mtp_chunk_hiddens, axis=1))
+                        self._mtp_chunk_hiddens = [merged[:, -limit:]]
+        else:
+            self._mtp_chunk_hiddens = [chunk_hidden]
+        mx.eval([c.state for c in self.prompt_cache] + [chunk_hidden]
+                + ([c.state for c in seed_ctx["kv"]] if streamed else []))
+        self._processed_prompt_columns += n
+        # The ckpt cursor rides the wrapped stock store (see
+        # _install_ckpt_checkpoint_store).
+        self._store_apc_exact_checkpoints()
+        self._inputs_embeds = self._inputs_embeds[:, n:]
+        self._input_ids = self._input_ids[:, n:]
+        for k in self._prompt_length_aware_keys:
+            self._prompt_kwargs[k] = self._prompt_kwargs[k][:, n:, ...]
+        mx.clear_cache()
+        return n
+
+    def _mtp_generate(
+        self, sampler, stop_criteria, compute_logprobs=True, top_logprobs_k=0
+    ):
+        if self.draft_kind == "mtp":
+            # Short prompts never enter prompt_step (chunked prefill is not
+            # needed), so the APC lookup/store arming runs here instead.
+            _mtp_prefill_init(self)
+        result = _orig_generate(
+            self,
+            sampler,
+            stop_criteria,
+            compute_logprobs=compute_logprobs,
+            top_logprobs_k=top_logprobs_k,
+        )
+        from mlx_vlm.generate.ar import SpeculativeGenerationBatch
+
+        if self.draft_kind != "mtp" or not isinstance(
+            result, SpeculativeGenerationBatch
+        ):
+            # Stock-path ckpt batches store the full prompt here, the
+            # moment the MTP path stores it at rounds entry: prefill just
+            # finished, the first token is out, its KV not yet appended.
+            if (
+                getattr(self, "_kq_ckpt_armed", False)
+                and getattr(self, "draft_kind", None) is None
+            ):
+                try:
+                    cache = getattr(result, "prompt_cache", None) or []
+                    stash = getattr(cache[0], "_kq_apc_retire", None) if cache else None
+                    if stash is not None and stash.get("mode") == "ckpt":
+                        from gmlx.cache.snapshot import (
+                            ckpt_full_store_redundant,
+                            ckpt_store,
+                        )
+                        m = stash.get("apc_meta")
+                        if ckpt_full_store_redundant(m):
+                            _log.info("APC ckpt post-prefill store "
+                                      "skipped: render-stable boundary "
+                                      "landed")
+                        elif ckpt_store(
+                                stash["manager"], stash["full_ids"], cache,
+                                extra_hash=int(stash.get("extra_hash", 0))):
+                            if m is not None:
+                                m.setdefault(
+                                    "ckpt_stored_boundaries", []
+                                ).append(len(stash["full_ids"]))
+                except Exception:
+                    _log.warning(
+                        "APC plain post-prefill store failed; continuing", exc_info=True
+                    )
+            return result
+        chunk_hiddens = getattr(self, "_mtp_chunk_hiddens", None)
+        full_ids = getattr(self, "_mtp_full_input_ids", None)
+        l1_prefix = int(getattr(self, "_mtp_l1_prefix_len", 0) or 0)
+        if not chunk_hiddens:
+            # No captured chunks: the whole (remaining) prompt went through
+            # the final generate forward, so stock prompt_tokens/hidden are
+            # already an aligned pair (suffix-only on an L1 hit) and
+            # result.hidden needs no rebuild; with seed streaming and no
+            # retention, result.hidden is already the residual unstreamed
+            # tail (retention accompanies an armed L0 store, so none can
+            # fire here). The L0 store below must still run for the
+            # single-shot case: arch prefill profiles can raise the step
+            # past typical prompt lengths (qwen4exp defaults to 8192), so
+            # sub-step prompts land here and still need their warm-start
+            # entry.
+            full_hidden = result.hidden
+        else:
+            parts = chunk_hiddens + [result.hidden]
+            full_hidden = mx.concatenate(parts, axis=1)
+        seed_ctx = getattr(self, "_mtp_seed_ctx", None)
+        seed_len = int(seed_ctx["len"]) if seed_ctx else 0
+        if chunk_hiddens:
+            if seed_len > 0:
+                # Columns [0, seed_len) are already teacher-forced into
+                # the streamed head KV; hand the owned round only the
+                # residual hidden so its seed call covers exactly the
+                # unstreamed tail at the adopted offset. full_hidden (the
+                # retained span) still feeds the L0 store below, which
+                # needs the whole prompt.
+                rfrom = int(seed_ctx.get("retained_from") or 0)
+                result.hidden = full_hidden[:, seed_len - rfrom:]
+            else:
+                result.hidden = full_hidden
+        if chunk_hiddens and full_ids is not None:
+            # On an L1 hit the captured hidden covers only the forwarded
+            # suffix, so hand the drafter the matching suffix tokens: the
+            # teacher-forcing (token, hidden) pair must stay positionally
+            # aligned. The missing prefix can only affect draft acceptance,
+            # never correctness -- verify catches every draft.
+            result.prompt_tokens = (
+                full_ids[:, l1_prefix:] if l1_prefix > 0 else full_ids
+            )
+
+        # APC L0 store: cache this request's target KV + hidden so a
+        # future request sharing this token prefix skips re-prefill.
+        # Uses result.prompt_cache (SpecBatch owns the cache now),
+        # not self.prompt_cache (empty after _orig_generate).
+        #
+        # B=1 only -- same prefill-granularity gate as the lookup.
+        # The stored single-row snapshot is valid for injection into
+        # a B>1 batch: SpecPrefixCache.restore writes into a fresh
+        # single-row prompt_cache, and BatchKVCache.extend merges
+        # it at the correct per-row offset.
+        #
+        # Skipped on an L1 hit: hidden covers only the suffix, and L0
+        # entries pair full-prompt keys with full-prompt hidden.
+        b = int(full_hidden.shape[0]) if full_ids is not None else 0
+        # With streaming, full_hidden covers the whole prompt only when
+        # retention ran from column 0: after a mid-request streaming stop
+        # the retained span starts past column 0, and with no retention at
+        # all full_hidden is just the residual tail. Neither must ever be
+        # stored as a full-prompt entry.
+        full_covers_prompt = seed_len == 0 or (
+            bool(chunk_hiddens)
+            and int(seed_ctx.get("retained_from") or 0) == 0)
+        spec_cache = (
+            _get_spec_prefix_cache(self.model)
+            if b == 1 and l1_prefix == 0 and full_covers_prompt
+            and not getattr(self, "_mtp_upstream_warm", False) else None
+        )
+        if spec_cache is not None and full_ids is not None:
+            # Window-limited heads only use the trailing capture window;
+            # chunked prefill already trimmed, single-shot must match (an
+            # uncapped entry pins the whole prompt's hidden for nothing).
+            limit = getattr(self.draft_model, "hidden_capture_limit", None)
+            store_hidden = (full_hidden if not limit
+                            else full_hidden[:, -int(limit):])
+            spec_cache.store(full_ids, result.prompt_cache, store_hidden)
+            _log.info(
+                "APC store: tokens=%d layers=%d",
+                int(full_ids.shape[1]),
+                len(result.prompt_cache),
+            )
+        else:
+            _log.debug(
+                "APC store skipped: b=%d l1_prefix=%d upstream_warm=%s "
+                "full_ids=%s",
+                b, l1_prefix,
+                getattr(self, "_mtp_upstream_warm", False),
+                "set" if full_ids is not None else "None",
+            )
+
+        return result
+
+    PromptProcessingBatch.__init__ = _mtp_init
+    PromptProcessingBatch.prompt_step = _mtp_prompt_step
+    PromptProcessingBatch.generate = _mtp_generate
+    setattr(PromptProcessingBatch, _FULL_PREFILL_FLAG, True)
+    if _SPEC_APC_DISABLED:
+        apc_status = "off"
+    elif _L1_BOUND[0]:
+        apc_status = "on: L0+L1"
+    else:
+        apc_status = "on: L0 only"
+    _debug_note(
+        f"[mtp] serve prefill: full-prompt hidden capture installed (APC {apc_status})"
+    )
diff --git a/gmlx/spec/speculative.py b/gmlx/spec/speculative.py
index 9cea07d2..21ef3aea 100644
--- a/gmlx/spec/speculative.py
+++ b/gmlx/spec/speculative.py
@@ -1520,7 +1520,7 @@ def _pop_drafter_warm(prompt_cache: list) -> list | None:
 
 def _pop_seed_stream(prompt_cache: list) -> dict | None:
     """Detach the streamed-seed context ({"kv", "len", ...}) stashed by the
-    prefill's seed streaming (engine._mtp_seed_stream_init) on the request's
+    prefill's seed streaming (mtp_prefill._mtp_seed_stream_init) on the request's
     first cache entry. Same request-scoped discipline and pop-before-buffer
     timing as the drafter warm sidecar above."""
     if not prompt_cache:
@@ -1825,7 +1825,7 @@ def _lift_injected_cache(cache, other):
 
         return kvarn_batch_row(cache, other)
     if _packed_single(other):
-        from gmlx.spec.engine import kvarn_lift_cache, lift_single_cache
+        from gmlx.spec.kv_quant import kvarn_lift_cache, lift_single_cache
 
         if getattr(other, "kv_quant_scheme", None) == "kvarn":
             # An fp16 batch host (the batched arm dropped): the kvarn row
@@ -1935,7 +1935,7 @@ def _lift_live_cache(cache):
             cache.caches = tuple(_lift_live_cache(m) for m in members)
         return cache
     if _packed_single(cache):
-        from gmlx.spec.engine import lift_single_cache
+        from gmlx.spec.kv_quant import lift_single_cache
 
         return lift_single_cache(cache)
     if _batch_capable(cache) or not callable(getattr(type(cache), "merge", None)):
diff --git a/gmlx/tui/chat.py b/gmlx/tui/chat.py
index 0299a5f3..05a75a0c 100644
--- a/gmlx/tui/chat.py
+++ b/gmlx/tui/chat.py
@@ -2891,7 +2891,7 @@ def _backend_mtp_text(args, kv_kwargs) -> _ChatBackend:
     mtp_kv_policy = None
     if kv_kwargs.get("kv_bits") is not None:
         from gmlx.cache.kv_policy import resolve_and_report
-        from gmlx.spec.engine import mtp_kv_decline
+        from gmlx.spec.kv_quant import mtp_kv_decline
         from gmlx.spec.speculative import use_owned_engine
 
         lm = b.model.language_model
diff --git a/gmlx/upstream/seams.py b/gmlx/upstream/seams.py
index 5f22a260..0570ab0d 100644
--- a/gmlx/upstream/seams.py
+++ b/gmlx/upstream/seams.py
@@ -169,7 +169,7 @@ class Seam:
     Seam("mlx_vlm.generate.ar", "_extend_cache",
          "cascade_sdpa.install_cascade_stamp (stamp carry across the "
          "B=1-to-batch merge lift on admission)", critical=True),
-    # --- speculative / AR batch engine (spec_engine owns these methods) ---
+    # --- speculative / AR batch engine (the spec package owns these methods) ---
     Seam("mlx_vlm.generate.ar", "BatchGenerator.__init__",
          "spec_engine._install_apc_manager_stash + kvarn_serve APC gate",
          critical=True),
@@ -178,32 +178,32 @@ class Seam:
          critical=True),
     Seam("mlx_vlm.generate.ar",
          "PromptProcessingBatch._store_apc_exact_checkpoints",
-         "spec_engine._install_ckpt_checkpoint_store (ckpt cursor rides "
+         "ckpt._install_ckpt_checkpoint_store (ckpt cursor rides "
          "the stock store)", critical=True),
     Seam("mlx_vlm.generate.ar", "PromptProcessingBatch.__init__",
-         "spec_engine.install_full_prompt_mtp_prefill (prefill-step "
+         "mtp_prefill.install_full_prompt_mtp_prefill (prefill-step "
          "restore + stock-path ckpt arming)", critical=True),
     Seam("mlx_vlm.generate.ar", "GenerationBatch._step",
-         "spec_engine._install_plain_ckpt_decode (token accounting + "
+         "ckpt._install_plain_ckpt_decode (token accounting + "
          "decode-time snapshots)", critical=True),
     Seam("mlx_vlm.generate.ar", "GenerationBatch.filter",
-         "spec_engine._install_plain_ckpt_decode (B=1 retirement at row "
+         "ckpt._install_plain_ckpt_decode (B=1 retirement at row "
          "exit)", critical=True),
     Seam("mlx_vlm.generate.ar", "PromptProcessingBatch.prompt_step",
-         "spec_engine.install_full_prompt_mtp_prefill", critical=True),
+         "mtp_prefill.install_full_prompt_mtp_prefill", critical=True),
     Seam("mlx_vlm.generate.ar", "PromptProcessingBatch.generate",
-         "spec_engine.install_full_prompt_mtp_prefill; "
+         "mtp_prefill.install_full_prompt_mtp_prefill; "
          "seed_rows.install_per_request_seed (row-uid publish); "
          "server_patches.mtp_thinking (thinking-hook transport, outermost)",
          critical=True),
     Seam("mlx_vlm.generate.ar", "SpeculativeGenerationBatch.next",
-         "spec_engine.install_continuous_batch_admission", critical=True),
+         "admission.install_continuous_batch_admission", critical=True),
     Seam("mlx_vlm.generate.ar", "SpeculativeGenerationBatch.filter",
-         "spec_engine._filter_with_release (per-row release rides the "
+         "admission._filter_with_release (per-row release rides the "
          "mark-finished contract; the rounds generator sheds via "
          "stop_check)", critical=True),
     Seam("mlx_vlm.generate.ar", "SpeculativeGenerationBatch.__len__",
-         "spec_engine._len_with_promotion (patched semantics are D-3's "
+         "admission._len_with_promotion (patched semantics are D-3's "
          "hazard; decision modules count rows via _orig_len only)",
          critical=True),
     Seam("mlx_vlm.generate.ar", "GenerationBatch._eval_pending_state",
@@ -215,7 +215,7 @@ class Seam:
     Seam("mlx_vlm.generate.ar", "run_speculative_server_rounds",
          "spec_engine.install_owned_spec_engine", critical=True),
     Seam("mlx_vlm.speculative.utils", "make_speculative_prompt_cache",
-         "spec_engine.install_spec_kv_quant (B=1 KV_BITS/kvarn)",
+         "kv_quant.install_spec_kv_quant (B=1 KV_BITS/kvarn)",
          critical=True),
     Seam("mlx_vlm.generate.ar", "BatchGenerator._apc_pick_for",
          "spec_engine._bind_l1_view (L1 APC helpers)"),
@@ -376,7 +376,7 @@ class Seam:
     Seam("mlx_vlm.models.cache", "BatchKVCache",
          "mtp_drafter / cache_snapshot row round-trip"),
     Seam("mlx_vlm.models.cache", "BatchKVCache.filter",
-         "spec_engine per-row release + governor retire (physical row "
+         "admission per-row release + governor retire (physical row "
          "drop through the cache's own filter)", critical=True),
     Seam("mlx_vlm.models.cache", "BatchKVCache.extract",
          "governor orange retire (contiguous single-row extract before "
diff --git a/tests/cache/test_apc_kv_quant_compose.py b/tests/cache/test_apc_kv_quant_compose.py
index 5a03b127..1f634e2a 100644
--- a/tests/cache/test_apc_kv_quant_compose.py
+++ b/tests/cache/test_apc_kv_quant_compose.py
@@ -8,7 +8,7 @@
 
 import mlx.core as mx
 
-import gmlx.spec.engine as spec_engine
+import gmlx.spec.engine as engine
 
 
 def _stamped_model(bits=8, group=32):
@@ -31,12 +31,12 @@ def _stamped_model(bits=8, group=32):
 def test_live_kv_quant_config_off_without_stamp(monkeypatch):
     # env alone never decides the warm merge
     monkeypatch.setenv("KV_BITS", "8")
-    assert spec_engine._live_kv_quant_config() is None
-    assert spec_engine._live_kv_quant_config(object()) is None
+    assert engine._live_kv_quant_config() is None
+    assert engine._live_kv_quant_config(object()) is None
 
 
 def test_live_kv_quant_config_reads_stamped_policy():
-    cfg = spec_engine._live_kv_quant_config(_stamped_model(8, 32))
+    cfg = engine._live_kv_quant_config(_stamped_model(8, 32))
     assert cfg is not None
     assert float(cfg["bits"]) == 8.0 and int(cfg["group_size"]) == 32
 
@@ -54,7 +54,7 @@ def test_warm_merge_requantizes_float_row():
         k = mx.random.normal((1, 2, 64, 64))
         c.update_and_fetch(k, k)
         row.append(c)
-    cfg = spec_engine._live_kv_quant_config(_stamped_model(8, 32))
+    cfg = engine._live_kv_quant_config(_stamped_model(8, 32))
     warm, n = apc.make_warm_batch_exact_cache_multi(
         [row], prefix_lens=[64], kv_quant_config=cfg)
     assert warm is not None and n == 64
diff --git a/tests/cache/test_apc_pooling.py b/tests/cache/test_apc_pooling.py
index c57b2385..72c68adc 100644
--- a/tests/cache/test_apc_pooling.py
+++ b/tests/cache/test_apc_pooling.py
@@ -151,7 +151,7 @@ def test_spec_apc_kill_switch_strips_manager_from_stock(monkeypatch):
     import importlib
     from types import SimpleNamespace
 
-    import gmlx.spec.engine as spec_engine
+    import gmlx.spec.engine as engine
 
     ar = importlib.import_module("mlx_vlm.generate.ar")
     seen = {}
@@ -161,16 +161,16 @@ def __init__(self, model, processor, **kwargs):
             seen.update(kwargs)
 
     monkeypatch.setattr(ar, "BatchGenerator", _ProbeBG)
-    spec_engine._install_apc_manager_stash()
+    engine._install_apc_manager_stash()
     mgr = object()
 
-    monkeypatch.setattr(spec_engine, "_SPEC_APC_DISABLED", True)
+    monkeypatch.setattr(engine, "_SPEC_APC_DISABLED", True)
     model = SimpleNamespace()
     ar.BatchGenerator(model, None, draft_model=object(), apc_manager=mgr)
     assert seen["apc_manager"] is None
     assert model._kq_apc_manager is None
 
-    monkeypatch.setattr(spec_engine, "_SPEC_APC_DISABLED", False)
+    monkeypatch.setattr(engine, "_SPEC_APC_DISABLED", False)
     seen.clear()
     model = SimpleNamespace()
     ar.BatchGenerator(model, None, draft_model=object(), apc_manager=mgr)
@@ -178,7 +178,7 @@ def __init__(self, model, processor, **kwargs):
     assert model._kq_apc_manager is mgr
 
     # Non-speculative batches are outside the spec kill switch's scope.
-    monkeypatch.setattr(spec_engine, "_SPEC_APC_DISABLED", True)
+    monkeypatch.setattr(engine, "_SPEC_APC_DISABLED", True)
     seen.clear()
     ar.BatchGenerator(SimpleNamespace(), None, apc_manager=mgr)
     assert seen["apc_manager"] is mgr
@@ -194,7 +194,7 @@ def test_kv_bits_apc_optout_warns_at_boot(monkeypatch, caplog):
     import logging
     from types import SimpleNamespace
 
-    import gmlx.spec.engine as spec_engine
+    import gmlx.spec.engine as engine
 
     ar = importlib.import_module("mlx_vlm.generate.ar")
 
@@ -206,8 +206,8 @@ def __init__(self, model, processor, **kwargs):
             self.apc_manager = mgr
 
     monkeypatch.setattr(ar, "BatchGenerator", _UpstreamLikeBG)
-    monkeypatch.setattr(spec_engine, "_SPEC_APC_DISABLED", False)
-    spec_engine._install_apc_manager_stash()
+    monkeypatch.setattr(engine, "_SPEC_APC_DISABLED", False)
+    engine._install_apc_manager_stash()
     mgr = object()
 
     with caplog.at_level(logging.WARNING, logger="gmlx.spec.engine"):
diff --git a/tests/cache/test_apc_schedule_invariant.py b/tests/cache/test_apc_schedule_invariant.py
index 6a4f3f30..7dcccfe6 100644
--- a/tests/cache/test_apc_schedule_invariant.py
+++ b/tests/cache/test_apc_schedule_invariant.py
@@ -16,7 +16,7 @@ class the apc-overhaul branch exists for (bug 1): every store position
 import pytest
 from mlx_vlm.apc import APCManager
 
-import gmlx.spec.engine as se
+import gmlx.spec.ckpt as ckpt
 import gmlx.cache.retire_key as retire_key
 from gmlx.cache.snapshot import (
     _ckpt_records,
@@ -89,11 +89,11 @@ def make(p, seed=0):
         _apc_meta=[meta], prompt_cache=None,
         model=SimpleNamespace(_kq_apc_ckpt_layout=tags),
         _row_real_tokens_processed=lambda i: 0)
-    se._ckpt_arm_schedule(batch, meta, len(ids), 0, 16)
+    ckpt._ckpt_arm_schedule(batch, meta, len(ids), 0, 16)
     for pos, _kind in list(meta["ckpt_boundaries"]):
         batch.prompt_cache = make(pos, seed=pos % 977)
         batch._row_real_tokens_processed = lambda i, b=pos: b
-        se._ckpt_mid_prefill_store(batch)
+        ckpt._ckpt_mid_prefill_store(batch)
     assert meta.get("checkpoint_done", not meta["ckpt_boundaries"]) or \
         meta.get("ckpt_boundaries") == []
     if not ckpt_full_store_redundant(meta):
diff --git a/tests/cache/test_ckpt_cursor.py b/tests/cache/test_ckpt_cursor.py
index 41843b3a..cb92c747 100644
--- a/tests/cache/test_ckpt_cursor.py
+++ b/tests/cache/test_ckpt_cursor.py
@@ -13,7 +13,9 @@
 import mlx.core as mx
 from mlx_vlm.apc import APCManager
 
+import gmlx.spec.ckpt as ckpt
 import gmlx.spec.engine as se
+import gmlx.spec.mtp_prefill as mtp_prefill
 from gmlx.cache.snapshot import ckpt_lookup
 from gmlx.spec.speculative import _sidecar_boundary
 
@@ -32,7 +34,7 @@ def _positions(bounds):
 
 
 def test_cursor_grid_and_terminal():
-    bounds, terminal, interval = se._ckpt_cursor_init(
+    bounds, terminal, interval = ckpt._ckpt_cursor_init(
         _batch(), guard=27000, restored=0, block_size=16)
     assert (terminal, interval) == (26624, 4096)
     assert _positions(bounds) == [4096, 8192, 12288, 16384, 20480,
@@ -40,17 +42,17 @@ def test_cursor_grid_and_terminal():
 
 
 def test_cursor_skips_restored_prefix():
-    bounds, terminal, interval = se._ckpt_cursor_init(
+    bounds, terminal, interval = ckpt._ckpt_cursor_init(
         _batch(), guard=27000, restored=8192, block_size=16)
     assert _positions(bounds)[0] == 12288
     # Restored past the terminal: nothing left to checkpoint.
-    assert se._ckpt_cursor_init(
+    assert ckpt._ckpt_cursor_init(
         _batch(), guard=27000, restored=26624, block_size=16) == ([], 0, 0)
 
 
 def test_cursor_interval_snaps_up_to_chunk_grid(monkeypatch):
     monkeypatch.setenv("GMLX_APC_CKPT_INTERVAL", "1000")
-    bounds, terminal, interval = se._ckpt_cursor_init(
+    bounds, terminal, interval = ckpt._ckpt_cursor_init(
         _batch(), guard=27000, restored=0, block_size=16)
     assert interval == 2048                       # never below one chunk
     assert _positions(bounds)[0] == 2048
@@ -58,14 +60,14 @@ def test_cursor_interval_snaps_up_to_chunk_grid(monkeypatch):
 
 def test_cursor_zero_interval_is_terminal_only(monkeypatch):
     monkeypatch.setenv("GMLX_APC_CKPT_INTERVAL", "0")
-    bounds, terminal, interval = se._ckpt_cursor_init(
+    bounds, terminal, interval = ckpt._ckpt_cursor_init(
         _batch(), guard=27000, restored=0, block_size=16)
     assert (terminal, interval) == (26624, 0)
     assert _positions(bounds) == [26624]
 
 
 def test_cursor_no_step_uses_block_grid():
-    bounds, terminal, interval = se._ckpt_cursor_init(
+    bounds, terminal, interval = ckpt._ckpt_cursor_init(
         _batch(step=None), guard=100, restored=0, block_size=16)
     assert (terminal, interval) == (96, 4096)
     assert _positions(bounds) == [96]
@@ -78,7 +80,7 @@ def _arm_meta(n, tags, restored=0, guard=None, step=2048):
         prefill_step_size=step,
         model=SimpleNamespace(_kq_apc_ckpt_layout=tuple(tags)))
     meta = {"full_input_ids": list(range(n))}
-    se._ckpt_arm_schedule(batch, meta, guard if guard is not None else n,
+    ckpt._ckpt_arm_schedule(batch, meta, guard if guard is not None else n,
                           restored, block_size=16)
     return meta
 
@@ -254,7 +256,7 @@ def test_cursor_advances_and_latches():
         batch.prompt_cache = make_hybrid_cache(boundary, seed=boundary)
         batch._row_real_tokens_processed = (
             lambda idx, b=boundary: b)
-        se._ckpt_mid_prefill_store(batch)
+        ckpt._ckpt_mid_prefill_store(batch)
         assert meta["ckpt_last_stored"] == boundary
     assert meta.get("checkpoint_done") is True
     # Strip-on-extend keeps the newest two boundaries hittable plus the
@@ -270,7 +272,7 @@ def test_cursor_off_boundary_chunk_is_a_noop():
     ids = list(range(500, 500 + 120))
     batch, meta = _armed_batch(man, ids, first=32, terminal=96, interval=32)
     batch._row_real_tokens_processed = lambda idx: 24
-    se._ckpt_mid_prefill_store(batch)
+    ckpt._ckpt_mid_prefill_store(batch)
     assert meta["checkpoint_len"] == 32 and meta["ckpt_last_stored"] == 0
 
 
@@ -283,11 +285,11 @@ def test_cursor_advances_past_failed_store():
     # Cache offset 48 != boundary 32: the store's offset guard declines.
     batch.prompt_cache = make_hybrid_cache(48)
     batch._row_real_tokens_processed = lambda idx: 32
-    se._ckpt_mid_prefill_store(batch)
+    ckpt._ckpt_mid_prefill_store(batch)
     assert meta["checkpoint_len"] == 64 and meta["ckpt_last_stored"] == 0
     batch.prompt_cache = make_hybrid_cache(64)
     batch._row_real_tokens_processed = lambda idx: 64
-    se._ckpt_mid_prefill_store(batch)
+    ckpt._ckpt_mid_prefill_store(batch)
     assert meta.get("checkpoint_done") and meta["ckpt_last_stored"] == 64
 
 
@@ -314,7 +316,7 @@ def test_stock_store_suppressed_by_advance():
         batch.prompt_cache = make_hybrid_cache(boundary, seed=boundary)
         batch._row_real_tokens_processed = (
             lambda idx, b=boundary: b)
-        se._ckpt_mid_prefill_store(batch)   # the two adjacent lines from
+        ckpt._ckpt_mid_prefill_store(batch)   # the two adjacent lines from
         stock_store(batch)                  # _mtp_prompt_step, same order
     assert calls == []
     assert meta.get("checkpoint_done") is True
@@ -353,14 +355,14 @@ def test_second_request_hits_first_requests_checkpoint():
         _inputs_embeds=mx.zeros((1, p_a, 4)),
         prompt_cache=make_hybrid_cache(p_a), _prompt_kwargs={},
         prefill_step_size=32)
-    se._mtp_prefill_init(batch_a)
+    mtp_prefill._mtp_prefill_init(batch_a)
     assert getattr(batch_a, "_kq_ckpt_armed", False)
     meta = batch_a._apc_meta[0]
     cl = int(meta["checkpoint_len"])
     assert cl == boundary                   # guard-trimmed terminal grid
     batch_a.prompt_cache = make_hybrid_cache(cl, seed=1)
     batch_a._row_real_tokens_processed = lambda idx: cl
-    se._ckpt_mid_prefill_store(batch_a)
+    ckpt._ckpt_mid_prefill_store(batch_a)
     assert meta["ckpt_last_stored"] == cl
 
     # Request B shares the first `boundary` tokens, diverges after.
@@ -370,7 +372,7 @@ def test_second_request_hits_first_requests_checkpoint():
         _inputs_embeds=mx.zeros((1, len(ids_b), 4)),
         prompt_cache=make_hybrid_cache(len(ids_b)), _prompt_kwargs={},
         prefill_step_size=32, _prompt_length_aware_keys=())
-    se._mtp_prefill_init(batch_b)
+    mtp_prefill._mtp_prefill_init(batch_b)
     assert batch_b._mtp_l1_prefix_len == boundary
 
 
diff --git a/tests/cache/test_ckpt_decode_lcp.py b/tests/cache/test_ckpt_decode_lcp.py
index ddbab3ba..f444d89b 100644
--- a/tests/cache/test_ckpt_decode_lcp.py
+++ b/tests/cache/test_ckpt_decode_lcp.py
@@ -339,7 +339,7 @@ def test_skeleton_disk_flag(monkeypatch):
 
 
 def test_cursor_skeleton_policy(monkeypatch):
-    import gmlx.spec.engine as spec_engine
+    import gmlx.spec.ckpt as ckpt
 
     seen = []
 
@@ -359,6 +359,6 @@ def rec_store(manager, ids, cache, *, extra_hash=0, skeleton_disk=True,
         _kq_ckpt_armed=True, _apc_manager=man, _apc_meta=[meta],
         prompt_cache=[], model=SimpleNamespace(_kq_apc_ckpt_layout=tags),
         _row_real_tokens_processed=lambda i: meta["checkpoint_len"])
-    spec_engine._ckpt_mid_prefill_store(batch)   # boundary 32: interval
-    spec_engine._ckpt_mid_prefill_store(batch)   # boundary 64: terminal
+    ckpt._ckpt_mid_prefill_store(batch)   # boundary 32: interval
+    ckpt._ckpt_mid_prefill_store(batch)   # boundary 64: terminal
     assert seen == [(32, False), (64, True)]
diff --git a/tests/cache/test_ckpt_kvarn_serve.py b/tests/cache/test_ckpt_kvarn_serve.py
index 30d36fa5..b5515d90 100644
--- a/tests/cache/test_ckpt_kvarn_serve.py
+++ b/tests/cache/test_ckpt_kvarn_serve.py
@@ -19,7 +19,8 @@
 import mlx.core as mx
 import pytest
 
-import gmlx.spec.engine as spec_engine
+import gmlx.spec.ckpt as ckpt
+import gmlx.spec.engine as engine
 from gmlx.cache.apc_manager import GmlxAPCManager
 from gmlx.cache.compat import runtime_cache_module
 from gmlx.cache.snapshot import ckpt_lookup, ckpt_store
@@ -152,16 +153,16 @@ def test_layout_live_ignores_model_memo():
     model = SimpleNamespace(_kq_apc_ckpt_layout=("kv", "arr"))
     b = SimpleNamespace(model=model,
                         prompt_cache=[KVarNKVCache(), ArraysCache(size=2)])
-    assert spec_engine._ckpt_layout_live(b) == (KVARN_TAG, "arr")
+    assert engine._ckpt_layout_live(b) == (KVARN_TAG, "arr")
     b2 = SimpleNamespace(model=model,
                          prompt_cache=[KVCache(), ArraysCache(size=2)])
-    assert spec_engine._ckpt_layout_live(b2) == ("kv", "arr")
+    assert engine._ckpt_layout_live(b2) == ("kv", "arr")
 
 
 def test_layout_live_falls_back_to_model_probe():
     model = SimpleNamespace(_kq_apc_ckpt_layout=("kv", "arr"))
     b = SimpleNamespace(model=model, prompt_cache=None)
-    assert spec_engine._ckpt_layout_live(b) == ("kv", "arr")
+    assert engine._ckpt_layout_live(b) == ("kv", "arr")
 
 
 def test_layout_live_unsupported_refuses_all_records():
@@ -170,8 +171,8 @@ def test_layout_live_unsupported_refuses_all_records():
         prompt_cache=[BatchKVarNKVCache(left_padding=[0]),
                       ArraysCache(size=2)],
     )
-    sig = spec_engine._ckpt_layout_live(b)
-    assert sig == spec_engine._LAYOUT_UNSUPPORTED
+    sig = engine._ckpt_layout_live(b)
+    assert sig == engine._LAYOUT_UNSUPPORTED
     man = GmlxAPCManager(num_blocks=8, block_size=16)
     ids = list(range(500, 532))
     assert ckpt_store(man, ids, [_hollow_kvarn(32), _arr()], extra_hash=0)
@@ -188,7 +189,7 @@ def test_store_lookup_signature_agreement_live():
                         prompt_cache=[_hollow_kvarn(p), _arr(seed=p)])
     assert ckpt_store(man, ids, b.prompt_cache, extra_hash=0)
     warm, got = ckpt_lookup(man, ids + [999], extra_hash=0,
-                            layout=spec_engine._ckpt_layout_live(b))
+                            layout=engine._ckpt_layout_live(b))
     assert got == p and type(warm[0]) is KVarNKVCache
 
 
@@ -217,13 +218,13 @@ def _plain_batch(man, ids, caches):
 def test_plain_init_adopts_kvarn_record():
     # End to end on the stock path: a kvarn-boot batch's live signature
     # matches a kvarn record and the warm adoption trims the prompt.
-    spec_engine._bind_l1_view()
+    engine._bind_l1_view()
     man = GmlxAPCManager(num_blocks=64, block_size=16)
     ids = list(range(300, 396))
     assert ckpt_store(man, ids[:32], [_hollow_kvarn(32), _arr(seed=5)],
                       extra_hash=0)
     b = _plain_batch(man, ids, [KVarNKVCache(), ArraysCache(size=2)])
-    spec_engine._plain_ckpt_init(b)
+    ckpt._plain_ckpt_init(b)
     assert b._processed_prompt_columns == 32
     assert type(b.prompt_cache[0]) is KVarNKVCache
     assert b.prompt_cache[0].offset == 32
@@ -233,13 +234,13 @@ def test_plain_init_adopts_kvarn_record():
 def test_plain_init_stock_batch_refuses_kvarn_record():
     # Same model, stock caches (the conversion declined this request):
     # the live signature refuses the kvarn record instead of adopting.
-    spec_engine._bind_l1_view()
+    engine._bind_l1_view()
     man = GmlxAPCManager(num_blocks=64, block_size=16)
     ids = list(range(300, 396))
     assert ckpt_store(man, ids[:32], [_hollow_kvarn(32), _arr(seed=5)],
                       extra_hash=0)
     b = _plain_batch(man, ids, [KVCache(), ArraysCache(size=2)])
-    spec_engine._plain_ckpt_init(b)
+    ckpt._plain_ckpt_init(b)
     assert b._processed_prompt_columns == 0
     assert b._kq_ckpt_armed
 
@@ -254,9 +255,9 @@ def test_production_wrap_chain_order():
 # importlib, not `import mlx_vlm.generate.ar as ...`: the package exports
 # a `generate` function that shadows the submodule attribute.
 ar = importlib.import_module("mlx_vlm.generate.ar")
-import gmlx.spec.engine as spec_engine
+import gmlx.spec.mtp_prefill as mtp_prefill
 from gmlx.cache import kvarn_serve as ks
-spec_engine.install_full_prompt_mtp_prefill()
+mtp_prefill.install_full_prompt_mtp_prefill()
 ks.install_kvarn_serve()
 
 def closure_names(fn):
diff --git a/tests/cache/test_ckpt_plain_path.py b/tests/cache/test_ckpt_plain_path.py
index 90ea6fc8..ce20a0bc 100644
--- a/tests/cache/test_ckpt_plain_path.py
+++ b/tests/cache/test_ckpt_plain_path.py
@@ -11,7 +11,8 @@
 from mlx_vlm.apc import APCManager
 
 import gmlx.cache.retire_key as retire_key
-import gmlx.spec.engine as spec_engine
+import gmlx.spec.ckpt as ckpt
+import gmlx.spec.engine as engine
 from gmlx.cache.snapshot import ckpt_lookup, ckpt_store
 from test_ckpt_tier import LAYOUT, make_hybrid_cache
 
@@ -46,11 +47,11 @@ def _plain_batch(man, ids, model=None):
 
 
 def test_plain_init_arms_cold_batch():
-    spec_engine._bind_l1_view()
+    engine._bind_l1_view()
     man = APCManager(num_blocks=64, block_size=16)
     ids = list(range(100, 148))
     b = _plain_batch(man, ids)
-    spec_engine._plain_ckpt_init(b)
+    ckpt._plain_ckpt_init(b)
     meta = b._apc_meta[0]
     assert b._kq_ckpt_armed and b._apc_harvest_enabled is False
     assert meta["ckpt_terminal"] > 0 and meta["checkpoint_len"] > 0
@@ -65,13 +66,13 @@ def test_plain_init_arms_cold_batch():
 
 
 def test_plain_init_restores_and_trims():
-    spec_engine._bind_l1_view()
+    engine._bind_l1_view()
     man = APCManager(num_blocks=64, block_size=16)
     ids = list(range(300, 396))
     warm_src = make_hybrid_cache(32, seed=5)
     assert ckpt_store(man, ids[:32], warm_src)
     b = _plain_batch(man, ids)
-    spec_engine._plain_ckpt_init(b)
+    ckpt._plain_ckpt_init(b)
     assert b._processed_prompt_columns == 32
     assert b._input_ids.shape[1] == 64
     assert b._inputs_embeds.shape[1] == 64
@@ -84,12 +85,12 @@ def test_plain_init_restores_and_trims():
 
 
 def test_plain_init_leaves_batched_rows_stock():
-    spec_engine._bind_l1_view()
+    engine._bind_l1_view()
     man = APCManager(num_blocks=64, block_size=16)
     ids = list(range(48))
     b = _plain_batch(man, ids)
     b._right_pad_per_row = [0]
-    spec_engine._plain_ckpt_init(b)
+    ckpt._plain_ckpt_init(b)
     assert not getattr(b, "_kq_ckpt_armed", False)
     assert b._apc_harvest_enabled is True
 
@@ -97,7 +98,7 @@ def test_plain_init_leaves_batched_rows_stock():
 def test_wrapped_stock_store_runs_cursor_and_suppresses_stock():
     from mlx_vlm.generate.ar import PromptProcessingBatch
 
-    spec_engine._install_ckpt_checkpoint_store()
+    ckpt._install_ckpt_checkpoint_store()
     man = APCManager(num_blocks=64, block_size=16)
     ids = list(range(400, 448))
     meta = {"full_input_ids": ids, "prefix_len": 0, "extra_hash": 0,
@@ -128,7 +129,7 @@ def _apc_prompt_cache_for_store(self, i):
 def test_gen_batch_filter_retires_lone_row(monkeypatch):
     from mlx_vlm.generate.ar import GenerationBatch
 
-    spec_engine._install_plain_ckpt_decode()
+    ckpt._install_plain_ckpt_decode()
     man = APCManager(num_blocks=64, block_size=16)
     full = list(range(500, 532))
     gen = list(range(900, 916))
@@ -151,7 +152,7 @@ def test_gen_batch_filter_retires_lone_row(monkeypatch):
 def test_gen_batch_retire_uses_decode_snap_on_divergence(monkeypatch):
     from mlx_vlm.generate.ar import GenerationBatch
 
-    spec_engine._install_plain_ckpt_decode()
+    ckpt._install_plain_ckpt_decode()
     man = APCManager(num_blocks=64, block_size=16)
     full = list(range(600, 632))
     gen = list(range(950, 966))
@@ -205,7 +206,7 @@ def _batched_hybrid_cache(lens, total):
 def test_gen_batch_filter_retires_leaving_row_batched():
     from mlx_vlm.generate.ar import GenerationBatch
 
-    spec_engine._install_plain_ckpt_decode()
+    ckpt._install_plain_ckpt_decode()
     man = APCManager(num_blocks=64, block_size=16)
     full = list(range(700, 724))
     gen = list(range(970, 986))                 # row 1: 24 + 16 = 40
@@ -228,7 +229,7 @@ def test_gen_batch_filter_retires_leaving_row_batched():
 def test_gen_batch_extend_carries_stash():
     from mlx_vlm.generate.ar import GenerationBatch
 
-    spec_engine._install_plain_ckpt_decode()
+    ckpt._install_plain_ckpt_decode()
     stash = {"full_ids": [1, 2], "mode": "ckpt", "gen": []}
     other = GenerationBatch.empty(model=None, sampler=None,
                                   stop_criteria=None)
@@ -242,12 +243,12 @@ def test_gen_batch_extend_carries_stash():
 
 
 def test_exact_anchor_init_arms_retire_stash():
-    spec_engine._bind_l1_view()
+    engine._bind_l1_view()
     man = APCManager(num_blocks=64, block_size=16)
     ids = list(range(800, 848))
     model = SimpleNamespace(_kq_apc_ckpt=False, config=SimpleNamespace())
     b = _plain_batch(man, ids, model=model)
-    spec_engine._plain_anchor_init(b)
+    ckpt._plain_anchor_init(b)
     stash = b.prompt_cache[0]._kq_apc_retire
     assert stash["mode"] == "exact" and stash["manager"] is man
     assert stash["full_ids"] == ids and stash["gen"] == []
@@ -257,7 +258,7 @@ def test_gen_batch_filter_retires_exact_row():
     from mlx_vlm.generate.ar import GenerationBatch
     from mlx_vlm.models.cache import KVCache
 
-    spec_engine._install_plain_ckpt_decode()
+    ckpt._install_plain_ckpt_decode()
     man = APCManager(num_blocks=64, block_size=16)
     full = list(range(850, 882))
     gen = list(range(990, 1006))                # 32 + 16 = 48
@@ -285,9 +286,9 @@ def test_plain_step_tick_disables_on_failure():
     stash = {"mode": "ckpt", "gen": 7}        # broken accounting slot
     gb = SimpleNamespace(
         uids=[3], prompt_cache=[SimpleNamespace(_kq_apc_retire=stash)])
-    spec_engine._plain_step_tick(gb, [[5]])
+    ckpt._plain_step_tick(gb, [[5]])
     assert stash["snap_ok"] is False and "gen" not in stash
-    spec_engine._plain_step_tick(gb, [[6]])   # gated off: silent no-op
+    ckpt._plain_step_tick(gb, [[6]])   # gated off: silent no-op
     assert "gen" not in stash
 
 
@@ -300,12 +301,12 @@ def test_plain_step_tick_reads_step_tuple_rows():
     s2 = {"mode": "exact", "gen": []}
     gb = SimpleNamespace(uids=[1, 2], prompt_cache=[SimpleNamespace()],
                          _kq_apc_retire_rows={1: s1, 2: s2})
-    spec_engine._plain_step_tick(gb, ([5, 9], None, None, None))
+    ckpt._plain_step_tick(gb, ([5, 9], None, None, None))
     assert s1["gen"] == [5] and s2["gen"] == [9]
-    spec_engine._plain_step_tick(gb, ([6, None], None, None, None))
+    ckpt._plain_step_tick(gb, ([6, None], None, None, None))
     assert s1["gen"] == [5, 6] and s2["gen"] == [9]
     assert s2.get("snap_ok") is not False
-    spec_engine._plain_step_tick(gb, (None, None, None, None))
+    ckpt._plain_step_tick(gb, (None, None, None, None))
     assert s1["gen"] == [5, 6] and s2["gen"] == [9]
 
 
diff --git a/tests/cache/test_ckpt_tier.py b/tests/cache/test_ckpt_tier.py
index 245f49f0..e59cc5e9 100644
--- a/tests/cache/test_ckpt_tier.py
+++ b/tests/cache/test_ckpt_tier.py
@@ -1110,7 +1110,7 @@ def test_ckpt_active_gating(monkeypatch):
 
 
 def test_mid_prefill_store_supersedes_stock(monkeypatch):
-    import gmlx.spec.engine as se
+    import gmlx.spec.ckpt as ckpt
 
     man = APCManager(num_blocks=64, block_size=16)
     ckpt_len = 32
@@ -1128,7 +1128,7 @@ def test_mid_prefill_store_supersedes_stock(monkeypatch):
         prompt_cache=cache,
         _row_real_tokens_processed=lambda idx: ckpt_len,
     )
-    se._ckpt_mid_prefill_store(batch)
+    ckpt._ckpt_mid_prefill_store(batch)
     # checkpoint_done set: the stock exact-clone store is now a no-op.
     assert batch._apc_meta[0]["checkpoint_done"] is True
     warm, got = ckpt_lookup(man, ids, extra_hash=5)
@@ -1142,7 +1142,7 @@ def test_mid_prefill_store_supersedes_stock(monkeypatch):
         prompt_cache=cache,
         _row_real_tokens_processed=lambda idx: ckpt_len,
     )
-    se._ckpt_mid_prefill_store(batch2)
+    ckpt._ckpt_mid_prefill_store(batch2)
     assert "checkpoint_done" not in batch2._apc_meta[0]
 
 
@@ -1174,7 +1174,9 @@ def test_spec_apc_master_disable_noops_store(monkeypatch):
     real owned-prefill APC entrypoint against a real APCManager with the master
     switch off and assert nothing is armed or stored; the switched-on control
     proves the same drive does store."""
+    import gmlx.spec.ckpt as ckpt
     import gmlx.spec.engine as se
+    import gmlx.spec.mtp_prefill as mtp_prefill
 
     p = 48
     ids = mx.array([list(range(100, 100 + p))])
@@ -1182,7 +1184,9 @@ def test_spec_apc_master_disable_noops_store(monkeypatch):
     def drive(disabled):
         for flag in ("_SPEC_APC_DISABLED", "_SPEC_APC_RETIRE_DISABLED",
                      "_SPEC_APC_SIDECAR_DISABLED", "_SPEC_APC_CKPT_DISABLED"):
-            monkeypatch.setattr(se, flag, disabled)
+            for mod in (se, ckpt, mtp_prefill):
+                if hasattr(mod, flag):
+                    monkeypatch.setattr(mod, flag, disabled)
         se._bind_l1_view()
         man = APCManager(num_blocks=64, block_size=16)
         model = SimpleNamespace(
@@ -1192,7 +1196,7 @@ def drive(disabled):
         batch = SimpleNamespace(
             model=model, _input_ids=ids, _inputs_embeds=mx.zeros((1, p, 4)),
             prompt_cache=make_hybrid_cache(p), _prompt_kwargs={})
-        se._mtp_prefill_init(batch)
+        mtp_prefill._mtp_prefill_init(batch)
         # The mid-prefill checkpoint moment fires either way; only an armed
         # batch stores.
         meta = (getattr(batch, "_apc_meta", None) or [{}])[0] or {}
@@ -1200,7 +1204,7 @@ def drive(disabled):
         if cl:
             batch.prompt_cache = make_hybrid_cache(cl)
             batch._row_real_tokens_processed = lambda idx: cl
-        se._ckpt_mid_prefill_store(batch)
+        ckpt._ckpt_mid_prefill_store(batch)
         return man, model, batch
 
     man_on, _model, batch_on = drive(disabled=False)     # switched-on control
diff --git a/tests/cache/test_exact_anchor.py b/tests/cache/test_exact_anchor.py
index c41ccf2b..d2f28e48 100644
--- a/tests/cache/test_exact_anchor.py
+++ b/tests/cache/test_exact_anchor.py
@@ -15,6 +15,7 @@
 from mlx_vlm.models.cache import CacheList, KVCache
 
 import gmlx.cache.snapshot as cs
+import gmlx.spec.ckpt as ckpt
 import gmlx.spec.engine as se
 from gmlx.cache.snapshot import anchor_exact_lookup, anchor_exact_store
 
@@ -123,33 +124,33 @@ def _bmeta(n=5000):
 def test_anchor_boundary_ungridded(monkeypatch):
     _stub_sys(monkeypatch, 2900)
     batch, meta = _bmeta()
-    assert se._exact_anchor_boundary(batch, meta, 4000, 0) == 2900
+    assert ckpt._exact_anchor_boundary(batch, meta, 4000, 0) == 2900
 
 
 def test_anchor_boundary_clamps_to_guard(monkeypatch):
     _stub_sys(monkeypatch, 4500)
     batch, meta = _bmeta()
-    assert se._exact_anchor_boundary(batch, meta, 4000, 0) == 4000
+    assert ckpt._exact_anchor_boundary(batch, meta, 4000, 0) == 4000
     # Guard 0 (stock checkpoint disabled): the divergence stands alone.
-    assert se._exact_anchor_boundary(batch, meta, 0, 0) == 4500
+    assert ckpt._exact_anchor_boundary(batch, meta, 0, 0) == 4500
 
 
 def test_anchor_boundary_floor_kill_restored(monkeypatch):
     _stub_sys(monkeypatch, 200)
     batch, meta = _bmeta()
-    assert se._exact_anchor_boundary(batch, meta, 4000, 0) is None
+    assert ckpt._exact_anchor_boundary(batch, meta, 4000, 0) is None
     monkeypatch.setenv("GMLX_APC_CKPT_SYS_MIN", "100")
-    assert se._exact_anchor_boundary(batch, meta, 4000, 0) == 200
-    assert se._exact_anchor_boundary(batch, meta, 4000, 200) is None
+    assert ckpt._exact_anchor_boundary(batch, meta, 4000, 0) == 200
+    assert ckpt._exact_anchor_boundary(batch, meta, 4000, 200) is None
     monkeypatch.setenv("GMLX_APC_CKPT_SYS", "0")
-    assert se._exact_anchor_boundary(batch, meta, 4000, 0) is None
+    assert ckpt._exact_anchor_boundary(batch, meta, 4000, 0) is None
 
 
 def test_anchor_boundary_no_render_ctx(monkeypatch):
     import gmlx.cache.retire_key as retire_key
     monkeypatch.setattr(retire_key, "lookup_render_ctx", lambda ids: None)
     batch, meta = _bmeta()
-    assert se._exact_anchor_boundary(batch, meta, 4000, 0) is None
+    assert ckpt._exact_anchor_boundary(batch, meta, 4000, 0) is None
 
 
 # -- two-stop schedule: anchor store, then the untouched stock guard --
@@ -163,7 +164,7 @@ def _armed_exact_batch(man, guard, monkeypatch, lcp):
         _apc_manager=man, _apc_meta=[meta], _apc_mode="exact",
         prompt_cache=None)
     batch._apc_prompt_cache_for_store = lambda idx: batch.prompt_cache
-    se._exact_anchor_arm(batch, meta, guard, 0)
+    ckpt._exact_anchor_arm(batch, meta, guard, 0)
     return batch, meta
 
 
@@ -184,7 +185,7 @@ def test_anchor_two_stop_schedule(monkeypatch):
     # and the stock body (running right after, as in the wrap) skips.
     batch.prompt_cache = make_kv_cache(64)
     batch._row_real_tokens_processed = lambda idx: 64
-    se._exact_anchor_store(batch)
+    ckpt._exact_anchor_store(batch)
     stock(batch)
     assert meta["anchor_done"] and meta["checkpoint_len"] == 96
     assert not meta.get("checkpoint_done") and calls == []
@@ -193,7 +194,7 @@ def test_anchor_two_stop_schedule(monkeypatch):
     # Guard stop: the hook is spent; the stock store fires and latches.
     batch.prompt_cache = make_kv_cache(96)
     batch._row_real_tokens_processed = lambda idx: 96
-    se._exact_anchor_store(batch)
+    ckpt._exact_anchor_store(batch)
     stock(batch)
     assert calls == [96] and meta.get("checkpoint_done")
 
@@ -212,7 +213,7 @@ def test_anchor_at_guard_single_stop(monkeypatch):
     assert meta["anchor_len"] == 96 and meta["checkpoint_len"] == 96
     batch.prompt_cache = make_kv_cache(96)
     batch._row_real_tokens_processed = lambda idx: 96
-    se._exact_anchor_store(batch)
+    ckpt._exact_anchor_store(batch)
     ar.PromptProcessingBatch._store_apc_exact_checkpoints(batch)
     assert anchor_exact_lookup(man, IDS, extra_hash=7)[1] == 96
     assert calls == [96] and meta.get("checkpoint_done")
@@ -235,7 +236,7 @@ def _pooling_model(man):
 def _pick_gen(man):
     from mlx_vlm.generate.ar import BatchGenerator
 
-    se._install_exact_anchor_pick()
+    ckpt._install_exact_anchor_pick()
     gen = SimpleNamespace(
         apc_manager=man, apc_mode="exact", model=_pooling_model(man),
         _apc_media_token_ids=lambda: set())
@@ -295,7 +296,7 @@ def test_plain_anchor_init_arms_on_a_right_padded_row(monkeypatch):
     # sibling. Upstream's checkpoint column handles it, so must we: this
     # is the common shape once any request is warm.
     batch, meta = _plain_batch(man, right_pad=[0])
-    se._plain_anchor_init(batch)
+    ckpt._plain_anchor_init(batch)
     assert batch._kq_anchor_armed and meta["checkpoint_len"] == 64
     # Nothing is trimmed here: restores come from the admission pick.
     assert batch._input_ids.shape[1] == len(IDS)
@@ -310,7 +311,7 @@ def test_plain_anchor_init_arms_above_a_shallow_warm_prefix(monkeypatch):
     # matches a token or two off an unrelated request; the divergence
     # still needs its clone.
     batch, meta = _plain_batch(man, right_pad=[0], prefix_len=1)
-    se._plain_anchor_init(batch)
+    ckpt._plain_anchor_init(batch)
     assert batch._kq_anchor_armed and meta["checkpoint_len"] == 64
 
 
@@ -323,6 +324,6 @@ def test_plain_anchor_init_skips_a_row_restored_past_the_divergence(
     # Restored at the divergence: the anchor it would store exists, so
     # no second stop and the stock guard runs alone.
     batch, meta = _plain_batch(man, right_pad=None, prefix_len=64)
-    se._plain_anchor_init(batch)
+    ckpt._plain_anchor_init(batch)
     assert not getattr(batch, "_kq_anchor_armed", False)
     assert meta["checkpoint_len"] == 96
diff --git a/tests/cache/test_kvarn_spec.py b/tests/cache/test_kvarn_spec.py
index 5f78b307..0c3a6129 100644
--- a/tests/cache/test_kvarn_spec.py
+++ b/tests/cache/test_kvarn_spec.py
@@ -21,7 +21,7 @@
 
 import mlx_kquant as kq  # noqa: E402
 
-import gmlx.spec.engine as spec_engine  # noqa: E402
+import gmlx.spec.kv_quant as kv_quant  # noqa: E402
 from gmlx.cache.kvarn_cache import KVarNKVCache  # noqa: E402
 from kvarn_testlib import Args, D, H, filled, needs_kvarn_ops, tokens  # noqa: E402
 
@@ -84,7 +84,7 @@ def _mk(lm=None, batch_size=1, make_cache=None):
 
 def test_params_kvarn_scheme_alone(restorable):
     restorable.setenv("KV_QUANT_SCHEME", "kvarn")
-    assert spec_engine._spec_kv_quant_params() == dict(
+    assert kv_quant._spec_kv_quant_params() == dict(
         scheme="kvarn", kv_bits=6, value_bits=6, tail_tokens=1024)
 
 
@@ -92,28 +92,28 @@ def test_params_kvarn_bits_and_tail(restorable):
     restorable.setenv("KV_QUANT_SCHEME", "kvarn")
     restorable.setenv("KV_BITS", "4")
     restorable.setenv("KV_TAIL_TOKENS", "256")
-    assert spec_engine._spec_kv_quant_params() == dict(
+    assert kv_quant._spec_kv_quant_params() == dict(
         scheme="kvarn", kv_bits=4, value_bits=4, tail_tokens=256)
 
 
 def test_params_kvarn_malformed(restorable):
     restorable.setenv("KV_QUANT_SCHEME", "kvarn")
     restorable.setenv("KV_BITS", "4.5")
-    assert spec_engine._spec_kv_quant_params() is None
+    assert kv_quant._spec_kv_quant_params() is None
 
 
 def test_params_kvarn_kill_switch(restorable):
     restorable.setenv("KV_QUANT_SCHEME", "kvarn")
     restorable.setenv("GMLX_SPEC_KV_QUANT", "0")
-    assert spec_engine._spec_kv_quant_params() is None
+    assert kv_quant._spec_kv_quant_params() is None
 
 
 def test_params_affine_unchanged(restorable):
     restorable.setenv("KV_BITS", "8")
-    assert spec_engine._spec_kv_quant_params() == dict(
+    assert kv_quant._spec_kv_quant_params() == dict(
         scheme="uniform", kv_bits=8, kv_group_size=64)
     restorable.setenv("KV_QUANT_SCHEME", "turboquant")
-    assert spec_engine._spec_kv_quant_params() is None
+    assert kv_quant._spec_kv_quant_params() is None
 
 
 def _stamp(lm, verdict="full", **single):
@@ -127,19 +127,19 @@ def _stamp(lm, verdict="full", **single):
 def test_stamped_params_rule_the_boot_env(restorable):
     # No stamp: None, so the boot env decides. A stamp: its scheme and
     # widths, {} when it quantizes nothing.
-    assert spec_engine._stamped_spec_params(_FakeLM()) is None
+    assert kv_quant._stamped_spec_params(_FakeLM()) is None
     kv = _stamp(_FakeLM(), scheme="kvarn", bits=4, value_bits=None,
                 tail_tokens=256)
-    assert spec_engine._stamped_spec_params(kv) == dict(
+    assert kv_quant._stamped_spec_params(kv) == dict(
         scheme="kvarn", kv_bits=4, value_bits=4, tail_tokens=256)
     kv = _stamp(_FakeLM(), scheme="kvarn", bits=6, value_bits=5,
                 tail_tokens=None)
-    assert spec_engine._stamped_spec_params(kv) == dict(
+    assert kv_quant._stamped_spec_params(kv) == dict(
         scheme="kvarn", kv_bits=6, value_bits=5, tail_tokens=1024)
     off = _stamp(_FakeLM(), scheme="uniform", bits=None, group_size=64)
-    assert spec_engine._stamped_spec_params(off) == {}
+    assert kv_quant._stamped_spec_params(off) == {}
     aff = _stamp(_FakeLM(), scheme="uniform", bits=8, group_size=32)
-    assert spec_engine._stamped_spec_params(aff) == dict(
+    assert kv_quant._stamped_spec_params(aff) == dict(
         scheme="uniform", kv_bits=8, kv_group_size=32)
 
 
@@ -150,17 +150,17 @@ def test_stamped_params_keep_a_declined_stamp_fp16():
     for scheme in ("uniform", "kvarn"):
         dropped = _stamp(_FakeLM(), verdict="dropped", scheme=scheme, bits=8,
                          group_size=64, value_bits=None, tail_tokens=None)
-        assert spec_engine._stamped_spec_params(dropped) == {}
+        assert kv_quant._stamped_spec_params(dropped) == {}
     err = _stamp(_FakeLM(), verdict="error", scheme="uniform", bits=8,
                  group_size=64)
-    assert spec_engine._stamped_spec_params(err) == {}
+    assert kv_quant._stamped_spec_params(err) == {}
 
 
 def test_stamped_params_honor_a_zero_tail():
     # tail 0 disables the fp16 tail; it is not the default's absence.
     kv = _stamp(_FakeLM(), scheme="kvarn", bits=6, value_bits=None,
                 tail_tokens=0)
-    assert spec_engine._stamped_spec_params(kv) == dict(
+    assert kv_quant._stamped_spec_params(kv) == dict(
         scheme="kvarn", kv_bits=6, value_bits=6, tail_tokens=0)
 
 
@@ -171,9 +171,9 @@ def test_stamped_params_honor_a_zero_tail():
 def test_b1_mtp_kvarn_converts(restorable, kvarn_ops_ok, caplog):
     import logging
 
-    caplog.set_level(logging.INFO, logger="gmlx.spec.engine")
+    caplog.set_level(logging.INFO, logger="gmlx.spec.kv_quant")
     restorable.setenv("KV_QUANT_SCHEME", "kvarn")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = _mk()
     # The shared carve-out holds the last layer of a deep stack fp16.
     assert type(caches[0]) is KVarNKVCache
@@ -189,7 +189,7 @@ def test_b1_mtp_kvarn_env_widths(restorable, kvarn_ops_ok):
     restorable.setenv("KV_QUANT_SCHEME", "kvarn")
     restorable.setenv("KV_BITS", "4")
     restorable.setenv("KV_TAIL_TOKENS", "256")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = _mk()
     assert caches[0].k_bits == 4 and caches[0].v_bits == 4
     assert caches[0].tail_cap == 256
@@ -199,7 +199,7 @@ def test_b1_mtp_kvarn_env_widths(restorable, kvarn_ops_ok):
 def test_b1_mtp_kvarn_from_the_stamp(restorable, kvarn_ops_ok):
     # Boot env says fp16; this model was loaded at kvarn k4 tail 256.
     restorable.delenv("KV_QUANT_SCHEME", raising=False)
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     lm = _stamp(_FakeLM(), scheme="kvarn", bits=4, value_bits=None,
                 tail_tokens=256)
     caches = _mk(lm=lm)
@@ -213,7 +213,7 @@ def test_b1_mtp_kvarn_from_the_stamp(restorable, kvarn_ops_ok):
 
 def test_readback_target_declines(restorable, kvarn_ops_ok):
     restorable.setenv("KV_QUANT_SCHEME", "kvarn")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = _mk(lm=_ReadbackLM())
     assert all(type(c) is not KVarNKVCache for c in caches)
 
@@ -223,7 +223,7 @@ def test_qwen35_arch_converts(restorable, kvarn_ops_ok):
     # The dispatch arm lifted the qwen3.5 bypass: the arch converts like
     # any other 128-dim stack (recurrent layers stay untouched).
     restorable.setenv("KV_QUANT_SCHEME", "kvarn")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = _mk(lm=_FakeLM(model_type="qwen3_5"))
     assert sum(type(c) is KVarNKVCache for c in caches) == 1
     assert type(caches[1]) is _SSMCache
@@ -237,7 +237,7 @@ def _batch_stack(lm, lp):
 
 def test_batch_without_kv_layers_passes_through(restorable, kvarn_ops_ok):
     restorable.setenv("KV_QUANT_SCHEME", "kvarn")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     sentinel = ["stock"]
     out = _mk(batch_size=2, make_cache=lambda lm, lp: sentinel)
     assert out is sentinel
@@ -254,9 +254,9 @@ def test_batch_converts_to_kvarn_rows(restorable, kvarn_ops_ok, caplog):
 
     from gmlx.cache.kvarn_cache import BatchKVarNKVCache
 
-    caplog.set_level(logging.INFO, logger="gmlx.spec.engine")
+    caplog.set_level(logging.INFO, logger="gmlx.spec.kv_quant")
     restorable.setenv("KV_QUANT_SCHEME", "kvarn")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = _mk(batch_size=2, make_cache=_batch_stack)
     assert type(caches[0]) is BatchKVarNKVCache
     assert isinstance(caches[1], _SSMCache)
@@ -272,9 +272,9 @@ def test_batch_declines_with_the_b1_reasons(restorable, kvarn_ops_ok, caplog):
     from gmlx.cache import kvarn_sdpa
     from gmlx.cache.kvarn_cache import BatchKVarNKVCache
 
-    caplog.set_level(logging.WARNING, logger="gmlx.spec.engine")
+    caplog.set_level(logging.WARNING, logger="gmlx.spec.kv_quant")
     restorable.setenv("KV_QUANT_SCHEME", "kvarn")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     # a target that reads K/V back declines batched as it does at B=1
     caches = _mk(lm=_ReadbackLM(), batch_size=2, make_cache=_batch_stack)
     assert all(type(c) is not BatchKVarNKVCache for c in caches)
@@ -288,7 +288,7 @@ def test_batch_declines_with_the_b1_reasons(restorable, kvarn_ops_ok, caplog):
 
 def test_rotating_stack_declines(restorable, kvarn_ops_ok):
     restorable.setenv("KV_QUANT_SCHEME", "kvarn")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
 
     class _RotatingLM(_FakeLM):
         def make_cache(self):
@@ -303,9 +303,9 @@ class _Owned(_ReadbackLM):
         def speculative_verify_hidden(self, x, pc):
             return x, {}
 
-    assert spec_engine._mtp_reads_kv_back(_ReadbackLM())
-    assert not spec_engine._mtp_reads_kv_back(_Owned())
-    assert not spec_engine._mtp_reads_kv_back(_FakeLM())
+    assert kv_quant._mtp_reads_kv_back(_ReadbackLM())
+    assert not kv_quant._mtp_reads_kv_back(_Owned())
+    assert not kv_quant._mtp_reads_kv_back(_FakeLM())
 
 
 def test_spec_probes_read_through_the_serve_wrapper(monkeypatch):
@@ -321,12 +321,12 @@ def rollback_speculative_cache(self, caches, gdn_states, accepted,
     inner = _LM()
     wrapper = SimpleNamespace(language_model=inner,
                               config={"model_type": "qwen3_5"})
-    assert spec_engine._mtp_reads_kv_back(wrapper)
-    spec_engine._harden_spec_target(wrapper)
+    assert kv_quant._mtp_reads_kv_back(wrapper)
+    kv_quant._harden_spec_target(wrapper)
     assert getattr(inner.rollback_speculative_cache, "_gmlx_kvarn_guard",
                    False)
     monkeypatch.setenv("GMLX_QWEN_OWNED", "0")
-    assert "stock fallback" in spec_engine.mtp_kv_decline(wrapper)
+    assert "stock fallback" in kv_quant.mtp_kv_decline(wrapper)
 
 
 # -- shared-KV readback guard ------------------------------------------------
@@ -547,7 +547,7 @@ def test_kvarn_lift_cache_recovers_original_domain():
     c.update_and_fetch(k, v)
     c._gmlx_cascade = "stamp"
 
-    lifted = spec_engine.kvarn_lift_cache(c)
+    lifted = kv_quant.kvarn_lift_cache(c)
     assert type(lifted) is BatchKVCache
     assert lifted.offset == 300
     assert lifted._gmlx_cascade == "stamp"
@@ -572,7 +572,7 @@ def test_kvarn_lift_matches_stock_attention():
     k, v = tokens(260, seed=11)
     c = KVarNKVCache(tail_tokens=256)
     c.update_and_fetch(k, v)
-    lifted = spec_engine.kvarn_lift_cache(c)
+    lifted = kv_quant.kvarn_lift_cache(c)
 
     q = mx.random.normal((1, H, 1, D)).astype(mx.float16)
     scale = D ** -0.5
@@ -595,16 +595,16 @@ def test_preemption_gate_admits_what_the_lift_handles():
 
     from gmlx.cache.kvarn_cache import KVarNRotatingKVCache
 
-    assert spec_engine.batch_liftable(KVarNKVCache(tail_tokens=256))
+    assert kv_quant.batch_liftable(KVarNKVCache(tail_tokens=256))
     # offset counts evicted tokens the rotating buffers no longer hold
-    assert not spec_engine.batch_liftable(KVarNRotatingKVCache(2048, tail_tokens=256))
-    assert spec_engine.batch_liftable(KVCache())
-    assert spec_engine.batch_liftable(QuantizedKVCache(group_size=64, bits=8))
+    assert not kv_quant.batch_liftable(KVarNRotatingKVCache(2048, tail_tokens=256))
+    assert kv_quant.batch_liftable(KVCache())
+    assert kv_quant.batch_liftable(QuantizedKVCache(group_size=64, bits=8))
 
     class _Opaque:
         pass
 
-    assert not spec_engine.batch_liftable(_Opaque())
+    assert not kv_quant.batch_liftable(_Opaque())
 
 
 @needs_kvarn_ops
diff --git a/tests/e2e/run_apc_depth_e2e.py b/tests/e2e/run_apc_depth_e2e.py
index c3600782..babf431a 100644
--- a/tests/e2e/run_apc_depth_e2e.py
+++ b/tests/e2e/run_apc_depth_e2e.py
@@ -15,7 +15,7 @@
 must retrieve the cold-proven probe facts or match the cold answer
 byte-for-byte. Reuse floors come from the ckpt cursor's own schedule
 arithmetic (the expected_* helpers mirror _ckpt_cursor_init and
-_ckpt_turn_boundaries in gmlx/spec/engine.py): an identical resend must
+_ckpt_turn_boundaries in gmlx/spec/ckpt.py): an identical resend must
 adopt the N-1 replay boundary, turns the render-stable grid floor, a
 divergent suffix the interval grid floor. The content field must stay
 markup-free everywhere.
@@ -120,7 +120,7 @@
     wait_disk_drained,
 )
 
-# Mirrors of the ckpt cursor's schedule arithmetic (gmlx/spec/engine.py):
+# Mirrors of the ckpt cursor's schedule arithmetic (gmlx/spec/ckpt.py):
 # boundaries sit on unit = lcm(prefill_step 2048, block 16); interval
 # points land every GMLX_APC_CKPT_INTERVAL (default 4096) snapped to that
 # grid; the replay boundary is N-1; turn boundaries are the unit-grid
diff --git a/tests/serve/test_batch_rows.py b/tests/serve/test_batch_rows.py
index e9a42432..e56e413e 100644
--- a/tests/serve/test_batch_rows.py
+++ b/tests/serve/test_batch_rows.py
@@ -61,7 +61,7 @@ def test_missing_or_none_batch_reads_zero():
 
 def test_no_len_on_generation_batch_in_decision_modules():
     # The loud tripwire for new call sites. The engine-side len() calls in
-    # spec_engine are the promotion mechanism and are exempt by design.
+    # spec.admission are the promotion mechanism and are exempt by design.
     root = Path(gmlx.__file__).parent
     offenders = []
     for name in DECISION_MODULES:
diff --git a/tests/serve/test_decode_batch.py b/tests/serve/test_decode_batch.py
index 9abcce85..78a6d3ae 100644
--- a/tests/serve/test_decode_batch.py
+++ b/tests/serve/test_decode_batch.py
@@ -38,7 +38,7 @@ def test_garbage_falls_back(monkeypatch):
 def test_stash_wrapper_injects(monkeypatch):
     from mlx_vlm.generate import ar
 
-    import gmlx.spec.engine as spec_engine
+    import gmlx.spec.engine as engine
 
     seen = {}
 
@@ -47,7 +47,7 @@ def __init__(self, model, processor, **kwargs):
             seen.update(kwargs)
 
     monkeypatch.setattr(ar, "BatchGenerator", _BG)
-    spec_engine._install_apc_manager_stash()
+    engine._install_apc_manager_stash()
 
     monkeypatch.setenv("GMLX_DECODE_BATCH", "5")
     ar.BatchGenerator(SimpleNamespace(), None)
@@ -65,7 +65,7 @@ def test_stash_wrapper_clamps_full_width_prefill_group(monkeypatch):
     # group to 1 only in that regime.
     from mlx_vlm.generate import ar
 
-    import gmlx.spec.engine as spec_engine
+    import gmlx.spec.engine as engine
 
     class _BG:
         def __init__(self, model, processor, **kwargs):
@@ -74,7 +74,7 @@ def __init__(self, model, processor, **kwargs):
             self.prefill_batch_size = kwargs.get("prefill_batch_size", 8)
 
     monkeypatch.setattr(ar, "BatchGenerator", _BG)
-    spec_engine._install_apc_manager_stash()
+    engine._install_apc_manager_stash()
 
     monkeypatch.setenv("GMLX_DECODE_BATCH", "8")
     gen = ar.BatchGenerator(SimpleNamespace(), None)
diff --git a/tests/serve/test_serve_apc_engagement.py b/tests/serve/test_serve_apc_engagement.py
index 970ffbd4..849e1be5 100644
--- a/tests/serve/test_serve_apc_engagement.py
+++ b/tests/serve/test_serve_apc_engagement.py
@@ -62,9 +62,9 @@ def family(request, gguf_index):
         pytest.skip(f"no {arch!r} GGUF under KQUANT_TEST_GGUF_DIR "
                     f"(have: {sorted(gguf_index)})")
     import gmlx.serve.bridge_vlm as serving
-    import gmlx.spec.engine as spec_engine
+    import gmlx.spec.mtp_prefill as mtp_prefill
 
-    spec_engine.install_full_prompt_mtp_prefill()   # the serve installs
+    mtp_prefill.install_full_prompt_mtp_prefill()   # the serve installs
     if scheme == "kvarn":
         from gmlx.cache.kvarn_apc import install_kvarn_apc
         from gmlx.cache.kvarn_serve import install_kvarn_serve
diff --git a/tests/serve/test_serve_mtp.py b/tests/serve/test_serve_mtp.py
index a9ac1cff..17894acd 100644
--- a/tests/serve/test_serve_mtp.py
+++ b/tests/serve/test_serve_mtp.py
@@ -728,7 +728,7 @@ def test_prompt_step_caps_mtp_hidden_capture():
     import mlx.core as mx
     from mlx_vlm.generate.ar import PromptProcessingBatch
 
-    from gmlx.spec.engine import install_full_prompt_mtp_prefill
+    from gmlx.spec.mtp_prefill import install_full_prompt_mtp_prefill
 
     install_full_prompt_mtp_prefill()
 
diff --git a/tests/spec/test_full_prompt_prefill.py b/tests/spec/test_full_prompt_prefill.py
index a95d2043..632c5eca 100644
--- a/tests/spec/test_full_prompt_prefill.py
+++ b/tests/spec/test_full_prompt_prefill.py
@@ -220,12 +220,10 @@ def _embed_prompts(model, ids_list):
 
 
 def _install_patches():
-    from gmlx.spec.engine import (
-        install_full_prompt_mtp_prefill,
-        install_owned_spec_engine,
-        install_continuous_batch_admission,
-        install_spec_kv_quant,
-    )
+    from gmlx.spec.admission import install_continuous_batch_admission
+    from gmlx.spec.engine import install_owned_spec_engine
+    from gmlx.spec.kv_quant import install_spec_kv_quant
+    from gmlx.spec.mtp_prefill import install_full_prompt_mtp_prefill
     from gmlx.cache.apc_pooling import install_pooling_apc_support
     from gmlx.cache.kvarn_apc import install_kvarn_apc
 
@@ -400,10 +398,10 @@ def test_prefill_step_env_override(monkeypatch):
 
     from mlx_vlm.generate.ar import PromptProcessingBatch
 
-    import gmlx.spec.engine as spec_engine
+    import gmlx.spec.mtp_prefill as mtp_prefill
 
-    spec_engine.install_full_prompt_mtp_prefill()
-    monkeypatch.setattr(spec_engine, "_mtp_prefill_init", lambda s: None)
+    mtp_prefill.install_full_prompt_mtp_prefill()
+    monkeypatch.setattr(mtp_prefill, "_mtp_prefill_init", lambda s: None)
 
     def fake_batch():
         return types.SimpleNamespace(
@@ -765,7 +763,7 @@ def test_l1_sidecar_warm_start(mtp_model):
     switch and must fall back to a plain L1 hit (acceptance-parity numbers
     are the D-run's job; this certifies plumbing + correctness)."""
     import gmlx.spec.speculative as _spec
-    import gmlx.spec.engine as _eng
+    import gmlx.spec.ckpt as _ckpt
     from mlx_vlm.apc import APCManager
 
     model, drafter, config, tokenizer = mtp_model
@@ -808,9 +806,9 @@ def test_l1_sidecar_warm_start(mtp_model):
 
         _clear_l0(model)
         old_spec, old_eng = _spec._SIDECAR_DISABLED, \
-            _eng._SPEC_APC_SIDECAR_DISABLED
+            _ckpt._SPEC_APC_SIDECAR_DISABLED
         _spec._SIDECAR_DISABLED = True
-        _eng._SPEC_APC_SIDECAR_DISABLED = True
+        _ckpt._SPEC_APC_SIDECAR_DISABLED = True
         try:
             with _capture_spec_log() as messages:
                 toks_off = _run_mtp(
@@ -818,7 +816,7 @@ def test_l1_sidecar_warm_start(mtp_model):
                     apc_manager=manager)[0]
         finally:
             _spec._SIDECAR_DISABLED = old_spec
-            _eng._SPEC_APC_SIDECAR_DISABLED = old_eng
+            _ckpt._SPEC_APC_SIDECAR_DISABLED = old_eng
         assert any("APC L1 hit" in m for m in messages)
         assert not any("APC sidecar hit" in m for m in messages), (
             "kill switch did not disable the sidecar lookup"
@@ -1041,7 +1039,9 @@ def test_l1_kill_switch(mtp_model, monkeypatch):
     """GMLX_SPEC_APC=0 must disable the L1 lookup even with a manager
     stashed. (The flag is read at import; patch the module constant.)"""
     from mlx_vlm.apc import APCManager
-    import gmlx.spec.engine as spec_engine
+    import gmlx.spec.ckpt as ckpt
+    import gmlx.spec.engine as engine
+    import gmlx.spec.mtp_prefill as mtp_prefill
 
     model, drafter, config, tokenizer = mtp_model
     _install_patches()
@@ -1049,7 +1049,8 @@ def test_l1_kill_switch(mtp_model, monkeypatch):
     manager = APCManager(num_blocks=2048, block_size=16)
 
     prefix_ids = _build_prompt(tokenizer, 3000, seed=_SEED_B)
-    monkeypatch.setattr(spec_engine, "_SPEC_APC_DISABLED", True)
+    for mod in (engine, ckpt, mtp_prefill):
+        monkeypatch.setattr(mod, "_SPEC_APC_DISABLED", True)
     try:
         with _capture_spec_log() as messages:
             _run_mtp(model, drafter, tokenizer, [prefix_ids], N_DECODE,
@@ -1390,7 +1391,7 @@ def test_apc_hit_on_injected_request(mtp_model):
 
     Also asserts that PromptProcessingBatch prefills single-request (B=1)
     -- if mlx-vlm ever coalesces prefills into B>1, the b==1 guard in
-    spec_engine silently disables APC and this test fails loudly via the
+    mtp_prefill silently disables APC and this test fails loudly via the
     APC-hit assertion rather than producing a silent degradation.
     """
     import logging
@@ -1448,7 +1449,7 @@ def test_apc_hit_on_injected_request(mtp_model):
     apc_hit_seen = False
 
     # Capture APC log to verify the hit actually fired.
-    apc_log = logging.getLogger("gmlx.spec.engine")
+    apc_log = logging.getLogger("gmlx.spec.ckpt")
     log_messages = []
     handler = logging.Handler()
     handler.emit = lambda record: log_messages.append(record.getMessage())
diff --git a/tests/spec/test_mrope_widen.py b/tests/spec/test_mrope_widen.py
index dd144c0b..e7476599 100644
--- a/tests/spec/test_mrope_widen.py
+++ b/tests/spec/test_mrope_widen.py
@@ -11,7 +11,7 @@
 
 import mlx.core as mx
 
-from gmlx.spec.engine import _widen_prompt_rope_state
+from gmlx.spec.mtp_prefill import _widen_prompt_rope_state
 
 
 def _batch(width, model):
diff --git a/tests/spec/test_mtp_preempt_resume.py b/tests/spec/test_mtp_preempt_resume.py
index 4721cc6c..9a0f9104 100644
--- a/tests/spec/test_mtp_preempt_resume.py
+++ b/tests/spec/test_mtp_preempt_resume.py
@@ -444,7 +444,7 @@ def fake_rounds(model, draft_model, prompt_cache, hidden, **kw):
 
 def test_preempt_rebuilds_scalar_for_waiters(monkeypatch):
     from mlx_vlm.generate import ar
-    from gmlx.spec.engine import install_continuous_batch_admission
+    from gmlx.spec.admission import install_continuous_batch_admission
 
     install_continuous_batch_admission()
     calls = []
@@ -480,7 +480,7 @@ def test_preempt_rebuilds_scalar_for_waiters(monkeypatch):
 
 def test_preempt_env_kill_switch(monkeypatch):
     from mlx_vlm.generate import ar
-    from gmlx.spec.engine import install_continuous_batch_admission
+    from gmlx.spec.admission import install_continuous_batch_admission
 
     install_continuous_batch_admission()
     monkeypatch.setenv("GMLX_MTP_PREEMPT", "0")
@@ -506,7 +506,7 @@ def test_preempt_waits_for_first_delivery(monkeypatch):
     """A batch that has not delivered its first tokens has no bonus to
     rebuild from; the preempt fires on the following next() instead."""
     from mlx_vlm.generate import ar
-    from gmlx.spec.engine import install_continuous_batch_admission
+    from gmlx.spec.admission import install_continuous_batch_admission
 
     install_continuous_batch_admission()
     calls = []
@@ -593,7 +593,7 @@ def test_preempt_delivers_captured_tail(monkeypatch):
     and the rebuild resumes from the tail's last token (the round's bonus,
     whose KV is not in the cache)."""
     from mlx_vlm.generate import ar
-    from gmlx.spec.engine import install_continuous_batch_admission
+    from gmlx.spec.admission import install_continuous_batch_admission
 
     install_continuous_batch_admission()
     calls = []
@@ -620,7 +620,7 @@ def test_preempt_captured_tail_finishes_row(monkeypatch):
     """A captured token that exhausts the budget finishes the row: the tail
     truncates at the finish, no rebuild happens, and the waiter promotes."""
     from mlx_vlm.generate import ar
-    from gmlx.spec.engine import install_continuous_batch_admission
+    from gmlx.spec.admission import install_continuous_batch_admission
 
     install_continuous_batch_admission()
     calls = []
diff --git a/tests/spec/test_seed_stream.py b/tests/spec/test_seed_stream.py
index 79596a8d..5c75a2b5 100644
--- a/tests/spec/test_seed_stream.py
+++ b/tests/spec/test_seed_stream.py
@@ -23,7 +23,7 @@
 import mlx.core as mx  # noqa: E402
 import mlx.nn as nn  # noqa: E402
 
-import gmlx.spec.engine as engine  # noqa: E402
+import gmlx.spec.mtp_prefill as mtp_prefill  # noqa: E402
 from gmlx.spec.mtp_drafter import QwenMTPDrafter  # noqa: E402
 
 D = 8
@@ -220,14 +220,14 @@ def __init__(self, drafter, b=1, chunk_hiddens=(), l1=0, processed=0,
 class TestEligibility:
 
     def _armed(self, batch):
-        engine._mtp_seed_stream_init(batch)
+        mtp_prefill._mtp_seed_stream_init(batch)
         return batch._mtp_seed_ctx is not None
 
     def test_eligible_cold_prefill_arms(self):
         assert self._armed(_StubBatch(FakeHead()))
 
     def test_env_kill_switch(self, monkeypatch):
-        monkeypatch.setattr(engine, "_SEED_STREAM_DISABLED", True)
+        monkeypatch.setattr(mtp_prefill, "_SEED_STREAM_DISABLED", True)
         assert not self._armed(_StubBatch(FakeHead()))
 
     def test_window_limited_head_defers(self):
@@ -259,7 +259,7 @@ def test_warm_sidecar_defers(self):
 
     def test_armed_ctx_rides_prompt_cache(self):
         b = _StubBatch(FakeHead())
-        engine._mtp_seed_stream_init(b)
+        mtp_prefill._mtp_seed_stream_init(b)
         assert b.prompt_cache[0]._kq_seed_stream is b._mtp_seed_ctx
         assert b._mtp_seed_ctx["len"] == 0
         assert b._mtp_seed_ctx["active"]
diff --git a/tests/spec/test_spec_engine_release.py b/tests/spec/test_spec_engine_release.py
index 94be62f3..5da77bef 100644
--- a/tests/spec/test_spec_engine_release.py
+++ b/tests/spec/test_spec_engine_release.py
@@ -18,12 +18,11 @@
 
 import mlx.core as mx
 
-from gmlx.spec.engine import (
-    _OWNED_MTP_ROUND_FLAG,
+from gmlx.spec.admission import (
     _RELEASED_FLAG,
     install_continuous_batch_admission,
-    install_owned_spec_engine,
 )
+from gmlx.spec.engine import _OWNED_MTP_ROUND_FLAG, install_owned_spec_engine
 
 
 class _Entry:
diff --git a/tests/spec/test_spec_kv_quant.py b/tests/spec/test_spec_kv_quant.py
index 9b666c0d..38666f69 100644
--- a/tests/spec/test_spec_kv_quant.py
+++ b/tests/spec/test_spec_kv_quant.py
@@ -14,7 +14,8 @@
 from mlx_vlm.speculative import utils as su  # noqa: E402
 
 import gmlx.models.qwen35.verify_fold as qwen35_verify_fold  # noqa: E402
-import gmlx.spec.engine as spec_engine  # noqa: E402
+import gmlx.spec.engine as engine  # noqa: E402
+import gmlx.spec.kv_quant as kv_quant  # noqa: E402
 
 
 class _SSMCache:
@@ -54,7 +55,7 @@ def _mk(batch_size=1, make_cache=None, lm=None):
 
 def test_b1_mtp_converts(restorable):
     restorable.setenv("KV_BITS", "4")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = _mk()
     assert isinstance(caches[0], QuantizedKVCache)
     # The last layer of a deep stack stays fp16. The MTP arm
@@ -65,14 +66,14 @@ def test_b1_mtp_converts(restorable):
     assert caches[0].offset == 0 and caches[0].is_trimmable()
     # idempotent: second install keeps the same wrapper
     wrapped = su.make_speculative_prompt_cache
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     assert su.make_speculative_prompt_cache is wrapped
 
 
 def test_group_size_env(restorable):
     restorable.setenv("KV_BITS", "8")
     restorable.setenv("KV_GROUP_SIZE", "32")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = _mk()
     assert caches[0].bits == 8 and caches[0].group_size == 32
 
@@ -82,7 +83,7 @@ def test_no_env_installs_and_stays_fp16(restorable):
     # stamped policy even when the boot env asked for nothing.
     restorable.delenv("KV_BITS", raising=False)
     before = su.make_speculative_prompt_cache
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     assert su.make_speculative_prompt_cache is not before
     assert all(type(c) in (KVCache, _SSMCache) for c in _mk())
 
@@ -100,7 +101,7 @@ def test_stamped_policy_rules_the_boot_env(restorable):
     from gmlx.cache.kv_policy import off_policy, resolve_kv_quant_policy
 
     restorable.setenv("KV_BITS", "8")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     fp16 = _stamp(_FakeLM(), off_policy("single"))
     assert all(type(c) in (KVCache, _SSMCache) for c in _mk(lm=fp16))
     narrow = _stamp(_FakeLM(), resolve_kv_quant_policy(
@@ -116,7 +117,7 @@ def test_kill_switch(restorable):
     restorable.setenv("KV_BITS", "4")
     restorable.setenv("GMLX_SPEC_KV_QUANT", "0")
     before = su.make_speculative_prompt_cache
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     assert su.make_speculative_prompt_cache is before
 
 
@@ -126,13 +127,13 @@ def test_kill_switch(restorable):
 def test_non_affine_stays_fp16(restorable, bits, scheme):
     restorable.setenv("KV_BITS", bits)
     restorable.setenv("KV_QUANT_SCHEME", scheme)
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     assert all(type(c) in (KVCache, _SSMCache) for c in _mk())
 
 
 def test_batch_passthrough(restorable):
     restorable.setenv("KV_BITS", "4")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     sentinel = ["stock"]
     out = _mk(batch_size=2, make_cache=lambda lm, lp: sentinel)
     assert out is sentinel
@@ -145,7 +146,7 @@ def test_batch_forces_fp16(restorable):
     from mlx_vlm.models.cache import BatchKVCache, BatchQuantizedKVCache
 
     restorable.setenv("KV_BITS", "8")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     lp = [0, 0]
     stock = [BatchQuantizedKVCache(lp, group_size=64, bits=8),
              _SSMCache(),
@@ -163,7 +164,7 @@ def test_batch_forces_fp16_nested(restorable):
                                       CacheList)
 
     restorable.setenv("KV_BITS", "8")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     lp = [0, 0]
     stock = [CacheList(BatchQuantizedKVCache(lp, group_size=64, bits=8),
                        _SSMCache()),
@@ -186,7 +187,7 @@ def test_dequantize_lift_cache():
     q = QuantizedKVCache(group_size=64, bits=8)
     q.update_and_fetch(k, v)
     q._gmlx_cascade = "stamp"
-    lifted = spec_engine.dequantize_lift_cache(q)
+    lifted = kv_quant.dequantize_lift_cache(q)
     assert type(lifted) is BatchKVCache
     assert lifted.offset == 41
     assert lifted._gmlx_cascade == "stamp"
@@ -332,7 +333,7 @@ def test_owned_off_gdn_declines_quantization(restorable, lm_cls):
     # guard keys on model_type, direct or config fallback.
     restorable.setenv("KV_BITS", "4")
     restorable.setenv("GMLX_QWEN_OWNED", "0")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = ar.make_speculative_prompt_cache(
         lm_cls(), draft_kind="mtp", batch_size=1, left_padding=[0],
         make_cache=lambda lm, lp: pytest.fail(
@@ -344,7 +345,7 @@ def test_owned_off_gdn_declines_quantization(restorable, lm_cls):
 def test_owned_on_gdn_still_converts(restorable):
     restorable.setenv("KV_BITS", "4")
     restorable.setenv("GMLX_QWEN_OWNED", "1")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = ar.make_speculative_prompt_cache(
         _GdnFakeLM(), draft_kind="mtp", batch_size=1, left_padding=[0],
         make_cache=lambda lm, lp: pytest.fail(
@@ -370,16 +371,16 @@ def test_warm_merge_config_follows_batched_policy(restorable):
                                      kv_group_size=64, mode="single")
     batched_drop = dropped_policy("mtp fp16 when batched", 8, 64, "batched")
     model._gmlx_kv_policy = ServeKvPolicy(single, batched_drop)
-    assert spec_engine._live_kv_quant_config(model) is None
+    assert engine._live_kv_quant_config(model) is None
 
     batched_full = resolve_kv_quant_policy([KVCache()], kv_bits=8,
                                            kv_group_size=64, mode="batched")
     model._gmlx_kv_policy = ServeKvPolicy(single, batched_full)
-    assert spec_engine._live_kv_quant_config(model) is not None
+    assert engine._live_kv_quant_config(model) is not None
 
     # no stamp: fail-safe None, never the environment
-    assert spec_engine._live_kv_quant_config(_Stamp()) is None
-    assert spec_engine._live_kv_quant_config(None) is None
+    assert engine._live_kv_quant_config(_Stamp()) is None
+    assert engine._live_kv_quant_config(None) is None
 
 
 # Hybrid arch shapes on the B=1 arm. CacheList members, opt-outs, nested
@@ -408,7 +409,7 @@ def make_cache(self):
 
 def test_b1_mtp_quantizes_the_kv_member_of_a_cache_list(restorable):
     restorable.setenv("KV_BITS", "8")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = ar.make_speculative_prompt_cache(
         _ListFakeLM(), draft_kind="mtp", batch_size=1, left_padding=[0],
         make_cache=lambda lm, lp: pytest.fail("B=1 mtp bypass"),
@@ -429,7 +430,7 @@ def make_cache(self):
 
 def test_b1_mtp_honors_kv_quant_unsupported(restorable):
     restorable.setenv("KV_BITS", "8")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = ar.make_speculative_prompt_cache(
         _OptOutFakeLM(), draft_kind="mtp", batch_size=1, left_padding=[0],
         make_cache=lambda lm, lp: pytest.fail("B=1 mtp bypass"),
@@ -450,7 +451,7 @@ def make_cache(self):
 def test_b1_mtp_sees_a_window_nested_in_a_cache_list(restorable):
     # A window-plus-state list classifies as state. Nothing converts.
     restorable.setenv("KV_BITS", "8")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = ar.make_speculative_prompt_cache(
         _NestedWindowFakeLM(), draft_kind="mtp", batch_size=1,
         left_padding=[0],
@@ -474,9 +475,9 @@ def test_b1_mtp_arms_the_pool_beside_the_kv_member(restorable, caplog):
     # layer kind. Every layer's pool packs, the held last layer included.
     import logging
 
-    caplog.set_level(logging.INFO, logger="gmlx.spec.engine")
+    caplog.set_level(logging.INFO, logger="gmlx.spec.kv_quant")
     restorable.setenv("KV_BITS", "8")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = ar.make_speculative_prompt_cache(
         _Glm5ShapeFakeLM(), draft_kind="mtp", batch_size=1,
         left_padding=[0],
@@ -504,9 +505,9 @@ def test_b1_mtp_notes_a_pool_only_engagement(restorable, caplog):
     # Nothing converts, so the note must key off the pools armed.
     import logging
 
-    caplog.set_level(logging.INFO, logger="gmlx.spec.engine")
+    caplog.set_level(logging.INFO, logger="gmlx.spec.kv_quant")
     restorable.setenv("KV_BITS", "8")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = ar.make_speculative_prompt_cache(
         _PoolOnlyFakeLM(), draft_kind="mtp", batch_size=1,
         left_padding=[0],
@@ -529,7 +530,7 @@ def test_b1_mtp_hybrid_stack_quantizes_full_attn_layers(restorable):
     # Top-level windows must not drop the whole stack: the verdict is
     # partial, windows stay fp16, full-attention layers convert.
     restorable.setenv("KV_BITS", "8")
-    spec_engine.install_spec_kv_quant()
+    kv_quant.install_spec_kv_quant()
     caches = ar.make_speculative_prompt_cache(
         _HybridFakeLM(), draft_kind="mtp", batch_size=1, left_padding=[0],
         make_cache=lambda lm, lp: pytest.fail("B=1 mtp bypass"),
@@ -547,16 +548,16 @@ def test_mtp_kv_decline_is_shared_by_serve_run_and_chat(restorable):
     kv_bits under MTP while serve quantized the same stack; the two
     stock verify walks are the only real declines."""
     restorable.setenv("GMLX_QWEN_OWNED", "1")
-    assert spec_engine.mtp_kv_decline(_FakeLM()) is None
-    assert spec_engine.mtp_kv_decline(_GdnFakeLM()) is None
+    assert kv_quant.mtp_kv_decline(_FakeLM()) is None
+    assert kv_quant.mtp_kv_decline(_GdnFakeLM()) is None
 
     restorable.setenv("GMLX_QWEN_OWNED", "0")
-    assert "stock fallback" in spec_engine.mtp_kv_decline(_GdnFakeLM())
-    assert "stock fallback" in spec_engine.mtp_kv_decline(_GdnConfigFakeLM())
-    assert spec_engine.mtp_kv_decline(_FakeLM()) is None
+    assert "stock fallback" in kv_quant.mtp_kv_decline(_GdnFakeLM())
+    assert "stock fallback" in kv_quant.mtp_kv_decline(_GdnConfigFakeLM())
+    assert kv_quant.mtp_kv_decline(_FakeLM()) is None
 
     restorable.setenv("GMLX_QWEN_OWNED", "1")
-    reason = spec_engine.mtp_kv_decline(_FakeLM(), owned_round=False)
+    reason = kv_quant.mtp_kv_decline(_FakeLM(), owned_round=False)
     assert "GMLX_OWNED_ROUND=0" in reason