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

[Justin Chu (justinchuby)]

Draft — opened early for feedback on approach. Performance numbers are still being measured against the correct baseline; see Benchmarks below.

Motivation

MlasLayerNormF32 dispatches LayerNormF32Kernel to a RISC-V RVV kernel where available and otherwise uses the scalar implementation in layernorm.cpp. There is no x86-64 kernel, so LayerNormalization and SimplifiedLayerNormalization run scalar on AVX2 hardware.

Verified against main (16b486a2): onnxruntime/core/mlas/lib/ contains only layernorm.cpp and riscv64/layernorm_kernel_rvv.cpp, and platform.cpp sets LayerNormF32Kernel only for RVV.

Change

Adds an 8-wide AVX2 + FMA3 two-pass kernel and registers it inside the existing AVX2 CPUID dispatch block, alongside the other AVX2 kernels selected there.

File
onnxruntime/core/mlas/lib/layernorm_kernel_avx2.cpp	new kernel
onnxruntime/core/mlas/lib/mlasi.h	declaration
onnxruntime/core/mlas/lib/platform.cpp	dispatch (AVX2 block)
cmake/onnxruntime_mlas.cmake	build wiring
onnxruntime/test/mlas/unittest/test_layernorm.cpp	tests

Numerics

The kernel keeps the same two-pass mean/variance formulation as the scalar path rather than switching to a one-pass sum-of-squares form, so accumulation behaviour matches the existing reference. Tail elements beyond the vector width use the scalar path, and Simplified (RMSNorm) mode skips mean subtraction exactly as the scalar kernel does.

Dispatch is fail-closed: the kernel is installed only inside the existing AVX2 feature check, so hardware without AVX2/FMA3 keeps using the scalar implementation.

Tests

[==========] 36 tests from 2 test suites ran.
[  PASSED  ] 36 tests.

• 32 parity tests: NormSize ∈ {1, 7, 8, 15, 16, 127, 128, 1024} × {simplified, full} × {bias, no-bias}, deliberately spanning non-multiples of the 8-wide vector to exercise the scalar tail.
• 4 edge cases: zero variance (all-equal input — the division risk in 1/sqrt(var+eps)), denormals, large magnitudes, NaN/Inf.
• Parity against an fp64-accumulated scalar reference. Tolerance 0.5% relative with a 1e-4 absolute floor, matching CloseEnough in test_util.h; zero-variance uses 2e-4. Worst observed divergence 0.02% relative at NormSize=1, from FMA contraction.
• Reachability: the tests assert a kernel was actually installed, so a dispatch regression fails the suite instead of silently exercising the scalar fallback. On AVX2 hardware there is no skip path.

Run on AMD EPYC 9V74 (AVX2, F16C; no AVX-512).

Benchmarks

Deliberately omitted for now. An initial measurement compared the kernel against an fp64 reference rather than the fp32 scalar path this actually replaces, which inflates the result, so those numbers are being redone against the correct same-binary baseline. I would rather post nothing than post a number I cannot defend — I will update this section with p50/p95 and variance once measured properly.

Happy to hear early feedback on the approach, dispatch placement, or whether a one-pass formulation would be preferred despite the numerics change.

[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.