最小用例

import torch
def rms_norm_weight_grad(grad_out_base, q, q_square_sum):
    grad_out = grad_out_base.permute(0, 2, 1, 3)
    inv_rms = torch.rsqrt(q_square_sum.unsqueeze(-1) / q.shape[-1] + 1e-6)
    grad_weight = (grad_out * q.float() * inv_rms).sum(dim=(0, 1, 2))
    return grad_weight.to(torch.bfloat16)

def main():
    batch = 2
    seq = 4096
    heads = 64
    head_dim = 128
    torch.manual_seed(0)
    device = "npu"

    grad_out_base = torch.randn(
        (batch, heads, seq, head_dim), device=device, dtype=torch.float32
    )
    q = torch.randn((batch, seq, heads, head_dim), device=device, dtype=torch.bfloat16)
    q_square_sum = torch.rand((batch, seq, heads), device=device, dtype=torch.float32) + 1.0

    compiled = torch.compile(rms_norm_weight_grad, backend="inductor", dynamic=False)
    out = compiled(grad_out_base, q, q_square_sum)
    torch_npu.npu.synchronize()
    print(f"compiled_out.shape={tuple(out.shape)}, dtype={out.dtype}")
    print(f"compiled_out[:8]={out[:8]}")

    expected = rms_norm_weight_grad(grad_out_base, q, q_square_sum)
    torch_npu.npu.synchronize()
    torch.testing.assert_close(out, expected, rtol=1e-2, atol=1e-2)
    print("check=passed")

问题DSL

@triton.jit
def triton_(in_ptr0, in_ptr1, in_ptr2, out_ptr1, x0_numel, r3_numel, r2_numel, r1_numel, X0BLOCK : tl.constexpr, X0BLOCK_SUB : tl.constexpr, R3BLOCK_SUB : tl.constexpr, R2BLOCK_SUB : tl.constexpr, R1BLOCK_SUB : tl.constexpr):
    x0_offset = tl.program_id(0) * X0BLOCK
    base_x0= tl.arange(0, X0BLOCK_SUB)
    loops_x0 = (X0BLOCK + X0BLOCK_SUB - 1) // X0BLOCK_SUB
    base_r3= tl.arange(0, R3BLOCK_SUB)
    loops_r3 = (r3_numel + R3BLOCK_SUB - 1) // R3BLOCK_SUB
    base_r2= tl.arange(0, R2BLOCK_SUB)
    loops_r2 = (r2_numel + R2BLOCK_SUB - 1) // R2BLOCK_SUB
    base_r1= tl.arange(0, R1BLOCK_SUB)
    loops_r1 = (r1_numel + R1BLOCK_SUB - 1) // R1BLOCK_SUB
    for loop_x0 in range(loops_x0):
        x0 = x0_offset + (loop_x0 * X0BLOCK_SUB) + base_x0[None,None,None,:]
        x0_mask = x0 < min(X0BLOCK+x0_offset, x0_numel)
        for loop_r3 in range(loops_r3):
            r3 = (loop_r3 * R3BLOCK_SUB) + base_r3[:,None,None,None]
            r3_mask = r3 < r3_numel
            for loop_r2 in range(loops_r2):
                r2_1 = (loop_r2 * R2BLOCK_SUB) + base_r2[None,:,None,None]
                r2 = (loop_r2 * R2BLOCK_SUB) + base_r2[None,None,:,None]
                r2_mask = r2 < r2_numel
                r2_1_mask = r2_1 < r2_numel
                _tmp13 = tl.full([X0BLOCK_SUB, R3BLOCK_SUB * R2BLOCK_SUB * R1BLOCK_SUB], 0, tl.float32)
                for loop_r1 in range(loops_r1):
                    r1_2 = (loop_r1 * R1BLOCK_SUB) + base_r1[None,None,:,None]
                    r1 = (loop_r1 * R1BLOCK_SUB) + base_r1[None,:,None,None]
                    r1_mask = r1 < r1_numel
                    r1_2_mask = r1_2 < r1_numel
                    tmp0 = tl.load(in_ptr0 + (x0 + 128*r2 + 524288*r1 + 33554432*r3), r1_mask & r2_mask & r3_mask & x0_mask, other=0.0)
                    tmp1 = tl.load(in_ptr1 + (x0 + 128*r1_2 + 8192*r2_1 + 33554432*r3), r1_2_mask & r2_1_mask & r3_mask & x0_mask, other=0.0).to(tl.float32)
                    tmp2 = tmp1.permute([0, 2, 1, 3])
                    tmp4 = tl.load(in_ptr2 + (r1 + 64*r2 + 262144*r3), r1_mask & r2_mask & r3_mask, other=0.0)
                    tmp3 = tmp2.to(tl.float32)
                    tmp5 = 128.0
                    tmp6 = (tmp4 / tmp5)
                    tmp7 = 9.999999974752427e-07
                    tmp8 = tmp6 + tmp7
                    tmp9 = tl.rsqrt(tmp8)
                    tmp10 = tmp3 * tmp9
                    tmp11 = tmp0 * tmp10
                    tmp12 = tl.reshape(tmp11, [X0BLOCK_SUB, R3BLOCK_SUB * R2BLOCK_SUB * R1BLOCK_SUB])
                    tmp14 = _tmp13 + tmp12
                    _tmp13 = tl.where((r1_mask & r2_mask & r3_mask & x0_mask).reshape([X0BLOCK_SUB, R3BLOCK_SUB * R2BLOCK_SUB * R1BLOCK_SUB]), tmp14, _tmp13)
                tmp13 = tl.sum(_tmp13, 0).reshape(1, 1, 1, R3BLOCK_SUB * R2BLOCK_SUB * R1BLOCK_SUB)
                tmp15 = tmp13.to(tl.float32)
                tl.store(out_ptr1 + (x0 + tl.arange(0,1) ), tmp15, x0_mask)

