avoid pointer mutation in add_rms_norm kernel (#1008)

## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Rewrite fused_add_rms_norm kernel to use explicit channel offsets
instead of mutating X/Y base pointers inside loops.

This improves Triton compiler optimization opportunities, enables more
predictable memory access patterns, and avoids loop-carried pointer
dependencies.

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<img width="1800" height="382" alt="image"
src="https://github.com/user-attachments/assets/b361d41e-1379-4835-8acb-b0234d5af22b"
/>


- Hardware Type: Ascend NPU 910B4
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

Author: TianHao324

src/liger_kernel/ops/fused_add_rms_norm.py

@@ -162,23 +162,21 @@ def _fused_add_rms_norm_backward_kernel(
 
     dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
 
-    dY_ptr += row_start * dY_row_stride
-    dX_ptr += row_start * dX_row_stride
-    if has_dS_out:
-        dS_out_ptr += row_start * dS_out_row_stride
-
-    X_ptr += row_start * X_row_stride
-    RSTD_ptr += row_start
-
     W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0.0)
     W_row = W_row + offset
 
-    for _ in range(row_start, row_end):
-        dY_row = tl.load(dY_ptr + col_offsets, mask=mask, other=0.0)
-        X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0)
+    for row_idx in range(row_start, row_end):
+        dy_base = dY_ptr + row_idx * dY_row_stride
+        dx_base = dX_ptr + row_idx * dX_row_stride
+
+        x_base = X_ptr + row_idx * X_row_stride
+        rstd_base = RSTD_ptr + row_idx * RSTD_row_stride
+
+        dY_row = tl.load(dy_base + col_offsets, mask=mask, other=0.0)
+        X_row = tl.load(x_base + col_offsets, mask=mask, other=0.0)
 
         # Get cached rms
-        rstd_row = tl.load(RSTD_ptr)
+        rstd_row = tl.load(rstd_base)
 
         X_row = X_row.to(tl.float32)
 
@@ -195,11 +193,11 @@ def _fused_add_rms_norm_backward_kernel(
         dX_row = rstd_row * m
 
         if has_dS_out:
-            dS_out_row = tl.load(dS_out_ptr + col_offsets, mask=mask, other=0.0)
+            ds_base = dS_out_ptr + row_idx * dS_out_row_stride
+            dS_out_row = tl.load(ds_base + col_offsets, mask=mask, other=0.0)
             dX_row += (rstd_row) * (
                 -(1 / n_cols) * rstd_row * rstd_row * tl.sum(m * X_row, axis=0) * X_row
             ) + dS_out_row
-            dS_out_ptr += dS_out_row_stride
         else:
             dX_row += (rstd_row) * (-(1 / n_cols) * rstd_row * rstd_row * tl.sum(m * X_row, axis=0) * X_row)
 
@@ -210,12 +208,7 @@ def _fused_add_rms_norm_backward_kernel(
             # here X_row is already in fp32 (see previous if block)
             dW_row += dY_row * (X_row * rstd_row)
 
-        tl.store(dX_ptr + col_offsets, dX_row.to(X_dtype), mask=mask)
-
-        dY_ptr += dY_row_stride
-        dX_ptr += dX_row_stride
-        X_ptr += X_row_stride
-        RSTD_ptr += RSTD_row_stride
+        tl.store(dx_base + col_offsets, dX_row.to(X_dtype), mask=mask)
 
     tl.store(dW_ptr + row_block_id * dW_row_stride + col_offsets, dW_row, mask=mask)