[XPU] tl.store 在向量计算之后写入标量时，实际写入的是过期的寄存器值

## 问题描述

在 Triton kernel 中，当一个标量通过 `tl.sum` +
算术运算计算得出后，如果中间穿插了向量操作，再对该标量执行
`tl.store`，实际写入的是算术运算之前的旧值（寄存器被复用），而非正确的结果。

具体表现：`rrms = 1 / tl.sqrt(var + eps)` 计算正确（前向输出使用 rrms 的结果是对的），但
`tl.store(INV_RMS + pid, rrms)` 写入的却是 `var` 的值。

## 环境

- FlagTree 版本：0.6.0+xpu.git0aa159bb
- 后端：KunlunXin XPU

## 复现代码

需要在 FlagGems KunlunXin 后端的**原始**（未修复）rms_norm 实现上运行：

```python
import numpy as np
import torch
import flag_gems

def main():
device = flag_gems.device
print(f"Device: {device}")

test_configs = [
(1, 2, torch.float16),
(1, 128, torch.float16),
(4, 1024, torch.float16),
(16, 4096, torch.float32),
]

for M, N, dtype in test_configs:
print(f"\nTest: shape=({M}, {N}), dtype={dtype}")

np.random.seed(0)
np_inp = np.random.uniform(-0.1, 0.1, (M, N)).astype(np.float32)
np_weight = np.random.uniform(-0.1, 0.1, (N,)).astype(np.float32)
np_grad = np.random.uniform(-0.01, 0.01, (M, N)).astype(np.float32)

inp = torch.tensor(np_inp, dtype=dtype, device=device, requires_grad=True)
weight = torch.tensor(np_weight, dtype=dtype, device=device, requires_grad=True)
out = flag_gems.rms_norm(inp, [N], weight=weight, eps=1e-5)

# 检查前向保存的 inv_rms
saved_inv_rms = out.grad_fn.saved_tensors[1]

# 期望的 inv_rms
x_f32 = inp.detach().float()
var = (x_f32**2).mean(dim=-1)
expected_inv_rms = 1.0 / torch.sqrt(var + 1e-5)

inv_rms_err = (saved_inv_rms.cpu().float() -
expected_inv_rms.cpu()).abs().max().item()
var_match = (saved_inv_rms.cpu().float() - var.cpu()).abs().max().item()

print(f" saved_inv_rms[0]: {saved_inv_rms[0].item():.8f}")
print(f" expected_inv_rms[0]: {expected_inv_rms[0].item():.8f}")
print(f" var[0]: {var[0].item():.8f}")
print(f" Error vs expected: {inv_rms_err:.6e}")
print(f" Error vs var: {var_match:.6e}")

if inv_rms_err > 0.01 and var_match < 1e-4:
print(" >>> BUG: forward stored var instead of inv_rms <<<")

# 检查反向梯度
grad_out = torch.tensor(np_grad, dtype=dtype, device=device)
dx, dw = torch.autograd.grad(out, (inp, weight), grad_out)

# PyTorch 参考实现
inp_ref = torch.tensor(np_inp, dtype=dtype, device=device, requires_grad=True)
w_ref = torch.tensor(np_weight, dtype=dtype, device=device, requires_grad=True)
x_f32_ref = inp_ref.float()
variance = x_f32_ref.pow(2).mean(-1, keepdim=True)
x_normed = (x_f32_ref * torch.rsqrt(variance + 1e-5)).to(dtype)
ref_out = x_normed * w_ref
dx_ref, dw_ref = torch.autograd.grad(ref_out, (inp_ref, w_ref), grad_out)

dx_err = (dx.float().cpu() - dx_ref.float().cpu()).abs().max().item()
dx_ref_max = dx_ref.float().cpu().abs().max().item()
print(f" dx max abs error: {dx_err:.6e}")
print(f" dx_ref max abs: {dx_ref_max:.6e}")

if dx_ref_max > 1e-6 and dx_err / dx_ref_max > 0.1:
print(" >>> BUG IMPACT: backward gradient is severely wrong <<<")

if __name__ == "__main__":
main()
```

```python
@triton.jit
def rms_norm_kernel(...):
...
x = tl.load(X + cols * x_stride_c, mask, other=0.0).to(tl.float32)

var = tl.sum(x * x, axis=0) / N # 标量，均方值
rrms = 1 / tl.sqrt(var + eps) # 标量，由 var 计算得到

# --- 中间执行了向量 load 和计算 ---
w = tl.load(W + tl.arange(0, BLOCK_SIZE), mask=mask, other=0.0)
y = (x * rrms).to(Y.dtype.element_ty) * w
tl.store(Y + cols * y_stride_c, y, mask=mask)

tl.store(INV_RMS + pid, rrms) # BUG：实际写入的是 var，不是 rrms
```

## 复现输出

```
Test: shape=(1, 2), dtype=torch.float16
saved_inv_rms[0]: 0.00097346 ← 实际存储的值（等于 var）
expected_inv_rms[0]: 31.88750839 ← 期望存储的值（1/sqrt(var+eps)）
var[0]: 0.00097346
Error vs expected: 3.188654e+01
Error vs var: 0.000000e+00
>>> BUG: forward stored var instead of inv_rms <<<
dx max abs error: 1.131952e-03
dx_ref max abs: 1.132011e-03
>>> BUG IMPACT: backward gradient is severely wrong <<<
```

相关PR：https://github.com/flagos-ai/FlagGems/pull/5771

Start with the original FlagGems KunlunXin rms_norm implementation and run the reproduction across the listed shapes and dtypes. Compare the behavior with related PR #5771; done means INV_RMS stores the expected inverse RMS rather than var and the backward gradients match the PyTorch reference.

Written by the indexing model from the issue text.