[snehalv2002]

What

Adds a Kimi-K3 tiktoken tokenizer to input_pipeline/tokenizer.py.

• KimiTikTokenTokenizer(TikTokenTokenizer) holds all the Kimi-specific logic:
  • KIMI_K3_TIKTOKEN_PAT_STR: the Han-aware pre-tokenizer regex, copied verbatim from the HF tokenization_kimi.py.
  • Special tokens come from added_tokens_decoder in the tokenizer_config.json next to the model file. Unnamed slots are <|reserved_token_{id}|>, same as HF. The config is read with tiktoken.load.read_file, the same reader load_tiktoken_bpe uses, so gs:// paths work too. If the config is missing or can't be parsed, the tokenizer logs a warning and falls back to KIMI_K3_DEFAULT_SPECIAL_TOKEN_OFFSETS, which mirrors the checkpoint's config.
  • bos/eos/pad/unk_id map to [BOS]/[EOS]/[PAD]/[UNK]. stop_tokens is {[EOS], <|end_of_msg|>, [EOT]}. Lookups are strict, so a missing token raises. pad_id is a real token id here, not -1. It flows into grain padding and ShiftData(ignored_ids=...).
• TikTokenTokenizer gets two small hooks (_build_special_tokens, _set_special_token_ids). Llama-3 behaviour doesn't change: bos/eos/pad/stop ids and encodings on tokenizer_llama3.tiktoken match before and after.
• build_tokenizer(..., "tiktoken") picks the Kimi subclass when the file name or its parent directory contains "kimi". This is only a convenience; KimiTikTokenTokenizer can also be built directly. No new TokenizerType value is added.

Why

Kimi-K3 uses a tiktoken vocabulary with a different regex and different special tokens. The Llama-3 tiktoken path tokenizes Kimi text incorrectly.

Size: 115 src / 194 test lines (vs. its stack base).

Tests

tests/unit/kimi_k3_tokenizer_test.py:

• Hermetic (tiny 256-byte BPE in a tmpdir, no downloads): Han/Latin regex split; special-token ids read from the config; warning plus fallback when the config is missing or malformed; Llama tiktoken unchanged; the path heuristic.
• Parity: compares encode/decode, special_tokens and bos/eos/pad/unk against HF AutoTokenizer(trust_remote_code=True), including chat special tokens and a copy of the files in a directory without "kimi" in its name. It is skipped unless the gitignored kimi-k3-hf-reference checkout and transformers are present, so CI runs only the hermetic tests.

JAX_PLATFORMS=cpu PYTHONPATH=./src python -m pytest -q tests/unit/kimi_k3_tokenizer_test.py
# 7 passed (with the reference checkout present)

pyink and pylint are clean (10.00/10).

Known limitations

• Choosing the Kimi tokenizer from config still relies on the path heuristic.
• TikTokenTokenizer.encode (existing code, not changed here) passes allowed_special through set(), so "all" doesn't work. Pass an explicit set instead.
• There is no chat-template support.

