Dsv4 attention opt on lb changes

docs/reference/core_concepts/moe_configuration.md

@@ -64,7 +64,12 @@ Dropping:
 
 `routed_bias`: If enabled, adds a learnable bias term to the gate logits to facilitate load balancing.
 
-`routed_bias_update_rate`: Defines the update rate to routed bias term above. Applicable only to the DeepSeek decoder block.
+`routed_bias_update_rate`: Defines the update rate to the routed bias term above. Applicable only to the DeepSeek decoder block. For DeepSeek V4, this enables a specialized, auxiliary-loss-free routing bias mechanism. This implementation utilizes a pure `nnx.Variable` (`MoEBiasVar`) instead of a standard `nnx.Param`, which completely isolates the bias update step from the global model optimizer state. The bias is updated directly at the end of the routing step to balance the token distribution mathematically across experts without compromising language modeling convergence.
+
+#### DeepSeek V4 Auxiliary-Loss-Free & Sequence-Wise Load Balancing
+MaxText implements an exact, paper-aligned version of DeepSeek V4's load balancing strategies (as specified in [the DeepSeek-V4 technical report](https://arxiv.org/html/2606.19348v1)). The architecture employs two distinct mechanisms:
+1. **Auxiliary-Loss-Free Strategy**: Handled via `MoEBiasVar`, this implementation utilizes a pure `nnx.Variable` instead of a standard `nnx.Param`. This strictly isolates the bias update step from the global model optimizer state. Unlike the Hugging Face reference implementation—which deviates from the paper by keeping the bias parameters coupled to the global optimizer—our implementation ensures the routing bias balances token distribution mathematically across experts without polluting the main model gradients.
+2. **Sequence-Wise Balance Loss**: An augmenting auxiliary loss (`load_balance_loss_weight`) applied to prevent extreme routing imbalance within individual sequences.
 
 `routed_score_func`: Defines the scoring function for the router.
 

src/maxtext/common/metric_logger.py

@@ -216,7 +216,7 @@ def _log_training_metrics(self, metrics, step):
 
     if self.config.num_experts > 1:
       moe_lb_loss = scalars.get("learning/moe_lb_loss", 0.0)
-      log_parts.append(f"moe_lb_loss: {moe_lb_loss:.3f}")
+      log_parts.append(f"moe_lb_loss: {moe_lb_loss:.6f}")
 
     if self.config.mtp_num_layers > 0:
       mtp_loss = scalars.get("learning/mtp_loss", 0.0)

src/maxtext/configs/models/deepseek4-284b.yml

@@ -51,6 +51,9 @@ shared_experts: 1
 routed_score_func: "sqrtsoftplus"
 norm_topk_prob: true
 routed_bias: true
+routed_bias_update_rate: 0.001
+load_balance_loss_weight: 0.0001
+adamw_mask: [".*gate.*bias.*"]
 routed_scaling_factor: 1.5
 
 
@@ -67,3 +70,14 @@ rope_max_timescale: 10000                  # Main RoPE theta
 compressed_rope_max_timescale: 160000      # Compressed RoPE theta
 max_position_embeddings: 1048576
 original_max_position_embeddings: 65536
+
+# --- MLA Sharding Fix ---
+# DeepSeek-V4 uses Multi-head Latent Attention with only 1 KV head.
+# We override the default sharding rules to prevent the partitioner from attempting
+# to slice the single KV head across tensor parallel chips.
+logical_axis_rules: [
+  ['activation_kv_heads', []],
+  ['kv_heads', []],
+  ['activation_kv', []],
+  ['kv', []]
+]

src/maxtext/configs/models/deepseek4-tiny.yml

@@ -11,16 +11,18 @@
 # 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.
-# Test Model config for DeepSeek-V4-Flash 284B (https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash)
 
-base_emb_dim: 4096
-base_num_query_heads: 64
+# Tiny model config for DeepSeek V4 for CPU execution and testing
+
+base_emb_dim: 64
+base_num_query_heads: 4
 base_num_kv_heads: 1
-base_num_decoder_layers: 7
-base_mlp_dim: 2048
-base_moe_mlp_dim: 2048
+base_num_decoder_layers: 43
+base_mlp_dim: 64
+base_moe_mlp_dim: 64
 vocab_size: 129280
-head_dim: 512
+head_dim: 32
+qk_rope_head_dim: 32
 
 # --- Standard Defaults ---
 enable_dropout: false
@@ -31,40 +33,38 @@ normalization_layer_epsilon: 1.0e-6
 decoder_block: "deepseek4"
 mhc_expansion_rate: 4
 first_num_hash_layers: 3
