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

🐛 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 #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 #23153, Half/BF16 at N ≥ 256 take 4–6x as long as FP32. Measured at N = 512, 1024 and 4096. The cause is that #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 #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.
• #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. #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 @ 00e50304c68e4bfa645ef2e8432f1795cfa2f6b7. The "#23153" columns have #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