debug(mac): root-cause MLX batched-recall break — upstream mlx_lm gemma-4 batch>1 decode bug (RotatingKVCache ruled out)

docs/adr/0014-agent-connection-capacity-and-cross-host-topology-tests.md

@@ -266,6 +266,29 @@ batching) — recall is the bottom line. Note also the Mac speedup ceiling is lo
 even nominally (M4 saturates at small batch — 2.7× vs CUDA's 8.45×). The
 validated batched scheduler is **CUDA-only** today.
 
+**Root-cause of the Mac batched-recall break (investigated).** The
+`RotatingKVCache` / sliding-window-mask line was investigated and **ruled out**:
+
+| Mac diagnostic | batched per-session recall | what it isolates |
+| --- | --- | --- |
+| short prompt (ctx 339 < sliding window → **no rotation**) | 0.25 | **not** rotation |
+| concat-based Kakeya `SinkWindowKVCache` (no in-place buffer assign) | 0.25 | **not** the cache (in-place or concat) |
+| per-row first decoded token vs serialized | **matches** (all rows) | batched **prefill is correct** |
+
+So the break is **cache-independent and prefill-correct**: batched **decode**
+diverges only *after* the first token, with both mlx_lm's in-place cache and
+Kakeya's concat cache. That localizes it to **mlx_lm 0.31.3's gemma-4
+batched (batch>1) decode forward** (RoPE-offset / mask / shared-KV path), an
+**upstream MLX limitation** — not a Kakeya cache or rotation issue, and not
+present on CUDA (HF transformers batched decode is correct → §3.5/§3.7's 8.04×
+/ 8.45×, recall 1.0). Evidence:
+`results/research/k3_mlx_batched_{diag_short,kakeya_cache}_mac.json`.
+
+**Status:** recall-safe Mac multi-tenant = **serialized** (recall 1.0). Mac
+**batched** throughput needs an upstream mlx_lm gemma-4 batched-decode fix (or a
+custom batched gemma decode kernel) — tracked as a follow-up; CUDA is the
+recall-safe batched path today.
+
 ## 4. Case 2 — cross-host proposer/verifier (FEASIBILITY VERDICT)
 
 ### 4.1 Verdict: the requested topology is not implementable today, and is architecturally bounded out

inference_engine/bridge/manifest.py

@@ -131,6 +131,48 @@ def _harness_preset(
             ),
             timeout_minutes=60,
         ),
