Register BFloat16 LayerNorm/RMSNorm kernels on the CPU EP

[Justin Chu (justinchuby)]

What this adds

BFloat16 registration on the CPU EP for five LayerNorm/RMSNorm op families already permitted by their schemas:

• LayerNormalization (core, opset 17)
• LayerNormalization (contrib, opsets 1–16)
• SimplifiedLayerNormalization
• SkipLayerNormalization
• SkipSimplifiedLayerNormalization

The CPU implementations widen narrow-float inputs to float for arithmetic and narrow only output tensors. No native BF16 instruction or new ISA-specific path is added.

Statistics and type contracts

The contrib LayerNormalization and SimplifiedLayerNormalization CPU registrations now consistently declare U = tensor(float) for every T, including double; runtime dispatch uniformly uses ComputeImpl<T, float>. This corrects the generated operator documentation and matches the kernel registrations.

This is a compatibility change for contrib models that set stash_type = DOUBLE: shape inference may produce double statistic outputs, but no CPU kernel now matches that signature. Such models fail kernel matching during session initialization (for example, “Could not find an implementation for LayerNormalization(1)”), rather than producing different statistic values at runtime.

SkipLayerNormalization now materializes its optional CPU mean and inv_std_var outputs for all registered input types. It uses Welford accumulation for stable variance. SkipSimplifiedLayerNormalization materializes a zero centering mean, the inverse RMS used by its simplified normalization, and the optional residual output. This follows the existing CUDA forward/gradient contract: simplified gradients use inverse RMS and have no centering term.

CUDA and WebGPU SkipLayerNorm kernels are unchanged by this PR and do not currently materialize optional statistic outputs 1/2. This PR scopes the correction to the CPU kernels and covers the CPU behavior directly.

Implementation cleanup

• Consolidated the MLFloat16 and BFloat16 narrow-float LayerNorm fast paths behind one helper with thin forwarders.
• Centralized narrow-float tensor prepack conversion in core/util/narrow_float_utils.h, including identical empty-tensor behavior.
• Avoided allocating the full-size FP32 SkipLayerNorm residual scratch unless optional output 3 is requested.
• Regenerated and verified the affected CPU OperatorKernels.md rows.

Validation

• Windows x64 Debug onnxruntime_provider_test build.
• 43 focused LayerNormBFloat16CpuTest.* and SkipLayerNormTest.* tests passed, including float large-mean/small-variance Welford coverage, BF16 statistics, all optional Skip outputs, prepack paths, and generic broadcasting.
• Windows x86 Debug onnxruntime_providers build.
• CPU operator documentation generated from the Python bindings and the affected rows matched the checked-in document.
• clang-format, lintrunner, and git diff --check.

The Python-enabled documentation build initially hit vendored Abseil C4702 warnings promoted to errors; rerunning it with the repository-supported --compile_no_warning_as_error option succeeded for documentation generation. No source warning from this PR was involved.

[Justin Chu (justinchuby)]

Internal review pass — findings addressed

Ran an adversarial internal review before asking for your time. No blocking findings; one substantive item, now fixed.

U type constraint

The contrib REGISTER_CONTRIB_KERNELS macro registered U=T for every type, but the contrib schema constrains U — the optional Mean/InvStdDev outputs — to tensor(float). So my new BFloat16 registration declared a U the schema does not permit. The same was already true of the existing MLFloat16 registration.

No runtime correctness problem: the contrib LayerNorm constructor does not set contrib_op=true, so SrcDispatcher always dispatches to ComputeImpl<T, float> and those outputs are emitted as float regardless of the declaration. It was declaration-only.

The macro now takes (T, U) and registers narrow float types with U=float. That also corrects the pre-existing MLFloat16 declaration rather than adding a second, differently-wrong registration next to it — and it matches the CUDA contrib kernels, which already register U=float for narrow types.

I widened scope by one line to fix the pre-existing case; leaving two adjacent registrations inconsistent seemed worse than the small extra diff. Happy to revert to bf16-only if you would rather keep this strictly additive.

Deliberately not done

