NVIDIA / NVIDIA/SOL-ExecBench

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

Open
#21 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
296
Forks
34
PR merge metrics
No merged PRs in 30d

Description

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.

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.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with examples/triton/rmsnorm/definition.json and compare its fixed epsilon with the launch in examples/triton/rmsnorm/kernel.py. Update the launch and the embedded source in examples/triton/rmsnorm/solution_triton.json, then add a small-magnitude BF16 correctness case; done means the bundled example matches the reference for valid BF16 inputs.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, testing-qa
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
74/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.