[Inductor] Fix NaN output in LayerNorm CPU by guarding Welford variance. - #173989

The committers listed above are authorized under a signed CLA.

• ✅ login: karthik-0306 / name: karthik-0306 (e6cfb7d)

@pytorchbot label "release notes: inductor"

The label oncall: cpu inductor is only applicable to issues and has been removed. Please only use this label on issues.

if it repros in cpu but not cuda it's probably not a lowering issue, just a codegen one.

if it repros in cpu but not cuda it's probably not a lowering issue, just a codegen one.

That makes sense, if CUDA isn't seeing this, it’s definitely a CPU codegen mismatch rather than a lowering bug. I'll dig into the generated C++ kernel for the Welford reduction and see where the inf handling is diverging from the Triton path. Will update once I've isolated the logic in the CPU backend.

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/173989

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Run pull request jobs on OSDC runners in shadow mode

❌ You can merge normally! (1 Unrelated Failure), 1 Unclassified Failure

As of commit 8b040c0 with merge base 868be1c ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@claude explain those CI failures, check if they are related, and review the changes in this PR

Claude finished @jansel's task in 2m 3s —— View job

CI Failures Analysis & Code Review

• Examine the PR code changes
• Analyze CI failures and determine if they're related to the PR
• Review the code changes for correctness and appropriateness

CI Failures

All 9 "New Failures" are unrelated to this PR. They all exit with code 127 ("command not found"), which indicates a missing binary or broken test environment — not a logic error from the code change. These are infrastructure failures (likely a test runner dependency issue on the linux-jammy-py3.10-clang18, linux-jammy-py3.14-clang18, linux-jammy-py3.14t-clang18, and linux-jammy-py3.10-gcc11 runners). The remaining flaky failures (GraphExecutorTest.runAsync_executor, test_open_device_registration) are pre-existing trunk issues with links to known flaky test tickets.

None of these failures are related to the 1-line change in triton_helpers.py.

Code Review

The change (torch/_inductor/runtime/triton_helpers.py:264-266):

@triton.jit
def welford_combine(mean_1, m2_1, weight_1, mean_2, m2_2, weight_2):
    # Handle potential NaN if both means are infinite
    delta = tl.where(mean_1 == mean_2, 0.0, mean_2 - mean_1)

Issues

