[BUG] RoPE's CPU reference and device kernel disagree on angle-row mapping for 1 < angle_rows < rows

[atassis]

Describe the bug

iron/operators/rope/reference.py's reference() implements a different angle_rows-to-row mapping than the device kernel it's meant to check.

design.py's core_body acquires one angle row and applies it to rows / angle_rows consecutive input rows before moving to the next angle row -- row r uses angle row r // (rows / angle_rows). reference() instead does cos = cos.repeat(rep, 1) (rep = rows // angle_rows), which tiles the whole angle_rows-row block rep times -- row r uses angle row r % angle_rows. The two mappings agree only when angle_rows is 1 or rows; for any 1 < angle_rows < rows they disagree on most rows, silently (no error, just a wrong expected value).

This is live-shaped, not a corner case invented for the report: llama_3.2_1b/llama_npu.py's prefill RoPE calls are exactly this shape (RoPE(rows=prompt_len*n_heads, angle_rows=prompt_len, ...)). It hasn't fired because nothing currently calls RoPE.reference() at that shape -- rope/test.py builds its golden through generate_golden_reference()/apply_rope() instead (a different code path, indexed straight off the full cos/sin table, with an explicit .transpose(0, 1) to get the position-major layout right) and llama's prefill RoPE only uses the compiled callable, never .reference(). So this is one PR away from firing: anyone who calls RoPE(...).reference(x, angles) to check a batched (1 < angle_rows < rows) run gets a wrong expected answer with no error.

To Reproduce

import torch
from iron.operators.rope.reference import reference

rows, angle_rows, cols = 6, 3, 4
half = cols // 2
x = torch.randn(rows, cols).to(torch.bfloat16)
angles = torch.zeros(angle_rows, cols, dtype=torch.bfloat16)
angles[:, 0::2] = torch.rand(angle_rows, half).to(torch.bfloat16)
angles[:, 1::2] = torch.rand(angle_rows, half).to(torch.bfloat16)

# What the device kernel actually does (design.py's core_body): row r uses
# angle row r // (rows // angle_rows) -- consecutive rows share an angle row.
device_angle_idx = [r // (rows // angle_rows) for r in range(rows)]
print("device convention:   ", device_angle_idx)   # [0, 0, 1, 1, 2, 2]

# What reference()'s cos.repeat(rep, 1) actually selects per row.
cos = torch.arange(angle_rows).reshape(angle_rows, 1).expand(angle_rows, half).float()
ref_angle_idx = cos.repeat(rows // angle_rows, 1)[:, 0].long().tolist()
print("reference() convention:", ref_angle_idx)    # [0, 1, 2, 0, 1, 2]

# 4 of 6 rows use a different angle row -> reference(x, angles, rows=rows, cols=cols)
# disagrees with the device on 4 of 6 output rows.

Output on devel (deb6e1e7cd1e3aefcf7530108a4647a1d967ce18):

device convention:    [0, 0, 1, 1, 2, 2]
reference() convention: [0, 1, 2, 0, 1, 2]

Running the full reference(x, angles, rows=rows, cols=cols) and comparing row-by-row against a ground truth built from the device convention: 4 of 6 rows differ.

Expected behavior

reference() should select the same angle row per input row as design.py's core_body: consecutive rows, r // (rows // angle_rows).

Additional context

I'm fairly confident the device kernel is the intended semantics, not reference(): it matches design.py's own docstring ("each row of angles will be reused for rows / angle_rows consecutive rows of the input tensor"), it's what llama_npu.py's real prefill usage assumes, and rope/test.py's own golden generation independently encodes the same consecutive-row convention via its .transpose(0, 1). On that basis I've also prepared a PR fixing reference() (cos.repeat -> cos.repeat_interleave) rather than filing this as an open question -- happy to drop it if the intended convention is actually the other way around.