[snehalv2002]

feat(kimi-k3): model config, schema and HF->MaxText param mapping

Foundation PR for the Kimi-K3 (moonshotai/Kimi-K3, 93-layer hybrid KDA/MLA, latent MoE, attention residuals) bring-up.

What

• Registration: kimi-k3 model name (types.py, pyconfig_deprecated.py), DecoderBlockType.KIMI_LINEAR, HF id in globals.py, flattened HF text_config + KimiK3Config in hf_model_configs.py.
• kimi-k3.yml: production dims, 0-indexed full_attn_layers, scan_layers: false (the stack is unscanned-only).
• KimiK3 schema mixin: only fields later PRs read (attn_res_block_size, full_attn_layers, routed_expert_hidden_size, shared_intermediate_size, moe_intermediate_size, num_shared_experts, short_conv_kernel_size). A validator rejects moe_intermediate_size / num_shared_experts values that differ from base_moe_mlp_dim / shared_experts.
• PARAM_MAPPING / HOOK_FNS["kimi-k3"]: HF->MaxText, unscanned. scan_layers=True raises NotImplementedError. The layer layout comes from HF linear_attn_config.full_attn_layers (1-indexed) or a top-level full_attn_layers (0-indexed), with no default: a missing layout raises. A_log padding (128 -> 96) is sliced only when the tail is all zeros. No vLLM hook is registered.
• Test helpers (tests/utils/kimi_k3_conversion_utils.py, kimi_k3_reference.py): only later PRs (PR3+) use them. Key matching respects . boundaries and prefers the longest match. Hooks can receive target_shapes. The reference dir can be overridden with KIMI_K3_HF_REFERENCE_DIR. .gitignore ignores the default kimi-k3-hf-reference/.

Size: 802 src / 742 test lines (vs. its stack base).

Testing

JAX_PLATFORMS=cpu PYTHONPATH=./src python -m pytest -q tests/unit/kimi_k3_config_test.py tests/unit/configs_test.py
# 93 passed (with HF reference present)
KIMI_K3_HF_REFERENCE_DIR=/nonexistent JAX_PLATFORMS=cpu PYTHONPATH=./src python -m pytest -q tests/unit/kimi_k3_config_test.py
# 13 passed, 1 skipped

kimi_k3_config_test.py loads the yml through pyconfig and checks it against kimi_k3_dict. It also checks the alias validator, scan/layout errors, and per-layer mapping keys for all 93 layers, and runs the hooks on random tensors. All of that runs in CI. The test that compares against the released model.safetensors.index.json needs the local reference, so CI skips it.

Known limitations

• The yml cannot be used for training yet: the kimi_linear decoder lands in PR8, and PR2 fixes mlp_activations.
• Only HF->MaxText is tested. MaxText->HF (saving_to_hf=True) is untested, and hf_shape.py has no Kimi-K3 entry, so to_huggingface export is unsupported.
• tokenizer_type: tiktoken has no path yet (PR7).

Stack: merge order PR0, PR1 (this), 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, adding its configuration, parameter mapping, and transformation hooks to enable checkpoint conversion and parity testing. Feedback on the changes suggests robustly handling PretrainedConfig objects by converting them to dictionaries using .to_dict() within the parameter mapping and hook functions to avoid potential AttributeErrors. Additionally, it is recommended to remove the @requires_kimi_k3_reference decorator from the coverage test so that it can run partially in CI/CD environments using fallback configurations when the local reference checkout is unavailable.

[gemini-code-assist]

If config is passed as a Hugging Face PretrainedConfig object (e.g., KimiK3Config) rather than a dictionary, calling .get() directly on it will raise an AttributeError since PretrainedConfig does not implement a .get() method.

To make this robust and support both dictionary and PretrainedConfig inputs, convert config to a dictionary using .to_dict() if the attribute is present. This is especially useful because PretrainedConfig.to_dict() recursively converts nested configurations (like linear_attn_config) into dictionaries as well.

def _resolve_kimi_k3_full_attn_layers(config) -> set[int]:
  """Returns the 0-indexed set of full-attention (MLA) layer indices.

  Two sources are accepted, with an explicit index base for each rather than sniffing:
    * config["linear_attn_config"]["full_attn_layers"] -- the Hugging Face field. It is
      1-indexed by definition (KimiLinearConfig.is_kda_layer tests layer_idx + 1) and
      is authoritative when present.
    * config["full_attn_layers"] at the top level -- the MaxText convention used by
      kimi-k3.yml and by test configs, 0-indexed.
  Falls back to the production 93-layer layout if neither is present.
  """
  if hasattr(config, "to_dict"):
    config = config.to_dict()
  num_layers = config.get("num_hidden_layers", 93)

