From 02ba6bf53e5f346532266a5e023dbbbcb61f5a3a Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 3 Aug 2026 21:11:07 -0700 Subject: [PATCH 1/7] feat(chat,run): Ctrl-T finishes thinking early (SIGINFO + raw-byte fallback) --- gmlx/chat.py | 80 ++++++++++++------- gmlx/cli.py | 3 + gmlx/generation.py | 20 ++++- gmlx/thinking_budget.py | 144 +++++++++++++++++++++++++++++++--- tests/test_thinking_budget.py | 75 ++++++++++++++++++ 5 files changed, 277 insertions(+), 45 deletions(-) diff --git a/gmlx/chat.py b/gmlx/chat.py index 54e2570..8fa9d87 100644 --- a/gmlx/chat.py +++ b/gmlx/chat.py @@ -30,7 +30,9 @@ spinner that resolves to the same payoff. **Ctrl-O** toggles expand<->collapse live during a reply; ``/reasoning [show|hide|raw]`` (startup ``--reasoning``) sets the default and reaches ``raw`` (verbatim). The stored turn keeps the raw - text, so this is display-only; + text, so this is display-only. **Ctrl-T** while the model is thinking closes + the thinking block early (as if the thinking budget had just run out) so the + answer starts now; * ``/load `` prefills the *next* prompt with a text file's contents (via the readline startup hook), so it can be edited before Enter sends it; Tab completes ``/command`` names, ``/history`` arguments, and file paths @@ -715,6 +717,7 @@ def _print_shim_help(state: ChatState) -> None: "- '/reasoning [show|hide|raw]' control how thinking is displayed " "(Ctrl-O collapses/expands it live during a reply)" ) + print("- Ctrl-T while the model is thinking: wrap up and answer now") print("- '/render [plain|lite|rich]' set reply markdown rendering") print("- '/theme [NAME] [cb]' set the color theme ('cb' = colorblind accents)") print("- '/model' show the loaded model card - '/stats' session totals") @@ -1676,8 +1679,10 @@ class _EscCancel: immediately, not echoed into the streaming text; Ctrl-C still signals). ``pressed()`` polls stdin without blocking and drains whatever arrived, so arrow-key escape tails don't leak into the next prompt. Esc cancels; - Ctrl-O (``\\x0f``) fires ``on_toggle`` (collapse/expand thinking). Inert - when stdin is not a tty. + Ctrl-O (``\\x0f``) fires ``on_toggle`` (collapse/expand thinking); Ctrl-T + (``\\x14``) finishes thinking early - normally it arrives as SIGINFO (the + tty's status character survives cbreak), the raw byte is the fallback for + terminals with the status character unset. Inert when stdin is not a tty. """ def __init__(self, on_toggle=None): @@ -1709,6 +1714,10 @@ def pressed(self) -> bool: hit = True elif ch == b"\x0f" and self._on_toggle is not None: self._on_toggle() + elif ch == b"\x14": + from .thinking_budget import finish_thinking_now + + finish_thinking_now() return hit def __exit__(self, *exc): @@ -3159,6 +3168,9 @@ def _apply_session(doc: dict, name: str) -> None: _apply_session(doc, spath.stem) from .cli import _DIFFUSION_MAX_TOKENS, _UNCAPPED_MAX_TOKENS + from .thinking_budget import install_finish_thinking_key + + install_finish_thinking_key() # ^T closes an open thinking block while True: parts = [] @@ -3512,33 +3524,41 @@ def _apply_session(doc: dict, name: str) -> None: # Thinking-token cap (fresh per turn - the processor is stateful). Honored # regardless of enable_thinking; seed the in-thinking state from whether # the rendered prompt actually opens a block (see generation.generate). - if state.thinking_budget is not None: - from .thinking_budget import ( - make_thinking_budget_processor, - prompt_opens_thinking, - ) + # Built even with no budget set (interruptible): it never trips on its + # own, but ^T can close the thinking block through it mid-reply. + from .thinking_budget import ( + clear_finish_key_target, + make_thinking_budget_processor, + prompt_opens_thinking, + set_finish_key_target, + ) - tbp = make_thinking_budget_processor( - tok, - state.thinking_budget, - start_in_thinking=prompt_opens_thinking(prompt_text), - ) - if tbp is not None: - logits_processors = list(logits_processors) + [tbp] - first_turn = False # the system prompt is in the cache either way - reply, canceled = _stream_reply( - stream_generate( - model, - tok, - prompt, - max_tokens=eff_max, - sampler=sampler, - logits_processors=logits_processors, - prompt_cache=cache, - **kv_kwargs, - ), - state, - stops=args.stop, - start_in_thinking=_opens_thinking(prompt_text), + tbp = make_thinking_budget_processor( + tok, + state.thinking_budget, + start_in_thinking=prompt_opens_thinking(prompt_text), + interruptible=True, ) + if tbp is not None: + logits_processors = list(logits_processors) + [tbp] + first_turn = False # the system prompt is in the cache either way + set_finish_key_target(tbp) + try: + reply, canceled = _stream_reply( + stream_generate( + model, + tok, + prompt, + max_tokens=eff_max, + sampler=sampler, + logits_processors=logits_processors, + prompt_cache=cache, + **kv_kwargs, + ), + state, + stops=args.stop, + start_in_thinking=_opens_thinking(prompt_text), + ) + finally: + clear_finish_key_target() _end_turn(state, reply, canceled, cache=cache) diff --git a/gmlx/cli.py b/gmlx/cli.py index f02540a..0650962 100644 --- a/gmlx/cli.py +++ b/gmlx/cli.py @@ -1416,6 +1416,9 @@ def _run_generate(args) -> int: # Cap-vs-EOS is not reported back on this path (generate returns text), so # no cap-hit note here; the MTP/VLM paths and chat print one. + from .thinking_budget import install_finish_thinking_key + + install_finish_thinking_key() # ^T closes an open thinking block print(f"[generate] max_tokens={max_tokens_label(args)} temp={args.temp}\n") generate( model, diff --git a/gmlx/generation.py b/gmlx/generation.py index 718b52e..34e51ba 100644 --- a/gmlx/generation.py +++ b/gmlx/generation.py @@ -250,8 +250,10 @@ def generate( thinking tokens are generated it forces ```` so the model answers. An explicit budget is honored even when ``enable_thinking`` is false (a model may still emit ````); it is a no-op only when no ```` is ever - generated or the tokenizer lacks a ```` token. Returns the generated - text. ``prefill_progress`` shows a stderr spinner during a long prefill + generated or the tokenizer lacks a ```` token. With or without a + budget, ^T (see ``thinking_budget.install_finish_thinking_key``, armed by + the CLI entry points) force-closes an open thinking block the same way. + Returns the generated text. ``prefill_progress`` shows a stderr spinner during a long prefill (TTY only; cleared before the first token). ``reasoning`` shapes how a *verbose* stream displays a thinking model's chain-of-thought: ``"show"`` styles it under a label (the chat REPL's @@ -307,17 +309,21 @@ def generate( # prompt* actually opens a block (a pre-fill model opens it in the # prompt; a generate model emits it, which the processor detects). The flag # only shapes the prompt via the template - it never gates the cap. - if thinking_budget is not None and thinking_budget >= 0: + tbp = None + if thinking_budget is None or thinking_budget >= 0: from .thinking_budget import ( make_thinking_budget_processor, prompt_opens_thinking, ) + # interruptible: with no budget set the processor never trips on its + # own, but stays armed for the ^T finish-thinking key. tbp = make_thinking_budget_processor( tokenizer, thinking_budget, start_in_thinking=prompt_opens_thinking(prompt, tokenizer=tokenizer), verbose=verbose, + interruptible=True, ) if tbp is not None: logits_processors = list(logits_processors) + [tbp] @@ -434,6 +440,10 @@ def generate( gen_kwargs["prompt_progress_callback"] = progress_cb try: + if tbp is not None: + from .thinking_budget import set_finish_key_target + + set_finish_key_target(tbp) stop = [s for s in (stop or []) if s] styled = verbose and reasoning in ("show", "hide") if not stop: @@ -508,6 +518,10 @@ def generate( ) return "".join(pieces) finally: + if tbp is not None: + from .thinking_budget import clear_finish_key_target + + clear_finish_key_target() if close_progress is not None: close_progress() diff --git a/gmlx/thinking_budget.py b/gmlx/thinking_budget.py index b41864a..01b43e2 100644 --- a/gmlx/thinking_budget.py +++ b/gmlx/thinking_budget.py @@ -23,6 +23,10 @@ from __future__ import annotations +import os +import signal +import sys + import mlx.core as mx @@ -48,6 +52,13 @@ "write the complete final answer now.\n" ) +# Forced when the user presses ^T (finish thinking now) - same rationale, +# but the model should see the user's request, not a phantom budget. +_SKIP_WRAP_PHRASE = ( + "\n\nI've been asked to wrap up, so I'll stop reasoning here and " + "write the complete final answer now.\n" +) + def _template_think_pair(tokenizer): """The ``_THINK_PAIRS`` spelling the model's chat template actually uses, @@ -186,6 +197,9 @@ class ThinkingBudgetProcessor: of generated thinking tokens exceeds ``budget``, the next steps' logits are forced (one-hot) through ``forced_ids``; generation then continues normally. A natural ``end_seq`` before the budget disarms the processor. + ``budget=None`` never trips on its own: the processor is armed only for an + external ``request_close()`` (the ^T finish-thinking key), which closes the + block through ``skip_ids`` and then behaves exactly like a spent budget. ``start_seq`` / ``end_seq`` are token-id tuples. They are multi-token for channel-style delimiters (``<|channel>thought`` / ````) and length @@ -223,14 +237,17 @@ def __init__( start_in_thinking=True, eos_ids=None, reclose_ids=None, + skip_ids=None, ): self.end_seq = tuple(end_seq) self.start_seq = tuple(start_seq) if start_seq else None self.forced_ids = list(forced_ids) self.reclose_ids = list(reclose_ids) if reclose_ids else self.forced_ids - self.budget = max(0, int(budget)) + self.skip_ids = list(skip_ids) if skip_ids else self.forced_ids + self.budget = None if budget is None else max(0, int(budget)) self.in_thinking = bool(start_in_thinking) self.count = 0 + self._close_requested = False self._baseline = None # len(tokens) at first call -> skip the prompt self._forcing = False self._forced_idx = 0 @@ -271,6 +288,12 @@ def _strike(self): if self._strikes >= 3: self.done = True + def request_close(self): + """Finish thinking now (^T): the next step force-closes the open + thinking block exactly as if the budget had just been exceeded. + Consumed as a no-op when the model is not thinking.""" + self._close_requested = True + def __call__(self, tokens, logits): n = tokens.shape[0] if self._baseline is None: @@ -292,9 +315,11 @@ def __call__(self, tokens, logits): tokens, self.start_seq ): self.in_thinking = True - # A reopen after the budget was spent trips the forcing + self.count = 0 + # A reopen after a forced close trips the forcing # threshold immediately: closed on the spot. - self.count = self.budget + 1 if self._spent else 0 + if self._spent: + self._close_requested = True elif self.in_thinking: self.count += 1 elif self._answer_pending: @@ -303,13 +328,21 @@ def __call__(self, tokens, logits): self._strikes = 0 if self.done: return logits - if not self._forcing and self.in_thinking and self.count > self.budget: + over = self.budget is not None and self.count > self.budget + if not self._forcing and self.in_thinking and ( + over or self._close_requested): self._forcing = True self._forced_idx = 0 - # The first close carries the wrap-up phrase (the model must SEE + # The first close carries a wrap-up phrase (the model must SEE # itself decide to answer, or the cut thought continues untagged # in the answer); spent-mode recloses are terse. - self._force_seq = self.reclose_ids if self._spent else self.forced_ids + if self._spent: + self._force_seq = self.reclose_ids + elif self._close_requested: + self._force_seq = self.skip_ids + else: + self._force_seq = self.forced_ids + self._close_requested = False if self._forcing: fid = self._force_seq[self._forced_idx] self._forced_idx += 1 @@ -330,10 +363,15 @@ def __call__(self, tokens, logits): def make_thinking_budget_processor( tokenizer, budget, *, start_in_thinking=True, verbose=False, - eos_floor=True, + eos_floor=True, interruptible=False, ): """Build a thinking-budget logits processor, or ``None`` if unsupported. + ``interruptible`` builds the processor even with no budget (``None``): it + then never trips on its own but stays armed for the ^T finish-thinking + key (``request_close``). Without it, ``budget=None`` returns ``None`` as + ever. + Resolves the model's thinking delimiters via mlx-lm's tokenizer inference (````, ````, and the ``<|channel>thought`` / ```` channel format), so the cap is not tied to one spelling. @@ -354,11 +392,21 @@ def make_thinking_budget_processor( self-review on MiniMax-M3). Seeing itself decide to finalize keeps the trajectory coherent. Spent-mode recloses stay terse (``\\n`` + close). """ - if budget is None or budget < 0: + if budget is None and not interruptible: return None - start_seq, end_seq = _thinking_token_seqs(tokenizer) + if budget is not None and budget < 0: + return None + if budget is None: + # An implicit (^T-only) processor is best-effort: a tokenizer the + # probe can't handle must not break plain generation. + try: + start_seq, end_seq = _thinking_token_seqs(tokenizer) + except Exception: # noqa: BLE001 + return None + else: + start_seq, end_seq = _thinking_token_seqs(tokenizer) if not end_seq: - if verbose: + if verbose and budget is not None: print( "[thinking-budget] no thinking-end token detected; " "ignoring thinking budget" @@ -367,11 +415,15 @@ def make_thinking_budget_processor( nl_id = _last_token_id(tokenizer, "\n") reclose_ids = ([nl_id] if nl_id is not None else []) + list(end_seq) forced_ids = reclose_ids - if budget > 0: + skip_ids = reclose_ids + if budget is None or budget > 0: wrap = _encode_ids(tokenizer, _BUDGET_WRAP_PHRASE) if wrap: forced_ids = list(wrap) + list(end_seq) - if verbose: + wrap = _encode_ids(tokenizer, _SKIP_WRAP_PHRASE) + if wrap: + skip_ids = list(wrap) + list(end_seq) + if verbose and budget is not None: # Decode the actual forced ids first: the wrapper's think_end attr can # disagree with the resolved sequence (template-preferred spelling). try: @@ -394,4 +446,72 @@ def make_thinking_budget_processor( if eos_floor else [], reclose_ids=reclose_ids, + skip_ids=skip_ids, ) + + +# --- ^T: finish thinking now ------------------------------------------------- +# +# The chat REPL and `gmlx run` route Ctrl-T here. In a cooked (or cbreak) tty +# ^T is the BSD status character: the terminal driver turns it into SIGINFO, +# which install_finish_thinking_key catches; chat's mid-reply key listener +# also forwards a raw \x14 byte in case the status character is unset. Either +# way the in-flight generation's processor gets request_close() and the open +# thinking block is closed as if the budget had just run out. + +_target = None # the in-flight generation's processor, if any + + +def set_finish_key_target(processor) -> None: + """Expose ``processor`` (may be None) to the ^T finish-thinking trigger + for the duration of one generation; pair with clear_finish_key_target.""" + global _target + _target = processor + + +def clear_finish_key_target() -> None: + global _target + _target = None + + +def finish_thinking_now() -> bool: + """Ask the in-flight generation to close its thinking block. True when a + generation was listening (the request may still be a no-op if the model + is not currently thinking).""" + p = _target + if p is None or p.done: + return False + p.request_close() + return True + + +def _quiet_kernel_status() -> None: + # Besides SIGINFO, ^T makes the kernel print a "load: ..." status line + # straight onto the tty, mid-stream; NOKERNINFO turns that off. Python's + # termios doesn't export the flag, so use the value. The + # shell restores its own tty state at the next prompt. + try: + import termios + nokerninfo = getattr(termios, "NOKERNINFO", 0x02000000) + fd = sys.stdin.fileno() + if not os.isatty(fd): + return + attrs = termios.tcgetattr(fd) + attrs[3] |= nokerninfo + termios.tcsetattr(fd, termios.TCSANOW, attrs) + except Exception: # noqa: BLE001 - cosmetic; never block generation + pass + + +def install_finish_thinking_key() -> bool: + """Route ^T (SIGINFO, Darwin/BSD) to ``finish_thinking_now``. Returns + whether the handler is installed; False off-platform or off the main + thread. Safe to call more than once.""" + if not hasattr(signal, "SIGINFO"): + return False + try: + signal.signal(signal.SIGINFO, lambda signum, frame: finish_thinking_now()) + except ValueError: + return False + _quiet_kernel_status() + return True diff --git a/tests/test_thinking_budget.py b/tests/test_thinking_budget.py index baa098d..df8d3a9 100644 --- a/tests/test_thinking_budget.py +++ b/tests/test_thinking_budget.py @@ -235,6 +235,81 @@ def test_natural_close_skips_spent_mode(): assert p.done and not p._spent +# --- ^T finish-thinking (request_close / uncapped interruptible mode) -------- + +SKIP = 7 # stand-in for the ^T wrap phrase's first token + + +def test_request_close_forces_skip_sequence_then_spent_mode(): + p = ThinkingBudgetProcessor( + end_seq=(END,), forced_ids=[SKIP, END], budget=None, start_seq=(START,), + skip_ids=[NL, END], reclose_ids=[NL, END], + ) + tokens = [1, 2, START] + p(mx.array(tokens), _logits()) + for tid in range(30, 80): # no budget: thinking runs free until ^T + tokens.append(tid) + assert _argmax_if_forced(p(mx.array(tokens), _logits())) is None + p.request_close() + tokens.append(80) + assert _argmax_if_forced(p(mx.array(tokens), _logits())) == NL # skip_ids, + tokens.append(NL) # not forced_ids + assert _argmax_if_forced(p(mx.array(tokens), _logits())) == END + tokens.append(END) + p(mx.array(tokens), _logits()) + assert p._spent and not p.in_thinking + tokens.append(60) # a real answer token lands + p(mx.array(tokens), _logits()) + # A reopen after ^T is closed on the spot, budget-0 style. + tokens.append(START) + assert _argmax_if_forced(p(mx.array(tokens), _logits())) == NL + + +def test_request_close_consumed_when_not_thinking(): + p = ThinkingBudgetProcessor( + end_seq=(END,), forced_ids=[NL, END], budget=None, start_seq=(START,), + start_in_thinking=False, + ) + tokens = [1, 2, 3] + p(mx.array(tokens), _logits()) + p.request_close() + tokens.append(50) + assert _argmax_if_forced(p(mx.array(tokens), _logits())) is None + # The stale request must not close a block opened later. + tokens.append(START) + assert _argmax_if_forced(p(mx.array(tokens), _logits())) is None + tokens.append(51) + assert _argmax_if_forced(p(mx.array(tokens), _logits())) is None + assert p.in_thinking + + +def test_factory_interruptible_without_budget(): + assert make_thinking_budget_processor(_FakeTok(), None) is None + p = make_thinking_budget_processor(_FakeTok(), None, interruptible=True) + assert p is not None and p.budget is None + assert p.skip_ids[-1] == END + + +def test_finish_key_target_round_trip(): + from gmlx.thinking_budget import ( + clear_finish_key_target, + finish_thinking_now, + set_finish_key_target, + ) + + p = ThinkingBudgetProcessor( + end_seq=(END,), forced_ids=[NL, END], budget=None, start_seq=(START,) + ) + assert finish_thinking_now() is False # nothing armed + set_finish_key_target(p) + try: + assert finish_thinking_now() is True + assert p._close_requested + finally: + clear_finish_key_target() + assert finish_thinking_now() is False + + def test_factory_wires_eos_ids_and_floor(): class _EosTok(_FakeTok): eos_token_ids = {77, 33} From a0293f91874912d1aa500f5f1929494e5be25c0a Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 3 Aug 2026 21:11:23 -0700 Subject: [PATCH 2/7] docs: changelog entry for Ctrl-T finish-thinking key --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 62d1222..516d7a1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,12 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Added + +- Ctrl-T during a `chat` or `run` reply closes the model's open thinking + block early (as if the thinking budget had just run out) so the answer + starts now. + ### Fixed - CLI polish: `--help`/`--help-all` now win even after a value-taking flag, From 57b9719452668a74e422709a8c48cff4176906c0 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 3 Aug 2026 21:58:22 -0700 Subject: [PATCH 3/7] fix(chat,run): Ctrl-T works on the VLM path, prints a notice on MTP paths --- gmlx/chat.py | 67 +++++++++++++++++++++++------------ gmlx/cli.py | 53 +++++++++++++++++++-------- gmlx/generation.py | 48 +++++++++++++++++++++++-- gmlx/thinking_budget.py | 23 +++++++++++- tests/test_thinking_budget.py | 18 ++++++++++ 5 files changed, 169 insertions(+), 40 deletions(-) diff --git a/gmlx/chat.py b/gmlx/chat.py index 8fa9d87..b570666 100644 --- a/gmlx/chat.py +++ b/gmlx/chat.py @@ -3370,31 +3370,52 @@ def _apply_session(doc: dict, name: str) -> None: ) # No cross-turn KV cache on the VLM path yet: each turn # re-prefills the whole conversation (and re-encodes media). - reply, _canceled = _stream_reply( - vlm_stream( - model, - processor, - prompt, - image=list(vlm_images) or None, - audio=list(vlm_audios) or None, - max_tokens=eff_max, - temperature=s["temp"], - top_p=s["top_p"], - top_k=s["top_k"], - min_p=s["min_p"], - repetition_penalty=None if rep in (0.0, 1.0) else rep, - repetition_context_size=s["repetition_context_size"], - presence_penalty=s["presence_penalty"] or None, - frequency_penalty=s["frequency_penalty"] or None, - logit_bias=logit_bias, - resize_shape=resize_shape, - thinking_budget=state.thinking_budget, - **kv_kwargs, - ), - state, - stops=args.stop, + # mlx-vlm's own criteria enforces thinking_budget here; the + # budget-less interruptible processor rides along only so ^T can + # close a thinking block (mlx-vlm's generate_step applies extra + # logits_processors with the same (tokens, logits) contract). + from .thinking_budget import ( + clear_finish_key_target, + make_thinking_budget_processor, + set_finish_key_target, + ) + + tbp = make_thinking_budget_processor( + getattr(processor, "tokenizer", processor), + None, start_in_thinking=_opens_thinking(prompt), + interruptible=True, ) + set_finish_key_target(tbp) + try: + reply, _canceled = _stream_reply( + vlm_stream( + model, + processor, + prompt, + image=list(vlm_images) or None, + audio=list(vlm_audios) or None, + max_tokens=eff_max, + temperature=s["temp"], + top_p=s["top_p"], + top_k=s["top_k"], + min_p=s["min_p"], + repetition_penalty=None if rep in (0.0, 1.0) else rep, + repetition_context_size=s["repetition_context_size"], + presence_penalty=s["presence_penalty"] or None, + frequency_penalty=s["frequency_penalty"] or None, + logit_bias=logit_bias, + resize_shape=resize_shape, + thinking_budget=state.thinking_budget, + logits_processors=[tbp] if tbp is not None else None, + **kv_kwargs, + ), + state, + stops=args.stop, + start_in_thinking=_opens_thinking(prompt), + ) + finally: + clear_finish_key_target() if reply: vlm_msgs.append(_vlm_message(model_type, reply, "assistant")) _end_turn(state, reply, _canceled) diff --git a/gmlx/cli.py b/gmlx/cli.py index 0650962..83ddf5e 100644 --- a/gmlx/cli.py +++ b/gmlx/cli.py @@ -1310,6 +1310,10 @@ def _run_generate(args) -> int: logit_bias = parse_logit_bias(args.logit_bias) preset_native_fp_wire_env(args) + from .thinking_budget import install_finish_thinking_key + + install_finish_thinking_key() # ^T closes an open thinking block + if args.seed is not None: import mlx.core as mx @@ -1416,9 +1420,6 @@ def _run_generate(args) -> int: # Cap-vs-EOS is not reported back on this path (generate returns text), so # no cap-hit note here; the MTP/VLM paths and chat print one. - from .thinking_budget import install_finish_thinking_key - - install_finish_thinking_key() # ^T closes an open thinking block print(f"[generate] max_tokens={max_tokens_label(args)} temp={args.temp}\n") generate( model, @@ -1545,21 +1546,43 @@ def _run_vlm(args) -> int: if args.prefill_step_size is not None: extra["prefill_step_size"] = args.prefill_step_size + # ^T closes an open thinking block: mlx-vlm's generate_step applies extra + # logits_processors with the mlx-lm (tokens, logits) contract, so the + # budget-less interruptible processor rides along (mlx-vlm's own criteria + # still enforces --thinking-budget). + from .thinking_budget import ( + clear_finish_key_target, + install_finish_thinking_key, + make_thinking_budget_processor, + set_finish_key_target, + ) + + install_finish_thinking_key() + tbp = make_thinking_budget_processor( + getattr(processor, "tokenizer", processor), None, interruptible=True + ) + if tbp is not None: + extra["logits_processors"] = [tbp] + print( f"[generate] max_tokens={max_tokens_label(args)} temp={args.temp} " f"images={len(images)} audios={len(audios)}\n" ) - result = generate( - model, - processor, - prompt, - image=images or None, - audio=audios or None, - max_tokens=args.max_tokens, - temperature=args.temp, - verbose=True, - **extra, - ) + set_finish_key_target(tbp) + try: + result = generate( + model, + processor, + prompt, + image=images or None, + audio=audios or None, + max_tokens=args.max_tokens, + temperature=args.temp, + verbose=True, + **extra, + ) + finally: + clear_finish_key_target() if getattr(result, "finish_reason", None) == "length": warn_cap_hit(args, getattr(result, "generation_tokens", None)) return 0 @@ -1576,7 +1599,9 @@ def _run_vlm_mtp(args) -> int: from .chat import fold_thinking_flag, parse_template_config from .generation import generate_speculative from .mtp_load import load_vlm_mtp_model + from .thinking_budget import install_finish_thinking_key + install_finish_thinking_key() # ^T prints the MTP-path notice here if args.seed is not None: import mlx.core as mx diff --git a/gmlx/generation.py b/gmlx/generation.py index 34e51ba..ca43197 100644 --- a/gmlx/generation.py +++ b/gmlx/generation.py @@ -723,7 +723,34 @@ def _chunked_prefill_cache(lm, input_ids, chunk): return c -def generate_speculative( +_MTP_FINISH_WHY = "on the MTP path (run with --no-mtp to use it)" + + +def _with_mtp_finish_key_notice(fn, *args, **kwargs): + """Run ``fn`` with the ^T finish-thinking target armed as "unsupported": + the MTP walks expose no forced-close seam, so the key explains itself + instead of silently doing nothing.""" + from .thinking_budget import ( + FinishKeyUnsupported, + clear_finish_key_target, + set_finish_key_target, + ) + + set_finish_key_target(FinishKeyUnsupported(_MTP_FINISH_WHY)) + try: + return fn(*args, **kwargs) + finally: + clear_finish_key_target() + + +def generate_speculative(model, drafter, tokenizer, prompt, **kwargs) -> dict: + """See :func:`_generate_speculative` (^T-notice shim).""" + return _with_mtp_finish_key_notice( + _generate_speculative, model, drafter, tokenizer, prompt, **kwargs + ) + + +def _generate_speculative( model, drafter, tokenizer, @@ -1173,7 +1200,24 @@ def _stream_generate_speculative_owned( ) -def stream_generate_speculative( +def stream_generate_speculative(model, drafter, tokenizer, prompt, **kwargs): + """See :func:`_stream_generate_speculative` (^T-notice shim).""" + from .thinking_budget import ( + FinishKeyUnsupported, + clear_finish_key_target, + set_finish_key_target, + ) + + set_finish_key_target(FinishKeyUnsupported(_MTP_FINISH_WHY)) + try: + yield from _stream_generate_speculative( + model, drafter, tokenizer, prompt, **kwargs + ) + finally: + clear_finish_key_target() + + +def _stream_generate_speculative( model, drafter, tokenizer, diff --git a/gmlx/thinking_budget.py b/gmlx/thinking_budget.py index 01b43e2..596429f 100644 --- a/gmlx/thinking_budget.py +++ b/gmlx/thinking_budget.py @@ -474,12 +474,33 @@ def clear_finish_key_target() -> None: _target = None +class FinishKeyUnsupported: + """Armed as the ^T target on generation paths the key can't reach (the + MTP walks expose no logits-processor seam): pressing it then explains + itself once instead of silently doing nothing.""" + + def __init__(self, why: str): + self.why = why + self._told = False + + def note(self) -> None: + if not self._told: + self._told = True + print(f"\n[^T] finish-thinking isn't available {self.why}", + file=sys.stderr) + + def finish_thinking_now() -> bool: """Ask the in-flight generation to close its thinking block. True when a generation was listening (the request may still be a no-op if the model is not currently thinking).""" p = _target - if p is None or p.done: + if p is None: + return False + if isinstance(p, FinishKeyUnsupported): + p.note() + return False + if p.done: return False p.request_close() return True diff --git a/tests/test_thinking_budget.py b/tests/test_thinking_budget.py index df8d3a9..f986657 100644 --- a/tests/test_thinking_budget.py +++ b/tests/test_thinking_budget.py @@ -310,6 +310,24 @@ def test_finish_key_target_round_trip(): assert finish_thinking_now() is False +def test_finish_key_unsupported_notes_once(capsys): + from gmlx.thinking_budget import ( + FinishKeyUnsupported, + clear_finish_key_target, + finish_thinking_now, + set_finish_key_target, + ) + + set_finish_key_target(FinishKeyUnsupported("on the MTP path")) + try: + assert finish_thinking_now() is False + assert finish_thinking_now() is False + finally: + clear_finish_key_target() + err = capsys.readouterr().err + assert err.count("finish-thinking isn't available on the MTP path") == 1 + + def test_factory_wires_eos_ids_and_floor(): class _EosTok(_FakeTok): eos_token_ids = {77, 33} From 6cd966149c58515cfa4cd5122b8e0ddd7f0a082c Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:13:04 -0700 Subject: [PATCH 4/7] fix(chat): warn that --thinking-budget and /thinking-budget don't reach MTP replies --- gmlx/chat.py | 5 +++++ gmlx/cli.py | 1 + 2 files changed, 6 insertions(+) diff --git a/gmlx/chat.py b/gmlx/chat.py index b570666..79ad143 100644 --- a/gmlx/chat.py +++ b/gmlx/chat.py @@ -1172,6 +1172,11 @@ def _slash_thinking_budget(cmd, arg, state): print(f"[chat] /thinking-budget needs an int or 'off', got {arg!r}") return None print(f"[chat] thinking-budget = {arg} (next reply)") + if (state.model_info or {}).get("drafter"): + print( + "[chat] note: not applied on MTP-decoded replies " + "(restart with --no-mtp to honor it)" + ) return None if arg == "off": state.thinking_budget = None diff --git a/gmlx/cli.py b/gmlx/cli.py index 83ddf5e..b36b24d 100644 --- a/gmlx/cli.py +++ b/gmlx/cli.py @@ -566,6 +566,7 @@ def mtp_dropped_chat_flags(args) -> list[str]: # mtp_dropped_run_flags). ("--quantized-kv-start", args.quantized_kv_start != 0), ("--max-kv-size", args.max_kv_size is not None), + ("--thinking-budget", getattr(args, "thinking_budget", None) is not None), ) return [name for name, on in pairs if on] From 6e02b22c7f67e9657966c51251d86c3fad315362 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:30:38 -0700 Subject: [PATCH 5/7] fix(vlm): resolve think tokens via mlx-lm wrapper inference so Ctrl-T works on VLM chat (gemma4 channel format) --- gmlx/chat.py | 8 ++++-- gmlx/cli.py | 3 ++- gmlx/thinking_budget.py | 51 +++++++++++++++++++++++++++++++---- tests/test_thinking_budget.py | 21 +++++++++++++++ 4 files changed, 75 insertions(+), 8 deletions(-) diff --git a/gmlx/chat.py b/gmlx/chat.py index 79ad143..7658224 100644 --- a/gmlx/chat.py +++ b/gmlx/chat.py @@ -3382,13 +3382,17 @@ def _apply_session(doc: dict, name: str) -> None: from .thinking_budget import ( clear_finish_key_target, make_thinking_budget_processor, + prompt_opens_thinking, set_finish_key_target, + think_tokenizer_for, ) + vtok = think_tokenizer_for(processor) tbp = make_thinking_budget_processor( - getattr(processor, "tokenizer", processor), + vtok, None, - start_in_thinking=_opens_thinking(prompt), + start_in_thinking=isinstance(prompt, str) + and prompt_opens_thinking(prompt, tokenizer=vtok), interruptible=True, ) set_finish_key_target(tbp) diff --git a/gmlx/cli.py b/gmlx/cli.py index b36b24d..852e1ea 100644 --- a/gmlx/cli.py +++ b/gmlx/cli.py @@ -1556,11 +1556,12 @@ def _run_vlm(args) -> int: install_finish_thinking_key, make_thinking_budget_processor, set_finish_key_target, + think_tokenizer_for, ) install_finish_thinking_key() tbp = make_thinking_budget_processor( - getattr(processor, "tokenizer", processor), None, interruptible=True + think_tokenizer_for(processor), None, interruptible=True ) if tbp is not None: extra["logits_processors"] = [tbp] diff --git a/gmlx/thinking_budget.py b/gmlx/thinking_budget.py index 596429f..e634c57 100644 --- a/gmlx/thinking_budget.py +++ b/gmlx/thinking_budget.py @@ -184,6 +184,11 @@ def _thinking_token_seqs(tokenizer): if end_id is None: return (None, None) start_id = _last_token_id(tokenizer, "") + if start_id == end_id: + # Both tags collapsed to a shared trailing piece (a bare '>'): the + # vocab has no think tokens at all, and arming on that piece would + # trip on every '>' the model emits. + return (None, None) return ((start_id,) if start_id is not None else None, (end_id,)) @@ -361,6 +366,46 @@ def __call__(self, tokens, logits): return logits +def _eos_id_list(tokenizer) -> list[int]: + """Sorted EOS ids; tolerates the attr being an int (raw GGUF-built + tokenizers), a set/list (mlx-lm wrapper), or absent.""" + eos = getattr(tokenizer, "eos_token_ids", None) + if eos is None: + return [] + if isinstance(eos, int): + return [eos] + return sorted(int(t) for t in eos) + + +def think_tokenizer_for(processor): + """A think-capable tokenizer for an mlx-vlm processor. + + The raw GGUF-built HF tokenizer carries no thinking markers (gemma4's + channel format lives neither in the template text nor as single-token + tag spellings), so resolve them by wrapping with mlx-lm's + ``TokenizerWrapper`` inference - once, cached on the processor (the + inference scans the vocab). Falls back to the raw tokenizer.""" + cached = getattr(processor, "_gmlx_think_tokenizer", None) + if cached is not None: + return cached + tok = getattr(processor, "tokenizer", processor) + wrapped = tok + if not getattr(tok, "think_end_tokens", None): + try: + from mlx_lm.tokenizer_utils import TokenizerWrapper + + wrapped = TokenizerWrapper( + tok, eos_token_ids=getattr(tok, "_gguf_eos_token_ids", None) + ) + except Exception: # noqa: BLE001 - best-effort; ^T just stays off + wrapped = tok + try: + processor._gmlx_think_tokenizer = wrapped + except Exception: # noqa: BLE001 - unsettable processor: skip the cache + pass + return wrapped + + def make_thinking_budget_processor( tokenizer, budget, *, start_in_thinking=True, verbose=False, eos_floor=True, interruptible=False, @@ -440,11 +485,7 @@ def make_thinking_budget_processor( budget=budget, start_seq=start_seq, start_in_thinking=start_in_thinking, - eos_ids=sorted( - int(t) for t in (getattr(tokenizer, "eos_token_ids", None) or []) - ) - if eos_floor - else [], + eos_ids=_eos_id_list(tokenizer) if eos_floor else [], reclose_ids=reclose_ids, skip_ids=skip_ids, ) diff --git a/tests/test_thinking_budget.py b/tests/test_thinking_budget.py index f986657..56ad31f 100644 --- a/tests/test_thinking_budget.py +++ b/tests/test_thinking_budget.py @@ -310,6 +310,27 @@ def test_finish_key_target_round_trip(): assert finish_thinking_now() is False +def test_fallback_rejects_shared_trailing_piece(): + # A vocab with no think tokens at all: both tag spellings BPE down to a + # shared trailing '>' piece. Arming on that would trip on every '>'. + class _NoThinkTok: + def encode(self, text, add_special_tokens=True): + return [7, 8, 42] if "think" in text else [] + + assert _thinking_token_seqs(_NoThinkTok()) == (None, None) + assert make_thinking_budget_processor( + _NoThinkTok(), None, interruptible=True) is None + + +def test_factory_tolerates_int_eos_attr(): + # Raw GGUF-built tokenizers carry eos_token_ids as a bare int. + class _IntEosTok(_FakeTok): + eos_token_ids = 77 + + p = make_thinking_budget_processor(_IntEosTok(), 4) + assert p.eos_ids == [77] + + def test_finish_key_unsupported_notes_once(capsys): from gmlx.thinking_budget import ( FinishKeyUnsupported, From 07717320dd1711e7b82c6f0b83c261abf6f3d883 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:30:38 -0700 Subject: [PATCH 6/7] fix(chat): default REPL MTP to the owned engine (clean mid-round close), GMLX_OWNED_ROUND=0 opts back --- gmlx/generation.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/gmlx/generation.py b/gmlx/generation.py index ca43197..4b876db 100644 --- a/gmlx/generation.py +++ b/gmlx/generation.py @@ -9,6 +9,7 @@ from __future__ import annotations +import os import sys import time @@ -1131,9 +1132,9 @@ def _stream_generate_speculative_owned( min_p: float = 0.05, draft_block_size: int | None = None, ): - """Owned-engine sibling of :func:`stream_generate_speculative` for - drafters with ``requires_owned_engine`` (deepseek_v4). Same - ``_MTPStreamResponse`` surface, but drives ``stream_speculative`` (which + """Owned-engine body of :func:`stream_generate_speculative` (the REPL + default; see the routing note there). Same ``_MTPStreamResponse`` + surface as the stock round, but drives ``stream_speculative`` (which does its own chunked prefill through the persistent ``prompt_cache``).""" from mlx_lm.sample_utils import make_sampler @@ -1246,10 +1247,15 @@ def _stream_generate_speculative( stop/penalty/bias hooks (same surface as :func:`generate_speculative`); the REPL's other ``/`` sampling controls don't reach this path. """ - # Stochastic acceptance lives in the owned walk, so sampled runs route - # there when it's requested; greedy stays on the stock round. + # The owned walk is the default for the REPL (as it already is for serve + # and bench): unlike mlx-vlm's stock round it rolls the cache back to + # exactly the delivered tokens when the consumer closes mid-round (Esc / + # EOS / a stop string), so the persistent chat cache stays clean for the + # next turn. GMLX_OWNED_ROUND=0 opts back to the stock round, except for + # drafters whose contract demands the owned engine. from .speculative import use_owned_engine - if use_owned_engine(drafter, temp): + if (os.environ.get("GMLX_OWNED_ROUND") != "0" + or use_owned_engine(drafter, temp)): yield from _stream_generate_speculative_owned( model, drafter, tokenizer, prompt, prompt_cache=prompt_cache, max_tokens=max_tokens, temp=temp, top_p=top_p, top_k=top_k, From 76bfa54aab5304eb70c57adfd17b0eff5792a258 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:30:50 -0700 Subject: [PATCH 7/7] docs: changelog entry for owned-engine chat MTP default --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 516d7a1..c6de72f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,11 +12,15 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). block early (as if the thinking budget had just run out) so the answer starts now. +### Changed + ### Fixed - CLI polish: `--help`/`--help-all` now win even after a value-taking flag, the short `-h` pages gained the typical sampling/config/streaming flags, tab completion offers flag values (choices, themes, profiles). +- Chat: cancelling or stopping mid-round no longer leaves rejected draft KV + in the persistent chat cache. ## [0.2.0] - 2026-08-02