RoPE kernel for latency test: issue #166 (issue #94), standalone 8-core (4x2)

[yuhaohaoyu]

Issues addressiog

• To close create rope_fwd_ktir.mlir #166
  • Part of epic [Docs] Add latency estimation demo notebook #94

What's done?

Add standalone 8-core (4x2) RoPE kernel (issue #166) aligning with Granite-8b scale

• Grid [4,2], TILE_SEQ=256: head loop (20 iter) × seq-tile loop (4 iter).
• Half-layout RoPE (LLaMA convention): y[0:D/2] = x[0:D/2]*cos - x[D/2:D]*sin; memory-bound at AI ~0.63 FLOPs/byte.
• Tests: correctness, latency scaling (memory/simd/comm), MLIR frontend adapt.

[yuhaohaoyu]

Notes on addressing a loop-interchange rewriting of the standalone RMSNorm kernel.

Lars left a comment somewhere about a potential poor-forming of the mlir file.
"
Claude has a performance suggestion about your PR 194:

The loop nest is head-outer (20 iterations) / seq-tile-inner (4 iterations). `cos`/`sin`
depend only on seq position, not on head — but because head is the outer loop, each
`[256, 64]` cos/sin tile gets reloaded once per `(head, seq-tile)` pair, i.e. 80 loads total,
instead of once per seq-tile (4 loads total) if the loop nest were flipped (seq-tile outer,
head inner), which would let one cos/sin load serve all 20 heads for that tile — a 20×
reduction in cos/sin HBM traffic with no algorithmic change.

Does that have any grounding in reality?
"

---

In Gist: It translates that 'loop interchange' is the obvious path to avoid unnecessary reload sin-cos constants for RMSNorm.

---

Actions: update the kernel with following design thoughts:

interchanging those 2 loops is convincing now

• now: loop-over attention-heads; loop-over sequence-blocks; (re)load sin/cos; ... endloop; endloop
• interchanged: loop-over seq-blocks; (re)load sin/cos; loop-over attn-blocks; ... endloop; endloop.

---

I wrongfully assumed that MLIR has a way to pass part of the array to inside the loop. There was 'preload the needed sin/cos to my core' outside the loop nest. Yet MLIR can not link that load to the 'slicing-relayout/reload' of cos/sin.

Now will follow your suggestion,
1. remove the preload, (not useful when we have the current inner-loop not inlined.)
2. interchange loop
3. hoist the sin/cos load out of the inner loop.

[lasch]

/lgtm

Optional item:

G1 (minor, cheap). test_rope_fwd_zero_input provides little independent signal: with x = 0, both y_first and y_second reduce to 0 * cos ± 0 * sin = 0 regardless of whether cos/sin are loaded or indexed correctly, so this test would still pass even if the cos/sin addressing were completely broken. The main test_rope_fwd_correctness test already covers the general case well via random input compared against a full-tensor NumPy reference across all 8 cores, so this isn't a large gap — but a more targeted edge case (e.g. an input that isolates the sin/cos cross-term, or a boundary-head/boundary-tile-specific check) would add more real signal than the zero-input check does today.