[BUG] Inference layer_norm / rms_norm kernels read and write past the last row when rows is not a multiple of the rows per block (hidden size <= 128) #8828

Describe the bug

For hidden sizes <= 128, the inference norm launchers use a "sub-block" schedule that packs several rows into one thread block, but the kernels have no row bound. When the number of rows is larger than the rows per block and not a multiple of it, the last block reads its input and writes its outputs for rows past the end of the tensors.

Affected launchers (master 5f266e26, csrc/transformer/inference/csrc/):

• launch_fused_ln (layer_norm.cu:128), kernel fused_ln (layer_norm.cu:32)
• launch_fused_residual_ln (layer_norm.cu:344) and launch_fused_residual_ln_store_pre_ln_res (layer_norm.cu:417), kernel fused_residual_ln (layer_norm.cu:224)
• launch_rms_norm (rms_norm.cu:183), kernels rms_norm (rms_norm.cu:20) and pre_rms_norm (rms_norm.cu:85)

These are reached through the transformer_inference op bindings layer_norm, _layer_norm_residual, layer_norm_residual_store_pre_ln_res, rms_norm and pre_rms_norm (pt_binding.cpp:2053-2060), and launch_fused_residual_ln is also used by the injected MLP path.

Where it comes from

In each launcher (layer_norm.cu:150-159, :368-377, :442-451, rms_norm.cu:203-212):

one_step_threads     = next_pow2(ceil(elems_per_row / T_per_load))     // T_per_load = 4 for fp32, 8 for fp16/bf16
threadsPerGroup      = min(one_step_threads, 256)
groups_per_block_max = is_subblock_schedule ? 256 / threadsPerGroup : 1  // rows per block
groups_per_block     = min(rows, groups_per_block_max)
grid                 = ceil(rows / groups_per_block),  block = (threadsPerGroup, groups_per_block)

and in each kernel (layer_norm.cu:45, :240, rms_norm.cu:28, :99):

block_offset = blockIdx.x * (256 / threadsPerGroup) * elems_per_row + threadIdx.y * elems_per_row

Loads and stores are predicated only on the column (thread_offset + i * stride < elems_per_row), never on the row. So with R = 256 / threadsPerGroup rows per block, when rows > R and rows % R != 0, the last block processes R - rows % R rows beyond the end: it reads them from the inputs and writes them to the outputs (and to res_output for the residual-storing variants).

Which shapes are affected (threadsPerGroup from the formula above; all of these use the sub-block schedule, elems_per_row <= 128):

