Preserve FP32 LayerNorm statistics for backward (#1498)

## Summary

LayerNorm computes its forward statistics in FP32 but saves them in the
input dtype in the Triton and legacy cuTile implementations. BF16/FP16
rounding then changes the statistics used by backward, producing
incorrect input and affine gradients for near-constant inputs.

For a BF16 `[1, 1024]` input with halves `1` and `1.0078125`, unit
weight, zero bias, `eps=1e-5`, and an all-ones upstream gradient, Triton
returns `dX=-240` in the second half instead of zero. After this change,
`dX` is zero and `dWeight` matches an independent FP64 reference.

Allocate both Mean and RSTD in FP32 and store the FP32 values directly
in both affected implementations. This adds four bytes of saved
statistics per row for BF16/FP16 inputs. The dispatcher cuTile, CuTe
DSL, and Ascend implementations already preserve FP32 statistics. The
diff contains only the two production files.

## Testing Done

- Hardware: NVIDIA RTX A1000 Laptop GPU; PyTorch 2.13.0+cu130; Triton
3.7.1.
- Existing LayerNorm transformer and op suites: 60 passed.
- External regression harness: 12 passed across BF16/FP16/FP32, shapes
`[1, 1024]` and `[3, 1000]`, and separate mean/rstd rounding cases. The
original implementation fails all eight low-precision cases. The
reference differentiates an independent FP64 expression from the same
quantized inputs. This harness is not included in the diff.
- Legacy cuTile: FP32 allocation/launch checks passed for three dtypes
and two shapes; all 18 forward, input-only backward, and combined
backward variants compiled offline for `sm_100`.
- cuTile numerical execution was not available on this `sm_86` GPU. The
installed compiler also does not support `sm_90`, so Hopper compilation
was not validated.
- [x] `make checkstyle` with the repository's Ruff 0.14.11 version.
- [ ] Full `make test` (scoped LayerNorm suites run instead).
- [ ] `make test-convergence` (not run).

src/liger_kernel/ops/cutile/ops/layer_norm.py

@@ -77,9 +77,9 @@ def _layer_norm_fwd_kernel_ct(
     var = total_sum_sq / n_cols - mean * mean
     rstd = ct.rsqrt(var + eps)  # scalar
 
-    # Cache mean and rstd for backward
-    ct.scatter(mean_out, row_idx, ct.astype(mean, mean_out.dtype))
-    ct.scatter(rstd_out, row_idx, ct.astype(rstd, rstd_out.dtype))
+    # Preserve the fp32 statistics used by forward for the backward pass.
+    ct.scatter(mean_out, row_idx, mean)
+    ct.scatter(rstd_out, row_idx, rstd)
 
     # ---- Pass 2: Y = (X - mean) * rstd * W + B ----
     for ci in range(n_chunks):
@@ -216,8 +216,8 @@ def _layer_norm_forward_ct(X, W, B, eps):
     aligned = (n_cols & (n_cols - 1)) == 0  # True when n_cols is a power of 2
 
     Y = torch.empty_like(X2d)
-    Mean = torch.empty(n_rows, dtype=X.dtype, device=X.device)
-    RSTD = torch.empty(n_rows, dtype=X.dtype, device=X.device)
+    Mean = torch.empty(n_rows, dtype=torch.float32, device=X.device)
+    RSTD = torch.empty(n_rows, dtype=torch.float32, device=X.device)
 
     grid = (n_rows, 1, 1)
     ct.launch(

src/liger_kernel/ops/layer_norm.py

@@ -74,9 +74,9 @@ def _layer_norm_forward_kernel(
     var = tl.sum(X_centered_masked * X_centered_masked, axis=0) / n_cols
     rstd = rsqrt(var + eps.to(tl.float32))
 
-    # Store statistics (convert back to original dtype only once)
-    tl.store(row_Mean_ptr, mean.to(X_row.dtype))
-    tl.store(row_RSTD_ptr, rstd.to(X_row.dtype))
+    # Preserve the fp32 statistics used by forward for the backward pass.
+    tl.store(row_Mean_ptr, mean)
+    tl.store(row_RSTD_ptr, rstd)
 
     # Fused normalization and affine transformation
     # Y = (X - mean) * rstd * W + B = X_centered * rstd * W + B
@@ -187,8 +187,8 @@ def layer_norm_forward(X, W, B, eps):
 
     # Allocate output tensors
     Y = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
-    Mean = torch.empty(n_rows, dtype=X.dtype, device=X.device)
-    RSTD = torch.empty(n_rows, dtype=X.dtype, device=X.device)
+    Mean = torch.empty(n_rows, dtype=torch.float32, device=X.device)
+    RSTD = torch.empty(n_rows, dtype=torch.float32, device=X.device)
 
     # Validate input dimensions
     if X.shape[1] != W.shape[0]: