From 6b36fa48f88169acb8569ec154c85404700e32fc Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 05:12:03 +0000 Subject: [PATCH 01/15] 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 02/15] 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 03/15] 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 04/15] =?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 05/15] 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 06/15] 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 07/15] =?UTF-8?q?fix(mac=20chat):=20stop=20at=20gemma's=20?= =?UTF-8?q?natural=20turn=20end=20(real=20EOS=20detection)=20=E2=80=94=20o?= =?UTF-8?q?utput=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 08/15] 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 09/15] 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_]*)\}$") From 728d694ffad77b23f5db967d4c4c2099a2c46084 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 07:36:31 +0000 Subject: [PATCH 10/15] =?UTF-8?q?feat(mac):=20FULL=20fused-engine=20chat?= =?UTF-8?q?=20(verifier+DFlash=20proposer+f=5F=CE=B8+S5)=20=E2=80=94=20not?= =?UTF-8?q?=20verifier-only?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds --chat/--chat-scripted to k3_integrated_niah_eval_mac.py: an interactive REPL that drives the EXACT validated fused spec-decode per-turn sequence (build_restoration -> S5 prefill -> aux capture -> fused_specdecode_generate_mlx_trim), reusing the harness's engine construction. The NIAH eval loop is untouched (zero risk to the evidence path). Each turn reports blocks / mean_accept_len to prove the proposer is live + bounded resident KV. Natural EOS stop (). Adds bridge preset mlx-kakeya-fused-chat-smoke (full fused flags + --chat-scripted) + manifest test. This is the verifier/PROPOSER bounded-memory engine the project is about — not the verifier-only AR path. Co-authored-by: FluffyAIcode --- inference_engine/bridge/manifest.py | 35 +++++ .../research/k3_integrated_niah_eval_mac.py | 143 ++++++++++++++++++ .../inference_engine/bridge/test_manifest.py | 17 +++ 3 files changed, 195 insertions(+) diff --git a/inference_engine/bridge/manifest.py b/inference_engine/bridge/manifest.py index 60df4479..4bfd3a37 100644 --- a/inference_engine/bridge/manifest.py +++ b/inference_engine/bridge/manifest.py @@ -678,6 +678,41 @@ def _harness_preset( params={"max_new_tokens": ("int:max_new_tokens", "64")}, validate_reports=False, ), + Preset( + name="mlx-kakeya-fused-chat-smoke", + description="Run gemma-4 on the FULL Kakeya fused engine (verifier + " + "DFlash proposer + f_θ + S5 bounded KV) via the harness " + "--chat --chat-scripted mode — NOT verifier-only. Verifies " + "the proposer is live (blocks>0, mean_accept_len>0) AND the " + "answer is correct AND KV is bounded, per turn. Writes a " + "transcript JSON.", + command_templates=( + ( + "python3", "scripts/research/k3_integrated_niah_eval_mac.py", + "--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}", + "--drafter-id", "${ENV:KAKEYA_MAC_DRAFTER_ID}", + "--f-theta-dir", "${ENV:KAKEYA_MAC_FTHETA_DIR}", + "--s5-exact-full-attn", "--fused-specdecode", + "--all-mlx-drafter", "--cuda-trim", + "--sink-size", "4", "--window-size", "64", + "--block-size", "{block_size}", + "--max-new-tokens", "{max_new_tokens}", + "--prefill-chunk-size", "512", + "--chat", + "--chat-scripted", + "What is the capital of France? Answer in one short sentence." + "||Name three primary colors.", + "--output", + "results/research/k3_mac_bridge_mlx_kakeya_fused_chat.json", + ), + ), + timeout_minutes=60, + params={ + "max_new_tokens": ("int:max_new_tokens", "64"), + "block_size": ("int:block_size", "4"), + }, + validate_reports=False, + ), ) } diff --git a/scripts/research/k3_integrated_niah_eval_mac.py b/scripts/research/k3_integrated_niah_eval_mac.py index 8025206e..17484a16 100644 --- a/scripts/research/k3_integrated_niah_eval_mac.py +++ b/scripts/research/k3_integrated_niah_eval_mac.py @@ -170,6 +170,16 @@ def parse_args() -> argparse.Namespace: "long-context one-shot forward OOM path. Set <=0 to " "use a single full prompt forward.") ap.add_argument("--output", default=None) + # ---- interactive chat on the FULL fused engine (verifier+proposer+f_θ+S5) ---- + ap.add_argument("--chat", action="store_true", + help="Interactive REPL on the FULL fused spec-decode engine " + "(verifier + DFlash proposer + f_θ + S5 bounded KV) — " + "NOT verifier-only. Requires the fused flags " + "(--fused-specdecode --force-fused-specdecode " + "--all-mlx-drafter --s5-exact-full-attn --cuda-trim).") + ap.add_argument("--chat-scripted", default=None, + help="Non-interactive chat: '||'-separated user turns " + "(for Mac-bridge verification); writes a transcript.") return ap.parse_args() @@ -690,6 +700,139 @@ def eval_fused_specdecode() -> Tuple[List[str], List[float], List[int]]: mlx_model, embed_scale=embed_scale, aux_layer_ids=aux_layer_ids, bridge_to_torch=bridge) + def _run_fused_chat() -> Tuple[List[str], List[float], List[int]]: + """Interactive/scripted chat on the FULL fused engine — reuses the + EXACT per-turn sequence of the eval loop below (build_restoration → + S5 prefill → aux capture → fused_specdecode_generate_mlx_trim), so the + verifier + DFlash proposer + f_θ + S5 bounded KV are all live. This is + NOT verifier-only: each turn drafts blocks and reports blocks / + mean_accept_len to prove the proposer ran.""" + if not (args.force_fused_specdecode and mlx_drafter is not None + and args.cuda_trim): + raise SystemExit( + "--chat needs the FULL fused engine flags: --fused-specdecode " + "--force-fused-specdecode --all-mlx-drafter --s5-exact-full-attn " + "--cuda-trim") + # Stop at gemma's natural turn end: + eos + # (convert_tokens_to_ids is the reliable special-token lookup). + chat_eos = set(end_ids) + unk = getattr(tokenizer, "unk_token_id", None) + native = getattr(tokenizer, "eos_token_ids", None) + if native: + chat_eos |= {int(x) for x in native} + for m in ("", ""): + try: + tid = tokenizer.convert_tokens_to_ids(m) + except Exception: + tid = None + if isinstance(tid, int) and tid >= 0 and tid != unk: + chat_eos.add(int(tid)) + + def _encode_chat(history: List[Dict[str, str]]) -> List[int]: + try: + cids = tokenizer.apply_chat_template( + history, add_generation_prompt=True, enable_thinking=False) + except TypeError: + cids = tokenizer.apply_chat_template( + history, add_generation_prompt=True) + return list(cids.tolist() if hasattr(cids, "tolist") else cids) + + def _gen_turn(pid: List[int]) -> Dict[str, Any]: + rk, rv, tsrc = build_restoration(pid, prefill_native_s5=True) + T = len(pid) + evicted = compute_evicted_positions( + T, args.sink_size, args.window_size) + aux_prompt = capture_aux_hidden( + mlx_model, pid, aux_layer_ids, embed_scale=embed_scale) + adapter.prefill( + pid, restored_k_per_layer=_pad(rk, tsrc, T), + restored_v_per_layer=_pad(rv, tsrc, T), + evicted_positions=evicted, + prefill_chunk_size=args.prefill_chunk_size, full_kv=True) + t0 = time.perf_counter() + res = fused_specdecode_generate_mlx_trim( + adapter, active_drafter, aux_prompt=aux_prompt, + embed_fn=embed_fn, lm_head_fn=lm_head_fn, + gen_tokens=args.max_new_tokens, block_size=args.block_size, + eos_ids=chat_eos, single_fused=args.single_fused) + res["decode_s"] = round(time.perf_counter() - t0, 3) + res["text"] = tokenizer.decode(res["tokens"]) + res["resident_kv_bytes"] = int( + sum(int(getattr(c, "nbytes", 0)) for c in (adapter._cache or []))) + return res + + print(f"[chat] FULL fused engine: verifier={args.verifier_path} " + f"drafter={args.drafter_id} f_theta={args.f_theta_dir} " + f"S5 sink={args.sink_size} window={args.window_size} " + f"block={args.block_size} | chat_eos={sorted(chat_eos)}", + file=sys.stderr, flush=True) + + history: List[Dict[str, str]] = [] + if args.chat_scripted is not None: + turns = [t for t in args.chat_scripted.split("||") if t.strip()] + transcript = [] + for u in turns: + history.append({"role": "user", "content": u}) + res = _gen_turn(_encode_chat(history)) + history.append({"role": "assistant", "content": res["text"]}) + tps = (res["decode_tokens"] / res["decode_s"] + if res["decode_s"] > 0 else 0.0) + transcript.append({ + "user": u, "text": res["text"], + "tokens": res["decode_tokens"], "blocks": res["blocks"], + "mean_accept_len": res["mean_accept_len"], + "decode_s": res["decode_s"], "decode_tps": round(tps, 2), + "resident_kv_bytes": res["resident_kv_bytes"]}) + print(f"[chat] USER {u!r}", file=sys.stderr, flush=True) + print(f"[chat] GEMMA-4 {res['text'][:200]!r} (blocks=" + f"{res['blocks']}, accept_len={res['mean_accept_len']}, " + f"{round(tps,2)} tok/s, kv={res['resident_kv_bytes']/1e6:.1f}MB)", + file=sys.stderr, flush=True) + report = { + "kind": "mac_gemma4_kakeya_fused_chat", "schema_version": 1, + "engine": ("Kakeya-for-Mac FULL fused spec-decode " + "(verifier + DFlash proposer + f_θ + S5 bounded KV)"), + "model_path": args.verifier_path, "drafter_id": args.drafter_id, + "f_theta_dir": args.f_theta_dir, "sink": args.sink_size, + "window": args.window_size, "block_size": args.block_size, + "exact_layers": full_attn_idx, "chat_eos": sorted(chat_eos), + "turns": transcript} + if args.output: + op = Path(args.output) + op.parent.mkdir(parents=True, exist_ok=True) + op.write_text(json.dumps(report, indent=2), encoding="utf-8") + print(f"[chat] wrote transcript -> {op}", file=sys.stderr) + else: + print(json.dumps(report, indent=2)) + raise SystemExit(0) + + print("[chat] ready. Type a message; blank line / Ctrl-D quits.", + file=sys.stderr, flush=True) + while True: + if sys.stdin.isatty(): + sys.stderr.write("\nyou> "); sys.stderr.flush() + line = sys.stdin.readline() + if not line: + break + u = line.strip() + if not u: + break + history.append({"role": "user", "content": u}) + res = _gen_turn(_encode_chat(history)) + history.append({"role": "assistant", "content": res["text"]}) + tps = (res["decode_tokens"] / res["decode_s"] + if res["decode_s"] > 0 else 0.0) + sys.stdout.write("gemma-4> " + res["text"] + "\n") + sys.stdout.flush() + print(f"[chat] blocks={res['blocks']} accept_len=" + f"{res['mean_accept_len']} {round(tps,2)} tok/s " + f"bounded-KV {res['resident_kv_bytes']/1e6:.1f}MB", + file=sys.stderr, flush=True) + raise SystemExit(0) + + if args.chat: + _run_fused_chat() + decoded, lats, toks = [], [], [] rows = [] for i, pid in enumerate(sample_ids): diff --git a/tests/inference_engine/bridge/test_manifest.py b/tests/inference_engine/bridge/test_manifest.py index 40745cd7..d493e28f 100644 --- a/tests/inference_engine/bridge/test_manifest.py +++ b/tests/inference_engine/bridge/test_manifest.py @@ -81,6 +81,7 @@ def test_allowlist_contains_exactly_the_documented_presets(): "mlx-batched-pad-decode", "mlx-env-probe", "mlx-kakeya-chat-smoke", + "mlx-kakeya-fused-chat-smoke", "mlx-multitenant-pressure", "mlx-upgrade", "mlx-upstream-batch-probe", @@ -148,6 +149,22 @@ def test_mlx_kakeya_chat_smoke_preset_resolves(): assert not [t for t in argv if t.startswith("{") and t.endswith("}")] +def test_mlx_kakeya_fused_chat_smoke_preset_resolves(): + request = parse_manifest(_manifest( + preset="mlx-kakeya-fused-chat-smoke", + params={"max_new_tokens": "64", "block_size": "4"})) + (argv,) = build_commands(request, HARNESS_ENV) + assert argv[1].endswith("k3_integrated_niah_eval_mac.py") + # full fused engine flags (verifier + proposer + f_θ + S5), not verifier-only + for flag in ("--fused-specdecode", "--all-mlx-drafter", "--s5-exact-full-attn", + "--cuda-trim", "--chat", "--chat-scripted"): + assert flag in argv, flag + assert HARNESS_ENV["KAKEYA_MAC_DRAFTER_ID"] in argv + assert HARNESS_ENV["KAKEYA_MAC_FTHETA_DIR"] in argv + 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 559daf8c6c60d924c4059114fe48881cec8becf5 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 07:44:04 +0000 Subject: [PATCH 11/15] polish(mac fused chat): strip stop-token markers from displayed text (skip_special_tokens) Co-authored-by: FluffyAIcode --- scripts/research/k3_integrated_niah_eval_mac.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/scripts/research/k3_integrated_niah_eval_mac.py b/scripts/research/k3_integrated_niah_eval_mac.py index 17484a16..9bc69a6c 100644 --- a/scripts/research/k3_integrated_niah_eval_mac.py +++ b/scripts/research/k3_integrated_niah_eval_mac.py @@ -756,7 +756,13 @@ def _gen_turn(pid: List[int]) -> Dict[str, Any]: gen_tokens=args.max_new_tokens, block_size=args.block_size, eos_ids=chat_eos, single_fused=args.single_fused) res["decode_s"] = round(time.perf_counter() - t0, 3) - res["text"] = tokenizer.decode(res["tokens"]) + try: + txt = tokenizer.decode(res["tokens"], skip_special_tokens=True) + except TypeError: + txt = tokenizer.decode(res["tokens"]) + for marker in ("", "", ""): + txt = txt.replace(marker, "") + res["text"] = txt.strip() res["resident_kv_bytes"] = int( sum(int(getattr(c, "nbytes", 0)) for c in (adapter._cache or []))) return res From a6dd2e2115162d5bcf48c24c2a5ccc01c1a0b1f8 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 07:51:18 +0000 Subject: [PATCH 12/15] =?UTF-8?q?docs(mac=20fused=20chat):=20accurate=20en?= =?UTF-8?q?gine=20label=20=E2=80=94=20proposer=20live;=20f=5F=CE=B8=20bypa?= =?UTF-8?q?ssed=20on=20gemma-4=20(S5=20free=20lunch)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: FluffyAIcode --- scripts/research/k3_integrated_niah_eval_mac.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/scripts/research/k3_integrated_niah_eval_mac.py b/scripts/research/k3_integrated_niah_eval_mac.py index 9bc69a6c..3643f455 100644 --- a/scripts/research/k3_integrated_niah_eval_mac.py +++ b/scripts/research/k3_integrated_niah_eval_mac.py @@ -704,9 +704,11 @@ def _run_fused_chat() -> Tuple[List[str], List[float], List[int]]: """Interactive/scripted chat on the FULL fused engine — reuses the EXACT per-turn sequence of the eval loop below (build_restoration → S5 prefill → aux capture → fused_specdecode_generate_mlx_trim), so the - verifier + DFlash proposer + f_θ + S5 bounded KV are all live. This is - NOT verifier-only: each turn drafts blocks and reports blocks / - mean_accept_len to prove the proposer ran.""" + gemma-4 verifier + DFlash proposer + S5 bounded KV are all live. NOT + verifier-only: each turn the proposer drafts blocks the verifier + accepts (reports blocks / mean_accept_len). On gemma-4 f_θ restoration + is bypassed via S5 native exact-layer prefill (the free lunch); + f_θ is load-bearing on full-attention models.""" if not (args.force_fused_specdecode and mlx_drafter is not None and args.cuda_trim): raise SystemExit( @@ -796,8 +798,11 @@ def _gen_turn(pid: List[int]) -> Dict[str, Any]: file=sys.stderr, flush=True) report = { "kind": "mac_gemma4_kakeya_fused_chat", "schema_version": 1, - "engine": ("Kakeya-for-Mac FULL fused spec-decode " - "(verifier + DFlash proposer + f_θ + S5 bounded KV)"), + "engine": ("Kakeya-for-Mac fused spec-decode (gemma-4 verifier " + "+ DFlash proposer + S5 bounded KV; f_θ restoration " + "bypassed on gemma-4 via S5 native exact-layer " + "prefill — the free lunch — and load-bearing on " + "full-attention models)"), "model_path": args.verifier_path, "drafter_id": args.drafter_id, "f_theta_dir": args.f_theta_dir, "sink": args.sink_size, "window": args.window_size, "block_size": args.block_size, From b73198bfba43d84840ad3639e085a4d63ce03513 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 08:05:12 +0000 Subject: [PATCH 13/15] =?UTF-8?q?feat(mac=20fused=20chat):=20--force-f-the?= =?UTF-8?q?ta=20=E2=80=94=20f=5F=CE=B8=20actually=20RUNS=20each=20turn=20(?= =?UTF-8?q?even=20if=20recall-irrelevant=20on=20gemma-4)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit User: f_θ must execute even if its output is discarded by the verifier attention. - build_restoration: --force-f-theta bypasses the S5 native-prefill short-circuit (the line that stopped f_θ running under --s5-exact-full-attn), so f_θ projects proposer hidden -> verifier K/V and injects it into the sliding layers. - _gen_turn now branches over all-mlx (proposer; f_θ bypassed) AND torch drafter+f_θ (proposer + f_θ runs) paths; reports f_theta_ran + f_theta_layers. - new preset mlx-kakeya-fused-chat-ftheta (torch path + --force-f-theta). On gemma-4 the restored sliding K/V are recall-irrelevant (exact layers carry recall) — f_θ still EXECUTES, exercising the full verifier/proposer/f_θ pipeline. Co-authored-by: FluffyAIcode --- inference_engine/bridge/manifest.py | 36 ++++++++++ .../research/k3_integrated_niah_eval_mac.py | 65 +++++++++++++++---- .../inference_engine/bridge/test_manifest.py | 15 +++++ 3 files changed, 103 insertions(+), 13 deletions(-) diff --git a/inference_engine/bridge/manifest.py b/inference_engine/bridge/manifest.py index 4bfd3a37..2cccf2ad 100644 --- a/inference_engine/bridge/manifest.py +++ b/inference_engine/bridge/manifest.py @@ -713,6 +713,42 @@ def _harness_preset( }, validate_reports=False, ), + Preset( + name="mlx-kakeya-fused-chat-ftheta", + description="Like mlx-kakeya-fused-chat-smoke but on the TORCH drafter " + "+ f_θ path with --force-f-theta: f_θ restoration ACTUALLY " + "RUNS each turn (projects proposer hidden → verifier K/V, " + "injected into the sliding layers) even though on gemma-4 " + "those K/V are recall-irrelevant (the exact layers carry " + "recall). Verifies the FULL verifier/proposer/f_θ pipeline: " + "report shows f_theta_ran=true + blocks>0. (No " + "--all-mlx-drafter; torch bridge path is slower.)", + command_templates=( + ( + "python3", "scripts/research/k3_integrated_niah_eval_mac.py", + "--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}", + "--drafter-id", "${ENV:KAKEYA_MAC_DRAFTER_ID}", + "--f-theta-dir", "${ENV:KAKEYA_MAC_FTHETA_DIR}", + "--s5-exact-full-attn", "--fused-specdecode", "--force-f-theta", + "--sink-size", "4", "--window-size", "64", + "--block-size", "{block_size}", + "--max-new-tokens", "{max_new_tokens}", + "--prefill-chunk-size", "512", + "--chat", + "--chat-scripted", + "What is the capital of France? Answer in one short sentence." + "||Name three primary colors.", + "--output", + "results/research/k3_mac_bridge_mlx_kakeya_fused_chat_ftheta.json", + ), + ), + timeout_minutes=90, + params={ + "max_new_tokens": ("int:max_new_tokens", "32"), + "block_size": ("int:block_size", "4"), + }, + validate_reports=False, + ), ) } diff --git a/scripts/research/k3_integrated_niah_eval_mac.py b/scripts/research/k3_integrated_niah_eval_mac.py index 3643f455..75a06c8c 100644 --- a/scripts/research/k3_integrated_niah_eval_mac.py +++ b/scripts/research/k3_integrated_niah_eval_mac.py @@ -180,6 +180,14 @@ def parse_args() -> argparse.Namespace: ap.add_argument("--chat-scripted", default=None, help="Non-interactive chat: '||'-separated user turns " "(for Mac-bridge verification); writes a transcript.") + ap.add_argument("--force-f-theta", action="store_true", + help="Run f_θ restoration even under --s5-exact-full-attn " + "(bypass the S5 native-prefill short-circuit). On gemma-4 " + "the restored sliding-layer K/V are recall-irrelevant " + "(the exact layers carry recall), but f_θ EXECUTES and " + "its output is injected — exercising the full verifier/" + "proposer/f_θ pipeline. Requires the torch drafter+f_θ " + "(do NOT combine with --all-mlx-drafter).") return ap.parse_args() @@ -367,7 +375,8 @@ def build_restoration(prompt_ids: List[int], *, prefill_native_s5: bool = False) restored bank for those layers lets mlx_lm store their own post-RoPE cache directly and avoids the extra clean verifier forward. """ - if prefill_native_s5 and args.s5_exact_full_attn and not args.identity_restore: + if (prefill_native_s5 and args.s5_exact_full_attn + and not args.identity_restore and not args.force_f_theta): return {}, {}, len(prompt_ids) if drafter is None or f_theta is None or fcfg is None: raise RuntimeError("drafter/f_theta are required for this restoration mode") @@ -709,12 +718,16 @@ def _run_fused_chat() -> Tuple[List[str], List[float], List[int]]: accepts (reports blocks / mean_accept_len). On gemma-4 f_θ restoration is bypassed via S5 native exact-layer prefill (the free lunch); f_θ is load-bearing on full-attention models.""" - if not (args.force_fused_specdecode and mlx_drafter is not None - and args.cuda_trim): + _allmlx_ok = mlx_drafter is not None and args.cuda_trim + _torch_ftheta_ok = drafter is not None and f_theta is not None + if not (args.force_fused_specdecode and (_allmlx_ok or _torch_ftheta_ok)): raise SystemExit( - "--chat needs the FULL fused engine flags: --fused-specdecode " - "--force-fused-specdecode --all-mlx-drafter --s5-exact-full-attn " - "--cuda-trim") + "--chat needs the FULL fused engine. Either:\n" + " (a) --fused-specdecode --all-mlx-drafter --s5-exact-full-attn " + "--cuda-trim (verifier + proposer; f_θ bypassed on gemma-4 " + "via S5), or\n" + " (b) --fused-specdecode --force-f-theta (torch DFlash drafter " + "+ f_θ that ACTUALLY RUNS; do NOT pass --all-mlx-drafter).") # Stop at gemma's natural turn end: + eos # (convert_tokens_to_ids is the reliable special-token lookup). chat_eos = set(end_ids) @@ -741,23 +754,45 @@ def _encode_chat(history: List[Dict[str, str]]) -> List[int]: def _gen_turn(pid: List[int]) -> Dict[str, Any]: rk, rv, tsrc = build_restoration(pid, prefill_native_s5=True) + # f_θ ran iff build_restoration produced restored banks via the + # torch drafter+f_θ (under --force-f-theta the S5 short-circuit is + # bypassed → rk holds f_θ-projected sliding-layer K/V). + f_theta_ran = bool(rk) and (drafter is not None and f_theta is not None) T = len(pid) evicted = compute_evicted_positions( T, args.sink_size, args.window_size) - aux_prompt = capture_aux_hidden( + aux_prompt_mx = capture_aux_hidden( mlx_model, pid, aux_layer_ids, embed_scale=embed_scale) + aux_prompt = (aux_prompt_mx if bridge is None + else [bridge(a) for a in aux_prompt_mx]) adapter.prefill( pid, restored_k_per_layer=_pad(rk, tsrc, T), restored_v_per_layer=_pad(rv, tsrc, T), evicted_positions=evicted, - prefill_chunk_size=args.prefill_chunk_size, full_kv=True) + prefill_chunk_size=args.prefill_chunk_size, full_kv=args.cuda_trim) t0 = time.perf_counter() - res = fused_specdecode_generate_mlx_trim( - adapter, active_drafter, aux_prompt=aux_prompt, - embed_fn=embed_fn, lm_head_fn=lm_head_fn, - gen_tokens=args.max_new_tokens, block_size=args.block_size, - eos_ids=chat_eos, single_fused=args.single_fused) + if mlx_drafter is not None and args.cuda_trim: + res = fused_specdecode_generate_mlx_trim( + adapter, active_drafter, aux_prompt=aux_prompt, + embed_fn=embed_fn, lm_head_fn=lm_head_fn, + gen_tokens=args.max_new_tokens, block_size=args.block_size, + eos_ids=chat_eos, single_fused=args.single_fused) + elif mlx_drafter is not None: + res = fused_specdecode_generate_mlx( + adapter, active_drafter, aux_prompt=aux_prompt, + embed_fn=embed_fn, lm_head_fn=lm_head_fn, + gen_tokens=args.max_new_tokens, block_size=args.block_size, + eos_ids=chat_eos) + else: + res = fused_specdecode_generate( + adapter, active_drafter, aux_prompt=aux_prompt, + embed_fn=embed_fn, lm_head_fn=lm_head_fn, + gen_tokens=args.max_new_tokens, block_size=args.block_size, + eos_ids=chat_eos, argmax_fn=argmax_fn, arange_fn=arange_fn, + cat_aux_fn=cat_aux_fn, allow_greedy_fallback=False) res["decode_s"] = round(time.perf_counter() - t0, 3) + res["f_theta_ran"] = f_theta_ran + res["f_theta_layers"] = sorted(rk.keys()) if rk else [] try: txt = tokenizer.decode(res["tokens"], skip_special_tokens=True) except TypeError: @@ -789,11 +824,15 @@ def _gen_turn(pid: List[int]) -> Dict[str, Any]: "user": u, "text": res["text"], "tokens": res["decode_tokens"], "blocks": res["blocks"], "mean_accept_len": res["mean_accept_len"], + "f_theta_ran": res["f_theta_ran"], + "f_theta_layers": res["f_theta_layers"], "decode_s": res["decode_s"], "decode_tps": round(tps, 2), "resident_kv_bytes": res["resident_kv_bytes"]}) print(f"[chat] USER {u!r}", file=sys.stderr, flush=True) print(f"[chat] GEMMA-4 {res['text'][:200]!r} (blocks=" f"{res['blocks']}, accept_len={res['mean_accept_len']}, " + f"f_theta_ran={res['f_theta_ran']} " + f"layers={res['f_theta_layers']}, " f"{round(tps,2)} tok/s, kv={res['resident_kv_bytes']/1e6:.1f}MB)", file=sys.stderr, flush=True) report = { diff --git a/tests/inference_engine/bridge/test_manifest.py b/tests/inference_engine/bridge/test_manifest.py index d493e28f..bf2e8f6b 100644 --- a/tests/inference_engine/bridge/test_manifest.py +++ b/tests/inference_engine/bridge/test_manifest.py @@ -81,6 +81,7 @@ def test_allowlist_contains_exactly_the_documented_presets(): "mlx-batched-pad-decode", "mlx-env-probe", "mlx-kakeya-chat-smoke", + "mlx-kakeya-fused-chat-ftheta", "mlx-kakeya-fused-chat-smoke", "mlx-multitenant-pressure", "mlx-upgrade", @@ -149,6 +150,20 @@ def test_mlx_kakeya_chat_smoke_preset_resolves(): assert not [t for t in argv if t.startswith("{") and t.endswith("}")] +def test_mlx_kakeya_fused_chat_ftheta_preset_runs_f_theta_path(): + request = parse_manifest(_manifest( + preset="mlx-kakeya-fused-chat-ftheta", + params={"max_new_tokens": "32", "block_size": "4"})) + (argv,) = build_commands(request, HARNESS_ENV) + assert argv[1].endswith("k3_integrated_niah_eval_mac.py") + # torch drafter + f_θ path: --force-f-theta, and NOT --all-mlx-drafter + assert "--force-f-theta" in argv + assert "--fused-specdecode" in argv + assert "--all-mlx-drafter" not in argv + assert HARNESS_ENV["KAKEYA_MAC_FTHETA_DIR"] in argv + assert HARNESS_ENV["KAKEYA_MAC_DRAFTER_ID"] in argv + + def test_mlx_kakeya_fused_chat_smoke_preset_resolves(): request = parse_manifest(_manifest( preset="mlx-kakeya-fused-chat-smoke", From f73f05712752812d332a60bf7ee87ad587b2b9e8 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 08:25:37 +0000 Subject: [PATCH 14/15] =?UTF-8?q?feat(mac=20chat):=20f=5F=CE=B8=20default-?= =?UTF-8?q?ON=20in=20interactive=20chat=20+=20ADR=200015=20update?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - --chat now auto-enables --force-f-theta (full verifier/proposer/f_θ pipeline) unless --all-mlx-drafter (fast, f_θ-bypassed) or native baseline is chosen. - mlx-kakeya-fused-chat-ftheta preset drops the explicit --force-f-theta to verify the DEFAULT-on behavior on the Mac. - ADR 0015: new section 'Mac (MLX) interactive engine — full pipeline, f_θ default-ON' documenting the engine, the gemma-4 recall-irrelevant-but-runs caveat, the forensic timeline (f_θ S5-bypassed 2026-06-12 b3a04d0; proposer blocks=0 caught by 0a6fb19), and the measured result (f_theta_ran=TRUE 25 sliding layers + proposer blocks>0). Co-authored-by: FluffyAIcode --- ...5-kakeya-attention-and-engine-substrate.md | 35 +++++++++++++++++++ inference_engine/bridge/manifest.py | 2 +- .../research/k3_integrated_niah_eval_mac.py | 9 +++++ .../inference_engine/bridge/test_manifest.py | 5 +-- 4 files changed, 48 insertions(+), 3 deletions(-) diff --git a/docs/adr/0015-kakeya-attention-and-engine-substrate.md b/docs/adr/0015-kakeya-attention-and-engine-substrate.md index c3536757..5090b2bf 100644 --- a/docs/adr/0015-kakeya-attention-and-engine-substrate.md +++ b/docs/adr/0015-kakeya-attention-and-engine-substrate.md @@ -93,6 +93,41 @@ full-attention fraction: *only* way to bound memory at full recall — and vLLM, having no restoration, **must keep full KV and cannot match it**. This is the engine's target regime. +## Mac (MLX) interactive engine — full verifier/proposer/f_θ pipeline, f_θ default-ON + +The Apple-Silicon interactive CLI (`scripts/research/k3_integrated_niah_eval_mac.py +--chat`) runs the **full Kakeya engine** — gemma-4 verifier (MLX) + **DFlash +proposer** (fused spec-decode) + **f_θ K/V restoration** + **S5 bounded KV** — not +verifier-only. It reuses the validated `fused_specdecode_generate_mlx_trim` +per-turn sequence; the NIAH evidence loop is untouched. + +- **f_θ runs by default in chat.** `--force-f-theta` is auto-enabled in `--chat` + unless the fast all-MLX path (`--all-mlx-drafter`, f_θ bypassed) is explicitly + chosen. It bypasses the S5 native-prefill short-circuit so f_θ **executes** each + turn: it projects the proposer's hidden states → verifier K/V for the **25 + sliding layers** and injects them. +- **gemma-4 caveat (honest).** On gemma-4 those restored sliding-layer K/V are + **recall-irrelevant** — the 5 exact full-attention layers carry recall (the "S5 + free lunch"), so f_θ's output is effectively *discarded by the recall path*. We + still run f_θ by default so the **full verifier/proposer/f_θ pipeline is + exercised end-to-end**; on **full-attention models** the same f_θ path is + load-bearing (it is the only way to bound memory at full recall). +- **Forensic — when f_θ stopped running.** f_θ was silently bypassed under + `--s5-exact-full-attn` on **2026-06-12** by the *"Optimize MLX adaptive S5 + native smoke path"* commits (`b3a04d0` / `1f6e58c`), which made + `build_restoration` short-circuit to `{}` under S5; the same *"adaptive S5 + native"* path also let the proposer go to `blocks=0` while keeping the fused + label — caught by the evidence gate (`0a6fb19`, *"enforce PR #109 review + constraints"*) which added `--force-fused-specdecode`. Both squashed into main + via #117. f_θ remained S5-bypassed until `--force-f-theta` (this ADR's change) + made it default-on in the interactive chat. +- **Measured (Mac M4, via the git-bus bridge).** Both chat turns: + **`f_theta_ran=TRUE`** restoring the **25 sliding layers** + **proposer + `blocks=2/4`, `mean_accept_len=4.0/3.5`** + correct answers ("Paris"; "red, + yellow, and blue") + natural `` stop + bounded resident KV + (12–18 MB). Torch-bridge path is slow (~0.5–6 tok/s); the all-MLX path + (proposer-only) is the fast option. + ## Feasibility probes so far (informed the design — NOT the product) These ran on the eager-transformers research bench; they validate correctness and diff --git a/inference_engine/bridge/manifest.py b/inference_engine/bridge/manifest.py index 2cccf2ad..221c6243 100644 --- a/inference_engine/bridge/manifest.py +++ b/inference_engine/bridge/manifest.py @@ -729,7 +729,7 @@ def _harness_preset( "--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}", "--drafter-id", "${ENV:KAKEYA_MAC_DRAFTER_ID}", "--f-theta-dir", "${ENV:KAKEYA_MAC_FTHETA_DIR}", - "--s5-exact-full-attn", "--fused-specdecode", "--force-f-theta", + "--s5-exact-full-attn", "--fused-specdecode", "--sink-size", "4", "--window-size", "64", "--block-size", "{block_size}", "--max-new-tokens", "{max_new_tokens}", diff --git a/scripts/research/k3_integrated_niah_eval_mac.py b/scripts/research/k3_integrated_niah_eval_mac.py index 75a06c8c..55f60c34 100644 --- a/scripts/research/k3_integrated_niah_eval_mac.py +++ b/scripts/research/k3_integrated_niah_eval_mac.py @@ -262,6 +262,15 @@ def main() -> int: args.fused_specdecode = True args.force_fused_specdecode = True adaptive_s5_native = args.native_baseline_bypass + # Interactive chat runs the FULL verifier/proposer/f_θ pipeline by DEFAULT: + # f_θ executes each turn (torch drafter + f_θ) unless the fast all-MLX path + # (--all-mlx-drafter, f_θ bypassed) or the native baseline is explicitly chosen. + if args.chat and not args.all_mlx_drafter and not args.native_baseline_bypass: + if not args.force_f_theta: + print("[chat] f_θ default-ON for interactive chat (torch drafter + f_θ); " + "pass --all-mlx-drafter for the fast f_θ-bypassed path.", + file=sys.stderr, flush=True) + args.force_f_theta = True if args.all_mlx_drafter and not args.s5_exact_full_attn: raise SystemExit( "--all-mlx-drafter requires --s5-exact-full-attn: the all-MLX " diff --git a/tests/inference_engine/bridge/test_manifest.py b/tests/inference_engine/bridge/test_manifest.py index bf2e8f6b..3241f3ca 100644 --- a/tests/inference_engine/bridge/test_manifest.py +++ b/tests/inference_engine/bridge/test_manifest.py @@ -156,8 +156,9 @@ def test_mlx_kakeya_fused_chat_ftheta_preset_runs_f_theta_path(): params={"max_new_tokens": "32", "block_size": "4"})) (argv,) = build_commands(request, HARNESS_ENV) assert argv[1].endswith("k3_integrated_niah_eval_mac.py") - # torch drafter + f_θ path: --force-f-theta, and NOT --all-mlx-drafter - assert "--force-f-theta" in argv + # torch drafter + f_θ path: --chat WITHOUT --all-mlx-drafter → f_θ default-ON + # (no explicit --force-f-theta needed; the harness enables it for chat). + assert "--chat" in argv assert "--fused-specdecode" in argv assert "--all-mlx-drafter" not in argv assert HARNESS_ENV["KAKEYA_MAC_FTHETA_DIR"] in argv From 496339012ac072eb244ab6a98eb2f271df1d149f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 17 Jun 2026 08:46:33 +0000 Subject: [PATCH 15/15] fix(mac chat): strip gemma-4 reasoning-channel bleed (\nthought) from displayed answer Cut the decoded text at the first channel marker (<|channel / --- scripts/chat_mlx_kakeya.py | 8 +++++++- scripts/research/k3_integrated_niah_eval_mac.py | 7 +++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/scripts/chat_mlx_kakeya.py b/scripts/chat_mlx_kakeya.py index e7adb1a8..3dddb541 100644 --- a/scripts/chat_mlx_kakeya.py +++ b/scripts/chat_mlx_kakeya.py @@ -192,8 +192,14 @@ def _iter(): stop_reason = "loop" break dt = max(time.time() - t0, 1e-9) + _txt = tok.decode(toks, skip_special_tokens=True) + # gemma-4 sometimes bleeds its reasoning channel after the answer; cut it. + for _cut in ("<|channel", " 0: + _txt = _txt[:_i] return { - "text": tok.decode(toks, skip_special_tokens=True), + "text": _txt.strip(), "n_tokens": len(toks), "stop_reason": stop_reason, "decode_tps": round(len(toks) / dt, 2), diff --git a/scripts/research/k3_integrated_niah_eval_mac.py b/scripts/research/k3_integrated_niah_eval_mac.py index 55f60c34..d39b535d 100644 --- a/scripts/research/k3_integrated_niah_eval_mac.py +++ b/scripts/research/k3_integrated_niah_eval_mac.py @@ -808,6 +808,13 @@ def _gen_turn(pid: List[int]) -> Dict[str, Any]: txt = tokenizer.decode(res["tokens"]) for marker in ("", "", ""): txt = txt.replace(marker, "") + # gemma-4 sometimes bleeds its reasoning channel after the answer + # (e.g. a trailing "\nthought ...") — cut at the first channel + # marker so the chat shows only the natural-language answer. + for cut in ("<|channel", " 0: + txt = txt[:idx] res["text"] = txt.strip() res["resident_kv_bytes"] = int( sum(int(getattr(c, "nbytes", 0)) for c in (adapter._cache or [])))