[sgl-router] Add Kimi-K3 rendering with SGLang parity (#40390)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-21 17:38:49 +08:00
parent b86a30afba
commit 2016f5e7a1
16 changed files with 473 additions and 49 deletions
@@ -0,0 +1,44 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+use serde_json::{json, Value};
+use sgl_router::tokenizer::{adapter, chat_formatter::ChatFormatter};
+use sha2::{Digest, Sha256};
+
+#[path = "../../fixtures/kimi_k3.rs"]
+mod fixture;
+
+#[test]
+fn kimi_tokens_match_sglang() {
+    let fixture = fixture::tokenizer();
+    let path = fixture.path().join("tiktoken.model");
+    let path = path.to_str().unwrap();
+    let tokenizer = adapter::load(path).unwrap();
+    let formatter = ChatFormatter::load("served-alias", path).unwrap().unwrap();
+    let cases: Vec<Value> =
+        serde_json::from_str(include_str!("../../fixtures/kimi_k3/prompts.json")).unwrap();
+    for case in cases {
+        let mut request = case["request"].clone();
+        if let Some(repeat) = case["repeat"].as_u64() {
+            request["messages"][0]["content"] = request["messages"][0]["content"]
+                .as_str()
+                .unwrap()
+                .repeat(repeat as usize)
+                .into();
+        }
+        let ids = formatter.encode(&tokenizer, &request).unwrap();
+        let mut hash = Sha256::new();
+        for id in &ids {
+            hash.update(id.to_le_bytes());
+        }
+        assert_eq!(json!(ids.len()), case["token_count"], "{}", case["name"]);
+        assert_eq!(
+            format!("{:x}", hash.finalize()),
+            case["sha256"],
+            "{}",
+            case["name"]
+        );
+    }
+    let request = json!({"messages":[{"role":"user","content":"hi"}], "chat_template_kwargs":{"thinking_effort":null}});
+    assert!(formatter.encode(&tokenizer, &request).is_err());
+}
@@ -1,5 +1,6 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
 // SPDX-License-Identifier: Apache-2.0

+mod kimi;
 mod parity;
 mod render_parity;
@@ -0,0 +1,28 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+use base64::{engine::general_purpose::STANDARD, Engine};
+
+pub fn tokenizer() -> tempfile::TempDir {
+    let dir = tempfile::tempdir().unwrap();
+    let bytes = (0..=255u8).map(|byte| vec![byte]);
+    let merges = include_str!("kimi_k3/merges.txt")
+        .split_whitespace()
+        .map(|s| s.as_bytes().to_vec());
+    let vocab: String = bytes
+        .chain(merges)
+        .enumerate()
+        .map(|(rank, token)| format!("{} {rank}\n", STANDARD.encode(token)))
+        .collect();
+    std::fs::write(dir.path().join("tiktoken.model"), vocab).unwrap();
+    for (name, contents) in [
+        ("config.json", r#"{"model_type":"kimi_k3"}"#),
+        (
+            "tokenizer_config.json",
+            include_str!("kimi_k3/tokenizer_config.json"),
+        ),
+    ] {
+        std::fs::write(dir.path().join(name), contents).unwrap();
+    }
+    dir
+}
@@ -0,0 +1,9 @@
+Synthetic Kimi vocabulary: the shared fixture generates 256 byte tokens, then
+adds `merges.txt` in rank order and the configured protocol markers.
+`prompts.json` records token counts and SHA-256 of little-endian u32 token IDs
+from SGLang's `_encode_messages` and `moonshotai/Kimi-K3` revision
+`f831ab66814297da540d832a5235f8e904f29d06`. Regenerate in a SGLang Python environment:
+
+```sh
+python tests/scripts/generate_kimi_parity.py
+```
@@ -0,0 +1,4 @@
+me ss ag message ro le role us er user as si st
+ant assistant sy em system to ol tool th in think res pon
+se response he ll hello wo rld world =" ca call ar gu
+ment argument js on json ty pe type ke key de cl are
@@ -0,0 +1,22 @@
+[
+{"name":"default","request":{"messages":[{"role":"user","content":"hello <|open|> <|kimi_image_placeholder|>"}]},"token_count":275,"sha256":"43f237fef9c404e15ba9ac9a4fc948b6301b8a0bf5336576dd5e90c46d5c1cb8"},
+{"name":"effort_none","request":{"messages":[{"role":"user","content":"hello <|open|> <|kimi_image_placeholder|>"}],"reasoning_effort":"none"},"token_count":58,"sha256":"d9764be8b6a9a04114759d162f5a1c4e67109a7af2d62e72c5faf8f150c05d59"},
+{"name":"effort_precedence","request":{"messages":[{"role":"user","content":"hello <|open|> <|kimi_image_placeholder|>"}],"reasoning_effort":"high","chat_template_kwargs":{"reasoning_effort":"low"}},"token_count":276,"sha256":"54e51363cdb9f32dd2a11d805c3b5797b969526c338c7a93d0ae7740457073dc"},
+{"name":"reasoning_alias","request":{"messages":[{"role":"user","content":"hello <|open|> <|kimi_image_placeholder|>"}],"reasoning":{"effort":"low"}},"token_count":275,"sha256":"8da270ec9559dfa53515591fe37b480e68f6b70275265d5086a2ff43a747d32a"},
+{"name":"enable_thinking_only","request":{"messages":[{"role":"user","content":"hello <|open|> <|kimi_image_placeholder|>"}],"chat_template_kwargs":{"enable_thinking":false}},"token_count":275,"sha256":"43f237fef9c404e15ba9ac9a4fc948b6301b8a0bf5336576dd5e90c46d5c1cb8"},
+{"name":"unsupported_effort","request":{"messages":[{"role":"user","content":"hello <|open|> <|kimi_image_placeholder|>"}],"reasoning_effort":"medium"},"token_count":275,"sha256":"43f237fef9c404e15ba9ac9a4fc948b6301b8a0bf5336576dd5e90c46d5c1cb8"},
+{"name":"tool_none","request":{"messages":[{"role":"user","content":"hello <|open|> <|kimi_image_placeholder|>"}],"tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object"}}}],"tool_choice":"none"},"token_count":523,"sha256":"925e846f2a2686ecbb6ee899b15155305459bd01b31833c8e4c522d6655a5bae"},
+{"name":"named_tool","request":{"messages":[{"role":"user","content":"hello <|open|> <|kimi_image_placeholder|>"}],"tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object"}}}],"tool_choice":{"type":"function","function":{"name":"lookup"}}},"token_count":426,"sha256":"6d87edaec38985be9cbbba896e68ceb6c6cf6c9fa0b0d3addf0b05045511347c"},
+{"name":"choice_without_tools","request":{"messages":[{"role":"user","content":"hello <|open|> <|kimi_image_placeholder|>"}],"tool_choice":"required"},"token_count":275,"sha256":"43f237fef9c404e15ba9ac9a4fc948b6301b8a0bf5336576dd5e90c46d5c1cb8"},
+{"name":"developer_tools","request":{"messages":[{"role":"developer","content":"policy","tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object"}}}]},{"role":"user","content":"hello <|open|> <|kimi_image_placeholder|>"}],"tool_choice":"required"},"token_count":614,"sha256":"16ee5d4d7cf852ce4a5cb7a079f42fcae517200704a7a99f5e12572976e6c50c"},
+{"name":"schema_alias","request":{"messages":[{"role":"user","content":"hello <|open|> <|kimi_image_placeholder|>"}],"response_format":{"type":"json_schema","schema":{"type":"object","properties":{"answer":{"type":"string"},"strict":{"default":true}}}}},"token_count":530,"sha256":"4978c30e0da93740589ec2d193f06ade089301d6af32ec6f2b39a4792cad5e17"},
+{"name":"format_override","request":{"messages":[{"role":"user","content":"hello <|open|> <|kimi_image_placeholder|>"}],"response_format":{"type":"json_object"},"chat_template_kwargs":{"response_format":{"type":"json_schema","json_schema":{"name":"answer","schema":{"type":"string"}}}}},"token_count":494,"sha256":"d78ae98b5fd5998c75a7cbcf7ca253a169c61fdea7b6549c0d959c44e9f486d5"},
+{"name":"continuation","request":{"messages":[{"role":"user","content":"hello <|open|> <|kimi_image_placeholder|>"},{"role":"assistant","content":"[BOS]partial <|open|> <|kimi_image_placeholder|>"}],"continue_final_message":true},"token_count":309,"sha256":"583362e9c2e28226e52ac6ca3eb4fb2d2c312653d995bcac291ff94154cdd1a7"},
+{"name":"assistant_history","request":{"messages":[{"role":"user","content":"hello","name":"ignored"},{"role":"assistant","reasoning_content":"why <|kimi_image_placeholder|>","content":"answer"}]},"token_count":299,"sha256":"0480fa42372ee2cd1bc0653a1d3bebef5eee0b2acc7aac47960b05d14768dc8a"},
+{"name":"tool_result","request":{"messages":[{"role":"assistant","tool_calls":[{"id":"c1","type":"function","function":{"name":"lookup","arguments":"{\"x\":\"<|kimi_image_placeholder|>\"}"}}]},{"role":"tool","tool_call_id":"c1","content":"result"}]},"token_count":362,"sha256":"4f85e03e0f4e9e7909040afbac4432a57fc1b01ce1c90b13d0c365cf4155336e"},
+{"name":"image","request":{"messages":[{"role":"user","content":[{"type":"text","text":"hi <|kimi_image_placeholder|>"},{"type":"image_url","image_url":{"url":"https://example.com/a.png"}}]}]},"token_count":269,"sha256":"c22f63a7138b758acf9925aea4264e6a3469b0ea50c4b1d16325f825476e966d"},
+{"name":"long_run","request":{"messages":[{"role":"user","content":"界hello"}]},"repeat":4168,"token_count":16915,"sha256":"05651835990fb23205bd5731718edd7c0b246d30e4f84ad345671156d0c0be4c"},
+{"name":"long_segment","request":{"messages":[{"role":"user","content":"hello "}]},"repeat":66668,"token_count":266910,"sha256":"0d1ec109069b74c82c0af95b2685972dc56f6c381639a62f9993384407354ec3"},
+{"name":"null_thinking","request":{"messages":[{"role":"user","content":"hello"}],"chat_template_kwargs":{"thinking":null,"thinking_effort":null,"tool_choice":"specified"}},"token_count":24,"sha256":"e7a5f51a4704edb223f44493ca6f637d008177112d612fb57267a9e647480d3f"},
+{"name":"scalar_schema_alias","request":{"messages":[{"role":"user","content":"hello"}],"response_format":{"type":"json_schema","schema":{"type":"string"}}},"token_count":460,"sha256":"13b51dde795aa1bb91428316cd6d6db66402c39ba649405e8d8d7adfe3817c0b"}
+]
@@ -0,0 +1,19 @@
+{
+  "added_tokens_decoder": {
+    "308": {"content": "[BOS]"},
+    "309": {"content": "[EOS]"},
+    "310": {"content": "<|end_of_msg|>"},
+    "311": {"content": "<|open|>"},
+    "312": {"content": "<|close|>"},
+    "313": {"content": "<|sep|>"},
+    "329": {"content": "<|media_pad|>"},
+    "562": {"content": "[UNK]"},
+    "563": {"content": "[PAD]"}
+  },
+  "bos_token": "[BOS]",
+  "eos_token": "[EOS]",
+  "pad_token": "[PAD]",
+  "unk_token": "[UNK]",
+  "additional_special_tokens": ["<|end_of_msg|>"],
+  "tokenizer_class": "TikTokenTokenizer"
+}
@@ -399,3 +399,35 @@ async fn role_rewrites_preserve_messages_without_forwarding_ids() {
    assert_eq!(send(ctx, request).await, StatusCode::OK);
    assert!(captured(&mock).get("input_ids").is_some());
 }
