From 6b36fa48f88169acb8569ec154c85404700e32fc Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 05:12:03 +0000 Subject: [PATCH 1/9] feat(mac): gemma-4 interactive CLI chat on the Kakeya-for-Mac (MLX) engine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit scripts/chat_mlx_kakeya.py — single-stream (B=1) gemma-4 chat over the Kakeya S5 bounded cache (SinkWindowKVCache): sliding layers bounded to sink+window, full- attn (exact) layers keep full context (S5 free lunch, no f_theta/proposer needed on gemma-4). B=1 sidesteps the MLX B>1,L=1 batched-decode bug. Interactive REPL + --scripted mode (for non-interactive bridge verification) + transcript JSON. Adds mac-bridge preset mlx-kakeya-chat-smoke (scripted 3-turn run) + manifest test (100% coverage, 29 presets). Operator runs the same script without --scripted for a real interactive chat on the Mac. Co-authored-by: FluffyAIcode --- inference_engine/bridge/manifest.py | 28 +++ scripts/chat_mlx_kakeya.py | 219 ++++++++++++++++++ .../inference_engine/bridge/test_manifest.py | 13 ++ 3 files changed, 260 insertions(+) create mode 100644 scripts/chat_mlx_kakeya.py diff --git a/inference_engine/bridge/manifest.py b/inference_engine/bridge/manifest.py index a93ef592..32e49150 100644 --- a/inference_engine/bridge/manifest.py +++ b/inference_engine/bridge/manifest.py @@ -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." + "||Now multiply 6 by 7 and give only the number." + "||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..5cbc2c2a --- /dev/null +++ b/scripts/chat_mlx_kakeya.py @@ -0,0 +1,219 @@ +#!/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: + """gemma stops a turn on ; also honor the tokenizer EOS.""" + eos = set() + if getattr(tok, "eos_token_id", None) is not None: + eos.add(int(tok.eos_token_id)) + for marker in ("", ""): + try: + ids = tok.encode(marker, add_special_tokens=False) + ids = ids.tolist() if hasattr(ids, "tolist") else list(ids) + if len(ids) == 1: + eos.add(int(ids[0])) + except Exception: + pass + return eos + + +def _content_marker(tok) -> List[int]: + """gemma-4 emits a <|channel>thought ... reasoning preamble by default; this + marker nudges it straight to the content channel for direct chat answers.""" + try: + ids = tok.encode("<|channel>content\n", add_special_tokens=False) + except TypeError: + ids = tok.encode("<|channel>content\n") + 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=256) + 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) + marker = [] if args.thinking else _content_marker(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).") + + 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]: + ids = list(tok.apply_chat_template(history, add_generation_prompt=True)) + return ids + list(marker) + + 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() + for tok_id, _ in generate_step( + mx.array(prompt_ids), model, prompt_cache=cache, + max_tokens=args.max_new_tokens, + ): + t = int(tok_id) + if t in 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 + dt = max(time.time() - t0, 1e-9) + return { + "text": tok.decode(toks, skip_special_tokens=True), + "n_tokens": len(toks), + "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'][:200]!r} " + f"({info['n_tokens']} tok, {info['decode_tps']} tok/s, " + f"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, + "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"]}) + _log(f"{info['n_tokens']} tok, {info['decode_tps']} tok/s, " + f"resident bounded-KV {info['resident_kv_bytes']/1e6:.1f} MB " + f"(sliding layers capped at sink+window={args.sink}+{args.window})") + 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"})) From bfcac004055a5710e75ba39d970ed734c28c70fd Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 05:18:17 +0000 Subject: [PATCH 2/9] fix(mac chat): use template enable_thinking instead of raw channel marker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Smoke run showed the injected <|channel>content marker leaked 'thought' text and caused greedy looping (turns 1/3 repeated). Replace with the chat template's enable_thinking flag (default off → clean direct answers), with a TypeError fallback for templates that don't accept the kwarg. Co-authored-by: FluffyAIcode --- scripts/chat_mlx_kakeya.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/scripts/chat_mlx_kakeya.py b/scripts/chat_mlx_kakeya.py index 5cbc2c2a..9582439f 100644 --- a/scripts/chat_mlx_kakeya.py +++ b/scripts/chat_mlx_kakeya.py @@ -52,13 +52,16 @@ def _resolve_eos(tok) -> set: return eos -def _content_marker(tok) -> List[int]: - """gemma-4 emits a <|channel>thought ... reasoning preamble by default; this - marker nudges it straight to the content channel for direct chat answers.""" +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.encode("<|channel>content\n", add_special_tokens=False) + ids = tok.apply_chat_template( + history, add_generation_prompt=True, enable_thinking=thinking) except TypeError: - ids = tok.encode("<|channel>content\n") + ids = tok.apply_chat_template(history, add_generation_prompt=True) return ids.tolist() if hasattr(ids, "tolist") else list(ids) @@ -98,7 +101,6 @@ def main() -> int: n_layers = len(text_model.layers) full_idx = set(mlx_full_attention_layer_indices(text_model)) eos = _resolve_eos(tok) - marker = [] if args.thinking else _content_marker(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)}") @@ -117,8 +119,7 @@ def new_cache() -> list: ] def build_prompt_ids(history: List[Dict[str, str]]) -> List[int]: - ids = list(tok.apply_chat_template(history, add_generation_prompt=True)) - return ids + list(marker) + 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.""" From 39c85d04de35718b07be8663d50ab90acfd17e7e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 05:24:41 +0000 Subject: [PATCH 3/9] fix(mac chat): add repetition_penalty (default 1.3) to stop greedy JSON loops enable_thinking=False gave correct answers but greedy decoding looped on a repeated JSON wrapper. Add mlx_lm repetition penalty via logits_processors (with a TypeError fallback for older mlx_lm signatures). Co-authored-by: FluffyAIcode --- scripts/chat_mlx_kakeya.py | 32 ++++++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/scripts/chat_mlx_kakeya.py b/scripts/chat_mlx_kakeya.py index 9582439f..f023451e 100644 --- a/scripts/chat_mlx_kakeya.py +++ b/scripts/chat_mlx_kakeya.py @@ -76,6 +76,8 @@ def main() -> int: 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=256) + 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.") @@ -107,6 +109,16 @@ def main() -> int: _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. @@ -127,10 +139,22 @@ def generate_turn(prompt_ids: List[int], on_delta=None) -> Dict[str, Any]: toks: List[int] = [] shown = "" t0 = time.time() - for tok_id, _ in generate_step( - mx.array(prompt_ids), model, prompt_cache=cache, - max_tokens=args.max_new_tokens, - ): + 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 + + for tok_id, _ in _iter(): t = int(tok_id) if t in eos: break From c2168eca15ae10c09fb45e0d1d702e68e8709064 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 05:30:16 +0000 Subject: [PATCH 4/9] =?UTF-8?q?fix(mac=20chat):=20generic=20loop=20guard?= =?UTF-8?q?=20=E2=80=94=20stop=20+=20trim=20when=20output=20degenerates=20?= =?UTF-8?q?into=20a=20repeat?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 4-bit gemma-4-it checkpoint appends a repeating json answer-wrapper under greedy; repetition_penalty alone didn't stop it. Add a generic guard: break when a 24-char tail recurs 3+ times, and _trim_loop() keeps only the first occurrence for the stored/displayed text. The (correct) answer is always first. Co-authored-by: FluffyAIcode --- scripts/chat_mlx_kakeya.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/scripts/chat_mlx_kakeya.py b/scripts/chat_mlx_kakeya.py index f023451e..7fefc757 100644 --- a/scripts/chat_mlx_kakeya.py +++ b/scripts/chat_mlx_kakeya.py @@ -52,6 +52,18 @@ def _resolve_eos(tok) -> set: return eos +def _trim_loop(s: str, win: int = 24) -> str: + """If the text degenerated into a repeated block (greedy loop), keep only up + to the end of the first occurrence of the repeating unit.""" + if len(s) < win * 3: + return s + tail = s[-win:] + first = s.find(tail) + if 0 <= first < len(s) - win: + return s[: first + win].rstrip() + return s + + 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 @@ -164,9 +176,12 @@ def _iter(): if delta and on_delta is not None: on_delta(delta) shown = full + # greedy loop guard: a 24-char tail recurring 3+ times = degenerate + if len(full) > 96 and full.count(full[-24:]) >= 3: + break dt = max(time.time() - t0, 1e-9) return { - "text": tok.decode(toks, skip_special_tokens=True), + "text": _trim_loop(tok.decode(toks, skip_special_tokens=True)), "n_tokens": len(toks), "decode_tps": round(len(toks) / dt, 2), "resident_kv_bytes": int(total_kv_bytes(cache)), From f55427ae090aaf32b0c21c1d9aac5b597877e6bf Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 05:35:45 +0000 Subject: [PATCH 5/9] fix(mac chat): strict consecutive-loop guard (no longer trims real answers) _trim_loop over-cut (dropped 'Paris' because the answer text recurs inside the model's json echo). Replace with _is_degenerate_loop: stop ONLY on a true back-to-back repeat (same 16-char block 3x consecutively); never trim stored text. Answer + a single echo is preserved. Co-authored-by: FluffyAIcode --- scripts/chat_mlx_kakeya.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/scripts/chat_mlx_kakeya.py b/scripts/chat_mlx_kakeya.py index 7fefc757..dab18d08 100644 --- a/scripts/chat_mlx_kakeya.py +++ b/scripts/chat_mlx_kakeya.py @@ -52,16 +52,14 @@ def _resolve_eos(tok) -> set: return eos -def _trim_loop(s: str, win: int = 24) -> str: - """If the text degenerated into a repeated block (greedy loop), keep only up - to the end of the first occurrence of the repeating unit.""" - if len(s) < win * 3: - return s - tail = s[-win:] - first = s.find(tail) - if 0 <= first < len(s) - win: - return s[: first + win].rstrip() - return s +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]: @@ -176,12 +174,11 @@ def _iter(): if delta and on_delta is not None: on_delta(delta) shown = full - # greedy loop guard: a 24-char tail recurring 3+ times = degenerate - if len(full) > 96 and full.count(full[-24:]) >= 3: + if _is_degenerate_loop(full): # true back-to-back repeat → stop break dt = max(time.time() - t0, 1e-9) return { - "text": _trim_loop(tok.decode(toks, skip_special_tokens=True)), + "text": tok.decode(toks, skip_special_tokens=True), "n_tokens": len(toks), "decode_tps": round(len(toks) / dt, 2), "resident_kv_bytes": int(total_kv_bytes(cache)), From bd2eaadcb3e52b3d195bfc5f214ee575bcb66b8c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 06:25:01 +0000 Subject: [PATCH 6/9] fix(mac chat): raise default --max-new-tokens 256->1024 + show stop reason (eos/max/loop) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit User hit truncation ('断掉') on a long PoW explanation: exactly 256 tokens = the old default cap. Raise default to 1024 and surface the stop reason each turn (WARN when it's the max-new-tokens cap), so truncation is obvious and tunable. Co-authored-by: FluffyAIcode --- scripts/chat_mlx_kakeya.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/scripts/chat_mlx_kakeya.py b/scripts/chat_mlx_kakeya.py index dab18d08..174fd6a5 100644 --- a/scripts/chat_mlx_kakeya.py +++ b/scripts/chat_mlx_kakeya.py @@ -85,7 +85,9 @@ def main() -> int: 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=256) + ap.add_argument("--max-new-tokens", type=int, default=1024, + help="Generation cap. Long explanations can need 1500+; raise " + "this if answers get cut off ('断掉').") 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", @@ -164,9 +166,11 @@ 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) @@ -175,11 +179,13 @@ def _iter(): 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)), @@ -246,9 +252,11 @@ def _iter(): 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 layers capped at sink+window={args.sink}+{args.window})") + f"(sliding capped at sink+window={args.sink}+{args.window}){warn}") return 0 From 3616e33946edefb91f56c9061da4038d9eed520d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 06:32:58 +0000 Subject: [PATCH 7/9] =?UTF-8?q?fix(mac=20chat):=20stop=20at=20gemma's=20na?= =?UTF-8?q?tural=20turn=20end=20(real=20EOS=20detection)=20=E2=80=94=20out?= =?UTF-8?q?put=20length=20must=20not=20be=20a=20user=20knob?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit User correctly noted: needing --max-new-tokens to avoid truncation = unusable. Root cause: was dropped from the stop set (old code used encode()+ single-token check on a special token), so the model ran PAST its turn (verbose JSON echoes / self-correction) until the max-tokens cap — that IS the '断掉' and the verbosity, same bug. Now resolve stops via tokenizer.eos_token_ids + convert_tokens_to_ids(''/''); the cap is only a backstop. Report eos_token_ids + per-turn stop_reason; smoke now includes a long-answer prompt to confirm it stops on eos, not max. Co-authored-by: FluffyAIcode --- inference_engine/bridge/manifest.py | 2 +- scripts/chat_mlx_kakeya.py | 32 ++++++++++++++++++++--------- 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/inference_engine/bridge/manifest.py b/inference_engine/bridge/manifest.py index 32e49150..79d112a9 100644 --- a/inference_engine/bridge/manifest.py +++ b/inference_engine/bridge/manifest.py @@ -668,7 +668,7 @@ def _harness_preset( "--max-new-tokens", "{max_new_tokens}", "--scripted", "What is the capital of France? Answer in one short sentence." - "||Now multiply 6 by 7 and give only the number." + "||Explain how proof-of-work works, step by step." "||Name three primary colors.", "--output", "results/research/k3_mac_bridge_mlx_kakeya_chat.json", diff --git a/scripts/chat_mlx_kakeya.py b/scripts/chat_mlx_kakeya.py index 174fd6a5..9332e3eb 100644 --- a/scripts/chat_mlx_kakeya.py +++ b/scripts/chat_mlx_kakeya.py @@ -37,18 +37,28 @@ def _log(msg: str) -> None: def _resolve_eos(tok) -> set: - """gemma stops a turn on ; also honor the tokenizer EOS.""" - eos = 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: - ids = tok.encode(marker, add_special_tokens=False) - ids = ids.tolist() if hasattr(ids, "tolist") else list(ids) - if len(ids) == 1: - eos.add(int(ids[0])) + tid = tok.convert_tokens_to_ids(marker) except Exception: - pass + tid = None + if isinstance(tid, int) and tid >= 0 and tid != unk: + eos.add(int(tid)) return eos @@ -206,15 +216,17 @@ def _iter(): history.append({"role": "assistant", "content": info["text"]}) transcript.append({"user": user, **info}) _log(f"USER: {user!r}") - _log(f"GEMMA-4: {info['text'][:200]!r} " - f"({info['n_tokens']} tok, {info['decode_tps']} tok/s, " - f"resident_kv={info['resident_kv_bytes']/1e6:.1f}MB)") + _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: From 9064d53cd63c05a1677685ed56a1678a3b49046c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 06:38:30 +0000 Subject: [PATCH 8/9] fix(mac chat): default --max-new-tokens=2048 as a generous backstop (natural EOS stops first) Co-authored-by: FluffyAIcode --- scripts/chat_mlx_kakeya.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/scripts/chat_mlx_kakeya.py b/scripts/chat_mlx_kakeya.py index 9332e3eb..e7adb1a8 100644 --- a/scripts/chat_mlx_kakeya.py +++ b/scripts/chat_mlx_kakeya.py @@ -95,9 +95,9 @@ def main() -> int: 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=1024, - help="Generation cap. Long explanations can need 1500+; raise " - "this if answers get cut off ('断掉').") + 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", From f1e648b441c35ccabf64ff94fb797c287c88f4af Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 06:40:35 +0000 Subject: [PATCH 9/9] chore(mac-bridge): raise MAX_NEW_TOKENS bound 512->2048 (chat answers need it; natural EOS stops first) Co-authored-by: FluffyAIcode --- inference_engine/bridge/manifest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/inference_engine/bridge/manifest.py b/inference_engine/bridge/manifest.py index 79d112a9..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_]*)\}$")