MLX port (Step 1): incremental restored decode to kill throughput collapse

Port CUDA Gap-A to MLX. The existing MLX restored-attention dispatch already
calls cache.update_and_fetch/cache.offset, so the per-token re-forward collapse
is fixed by prefilling WITH a cache then decoding incrementally:

- restored_prefill_cache: prefill once with restored-K/V injection into the
  model's native hybrid cache (full/global layers -> exact own K/V (S5);
  sliding -> f_theta-restored, window-bounded by RotatingKVCache).
- restored_incremental_generate: greedy decode via mlx_lm generate_step over the
  prefilled cache (O(L)/token, async-pipelined). Recall carried by S5 full-attn.
- k3_integrated_niah_eval_mac.py: --incremental flag selects the new path.
- docs/mlx-port-lessons.md: Step 1 marked implemented + Mac validation command.

Linux: compiles, funcs import (mlx lazy), MLX helper tests pass. End-to-end
decode requires Apple Silicon -> Mac validation.

Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

Author: cursoragent

docs/mlx-port-lessons.md

@@ -50,11 +50,24 @@ speed** (on CUDA: 1.3–2.8 tok/s re-forward → ~21 tok/s incremental = AR).
 
 ## MLX port plan (ordered; each gates the next)
 
-1. **Incremental decode (kills the collapse).** Add an MLX analog of
-   `CrossModelRestoredSinkWindowVerifier(incremental=True)`: prefill → capture
-   restored K/V into `SinkWindowKVCache` (full-attn = own/exact; sliding = f_θ or
-   window-masked) → decode via `generate_step(prompt_cache=…)`. **Gate: decode
-   tok/s ≈ native mlx_lm AR; recall 1.0** (carried by S5).
+1. **Incremental decode (kills the collapse). [IMPLEMENTED — needs Mac validation]**
+   `backends/mlx/cross_model_dlm_verifier.py`: `restored_prefill_cache` (prefill
+   once with injection **into the model's native hybrid cache** → full-attn/global
+   layers store exact own K/V, sliding store f_θ-restored + window-bounded) +
+   `restored_incremental_generate` (decode via `mlx_lm.generate_step` over that
+   cache, O(L)/token, async-pipelined). Wired into the Mac harness via
+   `--incremental`:
+   ```bash
+   PYTHONPATH=.:sdks/python python scripts/research/k3_integrated_niah_eval_mac.py \
+     --verifier-path models/gemma-4-26B-A4B-it-mlx-4bit \
+     --drafter-id z-lab/gemma-4-26B-A4B-it-DFlash \
+     --f-theta-dir results/research/f_theta_v5_s5_sliding \
+     --s5-exact-full-attn --incremental --n-samples 5 --max-new-tokens 32
+   ```
+   **Gate: decode tok/s ≫ the per-token re-forward (toward native mlx_lm AR);
+   recall == oracle (1.0)** (carried by S5). Mechanism mirrors CUDA Gap-A: the
+   existing MLX dispatch already calls `cache.update_and_fetch`, so prefill *with*
+   a cache populates it; decode then runs native incremental attention.
 2. **Drop the extra build forward.** Capture full-attn own K/V at prefill; do not
    re-run a clean verifier forward per request beyond prefill. **Gate:
    `build_restoration` from ~12s → ~prefill cost.**

inference_engine/backends/mlx/cross_model_dlm_verifier.py