+        Preset(
+            name="mlx-batched-kakeya-cache",
+            description="Fix test: MLX batched multi-tenant with Kakeya's "
+                        "concat-based SinkWindowKVCache (S5) instead of "
+                        "mlx_lm's in-place buffer cache — should restore "
+                        "per-session recall at batch>1 + bound memory.",
+            command_templates=(
+                (
+                    "python3", "scripts/research/mlx_batched_multitenant_bench.py",
+                    "--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}",
+                    "--sessions", "8",
+                    "--haystack-lines", "60",
+                    "--max-new-tokens", "24",
+                    "--kakeya-cache", "--sink", "4", "--window", "64",
+                    "--output",
+                    "results/research/k3_mac_bridge_mlx_batched_kakeya_cache.json",
+                ),
+            ),
+            timeout_minutes=90,
+            validate_reports=False,
+        ),
+        Preset(
+            name="mlx-batched-diag-short",
+            description="Diagnostic: MLX batched multi-tenant on SHORT prompts "
+                        "(haystack 8, below the sliding window so no "
+                        "RotatingKVCache rotation) — isolates whether the "
+                        "batched-recall bug is rotation-under-batch. Logs "
+                        "per-row batched-vs-serialized first tokens.",
+            command_templates=(
+                (
+                    "python3", "scripts/research/mlx_batched_multitenant_bench.py",
+                    "--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}",
+                    "--sessions", "4",
+                    "--haystack-lines", "15",
+                    "--max-new-tokens", "16",
+                    "--output",
+                    "results/research/k3_mac_bridge_mlx_batched_diag_short.json",
+                ),
+            ),
+            timeout_minutes=60,
+            validate_reports=False,
+        ),
         Preset(
             name="mlx-batched-multitenant",
             description="Mac analog of the §3.7 batched scheduler: N sessions "

results/research/k3_mlx_batched_diag_short_mac.json

@@ -0,0 +1,18 @@
+{
+  "kind": "mlx_batched_multitenant",
+  "config": {
+    "sessions": 4,
+    "modal_prompt_len": 339,
+    "max_new_tokens": 16,
+    "verifier_path": "/Users/fluffy314/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit"
+  },
+  "serialized": {
+    "aggregate_tps": 25.241,
+    "recall": 1.0
+  },
+  "batched": {
+    "aggregate_tps": 23.945,
+    "recall": 0.25
+  },
+  "batched_speedup_vs_serialized": 0.95
+}

results/research/k3_mlx_batched_kakeya_cache_mac.json

@@ -0,0 +1,18 @@
+{
+  "kind": "mlx_batched_multitenant",
+  "config": {
+    "sessions": 8,
+    "modal_prompt_len": 1149,
+    "max_new_tokens": 24,
+    "verifier_path": "/Users/fluffy314/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit"
+  },
+  "serialized": {
+    "aggregate_tps": 25.801,
+    "recall": 1.0
+  },
+  "batched": {
+    "aggregate_tps": 7.879,
+    "recall": 0.25
+  },
+  "batched_speedup_vs_serialized": 0.31
+}

scripts/research/mlx_batched_multitenant_bench.py

@@ -31,18 +31,48 @@ def main() -> int:
     ap.add_argument("--haystack-lines", type=int, default=60)
     ap.add_argument("--max-new-tokens", type=int, default=24)
     ap.add_argument("--prefill-chunk", type=int, default=512)
+    ap.add_argument("--kakeya-cache", action="store_true",
+                    help="Build the batched cache from Kakeya's concat-based "
+                         "SinkWindowKVCache (avoids mlx_lm's in-place "
+                         "buffer-assignment decode path that breaks at batch>1). "
+                         "S5: full-attn layers keep all, sliding bounded.")
+    ap.add_argument("--window", type=int, default=64,
+                    help="sliding-layer window when --kakeya-cache (S5).")
+    ap.add_argument("--sink", type=int, default=4)
+    ap.add_argument("--full-window", type=int, default=100000,
+                    help="full-attn-layer window when --kakeya-cache "
+                         "(large = keep all, exact recall).")
     ap.add_argument("--output", default=None)
     args = ap.parse_args()
 
     import mlx.core as mx
     import mlx_lm
     sys.path.insert(0, "sdks/python")
     from inference_engine.v04 import make_niah_dataset
+    from inference_engine.backends.mlx.cache import SinkWindowKVCache
+    from inference_engine.backends.mlx.cross_model_dlm_verifier import (
+        resolve_mlx_text_model, mlx_full_attention_layer_indices,
+    )
 
     print(f"[mlx-mt] loading {args.verifier_path}", flush=True)
     model, tok = mlx_lm.load(args.verifier_path)
     N = args.sessions
 
+    text_model = resolve_mlx_text_model(model)
+    full_idx = set(mlx_full_attention_layer_indices(text_model))
+
+    def new_cache():
+        if not args.kakeya_cache:
+            return model.make_cache()
+        # S5 hybrid via concat-based caches: full-attn layers keep all (large
+        # window = exact recall), sliding layers bounded to sink+window.
+        return [
+            SinkWindowKVCache(sink_size=args.sink,
+                              window_size=(args.full_window if li in full_idx
+                                           else args.window))
+            for li in range(len(text_model.layers))
+        ]
+
     def encode(text):
         # Match the working Mac NIAH harness: neutral filler + a direct-answer
         # instruction, and append Gemma-4's content-channel marker so short
@@ -79,7 +109,7 @@ def recall(toks, ans):
 
     def prefill_batched(ids_2d):
         """Chunked batched prefill -> (cache, last_logits[N,V])."""
-        cache = model.make_cache()
+        cache = new_cache()
         chunk = args.prefill_chunk
         T = len(ids_2d[0])
         last = None
@@ -134,6 +164,17 @@ def decode_batched(cache, logits, max_tokens):
     serial_tps = round((N * args.max_new_tokens) / ser_decode_s, 3) if ser_decode_s else 0.0
     serial_recall = sum(recall(g_s[i], answers[i]) for i in range(N)) / N
 
+    # Diagnostic: per-row batched-vs-serialized first token + recall, to
+    # localize whether batched PREFILL diverges from serialized (batch-1).
+    print("[mlx-mt][diag] row | serial_tok0 | batched_tok0 | match | "
+          "serial_recall | batched_recall", flush=True)
+    for i in range(N):
+        s0 = g_s[i][0] if g_s[i] else None
+        b0 = g_b[i][0] if g_b[i] else None
+        print(f"[mlx-mt][diag] {i:2d} | {s0} | {b0} | {s0 == b0} | "
+              f"{recall(g_s[i], answers[i])} | {recall(g_b[i], answers[i])}",
+              flush=True)
+
     speedup = round(batched_tps / serial_tps, 2) if serial_tps else None
     report = {
         "kind": "mlx_batched_multitenant",

tests/inference_engine/bridge/test_manifest.py

@@ -72,6 +72,8 @@ def test_allowlist_contains_exactly_the_documented_presets():
         "k3-step2-fused",
         "k3-step2-fused-allmlx",
         "mlx-backend-tests",
+        "mlx-batched-diag-short",
+        "mlx-batched-kakeya-cache",
         "mlx-batched-multitenant",
         "mlx-env-probe",
         "mlx-multitenant-pressure",