Add RMSNorm kernel to latency-demo roofline notebook (#165)

[yuhaohaoyu]

Goal: to close issue #165
Belonging Epic: #94

Commit message:

Integrate a parameterized RMSNorm generator and 3-config scaling study (4-core baseline, 32-core strong, 32-core weak) into the multi-kernel roofline sections alongside matmul/softmax/SDPA/paged-attention.

Design choices:

• Unfused standalone kernel — models the HBM pass-through between matmul stages as it appears in prefill.
• 1D grid [N,1] row-partition — prefill has abundant row parallelism; hidden-dim sharding adds allreduce for zero benefit.
• W is 1D [hidden_dim] per PyTorch convention — avoids inflated HBM traffic (8 KB stays LX-resident after one cold miss).

Artifacts for reviewing:

1. Notebook Section 5 and 6 (not cell number) are modified to add RMSNorm to the 2 sets of roofline plots.
2. Ask Claude to produce the pre-run and post-run notebook html files.
3. Attaching the pair of html files below.

latency_demo_after.html
latency_demo_before.html

[lasch]

Summary

• (cleanup) Maintainability M1 hand-rolled access-tile boilerplate instead of the existing _access_tile helper
• (not blocker) Finding F1, silent wrong output when hidden_dim doesn't divide evenly by block_size (this is the same for other kernels and the notebook doesn't exercise any odd cases)

Details

• M1 — reinvents the _access_tile helper the module explicitly exists to avoid. The module's own docstring states: "Private helpers (_mem_view, _access_tile, _indirect_kv_tile) factor out repeated boilerplate." Every sibling generator (gen_matmul_mlir, gen_softmax_mlir, gen_sdpa_mlir) builds its access tiles via _access_tile(name, view, offsets, tile_shape, view_shape). gen_rmsnorm_mlir instead hand-writes four access-tile blocks inline (x_acc, x_acc2, w_acc, y_acc), each reproducing _access_tile's exact access_tile_set/access_tile_order boilerplate by and.
  x_acc and x_acc2 are, byte-for-byte, the same template (_access_tile("x_acc*", "x_view", ["%row", "%col"], [1, bs], [n_rows, hd])) duplicated verbatim between pass 1 and pass 2. This is exactly the kind of divergence the helper was written to prevent — a maintainer extending _mem_view/_access_tile (e.g. to fix an affine-bound convention, or add a dtype) would silently miss this generator's four inline copies.

• F1 — silent wrong output when hidden_dim % block_size != 0 (confirmed, not exercised by the notebook). Both scf.for %col = %c0 to %c_hd step %BLOCK_SIZE loops (pass 1 and pass 2) iterate while col < hidden_dim, but each iteration unconditionally requests a fixed 1×block_size access tile via a hardcoded affine bound -d1 + {block_size - 1} >= 0. When hidden_dim isn't a multiple of block_size, the final iteration's tile extends past the memory view's declared hidden_dim bound. Unlike gen_paged_attention_mlir, which explicitly ceiling-divides (num_tiles = (context_len + block_size - 1) // block_size) to handle a non-divisible trailing block, gen_rmsnorm_mlir has no such guard, no assertion, and no docstring note of the precondition. Not a merge blocker for this PR (every call site uses hidden_dim=4096, block_size=1024, which divides evenly), but it's a latent correctness hazard in reusable generator code that the next person to change RMS_HIDDEN or block_size in the notebook — or reuse gen_rmsnorm_mlir elsewhere — will hit silently. This PR inherits/follows an existing repo-wide gap/pattern rather than introducing a new one.

[yuhaohaoyu]

@lasch Addressed your review feedbacks. Please take a look

---

The full messaging from the semi-auto Claude-session

Address PR #197 review: use _access_tile helper, add divisibility assert

• M1: Replace 4 inline construct_access_tile blocks with _access_tile() calls, consistent with gen_matmul/gen_softmax/gen_sdpa generators.
• F1: Assert hidden_dim % block_size == 0 — documents (in-line) the precondition without adding partial-tile logic to exemplary code.

[lasch]

thx for addressing the comments.
/lgtm

[WarningRan]

Review — RMSNorm notebook kernel (head e8a186c)

Adds gen_rmsnorm_mlir (+113) in notebooks/demo_gen_mlir.py, a
run_kernel_rmsnorm wrapper, and three notebook points. Both suites pass on
e8a186c and the branch merges cleanly. The reported numbers were re-derived by
hand and reconcile exactly — 17,414 compute and 8,192 memory cycles per core for
25,606 total, 4,457,216 FLOPs over 8,388,608 bytes for AI 0.5313 — so the kernel
computes and costs what it claims to.

