Add flag-controlled sliced MLA projections by dandragona · Pull Request #4630 · AI-Hypercomputer/maxtext

dandragona · 2026-07-27T20:46:51Z

Performance optimizations that avoids an unnecessary split operation.

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.

Tests

added unit tests
https://xmanager.corp.google.com/experiments/271579010
Before: https://xprof.corp.google.com/trace_viewer/chengnuojin-11765148423170897263?view_start=1400.245&view_end=1420.311
After: https://xprof.corp.google.com/trace_viewer/dandragona-1224952786495301116?view_start=2691.354&view_end=2701.054
After: https://xprof.corp.google.com/overview_page/dandragona-7513819543584336902 (isolation microbench)

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

gemini-code-assist · 2026-07-27T20:46:54Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

codecov · 2026-07-27T20:54:28Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

github-actions · 2026-07-28T17:11:58Z

🤖 Hi @NuojCheng, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions · 2026-07-28T17:16:52Z

🤖 I'm sorry @NuojCheng, but I was unable to process your request. Please see the logs for more details.

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

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

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

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

github-actions · 2026-07-29T22:56:54Z

🤖 Hi @NuojCheng, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions · 2026-07-29T22:59:54Z

🤖 I'm sorry @NuojCheng, but I was unable to process your request. Please see the logs for more details.

NuojCheng · 2026-07-31T17:16:27Z

      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"


we discourage using assertion error in general, since they can be skipped sometimes. We could use if ... raise ValueError() instead

NuojCheng · 2026-07-31T17:28:16Z

+        model_mode=MODEL_MODE_TRAIN,
+    )
+
+    self.assertTrue(jnp.allclose(out_normal_train, out_sliced_train, rtol=1e-05, atol=1e-05, equal_nan=False))


for training we want to compare gradients to protect backward pass correctness

github-actions · 2026-08-01T06:32:10Z

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions

## 📋 Review Summary

The Pull Request introduces an elegant performance optimization for Multi-Head Latent Attention (MLA) projection operations. By slicing weight kernels and biases along the output feature axis prior to contraction, it successfully avoids unnecessary jnp.split operations on sharded tensors, which reduces memory copies and layout transitions. The implementation is highly robust, and the author has written comprehensive unit tests that cover all model execution modes (train, prefill, autoregressive decode) along with rigorous gradient checks.

🔍 General Feedback

Performance-Oriented Design: Slicing the parameters before contraction is a great optimization that reduces network transfer and gather overhead on TPUs/GPUs.
Exemplary Test Coverage: The detailed test assertions for forward outputs, input gradients, parameter gradients, and KV cache states are excellent and ensure zero regressions.
Clear Documentation: The added constraint documentation in Run_DeepSeek.md is precise, well-written, and easy to understand.

github-actions · 2026-08-01T06:39:59Z

    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)


🟡 When `use_sliced_mla_projections` is enabled, the LoRA path (`q_lora_rank > 0`) uses sliced projections to avoid the split operation. However, the non-LoRA path (`q_lora_rank == 0`) still performs a full projection followed by `jnp.split`. We should extend this optimization to the non-LoRA path for completeness and performance consistency under all configurations.

Suggested change

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)

if self.q_lora_rank == 0:

if self.config.use_sliced_mla_projections and self.query.quant is None:

q_nope = self.query(

inputs_q,

out_sharding=query_sharding,

slice_bounds=(0, self.qk_nope_head_dim),

) # [B, L, n_heads, qk_nope_head_dim]

q_pe = self.query(

inputs_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.query(inputs_q, out_sharding=query_sharding)

q_nope, q_pe = jnp.split(q, [self.qk_nope_head_dim], axis=-1)

github-actions · 2026-08-01T06:40:04Z

    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]


🟡 For the `value` sliced projection of `self.wkv_b`, the code uses `wkva_out_sharding` (derived from `key_logical_name`). Although key and value currently share the same logical layout configuration, using the specific `value_logical_name` to construct `wkvb_out_sharding` is more idiomatically correct, prevents future-proofing bugs if key and value shardings diverge, and aligns with the explicit intention of `mla_get_key_value` where both logical names are fetched.

