Add rope_pairwise config flag to optimize YaRN RoPE by dandragona · Pull Request #4627 · AI-Hypercomputer/maxtext

dandragona · 2026-07-27T16:26:17Z

Description

Add rope_pairwise config flag to optimize YaRN RoPE computation in YarnRotaryEmbedding.

When enabled (rope_pairwise: true), it uses a pairwise interleaved implementation of RoPE in YarnRotaryEmbedding which keeps the rank-5 pair tensor intact, avoiding some data formatting during MLA.

Tests

Added test_pairwise_call to YarnRotaryEmbeddingTest in tests/unit/embeddings_test.py.
Added test_pairwise_requires_interleave validation test.
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/trace_viewer/dandragona-8379284804318708435?view_start=109105.685&view_end=109109.290 (isolation mircrobench)

Checklist

I have performed a self-review of my code.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests or unit tests.

gemini-code-assist · 2026-07-27T16:26:21Z

Caution

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

codecov · 2026-07-27T16:37:56Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

github-actions · 2026-07-28T17:13:10Z

🤖 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:17:35Z

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

github-actions · 2026-07-28T22:07:46Z

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

…f YarnRotaryEmbedding. This addresses the Codecov patch coverage warnings. CHANGE_STEWARD=true

Add documentation for the new rope_pairwise optimization flag in the DeepSeek runner guide. TAG=agy CONV=a262e292-a0d4-49f9-ad26-f42e697a5979 CHANGE_STEWARD=true

The test was incorrectly expecting concatenated output for pairwise=True, but YarnRotaryEmbedding with pairwise=True returns interleaved output. Update expectation to match interleaved layout. TAG=agy CONV=a0af6c58-a54e-4ea7-8d04-fac29de8aa41 CHANGE_STEWARD=true

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

🤖 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-29T23:00:28Z

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

NuojCheng

Could you add a performance comparison, e.g. before/after, in the PR description?

dandragona · 2026-07-30T16:47:29Z

Could you add a performance comparison, e.g. before/after, in the PR description?

Done. Also updated for #4630

…ing in pairwise RoPE Addresses GitHub review comment #discussion_r3684287727 by: - Explicitly passing interleave=True in test_pairwise_call - Adding test_pairwise_explicit_shard_mode_call for interleave=True, pairwise=True, shard_mode=EXPLICIT - Annotating broadcasted cos/sin with 5D explicit sharding in rope_pairwise when shard_mode=EXPLICIT CHANGE_STEWARD=true CONV=fd71583d-3a4c-418b-a131-d086fa8e5c9d

CHANGE_STEWARD=true CONV=fd71583d-3a4c-418b-a131-d086fa8e5c9d

Addresses GitHub review comment #discussion_r3684262252 by updating the return type annotation from "Rope" to "YarnRope". CHANGE_STEWARD=true CONV=70d61320-1190-44c5-a1c8-f64aadb81397

github-actions · 2026-08-01T07:07:07Z

🤖 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 the rope_pairwise configuration flag, which optimizes YaRN RoPE computations in YarnRotaryEmbedding by implementing a highly efficient, real-valued pairwise interleaved layout. This avoids complex number creation and redundant tensor reshaping, resulting in significant performance improvements. The overall design is sound, mathematically correct, and includes excellent unit test coverage for various edge cases.

🔍 General Feedback

Performance Optimization: Avoiding the creation of complex numbers and keeping the rank-5 pair tensor intact is a great performance enhancement for MLA computations.
Robust Verification: The unit tests comprehensively cover the new path, including checking validation limits (e.g., rope_pairwise requiring rope_interleave) and explicit sharding mode behaviors.
API Consistency: Suggesting a minor adjustment to the initialization error message to match the class parameters rather than configuration flags.

github-actions · 2026-07-28T22:07:33Z

+    outputs = layer(inputs, position=position)
+    self.assertEqual(outputs.shape, (1, 2, 1, 4))
+
+    # Output with pairwise=True should match default interleaved YaRN RoPE


🟢 The comment states that the output with `pairwise=True` should match "default interleaved YaRN RoPE". However, the default implementation actually returns a concatenated layout (as tested in `test_basic_call`), whereas `pairwise=True` returns an interleaved layout. This is why the expected outputs have their second and third elements swapped. Clarifying this in the comment will help prevent future developer confusion.

Suggested change

# Output with pairwise=True should match default interleaved YaRN RoPE

# Output with pairwise=True matches default YaRN RoPE but with interleaved output layout [real0, imag0, real1, imag1]

github-actions · 2026-07-28T22:07:33Z

+        cos = jnp.real(freqs)[..., jnp.newaxis]
+        sin = jnp.imag(freqs)[..., jnp.newaxis]
+        swapped = jnp.flip(pairs, axis=-1)
+        sign = jnp.asarray([-1.0, 1.0], dtype=jnp.float32)


🟢 For constructing fixed literal arrays in JAX, `jnp.array` is more conventional and concise than `jnp.asarray`. Additionally, matching the dtype dynamically with `cos.dtype` is more robust and prevents hardcoded dtype mismatches if precision changes in the future.

Suggested change

sign = jnp.asarray([-1.0, 1.0], dtype=jnp.float32)

sign = jnp.array([-1.0, 1.0], dtype=cos.dtype)

github-actions · 2026-07-28T22:07:33Z