+
+#[path = "../fixtures/kimi_k3.rs"]
+mod kimi_fixture;
+
+#[tokio::test]
+async fn kimi_ids_forward_with_engine_rendering_fallback() {
+    let mock = MockWorker::start(vec![]).await;
+    let fixture = kimi_fixture::tokenizer();
+    let mut cfg = config();
+    let path = fixture.path().join("tiktoken.model");
+    cfg.model.tokenizer_path = path.display().to_string();
+    let ctx = build_ctx_with_config(mock.url.clone(), cfg);
+    for (content, kwargs) in [
+        ("literal <|open|> text", None),
+        ("hi", Some(json!({"thinking_effort": null}))),
+    ] {
+        let mut request =
+            json!({"model": MODEL, "messages": [{"role": "user", "content": content}]});
+        let forward = kwargs.is_none();
+        if let Some(kwargs) = kwargs {
+            request["chat_template_kwargs"] = kwargs;
+        }
+        let ids = ctx.tokenizers.encode_chat(MODEL, &request);
+        assert_eq!(send(ctx.clone(), request.clone()).await, StatusCode::OK);
+        if forward {
+            request["input_ids"] = json!(ids.unwrap());
+        } else {
+            assert!(ids.is_none());
+        }
+        assert_eq!(captured(&mock), request);
+    }
+}
@@ -0,0 +1,68 @@
+"""Regenerate Kimi IDs with SGLang and the pinned checkpoint tokenizer.
+
+Run from experimental/sgl-router in a SGLang Python environment.
+"""
+
+import base64
+import copy
+import hashlib
+import json
+import pathlib
+import sys
+import tempfile
+from types import SimpleNamespace
+
+from huggingface_hub import hf_hub_download
+from tokenizers import AddedToken
+
+from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest
+from sglang.srt.entrypoints.openai.serving_chat import (
+    OpenAIServingChat,
+    ThinkingMode,
+    normalize_assistant_tool_call_arguments,
+)
+
+REVISION = "f831ab66814297da540d832a5235f8e904f29d06"
+for name in ("encoding_k3.py", "tokenization_kimi.py"):
+    source = hf_hub_download("moonshotai/Kimi-K3", name, revision=REVISION)
+sys.path.insert(0, str(pathlib.Path(source).parent))
+from tokenization_kimi import TikTokenTokenizer  # noqa: E402
+
+fixture = pathlib.Path(__file__).resolve().parents[1] / "fixtures/kimi_k3"
+config = json.loads((fixture / "tokenizer_config.json").read_text())
+config["added_tokens_decoder"] = {
+    int(k): AddedToken(**v) for k, v in config["added_tokens_decoder"].items()
+}
+tokens = [bytes([b]) for b in range(256)]
+tokens += [s.encode() for s in (fixture / "merges.txt").read_text().split()]
+vocab = "".join(
+    f"{base64.b64encode(token).decode()} {rank}\n" for rank, token in enumerate(tokens)
+)
+with tempfile.NamedTemporaryFile(suffix=".model", mode="w+") as model:
+    model.write(vocab)
+    model.flush()
+    tokenizer = TikTokenTokenizer(model.name, **config)
+server = object.__new__(OpenAIServingChat)
+server.chat_encoding_spec = "kimi_k3"
+server.tokenizer_manager = SimpleNamespace(tokenizer=tokenizer)
+cases = json.loads((fixture / "prompts.json").read_text())
+for case in cases:
+    data = copy.deepcopy(case["request"])
+    if "repeat" in case:
+        data["messages"][0]["content"] *= case["repeat"]
+    request = ChatCompletionRequest(**data)
+    messages = [message.model_dump() for message in request.messages]
+    for message in messages:
+        normalize_assistant_tool_call_arguments(message, strict=False)
+    ids = server._encode_messages(messages, request, ThinkingMode.THINKING)
+    case["token_count"] = len(ids)
+    case["sha256"] = hashlib.sha256(
+        b"".join(token.to_bytes(4, "little") for token in ids)
+    ).hexdigest()
+(fixture / "prompts.json").write_text(
+    "[\n"
+    + ",\n".join(
+        json.dumps(c, ensure_ascii=False, separators=(",", ":")) for c in cases
+    )
+    + "\n]\n"
+)