Support Kimi K3 Tokenizer

[yijingl-nvidia]

What does this PR do?

Changed the code's check on tokenizer to accept Kimi K3 tokenizer.

Kimi K3 tokenizer does not inherit from PreTrainedTokenizerFast but it is implemented in Rust using tiktoken.core.Encoding. So it should be accepted by Endpoints to run in the benchmark.

Type of change

• Bug fix
• New feature
• Documentation update
• Refactor/cleanup

Testing

• Tests added/updated
• All tests pass locally
• Manual testing completed

Checklist

• Code follows project style
• Pre-commit hooks pass
• Documentation updated (if needed)

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[yijingl-nvidia]

recheck

[codecov-commenter]

Codecov Report

❌ Patch coverage is 89.36170% with 5 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@1df1bb2). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
...utils/services/metrics_aggregator/token_metrics.py	87.80%	5 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #435   +/-   ##
=======================================
  Coverage        ?   81.71%           
=======================================
  Files           ?      146           
  Lines           ?    19373           
  Branches        ?        0           
=======================================
  Hits            ?    15831           
  Misses          ?     3542           
  Partials        ?        0           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[nv-alicheng]

Review Council — Multi-AI Code Review

Reviewed by: Claude + Code-Quality | Depth: standard

codex was unavailable in this environment — Claude + Code-Quality review. See summary comment for the full breakdown.

[nv-alicheng]

[Claude + Code-Quality] medium (performance): _TikTokenBackend.count_texts tokenizes a batch with a serial Python loop of single-string encode calls ([len(self._tokenizer.encode(text, allow_special_tokens=False)) for text in texts]). This backend runs on both the live in-process lane and inside each drain shard. The HF path (line 223-224) instead issues one encode_batch(...) rayon call that parallelizes across the whole pinned core block. Shards pin CORES_PER_WORKER = 8 cores and export RAYON_NUM_THREADS=8 on the assumption one encode call saturates ~8 cores — but single-string tiktoken encode is single-threaded, so each 8-core shard runs on ~1 core (roughly 8x per-shard throughput loss) plus per-text GIL overhead on the 1024-item live flushes. The docstring claims sharding "supplies parallelism across texts," but that's cross-process; the 8 cores within each block sit idle for tiktoken. This weakens the very "tokenizer keeps up with completions" guarantee that makes a slow backend a startup error.

Root cause: the two count_texts branches are also a copy-pasted comprehension differing only by one kwarg (DRY). The durable fix addresses both: factor the loop once (kwargs = {"allow_special_tokens": False} if self._supports_allow_special else {}), and either batch the encode or adjust the shard core-block sizing to reflect that tiktoken shards are effectively single-core, so allocated CPUs aren't wasted.

[hvagadia]

Verified against moonshotai/Kimi-K3 revision 9f62e4e9fffbd0a83ddd60e1c209d828994b3569. For 256 ordinary ~12k-character texts, the wrapper loop took 0.111s versus 0.0177s for tiktoken.Encoding.encode_batch(..., num_threads=8), with identical counts—approximately 6.3x in this single warm-process comparison.

One important caveat: replacing the wrapper loop with raw full-string encode_batch is not token-count equivalent for long inputs. For a 450,003-character input, the Kimi wrapper returned 150,003 tokens while raw encoding returned 150,002 because the wrapper applies its 400k outer split and 25k consecutive-run split.

Suggested fix: reproduce Kimi’s exact safe chunk boundaries, batch the resulting chunks with disallowed_special=(), retain an owner index for each chunk, and sum chunk lengths back into each original text. This preserves the PR’s allow_special_tokens=False policy while using all assigned cores.

Please add parity tests for ordinary text, literal special-token text, the 25k/400k boundaries, and over-boundary inputs, plus a small throughput regression benchmark.

[nv-alicheng]

[Claude] low (data-integrity): the tiktoken text path forces allow_special_tokens=False, so a special-token literal (e.g. <|...|>) in a prompt or output is counted as many literal-text tokens. The HF fast path (encode_batch(..., add_special_tokens=False), 223-224) and the chat-template path (_chat_template_token_count, tokenize=True) both keep registered special tokens recognized as one token. Result: for identical text with special-token literals, a tiktoken-backed model (Kimi K3) reports a different ISL/OSL than an HF-backed model, and diverges from the server's real count — a silent, tokenizer-family-dependent metric shift. Edge case (models rarely emit raw special-token strings) and allow_special=False is defensible for untrusted text, but document the intentional asymmetry or align the two text paths.

[nv-alicheng]

Review Council — Multi-AI Code Review

Reviewed by: Claude + Code-Quality | Depth: standard
(codex CLI unavailable in this environment — Claude + Code-Quality only.)

