[snehalv2002]

What

Adds the Kimi-K3 SiTU-GLU activations to maxtext.layers.linears and wires them into the Kimi-K3 config:

• situ_gate(x, beta=4.0) = beta * tanh(x / beta) * sigmoid(x) (gate branch)
• situ_linear(x, beta=25.0) = beta * tanh(x / beta) (soft-capped up branch)
• Both can be selected by name through _convert_to_activation_function("situ_gate" | "situ_linear").
• kimi-k3.yml: mlp_activations: ["situ_gate", "situ_linear"]. This replaces the unresolvable ["situ", "linear"] from PR1.

Why

These match HF SituAndMul. The default betas come from Kimi-K3's config.json (activation_situ_beta=4.0, activation_situ_linear_beta=25.0) and are stored in KIMI_K3_SITU_BETA / KIMI_K3_SITU_LINEAR_BETA. Like HF, both functions upcast to fp32 and cast the result back to the input dtype, so bf16 runs don't round inside the tanh/scale chain. Other models are not affected.

Size: 70 src / 247 test lines (vs. its stack base).

How tested

JAX_PLATFORMS=cpu PYTHONPATH=./src python -m pytest -q tests/unit/kimi_k3_situ_activation_test.py
# 5 passed (reference checkout present)

tests/unit/kimi_k3_situ_activation_test.py covers:

• name registration
• golden values computed independently with Python math
• asymptotes (±25, 4/0) and small-x identity
• finite gradients at large |x|
• bf16 output dtype, checked against fp32-then-cast
• parity against HF SituAndMul in fp32 and bf16, with inputs scaled ×20 so the caps are exercised

The parity test pulls only the SituAndMul class out of the reference via ast. Because it doesn't import the reference module, it leaves no sys.modules/sys.path/transformers side effects. It skips when torch or the reference checkout is missing. pyink and pylint are clean apart from existing R0917 warnings.

Known limitations

• The betas can't be changed from the config. They're bound as function defaults.
• In CI (no reference checkout) only the hermetic tests run.
• The gate and up activations are each cast back to the compute dtype before the multiply. HF does the multiply in fp32 and casts once, which explains the small bf16 differences.

Stack: merge order PR0, PR1, PR2..PR7, PR8, PR9, PR10 (#6014) (this is PR2).

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, implementing its configuration, parameter mapping, custom SiTU activation functions, and unit tests. The feedback suggests enhancing the robustness of parameter mapping and hook functions by safely retrieving attributes from config when it is a PretrainedConfig instance instead of a dictionary, thereby preventing potential runtime crashes. Additionally, using kwargs.setdefault is recommended for cleaner initialization in KimiK3Config.

[gemini-code-assist]

Using .get() directly on config will raise an AttributeError if config is a PretrainedConfig instance (which is the case during actual checkpoint conversion) rather than a dictionary. Introducing a robust helper _get that handles both dictionaries and object attributes prevents potential runtime crashes.

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.
  """
  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")
  if linear_full_attn is not None:
    full_attn_layers = {int(x) - 1 for x in linear_full_attn}
  elif _get(config, "full_attn_layers") is not None:
    full_attn_layers = {int(x) for x in _get(config, "full_attn_layers")}
  else:
    full_attn_layers = {3, 7, 11, 15, 19, 23, 27, 31, 35, 39, 43, 47, 51, 55, 59, 63, 67, 71, 75, 79, 83, 87, 91, 92}
  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

[gemini-code-assist]

Using .get() directly on config will raise an AttributeError if config is a PretrainedConfig instance. Introducing a robust helper _get that handles both dictionaries and object attributes prevents potential runtime crashes.

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)
	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)
	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."""
	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)
	full_attn_layers = _resolve_kimi_k3_full_attn_layers(config)

[gemini-code-assist]

