debug(mlx): root-cause Mac batched-decode bug to MLX core quantized/rope kernel (B>1, L=1) — full ablation

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

@@ -289,6 +289,35 @@ present on CUDA (HF transformers batched decode is correct → §3.5/§3.7's 8.0
 custom batched gemma decode kernel) — tracked as a follow-up; CUDA is the
 recall-safe batched path today.
 
+**Deep localization (per-layer logits + ablations).** A per-layer hidden-diff
+instrument (`mlx_batched_layer_diff_diag.py`) + targeted ablations narrowed the
+batch>1 decode bug to a single op class. At decode step 1, batched **row 0 is
+bit-exact** vs serialized while **row 1+ diverge starting at layer 0**, with
+**layer-0 input identical** (embedding/per-layer-input correct). Every
+Python-level cause was **ruled out** with evidence:
+
+| candidate | test | result |
+| --- | --- | --- |
+| sliding-window rotation | short prompt (no rotation) | still diverges |
+| in-place cache write | concat `SinkWindowKVCache` | still diverges |
+| KV-shared layers | model has **0** shared layers (`first_kv_shared_idx=30`) | N/A |
+| embedding / per-layer-input | layer-0 **input** diff = 0 for all rows | ruled out |
+| attention mask | decode mask is `None` at this offset | ruled out |
+| `mx.fast.scaled_dot_product_attention` | manual matmul-softmax SDPA | still diverges |
+
+With identical layer-0 input + concat cache + manual SDPA, row 1 *still* breaks,
+and **prefill (L=327) is correct while decode (L=1) breaks** → the residual is
+an **MLX core-kernel bug for 4-bit-quantized *batched single-token* decode**
+(`mx.quantized_matmul` / `mx.fast.rope` at `B>1, L=1`) — below the Python layer,
+**not patchable in this repo**. (CUDA is unaffected: HF transformers, no MLX
+quantized kernels — §3.5/§3.7 keep recall 1.0.)
+
+**Outcome:** the Mac batched path is **blocked upstream in MLX**, now precisely
+characterized for an upstream report. Recall-safe Mac multi-tenant remains
+**serialized**; a Python workaround (e.g. L≥2 padded decode, or de-quantized
+projections) is a possible future probe. Evidence:
+`results/research/k3_mlx_batched_manual_sdpa_mac.json` + the layer-diff logs.
+
 ## 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,56 @@ def _harness_preset(
             ),
             timeout_minutes=60,
         ),
+        Preset(
+            name="mlx-batched-layer-diff",
+            description="Localize the mlx_lm gemma-4 batch>1 decode bug: "
+                        "per-layer hidden-state diff (batched row-i vs "
+                        "serialized-i) at decode step 1; prints the first "
+                        "divergent layer + its type/shared-KV status.",
+            command_templates=(
+                (
+                    "python3", "scripts/research/mlx_batched_layer_diff_diag.py",
+                    "--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}",
+                    "--rows", "2", "--haystack-lines", "15",
+                ),
+            ),
+            timeout_minutes=60,
+            validate_reports=False,
+        ),
+        Preset(
+            name="mlx-batched-manual-sdpa",
+            description="Candidate fix: MLX batched multi-tenant with a manual "
+                        "matmul-softmax SDPA replacing mx.fast.scaled_dot_"
+                        "product_attention (works around the batch>1 + GQA "
+                        "fast-kernel bug). Expect per-session recall -> 1.0.",
+            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", "--manual-sdpa",
+                    "--output",
+                    "results/research/k3_mac_bridge_mlx_batched_manual_sdpa.json",
+                ),
+            ),
+            timeout_minutes=90,
+            validate_reports=False,
+        ),
+        Preset(
+            name="mlx-batched-layer-diff-concat",
+            description="Layer-diff with the concat SinkWindowKVCache (no "
+                        "in-place write) — if layer-0 output then matches, the "
+                        "in-place cache write is the batch>1 bug.",
+            command_templates=(
+                (
+                    "python3", "scripts/research/mlx_batched_layer_diff_diag.py",
+                    "--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}",
+                    "--rows", "2", "--haystack-lines", "15", "--kakeya-cache",
+                ),
+            ),
+            timeout_minutes=60,
+            validate_reports=False,
+        ),
         Preset(
             name="mlx-batched-kakeya-cache",
             description="Fix test: MLX batched multi-tenant with Kakeya's "

results/research/k3_mlx_batched_manual_sdpa_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": 16.115,
+    "recall": 1.0
+  },
+  "batched": {
+    "aggregate_tps": 0.955,
+    "recall": 0.125
+  },
+  "batched_speedup_vs_serialized": 0.06
+}

