diff --git a/docs/adr/0014-agent-connection-capacity-and-cross-host-topology-tests.md b/docs/adr/0014-agent-connection-capacity-and-cross-host-topology-tests.md index 84c0ea71..78447727 100644 --- a/docs/adr/0014-agent-connection-capacity-and-cross-host-topology-tests.md +++ b/docs/adr/0014-agent-connection-capacity-and-cross-host-topology-tests.md @@ -266,6 +266,29 @@ batching) — recall is the bottom line. Note also the Mac speedup ceiling is lo even nominally (M4 saturates at small batch — 2.7× vs CUDA's 8.45×). The validated batched scheduler is **CUDA-only** today. +**Root-cause of the Mac batched-recall break (investigated).** The +`RotatingKVCache` / sliding-window-mask line was investigated and **ruled out**: + +| Mac diagnostic | batched per-session recall | what it isolates | +| --- | --- | --- | +| short prompt (ctx 339 < sliding window → **no rotation**) | 0.25 | **not** rotation | +| concat-based Kakeya `SinkWindowKVCache` (no in-place buffer assign) | 0.25 | **not** the cache (in-place or concat) | +| per-row first decoded token vs serialized | **matches** (all rows) | batched **prefill is correct** | + +So the break is **cache-independent and prefill-correct**: batched **decode** +diverges only *after* the first token, with both mlx_lm's in-place cache and +Kakeya's concat cache. That localizes it to **mlx_lm 0.31.3's gemma-4 +batched (batch>1) decode forward** (RoPE-offset / mask / shared-KV path), an +**upstream MLX limitation** — not a Kakeya cache or rotation issue, and not +present on CUDA (HF transformers batched decode is correct → §3.5/§3.7's 8.04× +/ 8.45×, recall 1.0). Evidence: +`results/research/k3_mlx_batched_{diag_short,kakeya_cache}_mac.json`. + +**Status:** recall-safe Mac multi-tenant = **serialized** (recall 1.0). Mac +**batched** throughput needs an upstream mlx_lm gemma-4 batched-decode fix (or a +custom batched gemma decode kernel) — tracked as a follow-up; CUDA is the +recall-safe batched path 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 diff --git a/inference_engine/bridge/manifest.py b/inference_engine/bridge/manifest.py index b5e3cec6..717ab3b9 100644 --- a/inference_engine/bridge/manifest.py +++ b/inference_engine/bridge/manifest.py @@ -131,6 +131,48 @@ def _harness_preset( ), timeout_minutes=60, ), + Preset( + name="mlx-batched-kakeya-cache", + description="Fix test: MLX batched multi-tenant with Kakeya's " + "concat-based SinkWindowKVCache (S5) instead of " + "mlx_lm's in-place buffer cache — should restore " + "per-session recall at batch>1 + bound memory.", + 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", + "--kakeya-cache", "--sink", "4", "--window", "64", + "--output", + "results/research/k3_mac_bridge_mlx_batched_kakeya_cache.json", + ), + ), + timeout_minutes=90, + validate_reports=False, + ), + 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", "15", + "--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 " diff --git a/results/research/k3_mlx_batched_diag_short_mac.json b/results/research/k3_mlx_batched_diag_short_mac.json new file mode 100644 index 00000000..e5ce52b6 --- /dev/null +++ b/results/research/k3_mlx_batched_diag_short_mac.json @@ -0,0 +1,18 @@ +{ + "kind": "mlx_batched_multitenant", + "config": { + "sessions": 4, + "modal_prompt_len": 339, + "max_new_tokens": 16, + "verifier_path": "/Users/fluffy314/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit" + }, + "serialized": { + "aggregate_tps": 25.241, + "recall": 1.0 + }, + "batched": { + "aggregate_tps": 23.945, + "recall": 0.25 + }, + "batched_speedup_vs_serialized": 0.95 +} \ No newline at end of file diff --git a/results/research/k3_mlx_batched_kakeya_cache_mac.json b/results/research/k3_mlx_batched_kakeya_cache_mac.json new file mode 100644 index 00000000..412796d4 --- /dev/null +++ b/results/research/k3_mlx_batched_kakeya_cache_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": 25.801, + "recall": 1.0 + }, + "batched": { + "aggregate_tps": 7.879, + "recall": 0.25 + }, + "batched_speedup_vs_serialized": 0.31 +} \ No newline at end of file diff --git a/scripts/research/mlx_batched_multitenant_bench.py b/scripts/research/mlx_batched_multitenant_bench.py index cb06e1e2..9910e94e 100644 --- a/scripts/research/mlx_batched_multitenant_bench.py +++ b/scripts/research/mlx_batched_multitenant_bench.py @@ -31,6 +31,17 @@ def main() -> int: 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("--kakeya-cache", action="store_true", + help="Build the batched cache from Kakeya's concat-based " + "SinkWindowKVCache (avoids mlx_lm's in-place " + "buffer-assignment decode path that breaks at batch>1). " + "S5: full-attn layers keep all, sliding bounded.") + ap.add_argument("--window", type=int, default=64, + help="sliding-layer window when --kakeya-cache (S5).") + ap.add_argument("--sink", type=int, default=4) + 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("--output", default=None) args = ap.parse_args() @@ -38,11 +49,30 @@ def main() -> int: import mlx_lm sys.path.insert(0, "sdks/python") from inference_engine.v04 import make_niah_dataset + from inference_engine.backends.mlx.cache import SinkWindowKVCache + from inference_engine.backends.mlx.cross_model_dlm_verifier import ( + resolve_mlx_text_model, mlx_full_attention_layer_indices, + ) print(f"[mlx-mt] loading {args.verifier_path}", flush=True) model, tok = mlx_lm.load(args.verifier_path) N = args.sessions + text_model = resolve_mlx_text_model(model) + full_idx = set(mlx_full_attention_layer_indices(text_model)) + + def new_cache(): + if not args.kakeya_cache: + return model.make_cache() + # S5 hybrid via concat-based caches: full-attn layers keep all (large + # window = exact recall), sliding layers bounded to sink+window. + return [ + SinkWindowKVCache(sink_size=args.sink, + window_size=(args.full_window if li in full_idx + else args.window)) + for li in range(len(text_model.layers)) + ] + 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 @@ -79,7 +109,7 @@ def recall(toks, ans): def prefill_batched(ids_2d): """Chunked batched prefill -> (cache, last_logits[N,V]).""" - cache = model.make_cache() + cache = new_cache() chunk = args.prefill_chunk T = len(ids_2d[0]) last = None @@ -134,6 +164,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", diff --git a/tests/inference_engine/bridge/test_manifest.py b/tests/inference_engine/bridge/test_manifest.py index f78f75d7..abe7a610 100644 --- a/tests/inference_engine/bridge/test_manifest.py +++ b/tests/inference_engine/bridge/test_manifest.py @@ -72,6 +72,8 @@ def test_allowlist_contains_exactly_the_documented_presets(): "k3-step2-fused", "k3-step2-fused-allmlx", "mlx-backend-tests", + "mlx-batched-diag-short", + "mlx-batched-kakeya-cache", "mlx-batched-multitenant", "mlx-env-probe", "mlx-multitenant-pressure",