Support dynamic Tokamax Flash Attention for DeepSeek-V4 HCA and CSA layers

Author: octatrifan

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

@@ -0,0 +1,47 @@
+# Copyright 2026 Google LLC
+# Small model config for DeepSeek-V4 testing and compilation profiling
+
+base_emb_dim: 1024
+base_num_query_heads: 16
+base_num_kv_heads: 1
+base_num_decoder_layers: 4
+base_mlp_dim: 1024
+base_moe_mlp_dim: 1024
+vocab_size: 32000
+head_dim: 128
+
+# --- Standard Defaults ---
+enable_dropout: false
+logits_via_embedding: false
+normalization_layer_epsilon: 1.0e-6
+
+# --- V4 Specific Architectural Keys ---
+decoder_block: "deepseek4"
+mhc_expansion_rate: 4
+first_num_hash_layers: 1
+indexer_head_dim: 64
+indexer_n_heads: 16
+indexer_topk: 64
+
+compress_ratios: [0, 4, 8, 4]
+
+# --- MoE configuration ---
+mlp_activations: ["silu", "linear"]
+num_experts: 16
+num_experts_per_tok: 2
+mlp_activations_limit: 10
+shared_experts: 1
+routed_score_func: "sqrtsoftplus"
+
+# --- Attention configuration ---
+attention_type: 'compressed'
+q_lora_rank: 256
+o_groups: 2
+o_lora_rank: 256
+sliding_window_size: 128
+
+# --- RoPE ---
+rope_type: "default"
+rope_max_timescale: 10000                  # Main RoPE theta
+compressed_rope_max_timescale: 160000      # Compressed RoPE theta
+max_position_embeddings: 65536

src/maxtext/configs/types.py

@@ -229,6 +229,7 @@ class ProfilerType(str, Enum):
     "deepseek3-tiny",
     "deepseek3.2-671b",
     "deepseek4-284b",
+    "deepseek4-small",
     "deepseek-custom",
     "kimi-k2-1t",
     "gemma-7b",
@@ -3189,8 +3190,9 @@ 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.decoder_block == DecoderBlockType.DEEPSEEK4 and self.attention not in ("dot_product", "flash"):
+      raise ValueError("DeepSeek4 decoder block currently supports dot_product and flash attention.")
+
     if self.use_indexer:
       if self.q_lora_rank == 0:
         raise NotImplementedError("Sparse indexer has not implemented for q_lora_rank = 0.")

src/maxtext/layers/attention_compressed.py

