User-typed <|im_end|> yields token id 163840, past declared vocab_size — encode() and __call__ disagree

#65
by paulch11 - opened

AutoTokenizer.from_pretrained("moonshotai/Kimi-K3", trust_remote_code=True) registers seven added tokens at ids 163840–163846. config.json declares text_config.vocab_size: 163840, so the last valid id is 163839 — all seven are out of range.

They are reachable from ordinary text, and the two encode paths disagree:

tok = AutoTokenizer.from_pretrained("moonshotai/Kimi-K3", trust_remote_code=True)

tok.encode("<|im_end|>")                 # [27, 91, 348, 11129, 91, 29]  -> ordinary text
tok("<|im_end|>")["input_ids"]           # [163840]                      -> out of range
tok.convert_tokens_to_ids("<|im_end|>")  # 163840

__call__ routes through tokenize() and the added-token trie, which matches these surfaces. encode() is overridden to call tiktoken directly, and tiktoken has never heard of them — they are not in the Encoding's special set. Most serving and training code uses tokenizer(...), so that is the path that gets the out-of-range id.

The seven tokens

id surface
163840 <|im_end|>
163841 <|im_user|>
163842 <|im_assistant|>
163843 <|start_header_id|>
163844 <|end_header_id|>
163845 <|im_system|>
163846 <|im_middle|>

These come from the hardcoded default additional_special_tokens list at tokenization_kimi.py:78-88 — and they are Kimi K2 names. That list holds eight entries; [EOT] is the eighth, and it already exists in K3's vocabulary, so it resolves to its real id 163593 and only the other seven land past the end.

K3's own equivalents use bracket style ([start_header_id], [end_header_id]) and live inside the vocabulary at 163590/163591, alongside the other 254 entries the tiktoken Encoding is built with (num_reserved_special_tokens = 256, tokenization_kimi.py:52). The K2 list appears to have been carried over from the previous generation's tokenizer file.

Reproduction

from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("moonshotai/Kimi-K3", trust_remote_code=True)

print(len(tok))                                    # 163840
ghosts = {s: i for s, i in tok.get_added_vocab().items() if i >= len(tok)}
print(sorted(ghosts.values()))                     # [163840 ... 163846]

print(tok("<|im_end|>")["input_ids"])              # [163840]  -- >= vocab_size
print(tok.encode("<|im_end|>"))                    # ordinary BPE ids

End-to-end, against an embedding table sized exactly as the checkpoint declares:

import torch, torch.nn as nn
emb = nn.Embedding(163840, 8)                      # text_config.vocab_size rows

emb(torch.tensor([tok("hello there")["input_ids"]]))   # OK
emb(torch.tensor([tok("<|im_end|>")["input_ids"]]))    # IndexError: index out of range in self

Environment: transformers 5.14.1, tokenizers 0.22.2, tiktoken 0.13.0, torch 2.13.0, Python 3.12.12. Tokenizer artifacts and config only — no weights downloaded, so the embedding table above is a stand-in sized from config.json rather than the model's own.

Impact

An embedding table sized from config.json has 163,840 rows (0–163839). Feeding it 163840–163846 indexes out of bounds — an IndexError on CPU, as above; a device-side assert would be the expected CUDA equivalent, though I have not run it on GPU. Since <|im_end|> is a plausible string for a user to type (it is a well-known marker from the ChatML/K2 lineage and appears in plenty of public prompt text), any deployment that tokenises untrusted input with tokenizer(...) and feeds the result to the model can be crashed by a prompt.

Two things limit this:

  • These are not K3 control tokens. They cannot forge a turn boundary, because K3's actual protocol tokens are <|end_of_msg|> (163586), <|open|> (163587), <|close|> (163588) and <|sep|> (163589), and none of the seven appears in the tiktoken special set.
  • apply_chat_template is unaffected: encoding_k3.py segments user text with allow_special=False, which reaches tiktoken as disallowed_special=() (tokenization_kimi.py:183), so the seven never appear through the chat path.

So this is a crash/DoS shape, not a prompt-injection shape.

Notes toward a fix

I did not find the exact internal reason the K2 list survives. Worth knowing: neither passing the checkpoint's own additional_special_tokens nor passing an empty list suppresses it — both still register 163840–163846:

AutoTokenizer.from_pretrained("moonshotai/Kimi-K3", trust_remote_code=True,
                              additional_special_tokens=[])
# still registers 163840-163846

So the if additional_special_tokens is None: guard at tokenization_kimi.py:78 is not the mechanism, and dropping the default list alone may not be sufficient. Two things that would each be worth checking:

  1. Remove the K2 default list, or replace it with K3's own surfaces (which are already in added_tokens_decoder at their correct ids, so nothing needs adding).
  2. Assert in __init__ that no registered added token exceeds self.model.n_vocab - 1. That turns a silent out-of-range id into a load-time error, and would have caught this before release.

A regression check worth having either way: max(tok.get_added_vocab().values()) < len(tok), and tok(s)["input_ids"] == tok.encode(s) over a corpus that includes every added-token surface — the disagreement between those two paths is the actual defect here, independent of the id range.

How this was found

While running a protocol-surface audit over K3 — it counts tokens a user can mint from plain text, and these seven showed up as registered-but-unmintable through encode(), which prompted a look at the other path. Happy to answer questions or test a fix.

One unrelated observation from the same run, offered as feedback rather than a bug: K3 refuses specials when encoding user segments in its chat path (allow_special=False throughout encoding_k3.py), so 0 of its 256 protocol tokens are forgeable through apply_chat_template. That is a genuinely good default and worth stating in the model card — most integrators will not discover it. The counterpart is that TikTokenTokenizer.encode() defaults to allow_special_tokens=True (tokenization_kimi.py:155), so a caller who tokenises user text directly gets all 256 back. Flipping that default, or documenting it next to the template guarantee, would close the gap between the safe path and the easy one.

Sign up or log in to comment