[Memory] Implement QK attention head chunking to prevent HBM OOM

Currently, evaluating the QK dot product natively in MLA materializes
the full [batch, count_of_heads, q_len, kv_len] attention scores tensor.
This causes severe memory constraints (HBM OOM) for large context lengths.

This change introduces the Config variable `qk_head_chunk_size` alongside
a `jax.lax.scan` algorithm. Since the `heads` dimension is locally dense
and unsharded (unlike `q_len` which is Context Parallel), we scan over
groups of heads iteratively computing query-key projections and their softmax
partials, avoiding vast intermediate allocations.

Author: zcjhao

src/maxtext/configs/base.yml

@@ -1336,3 +1336,6 @@ elastic_enabled: false
 elastic_timeout_seconds: 300
 elastic_max_retries: 10
 elastic_min_slice_count: -1
+
+# Limits HBM footprint by sequentially evaluating the QK matrix across the unsharded local heads dimension natively. Size must evenly divide the number of attention heads.
+qk_head_chunk_size: 0

src/maxtext/configs/types.py

@@ -637,6 +637,14 @@ class Attention(BaseModel):
   force_q_layout: bool = Field(False, description="Force the Q layout")
   use_qk_clip: bool = Field(False, description="Whether to use QK-Clip (MuonClip) for training stability.")
   qk_clip_threshold: float = Field(100.0, description="Threshold for QK-Clip (tau).")
+  qk_head_chunk_size: int = Field(
+      0,
+      description=(
+          "Chunk size over heads dimension for QK attention dot product. "
+          "Default is 0 (no chunking). Must divide the number of heads evenly. "
+          "This reduces memory footprint at the cost of time, especially helpful in long context."
+      ),
+  )
 
 
 class MoBa(BaseModel):

src/maxtext/layers/attention_mla.py

