From 33ffda671f17c3c26ac5a206c8ad9f71d88f67b0 Mon Sep 17 00:00:00 2001 From: Daniel Mandragona Date: Mon, 27 Jul 2026 17:49:55 +0000 Subject: [PATCH 1/7] Add flag-controlled sliced MLA projections Port performance optimization from CL 947805696 into MaxText under a new configuration flag use_sliced_mla_projections (defaulting to False). - Add use_sliced_mla_projections flag to base.yml and types.py. - Support slice_bounds in DenseGeneral.__call__ to slice projection weights/bias prior to contraction for unquantized paths. - Update mla_query_projection and mla_get_key_value in attention_mla.py to use sliced projections when enabled. - Add unit test coverage in linears_test.py. TAG=agy CHANGE_STEWARD=true --- src/maxtext/configs/base.yml | 1 + src/maxtext/configs/types.py | 7 ++++++ src/maxtext/layers/attention_mla.py | 39 ++++++++++++++++++++++++----- src/maxtext/layers/linears.py | 22 ++++++++++++++-- tests/unit/linears_test.py | 25 ++++++++++++++++++ 5 files changed, 86 insertions(+), 8 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index db0b1a68a7..d3e58d7bbf 100644 --- a/src/maxtext/configs/base.yml +++ b/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_projections: 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. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 6416d0ae9f..438657866d 100644 --- a/src/maxtext/configs/types.py +++ b/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_projections: bool = Field( + False, + description=( + "Whether to slice projection kernel weights before contraction in MLA" + " instead of running full projection + jnp.split." + ), + ) class CompressedAttention(BaseModel): diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index dca15ae071..689def51c5 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/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_projections 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_projections 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) diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py index 36e69c8a07..32c259d32b 100644 --- a/src/maxtext/layers/linears.py +++ b/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,11 @@ 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: + assert self.quant is None, "sliced contraction is only supported when quant is None" + begin, end = slice_bounds + kernel = kernel[..., begin:end] + kernel = self._maybe_two_stage_all_gather(kernel) # out_sharding should be None for auto mesh axis @@ -294,13 +309,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 diff --git a/tests/unit/linears_test.py b/tests/unit/linears_test.py index d65f7e9e2e..b01d9741e9 100644 --- a/tests/unit/linears_test.py +++ b/tests/unit/linears_test.py @@ -103,6 +103,31 @@ 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 _run_dense_test(self, axis, in_feat_shape, expected_shape): batch_size = 2 seq_len = 3 From f65d64f4f30ed48817459c30ee425704a404cee2 Mon Sep 17 00:00:00 2001 From: Daniel Mandragona Date: Mon, 27 Jul 2026 20:49:27 +0000 Subject: [PATCH 2/7] Add unit test for sliced MLA projections Verify correctness of sliced MLA projections by comparing outputs against the default unsliced implementation in train, prefill, and autoregressive modes. CHANGE_STEWARD=true TAG=agy CONV=86839aa9-de96-414b-b896-1ac9aeb38589 --- tests/unit/attention_test.py | 99 ++++++++++++++++++++++++++++++++++++ 1 file changed, 99 insertions(+) diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 653e0a5d92..bcd482204a 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -2093,6 +2093,105 @@ 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)) + @pytest.mark.tpu_only + 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_projections"] = True + + cfg_normal, mla_normal = self.init_mla(config_arguments, rope_type="default") + cfg_sliced, mla_sliced = self.init_mla(config_arguments_sliced, rope_type="default") + + # Sync weights + nnx.update(mla_sliced, nnx.state(mla_normal)) + + # Test TRAIN mode + lnx, decoder_segment_ids, decoder_positions = self.get_structured_data(cfg_normal, cfg_normal.dtype) + + out_normal_train, _ = mla_normal( + lnx, + lnx, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + out_sliced_train, _ = mla_sliced( + lnx, + lnx, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + self.assertTrue( + jnp.allclose(out_normal_train, out_sliced_train, rtol=1e-05, atol=1e-05, equal_nan=False) + ) + + # 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") + cfg_sliced, 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 From 3375ac05aea3e5eab4770a8c97bbe8ee1d8a34f4 Mon Sep 17 00:00:00 2001 From: Daniel Mandragona Date: Mon, 27 Jul 2026 20:59:26 +0000 Subject: [PATCH 3/7] Document use_sliced_mla_projections flag Add documentation for the new use_sliced_mla_projections performance optimization flag in the DeepSeek runner guide. TAG=agy CONV=a262e292-a0d4-49f9-ad26-f42e697a5979 CHANGE_STEWARD=true --- tests/end_to_end/tpu/deepseek/Run_DeepSeek.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md b/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md index 0c98e730b4..5e6e8476e5 100644 --- a/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md +++ b/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_projections` (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`. From 52d94b0ff98d2c1e1b84f546e2e14eeb1b3e3218 Mon Sep 17 00:00:00 2001 From: Daniel Mandragona Date: Mon, 27 Jul 2026 22:23:32 +0000 Subject: [PATCH 4/7] Address codecov comments: improve test coverage Remove tpu_only decorator from test_sliced_mla_projections to run it on CPU. Add test_slice_bounds_with_quantization to cover assertion in DenseGeneral. Format with pyink and fix pylint warnings. TAG=agy CONV=e336d40c-d648-45cb-9abd-7b79a924be5e CHANGE_STEWARD=true --- tests/unit/attention_test.py | 17 +++++------------ tests/unit/linears_test.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 12 deletions(-) diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index bcd482204a..340bca17c4 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -2093,7 +2093,6 @@ 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)) - @pytest.mark.tpu_only def test_sliced_mla_projections(self): config_arguments = self.config_arguments.copy() @@ -2102,7 +2101,7 @@ def test_sliced_mla_projections(self): config_arguments_sliced["use_sliced_mla_projections"] = True cfg_normal, mla_normal = self.init_mla(config_arguments, rope_type="default") - cfg_sliced, mla_sliced = self.init_mla(config_arguments_sliced, rope_type="default") + _, mla_sliced = self.init_mla(config_arguments_sliced, rope_type="default") # Sync weights nnx.update(mla_sliced, nnx.state(mla_normal)) @@ -2128,9 +2127,7 @@ def test_sliced_mla_projections(self): model_mode=MODEL_MODE_TRAIN, ) - self.assertTrue( - jnp.allclose(out_normal_train, out_sliced_train, 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)) # Test PREFILL mode followed by AUTOREGRESSIVE mode to test caching prefill_length = cfg_normal.max_prefill_predict_length @@ -2138,7 +2135,7 @@ def test_sliced_mla_projections(self): # Re-initialize to ensure clean cache cfg_normal, mla_normal = self.init_mla(config_arguments, rope_type="default") - cfg_sliced, mla_sliced = self.init_mla(config_arguments_sliced, 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, :] @@ -2163,9 +2160,7 @@ def test_sliced_mla_projections(self): model_mode=MODEL_MODE_PREFILL, ) - self.assertTrue( - jnp.allclose(out_normal_prefill, out_sliced_prefill, rtol=1e-05, atol=1e-05, equal_nan=False) - ) + 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): @@ -2188,9 +2183,7 @@ def test_sliced_mla_projections(self): model_mode=MODEL_MODE_AUTOREGRESSIVE, ) - self.assertTrue( - jnp.allclose(out_normal_idx, out_sliced_idx, rtol=1e-05, atol=1e-05, equal_nan=False) - ) + 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.""" diff --git a/tests/unit/linears_test.py b/tests/unit/linears_test.py index b01d9741e9..6a34892721 100644 --- a/tests/unit/linears_test.py +++ b/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 @@ -128,6 +129,33 @@ def test_slice_bounds(self): 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.assertRaises(AssertionError): + layer(inputs, slice_bounds=(0, 3)) + def _run_dense_test(self, axis, in_feat_shape, expected_shape): batch_size = 2 seq_len = 3 From 0e53e1a33dc867dad865b5adf0da6ab4be62bc1c Mon Sep 17 00:00:00 2001 From: Daniel Mandragona Date: Fri, 31 Jul 2026 21:21:30 +0000 Subject: [PATCH 5/7] Address NuojCheng review comments: check gradients and avoid assertions --- src/maxtext/layers/linears.py | 3 ++- tests/unit/attention_test.py | 43 +++++++++++++++++++++-------------- tests/unit/linears_test.py | 2 +- 3 files changed, 29 insertions(+), 19 deletions(-) diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py index 32c259d32b..bc3b3c55ba 100644 --- a/src/maxtext/layers/linears.py +++ b/src/maxtext/layers/linears.py @@ -292,7 +292,8 @@ def __call__( kernel = jnp.asarray(kernel, self.dtype) if slice_bounds is not None: - assert self.quant is None, "sliced contraction is only supported when quant is None" + if self.quant is not None: + raise ValueError("sliced contraction is only supported when quant is None") begin, end = slice_bounds kernel = kernel[..., begin:end] diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 340bca17c4..9fa9e18cd5 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -2106,28 +2106,37 @@ def test_sliced_mla_projections(self): # Sync weights nnx.update(mla_sliced, nnx.state(mla_normal)) - # Test TRAIN mode + # Test TRAIN mode with gradient comparison lnx, decoder_segment_ids, decoder_positions = self.get_structured_data(cfg_normal, cfg_normal.dtype) - out_normal_train, _ = mla_normal( - lnx, - lnx, - decoder_segment_ids=decoder_segment_ids, - inputs_positions=decoder_positions, - deterministic=True, - model_mode=MODEL_MODE_TRAIN, - ) + 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 - out_sliced_train, _ = mla_sliced( - lnx, - lnx, - decoder_segment_ids=decoder_segment_ids, - inputs_positions=decoder_positions, - deterministic=True, - model_mode=MODEL_MODE_TRAIN, - ) + (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(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 diff --git a/tests/unit/linears_test.py b/tests/unit/linears_test.py index 6a34892721..0b37ebf12c 100644 --- a/tests/unit/linears_test.py +++ b/tests/unit/linears_test.py @@ -153,7 +153,7 @@ def test_slice_bounds_with_quantization(self): inputs = jax.random.normal(jax.random.PRNGKey(0), (batch_size, in_features)) - with self.assertRaises(AssertionError): + with self.assertRaisesRegex(ValueError, "sliced contraction is only supported when quant is None"): layer(inputs, slice_bounds=(0, 3)) def _run_dense_test(self, axis, in_feat_shape, expected_shape): From 3f987b831bcc22b0fc446dcc64cf4cd2be770fed Mon Sep 17 00:00:00 2001 From: Daniel Mandragona Date: Tue, 4 Aug 2026 01:35:51 +0000 Subject: [PATCH 6/7] Address RissyRan review comments on sliced MLA projections - Rename flag from use_sliced_mla_projections to use_sliced_mla_proj in base.yml, types.py, attention_mla.py, and docs. - Add early validation check in types.py to reject use_sliced_mla_proj when quantization is enabled. - Validate slice_bounds in DenseGeneral.__call__ to ensure indices are within [0, kernel.shape[-1]]. - Add unit tests for invalid slice_bounds and quantization incompatibility. TAG=agy CONV=ac3bc03e-1ae4-4049-a085-f49debe869f7 CHANGE_STEWARD=true --- src/maxtext/configs/base.yml | 2 +- src/maxtext/configs/types.py | 4 +++- src/maxtext/layers/attention_mla.py | 4 ++-- src/maxtext/layers/linears.py | 4 ++++ tests/end_to_end/tpu/deepseek/Run_DeepSeek.md | 2 +- tests/unit/attention_test.py | 2 +- tests/unit/configs_value_test.py | 12 ++++++++++ tests/unit/linears_test.py | 24 +++++++++++++++++++ 8 files changed, 48 insertions(+), 6 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index d3e58d7bbf..fd102bed67 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -444,7 +444,7 @@ kv_lora_rank: 512 qk_nope_head_dim: 128 qk_rope_head_dim: 64 v_head_dim: 128 -use_sliced_mla_projections: false # Whether to slice projection kernel weights before contraction in MLA instead of running full projection + jnp.split. +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. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 438657866d..9bdd658499 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -666,7 +666,7 @@ 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_projections: bool = Field( + use_sliced_mla_proj: bool = Field( False, description=( "Whether to slice projection kernel weights before contraction in MLA" @@ -3812,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`.") diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index 689def51c5..f89a8c94a5 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -924,7 +924,7 @@ def mla_query_projection( 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") - if self.config.use_sliced_mla_projections and self.wq_b.quant is None: + 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, @@ -964,7 +964,7 @@ 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) - if self.config.use_sliced_mla_projections and self.wkv_b.quant is None: + 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, diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py index bc3b3c55ba..e75c1ddaf2 100644 --- a/src/maxtext/layers/linears.py +++ b/src/maxtext/layers/linears.py @@ -295,6 +295,10 @@ def __call__( 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) diff --git a/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md b/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md index 5e6e8476e5..57f255b2b3 100644 --- a/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md +++ b/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md @@ -308,7 +308,7 @@ To run MMLU benchmarks and validate the model's performance, follow the instruct To optimize Multi-Head Latent Attention (MLA) performance, you can enable sliced projections. -* **Flag**: `use_sliced_mla_projections` (default: `False`) +* **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`). diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 9fa9e18cd5..aa54a06dc3 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -2098,7 +2098,7 @@ def test_sliced_mla_projections(self): # Enable sliced projections for one config config_arguments_sliced = config_arguments.copy() - config_arguments_sliced["use_sliced_mla_projections"] = True + 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") diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index 5dcd4a174a..06ea7f9ea7 100644 --- a/tests/unit/configs_value_test.py +++ b/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() diff --git a/tests/unit/linears_test.py b/tests/unit/linears_test.py index 0b37ebf12c..a20d29a0d9 100644 --- a/tests/unit/linears_test.py +++ b/tests/unit/linears_test.py @@ -156,6 +156,30 @@ def test_slice_bounds_with_quantization(self): 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 From 6ae6470936646085335edd758dbe843739d30cc4 Mon Sep 17 00:00:00 2001 From: Daniel Mandragona Date: Wed, 5 Aug 2026 17:48:45 +0000 Subject: [PATCH 7/7] Fix code_quality linter errors for sliced MLA projections - Remove superfluous parens around condition after 'not' keyword in src/maxtext/layers/linears.py to fix pylint C0325 and conform to pyink formatting. - Assert loss_normal and loss_sliced are close in tests/unit/attention_test.py to fix pylint W0612 unused variable warnings. TAG=agy CONV=bfd8ec6e-f21b-4a02-8a96-6319fee8ff89 CHANGE_STEWARD=true --- src/maxtext/layers/linears.py | 6 ++---- tests/unit/attention_test.py | 1 + 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py index e75c1ddaf2..915bc9c008 100644 --- a/src/maxtext/layers/linears.py +++ b/src/maxtext/layers/linears.py @@ -295,10 +295,8 @@ def __call__( 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]}]" - ) + 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) diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index aa54a06dc3..569984d6f7 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -2128,6 +2128,7 @@ def loss_fn(model, x): 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))