M1 (hand-rolled access tiles) and F1 (hidden_dim % block_size) are both
resolved at e8a186c, via _access_tile and the precondition assert at
demo_gen_mlir.py:208.

Two things below: one bug that is free to fix, and one theme with two symptoms.

F2 — f16 sum-of-squares overflows into a silent all-zero row

demo_gen_mlir.py:240,247 and :262-268. sum(x^2) over 4096 elements exceeds
f16 range once RMS(x) > sqrt(65504/4096), about 4.0. The sum becomes inf,
rsqrt(inf) is exactly 0, and the row is written as zeros — no NaN, no Inf, no
warning. The same failure shape as F1, one level deeper. At H=4096, x_std=4
zeroes 1 row in 4 and x_std=8 zeroes all of them; standard_normal keeps the
committed demo clear of it, but raising hidden_dim or substituting real
activations does not.

Keeping the accumulator and the reciprocal-sqrt on tensor<1xf32>, and dropping
the tensor.extract to a scalar, fixes it. Staying on the tile rather than
declaring the scalars f32 matters for two reasons: gen_paged_attention_mlir in
this same file already uses f32 splat accumulators behind arith.extf, so it is
the established shape here; and scalar float binops are evaluated in f16 by the
interpreter regardless of the declared MLIR type (#205), so a scalar f32 chain
would read as fixed without being fixed. Verified clean up to x_std=32 at
7.4e-4 max relative error, with cycles and AI bit-identical — extf and
truncf are zero-cost and the SIMD rule counts elements, not bytes. The fix is
free on both axes.

Worth pairing with it: this path has no numerical check, while the matmul path in
the same notebook prints a max-abs-error line against NumPy. Measured max relative
error here is about 1.0% at H=4096 even at std 1 — plausible f16 accumulation
rather than a defect, but nothing currently separates the two, and that assertion
is what would have surfaced the overflow.

F3 — the notebook's prose describes hardware the estimator does not model

One root cause, two sites. The estimator is a charge model, not a hardware model:
it has no cache, reuse or residency mechanism, and per-core bandwidth is chip
bandwidth divided by active cores. Both added points narrate hardware intuition
that the numbers do not come from, and in a file whose product is understanding,
that is the correctness bar — a reader currently leaves with two conclusions that
are true of neither the kernel nor the model.

The residency claim. demo_gen_mlir.py:279. The description states that W
being 1D means "8 KB stays LX-resident after one cold miss". Every ktdp.load
against an HBM view is charged in full, every time, so the kernel's traffic is
four exactly equal quarters of 2,097,152 bytes — pass-1 x, pass-2 x, W, and the
stores. The 8 KB weight vector, loaded inside both loops, is 25% of total HBM
traffic, and AI 0.5313 is about half of reachable: 0.7048 with the W load
hoisted above the row loop, 1.0544 also holding x in LX across the two passes,
against an analytic ~1.0 for a read-x/write-y normalization. Either hoist the load
or reword — "unfused" defends not fusing with neighbouring kernels, not re-reading
x within this one, and the LX sentence describes a mechanism that is not there.

The bottleneck label. Section 4 reports bottleneck=compute for a kernel that
is memory-bound in every textbook sense, unexplained. The SIMD ridge is 0.25 F/B
at 4 cores and 2.0 at 32 — section 6 already prints both — so the same code at the
same AI is compute-bound at 4 cores and memory-bound at 32. One or two sentences
of markdown turns the most confusing number this kernel produces into the most
instructive one; without them the takeaway is backwards.

Design and spec grounding

No design findings — the kernel follows gen_softmax_mlir's structure and the
wrapper mirrors its siblings. Spec-safe, and settled by the ktdp contract in
CLAUDE.md rather than the RFC text, since only existing ops are emitted: index
access-tile element types, 1:1 load/store tile shapes including the 1D W load, no
allocation in construct_memory_view, Arith/Math/LinAlg with SCF control flow,
and every linalg and tensor op used here has precedent in the sibling
generators. No hardware constant or cost coefficient is touched.

Verdict

COMMENT. F2 is the one thing standing between this and merge, and it costs nothing
in cycles or AI with precedent two functions down in the same file; F3's second
site is a markdown sentence. Two choices worth keeping on the next push: routing
the access tiles through _access_tile, which reads better than the sibling
generators do, and the 1D-W-plus-linalg.broadcast shape, which is the right IR
for a PyTorch-convention weight independent of the traffic question above.

Optional, no action needed: rms_2 and rms_3 land on the same roofline point
(identical AI and both throughputs), since 256 rows divide evenly across both 4
and 32 cores. A non-divisible count in one 32-core config — 250 rows over 32 gives
a 1.143 busiest-to-idlest ratio — would make it the notebook's only demonstration
of stride-partition imbalance.

[yuhaohaoyu]

The latest commit addresses the latest comments from @lasch and @WarningRan , with conflict resolution following the merging of sister notebook related PRs.

Changes to address comments:

Added f16-to-f32 and back conversions to address overflow in acc-reduce phase of rmsnorm kernel

1. fix(rmsnorm mlir): promote accumulation-reduction block to f32-based
2. feature(latency): charge SIMD cost for extf, truncf, and splat (reflecting Spyre Rapid Core specs)
3. LX-residency docstring claim
4. explanation for bottleneck=compute at 4 cores

[lasch]

The latest commit addresses the latest comments from @lasch and @WarningRan , with conflict resolution following the merging of sister notebook related PRs.

Changes to address comments:

Added f16-to-f32 and back conversions to address overflow in acc-reduce phase of rmsnorm kernel

1. fix(rmsnorm mlir): promote accumulation-reduction block to f32-based

2. feature(latency): charge SIMD cost for extf, truncf, and splat (reflecting Spyre Rapid Core specs)

3. LX-residency docstring claim

4. explanation for bottleneck=compute at 4 cores

F3, site 2 (WarningRan — the notebook reports bottleneck=compute at 4 cores for a kernel
that is memory-bound "in every textbook sense," unexplained): not resolved, despite the commit
message claiming otherwise. The commit message's item (4) states "explanation for
bottleneck=compute at 4 cores" was added. Diffing notebooks/latency_demo.ipynb between
34ff0b5 and c87cd52 shows exactly one change in the entire notebook: "Run five kernels..."
→ "Run six kernels..." in the Section 4 markdown cell (a leftover from the unrelated RoPE
sister-PR merge that produced the rebase, not new content from this commit). No bottleneck,
ridge-point, or AI explanation was added anywhere in the notebook. This is a genuine discrepancy
between the commit message and the actual diff, not a matter of interpretation.

Your item (2) introduces a new issue (sorry for blasting the full text without summary). The recommendation here would be to extract that into a separate PR. On the other hand, we wanted to wrap up for now and not introduce new loose ends.

• F4 — the bundled latency-model change (item 2) has real grounding and side-effect problems, and
  is out of scope for this PR. arith.extf/arith.truncf move from no latency category (zero
  cost) to LatencyCategory.COMPUTE_FLOAT, and tensor.splat moves from no category to the same.
  Checked against the Spyre knowledgebase (wiki/foundations/hardware/microarchitecture.md,
  wiki/concepts/core-functional-units.md, wiki/foundations/hardware/hardware-generations.md,
  wiki/concepts/mixed-precision-pe-array.md, wiki/concepts/triton-compiler-pipeline.md):
  • Mislabeled unit. The commit message and code comment call this "SIMD pipeline cost," but
    the KB consistently scopes "SIMD" to the PE/MPE systolic matmul array. Type conversions are
    documented as running on the SFP (Scalar Function Processor), a distinct unit —
    core-functional-units.md lists "type conversions" under SFP, and
    hardware-generations.md explicitly excludes SFP throughput from the SIMD/matmul critical
    path ("not included in the table above because it is not on the matmul critical path"). Real
    SFP cycles for extf/truncf are plausible; charging them to the SIMD unit is not
    supported by the KB and, in this codebase, is not cosmetic — ExecutionModel.unit_categories
    (ktir_cpu/latency.py:78-81) puts COMPUTE_FLOAT inside the "simd" unit bucket used for
    bottleneck classification, so this miscategorization can flip which unit a kernel is reported
    as bound by.
  • tensor.splat has no KB support as a compute-cost operation at all. The KB describes
    broadcast mechanisms (RING multicast, XRF weight broadcast) as data movement, and categorizes
    tt.splat/tt.broadcast as "Control — Shape manipulation," not compute
    (triton-compiler-pipeline.md). Charging it under COMPUTE_FLOAT — the same category as
    genuine elementwise arithmetic like addf/mulf — has no documented hardware basis found in
    the KB.
  • AI/roofline inflation, confirmed empirically. latency.py's _estimate() returns
    flops = n_elems for every op categorized COMPUTE_FLOAT (ktir_cpu/latency.py, the
    LC.COMPUTE_FLOAT branch), and chip_roofline() sums exactly the categories in the dominant
    unit's bucket into chip_flops (_roofline_common, chip_roofline). Isolated the effect by
    running the same RMSNorm configuration (hidden_dim=4096, num_cores=4,
    HardwareConfig(lx_size_mb=2, hbm_bandwidth_tb_s=1.024), matching the notebook's own setup
    cell) three ways in the sandbox: pre-fix code — compute=18439, memory=8192, AI=0.5626; F2 fix
    alone (new generator, old un-costed arith_ops.py/tensor_ops.py) — compute=17414,
    memory=8192, AI=0.5313, which matches WarningRan's cited numbers exactly, confirming
    WarningRan reviewed the fix before this costing change existed; F2 fix + costing change
    (the actual c87cd52 head) — compute=23563, memory=8192, AI=0.719. The AI shift from
    0.5313 to 0.719 is caused entirely by extf/truncf/tensor.splat elements now counting as
    "FLOPs" in the numerator, despite performing no floating-point arithmetic. This directly
    contradicts WarningRan's own stated rationale for endorsing the F2 fix pattern — "cycles and
    AI bit-identical... extf and truncf are zero-cost" — which was true when written and is no
    longer true after this same commit.
  • Partial precedent exists, which narrows but doesn't remove the concern. arith.cmpf and
    arith.select were already registered under COMPUTE_FLOAT before this PR
    (ktir_cpu/dialects/arith_ops.py, pre-existing), and a comparison/select is arguably no more
    a "FLOP" than a cast is — so counting non-arithmetic ops toward FLOPs under this category is
    an existing codebase convention, not one invented here. Extending it to tensor.splat
    specifically is the weakest extension of that convention: a broadcast does even less
    computational work than a comparison, and the KB gives it no compute-unit grounding at all.
  • Repo-wide, unreviewed blast radius. arith.extf/arith.truncf/tensor.splat are used
    pervasively outside the new RMSNorm kernel — gen_softmax_mlir, gen_sdpa_mlir, and
    gen_paged_attention_mlir all use tensor.splat for accumulator initialization, and the SDPA/
    paged-attention generators use arith.extf for f16→f32 promotion the same way RMSNorm now
    does. This commit silently changes the reported cycles/AI/bottleneck for all three of those
    already-merged kernels in the same notebook, with no diff, discussion, or before/after
    comparison shown for any of them — only the RMSNorm numbers were the stated subject of this
    push.
  • No test coverage. tests/test_ops.py::test_extf_promotes_f32 and test_truncf_passthrough
    cover numeric correctness only; grepped tests/ for extf/truncf/tensor.splat alongside
    COMPUTE_FLOAT/cycle assertions and found nothing that exercises the new categorization. No
    test would catch a regression or a further change to this costing behavior.

[yuhaohaoyu]

@lasch Namely F3 and F4 are addressed.

---

Suggested fixes:

1. F3 fixed — Re-added (half-sleeply removed from manual rebase) 4 sentences after the RMSNorm paragraph in the Section 4 markdown cell explaining why bottleneck=compute at 4 cores is physically correct (SIMD ridge = 0.25 F/B at 4 cores, AI ~0.53 exceeds it; at 32 cores ridge rises to 2.0 and kernel becomes memory-bound).

2. F4 fixed — Reverted the extf/truncf/splat the COMPUTE_FLOAT categorization:

   • arith_ops.py: merged _FLOAT_CAST_OPS back into _CAST_UNARY_OPS with no latency category
   • tensor_ops.py: removed latency_category=LC.COMPUTE_FLOAT from tensor.splat and the unused LC import

Additional Notebook Prose Fixes

#	Fix	Where
1	"six" → "seven" kernels	cell-15 (Section 4 intro)
2	Removed paged_attn from simd-dominant bullet (it's systolic)	cell-15 (Section 4 intro)
3	Ridge formatting → .2f everywhere (was mix of .0f/.1f/.2f)	demo_helpers.py:85, cell-25, cell-27
4	Section 10 opening: no longer claims "Sections 1–8 all parallelise the same way" (sdpa_decode_pv has non-zero comm)	cell-36 (Section 10)
5	Section 2: fixed "memroy" typo, "her"→"here", missing articles, comma-splices, subject-verb disagreement	cell-10 (Section 2)
6	Section 6: fixed "is most common scenario for a relative weak AI accelerators" grammar	cell-21 (Section 6)
7	Summary: "for verify the correctiness" → "for verifying the correctness"; "ktir mlir" → "KTIR MLIR"	cell-39 (Summary)
8	Section 8: replaced unparseable opening sentence with clear 3-sentence description of the two parser paths	cell-28 (Section 8)
9	TOC links: removed parenthesized text from section titles that broke markdown anchor hrefs	cell-0, cell-28, cell-32

[lasch]

/lgtm