feat(mac): full fused-engine chat (verifier+proposer+f_θ, f_θ default-ON) + reasoning-channel strip - #144
Merged
Conversation
…ngine 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 <FluffyAIcode@users.noreply.github.com>
…rker 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 <FluffyAIcode@users.noreply.github.com>
…ON 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 <FluffyAIcode@users.noreply.github.com>
…es into a repeat 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 <FluffyAIcode@users.noreply.github.com>
…swers) _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 <FluffyAIcode@users.noreply.github.com>
…eason (eos/max/loop)
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 <FluffyAIcode@users.noreply.github.com>
…— output length must not be a user knob
User correctly noted: needing --max-new-tokens to avoid truncation = unusable.
Root cause: <end_of_turn> 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('<end_of_turn>'/'<eos>'); 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 <FluffyAIcode@users.noreply.github.com>
…natural EOS stops first) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>
… need it; natural EOS stops first) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>
… not verifier-only 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 (<end_of_turn>). 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 <FluffyAIcode@users.noreply.github.com>
…(skip_special_tokens) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>
…ssed on gemma-4 (S5 free lunch) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>
…even if recall-irrelevant on gemma-4) 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 <FluffyAIcode@users.noreply.github.com>
- --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>
… displayed answer Cut the decoded text at the first channel marker (<|channel / <channel / \nthought / \nthink) so the chat shows only the natural-language answer. Applied to both the full fused-engine chat (harness _gen_turn) and the lightweight AR chat. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>
…-chat-ftheta-2815 # Conflicts: # inference_engine/bridge/manifest.py # scripts/chat_mlx_kakeya.py # tests/inference_engine/bridge/test_manifest.py Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow-up to #143 (merged). #143 squash-merged only the early subset (verifier-only
chat_mlx_kakeya.py+ EOS fix); this PR brings the full fused-engine chat + f_θ default-on + ADR + a one-command local launcher + the\nthoughtstrip.One-command local launcher (NEW):
scripts/run_kakeya_mac.shRuns the complete Kakeya engine on a Mac mini in one command — gemma-4 verifier (MLX) + DFlash proposer + f_θ + S5 bounded KV, interactive chat:
Resolves model/drafter/f_θ from
KAKEYA_MAC_*env vars (fallbacks included); preflights MLX + model dir; passes extra flags through.Verified end-to-end on Mac M4 (bridge run 27678751690): the launcher resolved the model from
$KAKEYA_MAC_VERIFIER_PATH, loaded the engine, and answered "The capital of France is Paris." (proposerblocks=2, accept_len=4.0, 15 tok/s).mac_kakeya_launcher_smoke.log
Underlying full engine (the harness
--chat)f_θ runs by DEFAULT in chat (
--force-f-thetaauto-on unless--all-mlx-drafter). Verified on Mac M4 (run 27675929900): both turnsf_theta_ran=TRUE(25 sliding layers) + proposerblocks>0+ correct answers + bounded KV. On gemma-4 the restored sliding K/V are recall-irrelevant (exact layers carry recall — S5 free lunch), yet f_θ executes; on full-attention models the same f_θ path is load-bearing.mac_gemma4_kakeya_ftheta_chat.log
Also
\nthoughtreasoning-channel strip in chat decode (full engine + lightweight AR).b3a04d0; proposerblocks=0bypass caught by0a6fb19).Changes
scripts/run_kakeya_mac.sh(new) — one-command local launcher.scripts/research/k3_integrated_niah_eval_mac.py—--chat/--chat-scripted/--force-f-thetafull-engine REPL (f_θ default-on; all-MLX + torch paths); reasoning-channel strip.scripts/chat_mlx_kakeya.py— reasoning-channel strip.inference_engine/bridge/manifest.py— presetsmlx-kakeya-fused-chat-{smoke,ftheta},mlx-kakeya-launcher-smoke.docs/adr/0015-…md— Mac engine + f_θ default-on + forensics.tests/inference_engine/bridge/test_manifest.py— preset coverage (32 presets, 100%).Testing
pytest tests/inference_engine/bridge/test_manifest.py— 31 passed, 100% coverage.unit tests + 100% coverage (3.12)green on this PR; merge-conflict with main resolved.f_theta_ran=TRUE, 25 layers);--fastpath.pytest -m integration on Mac M4red — pre-existing self-hosted-runner infra issue, unrelated.To show artifacts inline, enable in settings.