RMSNorm epsilon refine in the example

[pengxin99]

According to the RMSNorm define in Pytorch, the epsilon operation is in the rsqrt()

Summary by CodeRabbit

• Bug Fixes

  • Improved numerical stability in normalization calculations by adjusting epsilon placement within the normalization formula.

• Tests

  • Updated normalization tests to reflect improved stability calculations.

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

Walkthrough

The pull request adjusts epsilon placement in RMS normalization calculations across implementation and test files. The constant 1e-12 is moved from outside to inside the rsqrt function argument in both examples/norm/rms_norm.py and examples/norm/test_rms_norm.py, altering the numerical stability characteristics of the normalization factor.

Changes

Cohort / File(s)	Summary
RMS normalization implementation examples/norm/rms_norm.py	Modified epsilon placement in rsqrt calculations: changed from rsqrt(A_powsum[i] / N) + 1e-12 to rsqrt(A_powsum[i] / N + 1e-12) in both split-k and standard variants
RMS normalization tests examples/norm/test_rms_norm.py	Updated test computations to match implementation: changed epsilon placement in three locations using the same pattern as implementation (rsqrt(A_powsum[i] / N + 1e-12))

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

• Verify the mathematical correctness of moving epsilon inside rsqrt (affects numerical stability and final output values)
• Confirm the change applies consistently across both split-k and standard normalization paths
• Check that test assertions remain valid with the new calculation method

Poem

🐰 A constant moved within the square root's embrace,
Where epsilon finds its proper place,
From outer addition to inner care,
Normalization blooms with balanced flair! ✨

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'RMSNorm epsilon refine in the example' clearly describes the main change: refining epsilon handling in RMSNorm implementation within examples, matching the actual modifications to rsqrt argument placement.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 6882bd5 and 6a8cb53.

📒 Files selected for processing (2)

• examples/norm/rms_norm.py (2 hunks)
• examples/norm/test_rms_norm.py (2 hunks)

🧰 Additional context used

🧬 Code graph analysis (1)

examples/norm/test_rms_norm.py (1)

tilelang/language/tir/op.py (1)

• rsqrt (2399-2412)

🔇 Additional comments (3)

examples/norm/rms_norm.py (2)

24-24: Correct epsilon placement aligns with PyTorch RMSNorm standard.

The change from rsqrt(A_powsum[i] / N) + 1e-12 to rsqrt(A_powsum[i] / N + 1e-12) is correct and fixes a numerical stability issue. The old approach added epsilon after the reciprocal square root, which does not prevent division by zero. The new approach adds epsilon inside, ensuring numerical stability when the variance is near zero, consistent with PyTorch's RMSNorm implementation and the reference program at line 63.

Also applies to: 54-54

---

62-63: Reference implementation was already correct.

The reference program already had the correct epsilon placement inside rsqrt. The changes at lines 24 and 54 bring the tilelang implementation into consistency with this reference, which should improve test accuracy.

examples/norm/test_rms_norm.py (1)

25-25: Test implementations now match the reference and PyTorch standard.

The epsilon placement corrections in the test functions mirror the changes in examples/norm/rms_norm.py, ensuring consistency. These implementations now correctly match the reference program at line 63, which already had the proper epsilon placement inside rsqrt.

Also applies to: 54-54

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[LeiWang1999]

LGTM, Thanks, also cc @chengyupku

[LeiWang1999]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. You're on a roll.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".