Repository navigation
fix: apply DSv4 RoPE from row-aligned frequencies - #219
Merged
Merged
Conversation
_apply_mla_rope requires freqs row i to already hold token i's position, then forwarded cu_seqlens into Megatron's packed RoPE helper. That helper rebuilds positions from cu_seqlens and, when context parallel is on, assumes a zigzag split. DSv4 packs with a contiguous split, so those positions are wrong and the rotation is applied to the wrong rows. The packed lengths are no longer forwarded. The elementwise path uses the frequency row that was stored for that token. Callers can still pass cu_seqlens; a length mismatch still raises.
hjh0119
approved these changes
Oct 9, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
_apply_mla_roperequires frequency row i to already hold token i's position. It then forwardedcu_seqlensinto Megatron's packed RoPE helper. That helper rebuilds positions from the cumulative lengths and, when context parallel is greater than 1, assumes a zigzag split. DeepSeek-V4 packs with a contiguous split, so the rotation is applied to the wrong rows.The packed lengths are no longer forwarded. The elementwise path uses the frequency row stored for that token. Callers can still pass
cu_seqlens. A row-count mismatch still raises.Checked with
python3 -m pytest tests/test_dsv4_mla_rope.py(2 passed). The test loads the function and checks that a packedcu_seqlensreachesapply_rotary_pos_embas None, and that a shorter frequency table still fails the row-alignment check.flake8is clean on the changed files.