Add RoPE kernel to latency-demo roofline notebook (#167)

[yuhaohaoyu]

Close: issue #167
Belonging Epic: #94
Independent from sibling issues #164 #165 #166 and the addressing PRs.

---

Notes for Reviewing:

Integrate a parameterized RoPE forward generator and 3-config scaling study (4-core [2,2], 32-core strong [8,4], 32-core weak [8,4]) into the multi-kernel roofline sections alongside matmul/softmax/SDPA/paged-attn.

Design choices:

• 2D grid [seq_parts, head_groups] — aligns with in-projection matmul and SDPA in the prefill pipeline (no tensor redistribution at kernel boundaries).
• Half-layout rotation (LLaMA/Granite convention), cos/sin precomputed as inputs — loaded once per seq-tile, reused across all heads.
• Single pass, no reduction — AI ~0.73, memory-bound SIMD at 32 cores, embarrassingly parallel (zero communication).

---

1. Notebook Section 5 and 6: RoPE appears on per-core and chip-wise roofline plots for SIMD unit.
2. Kept pre- and post- run notebooks in html here.

latency_demo_rope_before.html
latency_demo_rope_after.html

[lasch]

Same issues detected as previously in #197

• reinvention of _access_tile helper
• silent wrong output when seq_len doesn't divide evenly by grid_s * tile_seq, or
  num_heads doesn't divide evenly by grid_h
  • optional, since all other kernels follow the same pattern, since the case is not exercised in the notebook

[WarningRan]

Verified against f07c54ed with the notebook's own HardwareConfig(lx_size_mb=2, hbm_bandwidth_tb_s=1.024) — 1024 B/cycle chip, 64 F/cycle per core.

Spec-safe, and correct where exercised. All six new access tiles use index element type, identity access_tile_order, affine sets, <HBM> memory space, and 1:1 load/store data-tile shapes — conforms to the checklist in CLAUDE.md, no new ops or attributes. Views are non-distributed and the tail is handled by loop-count truncation, so neither the overlapping-coordinate-set caveat nor unspecified out-of-bounds behaviour is in play. Driven against a NumPy reference at H=4, S=512, D=128, grid=[2,2]: max abs error 3.6e-3, zero untouched rows. h_abs*S + cos_row matches the [H,S,D]→[H*S,D] flattening and cos/sin are indexed by absolute position, so grid_s > 1 is right. 40 = 32 Q + 8 KV is the right head count for Granite GQA.

Agreeing with lasch on both of their points rather than restating them; two additions below.

---

1. demo_gen_mlir.py:204 — the docstring's AI ~0.77 contradicts this PR's own commit message (AI ~0.73), and 0.77 is unreachable. Per seq-tile per core with h = heads_per_core and E = tile_seq·D/2: bytes = (8h+4)·E, flops = 6h·E, so AI = 3h/(4h+2) with supremum 0.75 as h → ∞. Measured 0.7317 (4-core) and 0.7143 (32-core), matching to four decimals. The same line also drops the "at 32 cores" qualifier the PR body has: compute/memory = 16·AI / cores_active, so this kernel is compute-bound below ~12 active cores — including Section 4's own 4-core config. Suggest AI = 3h/(4h+2) → 0.71–0.74 here; memory-bound above ~12 active cores.

2. Adding to lasch's divisibility point — the cliff is one constant away from Section 4. ROPE_SEQ = 1024 (the value SM_WIDTH already uses two cells up) gives seq_per_core = 128 < tile_seq, so num_seq_tiles = 0, no core records work, and it dies as ZeroDivisionError: float division by zero at ktir_cpu/latency.py:733 — a message naming neither RoPE nor the grid. That the failure surfaces there is a library-side gap, not this PR's to fix; a ValueError (not assert — those compile out under -O) at the head of gen_rope_mlir is what makes it diagnosable. The non-divisible cases are worse than the crash: H=40, grid_h=3 leaves 512/20480 rows unwritten and still reports kernel_cycles = 30,336, which plots as a perfectly normal point.

3. The 32-core grid gives up the locality the section calls the point of the kernel — a question about intent, not a request. The cos/sin reuse group is exactly heads_per_core, so grid_h=4 cuts it from 40 heads to 10; and since AI is independent of tile_seq, seq_len and head_dim, the head axis is the only one that costs anything to shard:

32-core grid	heads_per_core	core_AI	cycles	chip DRAM
[8,4], tile_seq=256 (this PR)	10	0.7143	116,736	88.1 MB
[32,1], tile_seq=128	40	0.7407	113,664	84.9 MB

The stated reason for [seq, heads] is pipeline alignment, which is real but inter-kernel — and this notebook runs each kernel standalone, so Sections 4–6 can't show the redistribution being avoided, only the intra-kernel metric, on which [8,4] is the worse of the two. Which is intended? If it's the pipeline-realistic grid, the section text is where to say so, since a reader who takes "the key locality optimization" at face value won't expect the 32-core config to dismantle three quarters of it. The same effect explains why AI drops from 0.7317 at 4 cores to 0.7143 at 32 — a nice result for the roofline story if it's stated deliberately.

4. No test. CONTRIBUTING.md asks for tests with new functionality; run_kernel_rope returns only a LatencyReport, so out_ptr is never checked, and the notebook ships with outputs: [] — after merge the repo records neither these numbers nor a way to recompute them. ~15 lines would close it: a NumPy reference at small divisible parameters, plus the per-unit ledger ((8h+4)·E bytes, 6h·E flops per seq-tile per core) rather than the aggregate cycle count, which passes for the wrong reason as soon as two terms move in opposite directions — and is what would have caught the 0.77.

Smaller things:

• demo_gen_mlir.py:250-310 — the six inline construct_access_tile blocks are structurally what _access_tile emits (24 lines → six two-line calls), in a function that already calls _mem_view four times. Likely just that the helper emits at 4-space indent while these sit two scf.for levels deep at 8 — valid either way, so the helper still applies. (lasch's first point.)
• Section 4 defines ROPE_TILE_SEQ but doesn't pass it; only Section 5's rope_1 does. They agree today only because 256 is the default, and tile_seq is one of the two knobs in item 2.
• Section 4 markdown gives only y[0:D/2] = x[0:D/2]*cos - x[D/2:]*sin. The missing second half matters in a section about AI — from the one line shown a reader counts 3 flops per element-pair, not 6.
• The AI figure is only stick-independent at head_dim ≥ 128: at D=128 f16 a tile row is exactly one 128 B stick, so packed == stick-granular. At D=64/32 both halves of a row share a stick, x/out traffic doubles, and reported AI drops to 6h/(16h+4) = 0.3529. Correct modelling, but worth a word since the docstring reads as if the number were intrinsic to RoPE.
• A three-row before/after table in the PR body (AI, cycles, bottleneck) would save opening the HTML attachment to find them.

Verdict — COMMENT. No correctness defect on the exercised path, and nothing touching spec-defined semantics. The one line standing between this and merge is the docstring's AI ~0.77, which the commit message already contradicts with the right value; the guard and the test are the cheap follow-ons, and item 3 is a question about intent.

[yuhaohaoyu]

Notes for re-review :) 6fe6ef5 commit. Attn: @lasch @WarningRan

Purpose

• addressing Lars and Ran's review comments

Gists and rough Actions for review comments

• The Effort column is only of referencing use.
• L : Lars; R : Ran . total 7 review points.

#	Issue	Effort	Action
L1	Reinvention of _access_tile helper — 6 inline blocks should use existing helper	Small (~30 min)	Refactor gen_rope_mlir to call _access_tile()
L2	Silent wrong output on non-divisible dims	= R2	Covered by R2 below
R1	Docstring says AI ~0.77, correct is ~0.73 (supremum 0.75)	Trivial (2 min)	Fix string in demo_gen_mlir.py:204
R2	ZeroDivisionError crash when seq_len doesn't divide by grid_s * tile_seq or num_heads doesn't divide by grid_h	Small (~20 min)	Add ValueError guard at top of gen_rope_mlir
R3	Grid [8,4] vs [32,1] — question about intent	Zero code	Add docstring/comment explaining pipeline-alignment rationale
R4	No test — CONTRIBUTING.md requires tests	Medium (~45 min)	New tests/test_nb_run_kernels.py with NumPy reference + per-unit watermark
R5	Smaller: _access_tile reuse, ROPE_TILE_SEQ passthrough, incomplete markdown formula, stick-independence caveat	Small (~30 min)	4 minor edits; L1 overlap

Review Notes and Code Change Locations

Hope this helps the communication.

#	Review Point	File	Lines
R1	Fix AI docstring (~0.77 → formula)	notebooks/demo_gen_mlir.py:234-236	AI = 3h/(4h+2) ≈ 0.71–0.74...
R2/L2	Add ValueError divisibility guards	notebooks/demo_gen_mlir.py:244-250	if H % grid_h != 0 / if S % (grid_s * tile_seq) != 0
L1	Replace 6 inline access-tile blocks with _access_tile helper	notebooks/demo_gen_mlir.py:261-275	cos_acc, sin_acc, x_first_acc, x_second_acc, y_first_acc, y_second_acc
R3	Document grid [seq,heads] pipeline-alignment intent	notebooks/demo_gen_mlir.py:228-231	Docstring: "2D grid... aligns with in-projection matmul and SDPA..."
R5	tile_seq passthrough + formula completion	latency_demo.ipynb cell 16 (adds tile_seq=ROPE_TILE_SEQ) + cell 15 markdown (both halves of RoPE formula) + notebooks/demo_gen_mlir.py:236 (stick note)
R4	Tests: AI watermark, per-core structure, divisibility guard	tests/test_nb_run_kernels.py:1-85	3 test classes, 9 test cases

Test Groups from R4

1. AI watermark - at small divisible params (H=4, S=512, D=128, grid=[2,2]):
   • Assert AI matches formula: 3*h/(4*h+2) where h = heads_per_core = H // grid_h
2. verify per_core_summary consistency:
   • report.per_core_summary() gives compute_cycles and memory_cycles per core
   • Assert memory-boundness, load-balance across cores, no core-2-core comm.
3. Divisibility guard — assert ValueError raised for bad inputs.

Verification steps in the agentic process.

1. uv run pytest tests/test_nb_run_kernels.py -v
2. uv run pytest tests/ -v — full regression
3. Re-run notebook cells for Sections 4–6 to verify output unchanged (except corrected AI string)
4. Save the full notebook (pre-run) and executed notebook (post-run) as rope-nb-in.html and rope-nb-out.html

Additional steps of verification

1. reviewed the Rope related contents in the rope-nb-out.html file (attached here)
   rope-nb-in.html
   rope-nb-out.html

2. Q/A --> RoPE has no concern with numerical overflow

"
A4: Correct. RoPE has no accumulation — each output element is exactly a*b - c*d or a*b + c*d (two multiplies and one add/sub). There is no reduction, no softmax denominator, no running sum over a sequence dimension. The error is bounded by a few ULPs of fp16 arithmetic per element with no amplification path. No reason to add an accuracy check for RoPE.
"

[lasch]

/lgtm
thx for addressing the comments