[Perf] Optimize rmsnorm/layernorm by cschenjunlin · Pull Request #848 · ROCm/FlyDSL

cschenjunlin · 2026-07-14T08:23:54Z

Motivation

Optimize rmsnorm/layernorm

Technical Details

Test Plan

Test Result

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

…, simplify numeric checks

…ed codes

coderfeli · 2026-07-14T09:17:56Z

 def _to_elem_vec(dtype_str: str, elem_dtype, use_hw_cvt_bf16: bool, y):
    if const_expr(dtype_str == "bf16"):
        if const_expr(use_hw_cvt_bf16):
            return y.to(elem_dtype)


use this by default?

The RNE pack logic was pre-existing and follows the documented split in docs/prebuilt_kernels_guide.md: software RNE packing on gfx942 and hardware cvt_pk_bf16_f32 on gfx950+. So I just kept it.
But I also verified on gfx942 with the default y.to(elem_dtype) path, and didn't observe incorrectness and performance regression. Should I remove the RNE packing path?

coderfeli · 2026-07-14T09:19:07Z

            for _sh_exp in range_constexpr(int(math.log2(WARP_SIZE))):
                off = WARP_SIZE // (2 << _sh_exp)
                peer = w.shuffle_xor(off, WARP_SIZE)
                w = w.addf(peer, fastmath=fm_fast)


use fast math compile flag. not need to set math

sjfeng1999 · 2026-07-14T09:37:54Z

 def _load_scalar(copy_atom, elem_dtype, divided_tensor, index):
    view = fx.slice(divided_tensor, (None, index))
    r = fx.make_rmem_tensor(1, elem_dtype)
    fx.copy_atom_call(copy_atom, view, r)
    return fx.memref_load_vec(r)[0]


If there is only one element to load, use divided_tensor[index] is enough

sjfeng1999 · 2026-07-14T09:38:39Z

 def _store_scalar(copy_atom, elem_dtype, store_dtype, divided_tensor, index, val):
    r = fx.make_rmem_tensor(1, elem_dtype)


The same style applies for this

jhinpan

Multi-model adversarial review of #848

I ran four independent reviewers (Claude Opus 4.8, Claude Sonnet 5, Composer 2.5, GPT-5.5) over the full diff + post-PR files, then synthesized. Inline comments capture the actionable findings; the tags show how many models independently raised each.

Verdict: no CRITICAL correctness / OOB / race bug was found by any of the four models. The trickiest parts were checked and hold up: large-M sub-warp group_reduce_add (THREADS_PER_ROW is always 32 or 64, which divides the 64-lane wavefront, with rows on aligned lane segments), the bf16 RNE bit-packing + even/odd shuffle, the vec8->int8 quant store mapping (out_idx = tid*2 + tile_i*BLOCK_THREADS*2) with write-once per-row YScale, LDS reuse across add->max reductions with body-level barriers, tail masking, and n_float denominators.

Consensus themes (act on):

Test-coverage gaps (all 4 models): bf16-only RMSNorm matrix; untested f32/f16 and >2048 unaligned generic path; N=2048/N=8192 boundaries; f16 LayerNorm fast path never run.
LayerNorm eps not parameterized while RMSNorm now is (3 models).
LayerNorm vectorizes only at exactly N==8192 (2 models).

Behavior changes to confirm / document (2 models / verified):

Loss of the M>8192 dispatch heuristic -> all N<=2048 now uses the large-M kernel regardless of M.
RMSNorm smoothquant XScale dtype contract changed fp32 -> elem_dtype (breaking input ABI).

Lower-priority nits (not posted inline):

Stale RMSNorm module docstring - still says the fast path is N % tile_cols == 0; the actual condition is N % VEC_WIDTH == 0 (2 models).
Redundant arch check (arch == "gfx950") or str(arch).startswith("gfx95") - the first disjunct is subsumed by the second.
Fast paths call block_reduce_add2(dummy, sumsq) where a single block_reduce_add exists, doubling the wave-shuffle work.
Quant store has no explicit clamp before .to(int8); relies on |q| <= 127 from the exact block max (holds in practice, absorbed by the quant_error <= 1 tolerance).
_store_scalar's elem_dtype/store_dtype split is always called with identical arguments.
Weakened quant test assertions (dequant / reconstruction checks removed; only quant_error<=1 + scale_error remain).
layernorm_kernel.py dynamicquant launcher return indentation differs from the RMSNorm twin (works, but reads like a mis-indent).

Overall: careful, well-structured work; mergeable once test coverage is broadened and the two intentional-looking behavior changes are confirmed/documented.

jhinpan · 2026-07-14T17:45:16Z

