Skip to content

feat: add the K3 XTML chat core and REPL, without the bundled tokenizer - #55

Open
Avicennasis wants to merge 7 commits into
FareedKhan-dev:mainfrom
Avicennasis:pr/xtml-chat
Open

feat: add the K3 XTML chat core and REPL, without the bundled tokenizer#55
Avicennasis wants to merge 7 commits into
FareedKhan-dev:mainfrom
Avicennasis:pr/xtml-chat

Conversation

@Avicennasis

Copy link
Copy Markdown

What this changes

--chat: a terminal REPL that renders the checkpoint's own XTML chat format exactly -- message envelopes, the thinking_effort system preamble, the assistant generation prompt, reasoning_content on prior assistant turns -- with JSONL history load/save and a parser for one assistant completion. ROADMAP item 6 ("Chat template"). Greedy by default; --temperature / --top-p / --seed opt in to sampling.

Why

Without the template the engine completes the prompt instead of answering it. I hit this on the released checkpoint: a hand-built XTML prompt fed through --ids gave a correct, well-formed answer, tok.py chat (which looks for a chat_template key the checkpoint does not ship) silently gave a document continuation. The format is not guessable from ChatML; it is defined by encoding_k3.py in the checkpoint, and this C core reproduces it.

Provenance, and the change from #20

The chat core and REPL are @BlakeEvans22's, from #20, which was closed for one stated reason: it bundled tiktoken.model, tokenizer_config.json, encoding_k3.py and the Kimi license as a test fixture. This branch carries his three commits with authorship preserved (-x trailers), cherry-picked via @ScriptedAlchemy's fork where they lived on, and drops the bundle entirely -- no commit in this PR ever adds those files. test_chat now takes TOK_FILES, exactly like test_tok, and prints NOT RUN when the vocabulary is absent; under CMake it goes through cmake/run_chat_test.cmake, a copy of run_tok_test.cmake, so ctest keeps counting it and reports a skip rather than dropping it.

My two additions on top: the fixture removal and its wiring, and a commit making chat greedy by default -- the REPL as written sampled by default with --greedy to opt out, which inverts ROADMAP item 5 and would have made --chat the only path whose output cannot be reproduced. That commit also restores the Sampling item his ROADMAP edit had removed. Blake should be the CONTRIBUTORS.md name for the feature.

Verification

  • make test passes (all weightless gates); the chat gate reports NOT RUN without a vocabulary
  • make test TOK_FILES=/path/to/Kimi-K3 passes with the released tokenizer: roundtrip: 94402 bytes -> 28264 ids -> 94402 bytes : PASS, CHAT TESTS PASSED
  • make portable builds with no new warnings
  • If kernels changed: n/a
  • If the config or tokenizer path changed: make tok (the roundtrip above) passes; third_party/tok.h gains tok_encode_mode() so untrusted user text is encoded with added tokens disallowed -- a user typing <|end_of_msg|> cannot forge a control id -- and fixes a signed-overflow UB in the base64 rank decoder on long lines
  • If output could change: oracle gates still match exactly; nothing outside --chat changes

Differential parity against Moonshot's own encoder, which is the check this thread could never run: rendering [system "You are helpful.", user "Hello"] with add_generation_prompt=True, thinking=True, thinking_effort="max" through the checkpoint's encoding_k3.py and encoding with its tiktoken.model gives 104 ids; test_chat --emit on the C core gives the same 104 ids, identical at every position.

Numbers, if this is a performance change

Not a performance change.

Risk

The template always emits the thinking_effort=max preamble and opens a <think> channel, as the official encoder does by default; there is no flag yet to render thinking=False, so a chat turn pays for the model's reasoning tokens before its answer. On a streamed-trunk desktop preset that is real money per turn (this engine is ~95 s/token there) and is the obvious follow-up. The REPL rebuilds the full prefill from the JSONL history on every restart; retained in-process state across turns is listed on the ROADMAP behind an equivalence gate, not done here.

BlakeEvans22 and others added 5 commits September 5, 2026 02:07
Exact rendering and tokenisation of the checkpoint's XTML chat format:
message envelopes, the assistant generation prompt, reasoning_content on
prior assistant turns, JSONL history load/save, and a parser for one
assistant completion. Untrusted user text is encoded with added tokens
disallowed, so a user typing "<|end_of_msg|>" cannot forge a control id.

