【Triton Copilot】Fix rms_norm backward precision issue

[factnn]

Description

This PR fixes the numerical precision issue in rms_norm backward implementation reported in #1068.

Development Tool:

• This operator was developed with Triton-Copilot, an AI-powered tool for Triton kernel development.

Problem

The backward pass of rms_norm had a precision issue where:

• x.grad contained NaN/Inf values
• Gradient values were significantly different from native PyTorch (e.g., 12.318 vs normal range -0.8 to 0.2)
• Relative errors exceeded 1000%

Root Cause

The issue was caused by not ensuring dy is contiguous before passing it to the kernel. When dy is not contiguous, the stride parameters passed to the kernel are incorrect, causing the kernel to read from wrong memory locations.

Solution

1. Key fix: Added dy = dy.contiguous() in the backward function to ensure correct stride parameters
2. Auxiliary fix: Changed partial_buffer initialization from torch.empty() to torch.zeros() for better robustness

Changes

• Added dy = dy.contiguous() in RmsNorm.backward() (line 186)
• Changed partial_buffer = torch.empty(...) to partial_buffer = torch.zeros(...) (line 200)
• Added comments explaining masked row handling in kernel

Verification

• ✅ All pytest tests pass (18 rms_norm related tests)
• ✅ Gradient values match native PyTorch implementation
• ✅ Maximum relative error: 0.11% for x.grad, 0.05% for w.grad
• ✅ No NaN/Inf values
• ✅ Verified with reviewer's reproduction script and test data

Test Results

Using the same test data from issue #1068:

• Before fix: x.grad had values like 12.318 (abnormal), NaN/Inf present
• After fix: x.grad matches native PyTorch exactly, max relative error < 0.2%

Fixes #1068