fix(mlx): --manual-sdpa candidate fix (manual batched matmul SDPA, works around mx.fast.scaled_dot_product_attention batch>1+GQA bug)

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

Author: cursoragent

inference_engine/bridge/manifest.py

@@ -147,6 +147,25 @@ def _harness_preset(
             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 "

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

@@ -76,6 +76,7 @@ def test_allowlist_contains_exactly_the_documented_presets():
         "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",