-indexer_head_dim: 128
-indexer_n_heads: 64
-indexer_topk: 512
+indexer_head_dim: 32
+indexer_n_heads: 4
+indexer_topk: 16
 
 # Note: Layers (0, 1, 2) are prefix layers as `first_num_hash_layers=3`.
-# The 6th layer (MTP module with compress_ratio=0) has been explicitly dropped for now.
-# This leaves exactly 7 layers: 3 prefix [0,0,4] + 4 scanned.
+# The 44th layer (MTP module with compress_ratio=0) has been explicitly dropped for now.
+# This leaves exactly 43 layers: 3 prefix [0,0,4] + 40 scanned.
 # `compress_ratio=0` uses sliding window attention. In this case, layer (0, 1).
-# This is a tiny version of deepseek4 with fewer layers and less experts for debugging.
-compress_ratios: [0, 0, 4, 128, 4, 128, 4]
+compress_ratios: [0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4]
 
 # --- MoE configuration ---
 mlp_activations: ["silu", "linear"]
-num_experts: 8
-num_experts_per_tok: 3
+num_experts: 16
+num_experts_per_tok: 4
 mlp_activations_limit: 10
 shared_experts: 1
 routed_score_func: "sqrtsoftplus"
 routed_bias: true
-routed_scaling_factor: 1.5
-
+routed_bias_update_rate: 0.001
+load_balance_loss_weight: 0.0001
+adamw_mask: [".*gate.*bias.*"]
 
 # --- Attention configuration ---
 attention_type: 'compressed'
-attention: 'dot_product'
-q_lora_rank: 1024
-o_groups: 8
-o_lora_rank: 1024
-sliding_window_size: 128
+q_lora_rank: 16
+o_groups: 4
+o_lora_rank: 16
+sliding_window_size: 32
 
 # --- RoPE ---
+
 rope_type: "default"
 rope_max_timescale: 10000                  # Main RoPE theta
 compressed_rope_max_timescale: 160000      # Compressed RoPE theta
-max_position_embeddings: 1048576
-original_max_position_embeddings: 65536
-
+max_position_embeddings: 4096

src/maxtext/configs/types.py

@@ -3246,8 +3246,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         raise ValueError("`local_checkpoint_period` must be > 0 for emergency checkpointing.")
     if self.moba and self.attention not in ("dot_product"):
       raise ValueError("MoBA is only supported with dot_product attention.")
-    if self.decoder_block == DecoderBlockType.DEEPSEEK4 and self.attention != "dot_product":
-      raise ValueError("DeepSeek4 decoder block currently only supports dot_product attention.")
     if self.use_indexer:
       if self.q_lora_rank == 0:
         raise NotImplementedError("Sparse indexer has not implemented for q_lora_rank = 0.")
@@ -3309,8 +3307,17 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           )
       if self.decoder_block == DecoderBlockType.GPT_OSS and not self.sparse_matmul and self.capacity_factor != -1:
         raise ValueError("GPT-OSS MoE only supports dropless (capacity_factor=-1) with dense matmul.")
-      if self.routed_bias and self.routed_bias_update_rate > 0.0 and self.decoder_block != DecoderBlockType.DEEPSEEK:
+      if (
+          self.routed_bias
+          and self.routed_bias_update_rate > 0.0
+          and self.decoder_block not in (DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4)
+      ):
         raise ValueError("Loss-free load balancing is only supported for the DeepSeek decoder block.")
+      if not self.pure_nnx and self.routed_bias and self.decoder_block == DecoderBlockType.DEEPSEEK4:
+        raise ValueError(
+            "Auxiliary-loss-free routed bias for DeepSeek V4 is only supported in pure NNX mode. "
+            "Please set pure_nnx=True or disable routed_bias."
+        )
       if self.model_name.startswith("deepseek4") and self.first_num_hash_layers > 0 and self.use_ring_of_experts:
         raise ValueError("DeepSeek V4 hash routing is currently not supported with ring of experts.")
       self.validate_ragged_buffer_factor()

src/maxtext/layers/attention_compressed.py