NarrowToFloat/FloatToNarrow are duplicated between layer_norm_impl.cc and skip_layer_norm.cc. Deduplicating needs a shared header, which is scope creep for a registration change, and the duplicated code is short and obviously correct. Flagging as a follow-up rather than smuggling it in here.

Independently verified during review

• All four ops' schemas genuinely permit T=BFloat16 — checked per op and per opset range, not just one
• No fp16-specific clamping or saturation leaks into the bf16 path (bf16 has f32's exponent range, so an fp16-motivated clamp would have been a real bug)
• BFloat16(float) rounds to nearest-even
• Prepacked gamma/beta widen bf16→f32 correctly
• Anti-fallback construction holds: single EP with no graph optimizations yields "no kernel found" if a registration is missing

[  PASSED  ] 10 tests.

Still no performance claim — this is a reachability/functionality change.

[Justin Chu (justinchuby)]

MLFloat16 regression check (follow-up to the U=float macro change)

The macro change touches the pre-existing MLFloat16 registration, and my bf16 tests alone would not have caught a regression there. So I checked the existing fp16 coverage directly rather than assuming:

onnxruntime_provider_test --gtest_filter=*LayerNorm*
[  PASSED  ] 89 tests.

23 of those are fp16-named LayerNorm tests (LayerNorm_Scale_Float16Input, LayerNorm_Scale_Float16ScaleOutput, EmbedLayerNormBatch1_Float16, …), so the changed MLFloat16 path is genuinely exercised and unbroken.

Tracing the dispatch confirms why: the contrib LayerNorm constructor leaves contrib_op=false, so SrcDispatcher always calls ComputeImpl<T, float> — Mean/InvStdDev were already float regardless of the declared U. The change is declaration-only.

On kernel matching, the new declaration is strictly more correct: U=MLFloat16 was schema-violating and would mismatch a schema-compliant model whose Mean/InvStdDev edges are float. U=float matches those correctly, so no valid model regresses.

Still happy to split the MLFloat16 line into its own PR if you would prefer this one stay strictly additive.

[Justin Chu (justinchuby)]

Retracting my test claim, and fixing a real numerical bug

An internal review rejected this PR. Two findings are mine to own publicly.

I claimed "45 MLAS kernel tests". That was wrong.

onnxruntime/test/mlas/unittest/test_layernorm_bf16.cpp (1037 lines) called no MLAS function at all - every Mlas token in it was a harness base class (MlasTestBase, MlasTestFixture). Its own header referred to a "kernel hook (MlasLayerNormBF16)" that does not exist and is not part of this PR. It tested standalone BFloat16 rounding arithmetic, while this change is CPU EP registration, with compute going through layer_norm_impl.cc.

They were real tests, but they exercised nothing this PR touches. Deleted.

The honest number is 17 CPU EP operator tests, covering all five registered families: core LayerNormalization opset 17, contrib LayerNormalization opset 1-16, SimplifiedLayerNormalization, SkipLayerNormalization, SkipSimplifiedLayerNormalization.

A real bug: stats were degraded to bf16 despite U being float

Mean/InvStdDev are declared float, but both narrow-float ComputeJob overloads round-tripped them through the narrow type before writing - losing ~0.4% for bf16 and diverging from what the generic path produces for identical input.

That is a silent cross-EP numerical inconsistency, and it also contradicts this PR's own argument: the point of widen -> f32-accumulate -> narrow is to preserve precision, so quantising the statistics back down undoes it.

Both overloads now call WriteStat<U> directly. The dead MLFloat16/BFloat16 WriteStat branches are removed, and SrcDispatcher uses if constexpr so ComputeImpl<NarrowType, NarrowType> is never instantiated.

Tolerances now split by output type

output	tolerance	rationale
BFloat16 Y	2 bf16 ULP (0.016)	8-bit mantissa
float Mean / InvStdDev	1e-5	U is float - must hold to f32 grade

That split is what makes the stat tests meaningful: the old behaviour was ~780x the tolerance, so these tests fail against the pre-fix code.

Also addressed

• ~31 CI failures all traced to one dead static float BF16Ulp(float) under -Werror. Removed. My local builds had used --compile_no_warning_as_error, which is exactly the flag that hides it - verification since done with warnings-as-errors.
• docs/OperatorKernels.md updated for the new registrations. Regenerating it needs built Python bindings, which was not practical here, so the five affected rows were hand-edited to match the existing format - flagging that rather than passing it off as generated.
• Internal working notes and reviewer names that had been committed by mistake are removed.

Scope decision on the MLFloat16 U change

Keeping it. The contrib schema constrains U to float, the CUDA contrib kernels already register U=float for narrow types, and it is declaration-only at runtime (SrcDispatcher always calls ComputeImpl<T, float>). Leaving two adjacent registrations inconsistent seemed worse. Happy to split it out if you would rather this stay strictly additive.

Build clean with warnings-as-errors; 17 BFloat16 tests and 96 across the LayerNorm suite pass. Staying draft pending another internal review.

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

Worth noting from the same log: gen_opkernel_doc.py ran and reported no diff for docs/OperatorKernels.md, so the kernel-doc rows in this PR match what the generator produces.

[Justin Chu (justinchuby)]

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

Marking ready for review. Three checks are red; I looked into each rather than assuming.

Windows GPU Kernel Documentation Validation — inherited from main. The regenerated diff is in docs/ContribOperators.md, in the MRotaryEmbedding description, which arrived with #31728 (e415ef9afd). The schema docstring and the checked-in file disagree, so the generator produces a different result and the step exits 1. This branch does not touch that file — git diff $(git merge-base HEAD upstream/main)..HEAD -- docs/ContribOperators.md is empty. It should affect any PR merging against current main until it is regenerated. I have not "fixed" it here, since an unrelated contrib doc does not belong in this PR.

macOS jobs — network failures, not compile failures:

job	failure
coreml (arm64, arm64, Debug)	Downloading … gradle-8.7-bin.zip failed: timeout (10000ms)
webgpu (arm64, arm64, Debug)	pytorch_cpuinfo-populate-download Error 1 during cmake configure
coreml (arm64, arm64, Release)	urllib.error.URLError: <urlopen error [Errno 8] nodename nor servname provided…> in test_dynamic_quantization_subgraph, which fetches a model over HTTP

In the last of those the build itself succeeded and 340 of 341 tests passed — the single failure is DNS resolution. The same jobs are green on #31969–#31972.

Happy to rebase or re-run if a maintainer would rather see a clean board first.

[Justin Chu (justinchuby)]

Note on the two red checks

Both fail before any of this PR's code is compiled.

Build Linux arm64 Debug — the build reaches the final link stage and the container stops with no compiler or linker diagnostic, no FAILED:, and no non-zero exit code in the log. The remaining targets are the heavy test executables. This job has intermittently failed and then passed on re-run across several pushes here with no code change in between.

Objective-C-StaticAnalysis — fails during dependency fetch, before compilation:

error: downloading 'https://github.com/Maratyszcza/FXdiv/archive/....zip' failed
      status_code: 60
      status_string: "SSL peer certificate or SSH remote key was not OK."

This PR touches no Objective-C and no build configuration. Related Apple/arm64 lanes on sibling PRs have failed on the same class of dependency-download error (cpuinfo, XNNPACK, eigen3, and a vcpkg git fetch), all before compilation.

For contrast, a one-line documentation PR I opened from the same fork at the same time reached 86/86 green, so this is not a systematic problem with the fork or the branch.

Everything that actually exercises this change is green: the CPU EP tests were built and run locally with warnings-as-errors from a clean rebuild — 20 BFloat16 tests and 106 across the LayerNorm suite, plus upstream #31676's prepacked-length validation at 7/7.

Happy to rebase or re-push to retrigger if a maintainer would rather see a fully green board first — gh run rerun will not re-run jobs on a fork PR, so a push is the only lever available to me.

[Justin Chu (justinchuby)]

Corrections and a compatibility note

Three things updated in the description, plus one regression caught and fixed while doing so.

A regression I introduced and then caught. While adding the PrePack counter assertions, a change flipped is_packed's default to true in LayerNormImpl::PrePack. ConvertMLFloat16ToFloatIfNeeded only sets that flag inside its narrow-float branches — for float inputs it is a no-op — so the float path wrongly believed Scale/Bias were prepacked, skipped reading the inputs, and failed with Missing Input: Scale. Nine LayerNormTest cases broke.

It is fixed (one line, restoring the false default). Worth noting how it hid: the BF16-filtered run stayed green at 21/21 throughout, and only the wider suite showed it. Current state, rebuilt from clean:

• *LayerNormBFloat16* — 21 passed, 0 failed
• *LayerNorm* — 107 passed, 0 failed
• *SkipLayerNorm* — 26 passed, 0 failed

The stat round-trip figure was imprecise. I quoted ~0.4% generally; that is the bfloat16 number (7-bit stored mantissa). float16 is about 0.049% (10-bit stored mantissa). Both are now stated separately.

The tolerance margin was overstated. I wrote "roughly 780x the stat tolerance", which compared against the bare 1e-5 absolute and ignored the checker's relative term. The checker applies absolute + relative * |expected|, so at unit scale the effective stat tolerance is about 1.1e-4. The honest margins are about 36x for bf16 and 4.4x for fp16 — still enough that the tests fail against the pre-fix code in both cases, which is the property that matters, but an order of magnitude less than I claimed.

Compatibility disclosure. A model that declared Mean or InvStdDev as float16 or bfloat16 — which violates the schema, where those outputs are float — was previously accepted and will now produce different values, because the stats are written at float precision rather than round-tripped through the narrow type. That is the intended fix, but it is an observable change for such a model.

Also added: the A/B PrePack tests now assert number_of_pre_packed_weights_counter is 0 on the non-initializer leg and 2 on the initializer leg, so they prove the prepacked path actually ran rather than only that the outputs agree. I verified that assertion is non-vacuous by forcing is_packed = false and confirming the test fails. And there is now an MLFloat16 opset-17 stats test, so the fp16 side of this behaviour change is covered rather than inferred.

[Copilot]

Pull request overview

Adds CPU EP BFloat16 support for LayerNorm/RMSNorm operator families.

Changes:

• Registers BFloat16 kernels for five normalization operators.
• Adds shared narrow-float conversion utilities and float-precision statistics.
• Adds CPU tests and updates kernel documentation.

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/test/contrib_ops/layer_norm_bf16_cpu_test.cc	Adds BFloat16 CPU tests.
onnxruntime/core/util/narrow_float_utils.h	Adds narrow-float conversion helpers.
onnxruntime/core/providers/cpu/nn/layer_norm.cc	Registers core BFloat16 LayerNorm.
onnxruntime/core/providers/cpu/nn/layer_norm_impl.h	Updates narrow-type dispatch.
onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc	Implements BFloat16 normalization and precise statistics.
onnxruntime/core/providers/cpu/cpu_execution_provider.cc	Exposes the core kernel registration.
onnxruntime/contrib_ops/cpu/skip_layer_norm.cc	Supports BFloat16 skip normalization.
onnxruntime/contrib_ops/cpu/layer_norm.cc	Registers contrib narrow-float kernels with float statistics.
onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc	Exposes contrib BFloat16 registrations.
docs/OperatorKernels.md	Documents expanded kernel support.

Suppressed comments (1)

onnxruntime/contrib_ops/cpu/layer_norm.cc:30

• This changes the existing contrib MLFloat16 kernel from U=MLFloat16 to U=float, but the new tests exercise MLFloat16 statistics only for the core opset-17 operator. Add a contrib-op MLFloat16 case with float Mean/InvStdDev outputs so a registration/dispatch regression here cannot pass unnoticed.

REGISTER_CONTRIB_KERNELS(MLFloat16, float)

---

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

[Tianlei Wu (tianleiwu)]

Thanks for the careful BF16 widening, prepack coverage, and generic-broadcast tests. I found two output-contract issues that need correction before the registrations are complete: the contrib double-stat dispatch remains unreachable, and the BF16 skip kernels do not materialize the optional float statistics they advertise. I also left one hot-path allocation suggestion and a documentation nit inline.

[Tianlei Wu (tianleiwu)]

Verdict: Request changes

The core BF16 registration work is sound — the schema already permits bfloat16, arithmetic widens to f32/f64 with no native bf16 math, both the fast path and the generic broadcast path have coverage, and the PrePack A/B tests asserting pre_packed_counter == 2 vs 0 against a shared reference are a strong property test. Thanks also for fixing all four items from the previous round.

Three things to resolve before merge.

1. docs/OperatorKernels.md will fail the Doc Gen check. gen_opkernel_doc.py builds the type map from KernelDef::type_constraints() (tools/python/gen_opkernel_doc.py:183), not from the schema. After REGISTER_CONTRIB_KERNELS(double, float), all four kOnnxDomain 1-16 registrations declare U = tensor(float). The hand-edit removed tensor(float16) from the two U lists but left tensor(double), so both rows will disagree with the generated output. Details inline.

2. The contrib U narrowing is broader and harder than the description says. REGISTER_CONTRIB_KERNELS(double) -> (double, float) is not mentioned in the PR body at all (only MLFloat16 is). And the contrib schemas expose a stash_type attribute with U in {tensor(float), tensor(double)} — after this PR no CPU registration supports stash_type = DOUBLE, so such a model fails at session initialization with "Could not find an implementation for LayerNormalization(1)". That is a kernel-match failure, not "will now see different values". Both changes are arguably improvements over the previous latent MutableData<U>() throw, but please state them accurately.

3. Materializing SkipLayerNorm mean / inv_std_var is undisclosed scope. This affects float, double, MLFloat16 and BFloat16, not just the new BF16 path, and CUDA/WebGPU still ignore outputs 1 and 2 — so the same model now behaves differently across EPs. Details inline. Consider splitting it out.

Suggestions (non-blocking)

• The contrib_op plumbing is now dead code — with the if constexpr (std::is_same_v<T, float>) guard both arms resolve to ComputeImpl<float, float>.
• The BF16 ComputeJob is a ~60-line near-verbatim copy of the MLFloat16 overload.
• ConvertMLFloat16ToFloatIfNeeded is misnamed now and still exists in two copies that have drifted (only the skip_layer_norm.cc copy guards tensor_size > 0).
• Stale comment above the static_assert in ComputeImpl: still says "Currently only instantiated for T in {float, double, MLFloat16}"; only the assertion message was updated.
• WriteStat<U>(..., static_cast<double>(mean)) widens a float to double only for gsl::narrow_cast<float> to narrow it straight back.
• Test file naming: layer_norm_bf16_cpu_test.cc also covers core opset-17 LayerNormalization and the MLFloat16/double stat regressions, and test/contrib_ops/ already has layer_norm_op_test.cc and layer_norm_test.cc. Consider folding in or renaming.

Verified clean

• No new allocation in a hot path beyond the pre-existing FP16 pattern; the conditional skip_input_bias_add_output_fp32 allocation from the last round is in.
• Direct includes are correct in narrow_float_utils.h (core/common/float16.h for BFloat16ToFloat/FloatToBFloat16, mlas.h for the MLAS converters) — no transitive-include reliance.
• scale_float_ptr is unconditionally dereferenced in the BF16 fast path but is guaranteed non-null for narrow-float T by the if constexpr (is_narrow_float_v<T>) block in ComputeWithoutContext.
• Input rank is validated (2 or 3) before dims.back() in the new stat_shape lambda.
• No C API surface touched; offsets use SafeInt<ptrdiff_t>.

[Tianlei Wu (tianleiwu)]

Reviewed the current head after the follow-up fixes. The CPU registration constraints, shared narrow-float and prepack handling, optional statistics, Welford update, and focused coverage are consistent; I found no remaining actionable issues.

Local validation: rebuilt onnxruntime_provider_test and ran 21 tests from 2 suites (LayerNormBFloat16CpuTest.* plus the two float SkipLayerNorm statistic tests); all passed.