Triton RMSNorm example uses epsilon 1e-6 but definition requires 1e-5

[anguyen8]

Summary

The Triton RMSNorm example does not implement its own problem definition.

Observed mismatch for x = BF16(1e-3) and weight = 1: expected 0.30078125 vs. candidate output 0.70703125.

• definition.json fixes EPS = 1e-5 in the reference and says “Epsilon is fixed at 1e-5.”
• kernel.py launches _rmsnorm_fwd_kernel with 1e-6.
• The embedded kernel.py content in solution_triton.json also contains 1e-6.

This is present on current main at commit a9fa0804c793d438e70850c33fe34426e66d53dd.

Minimal reproducer

Use a valid official-shape BF16 row with small finite values and unit weights:

import torch

x = torch.full((1, 4096), 1e-3, dtype=torch.bfloat16, device="cuda")
weight = torch.ones(4096, dtype=torch.bfloat16, device="cuda")

def rmsnorm(x, eps):
    x32 = x.float()
    return (
        x32 * torch.rsqrt(x32.square().mean(-1, keepdim=True) + eps)
        * weight.float()
    ).bfloat16()

print(x[0, 0].item())
print(rmsnorm(x, 1e-5)[0, 0].item())
print(rmsnorm(x, 1e-6)[0, 0].item())

Input:                                      0.00099945068359375
Expected (reference, epsilon = 1e-5):       0.30078125
Observed (Triton candidate, epsilon = 1e-6): 0.70703125

So the concise comparison is: expected 0.30078125 vs. observed 0.70703125.

A randomized (randn * 1e-3).bfloat16() input with batch size 7 produced 26,631 mismatching elements out of 28,672.

This is not reduction-order noise: for a constant row, the output is analytically x / sqrt(x² + eps), and the observed ratio agrees with changing epsilon from 1e-5 to 1e-6.

Expected behavior

The bundled example solution should use the epsilon required by rmsnorm_h4096 for all valid BF16 inputs.

Suggested fix

1. Change the launch argument in examples/triton/rmsnorm/kernel.py from 1e-6 to 1e-5.
2. Make the same change in the embedded source inside examples/triton/rmsnorm/solution_triton.json.
3. Add a small-magnitude BF16 correctness case; unit-scale inputs make epsilon comparatively negligible and can hide this mismatch.