diag(mac): per-row batched-vs-serialized first-token logging + short-prompt preset (isolate RotatingKVCache rotation-under-batch)

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

Author: cursoragent

inference_engine/bridge/manifest.py

@@ -131,6 +131,27 @@ def _harness_preset(
             ),
             timeout_minutes=60,
         ),
+        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", "8",
+                    "--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 "

scripts/research/mlx_batched_multitenant_bench.py

@@ -134,6 +134,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,7 @@ def test_allowlist_contains_exactly_the_documented_presets():
         "k3-step2-fused",
         "k3-step2-fused-allmlx",
         "mlx-backend-tests",
+        "mlx-batched-diag-short",
         "mlx-batched-multitenant",
         "mlx-env-probe",
         "mlx-multitenant-pressure",