Address PR109 Mac validation review

Use fair e2e prefill+decode timing for cross/oracle comparisons, chunk long-context MLX prefill paths, and record ctx280 n=5/gen32 evidence showing Step 2 recall parity and speedup under the corrected gate.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 2 people

docs/pr109-mac-ctx280-validation.md

@@ -0,0 +1,47 @@
+# PR109 Mac ctx280 Validation
+
+This note records the review-driven rerun for PR #109 after fixing the
+measurement issues called out in review.
+
+## Review Corrections
+
+- Fair timing: cross and oracle now report the same `e2e_prefill_plus_decode`
+  scope, plus per-sample `prefill_s`, `decode_s`, and `e2e_s`.
+- Chunked prefill: MLX prompt prefill now uses `--prefill-chunk-size` to avoid
+  the long-context one-shot forward path that can OOM.
+- Adaptive native path: Step 2 adaptive S5 native skips `build_restoration`,
+  f_theta restoration, and aux capture.
+- Gemma4 stop tokens: `<turn|>` is treated as a generation stop token alongside
+  `<eos>`.
+- Gate scale: validation was rerun with `n=5`, `max_new_tokens=32`, and
+  haystack lines `238..322`, producing prompt lengths `4406..5810`.
+
+## Command
+
+```bash
+PYTHONPATH=.:sdks/python python scripts/research/k3_integrated_niah_eval_mac.py \
+  --verifier-path /Users/fluffy314/Documents/Kakeya-LLM-Inference-engine-pr94-resolve/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 --fused-specdecode --block-size 4 \
+  --n-samples 5 --haystack-min-lines 238 --haystack-max-lines 322 \
+  --max-new-tokens 32 --prefill-chunk-size 512 --decode-warmup-tokens 1 \
+  --output results/research/k3_mlx_fused_fair_ctx280_n5_gen32_20260612_105807.json
+```
+
+## Result
+
+- Recall: cross `5/5 = 1.0`, oracle `5/5 = 1.0`, delta `0pp`.
+- Prompt lengths: `4406..5810` tokens.
+- Timing scope: `e2e_prefill_plus_decode` for both cross and oracle.
+- Cross Step 2 throughput: `0.2217 tok/s` (`39 tok / 175.893s`).
+- Oracle AR throughput: `0.0858 tok/s` (`39 tok / 454.484s`).
+- Speedup vs oracle AR: `2.584x`.
+- KV memory: S5 `132.92 MB`, naive full KV `1308.88 MB`, savings `89.8%`.
+
+## Interpretation
+
+This validation supports Step 2 adaptive S5 native under the corrected e2e
+measurement scope at ctx280 scale on the tested Mac setup. It does not claim
+that Step 1 incremental is fixed; earlier evidence still shows Step 1 remains
+slow and should be treated as a separate optimization target.

inference_engine/backends/mlx/cross_model_dlm_verifier.py

@@ -375,6 +375,7 @@ def restored_prefill_cache(
     restored_k_per_layer: Dict[int, Any],
     restored_v_per_layer: Dict[int, Any],
     evicted_positions: Sequence[int],
+    prefill_chunk_size: int = 0,
 ):
     """Prefill ONCE with restoration, capturing the restored K/V into a
     persistent mlx_lm prompt cache; return ``(cache, last_logits)``.
@@ -396,11 +397,32 @@ def restored_prefill_cache(
     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)
-    touched = []
-    try:
+
+    def _slice_restored(a, start: int, end: int):
+        if a is None:
+            return None
+        try:
+            return a[:, start:end, :, :]
+        except Exception:
+            # Linux fake tests use sentinel objects rather than tensors.
+            return a
+
+    def _clear(touched):
+        for obj in touched:
+            for name in (
+                "_kakeya_inject",
+                "kakeya_evicted_mask",
+                "kakeya_restored_pre_keys",
+                "kakeya_restored_pre_values",
+            ):
+                if hasattr(obj, name):
+                    delattr(obj, name)
+
+    def _attach_chunk(start: int, end: int):
+        evicted_mask = mx.array([p in evicted for p in range(start, end)])
+        touched = []
+        needs_attention_patch = False
         for idx, layer in enumerate(text_model.layers):
             attn = layer.self_attn
             if idx >= len(cache) or not bool(getattr(attn, "has_kv", True)):
@@ -410,22 +432,47 @@ def restored_prefill_cache(
             if rk is None:
                 continue
             c = cache[idx]
-            c.kakeya_evicted_mask = evicted_mask
-            c.kakeya_restored_pre_keys = rk
-            c.kakeya_restored_pre_values = rv
-            touched.append(c)
-        ids = mx.array([list(input_ids)])
+            try:
+                c.kakeya_evicted_mask = evicted_mask
+                c.kakeya_restored_pre_keys = _slice_restored(rk, start, end)
+                c.kakeya_restored_pre_values = _slice_restored(rv, start, end)
+                touched.append(c)
+            except Exception:
+                attn._kakeya_inject = {
+                    "mode": "inject",
+                    "evicted_mask": evicted_mask,
+                    "restored_k": _slice_restored(rk, start, end),
+                    "restored_v": _slice_restored(rv, start, end),
+                }
+                touched.append(attn)
+                needs_attention_patch = True
+        return touched, needs_attention_patch
+
+    ids_list = list(input_ids)
+    chunk = int(prefill_chunk_size or 0)
+    if chunk <= 0 or T <= chunk:
+        chunks = [(0, T)]
+    else:
+        chunks = [(s, min(s + chunk, T)) for s in range(0, T, chunk)]
+
+    logits = None
+    for start, end in chunks:
+        touched, needs_attention_patch = _attach_chunk(start, end)
+        try:
+            ids = mx.array([ids_list[start:end]])
+            if needs_attention_patch:
+                with _patched_attention_class(text_model):
+                    logits = mlx_model(ids, cache=cache)
+                    mx.eval(logits)
+            else:
+                logits = mlx_model(ids, cache=cache)
+                mx.eval(logits)
+        finally:
+            _clear(touched)
+    if logits is None:
+        ids = mx.array([ids_list])
         logits = mlx_model(ids, cache=cache)
         mx.eval(logits)
-    finally:
-        for c in touched:
-            for name in (
-                "kakeya_evicted_mask",
-                "kakeya_restored_pre_keys",
-                "kakeya_restored_pre_values",
-            ):
-                if hasattr(c, name):
-                    delattr(c, name)
     # Subsequent decode steps run native incremental attention over this cache.
     return cache, logits[0, -1]
 

inference_engine/backends/mlx/fused_specdecode.py

@@ -169,6 +169,7 @@ def prefill(
         restored_k_per_layer: Dict[int, Any],
         restored_v_per_layer: Dict[int, Any],
         evicted_positions: Sequence[int],
+        prefill_chunk_size: int = 0,
     ) -> None:
         if not prompt_ids:
             raise ValueError("prompt_ids must be non-empty")
@@ -178,6 +179,7 @@ def prefill(
             restored_k_per_layer=restored_k_per_layer,
             restored_v_per_layer=restored_v_per_layer,
             evicted_positions=evicted_positions,
+            prefill_chunk_size=prefill_chunk_size,
         )
         self._past_len = len(prompt_ids)