Repository navigation
[Perf] Optimize rmsnorm/layernorm - #848
Conversation
jhinpan
left a comment
jhinpan
left a comment
There was a problem hiding this comment.
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/f16and>2048unaligned generic path;N=2048/N=8192boundaries; f16 LayerNorm fast path never run. - LayerNorm
epsnot 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>8192dispatch heuristic -> allN<=2048now uses the large-M kernel regardless of M. - RMSNorm smoothquant
XScaledtype 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 isN % 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 singleblock_reduce_addexists, doubling the wave-shuffle work. - Quant store has no explicit clamp before
.to(int8); relies on|q| <= 127from the exact block max (holds in practice, absorbed by thequant_error <= 1tolerance). _store_scalar'selem_dtype/store_dtypesplit is always called with identical arguments.- Weakened quant test assertions (dequant / reconstruction checks removed; only
quant_error<=1+scale_errorremain). layernorm_kernel.pydynamicquant launcherreturnindentation 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.
|
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:
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 |
Thread caller-provided epsilon through dynamic, smooth, and fused-add quant builders, with coverage that detects silently hard-coded defaults.
Vectorize partial large-N tiles and small-N prefixes with scalar tails while preserving mixed weights and rstd behavior. Add boundary and source-IR coverage for the new routing. Co-authored-by: Cursor <cursoragent@cursor.com>
Thread caller-provided epsilon through fused and quantized forward builders, with tests that detect silently hard-coded defaults. Co-authored-by: Cursor <cursoragent@cursor.com>
Use masked vector iterations for aligned f16 and bf16 row widths, including a partially active final iteration in the test matrix. Co-authored-by: Cursor <cursoragent@cursor.com>
3e569d9 to
0f2d86c
Compare
Isolate scalar, single-tile, and multi-tile XScale handling so Navi CI can identify the failing path without running the quarantined benchmark matrix. Co-authored-by: Cursor <cursoragent@cursor.com>
Run the full plain SmoothQuant matrix on Navi now that targeted scalar and vec8 diagnostics pass. Co-authored-by: Cursor <cursoragent@cursor.com>
Exercise the historical f32 shape and adjacent scalar boundaries so gfx1201 failures distinguish output, scale, and fused-residual corruption. Co-authored-by: Cursor <cursoragent@cursor.com>
Keep the exact regression matrix as expected-failure diagnostics while preventing the known incorrect path from blocking unrelated Navi coverage.
Document the activation-dtype XScale ABI and gfx1201 quarantine so callers do not rely on an unsupported path. Co-authored-by: Cursor <cursoragent@cursor.com>
Use vec8 for the aligned prefix and scalar copies only for the tail, avoiding whole-row scalar fallback. Co-authored-by: Cursor <cursoragent@cursor.com>
block_reduce_max (rmsnorm/layernorm quant) and block_reduce (softmax) store their per-wave partials into a scratch buffer that the preceding reduction is still broadcast-reading on exit. Nothing ordered that read against the new store, so a lagging wave could observe the next reduction's value and quantize a row with the wrong statistic. That race is the source of the intermittent gfx1201 RMSNorm/LayerNorm quant mismatches in CI: the window is visible in the generated ISA on both gfx1201 and gfx942, and a host-side model of the corrupted scale matched 539/539 observed bad rows. Add the missing barrier at the head of those helpers, which closes the window at under 1% cost, and lift the now-unnecessary gfx1201 quarantine from the SmoothQuant tests.
Co-authored-by: Cursor <cursoragent@cursor.com>
Select one-row or multi-row forward kernels from runtime m_in, preserving the M > 8192 crossover without specializing the build cache on M. Co-authored-by: Cursor <cursoragent@cursor.com>
Reduce wide-row latency with streaming accesses, gamma reuse, fast BF16 packing, and runtime geometry dispatch while keeping M dynamic. Add a three-backend benchmark to reproduce FlyDSL, HIP, and Triton comparisons. Co-authored-by: Cursor <cursoragent@cursor.com>
|
performance comparison on MI308: |
Co-authored-by: Cursor <cursoragent@cursor.com>
jhinpan
left a comment
jhinpan
left a comment
There was a problem hiding this comment.
APPROVE. Both items from my previous CHANGES_REQUESTED review (at bbdde41a41ae) are closed at this head, and I found no new material defect.
Prior items
[P1] BF16 round-to-nearest-away regression in the gfx942 fast path — fixed.
Commit cae767d63cfd ("restore the RNA semantic to RNE") deletes to_elem_vec_rna from kernels/norm/rmsnorm_common.py, drops the USE_FAST_BF16_OUTPUT gate in build_rmsnorm_module, and makes the vector store call _to_elem_vec(...) unconditionally (kernels/norm/rmsnorm_kernel.py:331). How I checked at this commit:
grep -rn "to_elem_vec_rna\|USE_FAST_BF16_OUTPUT"over the tree: no hits, so no dangling reference or dead gate remains.grep -rn "0x8000"underkernels/norm/: no hits. The only rounding bias left is the RNEbias = lsb + 0x7FFFform inrmsnorm_common.py:107andlayernorm_kernel.py:89, which is the conversion my earlier tie probe showed matching both_to_elem_vecand PyTorch.kernels/norm/rmsnorm_common.pyat this head is byte-identical tomain(diffagainstcontents?ref=main), i.e. the helper was added and removed entirely within the PR and this file now carries no PR-local delta.
[P2] Reversed speedup column labels in tests/kernels/benchmark_rmsnorm_backends.py — outdated.
Commit b98ff768d95f removes the file (+0/-288); it no longer exists under tests/kernels/, and grep -rn "benchmark_rmsnorm_backends" over the tree returns nothing, so no script, doc, or workflow references it. The misstated committed artifact is gone.
What else I verified at this commit
- Scope of change since my last review. I diffed all four PR files at
7f2b9aa3against theirbbdde41aversions:layernorm_kernel.py,test_rmsnorm.py, andtest_layernorm.pyare byte-identical, andrmsnorm_kernel.pydiffers only by the three RNA hunks above. The merge frommain(7f2b9aa3) left the PR-owned files unchanged. This matters for the test evidence: on the gfx950 host I used last round,USE_GFX942_BF16_FAST_PATHis False, so the code that actually executed there was already the now-unconditional_to_elem_vecpath — my previous run (45 passed / 7 skipped on the non-large RMSNorm + LayerNorm suites, both large-shape tests, the Softmax test, and the N=7/8/2047/2048/2049/3073 and M=8192/8193 boundary probes) applies unchanged to this head. - Static checks that run without a build.
python3 -m py_compileon the four changed files plusrmsnorm_common.py: clean.python3 scripts/check_repo.py:check_docs_apiOK (17 files, 869 symbols, 16 skills); thetyped-arithmeticcheck errored only because this pinned checkout has noorigin/mainref for its merge-base, so I ran its rule directly over the PR's added lines (maximumf|minimumf|maxsi|maxui|minsi|minui|maxnumf|minnumf|ceildivsi|ceildivui) — no hits. No conflict markers underkernels/ortests/. GitHub now reports the branch asmergeable: true, resolving the DIRTY state I noted last time. - CI at this head.
Check Python Code Style,Check C++ Code Style,CodeQL, andDocumentationare green. I could not reproduce the style gate locally (noblack/ruffin this environment) and I could not re-run the device suite here (this review checkout has nobuild-fly/and notorch), so I am relying on CI plus the byte-identity argument above for that. Note that theFly DSL testworkflow is stillin_progressat7f2b9aa3; please confirm it lands green before merging.
Non-blocking note
The MI308 and MI350 tables in the PR conversation were produced by the now-deleted script and still carry the reversed headers — e.g. the MI350 1x512 row prints 1.05x under FlyDSL/HIP for 1.89 us FlyDSL vs 1.99 us HIP, which is HIP/FlyDSL. The committed code no longer misstates anything, so this is only about how a future reader should interpret those posted tables.
Thanks for the thorough review. |
Compile both cache policies into one gfx95x launcher and pick between them on the runtime row count, since streaming only helps once the working set outgrows the last-level cache. Widen the default block for wide gfx95x rows and mirror it in the autotune default config. Rename _build_rmsnorm_large_m_small_n_module to _build_rmsnorm_small_n_module. Co-authored-by: Cursor <cursoragent@cursor.com>
|
performance comparison on MI350: |
Use masked vector iterations for vec8-aligned row widths instead of requiring an exact tile multiple, including a partially active final iteration in the test matrix. Co-authored-by: Cursor <cursoragent@cursor.com>
The quantized RMSNorm and LayerNorm launchers specialize every non-quantized operand (Gamma, Beta, the fused-add residuals, and SmoothQuant XScale) on the activation dtype, but that ABI was implicit and unchecked: an operand passed with any other dtype was reinterpreted bit-for-bit and silently produced corrupted scales instead of failing, and callers written against the older fp32-only SmoothQuant XScale ABI broke the same silent way. Add a validate_norm_operand_dtypes helper, assert the ABI in all eight quantized launchers so a mismatch raises ValueError at compile time, and document the operand dtype requirement in the RMSNorm section of docs/prebuilt_kernels_guide.md. Co-authored-by: Cursor <cursoragent@cursor.com>
The scalar paths moved one element through a copy atom, a (1, 1) logical divide and a one-element register-memory tensor, which is boilerplate for what buffer-tensor indexing already expresses. Replace the load_scalar/store_scalar helpers in rmsnorm_common.py and the _load_scalar/_store_scalar/_store_yscale helpers in layernorm_kernel.py with direct t[index] reads and writes at every call site, and drop the scalar copy atoms, (1, 1) divides and element-width locals that only fed them; this also retires the elem_dtype/store_dtype split _store_scalar needed. rmsnorm_bwd_kernel.py is touched because it imported the deleted helpers. The generated ISA is byte-identical to before the change (502 instructions, 36 VGPRs, no spills on the N=3073 and N=3001 bf16 configurations), so this is a source-level cleanup with no codegen impact. Co-authored-by: Cursor <cursoragent@cursor.com>
jhinpan
left a comment
jhinpan
left a comment
There was a problem hiding this comment.
Requesting changes for one reproducible silent-correctness hole in the new quantized-norm ABI guard; details are inline. I found no other material defect in the gfx95x cache-policy dispatch, partial vec8 ownership, or reduction/barrier ordering.
Non-blocking: the guide still describes the old LayerNorm N == 8192 vector gate and a fixed 256-thread RMSNorm default, while the current gfx95x wide-row default is 512. Also, the direct-indexing cleanup is numerically clean in focused tests, but it is not codegen-neutral for every touched path (for example, LayerNorm BF16 N=3001 uses 32 rather than 30 VGPRs), so the no-codegen-impact claim should be kept narrow.
…_str validate_norm_operand_dtypes compared every operand against Input.dtype but never against dtype_str, so a self-consistent operand set that did not match the builder specialization passed the guard and was reinterpreted bit-for-bit by the kernel: building a bf16 module and passing fp16 tensors produced corrupted residuals and int8 output with no error raised. Take dtype_str in the validator and reject Input.dtype != dtype_to_elem_type(dtype_str) ahead of the existing operand comparison, then pass dtype_str at all eight quantized launchers in rmsnorm_kernel.py and layernorm_kernel.py. Add negative tests covering the four RMSNorm and the four LayerNorm quantized launchers; both fail with DID NOT RAISE ValueError when the validator change is reverted. Update the RMSNorm quantized-variant paragraph in docs/prebuilt_kernels_guide.md, which stated the same Input-relative contract. Co-authored-by: Cursor <cursoragent@cursor.com>
…kernels Several claims in the norm sections predate the current kernels. The LayerNorm vector path is gated on a 16-bit element type and `N % VEC_WIDTH == 0` rather than the old `N == BLOCK_THREADS * VEC_WIDTH * 4` fast path, and a row that is not divisible falls back to a fully scalar two-pass body, so both the guide bullets and the module docstring are rewritten around `USE_VEC_N`. `VEC_ALIGN` and `USE_NONTEMPORAL` were dropped from layernorm_kernel.py and no longer exist anywhere under kernels/, so their table rows go away, and `WARP_SIZE` now reads as `get_warp_size()` resolves it. RMSNorm's `BLOCK_THREADS` defaults to `None` and resolves through `default_block_threads` to 512 on gfx95x rows of `N >= 8192`, which the section now states in a table matching LayerNorm and Softmax. Its forward is two passes with the row held in registers, not three with an LDS row cache; LDS carries only the cross-wave reduction slots, sized by wave count rather than by N. Both "Kernel signature" blocks described a GPU_MODULE_NAME / self / __call__ class API that no longer exists and are replaced by the actual @flyc.kernel and @flyc.jit signatures. Co-authored-by: Cursor <cursoragent@cursor.com>
|
/rerun-ci |
1 similar comment
|
/rerun-ci |
#1159 (336d6ff) dropped source_ir from CompiledArtifact.__getstate__, so it reads back as None on any artifact restored from disk and the five norm assertions on pre-lowering IR dereference it unconditionally. Only test_rmsnorm_large_shape trips on CI's cold cache: build_rmsnorm_module is a factory, so each call gets a fresh JitFunction with an empty in-memory cache, and M is not a builder argument, so it reloads the (8192, bf16) entry test_rmsnorm wrote earlier in the session. The other four survive only because their keys are unique to one test; a second run against the same build fails all five. Guarding loses nothing: manager_key hashes the kernel source through the launcher's closure, so any edit that could drop the wide copy invalidates the entry and the assertion runs against a real compile. Co-authored-by: Cursor <cursoragent@cursor.com>
The latest
|
| site | checks |
|---|---|
test_rmsnorm.py:197 |
forward copy width, via _assert_rmsnorm_forward_copy_width |
test_rmsnorm.py:1629 |
FP32 weight copy width, small-N tail |
test_rmsnorm.py:1720 |
FP32 weight copy width, staged backward (from #884) |
test_rmsnorm_autotune.py:76 |
known_block_size specialization |
test_rmsnorm_autotune.py:81 |
FP32 weight copy width |
The commit message has the mechanism, including why only
test_rmsnorm_large_shape trips on a cold cache.
Fix
A source_ir is not None guard at each site. No production code touched.
Skipping costs nothing: manager_key hashes the kernel source through the
launcher's closure, so any edit that could drop the wide copy invalidates the
cache entry and the assertion runs against a real compile anyway.
Verified on gfx950:
| unfixed | fixed | |
|---|---|---|
cold cache, test_rmsnorm + test_rmsnorm_large_shape |
1 failed, 1 passed | 2 passed |
| the four inline assertions (5 test cases), second run against the same build | 5 failed | 5 passed |
| rmsnorm + rmsnorm_autotune + layernorm, cold and warm | — | 58 passed, 3 skipped (both) |
On gfx1250 the RMSNorm kernels from #848 lose time in four places: - BF16 outputs use integer rounding, although gfx1250 has v_cvt_pk_bf16_f32. The integer rounding also turns NaN payloads >= 0x7FFF8000 into +-0. - The compiler cannot move a load above a store that may alias it. The plain and small-N kernels load gamma in pass 2, after the first output store, and the fused kernel interleaves its x and residual loads with the residual_out stores. Each such load waits one full memory latency. - At N = 32768 the bounds checks of full tiles raise the plain kernel to 258 VGPRs, which needs s_set_vgpr_msb bank switches. - For 16-bit rows with N <= 2048, every M <= 8192 takes one workgroup per row. At N = 512 only 64 of its 256 threads have work. On gfx1250 only: - has_hw_cvt_pk_bf16_f32() selects the hardware conversion in every RMSNorm builder (forward, small N, fused add, fused-add quant, backward). - The plain and small-N kernels load gamma in pass 1. The fused kernel issues every x, residual and gamma load before the first residual_out store. preload_weight_vec() keeps 16-bit gamma unconverted in registers. - The plain and fused kernels skip the bounds checks of full tiles. - The small-N module takes the multi-row kernel at every M when each lane handles at most three vectors of a row (N <= 768), and for M > N above that. The LLVM IR for gfx942 and gfx950 does not change. test_rmsnorm_bf16_output_keeps_nan checks that NaN payloads in FP32 weights stay NaN in BF16 outputs on targets with the hardware conversion. Signed-off-by: Jin Pan <jin.pan@amd.com>
4 participants
Motivation
Optimize rmsnorm/layernorm
Technical Details
Test Plan
Test Result
Submission Checklist