Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
6b36fa4
feat(mac): gemma-4 interactive CLI chat on the Kakeya-for-Mac (MLX) e…
cursoragent Jun 17, 2026
bfcac00
fix(mac chat): use template enable_thinking instead of raw channel ma…
cursoragent Jun 17, 2026
39c85d0
fix(mac chat): add repetition_penalty (default 1.3) to stop greedy JS…
cursoragent Jun 17, 2026
c2168ec
fix(mac chat): generic loop guard — stop + trim when output degenerat…
cursoragent Jun 17, 2026
f55427a
fix(mac chat): strict consecutive-loop guard (no longer trims real an…
cursoragent Jun 17, 2026
bd2eaad
fix(mac chat): raise default --max-new-tokens 256->1024 + show stop r…
cursoragent Jun 17, 2026
3616e33
fix(mac chat): stop at gemma's natural turn end (real EOS detection) …
cursoragent Jun 17, 2026
9064d53
fix(mac chat): default --max-new-tokens=2048 as a generous backstop (…
cursoragent Jun 17, 2026
f1e648b
chore(mac-bridge): raise MAX_NEW_TOKENS bound 512->2048 (chat answers…
cursoragent Jun 17, 2026
728d694
feat(mac): FULL fused-engine chat (verifier+DFlash proposer+f_θ+S5) —…
cursoragent Jun 17, 2026
559daf8
polish(mac fused chat): strip stop-token markers from displayed text …
cursoragent Jun 17, 2026
a6dd2e2
docs(mac fused chat): accurate engine label — proposer live; f_θ bypa…
cursoragent Jun 17, 2026
b73198b
feat(mac fused chat): --force-f-theta — f_θ actually RUNS each turn (…
cursoragent Jun 17, 2026
f73f057
feat(mac chat): f_θ default-ON in interactive chat + ADR 0015 update
cursoragent Jun 17, 2026
4963390
fix(mac chat): strip gemma-4 reasoning-channel bleed (\nthought) from…
cursoragent Jun 17, 2026
6911092
Merge remote-tracking branch 'origin/main' into AgentMemory/mac-fused…
cursoragent Jun 17, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions docs/adr/0015-kakeya-attention-and-engine-substrate.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 `<end_of_turn>` 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
Expand Down
71 changes: 71 additions & 0 deletions inference_engine/bridge/manifest.py
Original file line number Diff line number Diff line change
Expand Up @@ -678,6 +678,77 @@ 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,
),
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",
"--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,
),
)
}

Expand Down
8 changes: 7 additions & 1 deletion scripts/chat_mlx_kakeya.py
Original file line number Diff line number Diff line change
Expand Up @@ -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", "<channel", "\nthought", "\nthink"):
_i = _txt.find(_cut)
if _i > 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),
Expand Down
Loading
Loading