Cherry-picked from BlakeEvans22's PR FareedKhan-dev#20 with one change: the bundled
copies of tiktoken.model, tokenizer_config.json, encoding_k3.py and the
Kimi license are NOT carried over. That bundle is what closed FareedKhan-dev#20. The
chat gate now takes TOK_FILES and prints NOT RUN when the vocabulary is
absent, exactly like the tokenizer gate; under CMake it goes through
cmake/run_chat_test.cmake, a copy of run_tok_test.cmake, so ctest keeps
counting it and reports a skip rather than dropping it from the list.
Wire --chat, transcript restart, system-message validation, REPL commands, and deterministic sampling into the streamed CPU inference CLI. Keep batch completion and its memory-budget behavior unchanged.

(cherry picked from commit ae91df7)
Document the official text-only chat command, transcript privacy and restart behavior, sampling/reproducibility, inherited context limits, CPU/disk memory model, test coverage, and remaining scope.

(cherry picked from commit 8214583)
…/--seed

The REPL as landed sampled by default (temperature 1.0, top-p 0.95) and
offered --greedy to switch it off. That inverts docs/ROADMAP.md item 5,
which says sampling must be opt-in and off by default, and it would have
made --chat the one path in the engine whose output cannot be reproduced.
Greedy is now the default in chat as in batch; any sampling flag turns
sampling on, and --greedy still forces argmax over them.

Also: docs/TESTING.md described a tests/fixtures/chat bundle that this
branch does not carry (see the chat-core commit), and the REPL commit's
ROADMAP edit had dropped the Sampling item entirely. Both put back.
…<|end_of_msg|>