+  )
+
+  @model_validator(mode="after")
+  def validate_rope_pairwise(self) -> "Rope":


🟡 The return type of `validate_rope_pairwise` is annotated as `"Rope"`, but it actually returns `self` (an instance of `YarnRope`). Since `YarnRope` and `Rope` are distinct configuration classes, this annotation is incorrect and can trigger static analysis/type checking errors.

Suggested change

def validate_rope_pairwise(self) -> "Rope":

def validate_rope_pairwise(self) -> "YarnRope":

github-actions · 2026-08-01T07:11:53Z

+    self.pairwise = pairwise
+
+    if self.pairwise and not self.interleave:
+      raise ValueError("rope_pairwise=True requires rope_interleave=True.")


🟡 The parameters of the `YarnRotaryEmbedding` class are named `pairwise` and `interleave`. Raising an exception referencing the config flag names (`rope_pairwise` and `rope_interleave`) is confusing for developers using the Python API directly.

Suggested change

raise ValueError("rope_pairwise=True requires rope_interleave=True.")

raise ValueError("pairwise=True requires interleave=True.")

RissyRan

LGTM, minor comments. Thanks!

RissyRan · 2026-08-01T07:11:52Z

 rope_interleave: true # RoPE with sin/cos interleaved vs concatenated
 rope_truncate: true # Floor lower bound and ceil upper bound for correction range
 rope_attention_scaling: false # Scale the rotary embedding output
+rope_pairwise: false # Keep rank-5 pair tensor intact and return interleaved RoPE


Shall we make a comment to explain why is rank-5?

i.e. tensor shape with [Batch size, Sequence length, Attention heads, half_dim, 2]

RissyRan · 2026-08-01T07:19:15Z

+    self.assertEqual(outputs.shape, (1, 2, 1, 4))
+
+    # Output with pairwise=True should match default interleaved YaRN RoPE
+    expected = jnp.array([[[[1.0, 1.0, 1.0, 1.0]], [[-0.300781, 1.38281, 0.996094, 1.00781]]]])


Shall we compare with default implementation if possible, i.e. with pairwise=False? Similar comments for other tests if applies.

embeddings.YarnRotaryEmbedding( embedding_dims=4, mesh=self.mesh, max_position_embeddings=16384, original_max_position_embeddings=4096, interleave=True, pairwise=False, rngs=self.rngs, )

dandragona requested review from darisoy and dipannita08 as code owners July 27, 2026 16:26

dandragona force-pushed the rope-pairwise-opt branch from 3bd0253 to 05da824 Compare July 27, 2026 16:30

dandragona force-pushed the rope-pairwise-opt branch 4 times, most recently from 3b47905 to 61a5e35 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 rope_pairwise config flag to optimize YaRN RoPE

61da5f8

Add unit tests to cover non-interleaved and explicit sharding modes o…

dd2fd58

…f YarnRotaryEmbedding. This addresses the Codecov patch coverage warnings. CHANGE_STEWARD=true

Document rope_pairwise flag

d002698

Add documentation for the new rope_pairwise optimization flag in the DeepSeek runner guide. TAG=agy CONV=a262e292-a0d4-49f9-ad26-f42e697a5979 CHANGE_STEWARD=true

dandragona force-pushed the rope-pairwise-opt branch from 61a5e35 to fc5a91e Compare July 29, 2026 16:00

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

NuojCheng reviewed Jul 30, 2026

View reviewed changes

Comment thread src/maxtext/configs/types.py Outdated

NuojCheng reviewed Jul 30, 2026

View reviewed changes

Comment thread tests/unit/embeddings_test.py

NuojCheng reviewed Jul 30, 2026

View reviewed changes

dandragona added 3 commits July 30, 2026 17:01

style: fix pyink formatting in embeddings.py

78c3e3a

CHANGE_STEWARD=true CONV=fd71583d-3a4c-418b-a131-d086fa8e5c9d

Fix return type annotation of validate_rope_pairwise in YarnRope

d992980

Addresses GitHub review comment #discussion_r3684262252 by updating the return type annotation from "Rope" to "YarnRope". CHANGE_STEWARD=true CONV=70d61320-1190-44c5-a1c8-f64aadb81397

NuojCheng approved these changes Jul 30, 2026

View reviewed changes

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

github-actions Bot reviewed Aug 1, 2026

View reviewed changes

RissyRan reviewed Aug 1, 2026

View reviewed changes

	# Output with pairwise=True should match default interleaved YaRN RoPE
	# Output with pairwise=True matches default YaRN RoPE but with interleaved output layout [real0, imag0, real1, imag1]

	sign = jnp.asarray([-1.0, 1.0], dtype=jnp.float32)
	sign = jnp.array([-1.0, 1.0], dtype=cos.dtype)

	def validate_rope_pairwise(self) -> "Rope":
	def validate_rope_pairwise(self) -> "YarnRope":

	raise ValueError("rope_pairwise=True requires rope_interleave=True.")
	raise ValueError("pairwise=True requires interleave=True.")

Uh oh!

Conversation

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

Uh oh!

Description

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 28, 2026

Uh oh!

github-actions Bot commented Jul 29, 2026

Uh oh!

github-actions Bot commented Jul 29, 2026

Uh oh!

Uh oh!

Uh oh!

NuojCheng left a comment

Choose a reason for hiding this comment

Uh oh!

dandragona commented Jul 30, 2026

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 Jul 28, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 28, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 28, 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 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