pytorch - ✅(Solved) Fix CUDA layer_norm produces incorrect output when flattened tensor size exceeds 2**32, likely 32-bit offset overflow in vectorized kernel [1 pull requests, 1 comments, 2 participants]

Official PRs (…)
ON THIS PAGE

Recommended Tools

×6

Utilities matched from this issue’s tags and category — try them while you read without losing context.

GitHub issue graph ai analysis

Paste a GitHub issue URL. We fetch that issue, discover linked issues from bodies/comments/timeline, collect linked pull requests, and produce a structured English report.

The report is written in English Markdown for sharing and archival.

Helpful · Quick feedback

Loading…
GitHub stats
pytorch/pytorch#181555Fetched 2026-04-28 06:24:46
View on GitHub
Comments
1
Participants
2
Timeline
46
Reactions
0
Author
Participants
Timeline (top)
mentioned ×19subscribed ×19labeled ×5commented ×1

torch.layer_norm on CUDA appears to produce incorrect output for large 2D inputs when the total number of elements crosses 2**32.

For an input tensor feats with shape (M, C), I observed that when:

M * C >= 2**32

part of the output becomes all zeros or otherwise incorrect. The failure seems correlated with the flattened linear offset crossing the 32-bit boundary.

Error Message

Given a CUDA tensor:

Root Cause

I looked at the CUDA layer norm kernel and suspect the issue may be caused by 32-bit integer overflow in the vectorized forward kernel.

Fix Action

Fixed

PR fix notes

PR #181600: [CUDA] Fix int32 overflow in layer_norm for tensors with >2^32 elements

Description (problem / solution / changelog)

Cast blockIdx.x to int64_t before multiplying by N to prevent 32-bit overflow when computing row offsets. Affects vectorized forward and backward kernels where blockIdx.x * N wraps around for large M.

Fixes https://github.com/pytorch/pytorch/issues/181555

Changed files

  • aten/src/ATen/native/cuda/layer_norm_kernel.cu (modified, +17/-13)
  • test/test_nn.py (modified, +18/-0)

Code Example

M * C >= 2**32

---

feats.shape == (M, C)

---

out = torch.layer_norm(
    feats,
    normalized_shape=(C,),
    weight=None,
    bias=None,
    eps=1e-5,
)

---

import torch

device = "cuda"
dtype = torch.float16  # also worth testing float32 / bfloat16

C = 4096
boundary_row = (2**32) // C

# Important: make M larger than boundary_row, so that we can inspect rows
# after the flattened 2**32 offset boundary.
M = boundary_row + 64

feats = torch.randn((M, C), device=device, dtype=dtype)

out = torch.layer_norm(
    feats,
    normalized_shape=(C,),
    weight=None,
    bias=None,
    eps=1e-5,
)

torch.cuda.synchronize()

print("shape:", feats.shape)
print("numel:", feats.numel())
print("boundary_row:", boundary_row)

for r in [
    boundary_row - 2,
    boundary_row - 1,
    boundary_row,
    boundary_row + 1,
    boundary_row + 2,
    boundary_row + 16,
    boundary_row + 32,
]:
    if 0 <= r < M:
        ref = torch.layer_norm(
            feats[r:r + 1],
            normalized_shape=(C,),
            weight=None,
            bias=None,
            eps=1e-5,
        )

        print(
            r,
            "out_abs_max =", out[r].abs().max().item(),
            "out_abs_mean =", out[r].abs().mean().item(),
            "max_diff_vs_single_row_ref =", (out[r:r + 1] - ref).abs().max().item(),
        )

---

(conda-env-name) [xxx@yyy zzz]# python layernorm_min.py 
shape: torch.Size([1048640, 4096])
numel: 4295229440
boundary_row: 1048576
row = 1048574 out_abs_max = 4.4453125 out_abs_mean = 0.79638671875 max_diff_vs_single_row_ref = 0.0
row = 1048575 out_abs_max = 4.37109375 out_abs_mean = 0.79931640625 max_diff_vs_single_row_ref = 0.0
row = 1048576 out_abs_max = 0.0 out_abs_mean = 0.0 max_diff_vs_single_row_ref = 3.48046875
row = 1048577 out_abs_max = 0.0 out_abs_mean = 0.0 max_diff_vs_single_row_ref = 3.61328125
row = 1048578 out_abs_max = 0.0 out_abs_mean = 0.0 max_diff_vs_single_row_ref = 3.67578125
row = 1048592 out_abs_max = 0.0 out_abs_mean = 0.0 max_diff_vs_single_row_ref = 4.125
row = 1048608 out_abs_max = 0.0 out_abs_mean = 0.0 max_diff_vs_single_row_ref = 3.849609375

