rope: make the CPU reference match the device's angle-row convention - #191
Merged
andrej merged 2 commits intoSep 12, 2026
Merged
Conversation
…vention
design.py's core_body applies one angle row to rows/angle_rows CONSECUTIVE
input rows (row r uses angle row r // (rows/angle_rows)). reference()
used cos.repeat(rep, 1), which tiles the whole angle block instead
(row r uses angle row r % angle_rows) -- the interleaved convention. The
two agree only at angle_rows in {1, rows}, so this was invisible until
something calls reference() at 1 < angle_rows < rows, e.g. llama_npu.py's
prefill RoPE shape (rows=prompt_len*n_heads, angle_rows=prompt_len).
repeat_interleave is the one-line fix.
atassis
force-pushed
the
issue/rope-reference-convention-mismatch
branch
from
September 10, 2026 10:27
696fe2b to
f3ef772
Compare
andrej
requested changes
Sep 11, 2026
Comment on lines
+11
to
+14
| next angle row: row r uses angle row `r // (rows // angle_rows)`. A prior | ||
| version of reference() used `cos.repeat(rep, 1)`, which tiles the whole | ||
| angle block `rep` times (row r uses angle row `r % angle_rows`) -- the | ||
| interleaved convention. The two conventions agree only when angle_rows is |
The module and test docstrings framed reference() against its previous cos.repeat() form, and the inline comment restated the docstring above it. Also black-formats an assert the branch had left unformatted.
andrej
approved these changes
Sep 12, 2026
andrej
enabled auto-merge
September 12, 2026 00:42
Contributor
CI Test Results7f9a47a (2026_09_12_01_49_41) IRON - CI SummaryExamplesiron/applications/llama_3.2_1b
Smalliron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/flm/gemm
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/mha
iron/operators/relu
iron/operators/repeat
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/strided_copy
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Krackan - SmallIRONTested on iron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/flm/gemm
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/mha
iron/operators/relu
iron/operators/repeat
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/strided_copy
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Krackan - ExamplesIRONTested on iron/applications/llama_3.2_1b
Phoenix - SmallIRONTested on iron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/flm/gemm
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/relu
iron/operators/repeat
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/strided_copy
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Phoenix - ExamplesIRONTested on Trend tables omitted, the comment hit GitHub's size limit. Full report in the workflow run. |
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.
Fixes #188.
reference()usedcos.repeat(rep, 1)to stretch anangle_rows-row LUT torowsrows, which tiles the whole block (rowruses angle rowr % angle_rows). The device kernel (design.py'score_body) applies one angle row torows / angle_rowsconsecutive input rows instead (rowruses angle rowr // (rows / angle_rows)). The two conventions only agree atangle_rowsin{1, rows}, so this was invisible until something callsreference()at1 < angle_rows < rows-- exactly the shapellama_npu.py's prefill RoPE uses (rows=prompt_len*n_heads, angle_rows=prompt_len), though nothing currently routes throughreference()at that shape.repeat_interleaveis the direct fix: it repeats each row in place instead of tiling the block.I'm recommending the device kernel as canonical over
reference(), not the other way around -- see the linked issue for why.Added
iron/tests/operators/rope_reference_convention.py: builds a ground truth directly from the device's row-to-angle mapping and checksreference()against it across several(rows, angle_rows)shapes, including the decisive1 < angle_rows < rowscase.Changed
iron/operators/rope/reference.py:cos.repeat(rep, 1)/sin.repeat(rep, 1)->repeat_interleave(rep, dim=0), and the docstring now says which convention and why.Removed
Evidence
Before the fix, the first test's
rows=6, angle_rows=3case shows 4 of 6 rows mismatching the device-convention ground truth (see the linked issue for the exact numbers); afterrepeat_interleave, all shapes tested ((6,3), (8,2), (1024,1), (4,4), (13,13), (12,4)) match exactly.Not touched:
generate_golden_reference()/apply_rope(), whichrope/test.pyactually uses for its device-comparison golden -- that path already encodes the consecutive-row convention correctly via its.transpose(0, 1), so it isn't affected by this bug and this PR doesn't change it.PR Merge Checklist
develcommit and pointing todevel.