(On master, fp32 hidden 65-128 launches no kernel at all; that is #8826, fixed by #8827, which uses one row per block for that case.) The 16- and 8-thread rows of the table were checked with compute-sanitizer below; the 4/2/1-thread rows follow from the same code and were not run.

To Reproduce

Standalone script (direct calls to the op bindings, fp32, hidden 64, rows = 17): deepspeed__16thread_row_overrun.py (below). Run it under compute-sanitizer:

PYTORCH_NO_CUDA_MEMORY_CACHING=1 compute-sanitizer --tool memcheck --padding 65536 \
    --print-limit 4 --show-backtrace no python deepspeed__16thread_row_overrun.py

(PYTORCH_NO_CUDA_MEMORY_CACHING=1 gives every tensor its own allocation, and --padding makes overruns into neighbouring memory visible.)

Output without the sanitizer (master, 3 runs identical). The plain layer_norm output looks fine, but the residual-storing variant's res_output is overwritten by the overrun of its norm_output, which is allocated right before it:

rows=17 hidden=64 torch.float32: 16 threads/row, 16 rows/block, grid=2 -> last block covers rows 17..31 past the end
layer_norm               max|err| vs torch: 9.5367431640625e-07
residual store: res_output max|err| vs x + bias + residual: 17.698978424072266 (non-zero = overwritten by the overrun)
residual store: norm_output max|err| vs torch: 4.76837158203125e-07

Under compute-sanitizer the first kernel already faults (ERROR SUMMARY: 244 errors, then CUDA error: unspecified launch failure):

========= Invalid __global__ read of size 16 bytes
=========     at void fused_ln<float, (int)1, (int)16, (int)256>(T1 *, const T1 *, const T1 *, const T1 *, float, int)+0xf0
=========     by thread (0,4,0) in block (1,0,0)
=========     Address 0x7dfcce201400 is out of bounds
=========     and is 769 bytes after the nearest allocation at 0x7dfcce200000 of size 4,352 bytes

# DeepSpeed inference layer_norm / rms_norm: the sub-block launch (hidden size <= 128) can read and
# write past the last row when `rows` is not a multiple of the rows-per-block.
#
# fp32, hidden 64 -> 16 threads per row, 256 / 16 = 16 rows per block. With rows = 17 the launcher uses
# grid = 2 blocks of 16 rows, and the kernels have no row bound, so block 1 touches rows 17..31.
#
# Run it under compute-sanitizer to see the out-of-bounds accesses (the caching allocator is disabled so
# every tensor is its own allocation, and padding makes overruns into neighbouring memory visible):
#
#   PYTORCH_NO_CUDA_MEMORY_CACHING=1 compute-sanitizer --tool memcheck --padding 65536 \
#       --print-limit 4 --show-backtrace no python deepspeed__16thread_row_overrun.py
#
# Without the sanitizer, the store-pre-LN-residual variant usually shows the effect directly: its two
# outputs are allocated back to back, so the overrun of the normalized output lands in the residual output.
import torch

import deepspeed
from deepspeed.ops.op_builder import InferenceBuilder

print("deepspeed", deepspeed.__version__, "torch", torch.__version__, "cuda", torch.version.cuda,
      torch.cuda.get_device_name(0))
mod = InferenceBuilder().load()

torch.manual_seed(0)
rows, H, dtype = 17, 64, torch.float32
x = torch.randn(1, rows, H, device="cuda", dtype=dtype)
res = torch.randn(1, rows, H, device="cuda", dtype=dtype)
bias = torch.randn(H, device="cuda", dtype=dtype)
gamma = torch.randn(H, device="cuda", dtype=dtype)
beta = torch.rand(H, device="cuda", dtype=dtype)
eps = 1e-5

rows_per_block = 256 // 16
print(f"rows={rows} hidden={H} {dtype}: 16 threads/row, {rows_per_block} rows/block, "
      f"grid={-(-rows // rows_per_block)} -> last block covers rows {rows}..{-(-rows // rows_per_block) * rows_per_block - 1}"
      f" past the end")

y = mod.layer_norm(x, gamma, beta, eps)
y_ref = torch.nn.functional.layer_norm(x, (H, ), gamma, beta, eps)
print("layer_norm               max|err| vs torch:", (y - y_ref).abs().max().item())

norm_out, res_out = mod.layer_norm_residual_store_pre_ln_res(x, bias, res, gamma, beta, eps)
torch.cuda.synchronize()
pre = x + bias + res
print("residual store: res_output max|err| vs x + bias + residual:", (res_out - pre).abs().max().item(),
      "(non-zero = overwritten by the overrun)")
print("residual store: norm_output max|err| vs torch:",
      (norm_out - torch.nn.functional.layer_norm(pre, (H, ), gamma, beta, eps)).abs().max().item())

memcheck results on master (one process per case, same flags as above; the number is the count of invalid global accesses reported for that call, 0 = clean):

Rows tested: 9, 13, 16, 17, 32, 33 (all shapes (1, rows, H)). Counts are for the first faulting launch (the CUDA context is lost after it). All faults are in the <T, 1, 16, 256> instances (<float, 1, 8, 256> for fp32 H=32).

Expected behavior

No accesses outside the input/output tensors for any number of rows.

Possible fix

Either:

1. launch one row per block in the sub-block branches, as the > 128 path already does and as Fix inference layer_norm/rms_norm launching no kernel for fp32 hidden sizes 65-128 #8827 does for the 32-thread case (grid(rows), block(threadsPerGroup), maxThreads = threadsPerGroup in the template so the kernel's rows-per-block factor becomes 1), or
2. pass rows to the kernels and mask the loads/stores with row < rows (keeping every thread in the reductions).

Option 1 is local to the launchers but launches more, smaller blocks (e.g. 16 threads per block for hidden 64 in fp32), which may be slower for many short rows; option 2 keeps the current launch shape but touches all four kernels. I haven't measured the performance of either. Happy to send a PR with whichever you prefer.

System info

• GPU: 1x NVIDIA GeForce RTX 5080 Laptop GPU (sm_120), driver 595.99.02
• CUDA 12.8, compute-sanitizer 2025.1.0
• torch 2.11.0+cu128, Python 3.12.3, Ubuntu 24.04
• DeepSpeed master 5f266e26, transformer_inference op JIT-built with TORCH_CUDA_ARCH_LIST=12.0

Related: #8826 / #8827 (fp32 hidden 65-128 launched no kernel; #8827's 32-thread launch uses one row per block to avoid this overrun), #8742 / #8743 (same launchers, rows == 0).