test(mac): MLX batched multi-tenant run + honest finding (batched breaks recall on Mac)

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

@@ -246,6 +246,26 @@ the dominant multi-tenant case); dynamic mid-flight arrival + ragged-length
 continuous batching (and the async-gRPC futures glue that lets independent
 `Generate` RPC coroutines feed one batch loop) is the remaining productization.
 
+**Platform note — §3.7 is CUDA only.** The `BatchedDecodeScheduler` and its
+bench are torch/CUDA (they ran on H200); the Mac served-restored path is blocked
+by the MLX-gemma nested-config load gap (§6). A Mac analog
+(`scripts/research/mlx_batched_multitenant_bench.py`, preset
+`mlx-batched-multitenant`) was run on the Mac mini and surfaced a **correctness
+blocker**: batched MLX decode over gemma-4 at batch > 1 **breaks per-session
+recall** (`results/research/k3_mlx_batched_multitenant_mac.json`):
+
+| Mac (M4, 8 sessions) | aggregate tok/s | per-session recall |
+| --- | --- | --- |
+| serialized | 21.2 | **1.0** |
+| batched (MLX, batch>1) | 57.1 (2.7× *if* recall held) | **0.125** ✗ |
+
+So on Mac the **recall-safe** multi-tenant path is **serialized** (recall 1.0);
+batched throughput is *not* shippable there until the MLX batch>1 forward/cache
+correctness is fixed (likely the gemma hybrid/sliding `RotatingKVCache` under
+batching) — recall is the bottom line. Note also the Mac speedup ceiling is low
+even nominally (M4 saturates at small batch — 2.7× vs CUDA's 8.45×). The
+validated batched scheduler is **CUDA-only** 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,31 @@ def _harness_preset(
             ),
             timeout_minutes=60,
         ),
+        Preset(
+            name="mlx-batched-multitenant",
+            description="Mac analog of the §3.7 batched scheduler: N sessions "
+                        "decoded in one batched MLX forward over the gemma "
+                        "verifier vs serialized; reports aggregate tok/s, "
+                        "speedup, per-session recall (recall-preserving native "
+                        "cache).",
+            command_templates=(
+                (
+                    "python3", "scripts/research/mlx_batched_multitenant_bench.py",
+                    "--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}",
+                    "--sessions", "{n_samples}",
+                    "--haystack-lines", "60",
+                    "--max-new-tokens", "{max_new_tokens}",
+                    "--output",
+                    "results/research/k3_mac_bridge_mlx_batched_multitenant.json",
+                ),
+            ),
+            timeout_minutes=90,
+            params={
+                "n_samples": ("int:n_samples", "8"),
+                "max_new_tokens": ("int:max_new_tokens", "24"),
+            },
+            validate_reports=False,
+        ),
         Preset(
             name="agent-capacity-loadtest",
             description="Test case 1: ramp concurrent agent connections "

results/research/k3_mlx_batched_multitenant_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": 21.157,
+    "recall": 1.0
+  },
+  "batched": {
+    "aggregate_tps": 57.144,
+    "recall": 0.125
+  },
+  "batched_speedup_vs_serialized": 2.7
+}

scripts/research/mlx_batched_multitenant_bench.py