Using .get() directly on config will raise an AttributeError if config is a PretrainedConfig instance. Introducing a robust helper _get that handles both dictionaries and object attributes prevents potential runtime crashes.

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."""
  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)
  qk_head_dim = qk_nope_head_dim + qk_rope_head_dim

[gemini-code-assist]

Using .get() directly on config will raise an AttributeError if config is a PretrainedConfig instance. Use the _get helper defined at the top of the function to safely retrieve these attributes.

Suggested change

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

[gemini-code-assist]

Using kwargs.setdefault is a cleaner and more idiomatic way to set default configuration values before passing them to the parent class constructor. This ensures that the parent class's __init__ properly handles the attribute initialization.

Suggested change

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[AusarYao28]

Review notes on PR #5547:

Blocking

1. No combined SituAndMul parity or VJP check runs in PR CI, and bf16 tolerance 2e-2 passes broken betas
   

   maxtext/tests/unit/kimi_k3_situ_activation_test.py

   Lines 87 to 104 in a5e8c86

   	@unittest.skipIf(torch is None, "torch is not installed")
   	@requires_kimi_k3_reference
   	def test_parity_against_reference(self):
   	"""Numerical parity against HF SituAndMul, covering the soft-cap regime and bf16."""
   	situ_and_mul = _load_reference_situ_and_mul()(beta=4.0, linear_beta=25.0)
   	torch.manual_seed(42)
   	x_pt = torch.randn(4, 16, 64, dtype=torch.float32) * 20.0
   	for dtype_pt, dtype_jax, tol in ((torch.float32, jnp.float32, 1e-5), (torch.bfloat16, jnp.bfloat16, 2e-2)):
   	x = x_pt.to(dtype_pt)
   	y_pt = situ_and_mul(x).float().numpy()
   	x_jax = jnp.asarray(x.float().numpy()).astype(dtype_jax)
   	y_jax = situ_gate(x_jax[..., :32]) * situ_linear(x_jax[..., 32:])
   	self.assertEqual(y_jax.dtype, dtype_jax)
   	np.testing.assert_allclose(np.asarray(y_jax, np.float32), y_pt, rtol=tol, atol=tol)
   	if __name__ == "__main__":
   	unittest.main()

• test_parity_against_reference is skipped in GitHub PR CI (cpu-unit has neither torch nor KIMI_K3_HF_REFERENCE_DIR), and test_asymptotes_and_gradients (L74) only checks jnp.isfinite(grads) on 5 scalar points without checking VJP gradient values. Please add a hermetic float64 NumPy SituAndMul reference test covering both float32 and bfloat16 forward + VJP gradient values so CI always verifies the combined gate-times-up math.
• On L94, tol = 2e-2 (rtol=2e-2, atol=2e-2) in bfloat16 allows errors up to ~2.0 when situ_a * up reaches ~100. That silently passes a negative control with a +1% error in beta (4.04) or linear_beta (25.25, max_abs = 1.0), and hides the 3x bfloat16 rounding below. Also please read activation_situ_beta and activation_situ_linear_beta from KIMI_K3_REFERENCE_DIR/config.json in test_parity_against_reference (L91) instead of hardcoding (4.0, 25.0) on both sides.

Should-Fix

2. Three bfloat16 roundings vs HF SituAndMul's single final cast
   

   maxtext/src/maxtext/layers/linears.py

   Lines 53 to 63 in a5e8c86

   	def situ_gate(x: Array, beta: float = KIMI_K3_SITU_BETA) -> Array:
   	"""SiTU gate activation: beta * tanh(x / beta) * sigmoid(x), computed in fp32."""
   	x32 = x.astype(jnp.float32)
   	return (beta * jnp.tanh(x32 / beta) * jax.nn.sigmoid(x32)).astype(x.dtype)
   	def situ_linear(x: Array, beta: float = KIMI_K3_SITU_LINEAR_BETA) -> Array:
   	"""SiTU soft-capped linear (up-branch) activation: beta * tanh(x / beta), computed in fp32."""
   	x32 = x.astype(jnp.float32)
   	return (beta * jnp.tanh(x32 / beta)).astype(x.dtype)

   
   In modeling_kimi_linear.py:285-292, SituAndMul.forward upcasts both gate and up to float32, multiplies situ_a * up in float32, and casts back to x.dtype (bfloat16) once at the end (return (situ_a * up).to(x.dtype)). Because situ_gate (L56) and situ_linear (L62) each cast back to x.dtype before the caller multiplies them, bfloat16 callers (MlpBlock when activations_in_float32=False or fused_mlp=True, KimiDenseMLP, and RoutedMoE) round to bfloat16 three times instead of once. Across bfloat16 inputs in [-30, 30], that differs from HF SituAndMul on 26.85% of forward elements (max_abs = 0.5) and 30.09% of backward gradients (vs 0.0 mismatch when multiplied in float32 before casting to bfloat16, which lets test_parity_against_reference tighten bf16 tolerance to rtol=4e-3, atol=1e-4).

3. beta <= 0 guard and Kimi-K3 hardcoded beta defaults
   

   maxtext/src/maxtext/layers/linears.py

   Lines 49 to 63 in a5e8c86

   	KIMI_K3_SITU_BETA = 4.0
   	KIMI_K3_SITU_LINEAR_BETA = 25.0
   	def situ_gate(x: Array, beta: float = KIMI_K3_SITU_BETA) -> Array:
   	"""SiTU gate activation: beta * tanh(x / beta) * sigmoid(x), computed in fp32."""
   	x32 = x.astype(jnp.float32)
   	return (beta * jnp.tanh(x32 / beta) * jax.nn.sigmoid(x32)).astype(x.dtype)
   	def situ_linear(x: Array, beta: float = KIMI_K3_SITU_LINEAR_BETA) -> Array:
   	"""SiTU soft-capped linear (up-branch) activation: beta * tanh(x / beta), computed in fp32."""
   	x32 = x.astype(jnp.float32)
   	return (beta * jnp.tanh(x32 / beta)).astype(x.dtype)

   
   situ_gate and situ_linear are registered under generic names in _convert_to_activation_function, yet hardcode (4.0, 25.0) while KimiLinearConfig itself defaults to activation_situ_beta=1.0, activation_situ_linear_beta=None. Please guard beta > 0 (beta=0 currently produces NaN from 0.0 / 0.0) and either wire the betas from config or make the Kimi-K3 preset explicit.