---

auto i1 = blockIdx.x;
const T* block_row = X + i1 * N;
...
vec_t* Y_vec = reinterpret_cast<vec_t*>(Y + i1 * N);

---

blockIdx.x * N >= 2**32

---

const int64_t i = blockIdx.x;
const int64_t index = i * N + j;
RAW_BUFFERClick to expand / collapse

🐛 Describe the bug

Summary

torch.layer_norm on CUDA appears to produce incorrect output for large 2D inputs when the total number of elements crosses 2**32.

For an input tensor feats with shape (M, C), I observed that when:

M * C >= 2**32

part of the output becomes all zeros or otherwise incorrect. The failure seems correlated with the flattened linear offset crossing the 32-bit boundary.

Observed behavior

Given a CUDA tensor:

feats.shape == (M, C)

and applying layer norm over the last dimension:

out = torch.layer_norm(
    feats,
    normalized_shape=(C,),
    weight=None,
    bias=None,
    eps=1e-5,
)

when M * C >= 2**32, part of out becomes zero or incorrect. The issue is not observed for smaller tensors where M * C < 2**32.

The problem appears to affect the forward CUDA path.

Expected behavior

torch.layer_norm should produce correct output regardless of whether the flattened number of elements exceeds 2**32, as long as the tensor size is otherwise valid and supported.

Minimal reproduction

import torch

device = "cuda"
dtype = torch.float16  # also worth testing float32 / bfloat16

C = 4096
boundary_row = (2**32) // C

# Important: make M larger than boundary_row, so that we can inspect rows
# after the flattened 2**32 offset boundary.
M = boundary_row + 64

feats = torch.randn((M, C), device=device, dtype=dtype)

out = torch.layer_norm(
    feats,
    normalized_shape=(C,),
    weight=None,
    bias=None,
    eps=1e-5,
)

torch.cuda.synchronize()

print("shape:", feats.shape)
print("numel:", feats.numel())
print("boundary_row:", boundary_row)

for r in [
    boundary_row - 2,
    boundary_row - 1,
    boundary_row,
    boundary_row + 1,
    boundary_row + 2,
    boundary_row + 16,
    boundary_row + 32,
]:
    if 0 <= r < M:
        ref = torch.layer_norm(
            feats[r:r + 1],
            normalized_shape=(C,),
            weight=None,
            bias=None,
            eps=1e-5,
        )

        print(
            r,
            "out_abs_max =", out[r].abs().max().item(),
            "out_abs_mean =", out[r].abs().mean().item(),
            "max_diff_vs_single_row_ref =", (out[r:r + 1] - ref).abs().max().item(),
        )

In my case, rows after the 2**32 flattened offset boundary appear to become zero or incorrect, and differ from the single-row reference result.

(conda-env-name) [xxx@yyy zzz]# python layernorm_min.py 
shape: torch.Size([1048640, 4096])
numel: 4295229440
boundary_row: 1048576
row = 1048574 out_abs_max = 4.4453125 out_abs_mean = 0.79638671875 max_diff_vs_single_row_ref = 0.0
row = 1048575 out_abs_max = 4.37109375 out_abs_mean = 0.79931640625 max_diff_vs_single_row_ref = 0.0
row = 1048576 out_abs_max = 0.0 out_abs_mean = 0.0 max_diff_vs_single_row_ref = 3.48046875
row = 1048577 out_abs_max = 0.0 out_abs_mean = 0.0 max_diff_vs_single_row_ref = 3.61328125
row = 1048578 out_abs_max = 0.0 out_abs_mean = 0.0 max_diff_vs_single_row_ref = 3.67578125
row = 1048592 out_abs_max = 0.0 out_abs_mean = 0.0 max_diff_vs_single_row_ref = 4.125
row = 1048608 out_abs_max = 0.0 out_abs_mean = 0.0 max_diff_vs_single_row_ref = 3.849609375

Suspected cause

I looked at the CUDA layer norm kernel and suspect the issue may be caused by 32-bit integer overflow in the vectorized forward kernel.

In the vectorized layer norm implementation, the row pointer is computed roughly like this:

auto i1 = blockIdx.x;
const T* block_row = X + i1 * N;
...
vec_t* Y_vec = reinterpret_cast<vec_t*>(Y + i1 * N);

Here, blockIdx.x is a CUDA built-in with a 32-bit type, and N is passed as int in the vectorized path. Therefore i1 * N may be computed in 32-bit arithmetic before being used as a pointer offset.

