feat(renderer): add Kimi K3 chat renderer with thinking/disable-thinking modes

[yoyoyocmu]

Summary

Adds a first-class renderer for Kimi K3's chat format. K3 uses a structured, self-delimiting format (<|open|> / <|sep|> / <|close|> / <|end_of_msg|>) that is unrelated to the K2.x <|im_*|> / <think> families — so no existing renderer works. In particular kimi_k27_code crashes on K3 because it assumes <think> is a single special token, whereas K3 tokenizes <think> into ['<','think','>'] and expresses reasoning as a structured think section instead.

Registers two names:

• kimi_k3 — thinking enabled (HF default); generation suffix opens a think section.
• kimi_k3_disable_thinking — HF thinking=False; generation suffix opens a response section directly.

Format (validated byte-for-byte vs the tokenizer's apply_chat_template)

system / user:
  <|open|>message role="<role>"<|sep|><content><|close|>message<|sep|><|end_of_msg|>
assistant (history + SFT target; thinking stripped -> response only):
  <|open|>message role="assistant"<|sep|><|open|>response<|sep|><content><|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|>
generation suffix (thinking):        ...<|open|>message role="assistant"<|sep|><|open|>think<|sep|>
generation suffix (disable-thinking):...<|open|>message role="assistant"<|sep|><|open|>response<|sep|>

Thinking is uniformly stripped from history (K3 ignores reasoning_content on historical assistant turns), so every message renders identically regardless of position and each conversation prefix is a strict token prefix of the next. The default concatenating build_supervised_example / build_generation_prompt are therefore exact, has_extension_property=True, and multi-turn ALL_ASSISTANT_MESSAGES SFT is well defined.

Structural markers are encoded with allowed_special="all" so they map to their single special-token ids (matching apply_chat_template).

Type of change

• New feature (non-breaking)

Testing

training/tests/unit/test_kimi_k3_renderer.py — 24 tests, all passing (run in a K3 trainer pod with the deployed tokenizer):

• Token parity vs apply_chat_template (ID-for-ID): generation prompt (user-only, system+user, single/multi-turn, with-reasoning, unicode) and supervised example (single/multi-turn, reasoning-stripped, empty assistant, unicode).
• Thinking vs disable-thinking generation suffixes.
• Weight correctness: LAST_ASSISTANT_MESSAGE vs ALL_ASSISTANT_MESSAGES; response closers trained; headers / user / system never trained.
• Invariants: prefix-extension, parse_response roundtrip + unterminated, registration/factory wiring, stop sequences.

The tokenizer is not on the HF Hub, so the suite loads a local mirror (/shared/text-models/kimi-k3, override via KIMI_K3_TOKENIZER) and skips cleanly when absent — same pattern as the DeepSeek-V4 suite. The mirrored chat-format files (tokenization_kimi.py, encoding_k3.py, tokenizer_config.json, tiktoken.model) are MD5-identical to the upstream model bucket.

Scope

Intentionally out of scope (documented in the module docstring): tool declarations / tool_calls / role="tool" results (the nested <|open|>tools...call...argument... encoding) and multimodal/image content. The text/coding SFT path does not use them and the K3 tool/vision encodings are substantially more involved. Follow-up if needed.

Code overview

flowchart TD
    A["messages[]"] --> B{"role?"}
    B -->|assistant| C["header: <|open|>message role=assistant<|sep|><|open|>response<|sep|>\noutput: content + closers + <|end_of_msg|>"]
    B -->|system/user| D["header: <|open|>message role=X<|sep|>\noutput: content + <|close|>message<|sep|><|end_of_msg|>"]
    C --> E["build_supervised_example (base):\nconcat header+output per msg,\nweight output by train_on_what"]
    D --> E
    E --> F["tokens (== apply_chat_template)\n+ per-token weights (response spans only)"]
    B -.gen prompt.-> G["_get_generation_suffix:\nthink (kimi_k3) | response (disable)"]

Loading

Checklist

• Self-reviewed the diff (3 files, +439, minimum necessary)
• Tests added and passing (24/24)
• No secrets / debug artifacts
• Linter clean

Made with Cursor

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit ac5a17d. Configure here.}

[cursor]

Think-only completions leak reasoning

Medium Severity

For the default kimi_k3 renderer, parse_response only extracts the response channel when the decoded completion contains <|open|>response<|sep|>. Sampling starts inside the prefilled think section, so truncations or stop-before-response completions still in think are returned whole as assistant content, exposing chain-of-thought to RL grading and agent parsing.

 

^{Reviewed by Cursor Bugbot for commit ac5a17d. Configure here.}