import torch
from torch import nn

norm = nn.LayerNorm(128, eps=1e-5)
linear = nn.Linear(128, 128, bias=False)

def f(x):
    return linear(norm(x))

compiled_f = torch.compile(f, fullgraph=True)

with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    torch.manual_seed(42)
    x = torch.randn(4, 32, 32, 128)
    out_eager = f(x)
    out_compiled = compiled_f(x)
    torch.testing.assert_close(out_eager, out_compiled)

Before fix output:

AssertionError: Tensor-likes are not close!
Mismatched elements: 1 / 524288 (0.0%)
Greatest absolute difference: 0.00030517578125
Greatest relative difference: 0.032958984375

After fix output: (no error)

Worklog: Issue #168126 - LayerNorm bf16 divergence

Run 1: Diagnosis and Plan

Build setup

• Job venv uses Python 3.12; built pytorch CPU-only (USE_CUDA=0 USE_MKL=0 BLAS=Eigen)
• Required PYTHONPATH override due to stale /home/aorenste/local/pytorch in sys.path

Reproduction

• Confirmed minimal repro (CPU, no CUDA) fails: 2 mismatched elements out of 524288
• Max abs diff: 0.000275, max rel diff: 0.123 (exceeds bf16 tolerances)
• Confirmed: aot_eager PASS, aot_eager_decomp_partition FAIL, inductor FAIL
• Confirmed: fp32 PASS, bf16 (with or without autocast) FAIL

Key diagnostic results

1. aot_eager_decomp_partition without layer_norm decomp: PASS — removing the
   decomposition of native_layer_norm fixes the issue entirely

2. Decomp formula with EAGER stats: diff = 0.0 — the decomposition's
   (x - mean) * rstd * weight + bias formula is BIT-EXACT with the C++ kernel
   when given the same mean/rstd

3. Decomp formula with DECOMP stats: diff = 7.15e-7 — the ENTIRE error comes
   from the mean/variance computation being different

4. torch.var_mean vs RowwiseMoments: mean differs by up to 4.66e-08,
   variance by up to 5.5e-7 across 76% of rows. The C++ native_layer_norm kernel
   uses Welford's algorithm with cascade summation (RowwiseMoments), while the
   decomposition uses torch.var_mean which uses a different reduction algorithm.

5. rsqrt vs 1/sqrt: ZERO difference (not a factor)

6. bf16 rounding amplification: 7 mismatched bf16 elements after casting norm
   output → after matmul amplification: max diff 0.125, 23 mismatched elements

Root cause

The decomposition of aten.native_layer_norm in torch/_refs/__init__.py::_normalize
computes mean/variance using torch.var_mean, while the C++ kernel uses
RowwiseMoments (Welford + cascade sum). These give different fp32 results
(≤7e-7), which cause bf16 rounding divergences at ~0.001% of elements, then
amplified by subsequent ops (linears, matmuls).

Run 2

User: Verify if the cuda implementation has the same problem and would have the same solution

CUDA kernel analysis

Examined aten/src/ATen/native/cuda/layer_norm_kernel.cu. The CUDA native_layer_norm
kernel has two paths:

1. Vectorized path (line 1114, for float/half/bf16 with aligned data, N < 2^24,
   N % 4 == 0): Uses compute_stats() → cuWelfordOnlineSum() + cuWelfordCombine()
   — a custom Welford implementation with vectorized loads, warp shuffle reduction,
   and shared-memory inter-warp reduction.

2. Non-vectorized path (line 1117): Uses RowwiseMomentsCUDAKernel →
   WelfordOps with cuda_utils::BlockReduce — Welford's algorithm through the
   standard CUDA block reduction utility.

Both paths use Welford's online algorithm, but with CUDA-specific parallelism
(warp shuffles, shared memory).

CUDA torch.var_mean analysis

