Implement QK attention head chunking for CSA [Deepseek v4]

src/maxtext/configs/base.yml

@@ -204,7 +204,7 @@ logits_dot_in_fp32: false  # whether to use fp32 in logits_dense or shared_embed
 cast_logits_to_fp32: true # whether to cast the logits to fp32. the higher precision is generally beneficial, but it can vary slightly.
 float32_qk_product: false # in dot_product attention, whether to cast to fp32 the inputs to qk product
 float32_logits: false # in dot_product attention, whether to cast to fp32 the inputs to softmax
-mla_qk_head_chunk_size: 0 # Limits HBM footprint by sequentially evaluating the QK matrix in the Indexer across the unsharded local heads dimension natively.
+mla_qk_head_chunk_size: 0 # Limits HBM footprint by sequentially evaluating the QK matrix in MLA and CSA Indexers across the unsharded local heads dimension natively.
 float32_weight_sum: true # whether to use full fp32 precision to sum expert weights for numerical stability
 float32_gate_logits: false # whether to cast inputs to fp32 to compute MoE gate logits for numerical stability
 

src/maxtext/configs/types.py

@@ -700,7 +700,7 @@ class Attention(BaseModel):
       0,
       ge=0,
       description=(
-          "Chunk size over heads dimension for QK attention dot product in mla.  "
+          "Chunk size over heads dimension for QK attention dot product in MLA and CSA (DeepSeek-V4).  "
           "Default is 0 (no chunking). Reduces memory footprint at the cost of time."
       ),
   )

src/maxtext/layers/attention_compressed.py

@@ -883,20 +883,52 @@ def indexer_compressor_fn(buf_kv, buf_gate):
       return empty_indices, (jnp.zeros((batch_size, seq_len, 0), dtype=jnp.float32) if return_scores else None)
 
     # --- TOP-K ROUTING MATH (Executes in both Prefill and AR) ---
-    compressed_kv = jnp.expand_dims(compressed, axis=1)
-    compressed_kv = jnp.broadcast_to(compressed_kv, (batch_size, self.index_n_heads, compressed_len, self.index_head_dim))
-
     q = self.q_proj(q_latent).reshape((batch_size, seq_len, self.index_n_heads, self.index_head_dim))
     q = jnp.transpose(q, (0, 2, 1, 3))
     q = self.rotary_emb(q, position_ids, unsqueeze_dim=1)
 
-    q = q.astype(jnp.float32)
-    compressed_kv = compressed_kv.astype(jnp.float32)
-
-    scores = jnp.einsum("bhsd,bhwd->bhsw", q, compressed_kv)
-    scores = jax.nn.relu(scores) * self.softmax_scale
     weights = self.weights_proj(hidden_states).astype(jnp.float32) * self.weights_scaling
-    index_scores = jnp.einsum("bhsw,bsh->bsw", scores, weights)
+
+    head_chunk_size = getattr(self.config, "mla_qk_head_chunk_size", 0)
+    if head_chunk_size > 0:
+      num_chunks = self.index_n_heads // head_chunk_size
+      q_h = q.transpose(1, 0, 2, 3).reshape(num_chunks, head_chunk_size, batch_size, seq_len, self.index_head_dim)
+      w_h = weights.transpose(2, 0, 1).reshape(num_chunks, head_chunk_size, batch_size, seq_len)
+      compressed_fp32 = compressed.astype(jnp.float32)
+
+      def scan_body_indexer(carry, xs):
+        q_c = xs["q"].astype(jnp.float32)
+        w_c = xs["w"]
+
+        scores_inner = jnp.einsum(
+            "cbsd, bwd -> bcsw",
+            q_c,
+            compressed_fp32,
+            precision=self.config.matmul_precision,
+        )
+        scores_inner = jax.nn.relu(scores_inner) * self.softmax_scale
+        score_chunk = jnp.einsum(
+            "bcsw, cbs -> bsw",
+            scores_inner,
+            w_c,
+            precision=self.config.matmul_precision,
+        )
+        return carry + score_chunk, None
+
+      init_score = jnp.zeros((batch_size, seq_len, compressed_len), dtype=jnp.float32)
+      index_scores, _ = jax.lax.scan(jax.checkpoint(scan_body_indexer), init_score, {"q": q_h, "w": w_h})
+    else:
+      compressed_kv = jnp.expand_dims(compressed, axis=1)
+      compressed_kv = jnp.broadcast_to(
+          compressed_kv, (batch_size, self.index_n_heads, compressed_len, self.index_head_dim)
+      )
+
+      q = q.astype(jnp.float32)
+      compressed_kv = compressed_kv.astype(jnp.float32)
+
+      scores = jnp.einsum("bhsd,bhwd->bhsw", q, compressed_kv, precision=self.config.matmul_precision)
+      scores = jax.nn.relu(scores) * self.softmax_scale
+      index_scores = jnp.einsum("bhsw,bsh->bsw", scores, weights, precision=self.config.matmul_precision)
 
     k = min(self.index_topk, compressed_len)
 