The released checkpoint declares two end-of-sequence ids that disagree:
config.json / generation_config.json say <|end_of_msg|> (163586), while
tokenizer_config.json says [EOS] (163585). Run against the real weights,
the model closes its turn with <|close|>message<|sep|> followed by
[EOS] -- 163585 -- and never emits 163586, which the template inserts
between messages. The REPL stopped only on 163586 and the parser
accepted only a 163586 closure, so against the real model a --chat turn
would have run to --gen (4096 by default in chat: about four days at
this engine's speed) and then been rejected as malformed.

K3ChatTemplate gains eos_id ([EOS], checked to be 163585 like eom_id is
checked to be 163586); the REPL breaks on either id and prints which one
ended the turn; the parser accepts a closure followed by either. Two
test_chat cases cover the new behaviour: a turn ended by [EOS] parses,
and a closure followed by any other id is still rejected.

Verified after the change: test_chat passes against the released
tokenizer; make test green.
@Avicennasis

Copy link
Copy Markdown
Author

Pushed one more commit, 8bf2d04, after reading the REPL against what the released model actually emits.

The finding. The checkpoint declares two end-of-sequence ids that disagree: config.json / generation_config.json say <|end_of_msg|> (163586), tokenizer_config.json says [EOS] (163585). Driven on the real weights, the model closes its turn with <|close|>message<|sep|> followed by [EOS], 163585 -- it never emits 163586, which is the separator the template inserts between messages. The REPL broke only on 163586 (k3_run.c: if (next == tmpl->eom_id) break;) and the parser required a 163586 closure, so a --chat turn on the real model would have run to --gen -- 4096 by default in chat, roughly four days at this engine's speed -- and then been rejected as malformed. #20's author could not have seen this: the weights were not available to him, and nothing weightless can.

The fix. K3ChatTemplate gains eos_id ([EOS], checked to be 163585 the way eom_id is checked to be 163586); the REPL ends the turn on either id and prints which one did; the parser accepts a closure followed by either. Two test_chat cases cover it: an [EOS]-terminated turn parses, and a closure followed by any other id is still rejected. test_chat passes against the released tokenizer; make test green.

Pending. A real-checkpoint REPL smoke (one greedy turn, --preset laptop) is running now; it will record which end id the model emitted. I will post the result here either way. This is the last thing on this PR that I know of that no CI can check.

At the speeds a streamed trunk runs at, a minute or two per token, the
REPL printed nothing between reading the user line and printing the
finished turn. A six-hour run that hit its timeout left no record of
how far it had got, and a turn cut short by --gen looked identical to a
hung process. One unbuffered line per token: count against --gen, the
id, and seconds since the turn began.

Not exercised on the tiny checkpoint, whose 256-id vocabulary cannot
host the template's control ids (the REPL refuses it, correctly); the
scheduled real-checkpoint run is the first exercise of this line.
@Avicennasis

Copy link
Copy Markdown
Author

Real-model REPL smoke (dev, 2026-09-06)

bin/k3 /srv/models3/Kimi-K3 --trunk /opt/k3trunk --tok /srv/models3/Kimi-K3 --preset laptop --chat --greedy --gen 150 --incremental --history … at 7a4402b, one turn: Reply with exactly five words. (94 XTML ids).

Result: structurally complete turn, correct answer, cap reached one token before the end id.

Decoded output (150 ids, greedy; think block abbreviated):

The user wants exactly five words. Let me count carefully. … Let me go with something clean: "Here are five words exactly." - 5 words, clear.<|close|>think<|sep|><|open|>response<|sep|>Here are five words exactly.<|close|>response<|sep|><|close|>message<|sep|>

Token 150 was the <|sep|> that finishes <|close|>message<|sep|>, so the --gen cap fired before the model could emit an end id. The REPL printed the designed fallback (assistant did not complete an official turn within --gen 150; transcript is preserved) and exited 1. That is the intended cap behaviour, but it means this run did not exercise the chat: turn ended by … path — on the real model that path is covered only by the tiny-checkpoint unit tests in this PR ([EOS] 163585 and <|end_of_msg|> 163586 cases). In the earlier one-shot run on the same box the model emitted [EOS] immediately after this same closure.

Timing at laptop (3 GB trunk, single ring slot, experts on SATA): first token at 5497 s (94-id prefill), then 140 s/token. 119 of the 150 tokens were the think block — for short answers a thinking=False / --no-think template option (the follow-up named under Risk) would cut this turn to ~30 tokens, which is worth more than any I/O tuning.

Log: dev /tmp/k3chat-smoke.log.

@BlakeEvans22

BlakeEvans22 commented Sep 8, 2026

Copy link
Copy Markdown

Heads up before the workflow approval goes through — python tools lint will fail on this branch as it stands.

In tools/tok.py, line 201 unpacks cfg and nothing reads it: this PR replaces the jinja2 chat-template path, which was main()'s only reader, with the hand-built XTML segment builder. ruff.toml selects all of RUF and doesn't ignore RUF059, so ruff check tools/ exits 1:

tools/tok.py:201:10: RUF059 Unpacked variable `cfg` is never used
Found 1 error.

It's the only finding in tools/, and main is clean, so it arrives with this branch. Easy to miss — line 201 falls between two diff hunks and isn't shown in Files changed by default.

The fix is to rename cfg to _cfg on that line:

-    enc, cfg, special = load()
+    enc, _cfg, special = load()

load() keeps returning the config either way, and _cfg matches what tok_parity.py already uses. Same issue came up on #2083175e0 there was this exact rename.

Separately: the description mentions crediting me in CONTRIBUTORS.md, but the diff doesn't touch that file. Not a blocker, just flagging it since it'd be easy to lose at merge time. Thanks for picking this up and dropping the tokenizer bundle — that was the right call.

…int encoder

The checkpoint's own encoder (encoding_k3.py) takes thinking=True/False and
thinking_effort in {low, high, max}; the C template hard-coded thinking on at
max. Every real-model smoke so far spent most of its budget in the think
channel (119 of 150 tokens in the last one, at ~140 s each), so the option
that skips it is the largest speedup available without touching the engine.

- K3ChatOptions {thinking, thinking_effort}; k3_chat_render_opts and
  k3_chat_parse_assistant_opts take it, the old entry points wrap the
  defaults (thinking on, max) so existing callers and bytes are unchanged.
- thinking=False: no thinking-effort system message, stored assistant turns
  render without their think channel, the generation prompt opens
  <|open|>response<|sep|> directly, and the parser expects response-only
  output with reasoning_content NULL. All as the encoder does.
- thinking_effort: the effort message is one XTML segment (BPE merges cross
  "=max", so it cannot be split); "medium" is refused because the encoder's
  _VALID_THINKING_EFFORTS asserts on it even though its prompt text lists it.
- CLI: --no-think, --thinking-effort E; both refused outside --chat, together,
  or with a bad effort (rc 2), gated in make test with a fake model dir so
  the exit code cannot come from the loader.
- tests: no-think (37 ids), effort=low (105 ids) and two-turn no-think
  (64 ids) renders checked byte- and id-exact against the official tokenizer;
  parser cases for response-only turns, missing end id, and cross-mode
  rejection.

Reported by the real-model chat smoke on PR FareedKhan-dev#55.
@Avicennasis

Copy link
Copy Markdown
Author

Pushed dea3c5c: --no-think and --thinking-effort low|high|max, both taken from the checkpoint's own encoder rather than invented.

Why. The smoke above spent 119 of its 150 tokens in the think channel at ~140 s each. The encoder (encoding_k3.py) already exposes thinking=False and thinking_effort; the C template hard-coded thinking on at max. Skipping the channel is the largest speedup available without touching the engine, so it belongs in the chat PR.

What changed.

  • K3ChatOptions {thinking, thinking_effort} with k3_chat_render_opts / k3_chat_parse_assistant_opts. The old entry points wrap the defaults (thinking on, max), so existing callers and the 104-id default render are byte-for-byte unchanged.
  • thinking=False follows the encoder exactly: no thinking-effort system message, stored assistant turns render without their think channel, the generation prompt opens <|open|>response<|sep|> directly, and the parser expects a response-only turn with reasoning_content NULL.
  • thinking_effort keeps the effort message as one XTML segment (BPE merges cross =max, so it cannot be split). medium is refused: the encoder's _VALID_THINKING_EFFORTS asserts on it even though its prompt text lists it.
  • CLI: both flags are refused outside --chat, together, or with a bad effort (rc 2). make test gates all four misuses against a fake model dir and checks the message, so the exit code cannot come from the loader.

Verification.

  • tests/unit/test_chat.c: no-think render (37 ids), thinking_effort=low (105 ids) and a two-turn no-think render (64 ids) are checked byte- and id-exact against apply_chat_template from the checkpoint's tokenizer. Parser cases cover response-only turns, a missing end id, and each mode rejecting the other's completion. 36 checks pass.
  • Full make test passes on the branch (log summary below).
  • Prompt-size effect on the smoke above: 94 ids with thinking, 27 without (same message), so prefill drops ~70% as well as the generated tokens.

Next. Smoke #4 is queued on the same box: --preset desktop --chat --greedy --no-think --gen 60, starting after tonight's backup window. I'll post the outcome here either way.

make test TOK_FILES=/srv/models3/Kimi-K3 on dea3c5c
== ultra CLI contract ==
== op kernels ==
== streaming cache ==
CACHE TESTS PASSED
== safetensors ==
== model streaming ==
model stream parity: PASSED
== config reader ==
== config refusals ==
== tokenizer ==
== chat option contract ==
== chat template ==
CHAT TESTS PASSED
== real dimensions ==
SCALE TEST PASSED
== trunk streaming ==
TRUNK TESTS PASSED
== full-model oracle ==
ALL WEIGHTLESS TESTS PASSED

@Avicennasis

Copy link
Copy Markdown
Author

Smoke #4 result: with --no-think the model completed an official turn, and it did so in 45 minutes instead of not finishing in 7.3 hours.

Same box, same checkpoint, same message ("Reply with exactly five words."), binary dea3c5c, --preset desktop --chat --greedy --no-think --gen 60, started 2026-09-08 09:29 after the nightly backup window.

user>   XTML prompt: 27 ids, generation limit 60, greedy, thinking off
chat: token 1/60 id 40 (1510 s)
…
chat: token 13/60 id 163586 (2724 s)
chat: turn ended by <|end_of_msg|> (163586)
<response>I understand your request clearly.</response>
end 10:15:31 rc=0
smoke #3 (thinking, laptop, --gen 150) smoke #4 (--no-think, desktop, --gen 60)
prompt 94 ids 27 ids
first token 5497 s 1510 s
tokens generated 150 (cap) 13
decode rate ~140 s/token ~101 s/token
turn complete no: cap hit one id before the end id yes: `<
wall clock 26 422 s 2 724 s
answer "Here are five words exactly." "I understand your request clearly."

Generated ids: 40,4400,651,3003,14383,13,163588,12092,163589,163588,2778,163589,163586 — response body, <|close|>response<|sep|>, <|close|>message<|sep|>, end id. The parser accepted it, reasoning_content is absent, and the --history JSONL holds the user and assistant records.

Two things worth noting for the review:

  • The end id this time was <|end_of_msg|> (163586). The 2026-09-05 --ids run with the same thinking=False rendering (the HTML page, different prompt) ended with [EOS] (163585). So the model uses both ids, and the parser accepts either, which is why the chat core checks for both rather than picking one.
  • Both answers are exactly five words, so --no-think did not cost correctness on this prompt. It is not a general claim about quality; it is the encoder's own thinking=False mode, and the default stays thinking=True at max.

Nothing else on this PR is pending from my side.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants