Add flag-controlled sliced MLA projections

src/maxtext/configs/base.yml

@@ -444,6 +444,7 @@ kv_lora_rank: 512
 qk_nope_head_dim: 128
 qk_rope_head_dim: 64
 v_head_dim: 128
+use_sliced_mla_proj: false # Whether to slice projection kernel weights before contraction in MLA instead of running full projection + jnp.split.
 
 # Compressed Attention parameters
 o_lora_rank: 0 # Output LoRA rank for Compressed Attention.

src/maxtext/configs/types.py

@@ -666,6 +666,13 @@ class MlaAttention(BaseModel):
   qk_nope_head_dim: NonNegativeInt = Field(128, description="Dimension for non-RoPE part of QK heads in MLA.")
   qk_rope_head_dim: NonNegativeInt = Field(64, description="Dimension for RoPE part of QK heads in MLA.")
   v_head_dim: NonNegativeInt = Field(128, description="Dimension of V heads in MLA.")
+  use_sliced_mla_proj: bool = Field(
+      False,
+      description=(
+          "Whether to slice projection kernel weights before contraction in MLA"
+          " instead of running full projection + jnp.split."
+      ),
+  )
 
 
 class CompressedAttention(BaseModel):
@@ -3805,6 +3812,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       raise ValueError("`share_kv_projections` is not compatible with `fused_qkv`.")
     if self.share_kv_projections and self.attention_type == "mla":
       raise ValueError("`share_kv_projections` is not compatible with `attention_type='mla'`.")
+    if self.use_sliced_mla_proj and (self.quantization or self.use_qwix_quantization):
+      raise ValueError("`use_sliced_mla_proj` is not supported with quantization.")
 
     if self.use_manual_quantization and not self.use_batch_split_schedule:
       raise ValueError("manual quantization is only used when `use_batch_split_schedule=True`.")

src/maxtext/layers/attention_mla.py

