Add AVX2 LayerNorm/RMSNorm kernel for x86-64 (MLAS)

[Justin Chu (justinchuby)]

Summary

Adds an AVX2/FMA MLAS kernel for LayerNorm and RMSNorm on x86. LayerNorm uses a centered two-pass algorithm with an fp64 mean accumulator, improving accuracy over the existing scalar fp32 Welford path for poorly conditioned inputs. Short rows retain the scalar fallback, and existing non-x86 kernels are unchanged.

Performance

On an AMD EPYC 9V74, the corrected microbenchmark (whose scalar baseline mirrors production's per-output division) measured approximately these median speedups:

NormSize	LayerNorm	RMSNorm
128	9x	3x
256	9x	5x
768–4096	10–14x	4.3–6.3x

These are per-row microbenchmarks; the committed benchmark verifies that the AVX2 kernel is dispatched.

Validation

• MLAS tests: 41 passed, 2 disabled; 43/43 with disabled tests enabled
• fp64 parity sweep: 180 cases, 0 failures; worst normalized max error 0.022318
• Adversarial case relative error: 0.032976 for AVX2 vs. 0.93573 for the scalar path (28.4x better)
• Windows builds compile the kernel translation unit with /arch:AVX2

./onnxruntime_mlas_test --gtest_filter="*Fp64ParitySweep*" --gtest_also_run_disabled_tests
./onnxruntime_mlas_test --gtest_filter="*Benchmark*" --gtest_also_run_disabled_tests

[Justin Chu (justinchuby)]

Correction to the numerics claim, and benchmark results

Correcting something in the commit message before it misleads a reviewer.

The numerics claim was wrong

The commit says the kernel "keeps the same two-pass mean/variance formulation as the scalar path". That is not accurate. onnxruntime/core/mlas/lib/layernorm.cpp is dispatch only (41 lines) — MLAS has no scalar LayerNorm kernel. On x86-64 today MlasLayerNormF32() returns false, and the work is done by ComputeJob in onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc, which uses Welford's online algorithm (explicitly chosen there as numerically stable).

So this PR does change the numerical formulation for full LayerNorm on x86: Welford's one-pass → two-pass mean/variance. I should have said so up front. RMSNorm is unaffected — both sides use sum-of-squares.

This is a real trade-off and reviewers should decide it, not me. Two-pass avoids Welford's per-element delta / (h+1) division and is much faster, but Welford's is the more stable choice for large N. If you would prefer the kernel match Welford's semantics, I am happy to do that — it will cost most of the full-LayerNorm speedup below while keeping the RMSNorm gain.

Benchmarks

Now measured against the true baseline (the ComputeJob path that actually runs on x86 today), same binary, same flags. The baseline is confirmed not auto-vectorized beyond SSE2 (compiled -O3 -fno-fast-math, no -mavx2; zero ymm instructions in the object file).

Host: AMD EPYC 9V74 (AVX2, FMA, F16C; no AVX-512). p50/p95 over 1000 iterations after warmup.

RMSNorm — pure SIMD effect, same algorithm both sides:

NormSize	AVX2 p50	Scalar p50	p50 speedup	p95 speedup
128	0.07µs	0.18µs	2.6×	2.2×
768	0.21µs	0.91µs	4.3×	4.1×
1024	0.29µs	1.20µs	4.1×	4.0×
2048	0.56µs	2.38µs	4.3×	4.0×
4096	1.27µs	4.90µs	3.9×	3.8×

Full LayerNorm — SIMD plus the algorithmic change above, so read with that caveat:

NormSize	AVX2 p50	Scalar p50	p50 speedup	p95 speedup
128	0.08µs	0.81µs	10.1×	10.0×
768	0.22µs	4.76µs	21.5×	20.7×
1024	0.30µs	6.33µs	21.1×	20.8×
4096	1.30µs	25.16µs	19.3×	19.0×

Much of that is removing Welford's per-element division, not vectorization. I would not want the 20× figure quoted without the caveat.

Known regression

For NormSize ≤ 15, RMSNorm is 0.7–0.8× — AVX2 setup overhead exceeds the gain on very short rows. Disclosing rather than hiding it; happy to add a size threshold that falls back to scalar below the vector width if you would prefer.

Scope of the measurement

This is a single-row kernel microbenchmark. End-to-end model impact is unmeasured. LayerNorm is typically a small fraction of total inference time, so please do not read these as model-level numbers.

[Justin Chu (justinchuby)]

Update: Welford preserved, tiny rows excluded — both driven by measurement

Following up on my own correction above. Two changes; the PR is materially different and, I think, materially better.

1. The two-pass variance was not just less precise — it was wrong

I said the Welford → two-pass change was a trade-off for reviewers to weigh. Adversarial testing settled it against us:

case	two-pass	Welford
base=1e6, spread=1e-3, N=256	NaN	finite
base=1e6, spread=1e-3, N=1024	NaN	finite
base=1e7, spread=1e-2, N=256	100% rel. error	finite
base=1e7, spread=1e-2, N=1024	100% rel. error	finite

With mean ≈1e6, both terms of E[x²] - mean² are ≈1e12 and the subtraction consumes every significant fp32 digit. Not a tolerance question — a correctness cliff. Exactly why ComputeJob uses Welford's.

The reduction now uses Welford's with 8 parallel AVX2 accumulators combined by the standard pairwise merge, so the formulation matches the scalar baseline.

Against an fp64 reference it is more accurate than scalar Welford at every size, since the parallel accumulators shorten each dependent chain:

scenario	N	scalar Welford fp32	AVX2 Welford SIMD
LLM activations	4096	2.30e-05	5.51e-07
large-N benign	65536	2.97e-05	3.23e-07
high dynamic range	4096	1.73e-06	3.62e-07

RMSNorm keeps sum-of-squares — no mean subtraction, so no cancellation to avoid.

Cost: Welford's per-element division runs ~2.5–3× slower than two-pass, so full-LayerNorm drops from ~20× to 5–7×. That is the honest number. The earlier 20× was mostly the removed division, not vectorization, which is why I did not want it quoted.

2. Tiny rows now fall back instead of regressing

The ≤15 regression I disclosed is gone by construction. Measured crossover on AMD EPYC 9V74 (AVX2/FMA, no AVX-512): 3–22% regression for NormSize 1–7, parity at 8 — where the first 256-bit iteration executes. MlasLayerNormF32 now declines below 8 and the caller keeps its existing path.

Tests assert the contract on both sides: ≥8 the kernel must run, <8 it must decline. Neither a silent fallback nor an accidental re-enable for tiny rows can pass unnoticed.

Status

[  PASSED  ] 40 tests.

Includes the catastrophic-cancellation cases, asserting finiteness and exact parity with scalar Welford. Formatting fixed (core/mlas/** is excluded from clang-format per .lintrunner.toml, so only the test file was reformatted).

Still unmeasured: end-to-end model impact. This is a single-row kernel microbenchmark and I am not going to imply otherwise.

Keeping it draft — happy to take feedback on the Welford SIMD merge or the threshold value.

[Justin Chu (justinchuby)]

Internal review pass — findings addressed

Ran an adversarial internal review before asking for your time. No blocking findings; two substantive items, both now fixed.

Wasted work in the RMSNorm path. In Simplified mode the running sum feeds only the optional Mean output — normalization never subtracts the mean — so when MeanOut is null the accumulation and its horizontal reduce were dead work. Now skipped via a per-row check outside the inner loop (a branch inside the loop could cost more than the single vaddps it saves).

Worth being precise: the reviewer estimated ~15% from reading the code. Measured, it is 5–9% for NormSize 8–64 and under 1% for NormSize ≥ 256 — so at LLM-typical hidden sizes it is in the noise. The change stands on making the dead-code intent explicit, not on a perf claim.

fp64 reference uses two-pass. Reasonable thing to flag, given this PR removed two-pass from the kernel for causing NaN. It is deliberate and now documented in-code: at these magnitudes the cancellation cannot bite in fp64, and keeping a different algorithm in the reference is what makes it an independent oracle — if reference and kernel both used Welford, a shared conceptual error could produce matching wrong answers.

Also removed a dead statement in the tests.

Independently verified during review

• Welford pairwise merge formula derived and checked against the implementation
• Dispatch installed only under the AVX2 + FMA3 CPUID check
• NormSize < 8 decline leaves Output/Mean/InvStdDev untouched, so the caller's fallback is correct
• All loads unaligned; no buffer overread, no UB in the scalar tail
• Uses full-precision 1/sqrtf, not an rsqrt approximation
• Precision claim genuine: 2–40× lower error than scalar Welford

[  PASSED  ] 40 tests.

Welford reduction, the NormSize < 8 contract, and the full LayerNorm path are untouched. Still no end-to-end model claim — this remains a kernel microbenchmark.

[Justin Chu (justinchuby)]

Correction: the AVX2 reduction was less accurate than scalar. Replaced.

An internal review found a blocker that invalidates an accuracy claim I made earlier in this PR. Correcting it directly.

B1 — the lane-parallel Welford was ~1000× worse, not better

I previously claimed the AVX2 Welford was more accurate than scalar Welford. That was wrong on large-base/small-spread inputs. Reproduced against an fp64 oracle:

base=1e5, spread=1e-2, N=4096
  scalar Welford   rel err 3.35e-05
  AVX2 Welford     rel err 2.71e-01     ~8000x worse

Each of the 8 lanes accumulates its own mean over ~N/8 elements in fp32, so the rounding is already baked in before the pairwise merge runs. Merging in double cannot recover it. My earlier measurements simply did not sample this region — which is a fair criticism of the tests as much as the kernel.

B2 — replaced with centered two-pass (double first-pass sum)

                                 worst rel err   speed vs scalar
scalar Welford                      5.03e-02          1.0x
AVX2 Welford (removed)              2.82e-01          8.1x
centered two-pass, fp32 sum         1.00e+00         ~15x
centered two-pass, double sum       5.95e-03         14.3x   <-- now used

Worth distinguishing from earlier in this PR's history: the formulation that produced NaN was the uncentered Var = E[x²] − mean². Subtracting the mean before squaring removes that cancellation entirely, so centered two-pass is both more accurate than scalar Welford here and faster, since it avoids Welford's per-element division. The fp32-sum variant is not viable — double accumulation on the first pass is required.

Cross-platform bugs I introduced, now fixed

• The NormSize < 8 gate was in shared dispatch, so it also disabled the pre-existing RISC-V RVV kernel for short rows. Now scoped to x86; RVV behaves exactly as before this PR. Thanks for catching that — it was a regression for a platform I wasn't touching.
• Tests asserted AVX2 dispatch unconditionally, which would have failed CI on every non-AVX2 platform. Now capability-gated with a skip, while retaining the reachability assertion where a kernel exists so a silent fallback still fails.
• Zero-variance assertions assumed Welford semantics and conflicted with RVV's E[x²] − mean². Now accept both while still checking finiteness.
• Added MSVC /arch:AVX2; the kernel source sits outside the globbed AVX2 list.

Tests strengthened

The old precision tests were too weak to catch a 1000× regression — that is how B1 reached review. Added an fp64 parity sweep over base 1e3–1e6, spread 1–1e-3, eps 1e-5/1e-6/1e-12, NormSize 9–4096, plus an explicit B1 regression guard calibrated so the removed Welford (2.49e-01) fails and the current kernel (3.30e-02) passes.

Also fixed the sweep's own metric: per-element relative error returns exactly 1.0 whenever a near-zero normalized output rounds to zero, which is routine for LayerNorm. It now uses vector-normalized max error.

[  PASSED  ] 41 tests.

Staying in draft pending another internal review pass. Apologies for the churn — I would rather correct my own numbers here than have you find them.

[Justin Chu (justinchuby)]

Note on the Build Linux arm64 Debug failure, and a history rewrite

The arm64 Debug job

This job has been the only red check here, and I have not been able to attribute it to this change. What the log actually shows:

[1452/1458] Linking CXX executable onnxruntime_mlas_test    07:18:45
Post job cleanup.                                           07:19:14

No FAILED:, no ninja: build stopped, no compiler or linker diagnostic, no non-zero exit code. The build got to 1452 of 1458 targets and the container stopped 29 seconds later, during the remaining test-executable links (onnxruntime_test_all, onnxruntime_provider_test, onnxruntime_autoep_test — the heavy ones). Runner was Standard_D8pds_v5, CCache missed, so every object was a fresh compile.

Against that, everything this PR adds is x86-only and should not be reachable on arm64 at all:

• the AVX2 entry points are behind MLAS_TARGET_AMD64 (mlasi.h, platform.cpp)
• the new sources are in x86_64-only sections of cmake/onnxruntime_mlas.cmake
• the NormSize < 8 dispatch gate in layer_norm_impl.cc is x86-only — deliberately, so it cannot suppress the existing RISC-V RVV kernel
• arm64 Release passes; only Debug fails

I checked three other open PRs (#31972, #31971, #31970) and arm64 Debug is green on all of them, so I am not claiming a broken pipeline either.

My best hypothesis is memory pressure during parallel Debug linking, but I want to be straight that I cannot prove it — I have no access to the runner's kernel logs, there is no exit 137 or Killed in the output, and gh run rerun will not re-run jobs on a fork PR. So this is unresolved, not explained away. The force-push below re-runs it; if it goes green with no code change, that is the answer.

If a maintainer can see the runner-side logs for that job, I would appreciate a pointer.

History rewrite

I force-pushed to remove an internal working note (.squad/…) that an over-broad git add had swept into an early commit. Deleting the file in a later commit was not enough — the content was still reachable in history. The branch is now rebuilt without it.

The resulting tree is byte-identical to what was reviewed (de78f4f5… before and after), so no code changed; only the commit graph did. 9 commits became 7.

[Justin Chu (justinchuby)]

The Windows GPU Kernel Documentation Validation failure is inherited from main, not from this PR

The failing diff is in docs/ContribOperators.md, in the MRotaryEmbedding description:

-  (or omitting it) reduces this op to standard RoPE.
+  reduces this op to standard RoPE.

That text arrived with #31728 (e415ef9afd, "Add fused MRotaryEmbedding contrib op for Qwen mRoPE variants"). The schema docstring and the checked-in ContribOperators.md disagree, so gen_contrib_doc.py --domains com.microsoft regenerates a different file and the validation step exits 1.

This branch does not touch ContribOperators.md — git diff $(git merge-base HEAD upstream/main)..HEAD -- docs/ContribOperators.md is empty. I have not modified it here, since regenerating an unrelated contrib doc does not belong in this PR. It should affect any PR that merges against current main until it is regenerated upstream.

This PR changes no kernel registrations and no operator schemas, so the documentation validation result here is unrelated to its contents.

Separately, the Build Linux arm64 Debug failure I described earlier did not reproduce after re-running — it is green on the current run. That supports the resource/flake reading rather than anything architecture-specific in this change, which is consistent with every symbol here being behind MLAS_TARGET_AMD64.

[Justin Chu (justinchuby)]

CI status: the remaining red checks are not from this PR

Marking ready for review.

Windows GPU Kernel Documentation Validation — inherited from main. The regenerated diff is in docs/ContribOperators.md (MRotaryEmbedding description), which came in with #31728 (e415ef9afd). This branch changes no schemas and no kernel registrations, and does not touch that file.

coreml (arm64, …) — Downloading … gradle-8.7-bin.zip failed: timeout (10000ms). A Gradle CDN timeout during the Java build; the C++ compiler never reached our code. The same job is green on #31969–#31972.

Build Linux arm64 Debug, which I flagged earlier, did not reproduce on re-run and is now green. That is consistent with the code: every symbol added here is behind MLAS_TARGET_AMD64, and the new sources sit in x86_64-only sections of cmake/onnxruntime_mlas.cmake.

Recap of what this PR is now

The AVX2 kernel uses a centered two-pass formulation — mean = sum/n, then sum((x-mean)^2) — with the first-pass sum accumulated in double.

That replaced a lane-parallel Welford version I had originally proposed. Review found it was not merely imperfect but worse than the scalar baseline for large-base/small-spread inputs (base 1e5, spread 1e-2, N=1024: scalar relative error 2.54e-4 against AVX2 at 0.249), because per-lane means round in fp32 before the merge, and merging in double cannot recover what the lanes already lost. The current form measured both more accurate and ~4.7x faster than the Welford version.

Worth stating plainly since it is easy to conflate: this is centered two-pass, not the uncentered E[x^2] - mean^2 identity, which cancels catastrophically at large base.

Also in response to review: the small-NormSize dispatch gate is x86-only, so it cannot suppress the existing RISC-V RVV kernel; tests use explicit capability checks with GTEST_SKIP rather than asserting AVX2 dispatch on every platform; /arch:AVX2 is applied on Windows; and there is an fp64 parity sweep across base 1e3-1e6, spread 1-1e-3 and eps 1e-5/1e-6/1e-12, plus a regression guard for the case above.

[Hariharan Seshadri (hariharans29)]

Review — PR #31973: Add AVX2 LayerNorm/RMSNorm kernel for x86-64 (MLAS)

Scope (7 files, 13 commits)

• New kernel layernorm_kernel_avx2.cpp — centered two-pass LayerNorm with double-precision first-pass sum, single-pass sum-of-squares RMSNorm. 8-wide AVX2+FMA3 inner loops + scalar tails.
• platform.cpp — installs LayerNormF32Kernel = &MlasLayerNormKernelAvx2 inside the existing Cpuid1[2] & 0x1000 (FMA3) + Cpuid7[1] & 0x20 (AVX2) branch. One-line addition next to other FMA3 kernels.
• layernorm.cpp — adds #if MLAS_TARGET_AMD64 || MLAS_TARGET_IX86 guard around a NormSize < 8 → return false gate. Explicitly x86-scoped so it can't suppress the pre-existing RVV kernel.
• mlasi.h — forward-declares MlasLayerNormKernelAvx2 under the x86 gate.
• mlas.h — documents the NormSize < 8 decline contract.
• cmake/onnxruntime_mlas.cmake — adds the kernel source to both the Windows and GCC/Clang x86_64 source lists, with explicit /arch:AVX2 for the Windows source-properties block (needed because the file lives at MLAS root, not under intrinsics/avx2/* where the glob picks it up).
• test_layernorm.cpp — grows from ~150 to ~1000 lines: fp64 oracle, dispatch contract, fp64 parity sweep, adversarial precision scenarios, edge cases (denormals, NaN/Inf, zero variance, large magnitudes), disabled microbenchmark.

The dispatch and kernel construction are correct

I traced the CPUID gate: the AVX2 branch in platform.cpp line 477 checks (Cpuid1[2] & 0x1000) != 0 (FMA3, ECX.12 of function 1) and (Cpuid7[1] & 0x20) != 0 (AVX2, EBX.5 of function 7). The kernel uses _mm256_fmadd_ps (FMA3), _mm256_add_pd, _mm256_cvtps_pd, and 256-bit AVX2 loads/stores — every intrinsic is covered by the gate. Kernel install is at the same nesting level as MlasGemmFloatKernelFma3 and friends, so the dispatch semantics are identical to every other FMA3-required kernel in MLAS. ✓

The dispatch site in layernorm.cpp is 6 lines. After this PR it consults GetMlasPlatform().LayerNormF32Kernel first, then applies the new NormSize < 8 fallback on x86 only. Non-x86 platforms are untouched — RVV still dispatches for all N ≥ 1, matching its variable-length vector design. The x86-only guard is spelled with #if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_IX86), which matches how MlasLayerNormKernelAvx2 is declared in mlasi.h line 1416. ✓

The caller contract still holds

layer_norm_impl.cc line 48 already handles a false return from MlasLayerNormF32 by falling through to a scalar Welford path in the same function. That code shape is unchanged, so a NormSize < 8 call on x86 now returns false and falls into the existing scalar loop — no observable behavior change on the small-N path. Verified locally against the source at line 48–75. ✓

The Windows CMake edit is not redundant

I initially thought /arch:AVX2 was redundant because there's a set_source_files_properties(${mlas_platform_srcs_avx2} ...) at cmake/onnxruntime_mlas.cmake line 226. It isn't — mlas_platform_srcs_avx2 on Windows is a file(GLOB_RECURSE ... intrinsics/avx2/*.cpp), and layernorm_kernel_avx2.cpp lives at MLAS root, not under intrinsics/avx2/. Without the new per-source property, MSVC would fail to compile _mm256_fmadd_ps (or at best emit worse code). The GCC/Clang side folds the file into mlas_platform_srcs_avx2 (which on that side is an explicit list, not a glob), so it inherits -mavx2 -mfma -mf16c at line 862. Both branches are correct. ✓

Algorithmic change vs. scalar Welford — worth calling out explicitly

This isn't a bit-identical drop-in replacement. The scalar path in layer_norm_impl.cc uses Welford's online algorithm; this kernel uses centered two-pass with a double-precision first-pass sum. The author has walked through this in-thread and settled on the current form after two rejected variants (uncentered E[x²] - mean², then lane-parallel Welford — both worse than scalar). Concretely:

• Uncentered was wrong — catastrophic cancellation at base ≈ 1e6, producing NaN. Discarded.
• Lane-parallel Welford was ~1000× worse than scalar at (base=1e5, spread=1e-2, N=1024) because per-lane means round to fp32 before the pairwise merge, so the double-precision merge can't recover what the lanes already lost. Discarded.
• Centered two-pass with double first-pass sum is the current form. Measured on the author's box at that same (1e5, 1e-2, 1024) case: fp64 Welford ref → scalar Welford fp32 has ~9.4e-01 relative error (yes, ~1.0 — normalized max), this kernel has ~3.3e-02. About 28× more accurate than the scalar path it replaces, at that condition number.

The fp64 parity sweep in the test file exercises 180 cases from (base, spread) ∈ {1e3..1e6} × {1, 1e-1, 1e-2, 1e-3} filtered by condition < 1e6, three epsilons, ten NormSizes — worst 2.23e-02, threshold 3e-2. That's a real regression net, not a smoke test. The B1 regression guard at 5e-2 is calibrated so the removed lane-parallel Welford (which measured 0.249) would fail it and the current kernel (0.033) passes. Good check.

Users depending on bitwise LayerNorm equivalence between MlasLayerNormF32 → false and → true will see differences after this PR. In practice that would only matter for exact-repro test suites; the accuracy claim is well-defended and the ONNX Runtime CPU EP has never contracted bit-equivalence between the fast path and the fallback. Non-blocking, but worth surfacing in the merge commit body.

Kernel body — line-by-line spot checks

I read the 260-line kernel and checked the following:

• Simplified (RMSNorm) with MeanOut != nullptr: parallel accumulation of dsum (fp64, 4 lanes) and vsumsq (fp32, 8 lanes) in the same 8-wide loop. Uses _mm256_castps256_ps128 + _mm256_extractf128_ps to split the 8-wide float for _mm256_cvtps_pd. Horizontal reduces are the standard permute/add pattern. Scalar tail correctly continues both accumulators. ✓
• Simplified with MeanOut == nullptr: skips the double-precision sum accumulation entirely — no vsumd, no horizontal reduce for it, no scalar tail. This is the ~5–9% short-N win from the author's second internal review pass. ✓
• Full LayerNorm Pass 1: 4-wide _mm256_cvtps_pd loop for the double-precision sum. Loop invariant i + 4 <= n correctly leaves 0–3 elements for the scalar tail. ✓
• Full LayerNorm Pass 2: 8-wide vd = vx - vmean; vvar = fmadd(vd, vd, vvar). Correct centered form. Scalar tail matches. ✓
• Normalization pass: three variants (RMSNorm, LayerNorm no-bias, LayerNorm with-bias), 8-wide + scalar tail each. In the with-bias case, _mm256_fmadd_ps(vy, vs, vb) computes y * scale + bias in one FMA — matches the scalar (x - mean) * inv_denom * scale + bias. ✓
• All loads are _mm256_loadu_ps / _mm_loadu_ps — unaligned. Alignment is not a caller contract on MLAS f32 kernels. ✓
• inv_denom = 1.0f / sqrtf(var + eps) — full-precision reciprocal + sqrt, not _mm256_rsqrt_ps approximation. Matches the numerical claim in the docstring. ✓
• No use of _mm_free/_mm_malloc — kernel is stack-only. No allocation. Not swallowing exceptions. ✓

I did NOT find any of the classic AVX intrinsic pitfalls: no _mm256_set_m128 type-punning, no ambiguous _mm256_zeroupper placement, no scalar tail reading past Input + n.

Comments / follow-ups

1. The kKernelDispatchThreshold constant is duplicated. layernorm.cpp hard-codes 8 under the x86 gate; test_layernorm.cpp redefines the same constant under the same #if. Comment says "Keep in sync" — that's a human check that will fail eventually. Exporting from mlasi.h as constexpr size_t kMlasLayerNormF32MinNormSize = ...; under the x86 gate would make it a single source of truth. Non-blocking; a follow-up file the "TODO" comment would justify.

2. Windows-only trap — _mm256_fmadd_ps availability on MSVC without /arch:AVX2 is a source of build breaks. The added set_source_files_properties(... /arch:AVX2) handles this correctly for the new file. But: rotary_embedding_kernel_avx2.cpp and qkv_quant_kernel_avx2.cpp sit in the same root and I couldn't find /arch:AVX2 for them in the Windows block either. Either they don't use FMA3 intrinsics, or there's a pre-existing latent build problem. Not this PR's issue; worth a follow-up scan.

3. The fp64 reference is intentionally different from the kernel algorithm. The comment in the test file explains this clearly: an independent oracle rather than a mirror. Good practice, and the comment is the right length. Consider extracting ReferenceLayerNorm into a small internal test header if a future softmax/RMSNorm review wants to reuse the fp64 pattern — same shape appears in the softmax kernel tests.

4. The dispatch contract in mlas.h is now the authoritative documentation. Consider stating the single in-tree caller (layer_norm_impl.cc) and requiring new callers to test both branches of the bool return. Nit.

5. The 660-line test file is a lot. Given the author's development history (three algorithm iterations, each surfaced by adversarial testing that the earlier suite would have accepted), the tests are load-bearing. The DISABLED_AdversarialPrecisionReport and DISABLED_Benchmark are correctly gated behind --gtest_also_run_disabled_tests so they don't clog CI. Keep.

6. RMSNorm at NormSize ∈ [8, 15] — the author's benchmark shows parity or small win at N=8 and regression fixed at N<8 by the gate. If real workloads use RMSNorm with NormSize just above 8, monitor. Could revisit the threshold as a follow-up if micro-batched inference surfaces regressions.

7. Nit: assert(!Simplified || Bias == nullptr); inside the kernel — this is a caller-contract assertion (RMSNorm can't have a bias). It's an internal MLAS invariant. assert compiles out in Release, so it's a debug-only check, which is fine here. The dispatcher's layernorm.cpp doesn't re-check this before calling the kernel — but the caller in layer_norm_impl.cc explicitly passes (simplified || !bias_data) ? nullptr : bias_data + i, so the invariant holds. ✓

Non-issues I verified

• Force-push / history rewrite — removed an internal .squad/… note. Tree unchanged (de78f4f5… before and after). Not a code concern. ✓
• Arm64 Debug CI flake — author's analysis is right; every symbol here is behind MLAS_TARGET_AMD64 or MLAS_TARGET_IX86. Flake, not a real signal. ✓
• The unrelated docs/ContribOperators.md doc-validation failure is inherited from Add fused MRotaryEmbedding contrib op for Qwen mRoPE variants #31728, not this PR. Nothing to do here. ✓
• MLAS_LAYERNORM_F32_KERNEL signature at mlasi.h:741 matches the kernel definition exactly (input, scale, bias, output, meanout, invstdout, normsize, epsilon, simplified). ✓

Recommendation

Approve. The kernel is correct, correctly scoped to x86, correctly gated on CPUID (AVX2+FMA3), doesn't touch non-x86 dispatch, and comes with the strongest precision-testing suite I've seen on a LayerNorm PR to ORT. The centered two-pass algorithm is a genuine numerical improvement over the scalar Welford it replaces — up to ~30× more accurate at the large-mean/small-spread condition — and the author's development history (three algorithm iterations, each caught by adversarial testing) is exactly the sequence a reviewer wants to see: the final form is not the first thing that came to mind.

Merge commit message should note that this changes the LayerNorm numerical formulation on x86 from Welford to centered two-pass, so anyone reproducing outputs bit-for-bit against a pre-PR build will see small differences. Not a compat break, but worth surfacing.

Two follow-up items worth filing separately: (1) consolidate kKernelDispatchThreshold into a shared constexpr in mlasi.h; (2) audit whether other root-level *_kernel_avx2.cpp files (rotary_embedding, qkv_quant) need explicit /arch:AVX2 on Windows.

[Copilot]

Pull request overview

Adds AVX2/FMA LayerNorm and RMSNorm acceleration to MLAS on x86.

Changes:

• Implements centered two-pass LayerNorm and vectorized RMSNorm kernels.
• Adds runtime feature dispatch and short-row fallback.
• Expands correctness, precision, edge-case, and benchmark coverage.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
onnxruntime/test/mlas/unittest/test_layernorm.cpp	Adds extensive kernel tests and benchmarks.
onnxruntime/core/mlas/lib/platform.cpp	Registers kernels on AVX2/FMA-capable x86 CPUs.
onnxruntime/core/mlas/lib/mlasi.h	Declares the AVX2 kernel.
onnxruntime/core/mlas/lib/layernorm.cpp	Adds the x86 short-row dispatch gate.
onnxruntime/core/mlas/lib/layernorm_kernel_avx2.cpp	Implements AVX2 LayerNorm and RMSNorm.
onnxruntime/core/mlas/inc/mlas.h	Documents the dispatch contract.
cmake/onnxruntime_mlas.cmake	Compiles the kernel with AVX2/FMA flags.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Tianlei Wu (tianleiwu)]

Reviewed the current head. The AVX2/FMA3 kernel, feature dispatch, x86-specific short-row fallback, and cross-platform build integration look sound, and the prior review threads are resolved. I found three non-blocking test/comment maintenance items, included inline.