Repository navigation
[Bug] Fused adaptive LayerNorm returns all-NaN rows for some finite BF16 inputs #20016
Description
Summary. launchFusedAdaptiveLayerNormKernel (cpp/tensorrt_llm/kernels/fusedAdaptiveLayerNormKernel.cu) computes the variance in FP32 as E[x²] - mean². For a row whose mean is large relative to its spread, cancellation can make variance + eps negative, and the whole output row becomes NaN for finite BF16 input. We propose computing the mean first, then the variance from the cached values centered on it.
Reproduce
Width 5120, BF16: two values 15.6875 followed by 5,118 values 15.75. The exact variance is about 1.53e-6, so the output should be finite. Built from source (a1f562a; main e2852d9 has the same kernel) and called through its launcher on an A100, the kernel returns NaN in all 5,120 columns, in plain, affine and modulation modes, for eps 1e-6 and 1e-5. Control rows (constant 1, and balanced ±1) are correct.
Through the op, in builds that include it (TensorRT-LLM 1.2.1 does not yet, so we ran the launcher directly):
import torch
import tensorrt_llm # registers torch.ops.trtllm
x = torch.full((1, 5120), 15.75, dtype=torch.bfloat16, device="cuda")
x[:, :2] = 15.6875
y = torch.ops.trtllm.fused_adaptive_layernorm(x, None, None, None, None, 0, 1e-6)
x64 = x.cpu().double()
mean = x64.mean(-1, keepdim=True)
variance = ((x64 - mean) ** 2).mean(-1, keepdim=True)
eps = torch.tensor(1e-6, dtype=torch.float32).double()
reference = (x64 - mean) / torch.sqrt(variance + eps)
print("output_nan_count", torch.isnan(y).sum().item())
print("reference_is_finite", torch.isfinite(reference).all().item())Expected output_nan_count 0.
Suggested fix
--- a/cpp/tensorrt_llm/kernels/fusedAdaptiveLayerNormKernel.cu
+++ b/cpp/tensorrt_llm/kernels/fusedAdaptiveLayerNormKernel.cu
@@ -286,54 +286,72 @@
}
}
- // Phase 1: warp-reduce sum/sum-of-squares → mean and rstd (LayerNorm, not RMSNorm).
+ // Phase 1: mean, followed by centered variance from the cached values.
float localSum = 0.0f;
- float localSqSum = 0.0f;
#pragma unroll
for (int i = 0; i < ELEMS_PER_THREAD; ++i)
{
localSum += xVals[i];
- localSqSum += xVals[i] * xVals[i];
- }
-
- // Warp-level reduction.
+ }
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2)
{
localSum += __shfl_xor_sync(0xffffffff, localSum, offset);
- localSqSum += __shfl_xor_sync(0xffffffff, localSqSum, offset);
- }
-
+ }
if (laneId == 0)
{
warpSums[warpId] = localSum;
- warpSqSums[warpId] = localSqSum;
}
__syncthreads();
-
- // Cross-warp reduction in warp 0.
if (warpId == 0)
{
float s = (laneId < NUM_WARPS) ? warpSums[laneId] : 0.0f;
+#pragma unroll
+ for (int offset = 16; offset > 0; offset /= 2)
+ {
+ s += __shfl_xor_sync(0xffffffff, s, offset);
+ }
+ if (laneId == 0)
+ {
+ meanRstd[0] = s * (1.0f / static_cast<float>(D));
+ }
+ }
+ __syncthreads();
+
+ float const mean = meanRstd[0];
+ float localCenteredSqSum = 0.0f;
+#pragma unroll
+ for (int i = 0; i < ELEMS_PER_THREAD; ++i)
+ {
+ float const centered = xVals[i] - mean;
+ localCenteredSqSum += centered * centered;
+ }
+#pragma unroll
+ for (int offset = 16; offset > 0; offset /= 2)
+ {
+ localCenteredSqSum += __shfl_xor_sync(0xffffffff, localCenteredSqSum, offset);
+ }
+ if (laneId == 0)
+ {
+ warpSqSums[warpId] = localCenteredSqSum;
+ }
+ __syncthreads();
+ if (warpId == 0)
+ {
float s2 = (laneId < NUM_WARPS) ? warpSqSums[laneId] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2)
{
- s += __shfl_xor_sync(0xffffffff, s, offset);
s2 += __shfl_xor_sync(0xffffffff, s2, offset);
}
if (laneId == 0)
{
- float const invD = 1.0f / static_cast<float>(D);
- float const mean = s * invD;
- float const var = s2 * invD - mean * mean;
- meanRstd[0] = mean;
+ float const var = s2 * (1.0f / static_cast<float>(D));
meanRstd[1] = rsqrtf(var + p.eps);
}
}
__syncthreads();
- float const mean = meanRstd[0];
float const rstd = meanRstd[1];
// Pre-read sf_scale scalar (HAS_QUANT only).With this change, all 36 test rows are finite and match an FP64 reference (|error| <= 0.002 + 0.01*|reference|), and the control outputs are unchanged. It adds one reduction and two barriers; performance was not measured.
Scope
- Reach: the public op accepts any contiguous BF16
[M, 5120]input, and Wan's dense normalization calls this kernel at width 5120. We have not measured how often real activations reach this regime. - Environment: NVIDIA A100-SXM4-80GB (SM80), CUDA 12.9.86, GCC 11.5.
Activity
[NVIDIA#20016][test] Cover affine, modulation and NVFP4 in the large-…