1. 多个 reduction 轴 flatten 后，accumulator(_tmp13) 初始化、post-loop reduction/store 的作用域不正确导致只会存储部分结果
2. flatten 前的轴顺序没有保持一致，导致reshape失败。

做了以下修复：

1. 在 multi-reduction 场景下，将 reduction accumulator 的初始化和最终 reduce/store 放到完整 reduction loop 外侧，避免中间 tile 清零或提前 store。
2. 对 contiguous multi-reduction 的 value/mask 在 flatten 前按 [non-reduction axes, reduction axes] 调整轴顺序，保证 tl.sum 沿正确 reduction lane 执行。
3. 修正 contiguous multi-reduction 的 reduction dim 计算，使最终 tl.sum 维度与 flatten 后布局一致。
   修复后DSL：

def triton_red_fused__to_copy_add_div_mul_permute__0(in_ptr0, in_ptr1, in_ptr2, out_ptr1, x0_numel, r3_numel, r2_numel, r1_numel, X0BLOCK : tl.constexpr, X0BLOCK_SUB : tl.constexpr, R3BLOCK_SUB : tl.constexpr, R2BLOCK_SUB : tl.constexpr, R1BLOCK_SUB : tl.constexpr):
    x0_offset = tl.program_id(0) * X0BLOCK
    base_x0= tl.arange(0, X0BLOCK_SUB)
    loops_x0 = (X0BLOCK + X0BLOCK_SUB - 1) // X0BLOCK_SUB
    base_r3= tl.arange(0, R3BLOCK_SUB)
    loops_r3 = (r3_numel + R3BLOCK_SUB - 1) // R3BLOCK_SUB
    base_r2= tl.arange(0, R2BLOCK_SUB)
    loops_r2 = (r2_numel + R2BLOCK_SUB - 1) // R2BLOCK_SUB
    base_r1= tl.arange(0, R1BLOCK_SUB)
    loops_r1 = (r1_numel + R1BLOCK_SUB - 1) // R1BLOCK_SUB
    for loop_x0 in range(loops_x0):
        x0 = x0_offset + (loop_x0 * X0BLOCK_SUB) + base_x0[None,None,None,:]
        x0_mask = x0 < min(X0BLOCK+x0_offset, x0_numel)
        _tmp13 = tl.full([X0BLOCK_SUB, R3BLOCK_SUB * R2BLOCK_SUB * R1BLOCK_SUB], 0, tl.float32)
        for loop_r3 in range(loops_r3):
            r3 = (loop_r3 * R3BLOCK_SUB) + base_r3[:,None,None,None]
            r3_mask = r3 < r3_numel
            for loop_r2 in range(loops_r2):
                r2_1 = (loop_r2 * R2BLOCK_SUB) + base_r2[None,:,None,None]
                r2 = (loop_r2 * R2BLOCK_SUB) + base_r2[None,None,:,None]
                r2_mask = r2 < r2_numel
                r2_1_mask = r2_1 < r2_numel
                for loop_r1 in range(loops_r1):
                    r1_2 = (loop_r1 * R1BLOCK_SUB) + base_r1[None,None,:,None]
                    r1 = (loop_r1 * R1BLOCK_SUB) + base_r1[None,:,None,None]
                    r1_mask = r1 < r1_numel
                    r1_2_mask = r1_2 < r1_numel
                    tmp0 = tl.load(in_ptr0 + (x0 + 128*r2 + 524288*r1 + 33554432*r3), r1_mask & r2_mask & r3_mask & x0_mask, other=0.0)
                    tmp1 = tl.load(in_ptr1 + (x0 + 128*r1_2 + 8192*r2_1 + 33554432*r3), r1_2_mask & r2_1_mask & r3_mask & x0_mask, other=0.0).to(tl.float32)
                    tmp2 = tmp1.permute([0, 2, 1, 3])
                    tmp5 = tl.load(in_ptr2 + (r1 + 64*r2 + 262144*r3), r1_mask & r2_mask & r3_mask, other=0.0)
                    tmp3 = tmp2.to(tl.float32)
                    tmp4 = tmp0 * tmp3
                    tmp6 = 0.0078125
                    tmp7 = tmp5 * tmp6
                    tmp8 = 1e-06
                    tmp9 = tmp7 + tmp8
                    tmp10 = tl.rsqrt(tmp9)
                    tmp11 = tmp4 * tmp10
                    tmp12 = tl.reshape(tmp11.permute([3, 0, 1, 2]), [X0BLOCK_SUB, R3BLOCK_SUB * R2BLOCK_SUB * R1BLOCK_SUB])
                    tmp14 = _tmp13 + tmp12
                    _tmp13 = tl.where((r1_mask & r2_mask & r3_mask & x0_mask).permute([3, 0, 1, 2]).reshape([X0BLOCK_SUB, R3BLOCK_SUB * R2BLOCK_SUB * R1BLOCK_SUB]), tmp14, _tmp13)
        tmp13 = tl.sum(_tmp13, 1).reshape(1, 1, 1, X0BLOCK_SUB)
        tmp15 = tmp13.to(tl.float32)
        tl.store(out_ptr1 + (x0 + tl.arange(0,1) ), tmp15, x0_mask)