diff --git a/inference_engine/bridge/manifest.py b/inference_engine/bridge/manifest.py index a93ef592..60df4479 100644 --- a/inference_engine/bridge/manifest.py +++ b/inference_engine/bridge/manifest.py @@ -29,7 +29,7 @@ # bridge is for evidence runs and debugging, not for monopolizing the # single Mac with open-ended workloads. MAX_N_SAMPLES = 50 -MAX_NEW_TOKENS = 512 +MAX_NEW_TOKENS = 2048 # backstop for chat; natural EOS stops well before this MAX_BLOCK_SIZE = 16 _ENV_PLACEHOLDER = re.compile(r"^\$\{ENV:([A-Z][A-Z0-9_]*)\}$") @@ -650,6 +650,34 @@ def _harness_preset( }, validate_reports=False, ), + Preset( + name="mlx-kakeya-chat-smoke", + description="Run gemma-4 on the Kakeya-for-Mac (MLX) engine via the " + "interactive chat CLI in NON-interactive --scripted mode: " + "single-stream generation over the Kakeya S5 bounded " + "sink+window cache (sliding layers bounded; full-attn " + "layers full). Writes a transcript JSON so we can verify " + "gemma-4 responds coherently on the engine; the operator " + "runs the same script without --scripted for a real " + "interactive REPL on the Mac.", + command_templates=( + ( + "python3", "scripts/chat_mlx_kakeya.py", + "--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}", + "--sink", "4", "--window", "64", + "--max-new-tokens", "{max_new_tokens}", + "--scripted", + "What is the capital of France? Answer in one short sentence." + "||Explain how proof-of-work works, step by step." + "||Name three primary colors.", + "--output", + "results/research/k3_mac_bridge_mlx_kakeya_chat.json", + ), + ), + timeout_minutes=45, + params={"max_new_tokens": ("int:max_new_tokens", "64")}, + validate_reports=False, + ), ) } diff --git a/scripts/chat_mlx_kakeya.py b/scripts/chat_mlx_kakeya.py new file mode 100644 index 00000000..e7adb1a8 --- /dev/null +++ b/scripts/chat_mlx_kakeya.py @@ -0,0 +1,276 @@ +#!/usr/bin/env python3 +"""Interactive CLI chat with gemma-4 on the Kakeya-for-Mac engine (MLX). + +Runs the gemma-4 MLX verifier with **Kakeya Attention's bounded sink+window KV +cache (S5)**: the model's sliding-attention layers keep only ``sink + window`` +tokens resident, while gemma-4's native full-attention layers keep full context +(the "S5 free lunch" — recall is carried by the full layers, so no f_θ/proposer +restoration is needed on gemma-4). This is single-stream (B=1) generation, which +sidesteps the MLX ``B>1, L=1`` batched-decode kernel bug entirely. + +Usage (on the Mac, in the repo checkout): + + # interactive REPL — type a message, get gemma-4's reply, blank line/Ctrl-D quits + PYTHONPATH=. python3 scripts/chat_mlx_kakeya.py \ + --verifier-path /Users/fluffy314/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit + + # non-interactive smoke (used by the Mac-bridge preset): fixed turns -> JSON transcript + PYTHONPATH=. python3 scripts/chat_mlx_kakeya.py --verifier-path \ + --scripted "What is the capital of France?||Now multiply 6 by 7." \ + --output results/research/mac_gemma4_kakeya_chat.json + +mlx_lm / mlx are imported lazily inside ``main`` so ``--help`` works off-Mac. +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + + +def _log(msg: str) -> None: + print(f"[kakeya-chat] {msg}", file=sys.stderr, flush=True) + + +def _resolve_eos(tok) -> set: + """All token ids that end gemma's turn — so generation stops NATURALLY (the + user must never tune output length). gemma ends a turn on ```` + (and ````); the previous code used ``encode()`` + a single-token check + which silently DROPPED ```` (it's a special token), so the model + ran past its turn (verbose echoes) until the max-tokens cap — the real bug.""" + eos: set = set() + # mlx_lm's TokenizerWrapper exposes the authoritative stop set when present. + ids = getattr(tok, "eos_token_ids", None) + if ids: + eos |= {int(i) for i in ids} + if getattr(tok, "eos_token_id", None) is not None: + eos.add(int(tok.eos_token_id)) + # convert_tokens_to_ids is the reliable lookup for a KNOWN special token + # (unlike encode(), which may split it or add specials). + unk = getattr(tok, "unk_token_id", None) + for marker in ("", ""): + try: + tid = tok.convert_tokens_to_ids(marker) + except Exception: + tid = None + if isinstance(tid, int) and tid >= 0 and tid != unk: + eos.add(int(tid)) + return eos + + +def _is_degenerate_loop(s: str, unit: int = 16) -> bool: + """True only on a TRUE consecutive loop: the same ``unit``-char block + repeated 3x back-to-back at the tail. (Deliberately strict so an answer that + merely echoes itself once — e.g. text + a json wrapper — is NOT cut.)""" + if len(s) < unit * 3: + return False + a, b, c = s[-unit:], s[-2 * unit:-unit], s[-3 * unit:-2 * unit] + return a == b == c and a.strip() != "" + + +def _apply_template(tok, history, *, thinking: bool) -> List[int]: + """Encode the chat history. gemma-4 has a reasoning ("thought") channel; the + clean way to get direct answers is the template's ``enable_thinking`` flag + (NOT injecting a raw channel marker, which leaks 'thought' text and loops). + Falls back gracefully if the template doesn't accept the kwarg.""" + try: + ids = tok.apply_chat_template( + history, add_generation_prompt=True, enable_thinking=thinking) + except TypeError: + ids = tok.apply_chat_template(history, add_generation_prompt=True) + return ids.tolist() if hasattr(ids, "tolist") else list(ids) + + +def main() -> int: + ap = argparse.ArgumentParser(description="gemma-4 chat on the Kakeya-for-Mac (MLX) engine") + ap.add_argument("--verifier-path", required=True, + help="Local MLX gemma-4 model dir.") + ap.add_argument("--sink", type=int, default=4, help="Kakeya sink tokens.") + ap.add_argument("--window", type=int, default=64, + help="Kakeya sliding-window tokens (S5; sliding layers).") + ap.add_argument("--full-window", type=int, default=8192, + help="Resident window for the full-attention (exact) layers " + "— large = effectively full context (S5 recall carrier).") + ap.add_argument("--max-new-tokens", type=int, default=2048, + help="Backstop cap ONLY — generation stops naturally at gemma's " + ", so you should not need to tune this.") + ap.add_argument("--repetition-penalty", type=float, default=1.3, + help="Penalize repeated tokens to stop greedy loops (1.0 = off).") + ap.add_argument("--thinking", action="store_true", + help="Allow gemma-4's reasoning channel (default: direct answers).") + ap.add_argument("--system", default=None, help="Optional system prompt.") + ap.add_argument("--scripted", default=None, + help="Non-interactive: '||'-separated user turns; writes a transcript.") + ap.add_argument("--output", default=None, help="Transcript JSON (scripted mode).") + args = ap.parse_args() + + import mlx.core as mx # type: ignore + import mlx_lm # type: ignore + from mlx_lm.generate import generate_step # type: ignore + from inference_engine.backends.mlx.cache import ( + SinkWindowKVCache, total_kv_bytes, cache_seq_length, + ) + from inference_engine.backends.mlx.cross_model_dlm_verifier import ( + resolve_mlx_text_model, mlx_full_attention_layer_indices, + ) + + _log(f"loading MLX model: {args.verifier_path}") + t_load = time.time() + model, tok = mlx_lm.load(args.verifier_path) + text_model = resolve_mlx_text_model(model) + n_layers = len(text_model.layers) + full_idx = set(mlx_full_attention_layer_indices(text_model)) + eos = _resolve_eos(tok) + _log(f"loaded in {time.time()-t_load:.1f}s | layers={n_layers} " + f"exact(full-attn)={sorted(full_idx)} sink={args.sink} window={args.window} " + f"eos={sorted(eos)}") + _log("Kakeya Attention: sliding layers bounded to sink+window; " + "exact layers keep full context (S5).") + + logits_processors = None + if args.repetition_penalty and args.repetition_penalty != 1.0: + try: + from mlx_lm.sample_utils import make_logits_processors # type: ignore + logits_processors = make_logits_processors( + repetition_penalty=args.repetition_penalty) + _log(f"repetition_penalty={args.repetition_penalty} enabled") + except Exception as exc: # noqa: BLE001 + _log(f"repetition penalty unavailable ({exc}); greedy") + + def new_cache() -> list: + # S5 hybrid: exact (full-attn) layers get a large window (≈full context, + # the recall carrier); sliding layers get the tight Kakeya window. + return [ + SinkWindowKVCache( + sink_size=args.sink, + window_size=(args.full_window if li in full_idx else args.window), + ) + for li in range(n_layers) + ] + + def build_prompt_ids(history: List[Dict[str, str]]) -> List[int]: + return _apply_template(tok, history, thinking=args.thinking) + + def generate_turn(prompt_ids: List[int], on_delta=None) -> Dict[str, Any]: + """Single-stream greedy decode over a FRESH Kakeya bounded cache.""" + cache = new_cache() + toks: List[int] = [] + shown = "" + t0 = time.time() + gkw: Dict[str, Any] = dict(prompt_cache=cache, max_tokens=args.max_new_tokens) + if logits_processors is not None: + gkw["logits_processors"] = logits_processors + try: + stream = generate_step(mx.array(prompt_ids), model, **gkw) + first = next(stream) + except TypeError: # older mlx_lm without logits_processors kwarg + gkw.pop("logits_processors", None) + stream = generate_step(mx.array(prompt_ids), model, **gkw) + first = next(stream) + + def _iter(): + yield first + yield from stream + + stop_reason = "max" # generator exhausts at max_tokens unless we break + for tok_id, _ in _iter(): + t = int(tok_id) + if t in eos: + stop_reason = "eos" + break + toks.append(t) + full = tok.decode(toks, skip_special_tokens=True) + delta = full[len(shown):] + if delta and on_delta is not None: + on_delta(delta) + shown = full + if _is_degenerate_loop(full): # true back-to-back repeat → stop + stop_reason = "loop" + break + dt = max(time.time() - t0, 1e-9) + return { + "text": tok.decode(toks, skip_special_tokens=True), + "n_tokens": len(toks), + "stop_reason": stop_reason, + "decode_tps": round(len(toks) / dt, 2), + "resident_kv_bytes": int(total_kv_bytes(cache)), + "resident_kv_seq_len_first_layer": int(cache_seq_length(cache)), + "prompt_tokens": len(prompt_ids), + } + + history: List[Dict[str, str]] = [] + if args.system: + history.append({"role": "system", "content": args.system}) + + # ---- scripted (non-interactive) mode: for Mac-bridge verification ---- + if args.scripted is not None: + turns = [t for t in args.scripted.split("||") if t.strip()] + transcript: List[Dict[str, Any]] = [] + for user in turns: + history.append({"role": "user", "content": user}) + info = generate_turn(build_prompt_ids(history)) + history.append({"role": "assistant", "content": info["text"]}) + transcript.append({"user": user, **info}) + _log(f"USER: {user!r}") + _log(f"GEMMA-4: {info['text'][:160]!r} " + f"({info['n_tokens']} tok, stop={info['stop_reason']}, " + f"{info['decode_tps']} tok/s, resident_kv={info['resident_kv_bytes']/1e6:.1f}MB)") + report = { + "kind": "mac_gemma4_kakeya_chat", "schema_version": 1, + "model_path": args.verifier_path, + "engine": "Kakeya-for-Mac (MLX, S5 bounded sink+window, single-stream)", + "sink": args.sink, "window": args.window, "full_window": args.full_window, + "exact_layers": sorted(full_idx), "n_layers": n_layers, + "eos_token_ids": sorted(eos), + "max_new_tokens_cap": args.max_new_tokens, + "turns": transcript, + } + if args.output: + outp = Path(args.output) + outp.parent.mkdir(parents=True, exist_ok=True) + outp.write_text(json.dumps(report, indent=2), encoding="utf-8") + _log(f"wrote transcript -> {outp}") + else: + print(json.dumps(report, indent=2)) + return 0 + + # ---- interactive REPL ---- + _log("ready. Type a message and press Enter. Blank line or Ctrl-D to quit.") + while True: + try: + if sys.stdin.isatty(): + sys.stderr.write("\nyou> ") + sys.stderr.flush() + line = sys.stdin.readline() + except KeyboardInterrupt: + _log("interrupted") + break + if not line: + break + user = line.strip() + if not user: + break + history.append({"role": "user", "content": user}) + sys.stderr.write("gemma-4> ") + sys.stderr.flush() + info = generate_turn( + build_prompt_ids(history), + on_delta=lambda d: (sys.stdout.write(d), sys.stdout.flush()), + ) + sys.stdout.write("\n") + sys.stdout.flush() + history.append({"role": "assistant", "content": info["text"]}) + warn = (" [WARN: hit --max-new-tokens; raise it for longer answers]" + if info["stop_reason"] == "max" else f" [stopped: {info['stop_reason']}]") + _log(f"{info['n_tokens']} tok, {info['decode_tps']} tok/s, " + f"resident bounded-KV {info['resident_kv_bytes']/1e6:.1f} MB " + f"(sliding capped at sink+window={args.sink}+{args.window}){warn}") + 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 a287d708..40745cd7 100644 --- a/tests/inference_engine/bridge/test_manifest.py +++ b/tests/inference_engine/bridge/test_manifest.py @@ -80,6 +80,7 @@ def test_allowlist_contains_exactly_the_documented_presets(): "mlx-batched-multitenant", "mlx-batched-pad-decode", "mlx-env-probe", + "mlx-kakeya-chat-smoke", "mlx-multitenant-pressure", "mlx-upgrade", "mlx-upstream-batch-probe", @@ -135,6 +136,18 @@ def test_pad_decode_preset_carries_flag_and_forces_trimmable_cache(): assert HARNESS_ENV["KAKEYA_MAC_VERIFIER_PATH"] in argv +def test_mlx_kakeya_chat_smoke_preset_resolves(): + request = parse_manifest(_manifest( + preset="mlx-kakeya-chat-smoke", params={"max_new_tokens": "64"})) + (argv,) = build_commands(request, HARNESS_ENV) + assert argv[1].endswith("chat_mlx_kakeya.py") + assert HARNESS_ENV["KAKEYA_MAC_VERIFIER_PATH"] in argv + assert "--scripted" in argv + assert argv[argv.index("--max-new-tokens") + 1] == "64" + assert not [t for t in argv if t.startswith("${ENV:")] + assert not [t for t in argv if t.startswith("{") and t.endswith("}")] + + def test_drafter_parity_preset_resolves(): request = parse_manifest(_manifest( preset="k3-drafter-parity", params={"block_size": "8"}))