@@ -361,18 +361,57 @@ def __call__(
     if k.shape[1] <= self.indexer_topk:
       return None, None, None
 
-    # Compute Index Scores
-    # QK product: relu(q @ k.T), [b, t, s, h]
-    # Similar to MQA, each key is shared by h query head
-    logits = jnp.einsum("bthd, bsd -> btsh", q, k, precision=self.config.matmul_precision)
-    logits = jax.nn.relu(logits)
     # Compute head weights: project from input, [b, t, embed_dim] -> [b, t, h]
     weights = self.weights_proj(inputs_q)
     # Weights scaling affect indexer_score, but does not affect topk_indices. Keep scaling for numerical stability.
     # https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/87e509a2e5a100d221c97df52c6e8be7835f0057/inference/model.py#L478-L480
     weights = weights * (self.n_heads**-0.5) * self.softmax_scale
-    # Aggregate head-wise logits: logits @ weights
-    indexer_score = jnp.einsum("btsh, bth -> bts", logits, weights, precision=self.config.matmul_precision)  # [b, t, s]
+
+    # Compute Index Scores
+    # When qk_head_chunk_size > 0, compute Index Scores by chunking the 'heads' dimension to reduce memory
+    # The naive evaluation materializes [b, t, s, h].
+    # We use jax.lax.scan to compute the score iteratively over head chunks.
+    b, t, h, d = q.shape
+    # Control the HBM footprint of QK tensor: [batch, q_len, s_len, heads]
+    # If set to 0 (defaults), it falls back to native materialization.
+    head_chunk_size = getattr(self.config, "qk_head_chunk_size", 0)
+    if head_chunk_size > 0:
+      if head_chunk_size > h or h % head_chunk_size != 0:
+        raise ValueError(
+            f"qk_head_chunk_size ({head_chunk_size}) must be <= number of heads ({h}) "
+            f"and divide it evenly."
+        )
+      num_chunks = h // head_chunk_size
+      # q: [b, t, h, d] -> [h, b, t, d] -> [num_chunks, head_chunk_size, b, t, d]
+      q_h = q.transpose(2, 0, 1, 3).reshape(num_chunks, head_chunk_size, b, t, d)
+      # weights: [b, t, h] -> [h, b, t] -> [num_chunks, head_chunk_size, b, t]
+      w_h = weights.transpose(2, 0, 1).reshape(num_chunks, head_chunk_size, b, t)
+
+      def scan_body_indexer(carry, xs):
+        q_c = xs["q"]  # [h_chunk, b, t, d]
+        w_c = xs["w"]  # [h_chunk, b, t]
+
+        # Directly use the chunked shapes in einsum to avoid transposes inside the loop
+        logits = jnp.einsum("hbtd, bsd -> btsh", q_c, k, precision=self.config.matmul_precision)
+        logits = jax.nn.relu(logits)
+
+        score_chunk = jnp.einsum(
+            "btsh, hbt -> bts",
+            logits,
+            w_c,
+            precision=self.config.matmul_precision,
+        )
+        return carry + score_chunk.astype(jnp.float32), None
+
+      init_score = jnp.zeros((b, t, k.shape[1]), dtype=jnp.float32)
+      indexer_score, _ = jax.lax.scan(jax.checkpoint(scan_body_indexer), init_score, {"q": q_h, "w": w_h})
+      indexer_score = indexer_score.astype(q.dtype)
+
+    else:
+      # Aggregate head-wise logits: logits @ weights natively
+      logits = jnp.einsum("bthd, bsd -> btsh", q, k, precision=self.config.matmul_precision)
+      logits = jax.nn.relu(logits)
+      indexer_score = jnp.einsum("btsh, bth -> bts", logits, weights, precision=self.config.matmul_precision)
 
     internal_padding_mask = None
     if cached_s is not None:
@@ -1086,25 +1125,69 @@ def calculate_indexer_loss(
     query = jax.lax.stop_gradient(query)
     key = jax.lax.stop_gradient(key)
 
-    # Compute attention scores: [b, t, h, d] @ [b, s, h, d] -> [b, h, t, s]
-    attention_scores = jnp.einsum("bthd, bshd -> bhts", query, key, precision=self.config.matmul_precision)
-
+    # Ensure indexer_score updates identically in all branches
     if sparse_loss:
-      # indexer_mask is already pre-filtered with the attention_mask if any
-      attention_scores = attention_scores + indexer_mask[:, None, :, :]
       indexer_score = indexer_score + indexer_mask
-    elif attention_mask is not None:
-      # indexer_score already applies attention_mask; updating attention_scores only
-      attention_scores = attention_scores + attention_mask[:, None, :, :]
-
-    # Use float32 for softmax numerical stability.
-    attention_probs = jax.nn.softmax(attention_scores.astype(jnp.float32), axis=-1)
     indexer_probs = jax.nn.softmax(indexer_score.astype(jnp.float32), axis=-1)
 
-    # Aggregate heads: [b, h, t, s] -> [b, t, s]
-    attention_probs = jnp.sum(attention_probs, axis=1)
-    # Force materialization and prevent fusion across this point to reuse the intermediate tensor
-    attention_probs = jax.lax.optimization_barrier(attention_probs)
+    batch, q_len, heads, dim = query.shape
+
+    # Chunk across the 'heads' dimension manually using jax.lax.scan
+    # Control the HBM footprint of QK tensor: [batch, q_len, s_len, heads]
+    # If set to 0, it falls back to native implementation.
+    head_chunk_size = getattr(self.config, "qk_head_chunk_size", 0)
+    if head_chunk_size > 0:
+      if head_chunk_size > heads or heads % head_chunk_size != 0:
+        raise ValueError(
+            f"qk_head_chunk_size ({head_chunk_size}) must be <= number of heads ({heads}) "
+            f"and divide it evenly."
+        )
+      num_chunks = heads // head_chunk_size
+
+      # Transpose and reshape to put chunk dimension first for jax.lax.scan
+      # 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)
+      k_h = key.transpose(2, 0, 1, 3).reshape(num_chunks, head_chunk_size, batch, key.shape[1], dim)
+
+      def scan_body_heads(carry, xs):
+        q_c = xs["q"]  # [h_chunk, b, t, d]
+        k_c = xs["k"]  # [h_chunk, b, s, d]
+
+        # Directly use the chunked shapes in einsum to avoid transposes inside the loop
+        attn_chunk = jnp.einsum(
+            "hbtd, hbsd -> bhts",
+            q_c,
+            k_c,
+            precision=self.config.matmul_precision,
+        )
+
+        if sparse_loss:
+          attn_chunk = attn_chunk + indexer_mask[:, None, :, :]
+        elif attention_mask is not None:
+          attn_chunk = attn_chunk + attention_mask[:, None, :, :]
+
+        probs_chunk = jax.nn.softmax(attn_chunk.astype(jnp.float32), axis=-1)
+        probs_chunk_sum = jnp.sum(probs_chunk, axis=1)  # [b, t, s]
+
+        return carry + probs_chunk_sum, None
+
+      init_probs = jnp.zeros((batch, q_len, key.shape[1]), dtype=jnp.float32)
+      attention_probs, _ = jax.lax.scan(jax.checkpoint(scan_body_heads), init_probs, {"q": q_h, "k": k_h})
+
+    else:
+      # Native implementation (default) if chunking is disabled
+      attention_scores = jnp.einsum(
+          "bthd, bshd -> bhts",
+          query,
+          key,
+          precision=self.config.matmul_precision,
+      )
+      if sparse_loss:
+        attention_scores = attention_scores + indexer_mask[:, None, :, :]
+      elif attention_mask is not None:
+        attention_scores = attention_scores + attention_mask[:, None, :, :]
+      attention_probs = jnp.sum(jax.nn.softmax(attention_scores.astype(jnp.float32), axis=-1), axis=1)
+      attention_probs = jax.lax.optimization_barrier(attention_probs)
     # L1 normalize aggregated target distribution
     attention_probs = attention_probs / (jnp.sum(attention_probs, axis=-1, keepdims=True) + EPS)