Stack: PR7 of the Kimi-K3 series; merge order PR0, PR1, PR2..PR7, PR8, PR9, PR10 (#6014).

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces support for the Kimi-K3 model in MaxText, including its configuration, parameter mapping, transformation hooks, and a specialized tiktoken tokenizer with Han character splitting. Feedback on the changes highlights several potential runtime errors in the parameter mapping utility where the configuration object might be a PretrainedConfig rather than a dictionary, leading to AttributeError or TypeError when calling .get(). Additionally, suggestions were made to improve the robustness of the slice_a_log function when target_shape is missing, and to safely handle potentially null added_tokens_decoder fields in the tokenizer configuration.

[gemini-code-assist]

The config parameter can be a Hugging Face PretrainedConfig object rather than a raw dictionary during actual checkpoint conversion. Calling .get() or using bracket indexing [] directly on a PretrainedConfig object will raise an AttributeError or TypeError. Use a robust helper function to safely access keys and attributes.

Suggested change

	num_layers = config.get("num_hidden_layers", 93)
	linear_cfg = config.get("linear_attn_config") or {}
	if linear_cfg.get("full_attn_layers") is not None:
	full_attn_layers = {int(x) - 1 for x in linear_cfg["full_attn_layers"]}
	elif config.get("full_attn_layers") is not None:
	full_attn_layers = {int(x) for x in config["full_attn_layers"]}
	def _get(cfg, key, default=None):
	if isinstance(cfg, dict):
	return cfg.get(key, default)
	return getattr(cfg, key, default)
	num_layers = _get(config, "num_hidden_layers", 93)
	linear_cfg = _get(config, "linear_attn_config") or {}
	linear_full_attn = _get(linear_cfg, "full_attn_layers")
	config_full_attn = _get(config, "full_attn_layers")
	if linear_full_attn is not None:
	full_attn_layers = {int(x) - 1 for x in linear_full_attn}
	elif config_full_attn is not None:
	full_attn_layers = {int(x) for x in config_full_attn}

[gemini-code-assist]

Using .get() directly on config will fail when it is a PretrainedConfig object. Use a robust helper function to safely access attributes or keys.

Suggested change

	num_main_layers = config.get("num_hidden_layers", 93)
	first_num_dense_layers = config.get("first_k_dense_replace", 1)
	num_experts = config.get("num_experts", 896)
	def _get(cfg, key, default=None):
	if isinstance(cfg, dict):
	return cfg.get(key, default)
	return getattr(cfg, key, default)
	num_main_layers = _get(config, "num_hidden_layers", 93)
	first_num_dense_layers = _get(config, "first_k_dense_replace", 1)
	num_experts = _get(config, "num_experts", 896)

[gemini-code-assist]

Using .get() directly on config will fail when it is a PretrainedConfig object. Use a robust helper function to safely access attributes or keys.

  def _get(cfg, key, default=None):
    if isinstance(cfg, dict):
      return cfg.get(key, default)
    return getattr(cfg, key, default)

  hidden_size = _get(config, "hidden_size", 7168)
  num_heads = _get(config, "num_attention_heads", 96)
  q_lora_rank = _get(config, "q_lora_rank", 1536)
  kv_lora_rank = _get(config, "kv_lora_rank", 512)
  qk_nope_head_dim = _get(config, "qk_nope_head_dim", 128)
  qk_rope_head_dim = _get(config, "qk_rope_head_dim", 64)
  v_head_dim = _get(config, "v_head_dim", 128)

[gemini-code-assist]

Using .get() directly on config will fail when it is a PretrainedConfig object. Use the _get helper function defined above to safely access attributes or keys.

Suggested change

	num_main_layers = config.get("num_hidden_layers", 93)
	first_num_dense_layers = config.get("first_k_dense_replace", 1)
	num_main_layers = _get(config, "num_hidden_layers", 93)
	first_num_dense_layers = _get(config, "first_k_dense_replace", 1)

[gemini-code-assist]

When target_shape is None (such as during the convert_pytorch_module_to_maxtext_params test helper execution), slice_a_log currently skips slicing and returns the un-sliced tensor. Fall back to the number of attention heads from the config to ensure A_log is correctly sliced even when target_shape is not explicitly passed.

Suggested change

	if saving_to_hf or target_shape is None:
	return input_tensor
	n = int(target_shape[0])
	if saving_to_hf:
	return input_tensor
	n = int(target_shape[0]) if target_shape is not None else _get(config, "num_attention_heads", 96)

[gemini-code-assist]

If added_tokens_decoder is explicitly set to null in the JSON configuration, tok_cfg.get("added_tokens_decoder", {}) will return None, causing an AttributeError on .items(). Additionally, wrapping the loop body in a try...except block prevents a single malformed or invalid key/value from discarding all other successfully loaded added tokens.

Suggested change

	added_tokens = tok_cfg.get("added_tokens_decoder", {})
	for k_id, token_info in added_tokens.items():
	content = token_info.get("content") if isinstance(token_info, dict) else str(token_info)
	added_tokens_mapping[int(k_id)] = content
	added_tokens = tok_cfg.get("added_tokens_decoder") or {}
	for k_id, token_info in added_tokens.items():
	try:
	content = token_info.get("content") if isinstance(token_info, dict) else str(token_info)
	added_tokens_mapping[int(k_id)] = content
	except (ValueError, TypeError):
	pass

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!