Examined aten/src/ATen/native/cuda/ReduceMomentKernel.cu. The CUDA var_mean
kernel (std_var_kernel_impl) also uses WelfordOps, but goes through
gpu_reduce_kernel<scalar_t, out_t, 2> — the generic GPU reduction framework
with unrolling factor 2.

Same problem, different Welford implementations

Both the native_layer_norm CUDA kernel and torch.var_mean on CUDA use
Welford's algorithm, but through different code paths:

• native_layer_norm: Per-row Welford with vectorized loads (4 elements at a time),
  intra-warp shuffle reduction, inter-warp shared-memory tree reduction. Thread
  layout: blockDim.x * blockDim.y threads per row, one row per block.

• torch.var_mean (via gpu_reduce_kernel): Generic multi-dimensional reduction
  framework with unrolling factor 2, different thread/block decomposition and
  accumulation order.

The different accumulation orders produce different fp32 rounding, just like the
CPU case. The errors may be slightly different in magnitude due to different
parallelism granularity, but the fundamental issue is identical.

Decomposition guard analysis

In torch/_inductor/decomposition.py (line 206-217):

@register_decomposition(aten.native_layer_norm)
def _native_layer_norm(input, normalized_shape, weight, bias, eps):
    if input.is_mtia:
        return NotImplemented
    return decomp_native_layer_norm(input, normalized_shape, weight, bias, eps)

Only MTIA is excluded from decomposition. CUDA tensors ARE decomposed, so CUDA
hits the same _normalize() → torch.var_mean() path that causes the divergence.

Conclusion

Yes, the CUDA path has the same problem and the same solution would fix it.
The fix (making _native_layer_norm return NotImplemented unconditionally, or
at least for CUDA too) prevents the decomposition and keeps the native CUDA kernel,
which eliminates the statistics mismatch. This is device-agnostic — the root cause
is the decomposition replacing a specialized per-row Welford with a generic
reduction Welford, regardless of whether the device is CPU or CUDA.

Run 3

User: I agree with the reproduction and the proposed solution. Implement and verify it.

Your job:

• write the smallest fix that explains the reproduced failure
• add or update a focused regression test
• demonstrate that the test would fail before the fix
• demonstrate that the same test passes with the fix
• keep scope tight to the approved diagnosis and plan

Rules:

• no speculative cleanup or unrelated refactors
• if the test does not fail before the fix, stop and report that the evidence is insufficient
• if the fix requires a broader change than planned, stop and ask before expanding scope
• prefer the narrowest test that captures the reported failure mode

Output:

• patch summary
• regression test added or updated
• evidence that it fails before the fix
• evidence that it passes after the fix
• any remaining risks or gaps

Run 4

User: Continue with the build.

Run 5

User: Rebuild and then continue.

Build

• Rebuilt PyTorch successfully (CPU-only, incremental build)

Verification: test fails before fix, passes after

• Monkey-patched old decomposition back: test correctly fails (1/524288 mismatched,
  max_abs_diff=0.000305, max_rel_diff=0.033)
• With fix applied: test passes (0 mismatched)

In-tree regression test

• test_layer_norm_bf16_numerics_cpu passes (3.9s)

Full model repro: discovered additional inductor fusion issue

• The full TriangleMultiplicationOutgoing repro still fails with inductor
• Investigated: aot_eager_decomp_partition now matches eager perfectly (0/262144)
• The remaining divergence is purely from inductor code generation
• Root cause: inductor fuses sigmoid + multiply into one kernel, computing
  intermediates in fp32 without bf16 truncation between ops
• p_in(x) * g_in(x).sigmoid() with bf16: 138080/524288 (26%) mismatched
• This is a separate, pre-existing inductor behavior (not LayerNorm-related)
• TORCHINDUCTOR_EMULATE_PRECISION_CASTS=1 does not fix it

Linting

• All changed files pass linting (RUFF, PYFMT)

Artifacts generated

• fix.diff — 3 files changed: decomposition.py, lowering.py, test_torchinductor.py
• report.md — updated with full analysis including the inductor fusion gap