@@ -96,36 +96,36 @@ def csa_overlap_pooling(
 
   n_windows = chunk_kv.shape[1] // compress_rate
 
-  # Reshape flat sequence into discrete compression windows
-  # -> [batch, n_windows, compress_rate, 2 * head_dim]
-  chunk_kv = chunk_kv.reshape((batch_size, n_windows, compress_rate, 2 * head_dim))
-  chunk_gate = chunk_gate.reshape((batch_size, n_windows, compress_rate, 2 * head_dim)) + position_bias
+  # DIPAK - OPTIMIZATION
+  # Split features immediately on the 3D sequence tensor to keep memory contiguous.
+  # Pad the sequence at the beginning by exactly one compress_rate window, then truncate the end.
+  # This replaces the 4D block concatenation with a highly contiguous 1D sequence slice, improving layout.
 
-  # Split the projections into Ca and Cb components for overlapping
-  # 2x [batch, n_windows, compress_rate, head_dim]
+  # Split features immediately on the 3D sequence tensor to keep memory contiguous
   a_kv, b_kv = jnp.split(chunk_kv, 2, axis=-1)
   a_gate, b_gate = jnp.split(chunk_gate, 2, axis=-1)
 
-  # Shift Ca forward by one window to align with the next Cb
+  # Pad the sequence at the beginning by exactly one compress_rate window, then truncate the end
+  # This replaces the 4D block concatenation with a highly contiguous 1D sequence slice
   a_kv_shifted = jnp.concatenate(
-      [
-          jnp.zeros((batch_size, 1, compress_rate, head_dim), dtype=a_kv.dtype),
-          a_kv[:, :-1],
-      ],
-      axis=1,
+      [jnp.zeros((batch_size, compress_rate, head_dim), dtype=a_kv.dtype), a_kv[:, :-compress_rate]], axis=1
   )
   a_gate_shifted = jnp.concatenate(
-      [
-          jnp.full((batch_size, 1, compress_rate, head_dim), -jnp.inf, dtype=a_gate.dtype),
-          a_gate[:, :-1],
-      ],
-      axis=1,
+      [jnp.full((batch_size, compress_rate, head_dim), -jnp.inf, dtype=a_gate.dtype), a_gate[:, :-compress_rate]], axis=1
   )
 
-  # Concatenate shifted Ca and unshifted Cb to form the final overlapping window
+  # Reshape into windows
+  a_kv_windows = a_kv_shifted.reshape((batch_size, n_windows, compress_rate, head_dim))
+  b_kv_windows = b_kv.reshape((batch_size, n_windows, compress_rate, head_dim))
+
+  # Add position bias during reshape to fuse the operations
+  a_gate_windows = a_gate_shifted.reshape((batch_size, n_windows, compress_rate, head_dim)) + position_bias[:, :head_dim]
+  b_gate_windows = b_gate.reshape((batch_size, n_windows, compress_rate, head_dim)) + position_bias[:, head_dim:]
+
+  # Concatenate within the window
   # -> [batch, n_windows, 2 * compress_rate, head_dim]
-  new_kv = jnp.concatenate([a_kv_shifted, b_kv], axis=2)
-  new_gate = jnp.concatenate([a_gate_shifted, b_gate], axis=2)
+  new_kv = jnp.concatenate([a_kv_windows, b_kv_windows], axis=2)
+  new_gate = jnp.concatenate([a_gate_windows, b_gate_windows], axis=2)
 
   # Apply softmax gating and sum across the overlapping window dimension
   gate_weights = jax.nn.softmax(new_gate, axis=2).astype(new_kv.dtype)
@@ -482,30 +482,24 @@ def __call__(
           dtype=jnp.int32,
       )
 
-    # Broadcast the compressed KV representations across all indexer heads
-    # -> [batch, 1, n_windows, index_head_dim]
-    compressed_kv = jnp.expand_dims(compressed, axis=1)
-    # -> [batch, index_n_heads, n_windows, index_head_dim]
-    compressed_kv = jnp.broadcast_to(
-        compressed_kv,
-        (batch_size, self.index_n_heads, compressed_len, self.index_head_dim),
-    )
+    # OPTIMIZATION DIPAK: Removed jnp.expand_dims and jnp.broadcast_to on compressed_kv to save memory.
 
     # Project the latent query to match the Indexer's dimensions
     # [batch, seq_len, index_n_heads * index_head_dim] -> [batch, seq_len, index_n_heads, index_head_dim]
     q = self.q_proj(q_latent).reshape((batch_size, seq_len, self.index_n_heads, self.index_head_dim))
-    # -> [batch, index_n_heads, seq_len, index_head_dim]
-    q = jnp.transpose(q, (0, 2, 1, 3))
+
+    # OPTIMIZATION DIPAK: Removed jnp.transpose on q to maintain bshd layout for fused einsum.
 
     # Apply standard Rotary Positional Embeddings to queries
-    q = self.rotary_emb(q, position_ids, unsqueeze_dim=1)
+    # OPTIMIZATION DIPAK: Adjusted unsqueeze_dim for bshd layout
+    q = self.rotary_emb(q, position_ids, unsqueeze_dim=2)
 
     q = q.astype(jnp.float32)
-    compressed_kv = compressed_kv.astype(jnp.float32)
+    compressed_kv = compressed.astype(jnp.float32)
 
     # Compute dot product between Queries and Compressed KV Blocks
-    # -> [batch, index_n_heads, seq_len, n_windows]
-    scores = jnp.einsum("bhsd,bhwd->bhsw", q, compressed_kv)
+    # OPTIMIZATION DIPAK: Native aligned einsum. q is [b, s, h, d], compressed_kv is [b, w, d]
+    scores = jnp.einsum("bshd,bwd->bshw", q, compressed_kv)
     scores = jax.nn.relu(scores) * self.softmax_scale
 
     # Compute routing weights to combine scores across indexer heads
@@ -514,7 +508,7 @@ def __call__(
 
     # Combine individual head scores according to routing weights
     # -> [batch, seq_len, n_windows]
-    index_scores = jnp.einsum("bhsw,bsh->bsw", scores, weights)
+    index_scores = jnp.einsum("bshw,bsh->bsw", scores, weights)
 
     k = min(self.index_topk, compressed_len)
 
@@ -523,18 +517,22 @@ def __call__(
     entry_indices = jnp.arange(compressed_len)
     future_mask = entry_indices[None, None, :] >= jnp.expand_dims(causal_threshold, axis=-1)
 
-    index_scores = jnp.where(future_mask, jnp.full_like(index_scores, -jnp.inf), index_scores)
+    # OPTIMIZATION DIPAK: Additive causal mask to avoid full_like tensor allocation
+    causal_mask = jnp.where(future_mask, -jnp.inf, 0.0)
+    index_scores += causal_mask
 
     # Apply standard segment attention mask (additive 0 and -inf)
     if attention_mask is not None:
       index_scores += attention_mask[:, :, :compressed_len]
 
     # Retrieve the top-k highest scoring block indices for each token
-    top_k_indices = jax.lax.top_k(index_scores, k)[1]
+    # OPTIMIZATION DIPAK: Replaced slow top_k with hardware-friendly approx_max_k
+    top_k_indices = jax.lax.approx_max_k(index_scores, k)[1]
 
     # Invalidate any top-k selections that point to future blocks (edge case safety)
     invalid = top_k_indices >= jnp.expand_dims(causal_threshold, axis=-1)
-    top_k_indices = jnp.where(invalid, jnp.full_like(top_k_indices, -1), top_k_indices)
+    # OPTIMIZATION DIPAK: Removed full_like in top_k_indices invalidation
+    top_k_indices = jnp.where(invalid, -1, top_k_indices)
 
     return top_k_indices
 
@@ -654,12 +652,10 @@ def __call__(
 
     # Only compute and apply the complex block mask if top-k selections exist
     if k > 0:
-      valid = top_k_indices >= 0
-      entry_indices = jnp.arange(compressed_len)[None, None, :]
-      is_in_topk = jnp.expand_dims(top_k_indices, axis=-1) == entry_indices[None, ...]
-      is_valid_and_in_topk = is_in_topk & jnp.expand_dims(valid, axis=-1)
-
-      is_selected = jnp.any(is_valid_and_in_topk, axis=2)
+      # OPTIMIZATION DIPAK: Replaced memory-heavy 4D broadcast with boolean one_hot scatter.
+      # JAX natively maps -1 (padding) to all False in boolean one_hot, avoiding explicit validity checks.
+      one_hot = jax.nn.one_hot(top_k_indices, compressed_len, dtype=jnp.bool_)
+      is_selected = jnp.any(one_hot, axis=2)
       is_selected = jnp.expand_dims(is_selected, axis=1)
 
       compressed_mask = jnp.where(is_selected, 0.0, DEFAULT_MASK_VALUE).astype(self.dtype)
@@ -1053,8 +1049,20 @@ def __call__(
       )
 
     # Extend local KV tensors with the compressed blocks
+    decoder_segment_ids_q = decoder_segment_ids
+    decoder_segment_ids_kv = decoder_segment_ids
+
     if compressed_kv is not None:
       kv = jnp.concatenate([kv, compressed_kv], axis=1)
+      if decoder_segment_ids is not None:
+        # Pad segment IDs to match the new KV sequence length.
+        # We pad with 0 because the actual masking for compressed blocks is
+        # explicitly handled by the compressed_mask boolean logic.
+        segment_padding = jnp.zeros(
+            (decoder_segment_ids.shape[0], compressed_kv.shape[1]),
+            dtype=decoder_segment_ids.dtype,
+        )
+        decoder_segment_ids_kv = jnp.concatenate([decoder_segment_ids, segment_padding], axis=1)
 
     kv = checkpoint_name(kv, "kv_proj")
 
@@ -1072,7 +1080,8 @@ def __call__(
         q,
         kv,
         kv,
-        decoder_segment_ids,
+        decoder_segment_ids_q,
+        decoder_segment_ids_kv,
         inputs_positions,
         model_mode,
         sinks=self.sinks.value,