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 db35d512..84c0ea71 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 @@ -246,6 +246,26 @@ the dominant multi-tenant case); dynamic mid-flight arrival + ragged-length continuous batching (and the async-gRPC futures glue that lets independent `Generate` RPC coroutines feed one batch loop) is the remaining productization. +**Platform note — §3.7 is CUDA only.** The `BatchedDecodeScheduler` and its +bench are torch/CUDA (they ran on H200); the Mac served-restored path is blocked +by the MLX-gemma nested-config load gap (§6). A Mac analog +(`scripts/research/mlx_batched_multitenant_bench.py`, preset +`mlx-batched-multitenant`) was run on the Mac mini and surfaced a **correctness +blocker**: batched MLX decode over gemma-4 at batch > 1 **breaks per-session +recall** (`results/research/k3_mlx_batched_multitenant_mac.json`): + +| Mac (M4, 8 sessions) | aggregate tok/s | per-session recall | +| --- | --- | --- | +| serialized | 21.2 | **1.0** | +| batched (MLX, batch>1) | 57.1 (2.7× *if* recall held) | **0.125** ✗ | + +So on Mac the **recall-safe** multi-tenant path is **serialized** (recall 1.0); +batched throughput is *not* shippable there until the MLX batch>1 forward/cache +correctness is fixed (likely the gemma hybrid/sliding `RotatingKVCache` under +batching) — recall is the bottom line. Note also the Mac speedup ceiling is low +even nominally (M4 saturates at small batch — 2.7× vs CUDA's 8.45×). The +validated batched scheduler is **CUDA-only** 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 04bdf43f..b5e3cec6 100644 --- a/inference_engine/bridge/manifest.py +++ b/inference_engine/bridge/manifest.py @@ -131,6 +131,31 @@ def _harness_preset( ), timeout_minutes=60, ), + Preset( + name="mlx-batched-multitenant", + description="Mac analog of the §3.7 batched scheduler: N sessions " + "decoded in one batched MLX forward over the gemma " + "verifier vs serialized; reports aggregate tok/s, " + "speedup, per-session recall (recall-preserving native " + "cache).", + command_templates=( + ( + "python3", "scripts/research/mlx_batched_multitenant_bench.py", + "--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}", + "--sessions", "{n_samples}", + "--haystack-lines", "60", + "--max-new-tokens", "{max_new_tokens}", + "--output", + "results/research/k3_mac_bridge_mlx_batched_multitenant.json", + ), + ), + timeout_minutes=90, + params={ + "n_samples": ("int:n_samples", "8"), + "max_new_tokens": ("int:max_new_tokens", "24"), + }, + validate_reports=False, + ), Preset( name="agent-capacity-loadtest", description="Test case 1: ramp concurrent agent connections " diff --git a/results/research/k3_mlx_batched_multitenant_mac.json b/results/research/k3_mlx_batched_multitenant_mac.json new file mode 100644 index 00000000..e17f71c1 --- /dev/null +++ b/results/research/k3_mlx_batched_multitenant_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": 21.157, + "recall": 1.0 + }, + "batched": { + "aggregate_tps": 57.144, + "recall": 0.125 + }, + "batched_speedup_vs_serialized": 2.7 +} \ No newline at end of file diff --git a/scripts/research/mlx_batched_multitenant_bench.py b/scripts/research/mlx_batched_multitenant_bench.py new file mode 100644 index 00000000..cb06e1e2 --- /dev/null +++ b/scripts/research/mlx_batched_multitenant_bench.py @@ -0,0 +1,158 @@ +"""Mac analog of the PR-A3c batched scheduler (§3.7) — MLX, on Mac mini. + +The §3.7 BatchedDecodeScheduler is torch/CUDA (it ran on H200). On Apple Silicon +the equivalent is a batched MLX forward: N sessions decoded in one pass over the +MLX gemma verifier, each a batch row with its own KV-cache row. This bench +measures the served-path batching value on the Mac: + + * serialized — each session's decode run alone, summed (the §3.6 behaviour) + * batched — all N decoded in one batched forward per step + +reporting aggregate decode tok/s, the speedup, and per-session recall (recall is +the bottom line — uses the gemma-native cache, which preserves recall). Equal- +length prompts (modal NIAH bucket) keep the batch clean. +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from collections import Counter +from pathlib import Path +from typing import Any, Dict, List + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--verifier-path", required=True) + ap.add_argument("--sessions", type=int, default=8) + 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("--output", default=None) + 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 + + print(f"[mlx-mt] loading {args.verifier_path}", flush=True) + model, tok = mlx_lm.load(args.verifier_path) + N = args.sessions + + 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 + # completions don't spend tokens on the thought channel (else recall=0). + 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. " + "Do not explain, reason, or add any other text.") + ids = list(tok.apply_chat_template([{"role": "user", "content": text}], + add_generation_prompt=True)) + try: + marker = tok.encode("<|channel>content\n", + add_special_tokens=False) + except TypeError: + marker = tok.encode("<|channel>content\n") + if hasattr(marker, "tolist"): + marker = marker.tolist() + ids.extend(list(marker)) + return ids + + pool = make_niah_dataset(n_samples=N * 3, haystack_min_lines=args.haystack_lines, + haystack_max_lines=args.haystack_lines, seed=0) + enc = [(encode(s.prompt_text), s.answer_text) for s in pool] + modal = Counter(len(e[0]) for e in enc).most_common(1)[0][0] + bucket = [(i, a) for i, a in enc if len(i) == modal][:N] + while len(bucket) < N: + bucket += bucket[: N - len(bucket)] + prompts = [b[0] for b in bucket] + answers = [b[1] for b in bucket] + print(f"[mlx-mt] {N} sessions, modal prompt len={modal}", flush=True) + + def recall(toks, ans): + return ans in tok.decode(toks) + + def prefill_batched(ids_2d): + """Chunked batched prefill -> (cache, last_logits[N,V]).""" + cache = model.make_cache() + chunk = args.prefill_chunk + T = len(ids_2d[0]) + last = None + for s in range(0, T, chunk): + part = [row[s:s + chunk] for row in ids_2d] + last = model(mx.array(part), cache=cache) + mx.eval(last) + return cache, last[:, -1, :] + + def decode_batched(cache, logits, max_tokens): + B = logits.shape[0] + nxt = mx.argmax(logits, axis=-1) + gen = [[int(nxt[i].item())] for i in range(B)] + mx.eval(nxt) + t0 = time.perf_counter() + for _ in range(max_tokens - 1): + cur = nxt.reshape(B, 1) + out = model(cur, cache=cache) + mx.eval(out) + nxt = mx.argmax(out[:, -1, :], axis=-1) + for i in range(B): + gen[i].append(int(nxt[i].item())) + dt = time.perf_counter() - t0 + return gen, dt + + # warmup + try: + c, l = prefill_batched([prompts[0]] * min(2, N)) + decode_batched(c, l, 4) + except Exception as e: # noqa: BLE001 + print(f"[mlx-mt] warmup note: {e}", flush=True) + + # batched + cache, logits = prefill_batched(prompts) + g_b, dt_b = decode_batched(cache, logits, args.max_new_tokens) + batched_tps = round((N * args.max_new_tokens) / dt_b, 3) if dt_b > 0 else 0.0 + batched_recall = sum(recall(g_b[i], answers[i]) for i in range(N)) / N + + # serialized (one session at a time) + t0 = time.perf_counter() + g_s = [] + for i in range(N): + c, l = prefill_batched([prompts[i]]) + gg, _ = decode_batched(c, l, args.max_new_tokens) + g_s.append(gg[0]) + # serialized decode-only time: re-time decode alone (prefill excluded for fair tps) + ser_decode_s = 0.0 + for i in range(N): + c, l = prefill_batched([prompts[i]]) + _, dt = decode_batched(c, l, args.max_new_tokens) + ser_decode_s += dt + 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 + + speedup = round(batched_tps / serial_tps, 2) if serial_tps else None + report = { + "kind": "mlx_batched_multitenant", + "config": {"sessions": N, "modal_prompt_len": modal, + "max_new_tokens": args.max_new_tokens, + "verifier_path": args.verifier_path}, + "serialized": {"aggregate_tps": serial_tps, "recall": round(serial_recall, 3)}, + "batched": {"aggregate_tps": batched_tps, "recall": round(batched_recall, 3)}, + "batched_speedup_vs_serialized": speedup, + } + print(f"[mlx-mt] N={N}: serialized {serial_tps} tok/s (recall {serial_recall}) | " + f"batched {batched_tps} tok/s (recall {batched_recall}) | speedup {speedup}x", + flush=True) + if args.output: + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + Path(args.output).write_text(json.dumps(report, indent=2)) + print(f"[mlx-mt] wrote {args.output}", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/inference_engine/bridge/test_manifest.py b/tests/inference_engine/bridge/test_manifest.py index f7ed1c4e..f78f75d7 100644 --- a/tests/inference_engine/bridge/test_manifest.py +++ b/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-multitenant", "mlx-env-probe", "mlx-multitenant-pressure", "pytest-path",