[snehalv2002]

Done.

[gemini-code-assist]

Similar to _resolve_kimi_k3_full_attn_layers, config can be passed as a PretrainedConfig object, which will cause config.get to raise an AttributeError. Converting config to a dictionary at the start of KIMI_K3_MAXTEXT_TO_HF_PARAM_MAPPING ensures safe key access.

Suggested change

	def KIMI_K3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
	"""Generates parameter mapping from MaxText to Hugging Face for Kimi-K3."""
	num_main_layers = config.get("num_hidden_layers", 93)
	def KIMI_K3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
	"""Generates parameter mapping from MaxText to Hugging Face for Kimi-K3."""
	if hasattr(config, "to_dict"):
	config = config.to_dict()
	num_main_layers = config.get("num_hidden_layers", 93)

[snehalv2002]

Done.

[gemini-code-assist]

Similar to the mapping function, config can be passed as a PretrainedConfig object in KIMI_K3_MAXTEXT_TO_HF_PARAM_HOOK_FN. Converting config to a dictionary at the start of the function ensures safe key access.

Suggested change

	def KIMI_K3_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
	"""Transformation hooks for Kimi-K3 parameters."""
	hidden_size = config.get("hidden_size", 7168)
	def KIMI_K3_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
	"""Transformation hooks for Kimi-K3 parameters."""
	if hasattr(config, "to_dict"):
	config = config.to_dict()
	hidden_size = config.get("hidden_size", 7168)

[snehalv2002]

Done.

[gemini-code-assist]

The test_kimi_k3_param_mapping_100_percent_coverage test is currently decorated with @requires_kimi_k3_reference, which skips the entire test if the large reference checkout is not present locally.

However, the first part of this test (verifying that the parameter mapping and hooks contain all expected keys) runs perfectly fine using the fallback self.text_config dictionary defined in setUp. The only part requiring the reference checkout is the index_file check, which is already safely guarded by if os.path.exists(self.index_file):.

Removing the @requires_kimi_k3_reference decorator from this test allows it to run in CI/CD environments, providing valuable continuous integration coverage for the Kimi-K3 parameter mapping.

Suggested change

	@requires_kimi_k3_reference
	def test_kimi_k3_param_mapping_100_percent_coverage(self):
	def test_kimi_k3_param_mapping_100_percent_coverage(self):

[snehalv2002]

No longer applies. The test was split, and key coverage now runs without the reference checkout. Only the safetensors-index check is gated.

[codecov]

Codecov Report

❌ Patch coverage is 83.64486% with 35 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...xtext/checkpoint_conversion/utils/param_mapping.py	79.24%	23 Missing and 10 partials ⚠️
src/maxtext/checkpoint_conversion/to_maxtext.py	83.33%	1 Missing ⚠️
...rc/maxtext/checkpoint_conversion/utils/hf_shape.py	97.56%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[AusarYao28]

Review notes on PR #5548 (in addition to the open PretrainedConfig.get bot threads on param_mapping.py):

Blocking

