[BUG] `_rms_norm_fwd_fused` drops `x_stride1` in its normalize pass, so a non-unit last-dim stride gives wrong output

[truong-v]

Summary

lightllm/common/basemodel/triton_kernel/norm/rmsnorm.py takes x_stride1 and uses it in the variance pass, and the store uses y_stride1, but the normalize pass indexes the input as if it were contiguous:

    for off in range(0, N, BLOCK_SIZE):
        cols = off + tl.arange(0, BLOCK_SIZE)
        x = tl.load(X + cols * x_stride1, mask=cols < N, other=0.0).to(tl.float32)   # line 32, honours the stride
        _var += x * x
    ...
    for off in range(0, N, BLOCK_SIZE):
        cols = off + tl.arange(0, BLOCK_SIZE)
        mask = cols < N
        ...
        x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32)                   # line 42, does not
        ...
        tl.store(Y + cols * y_stride1, y.to(Y.dtype.element_ty), mask=mask)          # line 48, honours it

So whenever x.stride(-1) != 1 the row's RMS is computed from the correct elements while the values that get scaled by it are read from somewhere else — for a 2-D transposed input, straight down a column. The result is silently wrong; there is no error and nothing goes out of bounds.

rmsnorm_forward has no guard against this. It does x_arg = x.view(-1, x.shape[-1]), which for a 2-D input is shape-preserving and therefore succeeds whatever the strides are, and torch.empty_like(x) gives the output the same strides, so y_stride1 matches and only the load is wrong.

The three in-tree call sites (gemma4, deepseek3_2, NormWeight) all pass tensors whose last-dim stride is 1 today — deepseek3_2 slices cache_kv[:, :, :kv_lora_rank], which keeps stride 1 on the last dim — so this does not currently break a shipped model. But the parameter exists to support strided inputs, two of the three places that use it get it right, and any caller that passes one gets wrong numbers rather than a complaint.

Introduced with the kernel; the file was last touched in #1304.

Reproduction

import torch
from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward, torch_rms_norm

M, N = 64, 256
torch.manual_seed(0)
weight = torch.rand(N, dtype=torch.float32, device="cuda")

x = torch.randn(N, M, device="cuda").t()          # shape (M, N), stride (1, M)
out = rmsnorm_forward(x, weight, 1e-6)
ref = torch_rms_norm(x.float(), weight.float(), 1e-6)
print("strided   max|err|:", (out.float() - ref).abs().max().item())

xc = x.contiguous()                                # identical values, stride (N, 1)
out = rmsnorm_forward(xc, weight, 1e-6)
ref = torch_rms_norm(xc.float(), weight.float(), 1e-6)
print("contiguous max|err|:", (out.float() - ref).abs().max().item())

strided   max|err|: 5.832374572753906
contiguous max|err|: 0.0

15801 of 16384 output elements are wrong in the strided case (16230 with weight=None).

Expected Behavior

The normalize pass reads with x_stride1, like the variance pass above it and the store below it.

Environment

• ModelTC/lightllm at fe9bdabfc331b990124f1ec27daf6bb7945cf7ee
• NVIDIA B200 (sm_100), driver 595.71.05
• torch 2.13.0+cu130, triton 3.7.1, Python 3.12