Skip to content

Commit 92a8c5e

Browse files
FluffyAIcodecursoragent
FluffyAIcode
and
cursoragent
authored
debug(mac): root-cause MLX batched-recall break — upstream mlx_lm gemma-4 batch>1 decode bug (RotatingKVCache ruled out) (#133)
* diag(mac): per-row batched-vs-serialized first-token logging + short-prompt preset (isolate RotatingKVCache rotation-under-batch) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * diag: haystack 15 (>=10 min, < sliding window) for short-prompt no-rotation diagnostic Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * fix(mac): --kakeya-cache mode — concat-based SinkWindowKVCache (S5) for batched decode, sidestepping mlx_lm in-place buffer batch>1 bug Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * docs(adr0014): root-cause Mac batched-recall break — ruled out RotatingKVCache/rotation + cache type; it's an upstream mlx_lm gemma-4 batch>1 decode bug (prefill correct, decode diverges, cache-independent) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> --------- Co-authored-by: Cursor Agent <cursoragent@cursor.com> Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>
1 parent 7a96615 commit 92a8c5e

6 files changed

Lines changed: 145 additions & 1 deletion

File tree

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -266,6 +266,29 @@ batching) — recall is the bottom line. Note also the Mac speedup ceiling is lo
266266
even nominally (M4 saturates at small batch — 2.7× vs CUDA's 8.45×). The
267267
validated batched scheduler is **CUDA-only** today.
268268

269+
**Root-cause of the Mac batched-recall break (investigated).** The
270+
`RotatingKVCache` / sliding-window-mask line was investigated and **ruled out**:
271+
272+
| Mac diagnostic | batched per-session recall | what it isolates |
273+
| --- | --- | --- |
274+
| short prompt (ctx 339 < sliding window → **no rotation**) | 0.25 | **not** rotation |
275+
| concat-based Kakeya `SinkWindowKVCache` (no in-place buffer assign) | 0.25 | **not** the cache (in-place or concat) |
276+
| per-row first decoded token vs serialized | **matches** (all rows) | batched **prefill is correct** |
277+
278+
So the break is **cache-independent and prefill-correct**: batched **decode**
279+
diverges only *after* the first token, with both mlx_lm's in-place cache and
280+
Kakeya's concat cache. That localizes it to **mlx_lm 0.31.3's gemma-4
281+
batched (batch>1) decode forward** (RoPE-offset / mask / shared-KV path), an
282+
**upstream MLX limitation** — not a Kakeya cache or rotation issue, and not
283+
present on CUDA (HF transformers batched decode is correct → §3.5/§3.7's 8.04×
284+
/ 8.45×, recall 1.0). Evidence:
285+
`results/research/k3_mlx_batched_{diag_short,kakeya_cache}_mac.json`.
286+
287+
**Status:** recall-safe Mac multi-tenant = **serialized** (recall 1.0). Mac
288+
**batched** throughput needs an upstream mlx_lm gemma-4 batched-decode fix (or a
289+
custom batched gemma decode kernel) — tracked as a follow-up; CUDA is the
290+
recall-safe batched path today.
291+
269292
## 4. Case 2 — cross-host proposer/verifier (FEASIBILITY VERDICT)
270293

271294
### 4.1 Verdict: the requested topology is not implementable today, and is architecturally bounded out

‎inference_engine/bridge/manifest.py‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,48 @@ def _harness_preset(
131131
),
132132
timeout_minutes=60,
133133
),
134+
Preset(
135+
name="mlx-batched-kakeya-cache",
136+
description="Fix test: MLX batched multi-tenant with Kakeya's "
137+
"concat-based SinkWindowKVCache (S5) instead of "
138+
"mlx_lm's in-place buffer cache — should restore "
139+
"per-session recall at batch>1 + bound memory.",
140+
command_templates=(
141+
(
142+
"python3", "scripts/research/mlx_batched_multitenant_bench.py",
143+
"--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}",
144+
"--sessions", "8",
145+
"--haystack-lines", "60",
146+
"--max-new-tokens", "24",
147+
"--kakeya-cache", "--sink", "4", "--window", "64",
148+
"--output",
149+
"results/research/k3_mac_bridge_mlx_batched_kakeya_cache.json",
150+
),
151+
),
152+
timeout_minutes=90,
153+
validate_reports=False,
154+
),
155+
Preset(
156+
name="mlx-batched-diag-short",
157+
description="Diagnostic: MLX batched multi-tenant on SHORT prompts "
158+
"(haystack 8, below the sliding window so no "
159+
"RotatingKVCache rotation) — isolates whether the "
160+
"batched-recall bug is rotation-under-batch. Logs "
161+
"per-row batched-vs-serialized first tokens.",
162+
command_templates=(
163+
(
164+
"python3", "scripts/research/mlx_batched_multitenant_bench.py",
165+
"--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}",
166+
"--sessions", "4",
167+
"--haystack-lines", "15",
168+
"--max-new-tokens", "16",
169+
"--output",
170+
"results/research/k3_mac_bridge_mlx_batched_diag_short.json",
171+
),
172+
),
173+
timeout_minutes=60,
174+
validate_reports=False,
175+
),
134176
Preset(
135177
name="mlx-batched-multitenant",
136178
description="Mac analog of the §3.7 batched scheduler: N sessions "

‎results/research/k3_mlx_batched_diag_short_mac.json‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
{
2+
"kind": "mlx_batched_multitenant",
3+
"config": {
4+
"sessions": 4,
5+
"modal_prompt_len": 339,
6+
"max_new_tokens": 16,
7+
"verifier_path": "/Users/fluffy314/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit"
8+
},
9+
"serialized": {
10+
"aggregate_tps": 25.241,
11+
"recall": 1.0
12+
},
13+
"batched": {
14+
"aggregate_tps": 23.945,
15+
"recall": 0.25
16+
},
17+
"batched_speedup_vs_serialized": 0.95
18+
}

‎results/research/k3_mlx_batched_kakeya_cache_mac.json‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
{
2+
"kind": "mlx_batched_multitenant",
3+
"config": {
4+
"sessions": 8,
5+
"modal_prompt_len": 1149,
6+
"max_new_tokens": 24,
7+
"verifier_path": "/Users/fluffy314/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit"
8+
},
9+
"serialized": {
10+
"aggregate_tps": 25.801,
11+
"recall": 1.0
12+
},
13+
"batched": {
14+
"aggregate_tps": 7.879,
15+
"recall": 0.25
16+
},
17+
"batched_speedup_vs_serialized": 0.31
18+
}

‎scripts/research/mlx_batched_multitenant_bench.py‎

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,18 +31,48 @@ def main() -> int:
3131
ap.add_argument("--haystack-lines", type=int, default=60)
3232
ap.add_argument("--max-new-tokens", type=int, default=24)
3333
ap.add_argument("--prefill-chunk", type=int, default=512)
34+
ap.add_argument("--kakeya-cache", action="store_true",
35+
help="Build the batched cache from Kakeya's concat-based "
36+
"SinkWindowKVCache (avoids mlx_lm's in-place "
37+
"buffer-assignment decode path that breaks at batch>1). "
38+
"S5: full-attn layers keep all, sliding bounded.")
39+
ap.add_argument("--window", type=int, default=64,
40+
help="sliding-layer window when --kakeya-cache (S5).")
41+
ap.add_argument("--sink", type=int, default=4)
42+
ap.add_argument("--full-window", type=int, default=100000,
43+
help="full-attn-layer window when --kakeya-cache "
44+
"(large = keep all, exact recall).")
3445
ap.add_argument("--output", default=None)
3546
args = ap.parse_args()
3647

3748
import mlx.core as mx
3849
import mlx_lm
3950
sys.path.insert(0, "sdks/python")
4051
from inference_engine.v04 import make_niah_dataset
52+
from inference_engine.backends.mlx.cache import SinkWindowKVCache
53+
from inference_engine.backends.mlx.cross_model_dlm_verifier import (
54+
resolve_mlx_text_model, mlx_full_attention_layer_indices,
55+
)
4156

4257
print(f"[mlx-mt] loading {args.verifier_path}", flush=True)
4358
model, tok = mlx_lm.load(args.verifier_path)
4459
N = args.sessions
4560

61+
text_model = resolve_mlx_text_model(model)
62+
full_idx = set(mlx_full_attention_layer_indices(text_model))
63+
64+
def new_cache():
65+
if not args.kakeya_cache:
66+
return model.make_cache()
67+
# S5 hybrid via concat-based caches: full-attn layers keep all (large
68+
# window = exact recall), sliding layers bounded to sink+window.
69+
return [
70+
SinkWindowKVCache(sink_size=args.sink,
71+
window_size=(args.full_window if li in full_idx
72+
else args.window))
73+
for li in range(len(text_model.layers))
74+
]
75+
4676
def encode(text):
4777
# Match the working Mac NIAH harness: neutral filler + a direct-answer
4878
# instruction, and append Gemma-4's content-channel marker so short
@@ -79,7 +109,7 @@ def recall(toks, ans):
79109

80110
def prefill_batched(ids_2d):
81111
"""Chunked batched prefill -> (cache, last_logits[N,V])."""
82-
cache = model.make_cache()
112+
cache = new_cache()
83113
chunk = args.prefill_chunk
84114
T = len(ids_2d[0])
85115
last = None
@@ -134,6 +164,17 @@ def decode_batched(cache, logits, max_tokens):
134164
serial_tps = round((N * args.max_new_tokens) / ser_decode_s, 3) if ser_decode_s else 0.0
135165
serial_recall = sum(recall(g_s[i], answers[i]) for i in range(N)) / N
136166

167+
# Diagnostic: per-row batched-vs-serialized first token + recall, to
168+
# localize whether batched PREFILL diverges from serialized (batch-1).
169+
print("[mlx-mt][diag] row | serial_tok0 | batched_tok0 | match | "
170+
"serial_recall | batched_recall", flush=True)
171+
for i in range(N):
172+
s0 = g_s[i][0] if g_s[i] else None
173+
b0 = g_b[i][0] if g_b[i] else None
174+
print(f"[mlx-mt][diag] {i:2d} | {s0} | {b0} | {s0 == b0} | "
175+
f"{recall(g_s[i], answers[i])} | {recall(g_b[i], answers[i])}",
176+
flush=True)
177+
137178
speedup = round(batched_tps / serial_tps, 2) if serial_tps else None
138179
report = {
139180
"kind": "mlx_batched_multitenant",

‎tests/inference_engine/bridge/test_manifest.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,8 @@ def test_allowlist_contains_exactly_the_documented_presets():
7272
"k3-step2-fused",
7373
"k3-step2-fused-allmlx",
7474
"mlx-backend-tests",
75+
"mlx-batched-diag-short",
76+
"mlx-batched-kakeya-cache",
7577
"mlx-batched-multitenant",
7678
"mlx-env-probe",
7779
"mlx-multitenant-pressure",

0 commit comments

Comments
 (0)