scripts/research/mlx_batched_layer_diff_diag.py

@@ -0,0 +1,152 @@
+"""Localize the mlx_lm gemma-4 batch>1 decode bug: per-layer hidden-state diff,
+batched-row-i vs serialized-i, at decode step 1.
+
+Prefill is known correct (first decoded token matches per row); the divergence
+is in the batched decode forward. This dumps, per decoder layer, the max-abs
+difference between the batched forward's row-i output and the serialized
+single-row forward's output for the SAME session + SAME fed token. The first
+layer whose diff jumps locates the bug (layer 0 → RoPE/embed/per-layer-input;
+a sliding layer → sliding mask/cache; a full-attn layer → full path; the first
+KV-shared layer → shared-KV plumbing).
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from collections import Counter
+from typing import List
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--verifier-path", required=True)
+    ap.add_argument("--rows", type=int, default=2)
+    ap.add_argument("--haystack-lines", type=int, default=15)
+    ap.add_argument("--kakeya-cache", action="store_true",
+                    help="use concat-based SinkWindowKVCache (large window, no "
+                         "eviction) instead of mlx_lm's in-place cache, to test "
+                         "whether the in-place write is the batch>1 bug.")
+    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.cross_model_dlm_verifier import (
+        resolve_mlx_text_model,
+    )
+    from inference_engine.backends.mlx.cache import SinkWindowKVCache
+
+    model, tok = mlx_lm.load(args.verifier_path)
+    text_model = resolve_mlx_text_model(model)
+    layers = text_model.layers
+
+    def new_cache():
+        if not args.kakeya_cache:
+            return model.make_cache()
+        return [SinkWindowKVCache(sink_size=4, window_size=200000)
+                for _ in range(len(layers))]
+    n_layers = len(layers)
+    layer_types = [getattr(l, "layer_type", "?") for l in layers]
+    first_shared = n_layers - getattr(model.args, "num_kv_shared_layers", 0)
+    print(f"[diff] layers={n_layers} first_kv_shared_idx={first_shared}", flush=True)
+
+    R = args.rows
+    pool = make_niah_dataset(n_samples=R * 3, haystack_min_lines=args.haystack_lines,
+                             haystack_max_lines=args.haystack_lines, seed=0)
+
+    def encode(text):
+        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."
+        ids = list(tok.apply_chat_template([{"role": "user", "content": text}],
+                                           add_generation_prompt=True))
+        try:
+            m = tok.encode("<|channel>content\n<channel|>", add_special_tokens=False)
+        except TypeError:
+            m = tok.encode("<|channel>content\n<channel|>")
+        ids.extend(list(m if not hasattr(m, "tolist") else m.tolist()))
+        return ids
+
+    enc = [encode(s.prompt_text) for s in pool]
+    modal = Counter(len(e) for e in enc).most_common(1)[0][0]
+    prompts = [e for e in enc if len(e) == modal][:R]
+    while len(prompts) < R:
+        prompts += prompts[: R - len(prompts)]
+    print(f"[diff] {R} rows, prompt len={modal}", flush=True)
+
+    # capture per-layer output hidden by monkey-patching DecoderLayer.__call__
+    DecoderLayer = type(layers[0])
+    orig_call = DecoderLayer.__call__
+    captured: List = []
+
+    captured_in: List = []
+
+    def patched(self, *a, **k):
+        captured_in.append(a[0])     # layer INPUT h: [B, L, D]
+        out = orig_call(self, *a, **k)
+        captured.append(out[0])      # layer OUTPUT h: [B, L, D]
+        return out
+
+    def prefill(ids_2d):
+        cache = new_cache()
+        out = model(mx.array(ids_2d), cache=cache)
+        mx.eval(out)
+        return cache, out[:, -1, :]
+
+    def decode_capture(token_ids_2d, cache):
+        captured.clear()
+        captured_in.clear()
+        DecoderLayer.__call__ = patched
+        try:
+            out = model(mx.array(token_ids_2d), cache=cache)
+            mx.eval(out)
+        finally:
+            DecoderLayer.__call__ = orig_call
+        return list(captured), list(captured_in), out
+
+    # serialized: prefill + first token + capture decode step
+    ser_tok0, ser_out, ser_in = [], [], []
+    for i in range(R):
+        c, lg = prefill([prompts[i]])
+        t0 = int(mx.argmax(lg, axis=-1).item())
+        ser_tok0.append(t0)
+        co, ci, _ = decode_capture([[t0]], c)
+        ser_out.append(co)
+        ser_in.append(ci)
+
+    # batched: prefill + first tokens + capture decode step (same tokens)
+    cb, lgb = prefill(prompts)
+    bat_tok0 = [int(mx.argmax(lgb[i], axis=-1).item()) for i in range(R)]
+    caps_b, caps_in_b, _ = decode_capture([[t] for t in bat_tok0], cb)
+
+    print(f"[diff] tok0 serial={ser_tok0} batched={bat_tok0} "
+          f"match={ser_tok0 == bat_tok0}", flush=True)
+    # layer-0 INPUT diff isolates embedding/per-layer-input (pre-attention)
+    in0 = [round(float(mx.max(mx.abs(caps_in_b[0][i:i + 1] - ser_in[i][0])).item()), 4)
+           for i in range(R)]
+    print(f"[diff] layer-0 INPUT (embedding) max|Δ| per row = {in0}  "
+          f"(non-zero row>0 => embed/per-layer-input bug; zero => attention bug)",
+          flush=True)
+    print("[diff] layer | type | in max|Δ| | out max|Δ|  (per row)", flush=True)
+    first_div = None
+    for li in range(n_layers):
+        din = [round(float(mx.max(mx.abs(caps_in_b[li][i:i + 1] - ser_in[i][li]).astype(mx.float32)).item()), 3)
+               for i in range(R)]
+        dout = [round(float(mx.max(mx.abs(caps_b[li][i:i + 1] - ser_out[i][li]).astype(mx.float32)).item()), 3)
+                for i in range(R)]
+        mark = ""
+        if first_div is None and max(dout) > 1e-2:
+            first_div = li
+            mark = "  <-- FIRST OUT DIVERGENCE"
+        print(f"[diff] {li:2d} | {layer_types[li]:18s} | in={din} | out={dout}{mark}",
+              flush=True)
+    print(f"[diff] FIRST DIVERGENT LAYER = {first_div} "
+          f"(type={layer_types[first_div] if first_div is not None else None}, "
+          f"shared={first_div is not None and first_div >= first_shared})", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/research/mlx_batched_multitenant_bench.py

@@ -42,6 +42,11 @@ def main() -> int:
     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("--manual-sdpa", action="store_true",
+                    help="Replace mlx_lm gemma's mx.fast.scaled_dot_product_"
+                         "attention with a manual batched matmul-softmax SDPA "
+                         "(works around the suspected batch>1 + GQA fast-kernel "
+                         "bug). The candidate fix.")
     ap.add_argument("--output", default=None)
     args = ap.parse_args()
 
@@ -54,6 +59,36 @@ def main() -> int:
         resolve_mlx_text_model, mlx_full_attention_layer_indices,
     )
 
+    if args.manual_sdpa:
+        import mlx_lm.models.gemma4_text as g4
+
+        def _manual_sdpa(queries, keys, values, cache=None, scale=1.0, mask=None,
+                         sinks=None):
+            # queries [B, n_heads, L, D]; keys/values [B, n_kv, S, D] (GQA).
+            n_heads = queries.shape[1]
+            n_kv = keys.shape[1]
+            if n_kv != n_heads:
+                rep = n_heads // n_kv
+                keys = mx.repeat(keys, rep, axis=1)
+                values = mx.repeat(values, rep, axis=1)
+            scores = (queries * scale) @ mx.swapaxes(keys, -1, -2)  # [B,h,L,S]
+            if mask is not None:
+                if isinstance(mask, str):   # "causal"
+                    qL, kL = scores.shape[-2], scores.shape[-1]
+                    qi = mx.arange(kL - qL, kL)[:, None]
+                    ki = mx.arange(kL)[None]
+                    bmask = qi >= ki
+                    scores = mx.where(bmask, scores, mx.finfo(scores.dtype).min)
+                elif mask.dtype == mx.bool_:
+                    scores = mx.where(mask, scores, mx.finfo(scores.dtype).min)
+                else:
+                    scores = scores + mask
+            scores = mx.softmax(scores, axis=-1, precise=True)
+            return scores @ values
+
+        g4.scaled_dot_product_attention = _manual_sdpa
+        print("[mlx-mt] patched gemma SDPA -> manual batched matmul", flush=True)
+
     print(f"[mlx-mt] loading {args.verifier_path}", flush=True)
     model, tok = mlx_lm.load(args.verifier_path)
     N = args.sessions

tests/inference_engine/bridge/test_manifest.py

@@ -74,6 +74,9 @@ def test_allowlist_contains_exactly_the_documented_presets():
         "mlx-backend-tests",
         "mlx-batched-diag-short",
         "mlx-batched-kakeya-cache",
+        "mlx-batched-layer-diff",
+        "mlx-batched-layer-diff-concat",
+        "mlx-batched-manual-sdpa",
         "mlx-batched-multitenant",
         "mlx-env-probe",
         "mlx-multitenant-pressure",