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 78447727..98b59ec3 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 @@ -289,6 +289,35 @@ present on CUDA (HF transformers batched decode is correct → §3.5/§3.7's 8.0 custom batched gemma decode kernel) — tracked as a follow-up; CUDA is the recall-safe batched path today. +**Deep localization (per-layer logits + ablations).** A per-layer hidden-diff +instrument (`mlx_batched_layer_diff_diag.py`) + targeted ablations narrowed the +batch>1 decode bug to a single op class. At decode step 1, batched **row 0 is +bit-exact** vs serialized while **row 1+ diverge starting at layer 0**, with +**layer-0 input identical** (embedding/per-layer-input correct). Every +Python-level cause was **ruled out** with evidence: + +| candidate | test | result | +| --- | --- | --- | +| sliding-window rotation | short prompt (no rotation) | still diverges | +| in-place cache write | concat `SinkWindowKVCache` | still diverges | +| KV-shared layers | model has **0** shared layers (`first_kv_shared_idx=30`) | N/A | +| embedding / per-layer-input | layer-0 **input** diff = 0 for all rows | ruled out | +| attention mask | decode mask is `None` at this offset | ruled out | +| `mx.fast.scaled_dot_product_attention` | manual matmul-softmax SDPA | still diverges | + +With identical layer-0 input + concat cache + manual SDPA, row 1 *still* breaks, +and **prefill (L=327) is correct while decode (L=1) breaks** → the residual is +an **MLX core-kernel bug for 4-bit-quantized *batched single-token* decode** +(`mx.quantized_matmul` / `mx.fast.rope` at `B>1, L=1`) — below the Python layer, +**not patchable in this repo**. (CUDA is unaffected: HF transformers, no MLX +quantized kernels — §3.5/§3.7 keep recall 1.0.) + +**Outcome:** the Mac batched path is **blocked upstream in MLX**, now precisely +characterized for an upstream report. Recall-safe Mac multi-tenant remains +**serialized**; a Python workaround (e.g. L≥2 padded decode, or de-quantized +projections) is a possible future probe. Evidence: +`results/research/k3_mlx_batched_manual_sdpa_mac.json` + the layer-diff logs. + ## 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 717ab3b9..330b69f6 100644 --- a/inference_engine/bridge/manifest.py +++ b/inference_engine/bridge/manifest.py @@ -131,6 +131,56 @@ def _harness_preset( ), timeout_minutes=60, ), + Preset( + name="mlx-batched-layer-diff", + description="Localize the mlx_lm gemma-4 batch>1 decode bug: " + "per-layer hidden-state diff (batched row-i vs " + "serialized-i) at decode step 1; prints the first " + "divergent layer + its type/shared-KV status.", + command_templates=( + ( + "python3", "scripts/research/mlx_batched_layer_diff_diag.py", + "--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}", + "--rows", "2", "--haystack-lines", "15", + ), + ), + 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 " + "in-place write) — if layer-0 output then matches, the " + "in-place cache write is the batch>1 bug.", + command_templates=( + ( + "python3", "scripts/research/mlx_batched_layer_diff_diag.py", + "--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}", + "--rows", "2", "--haystack-lines", "15", "--kakeya-cache", + ), + ), + timeout_minutes=60, + validate_reports=False, + ), Preset( name="mlx-batched-kakeya-cache", description="Fix test: MLX batched multi-tenant with Kakeya's " diff --git a/results/research/k3_mlx_batched_manual_sdpa_mac.json b/results/research/k3_mlx_batched_manual_sdpa_mac.json new file mode 100644 index 00000000..41d642da --- /dev/null +++ b/results/research/k3_mlx_batched_manual_sdpa_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": 16.115, + "recall": 1.0 + }, + "batched": { + "aggregate_tps": 0.955, + "recall": 0.125 + }, + "batched_speedup_vs_serialized": 0.06 +} \ No newline at end of file diff --git a/scripts/research/mlx_batched_layer_diff_diag.py b/scripts/research/mlx_batched_layer_diff_diag.py new file mode 100644 index 00000000..f1a14564 --- /dev/null +++ b/scripts/research/mlx_batched_layer_diff_diag.py @@ -0,0 +1,152 @@ +"""Localize the mlx_lm gemma-4 batch>1 decode bug: per-layer hidden-state diff, +batched-row-i vs serialized-i, at decode step 1. + +Prefill is known correct (first decoded token matches per row); the divergence +is in the batched decode forward. This dumps, per decoder layer, the max-abs +difference between the batched forward's row-i output and the serialized +single-row forward's output for the SAME session + SAME fed token. The first +layer whose diff jumps locates the bug (layer 0 → RoPE/embed/per-layer-input; +a sliding layer → sliding mask/cache; a full-attn layer → full path; the first +KV-shared layer → shared-KV plumbing). +""" + +from __future__ import annotations + +import argparse +import sys +from collections import Counter +from typing import List + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--verifier-path", required=True) + ap.add_argument("--rows", type=int, default=2) + ap.add_argument("--haystack-lines", type=int, default=15) + ap.add_argument("--kakeya-cache", action="store_true", + help="use concat-based SinkWindowKVCache (large window, no " + "eviction) instead of mlx_lm's in-place cache, to test " + "whether the in-place write is the batch>1 bug.") + 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 + from inference_engine.backends.mlx.cross_model_dlm_verifier import ( + resolve_mlx_text_model, + ) + from inference_engine.backends.mlx.cache import SinkWindowKVCache + + model, tok = mlx_lm.load(args.verifier_path) + text_model = resolve_mlx_text_model(model) + layers = text_model.layers + + def new_cache(): + if not args.kakeya_cache: + return model.make_cache() + return [SinkWindowKVCache(sink_size=4, window_size=200000) + for _ in range(len(layers))] + n_layers = len(layers) + layer_types = [getattr(l, "layer_type", "?") for l in layers] + first_shared = n_layers - getattr(model.args, "num_kv_shared_layers", 0) + print(f"[diff] layers={n_layers} first_kv_shared_idx={first_shared}", flush=True) + + R = args.rows + pool = make_niah_dataset(n_samples=R * 3, haystack_min_lines=args.haystack_lines, + haystack_max_lines=args.haystack_lines, seed=0) + + def encode(text): + 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." + ids = list(tok.apply_chat_template([{"role": "user", "content": text}], + add_generation_prompt=True)) + try: + m = tok.encode("<|channel>content\n", add_special_tokens=False) + except TypeError: + m = tok.encode("<|channel>content\n") + ids.extend(list(m if not hasattr(m, "tolist") else m.tolist())) + return ids + + enc = [encode(s.prompt_text) for s in pool] + modal = Counter(len(e) for e in enc).most_common(1)[0][0] + prompts = [e for e in enc if len(e) == modal][:R] + while len(prompts) < R: + prompts += prompts[: R - len(prompts)] + print(f"[diff] {R} rows, prompt len={modal}", flush=True) + + # capture per-layer output hidden by monkey-patching DecoderLayer.__call__ + DecoderLayer = type(layers[0]) + orig_call = DecoderLayer.__call__ + captured: List = [] + + captured_in: List = [] + + def patched(self, *a, **k): + captured_in.append(a[0]) # layer INPUT h: [B, L, D] + out = orig_call(self, *a, **k) + captured.append(out[0]) # layer OUTPUT h: [B, L, D] + return out + + def prefill(ids_2d): + cache = new_cache() + out = model(mx.array(ids_2d), cache=cache) + mx.eval(out) + return cache, out[:, -1, :] + + def decode_capture(token_ids_2d, cache): + captured.clear() + captured_in.clear() + DecoderLayer.__call__ = patched + try: + out = model(mx.array(token_ids_2d), cache=cache) + mx.eval(out) + finally: + DecoderLayer.__call__ = orig_call + return list(captured), list(captured_in), out + + # serialized: prefill + first token + capture decode step + ser_tok0, ser_out, ser_in = [], [], [] + for i in range(R): + c, lg = prefill([prompts[i]]) + t0 = int(mx.argmax(lg, axis=-1).item()) + ser_tok0.append(t0) + co, ci, _ = decode_capture([[t0]], c) + ser_out.append(co) + ser_in.append(ci) + + # batched: prefill + first tokens + capture decode step (same tokens) + cb, lgb = prefill(prompts) + bat_tok0 = [int(mx.argmax(lgb[i], axis=-1).item()) for i in range(R)] + caps_b, caps_in_b, _ = decode_capture([[t] for t in bat_tok0], cb) + + print(f"[diff] tok0 serial={ser_tok0} batched={bat_tok0} " + f"match={ser_tok0 == bat_tok0}", flush=True) + # layer-0 INPUT diff isolates embedding/per-layer-input (pre-attention) + in0 = [round(float(mx.max(mx.abs(caps_in_b[0][i:i + 1] - ser_in[i][0])).item()), 4) + for i in range(R)] + print(f"[diff] layer-0 INPUT (embedding) max|Δ| per row = {in0} " + f"(non-zero row>0 => embed/per-layer-input bug; zero => attention bug)", + flush=True) + print("[diff] layer | type | in max|Δ| | out max|Δ| (per row)", flush=True) + first_div = None + for li in range(n_layers): + din = [round(float(mx.max(mx.abs(caps_in_b[li][i:i + 1] - ser_in[i][li]).astype(mx.float32)).item()), 3) + for i in range(R)] + dout = [round(float(mx.max(mx.abs(caps_b[li][i:i + 1] - ser_out[i][li]).astype(mx.float32)).item()), 3) + for i in range(R)] + mark = "" + if first_div is None and max(dout) > 1e-2: + first_div = li + mark = " <-- FIRST OUT DIVERGENCE" + print(f"[diff] {li:2d} | {layer_types[li]:18s} | in={din} | out={dout}{mark}", + flush=True) + print(f"[diff] FIRST DIVERGENT LAYER = {first_div} " + f"(type={layer_types[first_div] if first_div is not None else None}, " + f"shared={first_div is not None and first_div >= first_shared})", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/research/mlx_batched_multitenant_bench.py b/scripts/research/mlx_batched_multitenant_bench.py index 9910e94e..fa957e0a 100644 --- a/scripts/research/mlx_batched_multitenant_bench.py +++ b/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 diff --git a/tests/inference_engine/bridge/test_manifest.py b/tests/inference_engine/bridge/test_manifest.py index abe7a610..f286d1ee 100644 --- a/tests/inference_engine/bridge/test_manifest.py +++ b/tests/inference_engine/bridge/test_manifest.py @@ -74,6 +74,9 @@ def test_allowlist_contains_exactly_the_documented_presets(): "mlx-backend-tests", "mlx-batched-diag-short", "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",