Skip to content

Commit f73f057

Browse files
feat(mac chat): f_θ default-ON in interactive chat + ADR 0015 update
- --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 <FluffyAIcode@users.noreply.github.com>
1 parent b73198b commit f73f057

4 files changed

Lines changed: 48 additions & 3 deletions

File tree

docs/adr/0015-kakeya-attention-and-engine-substrate.md

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,41 @@ full-attention fraction:
9393
*only* way to bound memory at full recall — and vLLM, having no restoration,
9494
**must keep full KV and cannot match it**. This is the engine's target regime.
9595

96+
## Mac (MLX) interactive engine — full verifier/proposer/f_θ pipeline, f_θ default-ON
97+
98+
The Apple-Silicon interactive CLI (`scripts/research/k3_integrated_niah_eval_mac.py
99+
--chat`) runs the **full Kakeya engine** — gemma-4 verifier (MLX) + **DFlash
100+
proposer** (fused spec-decode) + **f_θ K/V restoration** + **S5 bounded KV** — not
101+
verifier-only. It reuses the validated `fused_specdecode_generate_mlx_trim`
102+
per-turn sequence; the NIAH evidence loop is untouched.
103+
104+
- **f_θ runs by default in chat.** `--force-f-theta` is auto-enabled in `--chat`
105+
unless the fast all-MLX path (`--all-mlx-drafter`, f_θ bypassed) is explicitly
106+
chosen. It bypasses the S5 native-prefill short-circuit so f_θ **executes** each
107+
turn: it projects the proposer's hidden states → verifier K/V for the **25
108+
sliding layers** and injects them.
109+
- **gemma-4 caveat (honest).** On gemma-4 those restored sliding-layer K/V are
110+
**recall-irrelevant** — the 5 exact full-attention layers carry recall (the "S5
111+
free lunch"), so f_θ's output is effectively *discarded by the recall path*. We
112+
still run f_θ by default so the **full verifier/proposer/f_θ pipeline is
113+
exercised end-to-end**; on **full-attention models** the same f_θ path is
114+
load-bearing (it is the only way to bound memory at full recall).
115+
- **Forensic — when f_θ stopped running.** f_θ was silently bypassed under
116+
`--s5-exact-full-attn` on **2026-06-12** by the *"Optimize MLX adaptive S5
117+
native smoke path"* commits (`b3a04d0` / `1f6e58c`), which made
118+
`build_restoration` short-circuit to `{}` under S5; the same *"adaptive S5
119+
native"* path also let the proposer go to `blocks=0` while keeping the fused
120+
label — caught by the evidence gate (`0a6fb19`, *"enforce PR #109 review
121+
constraints"*) which added `--force-fused-specdecode`. Both squashed into main
122+
via #117. f_θ remained S5-bypassed until `--force-f-theta` (this ADR's change)
123+
made it default-on in the interactive chat.
124+
- **Measured (Mac M4, via the git-bus bridge).** Both chat turns:
125+
**`f_theta_ran=TRUE`** restoring the **25 sliding layers** + **proposer
126+
`blocks=2/4`, `mean_accept_len=4.0/3.5`** + correct answers ("Paris"; "red,
127+
yellow, and blue") + natural `<end_of_turn>` stop + bounded resident KV
128+
(12–18 MB). Torch-bridge path is slow (~0.5–6 tok/s); the all-MLX path
129+
(proposer-only) is the fast option.
130+
96131
## Feasibility probes so far (informed the design — NOT the product)
97132

98133
These ran on the eager-transformers research bench; they validate correctness and

inference_engine/bridge/manifest.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -729,7 +729,7 @@ def _harness_preset(
729729
"--verifier-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}",
730730
"--drafter-id", "${ENV:KAKEYA_MAC_DRAFTER_ID}",
731731
"--f-theta-dir", "${ENV:KAKEYA_MAC_FTHETA_DIR}",
732-
"--s5-exact-full-attn", "--fused-specdecode", "--force-f-theta",
732+
"--s5-exact-full-attn", "--fused-specdecode",
733733
"--sink-size", "4", "--window-size", "64",
734734
"--block-size", "{block_size}",
735735
"--max-new-tokens", "{max_new_tokens}",

scripts/research/k3_integrated_niah_eval_mac.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -262,6 +262,15 @@ def main() -> int:
262262
args.fused_specdecode = True
263263
args.force_fused_specdecode = True
264264
adaptive_s5_native = args.native_baseline_bypass
265+
# Interactive chat runs the FULL verifier/proposer/f_θ pipeline by DEFAULT:
266+
# f_θ executes each turn (torch drafter + f_θ) unless the fast all-MLX path
267+
# (--all-mlx-drafter, f_θ bypassed) or the native baseline is explicitly chosen.
268+
if args.chat and not args.all_mlx_drafter and not args.native_baseline_bypass:
269+
if not args.force_f_theta:
270+
print("[chat] f_θ default-ON for interactive chat (torch drafter + f_θ); "
271+
"pass --all-mlx-drafter for the fast f_θ-bypassed path.",
272+
file=sys.stderr, flush=True)
273+
args.force_f_theta = True
265274
if args.all_mlx_drafter and not args.s5_exact_full_attn:
266275
raise SystemExit(
267276
"--all-mlx-drafter requires --s5-exact-full-attn: the all-MLX "

tests/inference_engine/bridge/test_manifest.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -156,8 +156,9 @@ def test_mlx_kakeya_fused_chat_ftheta_preset_runs_f_theta_path():
156156
params={"max_new_tokens": "32", "block_size": "4"}))
157157
(argv,) = build_commands(request, HARNESS_ENV)
158158
assert argv[1].endswith("k3_integrated_niah_eval_mac.py")
159-
# torch drafter + f_θ path: --force-f-theta, and NOT --all-mlx-drafter
160-
assert "--force-f-theta" in argv
159+
# torch drafter + f_θ path: --chat WITHOUT --all-mlx-drafter → f_θ default-ON
160+
# (no explicit --force-f-theta needed; the harness enables it for chat).
161+
assert "--chat" in argv
161162
assert "--fused-specdecode" in argv
162163
assert "--all-mlx-drafter" not in argv
163164
assert HARNESS_ENV["KAKEYA_MAC_FTHETA_DIR"] in argv

0 commit comments

Comments
 (0)