[cpu kernels] native_layer_norm: layer_norm_scalar returns NaN on large-mean rows; Half/BF16 at N>=256 slow after #23153

[mergennachin]

🐛 Describe the bug

Follow-up to #23153. There are two problems in native_layer_norm, and they have independent fixes:

1. layer_norm_scalar returns NaN or badly wrong output on large-mean rows.
   • It is the whole portable kernel.
   • It has been the optimized kernel's path for N < 256 since Use scalar fast path in optimized layer_norm for small tensors (#18636) #18636, which shipped in v1.3.0.
   • Example: FP32 nn.LayerNorm(192) on rows with mean 100 and std 0.01 returns NaN for about half of the rows.
   • All four dtypes are affected.
2. With [CPU] Fix FP16/BF16 convolution and layer normalization #23153, Half/BF16 at N ≥ 256 take 4–6x as long as FP32. Measured at N = 512, 1024 and 4096. The cause is that [CPU] Fix FP16/BF16 convolution and layer normalization #23153 routes them to layer_norm_scalar.

Fix (2) below restores Half/BF16 speed and accuracy at N ≥ 256 on its own. Fix (1) is a local change to layer_norm_scalar. It fixes the FP32 regression and the portable kernel. Under #23153's routing it also fixes Half/BF16 accuracy at N ≥ 256 by itself. The two can land separately; (1) is more urgent.

Problem 1: one-pass float variance in layer_norm_scalar

normalization_ops_util.h#L38-L46:

float sum = std::accumulate(x, x + N, 0.0f);
...  sq_sum += static_cast<float>(x[j]) * x[j];
float mean_value = sum / N;
float variance = sq_sum / N - mean_value * mean_value;
float std = std::sqrt(variance + eps);

E[x²] − E[x]² cancels when the mean is large relative to the std. The variance is not clamped, and the error can go either way:

• Below -eps: rstd and the whole output row are NaN. For a constant FP32 row of 3000.0 at N=192, the float Σx² is 1727995392 instead of 1728000000, so the computed variance is −24.
• Above -eps: rstd is finite but wrong. A constant FP32 row of 1000.1 at N=255 gives rstd 0.949 instead of 316.2.

Double inputs are also accumulated in float.

Where layer_norm_scalar is used (all four dtypes: Float, Double, Half, BFloat16):

