From 192491ba94c8834f2299272b4aa9a442ff3d78c2 Mon Sep 17 00:00:00 2001 From: Jiahao Chen Zhou Date: Wed, 22 Jul 2026 06:46:19 +0000 Subject: [PATCH] [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. --- src/maxtext/configs/base.yml | 3 + src/maxtext/configs/types.py | 23 ++++ src/maxtext/layers/attention_mla.py | 123 +++++++++++++++++---- tests/unit/deepseek32_vs_reference_test.py | 19 +++- 4 files changed, 142 insertions(+), 26 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index fe49545bdb..33944c0eb0 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -192,6 +192,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 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 @@ -1336,3 +1337,5 @@ elastic_enabled: false elastic_timeout_seconds: 300 elastic_max_retries: 10 elastic_min_slice_count: -1 + + diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 22a2a5ff0b..4d99af7c31 100644 --- a/src/maxtext/configs/types.py +++ b/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).") + mla_qk_head_chunk_size: int = Field( + 0, + ge=0, + description=( + "Chunk size over heads dimension for QK attention dot product in mla. " + "Default is 0 (no chunking). Reduces memory footprint at the cost of time." + ), + ) class MoBa(BaseModel): @@ -3248,6 +3256,21 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de 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.mla_qk_head_chunk_size > 0: + if self.attention != "dot_product": + raise ValueError("`mla_qk_head_chunk_size` is only supported with `dot_product` attention.") + if self.mla_qk_head_chunk_size > self.num_query_heads or self.num_query_heads % self.mla_qk_head_chunk_size != 0: + raise ValueError( + f"`mla_qk_head_chunk_size` ({self.mla_qk_head_chunk_size}) must cleanly divide exactly into " + f"`num_query_heads` ({self.num_query_heads})." + ) + if self.use_indexer and ( + self.mla_qk_head_chunk_size > self.indexer_n_heads or self.indexer_n_heads % self.mla_qk_head_chunk_size != 0 + ): + raise ValueError( + f"`mla_qk_head_chunk_size` ({self.mla_qk_head_chunk_size}) must cleanly divide exactly into " + f"`indexer_n_heads` ({self.indexer_n_heads})." + ) if self.use_indexer: if self.q_lora_rank == 0: raise NotImplementedError("Sparse indexer has not implemented for q_lora_rank = 0.") diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index 9fcc76f2ea..dca15ae071 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/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, "mla_qk_head_chunk_size", 0) + + def chunked_head_loop(chunk_size): + + num_chunks = h // 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, 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, 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_inner = jnp.einsum("hbtd, bsd -> btsh", q_c, k, precision=self.config.matmul_precision) + logits_inner = jax.nn.relu(logits_inner) + + score_chunk = jnp.einsum( + "btsh, hbt -> bts", + logits_inner, + 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) + idx_score, _ = jax.lax.scan(jax.checkpoint(scan_body_indexer), init_score, {"q": q_h, "w": w_h}) + return idx_score.astype(q.dtype) + + if head_chunk_size > 0: + indexer_score = chunked_head_loop(head_chunk_size) + + 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,65 @@ 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, "mla_qk_head_chunk_size", 0) + if head_chunk_size > 0: + + 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(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) diff --git a/tests/unit/deepseek32_vs_reference_test.py b/tests/unit/deepseek32_vs_reference_test.py index a16e3380a3..084266c9b9 100644 --- a/tests/unit/deepseek32_vs_reference_test.py +++ b/tests/unit/deepseek32_vs_reference_test.py @@ -751,7 +751,7 @@ def get_jax_mla_weights(pt_mla, cfg): } -def get_cfg_and_mesh(config, run_name, dtype, batch_size, seq_len, attention, indexer_topk): +def get_cfg_and_mesh(config, run_name, dtype, batch_size, seq_len, attention, indexer_topk, mla_qk_head_chunk_size=0): """Returns MaxText configuration and mesh.""" cfg = pyconfig.initialize( [None, get_test_config_path()], @@ -769,6 +769,7 @@ def get_cfg_and_mesh(config, run_name, dtype, batch_size, seq_len, attention, in max_prefill_predict_length=seq_len, attention=attention, indexer_topk=indexer_topk, + mla_qk_head_chunk_size=mla_qk_head_chunk_size, **asdict(config), ) devices_array = maxtext_utils.create_device_mesh(cfg) @@ -831,8 +832,8 @@ def get_data(self, seq_len): class DeepseekV32IndexerTest(DeepseekTestBase): """Tests for the Sparse Indexer (Top-K Selection).""" - # indexer_topk=4 - def test_indexer_match(self, seq_len=8): + @parameterized.parameters(0, 2, 4) + def test_indexer_match(self, mla_qk_head_chunk_size, seq_len=8): """Verifies Indexer output matches PyTorch output.""" torch_inputs, jax_inputs = self.get_data(seq_len) pt_mask = torch_inputs["mask"] @@ -865,6 +866,7 @@ def test_indexer_match(self, seq_len=8): seq_len=self.seq_len, attention="dot_product", indexer_topk=4, + mla_qk_head_chunk_size=mla_qk_head_chunk_size, ) # Indexer specific RoPE (interleave=False) @@ -936,6 +938,14 @@ class DeepseekV32MLATest(DeepseekTestBase): "indexer_topk": 128, "check_norm": True, }, + { + "testcase_name": "dot_product_s128_k128_c4", + "attention": "dot_product", + "seq_len": 128, + "indexer_topk": 128, + "check_norm": True, + "mla_qk_head_chunk_size": 4, + }, { "testcase_name": "flash_s128_k4", "attention": "flash", @@ -951,7 +961,7 @@ class DeepseekV32MLATest(DeepseekTestBase): "check_norm": True, }, ) - def test_mla_parity(self, attention, seq_len, indexer_topk, check_norm=False): + def test_mla_parity(self, attention, seq_len, indexer_topk, check_norm=False, mla_qk_head_chunk_size=0): """Verifies JAX MLA output against the PyTorch reference implementation.""" torch_inputs, jax_inputs = self.get_data(seq_len) @@ -978,6 +988,7 @@ def test_mla_parity(self, attention, seq_len, indexer_topk, check_norm=False): seq_len=self.seq_len, attention=attention, indexer_topk=indexer_topk, + mla_qk_head_chunk_size=mla_qk_head_chunk_size, ) jax_mla = attention_mla.MLA(