When the logical row offset satisfies:

blockIdx.x * N >= 2**32

the offset may wrap around, causing reads/writes to the wrong location. This matches the observed threshold where M * C >= 2**32.

The non-vectorized fallback path appears to use int64_t for the row index / linear index, for example:

const int64_t i = blockIdx.x;
const int64_t index = i * N + j;

so the issue may only affect the vectorized fast path.

The fast path selection appears to check that the per-row normalized dimension N is small enough and vectorizable, but it does not appear to check whether the total flattened offset M * N can exceed the 32-bit range.

Versions

Collecting environment information... PyTorch version: 2.7.1+cu126 Is debug build: False CUDA used to build PyTorch: 12.6 ROCM used to build PyTorch: N/A

OS: CentOS GCC version: (GCC) 13.3.1 20240611 (Red Hat 13.3.1-2) Clang version: 19.1.7 (https://github.com/conda-forge/clangdev-feedstock 966d7cbbe6cf2286c2ec30c49d8e957ba220cf18) CMake version: version 3.28.3 Libc version: glibc-2.28

Python version: 3.11.8 | packaged by conda-forge | (main, Feb 16 2024, 20:53:32) [GCC 12.3.0] (64-bit runtime) Python platform: Linux-5.4.241-1-tlinux4-0017.7-x86_64-with-glibc2.28 Is CUDA available: True CUDA runtime version: 12.8.93 CUDA_MODULE_LOADING set to: LAZY GPU models and configuration: CANNOT-DISCLOSE

Nvidia driver version: 535.161.08 cuDNN version: Probably one of the following: /usr/lib64/libcudnn.so.9.7.0 /usr/lib64/libcudnn_adv.so.9.7.0 /usr/lib64/libcudnn_cnn.so.9.7.0 /usr/lib64/libcudnn_engines_precompiled.so.9.7.0 /usr/lib64/libcudnn_engines_runtime_compiled.so.9.7.0 /usr/lib64/libcudnn_graph.so.9.7.0 /usr/lib64/libcudnn_heuristic.so.9.7.0 /usr/lib64/libcudnn_ops.so.9.7.0 Is XPU available: False HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True Caching allocator config: N/A

CPU: Architecture: x86_64 CPU op-mode(s): 32-bit, 64-bit Byte Order: Little Endian CPU(s): 384 On-line CPU(s) list: 0-383 Thread(s) per core: 2 Core(s) per socket: 96 Socket(s): 2 NUMA node(s): 2 Vendor ID: AuthenticAMD BIOS Vendor ID: Advanced Micro Devices, Inc. CPU family: 25 Model: 17 Model name: AMD EPYC 9K84 96-Core Processor BIOS Model name: AMD EPYC 9K84 96-Core Processor
Stepping: 1 CPU MHz: 3700.017 CPU max MHz: 2600.0000 CPU min MHz: 1500.0000 BogoMIPS: 5200.11 Virtualization: AMD-V L1d cache: 32K L1i cache: 32K L2 cache: 1024K L3 cache: 32768K NUMA node0 CPU(s): 0-95,192-287 NUMA node1 CPU(s): 96-191,288-383 Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local avx512_bf16 clzero irperf xsaveerptr wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid overflow_recov succor smca fsrm flush_l1d

Versions of relevant libraries: [pip3] mypy_extensions==1.1.0 [pip3] numpy==2.2.0 [pip3] nvidia-cublas-cu12==12.6.4.1 [pip3] nvidia-cuda-cupti-cu12==12.6.80 [pip3] nvidia-cuda-nvrtc-cu12==12.6.77 [pip3] nvidia-cuda-runtime-cu12==12.6.77 [pip3] nvidia-cudnn-cu12==9.5.1.17 [pip3] nvidia-cufft-cu12==11.3.0.4 [pip3] nvidia-curand-cu12==10.3.7.77 [pip3] nvidia-cusolver-cu12==11.7.1.2 [pip3] nvidia-cusparse-cu12==12.5.4.2 [pip3] nvidia-cusparselt-cu12==0.6.3 [pip3] nvidia-nccl-cu12==2.26.2 [pip3] nvidia-nvjitlink-cu12==12.6.85 [pip3] nvidia-nvtx-cu12==12.6.77 [pip3] optree==0.16.0 [pip3] pytorch-lightning==1.9.5 [pip3] pytorch-msssim==1.0.0 [pip3] torch==2.7.1 [pip3] torch_cluster==1.6.3 [pip3] torch_scatter==2.1.2 [pip3] torchdata==0.11.0 [pip3] torchdiffeq==0.2.5 [pip3] torchmetrics==1.7.4 [pip3] torchsde==0.2.6 [pip3] torchsparse==2.1.0 [pip3] torchvision==0.22.0+9eb57cd [pip3] triton==3.3.1 [conda] numpy 2.2.0 pypi_0 pypi [conda] nvidia-cublas-cu12 12.6.4.1 pypi_0 pypi [conda] nvidia-cuda-cupti-cu12 12.6.80 pypi_0 pypi [conda] nvidia-cuda-nvrtc-cu12 12.6.77 pypi_0 pypi [conda] nvidia-cuda-runtime-cu12 12.6.77 pypi_0 pypi [conda] nvidia-cudnn-cu12 9.5.1.17 pypi_0 pypi [conda] nvidia-cufft-cu12 11.3.0.4 pypi_0 pypi [conda] nvidia-curand-cu12 10.3.7.77 pypi_0 pypi [conda] nvidia-cusolver-cu12 11.7.1.2 pypi_0 pypi [conda] nvidia-cusparse-cu12 12.5.4.2 pypi_0 pypi [conda] nvidia-cusparselt-cu12 0.6.3 pypi_0 pypi [conda] nvidia-nccl-cu12 2.26.2 pypi_0 pypi [conda] nvidia-nvjitlink-cu12 12.6.85 pypi_0 pypi [conda] nvidia-nvtx-cu12 12.6.77 pypi_0 pypi [conda] optree 0.16.0 pypi_0 pypi [conda] pytorch-lightning 1.9.5 pypi_0 pypi [conda] pytorch-msssim 1.0.0 pypi_0 pypi [conda] tbb 2022.1.0 h4ce085d_0 conda-forge [conda] torch 2.7.1 pypi_0 pypi [conda] torch-cluster 1.6.3 pypi_0 pypi [conda] torch-scatter 2.1.2 pypi_0 pypi [conda] torchdata 0.11.0 pypi_0 pypi [conda] torchdiffeq 0.2.5 pypi_0 pypi [conda] torchmetrics 1.7.4 pypi_0 pypi [conda] torchsde 0.2.6 pypi_0 pypi [conda] torchsparse 2.1.0 pypi_0 pypi [conda] torchvision 0.22.0+9eb57cd pypi_0 pypi [conda] triton 3.3.1 pypi_0 pypi

cc @ptrblck @msaroufim @eqy @jerryzh168 @tinglvv @nWEIdia

extent analysis

TL;DR

The issue can be worked around by using the non-vectorized path for torch.layer_norm or by splitting the input into smaller chunks to avoid the 32-bit integer overflow.

Guidance

  • The suspected cause is a 32-bit integer overflow in the vectorized forward kernel of torch.layer_norm when the total number of elements exceeds 2**32.
  • To verify the issue, run the provided minimal reproduction code and check the output for incorrect values after the 2**32 boundary.
  • To mitigate the issue, consider using the non-vectorized path for torch.layer_norm or split the input into smaller chunks to avoid the overflow.
  • Alternatively, you can try updating to a newer version of PyTorch, if available, to see if the issue has been fixed.

Example

# Split the input into smaller chunks
chunk_size = 1024
for i in range(0, M, chunk_size):
    chunk = feats[i:i+chunk_size]
    out_chunk = torch.layer_norm(chunk, normalized_shape=(C,), weight=None, bias=None, eps=1e-5)
    # Process the chunk

Notes

  • The issue appears to be specific to the vectorized path of torch.layer_norm and may not affect the non-vectorized path.
  • The workaround of splitting the input into smaller chunks may impact performance.

Recommendation

Apply workaround: split the input into smaller chunks to avoid the 32-bit integer overflow, as this is a reliable and straightforward solution to the issue.

Vote matrix · Quick signals

Works
Did the solution work? Tap to confirm.
Easy Fix
Was it a quick fix?
Time Saver
Did it save you time?
Blocking
Was it severely blocking?
Common Issue
Are others likely hitting this too?
Flaky / Intermittent
Is it intermittent?
Verified / Reproducible
Can you reproduce it reliably?
Loading…

FAQ

Expected behavior

torch.layer_norm should produce correct output regardless of whether the flattened number of elements exceeds 2**32, as long as the tensor size is otherwise valid and supported.

Still need to ship something?

×6

Another batch ranked right after the header list — different links, same matching logic.

Back to top recommendations

TRENDING