Suggested change

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]

wkva_out_sharding = create_sharding(self.mesh, key_logical_name)

wkvb_out_sharding = create_sharding(self.mesh, value_logical_name)

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=wkvb_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]

github-actions · 2026-08-01T06:40:08Z

+    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
+      kernel = kernel[..., begin:end]


🟢 Adding a validation check to ensure `slice_bounds` are valid (i.e., non-negative, `begin < end`, and within the kernel last dimension boundary) prevents silent dimension mismatch errors or empty sliced tensors down the line if invalid indices are passed.

Suggested change

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

kernel = kernel[..., begin:end]

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]

Please consider.

RissyRan

Thanks! LGTM, please address comments before merge.

RissyRan · 2026-08-01T06:57:55Z

      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:


Shall we add assertion in types.py file to indicate that use_sliced_mla_projections doesn't work with quantization config for early check?

RissyRan · 2026-08-01T07:00:35Z

+
+To optimize Multi-Head Latent Attention (MLA) performance, you can enable sliced projections.
+
+*   **Flag**: `use_sliced_mla_projections` (default: `False`)


I see we put proj as short name in base.yml. Maybe we could make it as use_sliced_mla_proj?

dandragona requested review from RissyRan, bvandermoon, gagika, gobbleturk, jiangjy1982, parambole, shuningjin and suexu1025 as code owners July 27, 2026 20:46

dandragona requested review from A9isha, NuojCheng, SurbhiJainUSC, abhinavclemson, aireenmei, darisoy, dipannita08, hengtaoguo, huytransformer, igorts-git, khatwanimohit, richjames0, shralex, vipannalla and xibinliu as code owners July 27, 2026 20:46

dandragona force-pushed the sliced-mla-projections branch 2 times, most recently from 5a039e3 to 2a881d3 Compare July 28, 2026 16:20

NuojCheng added the gemini-review label Jul 28, 2026

dandragona added 4 commits July 29, 2026 16:00

Add unit test for sliced MLA projections

6fb9d83

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

Document use_sliced_mla_projections flag

650726c

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

dandragona force-pushed the sliced-mla-projections branch from 2a881d3 to fc728c5 Compare July 29, 2026 16:00

NuojCheng added gemini-review and removed gemini-review labels Jul 29, 2026

dandragona mentioned this pull request Jul 30, 2026

Add rope_pairwise config flag to optimize YaRN RoPE #4627

Open

3 tasks

NuojCheng reviewed Jul 31, 2026

View reviewed changes

Address NuojCheng review comments: check gradients and avoid assertions

6cd6dfc

RissyRan added gemini-review and removed gemini-review labels Aug 1, 2026

github-actions Bot reviewed Aug 1, 2026

View reviewed changes

RissyRan approved these changes Aug 1, 2026

View reviewed changes


		To optimize Multi-Head Latent Attention (MLA) performance, you can enable sliced projections.

		* Flag: `use_sliced_mla_projections` (default: `False`)

Uh oh!

Conversation

dandragona commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Tests

Checklist

Uh oh!

gemini-code-assist Bot commented Jul 27, 2026

Uh oh!

codecov Bot commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions Bot commented Jul 28, 2026

Uh oh!

github-actions Bot commented Jul 28, 2026

Uh oh!

github-actions Bot commented Jul 29, 2026

Uh oh!

github-actions Bot commented Jul 29, 2026

Uh oh!

NuojCheng Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

NuojCheng Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Aug 1, 2026

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🔍 General Feedback

Uh oh!

github-actions Bot Aug 1, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 1, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 1, 2026

Choose a reason for hiding this comment

Uh oh!

RissyRan Aug 1, 2026

Choose a reason for hiding this comment

Uh oh!

RissyRan left a comment

Choose a reason for hiding this comment

Uh oh!

RissyRan Aug 1, 2026

Choose a reason for hiding this comment

Uh oh!

RissyRan Aug 1, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dandragona commented Jul 27, 2026 •

edited

Loading

codecov Bot commented Jul 27, 2026 •

edited

Loading