Skip to content

Add rope_pairwise config flag to optimize YaRN RoPE - #4627

Open
dandragona wants to merge 3 commits into
mainfrom
rope-pairwise-opt
Open

Add rope_pairwise config flag to optimize YaRN RoPE#4627
dandragona wants to merge 3 commits into
mainfrom
rope-pairwise-opt

Conversation

@dandragona

@dandragona dandragona commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

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

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

Copy link
Copy Markdown

Caution

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

@codecov

codecov Bot commented Jul 27, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 46.15385% with 14 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/layers/embeddings.py 46.15% 12 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

@dandragona
dandragona force-pushed the rope-pairwise-opt branch 2 times, most recently from 5f37896 to a7aaa78 Compare July 27, 2026 17:09
…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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant