Add single-fused probe to classify the Metal two-phase instability

fused_specdecode_generate_mlx_trim(single_fused=True): skip the two-phase eval so
drafter+26B fuse into ONE graph (the b876-pathological path); report per-block eval
times (first8/max/mean). Harness --single-fused + preset k3-fused-singlefused-probe
(n=2,gen=16 so a pathological block is bounded). Classifies fundamental command-buffer
limit (eval scales w/ graph) vs fixable SDPA fallback (eval huge even at small scale).

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

Author: cursoragent

inference_engine/backends/mlx/fused_specdecode.py

@@ -373,6 +373,7 @@ def fused_specdecode_generate_mlx_trim(
     gen_tokens: int,
     block_size: int,
     eos_ids: Sequence[int] = (),
+    single_fused: bool = False,
 ) -> Dict[str, Any]:
     """CUDA-parity fused spec decode: KEEP accepted K/V, TRIM only the rejected
     tail (no rollback, no carry re-forward). Requires the adapter to be
@@ -397,6 +398,7 @@ def fused_specdecode_generate_mlx_trim(
 
     generated: List[int] = []
     accepts: List[int] = []
+    block_evals: List[float] = []
     ctx_len = C
     try:
         while len(generated) < gen_tokens:
@@ -410,7 +412,11 @@ def fused_specdecode_generate_mlx_trim(
                     ctx_kv, bonus_id, embed_fn, lm_head_fn,
                     n_masks=n_draft, context_len=base)
                 check_ids = mx.concatenate([bonus_id[None], drafts])   # [L]
-                mx.eval(check_ids)   # two-phase (drafter graph before 26B graph)
+                if not single_fused:
+                    mx.eval(check_ids)   # two-phase (drafter graph before 26B)
+                # single_fused=True → leave check_ids LAZY so the drafter and
+                # 26B verify fuse into ONE graph (the path b876 found Metal-
+                # pathological); this probe times it to classify the instability.
             else:
                 check_ids = bonus_id[None]
             block_logits = adapter.forward_block_lazy(check_ids[None])  # [L, V]
@@ -426,7 +432,9 @@ def fused_specdecode_generate_mlx_trim(
             timing["build_s"] += time.perf_counter() - t_build
             t_eval = time.perf_counter()
             mx.eval(accepted_mx, check_ids)
-            timing["eval_s"] += time.perf_counter() - t_eval
+            blk_eval = time.perf_counter() - t_eval
+            timing["eval_s"] += blk_eval
+            block_evals.append(round(blk_eval, 4))
             accepted = int(accepted_mx.item())
             check = [int(x) for x in check_ids.tolist()]
             commit = check[:accepted]
@@ -461,7 +469,13 @@ def fused_specdecode_generate_mlx_trim(
         "mean_accept_len": (round(sum(accepts) / len(accepts), 3)
                             if accepts else 0.0),
         "decode_tokens": len(generated),
-        "loop": "mlx_trim_keep_accepted_cuda_parity",
+        "loop": ("mlx_trim_single_fused_probe" if single_fused
+                 else "mlx_trim_keep_accepted_cuda_parity"),
+        "single_fused": bool(single_fused),
+        "block_eval_s_first8": block_evals[:8],
+        "block_eval_s_max": (round(max(block_evals), 4) if block_evals else None),
+        "block_eval_s_mean": (round(sum(block_evals) / len(block_evals), 4)
+                              if block_evals else None),
         "time_breakdown_s": {k: round(v, 3) for k, v in timing.items()},
     }
 

inference_engine/bridge/manifest.py

@@ -235,6 +235,37 @@ def _harness_preset(
             timeout_minutes=45,
             params={"path": ("path:tests", None)},
         ),
+        Preset(
+            name="k3-fused-singlefused-probe",
+            description="PROBE: single-fused (one drafter+26B graph) vs two-phase, "
+                        "to classify the Metal instability. Small (n=2, gen=16) so a "
+                        "pathological per-block eval is bounded. Compare block_eval_s "
+                        "vs k3-fused-allmlx-code-trim (two-phase).",
+            command_templates=(
+                (
+                    "python3", "scripts/research/k3_integrated_niah_eval_mac.py",
+                    "--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}",
+                    "--drafter-id", "${ENV:KAKEYA_MAC_DRAFTER_ID}",
+                    "--f-theta-dir", "${ENV:KAKEYA_MAC_FTHETA_DIR}",
+                    "--s5-exact-full-attn", "--fused-specdecode",
+                    "--all-mlx-drafter", "--code-prompts", "--cuda-trim",
+                    "--single-fused",
+                    "--n-samples", "{n_samples}",
+                    "--max-new-tokens", "{max_new_tokens}",
+                    "--block-size", "{block_size}",
+                    "--prefill-chunk-size", "512",
+                    "--output",
+                    "results/research/k3_mac_bridge_k3_fused_singlefused_probe.json",
+                ),
+            ),
+            timeout_minutes=60,
+            params={
+                "n_samples": ("int:n_samples", "2"),
+                "max_new_tokens": ("int:max_new_tokens", "16"),
+                "block_size": ("int:block_size", "4"),
+            },
+            validate_reports=False,
+        ),
         Preset(
             name="k3-fused-allmlx-code-trim",
             description="CUDA-parity rollback test: all-MLX fused + --cuda-trim "

scripts/research/k3_integrated_niah_eval_mac.py

@@ -112,6 +112,11 @@ def parse_args() -> argparse.Namespace:
                          "crossings per block. Requires --s5-exact-full-attn "
                          "(the all-MLX path uses native-S5 injection; the "
                          "f_theta sliding restoration path stays torch).")
+    ap.add_argument("--single-fused", action="store_true",
+                    help="PROBE: with --cuda-trim, fuse drafter+verifier into ONE "
+                         "graph (skip the two-phase eval) to classify the Metal "
+                         "instability (fundamental command-buffer vs fixable SDPA "
+                         "fallback). Reports per-block eval times.")
     ap.add_argument("--cuda-trim", action="store_true",
                     help="All-MLX fused with the CUDA-parity rollback: all-KVCache "
                          "verifier layout + native trim_prompt_cache (keep accepted "
@@ -731,7 +736,8 @@ def eval_fused_specdecode() -> Tuple[List[str], List[float], List[int]]:
                         adapter, active_drafter, aux_prompt=aux_prompt,
                         embed_fn=embed_fn, lm_head_fn=lm_head_fn,
                         gen_tokens=args.max_new_tokens,
-                        block_size=args.block_size, eos_ids=end_ids)
+                        block_size=args.block_size, eos_ids=end_ids,
+                        single_fused=args.single_fused)
                 elif mlx_drafter is not None:
                     # Single-sync all-MLX loop (levers ①②③) + v3 carry rollback.
                     res = fused_specdecode_generate_mlx(