[snehalv2002]

Description

Combines the Kimi-K3 layers from PR3–PR6 into a native KimiK3DecoderLayer and wires up decoder_block: kimi_linear. With this PR, Kimi-K3 trains in MaxText (bf16).

• models/kimi_k3.py: KimiK3DecoderLayer. Each layer uses KDA or gated NoPE MLA attention (chosen by full_attn_layers) and a dense MLP (layer 0) or the latent MoE, joined by the AttnRes highway. layer_idx, is_linear_attn and is_moe are required keyword arguments.
• layers/nnx_decoders.py: an unscanned kimi_linear layer loop. After the last layer comes the output AttnRes pooling, then decoder_norm.
• Unsupported modes now raise an error instead of giving wrong results:
  • autoregressive decode raises NotImplementedError on KDA layers, because KDA state isn't cached yet;
  • scan_layers and pipeline parallelism raise ValueError;
  • configs/types.py rejects packing=True for kimi_linear, because KDA ignores segment ids;
    kimi-k3.yml therefore sets packing: false.
• tests/utils/kimi_k3_parity_utils.py: loads the HF reference with a CPU-faithful KDA kernel, and undoes its sys.path/sys.modules/kernel patches in unpatch_hf_reference(). Also adds shared layer helpers for building and running a layer. assert_logit_parity is used only by PR9.

Size: 504 src / 1540 test lines (vs. its stack base).

• MLA LoRA norms (q_norm/kv_norm) use eps 1e-6 to match HF KimiRMSNorm (new optional MLA(lora_norm_epsilon=...); default unchanged for other models).
• Context parallelism is rejected for kimi_linear (KDA's conv/recurrence can't run on load-balanced CP shards).

Tests

• kimi_k3_decoder_test.py (runs in CI, needs no torch): config pipeline, released 93-layer layout, layer composition, NoPE, forward and jax.grad finiteness, remat full vs none equivalence, param-tree contract, and all four guards.
• kimi_k3_decoder_vs_reference_test.py and kimi_k3_full_depth_vs_reference_test.py: compare against Moonshot's PyTorch code (8-layer toy model, and the 93-layer layout at toy width). They skip unless torch, transformers and einops are installed and kimi-k3-hf-reference/ is cloned, so CI doesn't run them.

JAX_PLATFORMS=cpu PYTHONPATH=./src python -m pytest -q tests/unit/kimi_k3_decoder_test.py
# 19 passed (~1.5 min)
JAX_PLATFORMS=cpu PYTHONPATH=./src python -m pytest -q tests/unit/kimi_k3_decoder_vs_reference_test.py \
  tests/unit/kimi_k3_full_depth_vs_reference_test.py
# 9 passed, 194 subtests, 0 skipped (~23 min; the full-depth tests take ~10 min each)
JAX_PLATFORMS=cpu PYTHONPATH=./src python -m pytest -q tests/unit/nnx_decoders_test.py tests/unit/decoder_layer_model_mode_test.py
# 87 passed, 2 skipped

End to end: the full 93-layer model ran 3 training steps on C4 with lm_loss < 2.5.

• v5p-256, JAX KDA: lm_loss 2.105 / 1.973 / 2.092.
• TPU7x, tokamax KDA from a later commit: lm_loss 2.061 / 2.073 / 1.995.

Known limitations

• Only train and prefill on unpacked data (packing=False) are supported.
• Layers are unscanned, so compile time grows with depth.
• No pipeline parallelism, and no autoregressive decode until KDA state caching lands.

Stack: merge order PR0, PR1, PR2..PR7, PR8 (this PR), PR9, PR10 (#6014). This PR depends on PR3–PR6; its base is kimi-k3-stack/pr8-base until they merge.

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 adds support for the Kimi-K3 model in MaxText, implementing Kimi Delta Attention (KDA), Gated Multi-Head Latent Attention (MLA), the Attention Residuals (AttnRes) highway subsystem, and Latent MoE layers, along with configuration files, parameter mappings, and unit tests. The review feedback identifies critical issues: a missing is_nope_layer parameter and initialization in the MLA layer causing runtime errors, a severe memory bottleneck in the Latent MoE layer that could trigger Out-Of-Memory (OOM) errors, and an initialization order issue in KimiK3Config where configuration values might be overwritten.

[gemini-code-assist]

The is_nope_layer parameter is passed to attention_mla.MLA in kimi_k3.py but is missing from the __init__ signature. This will raise a TypeError at runtime. Please add is_nope_layer: bool = False to the signature.

      rope_factor: float = 40.0,  # rotary embedding factor
      name: str | None = None,
      rngs: Optional[nnx.Rngs] = None,
      mla_use_output_gate: bool = False,
      is_nope_layer: bool = False,
  ):

[gemini-code-assist]

Initialize self.is_nope_layer in __init__ so it can be accessed in mla_query_projection and mla_kv_projection without raising an AttributeError.

Suggested change

	self.original_max_position_embeddings = original_max_position_embeddings
	self.mscale = mscale
	self.rope_factor = rope_factor
	self.use_output_gate = mla_use_output_gate or getattr(config, "mla_use_output_gate", False)
	self.original_max_position_embeddings = original_max_position_embeddings
	self.mscale = mscale
	self.rope_factor = rope_factor
	self.use_output_gate = mla_use_output_gate or getattr(config, "mla_use_output_gate", False)
	self.is_nope_layer = is_nope_layer

[gemini-code-assist]

Gathering the expert weights for each token individually via self._selected materializes extremely large intermediate tensors of shape [N, top_k, in_features, intermediate_dim]. For a standard sequence length of $N = 2048$ with $top_k = 16$, $in_features = 3584$, and $intermediate_dim = 3072$, storing wi_0_selected alone in bfloat16 requires over 720 GB of memory, causing immediate Out-Of-Memory (OOM) errors on any accelerator. Avoid materializing the gathered weights per token. Instead, group/sort the tokens by their selected expert indices and perform batched matrix multiplications on the grouped tokens, or leverage MaxText's existing highly-optimized RoutedMoE implementation from maxtext/layers/moe.py which handles sharded expert computation efficiently without weight replication.

[gemini-code-assist]

Setting self.max_position_embeddings before calling super().__init__(**kwargs) can lead to it being overwritten by default values defined in the base class constructor if it is not passed in kwargs. It is safer to update kwargs directly before calling super().__init__.

Suggested change

	def __init__(self, **kwargs):
	self.max_position_embeddings = kwargs.get("max_position_embeddings", 1048576)
	super().__init__(**kwargs)
	def __init__(self, **kwargs):
	kwargs["max_position_embeddings"] = kwargs.get("max_position_embeddings", 1048576)
	super().__init__(**kwargs)

[codecov]

Codecov Report

❌ Patch coverage is 96.49123% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/nnx_decoders.py	93.93%	1 Missing and 1 partial ⚠️
src/maxtext/models/kimi_k3.py	97.50%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!