+            m_s, n_s, dt = [x.strip() for x in p.split(",")]
+            configs.append((int(m_s), int(n_s), dt))
+    else:
+        configs = [


[WARNING - flagged by all 4 reviewers] The default RMSNorm config matrix is bf16-only, and combined with the N<=2048->large-M vs N>2048->main-module dispatch it leaves several rewritten paths unexercised:

No f16/f32 configs, so those branches of every builder (normal, quant, smoothquant, fused-add) never run.

Both N>2048 shapes (3072, 8192) are multiples of VEC_WIDTH, so build_rmsnorm_module's generic scalar fallback (N>2048 and N%8!=0) is never hit.

The quant/fused vec fast path needs N%2048==0, so it is only exercised at N=8192.

The N=2048 dispatch boundary is untested.

Suggest adding an unaligned >2048 shape (e.g. (1024, 3073, "bf16")), one f32 and one f16 shape, and N=2048.

jhinpan · 2026-07-14T17:45:16Z

+    else:
+        configs = [
+            (64, 256, "f32"),  # f32 aligned
+            (32, 128, "f16"),  # f16 aligned


[WARNING - 2 reviewers] The LayerNorm vec fast path requires N == BLOCK_THREADS*VEC_WIDTH*4 (== 8192). This f16 config is N=128, and the only N=8192 config is bf16, so the f16 vectorized fast path (and its store path) is never exercised - only the bf16 RNE-packing branch is. Consider adding (64, 8192, "f16") for the normal / fused / quant / smoothquant variants.

jhinpan · 2026-07-14T17:45:17Z

+    raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')")
+
+
+def build_layernorm_module(N: int, dtype_str: str):


[WARNING - 3 reviewers] This PR makes eps a build-time parameter for the RMSNorm builders (build_rmsnorm_module(N, dtype_str, eps=DEFAULT_EPS)), but every LayerNorm builder still hardcodes the module-level EPS=1e-5 (eps_c = EPS inside each kernel) with no eps argument. That's an inconsistent API for the two families this PR otherwise unifies, and LayerNorm callers cannot honor a non-default eps. Add eps: float = EPS through the LayerNorm builders, or note why LayerNorm is intentionally excluded.

The LayerNorm side is not intentionally excluded. I will also makes eps a build-time parameter for all the LayerNorm builders in the future updates.

jhinpan · 2026-07-14T17:45:17Z

+        # ==================================================================
+        # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4
+        # ==================================================================
+        if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16):


[WARNING - 2 reviewers] All LayerNorm variants (normal / fused / quant) gate the vectorized path on N == BLOCK_THREADS*VEC_WIDTH*4 - i.e. exactly N==8192. Every other aligned width (2048/4096/5120/6144/...) silently falls to the scalar generic path. This is inconsistent with the RMSNorm rewrite (which vectorizes any N % VEC_WIDTH == 0 with tail masking) and undercuts the PR's stated 'vectorized fast paths when N is aligned' goal. Consider generalizing to N % VEC_WIDTH == 0 with a tile loop + masking, mirroring build_rmsnorm_module.

This is also in progress, same as the eps issue, not intentionally designed inconsistent with RMSNorm.

jhinpan · 2026-07-14T17:45:17Z

+
+def build_rmsnorm_module(N: int, dtype_str: str, eps: float = DEFAULT_EPS):
+    if N <= 2048:
+        return _build_rmsnorm_large_m_small_n_module(N, dtype_str, eps=eps)


[WARNING - 2 reviewers] Dropping the M build parameter changed this dispatch from the previous M > 8192 and N <= 2048 to just N <= 2048. Because M is no longer known at build time, every N<=2048 shape now uses the multi-row-per-block _build_rmsnorm_large_m_small_n_module, including small-M. It is functionally correct (rows are masked via row < MIn), but for small-batch decode (e.g. M=1, N=512 -> BLOCK_M=32) most launched rows are dead weight - a plausible perf regression in a perf PR. Please confirm this is intentional, or move the size-aware dispatch to the launcher (m_in is known at launch time).

jhinpan · 2026-07-14T17:45:17Z

-                y = (x * rrms) * g
-                out_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y)
+                    y = (x * rrms) * g
+                    out_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y)


[WARNING - lone reviewer, worth checking] The vectorized paths round f32->bf16 with the manual round-to-nearest-even helper _to_elem_vec (on non-gfx95), while the scalar generic paths use _to_elem_scalar -> y.to(elem_dtype). If .to(bf16) truncates (RTZ) rather than RNE, the fast and generic paths produce systematically different (biased) results for identical input - the exact fast-vs-generic divergence class to avoid. It stays within the bf16 atol=2e-2, so tests won't catch it. Either route the scalar store through the same RNE helper, or confirm/document that .to(bf16) is RNE on the target arches.

jhinpan · 2026-07-14T17:45:17Z

            copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits)
            if const_expr(is_smooth):
-                copy_atom_xs = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), 32)
+                copy_atom_xs = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits)


[WARNING - lone reviewer, verified against the diff] This changes the RMSNorm smoothquant XScale input contract: previously XScale was always read as fp32 (xscale_vec_width=4, BufferCopy128b(...,32) / BufferCopy32b(...,32), idx*2/idx*2+1); now it is read using elem_dtype/elem_bits (same dtype as Input). The tests were updated in lockstep (xscale_dev = (...).to(torch_dtype)) and this aligns RMSNorm with LayerNorm's existing convention, so it is internally consistent - but it is a silent breaking change to the kernel ABI: any out-of-tree caller still passing an fp32 XScale for a bf16/f16 norm will misread bytes (garbage scales, not a crash). Please call this out as a breaking change and/or assert the dtype at the launcher boundary.