1. Passing the released moonshotai/Kimi-K3/config.json crashes PARAM_MAPPING and HOOK_FNS
   

   maxtext/src/maxtext/checkpoint_conversion/utils/param_mapping.py

   Lines 4737 to 4768 in 38df8a0

   	if config.get("num_hidden_layers") is None:
   	raise ValueError("kimi-k3 param mapping requires `num_hidden_layers` in the HF config")
   	num_layers = int(config["num_hidden_layers"])
   	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"]}
   	else:
   	raise ValueError(
   	"kimi-k3 param mapping requires `linear_attn_config.full_attn_layers` (HF, 1-indexed) "
   	"or a top-level `full_attn_layers` (0-indexed)"
   	)
   	bad = sorted(i for i in full_attn_layers if not 0 <= i < num_layers)
   	if bad:
   	raise ValueError(f"kimi-k3 full_attn_layers resolved to out-of-range 0-indexed layers {bad} for {num_layers} layers")
   	return full_attn_layers
   	def _check_kimi_k3_unscanned(scan_layers):
   	"""Kimi-K3's hybrid KDA/MLA stack is unscanned-only; emit per-layer `layers_{i}` keys."""
   	if scan_layers:
   	raise NotImplementedError("kimi-k3 checkpoint conversion only supports scan_layers=False")
   	def KIMI_K3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
   	"""Generates parameter mapping from MaxText to Hugging Face for Kimi-K3 (unscanned layers only)."""
   	_check_kimi_k3_unscanned(scan_layers)
   	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)
   	full_attn_layers = _resolve_kimi_k3_full_attn_layers(config)

   
   In the released config.json (architectures: ["KimiK3ForConditionalGeneration"], model_type: "kimi_k3"), all language-model fields (num_hidden_layers, hidden_size, num_experts, linear_attn_config, etc.) live under config["text_config"], not at top level. Right now, passing a dict loaded directly from config.json into PARAM_MAPPING["kimi-k3"] raises ValueError: kimi-k3 param mapping requires num_hidden_layers in the HF config, and KIMI_K3_MAXTEXT_TO_HF_PARAM_HOOK_FN (L5006-L5013) silently falls back to hardcoded defaults when text_config is not unwrapped. kimi_k3_config_test.py missed this because it only tests against the flat hf_model_configs.kimi_k3_dict.
   Please unwrap text_cfg = config.get("text_config") or config (after normalizing PretrainedConfig to a dict) across _resolve_kimi_k3_full_attn_layers, KIMI_K3_MAXTEXT_TO_HF_PARAM_MAPPING, and KIMI_K3_MAXTEXT_TO_HF_PARAM_HOOK_FN, and add a test that loads KIMI_K3_REFERENCE_DIR/config.json directly. Also note hf_model_configs.py:1282 declares KimiLinearForCausalLM (whose state_dict uses model.* without language_model.), whereas param_mapping.py maps language_model.model.* from KimiK3ForConditionalGeneration.

2. to_maxtext.py --save_dtype=bfloat16 truncates all 414 float32 KDA tensors to bfloat16
   

   maxtext/src/maxtext/checkpoint_conversion/utils/param_mapping.py

   Lines 4835 to 4841 in 38df8a0

   	f"{prefix}-self_attention-A_log": f"language_model.model.layers.{i}.self_attn.A_log",
   	f"{prefix}-self_attention-dt_bias": f"language_model.model.layers.{i}.self_attn.dt_bias",
   	f"{prefix}-self_attention-b_proj-kernel": f"language_model.model.layers.{i}.self_attn.b_proj.weight",
   	f"{prefix}-self_attention-g_proj-kernel": f"language_model.model.layers.{i}.self_attn.g_proj.weight",
   	f"{prefix}-self_attention-o_norm-scale": f"language_model.model.layers.{i}.self_attn.o_norm.weight",
   	f"{prefix}-self_attention-out-kernel": f"language_model.model.layers.{i}.self_attn.o_proj.weight",
   	}

   
   In the released moonshotai/Kimi-K3 safetensors (and modeling_kimi_linear.py:791-818, 834), self_attn.A_log ([128]), self_attn.dt_bias ([12288]), self_attn.{q,k,v}_conv1d.weight ([12288, 1, 4]), and self_attn.o_norm.weight ([128]) are stored in F32 across all 69 KDA layers (414 tensors total). Under to_maxtext.py --save_dtype=bfloat16, _param_save_dtype only keeps float32 for ("MoEBiasVar-", "Tid2EidVar-"), so it downcasts A_log and dt_bias to bfloat16, shifting exp(A_log) by up to 3.12e-2 on Layer 0's real weights and compounding state decay drift over long context. Please update _param_save_dtype in to_maxtext.py to keep float32 for these KDA parameters.

3. max_target_length: 1048576 in kimi-k3.yml breaks standard CLI overrides and defaults to 1M-token allocation
   

   maxtext/src/maxtext/configs/models/kimi-k3.yml

   Line 73 in 38df8a0

   max_target_length: 1048576

   
   Setting max_target_length: 1048576 in the model YAML causes pyconfig's validate_no_keys_overridden_twice to reject any CLI override max_target_length=<N> with ValueError unless override_model_config=True is also passed, and omitting the override defaults sequence length to 1M tokens and OOMs immediately. Other long-context configs (kimi-k2-1t.yml, deepseek3-671b.yml, qwen3*.yml) set max_position_embeddings in the model YAML and leave max_target_length to base.yml / CLI. Also please set float32_logits: true in kimi-k3.yml, since KimiMLAAttention (modeling_kimi_linear.py:592) computes attention softmax in float32.

Should-Fix / Minor

