Merge kimi-k3-stack/pr3..pr6 (integration base for PR8)

Author: snehalv2002

src/maxtext/checkpoint_conversion/utils/param_mapping.py

@@ -4856,8 +4856,9 @@ def KIMI_K3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
       mapping.update(
           {
               f"{prefix}-mlp-routed_experts-gate-kernel": f"language_model.model.layers.{i}.block_sparse_moe.gate.weight",
-              # moe.RoutedMoE's router bias param is `gate/bias`.
-              f"{prefix}-mlp-routed_experts-gate-bias": (
+              # The router bias is `gate/bias`; for kimi-k3 it is a `moe.MoEBiasVar`, so (as for
+              # DeepSeek) it lives in the `MoEBiasVar` collection.
+              f"{prefix.replace('params-', 'MoEBiasVar-')}-mlp-routed_experts-gate-bias": (
                   f"language_model.model.layers.{i}.block_sparse_moe.gate.e_score_correction_bias"
               ),
               f"{prefix}-mlp-routed_expert_down_proj-kernel": (

src/maxtext/configs/base.yml

@@ -308,6 +308,7 @@ warmup_programs_in_init: false # before init_stop/run_start, run each precompile
 warmup_skip_train_dropless: false # requires warmup_programs_in_init=true; do not warm the train_dropless program (its warmup holds a second copy of the train state in HBM and can OOM). Use only when that warmup OOMs; the first dropless replay then pays the program's first load.
 log_step_diagnostics: false # append grad/param norms, max |grad|, routed-bias checksum, nonzero bias updates and the overflow flag to every training step log line
 moe_expert_input_dim: -1 # feature dimension of the tokens entering the MoE expert blocks.
+moe_router_input_dim: -1 # feature dimension of the router (gate) input, passed as `gate_inputs`; defaults to the expert input dim.
 base_moe_mlp_dim: -1 # intermediate dimension at MoE layer. For a fully MoE model, base_mlp_dim must be equal to base_moe_mlp_dim.
 load_balance_loss_weight: 0.0 # weight for the load balance loss
 moe_use_megatron_seq_aux_loss: false # when true, sigmoid and sqrtsoftplus routers use Megatron-LM's sequence-wise aux loss summed over MoE layers; false (default) uses the Switch-style loss averaged over MoE layers.

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

@@ -46,7 +46,8 @@ attn_res_block_size: 12
 
 # MLP and MoE Specifications
 first_num_dense_layers: 1
-mlp_activations: ["situ", "linear"]
+# SiTU on both branches: situ_gate(wi_0 x) * situ_linear(wi_1 x).
+mlp_activations: ["situ_gate", "situ_linear"]
 # `base_mlp_dim` is the layer-0 dense MLP width (HF `intermediate_size`); the MoE
 # shared expert uses `shared_intermediate_size` = moe_intermediate_size * num_shared_experts.
 base_mlp_dim: 33792
@@ -63,6 +64,11 @@ num_shared_experts: 2
 routed_scaling_factor: 1.0
 routed_score_func: "sigmoid"
 routed_bias: true
+# Routed experts run through moe.RoutedMoE on the 3584-wide latent, while the router reads the full
+# 7168-wide hidden state (passed as `gate_inputs`); router math is fp32 as in the reference.
+moe_expert_input_dim: 3584
+moe_router_input_dim: 7168
+float32_gate_logits: true
 
 # Sequence Length & RoPE
 # Kimi-K3 uses NoPE for its MLA layers (HF `mla_use_nope: true`) and ships no

src/maxtext/configs/types.py

@@ -1273,6 +1273,14 @@ class MoEGeneral(BaseModel):
       -1,
       description="Dimension of tokens entering the MoE layer. If < 0, defaults to emb_dim.",
   )
+  moe_router_input_dim: int = Field(
+      -1,
+      description=(
+          "Dimension of the router (gate) input when it differs from the expert input, e.g. a latent MoE that"
+          " routes on the full hidden state but runs experts on a down-projected latent (pass the router input"
+          " as `gate_inputs`). If < 0, defaults to the expert input dim."
+      ),
+  )
   base_moe_mlp_dim: int = Field(-1, description="Intermediate dimension at MoE layer.")
   padded_base_moe_mlp_dim: Optional[int] = Field(
       None,

src/maxtext/layers/attention_mla.py

@@ -522,13 +522,18 @@ def __init__(
       rope_factor: float = 40.0,  # rotary embedding factor
       name: str | None = None,
       rngs: Optional[nnx.Rngs] = None,
+      mla_use_output_gate: bool = False,
   ):
     """Initializes the MLA module.
 
     Args:
       config: The model configuration.
       ... and other configuration parameters for MLA attention.
+      is_nope_layer: If True, the rotary transform is skipped for the `qk_rope_head_dim`
+        slices of the query and key (the slices are still projected and attended over).
       rngs: The random number generators for initialization, passed by the nnx.to_linen wrapper.
+      mla_use_output_gate: If True, adds a per-head sigmoid output gate before the output
+        projection: `out = out_proj(attn * sigmoid(g_proj(inputs_q)))`.
     """
     base_kv_cache = config.attention != "paged" and config.mla_naive_kvcache
 
@@ -543,6 +548,7 @@ def __init__(
     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
 
     self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
 
@@ -765,6 +771,23 @@ def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> No
       mscale = 0.1 * self.mscale * math.log(self.rope_factor) + 1.0
       self.softmax_scale = self.softmax_scale * mscale * mscale
 
+    if self.use_output_gate:
+      self.g_proj = DenseGeneral(
+          in_features_shape=self.config.emb_dim,
+          out_features_shape=(self.num_query_heads, self.v_head_dim),
+          axis=-1,
+          kernel_init=self.kernel_init,
+          kernel_axes=("embed_attn", "q_heads", "kv"),
+          dtype=self.dtype,
+          weight_dtype=self.weight_dtype,
+          quant=self.quant,
+          matmul_precision=self.config.matmul_precision,
+          shard_mode=self.config.shard_mode,
+          rngs=self.rngs,
+      )
+    else:
+      self.g_proj = None
+
     self.out = self.init_out_w(output_dim=inputs_q_shape[-1])
 
   def _attention_op_num_kv_heads(self) -> int:
@@ -866,7 +889,10 @@ def mla_query_projection(
     # Partial RoPE: Split into non-positional and rotary parts.
     # last dimension: qk_nope_head_dim, qk_rope_head_dim
     q_nope = self._maybe_shard_with_logical(q_nope, query_logical_name)
-    q_pe = self.apply_rotary_embedding(q_pe, inputs_positions=inputs_positions)
+    # NoPE layers (e.g. Kimi-K3) still carry a `qk_rope_head_dim` slice through the
+    # projections, but never apply the rotary transform to it.
+    if not self.is_nope_layer:
+      q_pe = self.apply_rotary_embedding(q_pe, inputs_positions=inputs_positions)
     q_pe = self._maybe_shard_with_logical(q_pe, query_logical_name)
     # Query projection is scaled by self.softmax_scale to be consistent MaxText implementation.
     # DeepSeek v3 was doing it in attention score computation.
@@ -1024,9 +1050,10 @@ def mla_kv_projection(self, inputs: Array, inputs_positions: Array, decoder_segm
     low_rank_main, low_rank_rope = jnp.split(low_rank, [self.kv_lora_rank], axis=-1)
     low_rank_main = self.kv_norm(low_rank_main)
     low_rank_main = checkpoint_name(low_rank_main, "mla_kv")
-    # Apply rotary embedding to key_rope.
+    # Apply rotary embedding to key_rope (skipped entirely for NoPE layers).
     key_rope = jnp.expand_dims(low_rank_rope, axis=2)
-    key_rope = self.apply_rotary_embedding(key_rope, inputs_positions=inputs_positions)
+    if not self.is_nope_layer:
+      key_rope = self.apply_rotary_embedding(key_rope, inputs_positions=inputs_positions)
 
     key, value = self._build_key_value(low_rank_main, key_rope, model_mode)
     cached_values = [None, None]
@@ -1300,6 +1327,12 @@ def __call__(
     out = self._maybe_shard_with_logical(out, self.out_axis_names)
     out = jax.ad_checkpoint.checkpoint_name(out, "attention_out")
 
+    if self.use_output_gate:
+      # Matches the HF reference: the sigmoid is computed in the activation dtype.
+      g = jax.nn.sigmoid(self.g_proj(inputs_q))
+      g = self._maybe_shard_with_logical(g, self.out_axis_names)
+      out = out * g
+
     out_sharding = create_sharding(self.mesh, out_logical_name)
     out = self.out_projection(out, out_sharding=out_sharding)
     out = checkpoint_name(out, "out_proj")

src/maxtext/layers/attn_res.py

@@ -0,0 +1,76 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Attention Residuals (AttnRes) softmax pooling for Kimi-K3.
+
+Kimi-K3 replaces the plain residual stream with a learned softmax pooling over
+a list of block residuals plus the running prefix sum. This module provides only
+the pooling primitive; the decoder layer owns the per-site weights (RMSNorm scale
+and a [D, 1] projection) and decides when a new block residual is committed.
+
+  V = [block_residual_0, ..., block_residual_{K-1}, prefix_sum]
+  k = V * rsqrt(mean(V**2, -1) + eps)          (RMSNorm without scale, fp32)
+  scores = sum(k * (norm_weight * proj_weight), -1)
+  out = sum(softmax(scores) * V, over K+1)
+"""
+
+from typing import Optional
+
+import jax
+from jax import lax
+import jax.numpy as jnp
+
+
+def _apply_attn_res(
+    prefix_sum: jnp.ndarray,
+    block_residual: Optional[jnp.ndarray],
+    proj_weight: jnp.ndarray,
+    norm_weight: jnp.ndarray,
+    epsilon: float = 1e-5,
+    num_blocks: None = None,
+) -> jnp.ndarray:
+  """Softmax-pools `block_residual` and `prefix_sum` (port of HF `_apply_attn_res`).
+
+  Args:
+    prefix_sum: Running prefix sum of shape [..., D].
+    block_residual: Committed block residuals of shape [..., K, D], or None.
+      With no blocks (None or K == 0) the pooling is the identity and
+      `prefix_sum` is returned unchanged.
+    proj_weight: Projection kernel of shape [D, 1], [1, D] or [D].
+    norm_weight: RMSNorm scale of shape [D].
+    epsilon: RMSNorm variance epsilon.
+    num_blocks: Must be None. Accepted only for call-site compatibility; the
+      blocks are always taken from the (dynamically grown) `block_residual`.
+
+  Returns:
+    Pooled hidden states of shape [..., D] in `prefix_sum.dtype`.
+  """
+  if num_blocks is not None:
+    raise NotImplementedError("Static-padded AttnRes (num_blocks != None) is not supported.")
+  # No blocks: softmax over a single slot is 1, so the result is prefix_sum itself.
+  if block_residual is None or block_residual.shape[-2] == 0:
+    return prefix_sum
+
+  orig_dtype = prefix_sum.dtype
+  v = jnp.concatenate([block_residual, jnp.expand_dims(prefix_sum, axis=-2)], axis=-2)  # [..., K+1, D]
+
+  # fp32 materialization of V is deliberate, to match the HF reference numerics.
+  v_float = v.astype(jnp.float32)
+  k = v_float * lax.rsqrt(jnp.mean(jnp.square(v_float), axis=-1, keepdims=True) + epsilon)
+
+  score_weight = norm_weight.astype(jnp.float32).reshape(-1) * proj_weight.astype(jnp.float32).reshape(-1)  # [D]
+  scores = jnp.sum(k * score_weight, axis=-1)  # [..., K+1]
+  probs = jax.nn.softmax(scores, axis=-1)
+  hidden_states = jnp.sum(jnp.expand_dims(probs, axis=-1) * v_float, axis=-2)
+  return hidden_states.astype(orig_dtype)