Kernel	Uses layer_norm_scalar for
portable native_layer_norm (#L61-L70)	every N
optimized native_layer_norm, main (#L75-L91)	N < 256 (since #18636; before that, Welford for every N)
optimized with #23153	N < 256, plus Half/BF16 at every N

Repro. No delegation is involved; the pybindings run the optimized kernel. The last two lines show main's separate Half/BF16 bug at N ≥ 256, which #23153 fixes.

import warnings

import torch
from executorch.exir import to_edge
from executorch.runtime import Runtime

warnings.filterwarnings("ignore")
rt = Runtime.get()


def check(dtype, N, mean, std, rows=64):
    torch.manual_seed(0)
    ln = torch.nn.LayerNorm(N).eval()
    x = (mean + std * torch.randn(rows, N, dtype=torch.float64)).to(dtype)
    ref = ln.double()(x.double())
    ln = ln.to(dtype)
    pte = to_edge(torch.export.export(ln, (x,))).to_executorch().buffer
    y_et = rt.load_program(pte).load_method("forward").execute([x])[0]
    res = []
    for y in (y_et, ln(x)):
        d = (y.double() - ref).abs()
        err = d[d.isfinite()].max().item() if d.isfinite().any() else float("nan")
        res += [err, y.isnan().sum().item()]
    print(f"{str(dtype)[6:]:8} N={N:<5} mean={mean:<5} std={std:<5} "
          f"ET: max_err={res[0]:<8.3g} nan={res[1]:>5}/{y.numel():<6} "
          f"eager: max_err={res[2]:.3g}")


check(torch.float32, 1024, 100, 0.01)  # N >= 256: vectorized Welford
check(torch.float32, 192, 100, 0.01)  # N < 256: layer_norm_scalar
check(torch.float32, 192, 1000, 0.1)
check(torch.float32, 192, 3000, 0.0)  # constant rows
check(torch.float16, 192, 100, 0.1)
check(torch.bfloat16, 192, 100, 0.1)
check(torch.float16, 4096, 100, 1.0)  # N >= 256: RowwiseMoments reads Half as float
check(torch.bfloat16, 4096, 100, 1.0)

float32  N=1024  mean=100   std=0.01  ET: max_err=0.00157  nan=    0/65536  eager: max_err=0.000767
float32  N=192   mean=100   std=0.01  ET: max_err=2.82     nan= 6336/12288  eager: max_err=0.000997
float32  N=192   mean=1000  std=0.1   ET: max_err=3.17     nan= 5952/12288  eager: max_err=0.000759
float32  N=192   mean=3000  std=0.0   ET: max_err=nan      nan=12288/12288  eager: max_err=0
float16  N=192   mean=100   std=0.1   ET: max_err=3.37     nan=    0/12288  eager: max_err=0.00105
bfloat16 N=192   mean=100   std=0.1   ET: max_err=3.86     nan=    0/12288  eager: max_err=0.0292
float16  N=4096  mean=100   std=1.0   ET: max_err=21.8     nan=    0/262144 eager: max_err=0.00189
bfloat16 N=4096  mean=100   std=1.0   ET: max_err=3.08     nan=    0/262144 eager: max_err=0.0148

A C++ harness covers cases the Python repro doesn't. Settings: M=64, affine, eps 1e-5, max abs error of out against a double reference. It uses a different seed from the repro, so NaN counts differ.

case	main optimized	#23153 optimized	main portable
fp32 N=1024 mean 100 std 0.01	0.0024	0.0024	4.4 (30720/65536 NaN)
fp16 N=4096 mean 100 std 0.1	265 (wrong moments)	4.77 (rstd err 75%)	4.77 (rstd err 75%)
bf16 N=4096 constant 300	0 (by accident: the misread mean ≈ 300.5 rounds to 300 in BF16)	all NaN	all NaN

Related: both kernels convert eps to CTYPE (optimized #L32, portable #L29). In Half, eps=1e-12 rounds to 0, so a constant Half row (e.g. all ones) produces NaN where PyTorch returns 0.

Problem 2: Half/BF16 performance with #23153

Setup: median ms on an Apple M1 Pro, single thread (the kernel has no parallel_for), M=4096, randn input, gamma and beta present. The A/B runs were interleaved on a loaded machine, so the ratios are more reliable than the absolute times. main's Half/BF16 results at N ≥ 256 are wrong and are shown only for comparison.

dtype	N	main	#23153	(1)+(2)
fp32	192	1.66	1.66	0.79
fp16	192	1.61	1.62	0.91
bf16	192	1.92	1.93	1.22
fp32	512 / 1024 / 4096	1.32 / 2.36 / 9.79	1.30 / 2.33 / 9.87	1.29 / 2.33 / 9.85
fp16	512 / 1024 / 4096	0.76 / 1.30 / 5.69 (wrong)	5.26 / 11.43 / 48.18	1.17 / 2.15 / 8.61
bf16	512 / 1024 / 4096	2.53 / 4.78 / 19.28 (wrong)	6.15 / 13.06 / 53.52	1.40 / 2.59 / 9.82

Most of the cost is the scalar loop, not the Half conversion: FP32 through layer_norm_scalar takes 41.2 ms at N=4096. With (1) alone, #23153's Half/BF16 path at N ≥ 256 gets 2.0–2.6x faster (e.g. 18.72 ms FP16 and 24.55 ms BF16 at N=4096). That is still 1.7–2.7x FP32. (2) brings Half/BF16 to 0.87–1.11x FP32.

Why main's vectorized Half/BF16 path is wrong

The load in UpdateMomentsVec is broken (moments_utils.h#L74-L78):

• It calls Vectorized<acc_t<T>>::loadu(X_ptr + j * Vec::size()), where X_ptr is a const Half* or const BFloat16* and acc_t<T> is float. This reinterprets pairs of 16-bit values as float bit patterns.
• It advances by Vectorized<float>::size() elements per load, but RowwiseMomentsImpl sizes chunks by Vectorized<T>::size(), which is twice as large (4 vs 8 on NEON). Part of each chunk is read twice and the rest is never read.

The header says ATen's BF16 specializations "are excluded" (#L11-L13), but the kernel has dispatched Half/BF16 since #7752.

A second problem shows up once the loads are fixed. The map3 lambda is generic (auto x, #L112-L121).

• Where Vectorized<Half/BFloat16> is specialized (arm64 builds without C10_MOBILE, and Buck Linux builds, which define AVX2): the lambda computes in the reduced type, with scale and offset rounded to it. This alone leaves outputs up to 77x above the rounding floor.
• Android/iOS (C10_MOBILE), not measured: the NEON Half/BF16 specializations are disabled there, so the map3 problem does not apply. But convert_to_float is a scalar fallback there, so the speedup from (2) needs to be measured on device.
• x86: OSS CMake builds use the generic Vectorized. x86 was not measured.

The header hazard remains after #23153. moments_utils is an exported, PUBLIC target (targets.bzl#L79-L89). RowwiseMoments<Half> and RowwiseMoments<BFloat16> still compile and still return wrong moments for any other caller. For example, on {2, 3, 4, 5, 9, 10, 12, 13} (mean 7.25), Half at N=8 on arm64 gives a mean of 1.18e6.

The existing short case in moments_utils_test.cpp (#L15-L18) fails on every architecture:

• arm64: it hits the same reinterpretation (mean 196610).
• x86: acc_t<short> is int32_t, so the math is integer (mean 6, variance 26).

Nobody noticed because moments_utils_test_bin is a Buck cxx_binary, not a test target (test/targets.bzl#L41), and CMake doesn't build it.

Proposed fix

1. Make the variance in layer_norm_scalar stable. Use a corrected two-pass computation with 8 independent lane accumulators, so that clang vectorizes it:

   float mean_value = sum / N;
   // per lane: d = x[j] - mean_value; d_sum += d; sq_sum += d * d;
   float variance = std::max((sq_sum - d_sum * d_sum / N) / N, 0.0f);
   mean_value += d_sum / N;

   Three related changes were not prototyped: the std::max clamp, accumulating Double inputs in double, and passing eps without rounding it to CTYPE.
2. Port ATen's reduced-precision UpdateMomentsVec overload (ATen moments_utils.h#L81-L111). It does Vectorized<T>::loadu, then convert_to_float, then accumulates into two float vectors. RowwiseMomentsImpl already computes kVecSize and m0_add correctly for this. Then:
   • Go back to routing on N < kSmallNThreshold only.
   • Type the map3 lambda on Vectorized<float>, so that ATen's convert-through-float map3 (in functional_bfloat16.h) is selected. Alternatively, drop map3 and use the scalar normalize loop, which clang auto-vectorizes. On arm64 that changed time by −4% to +6%, with the same accuracy.
   • In RowwiseMoments, static_assert that T is float, double, Half, or BFloat16. Replace the short test case with Half/BF16 cases.

Prototype accuracy. The prototypes are a standalone harness, not gtest. Timings are in the Problem 2 table. The table shows max abs error of out with M=64; "floor" is the error of the correctly rounded result.

case	floor	main	#23153	(1)+(2)
fp32 N=192 mean 100 std 0.01	1.9e-7	7.8 (5568/12288 NaN)	7.8 (5568/12288 NaN)	5.6e-4
fp16 N=192 mean 100 std 1	1.9e-3	0.010	0.010	1.9e-3
bf16 N=512 randn	0.015	3.0	0.015	0.015
bf16 N=4096 mean 100 std 1	0.016	3.6	0.49 (rstd err 9.3%)	0.016 (rstd err 0.25%)
fp16 N=4096 mean 100 std 1	1.9e-3	32	0.14 (rstd err 2.7%)	1.9e-3 (rstd err 0.046%)

• Each fix alone:
  • (1) alone, with [CPU] Fix FP16/BF16 convolution and layer normalization #23153's routing, matches (1)+(2) on every Half/BF16 case at N ≥ 256.
  • (2) alone matches (1)+(2) on those cases too, but leaves the N < 256 cases unchanged.
• FP32 row: the remaining error comes from float rounding of the mean itself. Eager PyTorch shows about 1e-3 on a similar input in the repro above.

Notes:

• Alternatives to (1), measured on the FP32 row above:
  • Scalar Welford: 6.3e-3 error, and 2.3x (FP32) to 2.7x (FP16) slower than the current one-pass loop at N=192.
  • Plain two-pass without the correction term: 8.0e-3 error.
  • The corrected two-pass without lane accumulators: as accurate as (1), but not vectorized (2.2 ms vs 0.79 ms).
• Low-bit changes. (1) changes the summation order, so portable FP32 results shift in the low bits.
• [CPU] Fix FP16/BF16 convolution and layer normalization #23153's exact-equality test. With (2), the BFloat16LargeRows test fails on the width-514 row: the mean is −8.67e-9 instead of 0, while out and rstd stay exact. That assertion uses EXPECT_TENSOR_EQ and needs a tolerance. (1) alone passes it.

Suggested tests

• FP32/Double at N < 256 (portable and optimized):
  • Rows with mean ≫ std (mean 100, std 0.01, N=192).
  • Constant rows (3000.0, N=192).
  • Expect finite outputs and rstd close to a double reference.
• Half/BF16 at N ∈ {257, 512, 4096}:
  • Mean 100, std 1, plus a constant BF16 row of 300 at N=4096. [CPU] Fix FP16/BF16 convolution and layer normalization #23153's LargeRows tests use small means, so they don't exercise precision.
  • Use EXPECT_TENSOR_CLOSE_WITH_TOL with atol at the rounding floor. The default Half/BF16 atol (1e-3 / 1e-2) is tighter than one output ulp at |out| ≈ 2–4.
• Half constant row with eps=1e-12.
• moments_utils_test.cpp:
  • Add Half and BF16 cases with N large enough to reach the vector path on every architecture (e.g. 512 and 514).
  • Make it a real test target and build it in CMake, so that it runs in OSS CI.

The investigation, the C++ harness and prototypes, and this issue were done with Claude Code.

Versions

• ExecuTorch main @ 00e5030. The "[CPU] Fix FP16/BF16 convolution and layer normalization #23153" columns have [CPU] Fix FP16/BF16 convolution and layer normalization #23153 applied on top.
• Python repro: a locally built executorch 1.5.0 wheel (5b3da18), torch 2.13.0, Python 3.11. Released wheels 1.3.0+ should reproduce the FP32 N < 256 rows.
• C++ harness: Apple clang 17, -O3, ATen vec headers from torch 2.13.0.
• macOS 26 arm64, Apple M1 Pro, single thread.

cc @larryliu0820 @manuelcandales @JakeStevens

[mergennachin]

The scalar correctness part is now addressed in #23153 by 0769928: corrected two-pass moments in the shared helper, FP64 accumulation for Double, and epsilon preserved in compute precision. Six new regressions fail on the previous portable and optimized kernels and pass with the fix, including the reported FP16 NaN and tiny-epsilon case. Portable, optimized, ATen, and UBSan suites pass, along with an 800-case local numerical grid.

The reduced-precision SIMD moments/normalization work, exported RowwiseMoments header hazard, and moments-test registration remain follow-ups here. This issue should stay open for that work.

[mergennachin]

For SIMD work, not sure if it's necessary or hi-pri since we'll have it on XNNPACK anyway