Well-written, well-documented PR. Findings cluster entirely in the new tiktoken backend in token_metrics.py. 2 issues posted (4 found; 1 quality nit merged into #1 as its root cause, 1 taste-nit dropped at standard depth — run --depth thorough to see all).

#	File	Line	Severity	Category	Reviewer	Summary
1	metrics_aggregator/token_metrics.py	154	medium	performance	Both	tiktoken count_texts serial per-string loop → each 8-core shard runs ~1-core (~8x throughput loss vs HF encode_batch); the DRY duplication on the same method shares the root cause
2	metrics_aggregator/token_metrics.py	157	low	data-integrity	Claude	allow_special_tokens=False makes ISL/OSL for special-token literals diverge from HF-backed models and the server — silent, tokenizer-family-dependent metric shift

Dependency: tiktoken promoted from a transitive to a direct dep, pinned ==0.13.0 (real, locked with sdist hash); aiohttp 3.14.1→3.14.3, transformers→5.5.0. Pinning hygiene ✓. pip-audit unavailable in this environment — vulnerability scan not run; suggest uv run pip-audit before merge.

Dropped at standard depth: a code-quality nit on encode_lengths (duck-typed getattr(backend, "count_texts") dispatch — a shared Protocol would remove the runtime branch). Design taste, not material debt.

⚠️ Commit hygiene: 12 commits including 4 apparent fixups (fix documentation, etc.). Consider squashing before merge.

[viraatc]

Review Council — follow-up findings (verified)

2 additional issues in the tiktoken/chat-template path, both empirically verified against the pinned transformers==5.5.0 before posting. #1 is a merge blocker: the new tokenize=True chat-template path returns a BatchEncoding whose len() is 2, silently zeroing all tool-call token metrics for every tokenizer.

[viraatc]

[Claude, orchestrator-verified] critical (bug): apply_chat_template(..., tokenize=True) in the pinned transformers 5.5.0 defaults return_dict=True and returns a BatchEncoding — a UserDict, not a dict subclass — so this isinstance(encoded, dict) never fires and len(encoded) returns the mapping's key count (2: input_ids + attention_mask).

Reproduced end-to-end against the pinned install with this exact function and a real chat tokenizer (TinyLlama-1.1B-Chat): the function returns 2 where the true token count is 28 (the pre-PR tokenize=False + tok.tokenize(rendered) path also gives 28, so this is a regression).

Downstream, no exception is raised so nothing falls back: _load_tokenizer computes _prefix_len = 2, with_assistant = 2 → _baseline = 0, and _token_count_message gets full = 2 → max(0, 2 - 2 - 0) = 0. Every assistant/tool-call message silently records 0 tokens — for every tokenizer, not just Kimi K3.

The unit test (test_chat_template_uses_native_tokenize_path) passes only because _FakeTokenizerWithTemplate.apply_chat_template returns a plain list for tokenize=True, so the real transformers return contract is never exercised.

Fix: pass return_dict=False explicitly (returns the flat id list and skips building the attention mask), or replace the guard with isinstance(encoded, Mapping). Either way, make the fake return a BatchEncoding-shaped mapping so the actual contract is pinned by the test.

[viraatc]

[Codex + Claude] medium (data-integrity): this fallback branch calls encode(text) with the wrapper's defaults — for a tiktoken-backed wrapper that inherits PreTrainedTokenizerBase.encode, that means add_special_tokens=True (verified against pinned transformers 5.5.0; e.g. a Llama-family encode("count my tokens please") gives 5 vs 4 with add_special_tokens=False), so every text is over-counted by its BOS/EOS frame.

Every other counting path in this module is deliberately no-specials: the HF fast path uses encode_batch(texts, add_special_tokens=False) (line 224) and the in-proc fallback uses tokenize() (line 431). Result: systematic ISL/OSL/TPOT inflation for exactly the "other tiktoken-backed wrappers" case this branch exists to support (Kimi K3 exposes allow_special_tokens and takes the first branch, so K3 itself is unaffected).

Fix: probe add_special_tokens in __init__ the same way allow_special_tokens is probed and build one kwargs dict — which also collapses the two duplicated comprehensions flagged in the earlier DRY comment:

# __init__
params = inspect.signature(tokenizer.encode).parameters
self._encode_kwargs = {}
if "allow_special_tokens" in params:
    self._encode_kwargs["allow_special_tokens"] = False
if "add_special_tokens" in params:
    self._encode_kwargs["add_special_tokens"] = False

# count_texts
return [len(self._tokenizer.encode(t, **self._encode_kwargs)) for t in texts]

[hvagadia]

Any value added by adding this flag? Any situation where we would want this false? We may want to avoid bloating the configs too much.
cc: @arekay-nv