@@ -1058,6 +1058,23 @@ def __call__(
 
     kv = checkpoint_name(kv, "kv_proj")
 
+    # Pad total KV length to tile size multiple for Tokamax block alignment
+    if self.attention_kernel == "flash":
+      block_size = self.config.sa_block_kv
+      pad_kv_total = (block_size - (kv.shape[1] % block_size)) % block_size
+
+      if pad_kv_total > 0:
+        c_len = compressed_kv.shape[1] if compressed_kv is not None else 0
+        if c_len > 0:
+          # Prepend padding to the compressed blocks so they remain at the end of the sequence
+          local_kv = kv[:, :-c_len]
+          comp_kv = kv[:, -c_len:]
+          comp_kv_padded = jnp.pad(comp_kv, ((0, 0), (pad_kv_total, 0), (0, 0), (0, 0)))
+          kv = jnp.concatenate([local_kv, comp_kv_padded], axis=1)
+        else:
+          # Fallback: Pad at the end if no compressed blocks exist
+          kv = jnp.pad(kv, ((0, 0), (0, pad_kv_total), (0, 0), (0, 0)))
+
     # Prepare the mask shape for the underlying AttentionOp
     if compressed_mask is not None:
       compressed_mask = jnp.expand_dims(compressed_mask, axis=2)
@@ -1066,6 +1083,17 @@ def __call__(
     if self.query_pre_attn_scalar and self.query_pre_attn_scalar != 1.0:
       q = q * self.query_pre_attn_scalar
 
+    # Build indexer mask explicitly for tokamax splash kernel
+    indexer_mask = None
+    if self.attention_kernel == "flash" and compressed_mask is not None:
+      indexer_mask = self.attention_op.generate_attention_mask(
+          q, kv, decoder_segment_ids, model_mode, compressed_mask=compressed_mask
+      )
+
+      if indexer_mask is not None:
+        # Robustly extract the first head & first key dimension slices to match splash expectations
+        indexer_mask = indexer_mask[:, 0, 0, :, :]
+
     # Compute Attention
     # -> [batch, q_length, num_query_heads, head_dim]
     attn_out = self.attention_op(
@@ -1077,6 +1105,7 @@ def __call__(
         model_mode,
         sinks=self.sinks.value,
         compressed_mask=compressed_mask,
+        indexer_mask=indexer_mask,
     )
 
     # Reverse RoPE on Values

src/maxtext/layers/attention_op.py

@@ -1033,6 +1033,7 @@ def apply_attention(
             decoder_segment_ids,
             self.attn_logits_soft_cap,
             sinks,
+            indexer_mask,
             record_max_logits=record_max_logits,
         )
         if max_logits is not None:
@@ -1297,8 +1298,18 @@ def create_sa_config(config, query, key, attn_logits_soft_cap):
       return sa_config
 
     sa_config = create_sa_config(self.config, query, key, attn_logits_soft_cap)
-    mask_shape = (query.shape[2], key.shape[2])  # (q_seq_len, kv_seq_len)
+    block_q = sa_config.block_q
+    block_kv = sa_config.block_kv
+    if self.attention_type == AttentionType.COMPRESSED and (
+        (query.shape[2] % block_q != 0) or (key.shape[2] % block_kv != 0)
+    ):
+      padded_q_len = ((query.shape[2] + block_q - 1) // block_q) * block_q
+      padded_kv_len = ((key.shape[2] + block_kv - 1) // block_kv) * block_kv
+      mask_shape = (padded_q_len, padded_kv_len)
+    else:
+      mask_shape = (query.shape[2], key.shape[2])  # (q_seq_len, kv_seq_len)
     mask_module = tokamax_splash_mask if self.config.use_tokamax_splash else splash_attention_mask
+
     if self.attention_type == AttentionType.FULL:
       mask = mask_module.FullMask(mask_shape)
     else:
@@ -1315,14 +1326,14 @@ def create_sa_config(config, query, key, attn_logits_soft_cap):
       local_window_size = (self.sliding_window_size - 1, self.sliding_window_size)
       if use_load_balanced_cp:
         mask &= LoadBalancedLocalMask(
-            shape=(query.shape[2], key.shape[2]),
+            shape=mask_shape,
             window_size=local_window_size,
             offset=0,
             cp_size=cp_size,
         )
       else:
         mask &= mask_module.LocalMask(
-            shape=(query.shape[2], key.shape[2]),
+            shape=mask_shape,
             window_size=local_window_size,
             offset=0,
         )
@@ -1332,16 +1343,15 @@ def create_sa_config(config, query, key, attn_logits_soft_cap):
 
       if use_load_balanced_cp:
         mask &= LoadBalancedChunkedCausalMask(
-            shape=(query.shape[2], key.shape[2]),
+            shape=mask_shape,
             chunk_size=self.chunk_attn_window_size,
             cp_size=cp_size,
         )
       else:
         mask &= ChunkedCausalMask(
-            shape=(query.shape[2], key.shape[2]),
+            shape=mask_shape,
             chunk_size=self.chunk_attn_window_size,
         )
-
     max_logit_value = None
     if self.config.use_tokamax_splash:
       # Create mask
@@ -1364,9 +1374,12 @@ def wrap_splash_kernel(single_head_mask):
         )
         return splash_kernel
 
-      splash_kernel = wrap_splash_kernel(single_head_mask)
       segment_axis_names_splash_kernel = self._logical_to_mesh_axes((Q_LENGTH,))
-      splash_kernel = self._maybe_shard_with_pspec(splash_kernel, segment_axis_names_splash_kernel)
+      if indexer_mask is None:
+        splash_kernel = wrap_splash_kernel(single_head_mask)
+        splash_kernel = self._maybe_shard_with_pspec(splash_kernel, segment_axis_names_splash_kernel)
+      else:
+        splash_kernel = None
     elif self.config.use_jax_splash:
       if self.config.use_max_logit_estimate > 0:
         sa_config = dataclasses.replace(sa_config, max_logit_const=self.config.use_max_logit_estimate)
@@ -1489,7 +1502,7 @@ def wrap_flash_attention(
         decoder_segment_ids_tuple = None
 
       if self.config.use_tokamax_splash:
-        if self.config.use_indexer and indexer_mask is not None:
+        if indexer_mask is not None:
           # Construct the splash kernel call with dynamic mask
           def dynamic_mask_splash_kernel(q, k, v, segment, sinks, indexer_mask):
             splash_kernel = tokamax_splash_kernel.make_dynamic_splash_mha(
@@ -1533,6 +1546,7 @@ def kernel_fn(q, k, v, d, s):
                 query, key, value, decoder_segment_ids_tuple, sinks
             )
             return attention_output, None
+
       elif self.config.use_jax_splash:
         materialized_mask = jnp.asarray(mask[:, :])
         attention_output = jax_flash_attention.flash_attention_block_masked(

src/maxtext/layers/nnx_decoders.py

@@ -1272,6 +1272,7 @@ def get_norm_layer(self, num_features: int, rngs: nnx.Rngs):
         DecoderBlockType.MISTRAL,
         DecoderBlockType.MIXTRAL,
         DecoderBlockType.DEEPSEEK,
+        DecoderBlockType.DEEPSEEK4,
         DecoderBlockType.GEMMA,
         DecoderBlockType.GEMMA2,
         DecoderBlockType.GEMMA3,