diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 5669e012ae..389e0808f1 100644 --- a/src/maxtext/configs/base.yml +++ b/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 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index b96823ffe7..a3e89f6ea6 100644 --- a/src/maxtext/configs/types.py +++ b/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." ), ) diff --git a/src/maxtext/layers/attention_compressed.py b/src/maxtext/layers/attention_compressed.py index 3cf7d26230..a9b0e08330 100644 --- a/src/maxtext/layers/attention_compressed.py +++ b/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) diff --git a/tests/unit/deepseek_v4_indexer_loss_test.py b/tests/unit/deepseek_v4_indexer_loss_test.py index 9f83d2c122..e60600c278 100644 --- a/tests/unit/deepseek_v4_indexer_loss_test.py +++ b/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):