@@ -917,17 +917,30 @@ def mla_query_projection(
 
     if self.q_lora_rank == 0:
       q = self.query(inputs_q, out_sharding=query_sharding)
+      q_nope, q_pe = jnp.split(q, [self.qk_nope_head_dim], axis=-1)
     else:
       # LoRA path
       low_rank_q = self.wq_a(inputs_q, out_sharding=wqa_out_sharding)  # [B, L, q_lora_rank]
       low_rank_q = checkpoint_name(low_rank_q, "query_wa_proj")
       low_rank_q = self.q_norm(low_rank_q)  # RMSNorm on low rank
       low_rank_q = checkpoint_name(low_rank_q, "mla_q")
-      q = self.wq_b(low_rank_q, out_sharding=query_sharding)  # [B, L, n_heads, qk_head_dim]
+      if self.config.use_sliced_mla_proj and self.wq_b.quant is None:
+        q_nope = self.wq_b(
+            low_rank_q,
+            out_sharding=query_sharding,
+            slice_bounds=(0, self.qk_nope_head_dim),
+        )  # [B, L, n_heads, qk_nope_head_dim]
+        q_pe = self.wq_b(
+            low_rank_q,
+            out_sharding=query_sharding,
+            slice_bounds=(self.qk_nope_head_dim, self.qk_head_dim),
+        )  # [B, L, n_heads, qk_rope_head_dim]
+      else:
+        q = self.wq_b(low_rank_q, out_sharding=query_sharding)  # [B, L, n_heads, qk_head_dim]
+        q_nope, q_pe = jnp.split(q, [self.qk_nope_head_dim], axis=-1)
 
     # Partial RoPE: Split into non-positional and rotary parts.
     # last dimension: qk_nope_head_dim, qk_rope_head_dim
-    q_nope, q_pe = jnp.split(q, [self.qk_nope_head_dim], axis=-1)
     q_nope = self._maybe_shard_with_logical(q_nope, query_logical_name)
     q_pe = self.apply_rotary_embedding(q_pe, inputs_positions=inputs_positions)
     q_pe = self._maybe_shard_with_logical(q_pe, query_logical_name)
@@ -951,10 +964,24 @@ def mla_get_key_value(self, low_rank_main, key_rope, model_mode):
       value_logical_name = self.value_axis_names
 
     wkva_out_sharding = create_sharding(self.mesh, key_logical_name)
-    kv_out = self.wkv_b(low_rank_main, out_sharding=wkva_out_sharding)
-
-    # Split kv_out into key_nope and value parts.
-    key_nope, value = jnp.split(kv_out, [self.qk_nope_head_dim], axis=-1)
+    if self.config.use_sliced_mla_proj and self.wkv_b.quant is None:
+      key_nope = self.wkv_b(
+          low_rank_main,
+          out_sharding=wkva_out_sharding,
+          slice_bounds=(0, self.qk_nope_head_dim),
+      )  # [B, L, n_heads, qk_nope_head_dim]
+      value = self.wkv_b(
+          low_rank_main,
+          out_sharding=wkva_out_sharding,
+          slice_bounds=(
+              self.qk_nope_head_dim,
+              self.qk_nope_head_dim + self.v_head_dim,
+          ),
+      )  # [B, L, n_heads, v_head_dim]
+    else:
+      kv_out = self.wkv_b(low_rank_main, out_sharding=wkva_out_sharding)
+      # Split kv_out into key_nope and value parts.
+      key_nope, value = jnp.split(kv_out, [self.qk_nope_head_dim], axis=-1)
     key_rope = jnp.broadcast_to(key_rope, (key_nope.shape[0], key_nope.shape[1], self.num_query_heads, key_rope.shape[3]))
     key_nope = self._maybe_shard_with_logical(key_nope, key_logical_name)
     key_rope = self._maybe_shard_with_logical(key_rope, key_logical_name)

src/maxtext/layers/linears.py

@@ -249,11 +249,21 @@ def _maybe_two_stage_all_gather(self, kernel):
     kernel = shard(kernel, stage2)
     return kernel
 
-  def __call__(self, inputs: Array, _initializing: bool = False, out_sharding: NamedSharding | None = None) -> Array:
+  def __call__(
+      self,
+      inputs: Array,
+      _initializing: bool = False,
+      out_sharding: NamedSharding | None = None,
+      slice_bounds: tuple[int, int] | None = None,
+  ) -> Array:
     """Applies a linear transformation to the inputs along multiple dimensions.
 
     Args:
       inputs: The nd-array to be transformed.
+      _initializing: Whether the module is initializing.
+      out_sharding: Optional sharding for the output.
+      slice_bounds: Optional tuple (begin, end) to slice the kernel and bias on
+        the last (output-feature) axis before contraction. Unquantized only.
 
     Returns:
       The transformed input.
@@ -281,6 +291,14 @@ def __call__(self, inputs: Array, _initializing: bool = False, out_sharding: Nam
         kernel = jax.device_put(kernel, max_utils.device_space())
       kernel = jnp.asarray(kernel, self.dtype)
 
+    if slice_bounds is not None:
+      if self.quant is not None:
+        raise ValueError("sliced contraction is only supported when quant is None")
+      begin, end = slice_bounds
+      if not 0 <= begin < end <= kernel.shape[-1]:
+        raise ValueError(f"slice_bounds {slice_bounds} must be valid and within [0, {kernel.shape[-1]}]")
+      kernel = kernel[..., begin:end]
+
     kernel = self._maybe_two_stage_all_gather(kernel)
 
     # out_sharding should be None for auto mesh axis
@@ -294,13 +312,16 @@ def __call__(self, inputs: Array, _initializing: bool = False, out_sharding: Nam
         norm_axis,
         contract_ind,
         self.matmul_precision,
-        self.quant_dot_general,
+        self.quant_dot_general if slice_bounds is None else None,
         _initializing,
         out_sharding,
     )
 
     if self.bias is not None:
       bias = jnp.asarray(self.bias[...], self.dtype)
+      if slice_bounds is not None:
+        begin, end = slice_bounds
+        bias = bias[..., begin:end]
       output += bias
     return output
 

tests/end_to_end/tpu/deepseek/Run_DeepSeek.md

@@ -304,6 +304,14 @@ python3 -m tests.utils.forward_pass_logit_checker \
 
 To run MMLU benchmarks and validate the model's performance, follow the instructions provided [here](https://github.com/AI-Hypercomputer/maxtext/blob/main/benchmarks/api_server/README.md).
 
+## MLA Optimization
+
+To optimize Multi-Head Latent Attention (MLA) performance, you can enable sliced projections.
+
+*   **Flag**: `use_sliced_mla_proj` (default: `False`)
+*   **Description**: When set to `True`, it slices the projection kernel weights before contraction in MLA, instead of running the full projection followed by `jnp.split`. This can improve performance.
+*   **Constraint**: Sliced contraction is only supported when quantization is disabled (`quant=None`).
+
 ## Supported MoE strategy
 * Dropless
   * [MegaBlocks](https://arxiv.org/abs/2211.15841) implementation with flag `sparse_matmul=True megablox=True`.

tests/unit/attention_test.py

@@ -2093,6 +2093,108 @@ def test_indexer_autoregression(self, prefill_len, target_len):
       self.assertEqual(mla_full_this_idx.shape, mla_idx.shape)
       self.assertTrue(jax.numpy.allclose(mla_full_this_idx, mla_idx, rtol=2e-02, atol=2e-02, equal_nan=False))
 
+  def test_sliced_mla_projections(self):
+    config_arguments = self.config_arguments.copy()
+
+    # Enable sliced projections for one config
+    config_arguments_sliced = config_arguments.copy()
+    config_arguments_sliced["use_sliced_mla_proj"] = True
+
+    cfg_normal, mla_normal = self.init_mla(config_arguments, rope_type="default")
+    _, mla_sliced = self.init_mla(config_arguments_sliced, rope_type="default")
+
+    # Sync weights
+    nnx.update(mla_sliced, nnx.state(mla_normal))
+
+    # Test TRAIN mode with gradient comparison
+    lnx, decoder_segment_ids, decoder_positions = self.get_structured_data(cfg_normal, cfg_normal.dtype)
+
+    def loss_fn(model, x):
+      out, _ = model(
+          x,
+          x,
+          decoder_segment_ids=decoder_segment_ids,
+          inputs_positions=decoder_positions,
+          deterministic=True,
+          model_mode=MODEL_MODE_TRAIN,
+      )
+      return jnp.mean(out.astype(jnp.float32) ** 2), out
+
+    (loss_normal, out_normal_train), (grad_model_normal, grad_x_normal) = nnx.value_and_grad(
+        loss_fn, argnums=(0, 1), has_aux=True
+    )(mla_normal, lnx)
+
+    (loss_sliced, out_sliced_train), (grad_model_sliced, grad_x_sliced) = nnx.value_and_grad(
+        loss_fn, argnums=(0, 1), has_aux=True
+    )(mla_sliced, lnx)
+
+    self.assertTrue(jnp.allclose(loss_normal, loss_sliced, rtol=1e-05, atol=1e-05, equal_nan=False))
+    self.assertTrue(jnp.allclose(out_normal_train, out_sliced_train, rtol=1e-05, atol=1e-05, equal_nan=False))
+    self.assertTrue(jnp.allclose(grad_x_normal, grad_x_sliced, rtol=1e-05, atol=1e-05, equal_nan=False))
+
+    grad_model_close = jax.tree_util.tree_map(
+        lambda x, y: jnp.allclose(x, y, rtol=1e-05, atol=1e-05, equal_nan=False),
+        grad_model_normal,
+        grad_model_sliced,
+    )
+    self.assertTrue(jax.tree_util.tree_all(grad_model_close))
+
+    # Test PREFILL mode followed by AUTOREGRESSIVE mode to test caching
+    prefill_length = cfg_normal.max_prefill_predict_length
+    decode_total_length = cfg_normal.max_target_length
+
+    # Re-initialize to ensure clean cache
+    cfg_normal, mla_normal = self.init_mla(config_arguments, rope_type="default")
+    _, mla_sliced = self.init_mla(config_arguments_sliced, rope_type="default")
+    nnx.update(mla_sliced, nnx.state(mla_normal))
+
+    lnx_prefill = lnx[:, 0:prefill_length, :]
+    decoder_segment_ids_prefill = decoder_segment_ids[:, 0:prefill_length]
+    decoder_positions_prefill = decoder_positions[:, 0:prefill_length]
+
+    out_normal_prefill, _ = mla_normal(
+        lnx_prefill,
+        lnx_prefill,
+        decoder_segment_ids=decoder_segment_ids_prefill,
+        inputs_positions=decoder_positions_prefill,
+        deterministic=True,
+        model_mode=MODEL_MODE_PREFILL,
+    )
+
+    out_sliced_prefill, _ = mla_sliced(
+        lnx_prefill,
+        lnx_prefill,
+        decoder_segment_ids=decoder_segment_ids_prefill,
+        inputs_positions=decoder_positions_prefill,
+        deterministic=True,
+        model_mode=MODEL_MODE_PREFILL,
+    )
+
+    self.assertTrue(jnp.allclose(out_normal_prefill, out_sliced_prefill, rtol=1e-05, atol=1e-05, equal_nan=False))
+
+    # Run autoregressive steps
+    for idx in range(prefill_length, decode_total_length):
+      lnx_idx = lnx[:, idx : idx + 1, :]
+      decoder_positions_idx = decoder_positions[:, idx : idx + 1]
+
+      out_normal_idx, _ = mla_normal(
+          lnx_idx,
+          lnx_idx,
+          inputs_positions=decoder_positions_idx,
+          deterministic=True,
+          model_mode=MODEL_MODE_AUTOREGRESSIVE,
+      )
+
+      out_sliced_idx, _ = mla_sliced(
+          lnx_idx,
+          lnx_idx,
+          inputs_positions=decoder_positions_idx,
+          deterministic=True,
+          model_mode=MODEL_MODE_AUTOREGRESSIVE,
+      )
+
+      self.assertTrue(jnp.allclose(out_normal_idx, out_sliced_idx, rtol=1e-05, atol=1e-05, equal_nan=False))
+
   def test_projection_initialization(self):
     """Tests that MLA and Attention layers initialize the correct projection weights."""
     # 1. Initialize a standard Attention layer for comparison

tests/unit/configs_value_test.py

@@ -362,6 +362,18 @@ def test_elastic_backup_kind_validation(self):
     with self.assertRaises(pydantic.ValidationError):
       pyconfig.initialize(argv)
 
+  def test_sliced_mla_proj_disallows_quantization(self):
+    """Tests that use_sliced_mla_proj=True is incompatible with quantization."""
+    argv = [
+        "",
+        _BASE_CONFIG_PATH,
+        "run_name=test",
+        "use_sliced_mla_proj=true",
+        "quantization=int8",
+    ]
+    with self.assertRaises(pydantic.ValidationError):
+      pyconfig.initialize(argv)
+
 
 if __name__ == "__main__":
   absltest.main()

tests/unit/linears_test.py

@@ -16,6 +16,7 @@
 
 import sys
 import unittest
+from unittest.mock import MagicMock
 from flax import nnx
 from flax.linen import partitioning as nn_partitioning
 import jax
@@ -103,6 +104,82 @@ def test_bias(self):
     self.assertEqual(outputs.shape, (batch_size, out_features))
     self.assertIsNotNone(layer.bias)
 
+  def test_slice_bounds(self):
+    batch_size = 2
+    in_features = 4
+    n_heads = 2
+    head_dim = 8
+
+    layer = linears.DenseGeneral(
+        in_features_shape=in_features,
+        out_features_shape=(n_heads, head_dim),
+        use_bias=True,
+        rngs=self.rngs,
+    )
+
+    inputs = jax.random.normal(jax.random.PRNGKey(0), (batch_size, in_features))
+    full_output = layer(inputs)
+
+    slice1 = layer(inputs, slice_bounds=(0, 3))
+    slice2 = layer(inputs, slice_bounds=(3, head_dim))
+
+    self.assertEqual(slice1.shape, (batch_size, n_heads, 3))
+    self.assertEqual(slice2.shape, (batch_size, n_heads, head_dim - 3))
+
+    np.testing.assert_allclose(slice1, full_output[..., :3], rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(slice2, full_output[..., 3:], rtol=1e-5, atol=1e-5)
+
+  def test_slice_bounds_with_quantization(self):
+    batch_size = 2
+    in_features = 4
+    n_heads = 2
+    head_dim = 8
+
+    mock_quant = MagicMock()
+    mock_quant.quant_mode = None
+    # Configure mock for ToNNX initialization
+    mock_cls = mock_quant.dot_general_cls.return_value
+    mock_instance = mock_cls.return_value
+    mock_instance.init_with_output.return_value = (MagicMock(), {})
+    mock_instance.apply.return_value = (MagicMock(), {})
+
+    layer = linears.DenseGeneral(
+        in_features_shape=in_features,
+        out_features_shape=(n_heads, head_dim),
+        use_bias=True,
+        quant=mock_quant,
+        rngs=self.rngs,
+    )
+
+    inputs = jax.random.normal(jax.random.PRNGKey(0), (batch_size, in_features))
+
+    with self.assertRaisesRegex(ValueError, "sliced contraction is only supported when quant is None"):
+      layer(inputs, slice_bounds=(0, 3))
+
+  def test_slice_bounds_invalid(self):
+    batch_size = 2
+    in_features = 4
+    n_heads = 2
+    head_dim = 8
+
+    layer = linears.DenseGeneral(
+        in_features_shape=in_features,
+        out_features_shape=(n_heads, head_dim),
+        use_bias=True,
+        rngs=self.rngs,
+    )
+
+    inputs = jax.random.normal(jax.random.PRNGKey(0), (batch_size, in_features))
+
+    with self.assertRaisesRegex(ValueError, "slice_bounds .* must be valid and within"):
+      layer(inputs, slice_bounds=(3, 2))
+
+    with self.assertRaisesRegex(ValueError, "slice_bounds .* must be valid and within"):
+      layer(inputs, slice_bounds=(-1, 4))
+
+    with self.assertRaisesRegex(ValueError, "slice_bounds .* must be valid and within"):
+      layer(inputs, slice_bounds=(0, 100))
+
   def _run_dense_test(self, axis, in_feat_shape, expected_shape):
     batch_size = 2
     seq_len = 3