4. Conversion hook edge cases and test coverage (param_mapping.py & kimi_k3_conversion_utils.py)
   

   maxtext/src/maxtext/checkpoint_conversion/utils/param_mapping.py

   Lines 5052 to 5063 in 38df8a0

   	def slice_a_log(input_tensor, target_shape=None):
   	"""Drops the zero padding on the released KDA `A_log`.
   	`modeling_kimi_linear.py` declares `A_log` as `[num_heads]` (96), but the moonshotai/Kimi-K3
   	safetensors store it as `[head_dim]` (128): the 96 per-head log-decays followed by 32 exact
   	zeros. Slice to the MaxText shape and refuse to silently drop anything that is not padding.
   	HF -> MaxText only; on the way back the HF modeling code's own `[num_heads]` shape is kept.
   	"""
   	if saving_to_hf or target_shape is None:
   	return input_tensor
   	n = int(target_shape[0])
   	if input_tensor.shape[0] == n:

• slice_a_log (L5060) returns input_tensor unchanged when target_shape is None, while convert_pytorch_module_to_maxtext_params (tests/utils/kimi_k3_conversion_utils.py:53, 77-83) defaults target_shapes=None. On a safetensors slice with [128]-padded A_log, calling it without target_shapes leaves A_log at (128,) instead of (num_heads,). slice_a_log already has num_heads in scope (L5007) and can fall back to int(target_shape[0]) if target_shape is not None else int(num_heads).
• For saving_to_hf=True, "kimi-k3" is missing from HF_SHAPE (hf_shape.py), and slice_a_log does not zero-pad (96,) -> (128,) when target_shape is [128].
• kimi_k3_conversion_utils._to_numpy (L48) converts PyTorch tensors via .float().numpy(), losing bfloat16 weight quantization when testing bfloat16 modules, and kimi_k3_conversion_utils.py currently has no unit tests in this PR.

5. Alias validation edge cases (types.py:1866-1879) and KimiK3Config.__init__ nit
   

   maxtext/src/maxtext/configs/types.py

   Line 1877 in 38df8a0

   if kimi_value and generic_value is not None and kimi_value != generic_value:

• if kimi_value and generic_value is not None treats 0 as falsy, so num_shared_experts=0 silently leaves shared_experts == 2 without raising (kimi_value is not None should be used).
• shared_intermediate_size (6144 in kimi-k3.yml) is not checked against moe_intermediate_size * num_shared_experts (for example, _kimi_k3_config(base_moe_mlp_dim=64, moe_intermediate_size=64) in kimi_k3_config_test.py:90 leaves cfg.shared_intermediate_size == 6144 instead of 128).
• Nit (hf_model_configs.py:1369-1371): prefer kwargs.setdefault("max_position_embeddings", 1048576) before super().__init__(**kwargs), and consider adding "models/kimi-k3.yml" to KIMI_CONFIGS in tests/unit/configs_test.py.

[snehalv2002]

@AusarYao28 thanks for the review!

1. Done. On the KimiLinearForCausalLM vs language_model.* note: kept as-is. That's the released text_config's own value, and the language_model. prefix comes from the KimiK3ForConditionalGeneration wrapper the weights are saved from. Documented this in hf_model_configs.py.
2. Done.
3. Done.
4. Done. HF_SHAPE is registered and A_log is padded on export, but end-to-end to_huggingface export for kimi-k3 isn't tested yet.
5. Done. Kept shared_intermediate_size as a validated alias rather than deriving it, to match the existing alias-pair pattern.

[AusarYao28]

Thanks for addressing all the config, float32 KDA dtype preservation, HF_SHAPE, and slice_a_log comments!

One quick import issue to avoid breaking the downstream sync when this merges:

maxtext/src/maxtext/checkpoint_conversion/utils/hf_shape.py

Lines 1676 to 1681 in d8b471c

	# Imported here so this module stays free of the jax import that `param_mapping` pulls in.
	from maxtext.checkpoint_conversion.utils.param_mapping import ( # pylint: disable=import-outside-toplevel
	_kimi_k3_require,
	_kimi_k3_text_config,
	_resolve_kimi_k3_full_attn_layers,
	)

Every other shape helper in hf_shape.py is pure stdlib with zero maxtext.* imports, and the hf_shape library target in the internal BUILD file has deps = [] (BUILD files are excluded from automated sync). Even a function-scoped from maxtext.checkpoint_conversion.utils.param_mapping import ... inside hf_shape.py triggers a static strict-deps build failure on hf_shape during import. Keeping small local copies of _kimi_k3_text_config, _kimi_k3_require, and _resolve_kimi_k3_full_attn_layers inside hf_shape.py keeps hf_shape.py self-contained and avoids breaking the import presubmit.

[snehalv2002]

@AusarYao28 Done.