@@ -1796,7 +1828,7 @@ def calculate_csa_indexer_loss(
     if compressed_kv is None or indexer_score is None:
       return jnp.array(0.0, dtype=jnp.float32)
 
-    batch, q_len, _, _ = query.shape
+    batch, q_len, heads, dim = query.shape
     compressed_len = compressed_kv.shape[1]
     if compressed_len == 0:
       return jnp.array(0.0, dtype=jnp.float32)
@@ -1859,15 +1891,37 @@ def calculate_csa_indexer_loss(
     log_indexer_probs = jnp.where(valid_tokens_mask[:, :, None], log_indexer_probs, 0.0)
 
     # Query is already scaled by softmax_scale in compressed_query_projection; do not scale again
-    attention_scores = jnp.einsum("bthd, bwd -> bhtw", query, k_vec, precision=self.config.matmul_precision)
-    attention_scores = attention_scores + c_mask
-
-    # Apply NaN shielding for pre-block tokens
-    safe_scores = jnp.where(valid_tokens_mask[:, None, :, None], attention_scores, 0.0)
-    raw_probs = jax.nn.softmax(safe_scores.astype(jnp.float32), axis=-1)
-    raw_probs = jnp.where(valid_tokens_mask[:, None, :, None], raw_probs, 0.0)
-    target_probs = jnp.sum(raw_probs, axis=1)
-    target_probs = jax.lax.optimization_barrier(target_probs)
+    head_chunk_size = getattr(self.config, "mla_qk_head_chunk_size", 0)
+    if head_chunk_size > 0:
+      num_chunks = heads // head_chunk_size
+      # query: [b, t, h, d] -> [h, b, t, d] -> [num_chunks, head_chunk_size, b, t, d]
+      q_h = query.transpose(2, 0, 1, 3).reshape(num_chunks, head_chunk_size, batch, q_len, dim)
+
+      def scan_body_heads(carry, xs):
+        q_c = xs["q"]  # [h_chunk, b, t, d]
+        attn_chunk = jnp.einsum("hbtd, bwd -> bhtw", q_c, k_vec, precision=self.config.matmul_precision)
+        attn_chunk = attn_chunk + c_mask
+
+        # Apply NaN shielding for pre-block tokens
+        safe_attn = jnp.where(valid_tokens_mask[:, None, :, None], attn_chunk, 0.0)
+        probs_chunk = jax.nn.softmax(safe_attn.astype(jnp.float32), axis=-1)
+        probs_chunk = jnp.where(valid_tokens_mask[:, None, :, None], probs_chunk, 0.0)
+        probs_chunk_sum = jnp.sum(probs_chunk, axis=1)  # [b, t, w]
+
+        return carry + probs_chunk_sum, None
+
+      init_probs = jnp.zeros((batch, q_len, compressed_len), dtype=jnp.float32)
+      target_probs, _ = jax.lax.scan(scan_body_heads, init_probs, {"q": q_h})
+    else:
+      attention_scores = jnp.einsum("bthd, bwd -> bhtw", query, k_vec, precision=self.config.matmul_precision)
+      attention_scores = attention_scores + c_mask
+
+      # Apply NaN shielding for pre-block tokens
+      safe_scores = jnp.where(valid_tokens_mask[:, None, :, None], attention_scores, 0.0)
+      raw_probs = jax.nn.softmax(safe_scores.astype(jnp.float32), axis=-1)
+      raw_probs = jnp.where(valid_tokens_mask[:, None, :, None], raw_probs, 0.0)
+      target_probs = jnp.sum(raw_probs, axis=1)
+      target_probs = jax.lax.optimization_barrier(target_probs)
 
     # L1 normalize aggregated target distribution across compressed blocks
     target_probs = jnp.where(valid_tokens_mask[:, :, None], target_probs, 0.0)

tests/unit/deepseek_v4_indexer_loss_test.py

@@ -73,6 +73,7 @@ def _get_config(
       use_indexer=True,
       indexer_loss_scaling_factor=0.5,
       indexer_sparse_training=False,
+      mla_qk_head_chunk_size=0,
       indexer_topk=None,
   ):
     """Constructs a test MaxTextConfig with CSA indexer configuration."""
@@ -87,6 +88,7 @@ def _get_config(
         f"use_indexer={use_indexer}",
         f"indexer_loss_scaling_factor={indexer_loss_scaling_factor}",
         f"indexer_sparse_training={indexer_sparse_training}",
+        f"mla_qk_head_chunk_size={mla_qk_head_chunk_size}",
         f"max_target_length={self.seq_len}",
         f"indexer_topk={topk}",
         f"indexer_n_heads={self.indexer_n_heads}",
@@ -100,6 +102,7 @@ def _get_config(
         "o_groups=2",
         "o_lora_rank=16",
         "enable_checkpointing=False",
+        "vocab_size=32",
     ]
     return pyconfig.initialize(argv)
 
@@ -203,6 +206,123 @@ def test_csa_indexer_loss_kl_divergence_zero(self):
     )
     np.testing.assert_allclose(float(loss), 0.0, atol=1e-5)
 
+  def test_csa_indexer_loss_head_chunking_parity(self):
+    """Test that head chunking scan produces loss mathematically equivalent within FP tolerance to native einsum."""
+    config_chunked = self._get_config(mla_qk_head_chunk_size=2)
+    config_native = self._get_config(mla_qk_head_chunk_size=0)
+    attn_chunked = self._init_csa_attention(config_chunked)
+    attn_native = self._init_csa_attention(config_native)
+
+    n_windows = self.seq_len // self.compress_ratio
+    rng = jax.random.PRNGKey(42)
+    k1, k2, k3 = jax.random.split(rng, 3)
+    query = jax.random.normal(
+        k1, (self.batch_size, self.seq_len, config_chunked.num_query_heads, config_chunked.head_dim)
+    )
+    compressed_kv = jax.random.normal(
+        k2, (self.batch_size, n_windows, config_chunked.num_kv_heads, config_chunked.head_dim)
+    )
+    indexer_score = jax.random.normal(k3, (self.batch_size, self.seq_len, n_windows))
+    compressed_mask = jnp.zeros((self.batch_size, 1, self.seq_len, n_windows))
+
+    loss_chunked = attn_chunked.calculate_csa_indexer_loss(
+        indexer_score=indexer_score,
+        query=query,
+        compressed_kv=compressed_kv,
+        compressed_mask=compressed_mask,
+        segment_mask=None,
+        position_ids=None,
+        sparse_loss=False,
+        scaling_factor=1.0,
+    )
+    loss_native = attn_native.calculate_csa_indexer_loss(
+        indexer_score=indexer_score,
+        query=query,
+        compressed_kv=compressed_kv,
+        compressed_mask=compressed_mask,
+        segment_mask=None,
+        position_ids=None,
+        sparse_loss=False,
+        scaling_factor=1.0,
+    )
+    np.testing.assert_allclose(float(loss_chunked), float(loss_native), rtol=1e-5, atol=1e-5)
+
+  def test_csa_indexer_scoring_head_chunking_parity(self):
+    """Test that indexer forward scoring produces identical top-k indices and scores with chunking."""
+    config_chunked = self._get_config(mla_qk_head_chunk_size=2)
+    config_native = self._get_config(mla_qk_head_chunk_size=0)
+    attn_chunked = self._init_csa_attention(config_chunked)
+    attn_native = self._init_csa_attention(config_native)
+
+    idx_chunked = attn_chunked.csa_compressor.indexer
+    idx_native = attn_native.csa_compressor.indexer
+    nnx.update(idx_chunked, nnx.state(idx_native))
+
+    hidden_states = jax.random.normal(jax.random.PRNGKey(10), (self.batch_size, self.seq_len, config_native.emb_dim))
+    q_latent = jax.random.normal(jax.random.PRNGKey(11), (self.batch_size, self.seq_len, config_native.q_lora_rank))
+    positions = jnp.broadcast_to(jnp.arange(self.seq_len)[None, :], (self.batch_size, self.seq_len))
+
+    topk_chunked, scores_chunked = idx_chunked(
+        hidden_states=hidden_states,
+        q_latent=q_latent,
+        position_ids=positions,
+        model_mode=MODEL_MODE_TRAIN,
+        return_scores=True,
+    )
+    topk_native, scores_native = idx_native(
+        hidden_states=hidden_states,
+        q_latent=q_latent,
+        position_ids=positions,
+        model_mode=MODEL_MODE_TRAIN,
+        return_scores=True,
+    )
+    np.testing.assert_array_equal(np.array(topk_chunked), np.array(topk_native))
+    np.testing.assert_allclose(np.array(scores_chunked), np.array(scores_native), rtol=1e-5, atol=1e-5)
+
+  def test_csa_indexer_chunked_gradients_flow(self):
+    """Test that gradients flow through indexer under head chunking with jax.checkpoint rematerialization."""
+    config = self._get_config(indexer_loss_scaling_factor=1.0, indexer_sparse_training=False, mla_qk_head_chunk_size=2)
+    attn = self._init_csa_attention(config)
+
+    inputs_q = jax.random.normal(jax.random.PRNGKey(1), (self.batch_size, self.seq_len, config.emb_dim))
+    inputs_kv = jax.random.normal(jax.random.PRNGKey(2), (self.batch_size, self.seq_len, config.emb_dim))
+    positions = jnp.broadcast_to(jnp.arange(self.seq_len)[None, :], (self.batch_size, self.seq_len))
+    segment_ids = jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32)
+
+    def loss_fn(attn_model, q, kv, seg=segment_ids, pos=positions):
+      attn_model(
+          inputs_q=q,
+          inputs_kv=kv,
+          decoder_segment_ids=seg,
+          inputs_positions=pos,
+          deterministic=True,
+          model_mode=MODEL_MODE_TRAIN,
+      )
+      return attn_model.indexer_loss.get_value()
+
+    grad_model_fn = nnx.grad(loss_fn, argnums=0)
+    grads = grad_model_fn(attn, inputs_q, inputs_kv)
+
+    self.assertIsNotNone(grads.csa_compressor.indexer.q_proj.kernel)
+    self.assertIsNotNone(grads.csa_compressor.indexer.kv_proj.kernel)
+    self.assertIsNotNone(grads.csa_compressor.indexer.gate_proj.kernel)
+    self.assertIsNotNone(grads.csa_compressor.indexer.weights_proj.kernel)
+
+    q_grad_norm = jnp.linalg.norm(grads.csa_compressor.indexer.q_proj.kernel.get_value())
+    self.assertGreater(float(q_grad_norm), 0.0)
+    self.assertGreater(float(jnp.linalg.norm(grads.csa_compressor.indexer.weights_proj.kernel.get_value())), 0.0)
+
+    # Gradients must not leak into main model projections
+    self.assertAlmostEqual(float(jnp.linalg.norm(grads.wq_a.kernel.get_value())), 0.0)
+    self.assertAlmostEqual(float(jnp.linalg.norm(grads.wq_b.kernel.get_value())), 0.0)
+    self.assertAlmostEqual(float(jnp.linalg.norm(grads.wkv.kernel.get_value())), 0.0)
+
+    # Gradients with respect to inputs must be zero
+    grad_inputs_fn = nnx.grad(loss_fn, argnums=(1, 2))
+    grad_q, grad_kv = grad_inputs_fn(attn, inputs_q, inputs_kv)
+    self.assertAlmostEqual(float(jnp.linalg.norm(grad_q)), 0.0)
+    self.assertAlmostEqual(float(jnp.linalg.norm(grad_kv)), 0.0)
+
   def test_csa_indexer_gradients_flow(self):
     """Test that gradients flow to indexer parameters and do not leak into main projections or inputs."""
     for is_sparse in (False, True):