@@ -0,0 +1,158 @@
+"""Mac analog of the PR-A3c batched scheduler (§3.7) — MLX, on Mac mini.
+
+The §3.7 BatchedDecodeScheduler is torch/CUDA (it ran on H200). On Apple Silicon
+the equivalent is a batched MLX forward: N sessions decoded in one pass over the
+MLX gemma verifier, each a batch row with its own KV-cache row. This bench
+measures the served-path batching value on the Mac:
+
+  * serialized — each session's decode run alone, summed (the §3.6 behaviour)
+  * batched    — all N decoded in one batched forward per step
+
+reporting aggregate decode tok/s, the speedup, and per-session recall (recall is
+the bottom line — uses the gemma-native cache, which preserves recall). Equal-
+length prompts (modal NIAH bucket) keep the batch clean.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from collections import Counter
+from pathlib import Path
+from typing import Any, Dict, List
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--verifier-path", required=True)
+    ap.add_argument("--sessions", type=int, default=8)
+    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("--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
+
+    print(f"[mlx-mt] loading {args.verifier_path}", flush=True)
+    model, tok = mlx_lm.load(args.verifier_path)
+    N = args.sessions
+
+    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
+        # completions don't spend tokens on the thought channel (else recall=0).
+        text = text.replace("and does not contain the answer.",
+                            "and is unrelated filler.")
+        text = (text + "\n\nReturn only the secret code in PREFIX-NNNN format. "
+                       "Do not explain, reason, or add any other text.")
+        ids = list(tok.apply_chat_template([{"role": "user", "content": text}],
+                                           add_generation_prompt=True))
+        try:
+            marker = tok.encode("<|channel>content\n<channel|>",
+                                add_special_tokens=False)
+        except TypeError:
+            marker = tok.encode("<|channel>content\n<channel|>")
+        if hasattr(marker, "tolist"):
+            marker = marker.tolist()
+        ids.extend(list(marker))
+        return ids
+
+    pool = make_niah_dataset(n_samples=N * 3, haystack_min_lines=args.haystack_lines,
+                             haystack_max_lines=args.haystack_lines, seed=0)
+    enc = [(encode(s.prompt_text), s.answer_text) for s in pool]
+    modal = Counter(len(e[0]) for e in enc).most_common(1)[0][0]
+    bucket = [(i, a) for i, a in enc if len(i) == modal][:N]
+    while len(bucket) < N:
+        bucket += bucket[: N - len(bucket)]
+    prompts = [b[0] for b in bucket]
+    answers = [b[1] for b in bucket]
+    print(f"[mlx-mt] {N} sessions, modal prompt len={modal}", flush=True)
+
+    def recall(toks, ans):
+        return ans in tok.decode(toks)
+
+    def prefill_batched(ids_2d):
+        """Chunked batched prefill -> (cache, last_logits[N,V])."""
+        cache = model.make_cache()
+        chunk = args.prefill_chunk
+        T = len(ids_2d[0])
+        last = None
+        for s in range(0, T, chunk):
+            part = [row[s:s + chunk] for row in ids_2d]
+            last = model(mx.array(part), cache=cache)
+            mx.eval(last)
+        return cache, last[:, -1, :]
+
+    def decode_batched(cache, logits, max_tokens):
+        B = logits.shape[0]
+        nxt = mx.argmax(logits, axis=-1)
+        gen = [[int(nxt[i].item())] for i in range(B)]
+        mx.eval(nxt)
+        t0 = time.perf_counter()
+        for _ in range(max_tokens - 1):
+            cur = nxt.reshape(B, 1)
+            out = model(cur, cache=cache)
+            mx.eval(out)
+            nxt = mx.argmax(out[:, -1, :], axis=-1)
+            for i in range(B):
+                gen[i].append(int(nxt[i].item()))
+        dt = time.perf_counter() - t0
+        return gen, dt
+
+    # warmup
+    try:
+        c, l = prefill_batched([prompts[0]] * min(2, N))
+        decode_batched(c, l, 4)
+    except Exception as e:  # noqa: BLE001
+        print(f"[mlx-mt] warmup note: {e}", flush=True)
+
+    # batched
+    cache, logits = prefill_batched(prompts)
+    g_b, dt_b = decode_batched(cache, logits, args.max_new_tokens)
+    batched_tps = round((N * args.max_new_tokens) / dt_b, 3) if dt_b > 0 else 0.0
+    batched_recall = sum(recall(g_b[i], answers[i]) for i in range(N)) / N
+
+    # serialized (one session at a time)
+    t0 = time.perf_counter()
+    g_s = []
+    for i in range(N):
+        c, l = prefill_batched([prompts[i]])
+        gg, _ = decode_batched(c, l, args.max_new_tokens)
+        g_s.append(gg[0])
+    # serialized decode-only time: re-time decode alone (prefill excluded for fair tps)
+    ser_decode_s = 0.0
+    for i in range(N):
+        c, l = prefill_batched([prompts[i]])
+        _, dt = decode_batched(c, l, args.max_new_tokens)
+        ser_decode_s += dt
+    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
+
+    speedup = round(batched_tps / serial_tps, 2) if serial_tps else None
+    report = {
+        "kind": "mlx_batched_multitenant",
+        "config": {"sessions": N, "modal_prompt_len": modal,
+                   "max_new_tokens": args.max_new_tokens,
+                   "verifier_path": args.verifier_path},
+        "serialized": {"aggregate_tps": serial_tps, "recall": round(serial_recall, 3)},
+        "batched": {"aggregate_tps": batched_tps, "recall": round(batched_recall, 3)},
+        "batched_speedup_vs_serialized": speedup,
+    }
+    print(f"[mlx-mt] N={N}: serialized {serial_tps} tok/s (recall {serial_recall}) | "
+          f"batched {batched_tps} tok/s (recall {batched_recall}) | speedup {speedup}x",
+          flush=True)
+    if args.output:
+        Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.output).write_text(json.dumps(report, indent=2))
+        print(f"[mlx-mt] wrote {args.output}", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/inference_engine/bridge/test_manifest.py

@@ -72,6 +72,7 @@ def test_allowlist_contains_exactly_the_documented_presets():
         "k3-step2-fused",
         "k3-step2-fused-allmlx",
         "mlx-backend-tests",
+        "mlx-batched-multitenant",
         "mlx-env-probe",
         "mlx-multitenant-pressure",
         "pytest-path",