@@ -360,3 +360,96 @@ def restored_logits(
         logits = mlx_model(ids)            # full Model.__call__ → tied embed + softcap
         mx.eval(logits)
     return logits[0] if return_all else logits[0, -1]
+
+
+# ---------------------------------------------------------------------------
+# Incremental decode (MLX port of CUDA Gap-A) — kills the per-token re-forward
+# throughput collapse. See docs/mlx-port-lessons.md.
+# ---------------------------------------------------------------------------
+
+
+def restored_prefill_cache(
+    mlx_model: Any,
+    input_ids: Sequence[int],
+    *,
+    restored_k_per_layer: Dict[int, Any],
+    restored_v_per_layer: Dict[int, Any],
+    evicted_positions: Sequence[int],
+):
+    """Prefill ONCE with restoration, capturing the restored K/V into a
+    persistent mlx_lm prompt cache; return ``(cache, last_logits)``.
+
+    Same injection as :func:`restored_logits`, but run **with a cache** so the
+    patched attention's ``cache.update_and_fetch`` stores the post-injection
+    K/V (full-attention/S5 layers → exact own K/V; sliding → f_θ-restored,
+    window-bounded by the model's native RotatingKVCache). After this the
+    verifier can decode new tokens incrementally over the cache — O(L)/step —
+    instead of re-forwarding the whole sequence each token.
+
+    Returns the model's native hybrid cache (full `KVCache` for global layers,
+    `RotatingKVCache(sliding_window)` for sliding layers) populated to the
+    prompt, plus the last-row logits (``mx [V]``) predicting the first token.
+    """
+    import mlx.core as mx  # type: ignore
+    from mlx_lm.models.cache import make_prompt_cache  # type: ignore
+
+    text_model = resolve_mlx_text_model(mlx_model)
+    T = len(list(input_ids))
+    evicted = set(int(p) for p in evicted_positions if 0 <= int(p) < T)
+    evicted_mask = mx.array([p in evicted for p in range(T)])
+
+    cache = make_prompt_cache(mlx_model)
+    with _patched_attention_class(text_model):
+        for idx, layer in enumerate(text_model.layers):
+            attn = layer.self_attn
+            if not bool(getattr(attn, "has_kv", True)):
+                continue  # sharers inherit injected K/V via shared_kv
+            rk = restored_k_per_layer.get(idx)
+            rv = restored_v_per_layer.get(idx)
+            if rk is None:
+                continue
+            attn._kakeya_inject = {
+                "mode": "inject",
+                "evicted_mask": evicted_mask,
+                "restored_k": rk,
+                "restored_v": rv,
+            }
+        ids = mx.array([list(input_ids)])
+        logits = mlx_model(ids, cache=cache)
+        mx.eval(logits)
+    # Context manager restored the original Attention.__call__ → subsequent
+    # decode steps run NATIVE incremental attention over this cache.
+    return cache, logits[0, -1]
+
+
+def restored_incremental_generate(
+    mlx_model: Any,
+    cache: Any,
+    first_logits: Any,
+    *,
+    max_tokens: int,
+    eos_ids: Sequence[int] = (),
+) -> List[int]:
+    """Greedy-decode up to ``max_tokens`` tokens over a restored prefill cache
+    using mlx_lm's native ``generate_step`` (chunked + async-pipelined) — the
+    throughput-critical incremental loop. Recall is carried by the cache's
+    full-attention (S5) layers.
+    """
+    import mlx.core as mx  # type: ignore
+    from mlx_lm.generate import generate_step  # type: ignore
+
+    eos = set(int(t) for t in eos_ids)
+    nxt = int(mx.argmax(first_logits).item())
+    out: List[int] = [nxt]
+    if nxt in eos or max_tokens <= 1:
+        return out
+    # generate_step with a 1-token prompt + prefilled cache skips re-prefill
+    # (its chunked-prefill loop needs >1 prompt token) and decodes incrementally.
+    for tok, _ in generate_step(
+        mx.array([nxt]), mlx_model, prompt_cache=cache, max_tokens=max_tokens - 1,
+    ):
+        t = int(tok)
+        out.append(t)
+        if t in eos:
+            break
+    return out

scripts/research/k3_integrated_niah_eval_mac.py

@@ -63,6 +63,11 @@ def parse_args() -> argparse.Namespace:
     ap.add_argument("--sink-size", type=int, default=4)
     ap.add_argument("--window-size", type=int, default=64)
     ap.add_argument("--max-new-tokens", type=int, default=16)
+    ap.add_argument("--incremental", action="store_true",
+                    help="Use the INCREMENTAL restored decode (MLX Gap-A): "
+                         "prefill captures restored K/V into a persistent cache, "
+                         "decode via mlx_lm generate_step (O(L)/token). Fixes the "
+                         "per-token re-forward throughput collapse. Free-gen only.")
     ap.add_argument("--teacher-forced", action="store_true",
                     help="DIAGNOSTIC ONLY (under-measures retrieval): single "
                          "restored forward per sample, check argmax at the "
@@ -108,6 +113,7 @@ def main() -> int:
         resolve_mlx_text_model, mlx_full_attention_layer_indices,
         kv_source_layer_map, capture_own_kv, restored_logits,
         per_layer_kv_geometry, kv_memory_report,
+        restored_prefill_cache, restored_incremental_generate,
     )
     from inference_engine.v04.kv_compressor import make_default_compressor
     from scripts.research.k3_dflash_mlx_bridge import (
@@ -326,6 +332,37 @@ def eval_free_gen_cross() -> Tuple[List[str], List[float], List[int]]:
                   file=sys.stderr)
         return decoded, lats, toks
 
+    def eval_free_gen_cross_incremental() -> Tuple[List[str], List[float], List[int]]:
+        """INCREMENTAL restored free generation (MLX port of CUDA Gap-A):
+        prefill ONCE capturing restored K/V into a persistent cache, then
+        decode with mlx_lm's native incremental step (O(L)/token). Fixes the
+        per-token re-forward throughput collapse. Recall via S5 full-attn."""
+        decoded, lats, toks = [], [], []
+        for i, pid in enumerate(sample_ids):
+            rk, rv, tsrc = build_restoration(pid)
+            T = len(pid)
+            evicted = compute_evicted_positions(T, args.sink_size, args.window_size)
+            t0 = time.perf_counter()
+            if not evicted:
+                cache = (getattr(mlx_model, "make_cache", lambda: None)())
+                out = mlx_model(mx.array([pid]), cache=cache); mx.eval(out)
+                first = out[0, -1]
+            else:
+                cache, first = restored_prefill_cache(
+                    mlx_model, pid,
+                    restored_k_per_layer=_pad(rk, tsrc, T),
+                    restored_v_per_layer=_pad(rv, tsrc, T),
+                    evicted_positions=evicted)
+            gen = restored_incremental_generate(
+                mlx_model, cache, first,
+                max_tokens=args.max_new_tokens,
+                eos_ids=([eos_id] if eos_id is not None else ()))
+            lats.append(time.perf_counter() - t0)
+            decoded.append(tokenizer.decode(gen)); toks.append(len(gen))
+            print(f"[mac]   incr {i}: T={seq_lens[i]} -> {decoded[-1][:48]!r}",
+                  file=sys.stderr)
+        return decoded, lats, toks
+
     def eval_free_gen_oracle() -> Tuple[List[str], List[float], List[int]]:
         """Oracle free generation using mlx's NATIVE incremental KV cache
         (fast + correct reference; confirms the metric/dataset)."""
@@ -356,11 +393,14 @@ def oracle_logits_all(prompt_ids, full_ids):
 
     label = "identity" if args.identity_restore else (
         "s5" if args.s5_exact_full_attn else "f_theta_all")
-    eval_mode = "teacher_forced" if args.teacher_forced else "free_gen"
+    eval_mode = ("teacher_forced" if args.teacher_forced
+                 else "free_gen_incremental" if args.incremental else "free_gen")
     print(f"[mac] running restored cross-model verifier ({label}, {eval_mode})",
           file=sys.stderr, flush=True)
     if args.teacher_forced:
         cross_dec, cross_lat, cross_tok = eval_teacher_forced(cross_logits_all)
+    elif args.incremental:
+        cross_dec, cross_lat, cross_tok = eval_free_gen_cross_incremental()
     else:
         cross_dec, cross_lat, cross_tok = eval_free_gen_cross()
     cross_res = aggregate_recall("k3_cross_model_mac", samples, cross_dec, cross_lat, cross_tok)