1. Wrong fix location for the reported bug. The issue ([Inductor][CPU][float16] LayerNorm outputs NaN when input contains Inf #173885) reports NaN on CPU with float16. The CPU codegen path goes through the C++ welford_combine in torch/csrc/inductor/cpp_prefix.h:194, which has the same unguarded delta = b.mean - a.mean. This PR only fixes the Triton runtime helper, which is used for GPU kernels. As @eellison noted, if it repros on CPU but not CUDA, the fix needs to be in the CPU codegen path, not the Triton path.

2. The IR-level welford_combine_fn in torch/_inductor/ir.py:1239 also has unguarded delta = b_mean - a_mean. This is used for non-Triton backends. If the intent is to fix the general Welford algorithm, this should also be addressed.

3. Semantic correctness concern. The guard mean_1 == mean_2 → delta = 0 is mathematically correct for inf == inf and -inf == -inf, but it also fires for any two equal finite means, where the original mean_2 - mean_1 would already be exactly 0.0. In the finite case this is harmless, but the equality comparison on floats adds a branch to every Welford combine step in every Triton reduction kernel — this is a hot path. A more targeted guard like tl.where((mean_1 == mean_2) & ~tl.math.isfinite(mean_1), 0.0, mean_2 - mean_1) would limit the semantic change, but tl.where with the simpler condition is probably fine for Triton since the compiler can optimize it. Still, worth discussing.

4. Missing test. There's no regression test demonstrating the fix works. A test with float16 input containing inf values passed through a compiled LayerNorm would both validate the fix and prevent regression.

Summary

The fix logic itself is sound (guarding inf - inf → NaN in Welford), but it's applied to the wrong code path for the reported bug. To actually fix #173885 on CPU, the same guard needs to be in torch/csrc/inductor/cpp_prefix.h (and potentially torch/_inductor/ir.py). The Triton helper change may still be worthwhile as a defensive improvement, but it doesn't address the CPU issue reported in the linked bug.

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

• ciflow/inductor
• ciflow/torchtitan

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

Thanks for the review @jansel and @claude — the bot's analysis was spot on.
I've updated the PR to address all the points raised:

cpp_prefix.h — added the isinf guard to the CPU C++ Welford combine (scalar path + vectorized blendv path), which is the actual root cause for the CPU bug
ir.py — added the same guard to welford_combine_fn for non-Triton backends
triton_helpers.py — tightened the existing guard to also check isinf instead of just equality, for consistency
test_cpu_repro.py — added a regression test covering both float16 and bfloat16 with inf inputs

The CI failures (exit code 127) are all pre-existing infrastructure issues unrelated to this change.
@eellison @EikanWang — happy to hear any feedback on the approach.

@claude explain those CI failures

Claude finished @jansel's task in 5m 12s —— View job

CI Failure Analysis

• Gather current CI status
• Download and analyze failed job logs
• Determine if failures are related to this PR
• Summarize findings

Findings

There are 2 failed workflow runs: Lint and inductor. Here's the breakdown:

1. Lint failure (lintrunner-noclang-partial) — Caused by this PR

The PYFMT linter flags a formatting issue in torch/_inductor/runtime/triton_helpers.py. The multi-line tl.where(...) call should be collapsed to a single line:

# Current (fails lint):
    delta = tl.where(
        tl.math.isinf(mean_1) & (mean_1 == mean_2), 0.0, mean_2 - mean_1
    )

# Expected (passes lint):
    delta = tl.where(tl.math.isinf(mean_1) & (mean_1 == mean_2), 0.0, mean_2 - mean_1)

Fix: run lintrunner -a to auto-format.

2. GPU inductor benchmark failures (5 jobs) — Caused by this PR

All 5 GPU inductor test failures (inductor_huggingface, inductor_timm x2, inductor_torchbench x2) fail with the same root cause:

triton.compiler.errors.CompilationError: at 5:8:
def welford_combine(mean_1, m2_1, weight_1, mean_2, m2_2, weight_2):
    delta = tl.where(
        tl.math.isinf(mean_1) & (mean_1 == mean_2), 0.0, mean_2 - mean_1
        ^
AttributeError("module 'triton.language.math' has no attribute 'isinf'")

tl.math.isinf does not exist in the version of Triton used in CI. The correct Triton API is likely one of:

• tl.extra.cuda.libdevice.isinf(x) (for GPU)
• Using (x == float('inf')) | (x == float('-inf')) as a portable alternative
• Or simply (mean_1 == mean_2) as the original simpler guard (which was already correct for the inf case since inf == inf is True)

Every model that uses a Welford reduction (LayerNorm, BatchNorm, etc.) fails to compile on GPU because of this.

3. Unclassified failures (7 jobs) — Not caused by this PR

These ran on the merge base and DrCI couldn't classify them. Based on the log content:

• test_numpy_ref_linalg_tensorinv_cpu_float64 (AMX) — unrelated numerics test
• test_inductor_sequence_nr (cpp_wrapper) — segfault in an unrelated test
• test_tmp_not_defined_issue3_cuda, test_batch_norm_2d_2_cuda, test_issue_140766_cuda — pre-existing GPU test issues
• test_ddp_optimizer_cudagraph (distributed) — unrelated distributed test
• test_2d_welford_reduction... (triton-cpu) — this one could be related if the triton-cpu backend also hits the isinf issue, but the test name suggests it's a strided blocks test that may be pre-existing

Summary

6 of the 6 "New Failures" are caused by this PR — the tl.math.isinf() call in triton_helpers.py does not exist in Triton's API. The fix should either use a portable isinf check or revert to the simpler mean_1 == mean_2 guard (which already handles the inf-inf case correctly since inf == inf evaluates to True in IEEE 754). The unclassified failures appear to be pre-existing trunk issues.

Updated based on @claude bot's CI analysis:

Lint fix — collapsed the tl.where(...) to a single line as flagged by PYFMT
Triton API fix — replaced tl.math.isinf(mean_1) with just mean_1 == mean_2, since tl.math.isinf doesn't exist in the Triton version used in CI. The simpler guard still correctly handles the inf - inf case since inf == inf is True in IEEE 754

Both CI failures were caused by the same tl.math.isinf call.

@pytorchbot merge

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Merge failed

Reason: 1 jobs have failed, first few of them are: inductor / unit-test / inductor-cpu-test / test (inductor_amx, 2, 2, linux.2xlarge.amx)

@jansel the merge failed due to the inductor_amx (2, 2) failure which is a pre-existing infra issue on the linux.2xlarge.amx runner — the annotation on the merge job itself shows it's a Node.js 20 deprecation issue on PyTorch's CI infrastructure. Note that inductor_amx (1, 2) passed on the same runner type, confirming this is not caused by this PR. The distributed job was also auto-cancelled after 245 minutes due to a runner timeout. Could you rerun the inductor_amx job or use @pytorchbot merge -f to merge past the infra failure?

@pytorchbot merge -i

Merge started

Your change will be merged while ignoring the following 2 checks: pull / linux-jammy-py3.10-gcc11 / test-osdc (distributed, 1, 3, mt-l-x86iamx-8-64), inductor / unit-test / inductor-cpu-test / test (inductor_amx, 2, 2, linux.2xlarge.amx)

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

My oncall agent has flagged this PR to be responsible for compile time increase and speedup decreases. See the summary at internalfb.com/phabricator/paste/view/P2340863849

My oncall agent has flagged this PR to be responsible for compile time increase and speedup decreases. See the summary at internalfb.com/phabricator/paste/view/P2340863849

@anijain2305 would you prefer I revert the triton_helpers.py guard entirely since the original bug was CPU-only, or is there a preferred pattern in the codebase for guarding non-finite inputs in Triton reductions without hot path overhead?