[snehalv2002]

MLA: optional sigmoid output gate and NoPE (skip-RoPE) support

What

Two opt-in changes to the shared MLA layer (attention_mla.py), both needed by Kimi-K3 full-attention layers:

1. Output gate. New constructor arg MLA(mla_use_output_gate=False). When it is True, the layer creates g_proj (embed_attn → q_heads × v_head_dim) and computes out = out_proj(attn * sigmoid(g_proj(inputs_q))). The gate gets the same logical sharding as the attention output. As in the HF reference, the sigmoid runs in the activation dtype. The constructor arg is the only switch: there is no config field, which is the same pattern as is_nope_layer. The Kimi-K3 decoder layer (PR8) passes mla_use_output_gate=True.
2. NoPE. MLA now honours the existing is_nope_layer flag. It skips apply_rotary_embedding on the q_pe / k_pe slices in mla_query_projection and mla_kv_projection. The slices are still projected and still take part in attention. Kimi-K3 sets HF mla_use_nope: true. Before this PR, MLA silently ignored the flag.

Both features are off by default. No existing MLA caller (DeepSeek, Kimi-K2, …) passes either flag, so their behaviour and params are unchanged.

Tests

New tests/unit/kimi_k3_mla_gate_nope_test.py:

• Hermetic (runs in CI, no torch needed):
  • The gate is off by default.
  • With zero g_proj weights the output is exactly 0.5× the ungated output.
  • With is_nope_layer=True at non-zero positions, the output doesn't depend on position and equals the RoPE layer's output at position 0. RoPE at non-zero positions gives a different output.
• HF parity (small dims): gated NoPE MLA vs HF KimiMLAAttention at positions arange(seq), with weights loaded through PR1's kimi-k3 param mapping. Every expected key must convert. A negative check confirms that the RoPE variant diverges. This test skips unless torch, transformers, einops and the untracked kimi-k3-hf-reference/ checkout are all present, and it restores sys.path / sys.modules afterwards.

JAX_PLATFORMS=cpu PYTHONPATH=./src:. python -m pytest -q tests/unit/kimi_k3_mla_gate_nope_test.py
# 4 passed with the reference; 3 passed, 1 skipped without it
JAX_PLATFORMS=cpu PYTHONPATH=./src:. python -m pytest -q tests/unit/attention_test.py -k "MLA and not tpu"
# 33 passed

Size: 36 src / 289 test lines (vs. its stack base).

Known limitations

• The gate is tested in train mode only. Prefill/decode share the same code path.
• The gated activation has no remat checkpoint name of its own. It sits between the existing attention_out and out_proj names.

Stack: Kimi-K3 series, this is PR4; 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, adding the model configuration, parameter mapping, transformation hooks, and corresponding unit tests. The reviewer feedback highlights several critical issues: in param_mapping.py, the config parameter needs to be safely converted to a dictionary to prevent AttributeError when a PretrainedConfig object is passed. In the unit tests, there is duplicated code for convert_pytorch_module_to_maxtext_params that should be imported from the utility module, and the tests should pass self.hf_config["text_config"] instead of self.hf_config to avoid falling back to default configuration values.

[gemini-code-assist]

The config parameter can be either a dictionary or a Hugging Face PretrainedConfig object (such as KimiK3Config). Since PretrainedConfig does not have a .get() method, calling config.get(...) directly will raise an AttributeError during actual checkpoint conversion. Converting config to a dictionary at the beginning of the function ensures compatibility with both dictionary and configuration object inputs.

  if not isinstance(config, dict):\n    config = config.to_dict() if hasattr(config, "to_dict") else getattr(config, "__dict__", {})\n  num_layers = config.get("num_hidden_layers", 93)

[gemini-code-assist]

The config parameter can be a Hugging Face PretrainedConfig object which does not support the .get() method. Converting config to a dictionary at the beginning of the function ensures compatibility with both dictionary and configuration object inputs.

  if not isinstance(config, dict):\n    config = config.to_dict() if hasattr(config, "to_dict") else getattr(config, "__dict__", {})\n  num_main_layers = config.get("num_hidden_layers", 93)

[gemini-code-assist]

The config parameter can be a Hugging Face PretrainedConfig object which does not support the .get() method. Converting config to a dictionary at the beginning of the function ensures compatibility with both dictionary and configuration object inputs.

  if not isinstance(config, dict):\n    config = config.to_dict() if hasattr(config, "to_dict") else getattr(config, "__dict__", {})\n  hidden_size = config.get("hidden_size", 7168)

[gemini-code-assist]

The convert_pytorch_module_to_maxtext_params function is duplicated here and in tests/utils/kimi_k3_conversion_utils.py. To improve maintainability and avoid code duplication, import the canonical version from the utility module instead.

from tests.utils.kimi_k3_conversion_utils import convert_pytorch_module_to_maxtext_params

[gemini-code-assist]

The PARAM_MAPPING functions expect the flat configuration dictionary (or configuration object). Passing self.hf_config directly (which contains "text_config" as a nested key) causes the mapping functions to fall back to default values (e.g., defaulting num_hidden_layers to 93 instead of 4). Pass self.hf_config["text_config"] instead.

    converted_params = convert_pytorch_module_to_maxtext_params(\n        pt_mla.state_dict(),\n        self.hf_config["text_config"],\n        None,\n        prefix_filter="params-decoder-layers_3-self_attention",\n    )

[gemini-code-assist]

The HOOK_FNS factory expects the flat configuration dictionary (or configuration object). Passing self.hf_config directly (which contains "text_config" as a nested key) causes the hook functions to fall back to default values. Pass self.hf_config["text_config"] instead.

Suggested change

	hooks = HOOK_FNS["kimi-k3"](self.hf_config, None, scan_layers=False, saving_to_hf=False)
	hooks = HOOK_FNS["kimi-k3"](self.hf_config["text_config"], None, scan_layers=False, saving_to_hf=False)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!