jhinpan · 2026-07-14T17:45:17Z

-            Input_buf = fx.rocdl.make_buffer_tensor(Input)
-            Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma)
-            Output_buf = fx.rocdl.make_buffer_tensor(Output)
+        def block_reduce_add2(val0, val1):


[NIT - 3 reviewers] In rmsnorm_large_m_small_n_kernel, block_reduce_add2 (and wave_reduce_add plus the SharedStorage/lds/RED_SLOTS allocation it references) is never called - both the vec and scalar branches reduce via the locally-defined group_reduce_add. This is leftover reduction boilerplate; removing it avoids implying the LDS path is load-bearing.

jhinpan · 2026-07-14T17:45:17Z

+    VEC_TILES = N // VEC_WIDTH if USE_VEC_SMALL_N else 0
+
+    if USE_VEC_SMALL_N:
+        BLOCK_N = 1 << (N - 1).bit_length()


[NIT - 2 reviewers] BLOCK_N, BLOCK_M, THREADS_PER_ROW, and BLOCK_THREADS_SPECIAL are computed identically in both the if USE_VEC_SMALL_N: and else: branches - only NUM_VEC_ROW_ITERS differs. Hoist the shared sizing above the branch and only special-case NUM_VEC_ROW_ITERS.

cschenjunlin · 2026-07-17T15:41:51Z

Thanks much for all the reviews. A quick status/scope update:

This PR is still WIP. The current change set mainly focuses on the RMSNorm path updates; the LayerNorm side is not intentionally excluded, but the corresponding eps parameterization and broader vector-path generalization are still in progress and will be handled in follow-up updates unless reviewers prefer folding them into this PR.

For this PR, I’ll focus on:

Broadening the RMSNorm test matrix to cover f32/f16, N=2048, and an unaligned >2048 generic fallback case.
Cleaning up the RMSNorm large-M small-N path, including the unused reduction boilerplate and duplicated sizing logic.
Documenting/confirming the RMSNorm smoothquant XScale dtype contract.
Fixing stale comments and small low-risk nits.

For the bf16 conversion helper: the manual RNE pack path predates this PR and follows the documented gfx942/gfx950 split. I also verified on gfx942 with cache disabled that the default y.to(elem_dtype) path compiles and passes the RMSNorm bf16 vector-path tests, so I’m checking whether we should simplify that helper or keep the existing documented non-gfx95 software path.

cschenjunlin and others added 20 commits May 28, 2026 22:39

rmsnorm tests: reorganize variants, use helpers to reduce duplication…

d313f43

…, simplify numeric checks

add vector fast paths for layernorm variants

2b28750

refactor: extract scalar/vector load-store helpers to reduce duplicat…

d670743

…ed codes

fix layernorm shared storage name error, use helpers to simplify tests

f6c474c

fix norm xscale dtype issue and extract test helpers

d48c011

fix python code style issue

cb7d5ed

Merge branch 'main' into cjl/norm_optimization

a3800e1

Merge branch 'main' into cjl/norm_optimization

7dd9982

Merge branch 'main' into cjl/norm_optimization

cfa9acd

replace rsqrt API with fmath.rsqrt

947d00b

Merge branch 'main' into cjl/norm_optimization

328e78a

remove M from builder

c1b1211

use fly.compile, and split large shape from default test configs

24f5dd3

fix python code style issue

084527f

Merge branch 'main' into cjl/norm_optimization

a4eafd8

fix python style issue

0e740cf

make eps as build function parameter

0b2e535

WIP rmsnorm: use vec8 loads for aligned bf16/f16 cases

536b0ba

each block process multiple rows in aligned small-N vec8 path

6a55fba

test_rmsnorm: add unaligned cases for tuning

fa4ba4c

coderfeli reviewed Jul 14, 2026

View reviewed changes

coderfeli requested review from jhinpan and sjfeng1999 July 14, 2026 09:20

sjfeng1999 reviewed Jul 14, 2026

View reviewed changes

jhinpan reviewed Jul 14, 2026

View reviewed changes

cschenjunlin changed the title ~~Cjl/norm perf~~ [Perf] Optimize rmsnorm/layernorm Jul 17, 2026

		def _store_scalar(copy_atom, elem_dtype, store_dtype, divided_tensor, index, val):
		r = fx.make_rmem_tensor(1, elem_dtype)

		raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')")


		def build_layernorm_module(N: int, dtype_str: str):

Uh oh!

Conversation

cschenjunlin commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jhinpan left a comment

Choose a reason for hiding this comment

Multi-model adversarial review of #848

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cschenjunlin commented Jul 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

cschenjunlin commented Jul 14, 2026 •

edited

Loading