feat(triton): add bwd_dadt_cumsum_fused_kernel operator for Mamba backward pass-MindSpeed-Ops-AtomGit

Amindspeed_ops/api/triton/bwd_dadt_cumsum_fused.py+62-0

@@ -0,0 +1,62 @@
	1	+# Copyright (c) 2023, Tri Dao, Albert Gu
	2	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	3	+
	L L LinShuaSep 15 需要补充GPU上pytorch小算子与GPU上开源triton算子的精度对齐截图（以确保pytorch小算子标杆的准确性）
	4	+import torch
	5	+
	6	+from mindspeed_ops.arch32.triton.bwd_dadt_cumsum_fused import bwd_dadt_cumsum_fused_kernel
	7	+from mindspeed_ops.utils import is_arch35
	8	+
	9	+__all__ = ["bwd_dadt_cumsum_fused"]
	10	+
	11	+
	12	+def bwd_dadt_cumsum_fused(
	13	+ ddA_cs: torch.Tensor,
	14	+ ddA_cs_rev: torch.Tensor,
	15	+ dA_cs: torch.Tensor,
	16	+ dA_cs_rev: torch.Tensor,
	17	+ ddt_out: torch.Tensor,
	18	+ chunk_size: int,
	19	+) -> torch.Tensor:
	20	+ if is_arch35():
	21	+ raise NotImplementedError("bwd_dadt_cumsum_fused is not supported on arch35")
	atomgit-botJul 2 🟡 Medium Priority API 函数 `bwd_dadt_cumsum_fused` 接受 `chunk_size: int` 参数（第 18 行），但未对其进行有效性验证。当 `chunk_size <= 0` 时，第 36 行 `nchunks = (S + chunk_size - 1) // chunk_size` 会：建议：在 arch35 检查之后、contiguous 调用之前增加 chunk_size 正值校验，例如：`if chunk_size <= 0: raise ValueError(f"chunk_size must be positive, got {chunk_size}")` Incorrect?
	22	+
	23	+ if chunk_size <= 0:
	24	+ raise ValueError(f"chunk_size must be positive, got {chunk_size}")
	25	+
	26	+ ddA_cs = ddA_cs.contiguous()
	27	+ ddA_cs_rev = ddA_cs_rev.contiguous()
	28	+ dA_cs = dA_cs.contiguous()
	29	+ dA_cs_rev = dA_cs_rev.contiguous()
	30	+ ddt_out = ddt_out.contiguous()
	31	+
	32	+ B, H, S = ddA_cs.shape
	33	+
	34	+ if ddA_cs_rev.shape != (B, H, S):
	35	+ raise ValueError(f"ddA_cs_rev shape mismatch: expected ({B}, {H}, {S}), got {ddA_cs_rev.shape}")
	36	+ if dA_cs.shape != (B, H, S):
	37	+ raise ValueError(f"dA_cs shape mismatch: expected ({B}, {H}, {S}), got {dA_cs.shape}")
	38	+ if dA_cs_rev.shape != (B, H, S):
	39	+ raise ValueError(f"dA_cs_rev shape mismatch: expected ({B}, {H}, {S}), got {dA_cs_rev.shape}")
	40	+ if ddt_out.shape != (B, H, S):
	41	+ raise ValueError(f"ddt_out shape mismatch: expected ({B}, {H}, {S}), got {ddt_out.shape}")
	42	+
	43	+ nchunks = (S + chunk_size - 1) // chunk_size
	atomgit-botJul 2 🟡 Medium Priority API 函数 `bwd_dadt_cumsum_fused` (api 文件第 18 行) 接受 `chunk_size: int` 参数，在第 36 行直接用于 `nchunks = (S + chunk_size - 1) // chunk_size` 计算。当 `chunk_size = 0` 时触发 `ZeroDivisionError`；当 `chunk_size < 0` 时产生无意义的负数 grid 维度，kernel 行为未定义。该参数来源于调用方，无任何边界保护。建议：在函数开头添加 `chunk_size` 的正整数校验，例如：`if chunk_size <= 0: raise ValueError(f"chunk_size must be positive, got {chunk_size}")` Incorrect?
	44	+
	45	+ grid = (B * H, nchunks)
	46	+
	47	+ bwd_dadt_cumsum_fused_kernel[grid](
	48	+ ddA_cs,
	49	+ ddA_cs_rev,
	50	+ dA_cs,
	51	+ dA_cs_rev,
	52	+ ddt_out,
	53	+ ddA_cs.stride(0),
	54	+ ddA_cs.stride(1),
	55	+ ddA_cs.stride(2),
	56	+ B=B,
	57	+ H=H,
	58	+ S=S,
	59	+ CHUNK_SIZE=chunk_size,
	60	+ )
	61	+
	62	+ return ddt_out

Amindspeed_ops/arch32/triton/bwd_dadt_cumsum_fused.py+48-0

@@ -0,0 +1,48 @@
	1	+# Copyright (c) 2023, Tri Dao, Albert Gu
	2	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	3	+# pylint: disable=import-error
	4	+import triton
	5	+import triton.language as tl # pylint: disable=E0611
	6	+
	7	+
	8	+@triton.jit # pylint: disable=E1101
	ascend-robotJul 2 This code comment section is from+4To+8 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。
	9	+def bwd_dadt_cumsum_fused_kernel(
	ascend-robotJul 2 This code comment section is from+5To+9 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。
	10	+ ddA_cs_ptr,
	11	+ ddA_cs_rev_ptr,
	12	+ dA_cs_ptr,
	13	+ dA_cs_rev_ptr,
	14	+ ddt_out_ptr,
	15	+ stride_batch,
	16	+ stride_head,
	17	+ stride_seq,
	18	+ B: tl.constexpr,
	19	+ H: tl.constexpr,
	20	+ S: tl.constexpr,
	21	+ CHUNK_SIZE: tl.constexpr,
	22	+):
	23	+ pid_bh = tl.program_id(0)
	24	+ pid_chunk = tl.program_id(1)
	25	+
	26	+ pid_batch = pid_bh // H
	27	+ pid_head = pid_bh % H
	28	+
	29	+ chunk_start = pid_chunk * CHUNK_SIZE
	30	+ offs_seq = chunk_start + tl.arange(0, CHUNK_SIZE)
	31	+ mask = offs_seq < S
	32	+
	33	+ base_offset = pid_batch * stride_batch + pid_head * stride_head
	34	+
	35	+ ddA_cs = tl.load(ddA_cs_ptr + base_offset + offs_seq * stride_seq, mask=mask, other=0.0)
	36	+ ddA_cs_rev = tl.load(ddA_cs_rev_ptr + base_offset + offs_seq * stride_seq, mask=mask, other=0.0)
	37	+ dA_cs = tl.load(dA_cs_ptr + base_offset + offs_seq * stride_seq, mask=mask, other=0.0)
	38	+ dA_cs_rev = tl.load(dA_cs_rev_ptr + base_offset + offs_seq * stride_seq, mask=mask, other=0.0)
	39	+
	40	+ scaled_ddA_cs = tl.exp(dA_cs) * ddA_cs
	41	+ ddt_cs = tl.cumsum(scaled_ddA_cs, axis=0, reverse=True)
	42	+
	43	+ scaled_ddA_cs_rev = tl.exp(dA_cs_rev) * ddA_cs_rev
	44	+ ddt_cs_rev_inclusive = tl.cumsum(scaled_ddA_cs_rev, axis=0)
	45	+ ddt_cs_rev_exclusive = ddt_cs_rev_inclusive - scaled_ddA_cs_rev
	46	+
	47	+ ddt_total = ddt_cs + ddt_cs_rev_exclusive
	48	+ tl.store(ddt_out_ptr + base_offset + offs_seq * stride_seq, ddt_total, mask=mask)

Atests/atk_tests/triton/bwd_dadt_cumsum_fused/bwd_dadt_cumsum_fused.yaml+95-0

@@ -0,0 +1,95 @@
	1	+# Copyright (c) 2023, Tri Dao, Albert Gu
	2	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	3	+
	4	+api: pytorch
	5	+version: v2.1
	6	+name: torch_bwd_dadt_cumsum_fused
	7	+triton_name: triton_bwd_dadt_cumsum_fused.TritonBwdDadtCumsumFusedApi
	8	+api_type: torch_bwd_dadt_cumsum_fused
	9	+triton_api_type: triton_bwd_dadt_cumsum_fused
	10	+generate: generate_bwd_dadt_cumsum_fused
	11	+dtype_numbers: 30
	12	+standard:
	13	+ acc: single_bm
	14	+ perf: not_key
	15	+inputs:
	16	+ - name: ddA_cs
	17	+ type: tensor
	18	+ required: true
	19	+ dtypes:
	20	+ values: [ fp32 ]
	21	+ ranges:
	22	+ valid:
	23	+ values: [ [-1,1] ]
	24	+ invalid:
	25	+ values: [ [-1,1] ]
	26	+ shapes:
	27	+ dim_numbers:
	28	+ values: [ 3 ]
	29	+ dim_values:
	30	+ values: [ 1,2,4,8,16,32,64 ]
	31	+ max_length: 262144
	32	+ - name: ddA_cs_rev
	33	+ type: tensor
	34	+ required: true
	35	+ dtypes:
	36	+ values: [ fp32 ]
	37	+ ranges:
	38	+ valid:
	39	+ values: [ [-1,1] ]
	40	+ invalid:
	41	+ values: [ [-1,1] ]
	42	+ shapes:
	43	+ dim_numbers:
	44	+ values: [ 3 ]
	45	+ dim_values:
	46	+ values: [ 1,2,4,8,16,32,64 ]
	47	+ max_length: 262144
	48	+ - name: dA_cs
	49	+ type: tensor
	50	+ required: true
	51	+ dtypes:
	52	+ values: [ fp32 ]
	53	+ ranges:
	54	+ valid:
	55	+ values: [ [-1,1] ]
	56	+ invalid:
	57	+ values: [ [-1,1] ]
	58	+ shapes:
	59	+ dim_numbers:
	60	+ values: [ 3 ]
	61	+ dim_values:
	62	+ values: [ 1,2,4,8,16,32,64 ]
	63	+ max_length: 262144
	64	+ - name: dA_cs_rev
	65	+ type: tensor
	66	+ required: true
	67	+ dtypes:
	68	+ values: [ fp32 ]
	69	+ ranges:
	70	+ valid:
	71	+ values: [ [-1,1] ]
	72	+ invalid:
	73	+ values: [ [-1,1] ]
	74	+ shapes:
	75	+ dim_numbers:
	76	+ values: [ 3 ]
	77	+ dim_values:
	78	+ values: [ 1,2,4,8,16,32,64 ]
	79	+ max_length: 262144
	80	+ - name: ddt_out
	81	+ type: tensor
	82	+ required: true
	83	+ dtypes:
	84	+ values: [ fp32 ]
	85	+ ranges:
	86	+ valid:
	87	+ values: [ [-1,1] ]
	88	+ invalid:
	89	+ values: [ [-1,1] ]
	90	+ shapes:
	91	+ dim_numbers:
	92	+ values: [ 3 ]
	93	+ dim_values:
	94	+ values: [ 1,2,4,8,16,32,64 ]
	95	+ max_length: 262144

Atests/atk_tests/triton/bwd_dadt_cumsum_fused/generate_bwd_dadt_cumsum_fused.py+35-0

@@ -0,0 +1,35 @@
	1	+# Copyright (c) 2023, Tri Dao, Albert Gu
	2	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	3	+
	4	+from atk.case_generator.generator.base_generator import CaseGenerator
	5	+from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
	6	+from atk.configs.case_config import CaseConfig
	7	+
	8	+
	9	+CASES = (
	10	+ (1, 1, 8),
	11	+ (1, 1, 16),
	12	+ (1, 2, 32),
	13	+ (2, 4, 64),
	14	+ (1, 1, 128),
	15	+ (1, 8, 64),
	16	+ (1, 1, 256),
	17	+ (1, 4, 256),
	18	+ (2, 8, 128),
	19	+ (1, 32, 512),
	20	+)
	21	+
	22	+
	23	+@GENERATOR_REGISTRY.register("generate_bwd_dadt_cumsum_fused")
	24	+class BwdDadtCumsumFusedGenerator(CaseGenerator):
	25	+ _case_index = 0
	26	+
	27	+ def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
	28	+ B, H, S = CASES[BwdDadtCumsumFusedGenerator._case_index % len(CASES)]
	29	+ BwdDadtCumsumFusedGenerator._case_index += 1
	30	+
	31	+ for input_idx in range(5):
	32	+ case_config.inputs[input_idx].shape = [B, H, S]
	33	+ case_config.inputs[input_idx].dtype = "fp32"
	34	+
	35	+ return case_config

Atests/atk_tests/triton/bwd_dadt_cumsum_fused/triton_bwd_dadt_cumsum_fused.py+80-0

@@ -0,0 +1,80 @@
	1	+# Copyright (c) 2023, Tri Dao, Albert Gu
	2	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	3	+# pylint: disable=duplicate-code
	4	+# - duplicate-code: UT and ATK share the same reference implementation to keep
	5	+# both checks aligned with the operator formula.
	6	+
	7	+import torch
	8	+
	9	+from atk.configs.dataset_config import InputDataset
	10	+from atk.tasks.api_execute import register
	11	+from atk.tasks.api_execute.base_api import BaseApi
	12	+from atk.tasks.api_execute.triton_base_api import TritonBaseApi
	13	+from mindspeed_ops.api.triton.bwd_dadt_cumsum_fused import bwd_dadt_cumsum_fused
	14	+
	15	+
	16	+def _bwd_dadt_cumsum_fused_reference(
	17	+ ddA_cs: torch.Tensor,
	18	+ ddA_cs_rev: torch.Tensor,
	19	+ dA_cs: torch.Tensor,
	20	+ dA_cs_rev: torch.Tensor,
	21	+ chunk_size: int,
	22	+) -> torch.Tensor:
	23	+ """CPU reference implementation with chunked computation."""
	24	+ ddA_cs_f = ddA_cs.float()
	25	+ ddA_cs_rev_f = ddA_cs_rev.float()
	26	+ dA_cs_f = dA_cs.float()
	27	+ dA_cs_rev_f = dA_cs_rev.float()
	28	+
	29	+ B, H, S = ddA_cs_f.shape
	30	+
	31	+ pad_len = (chunk_size - S % chunk_size) % chunk_size
	32	+ if pad_len > 0:
	33	+ ddA_cs_f = torch.nn.functional.pad(ddA_cs_f, (0, pad_len))
	34	+ ddA_cs_rev_f = torch.nn.functional.pad(ddA_cs_rev_f, (0, pad_len))
	35	+ dA_cs_f = torch.nn.functional.pad(dA_cs_f, (0, pad_len))
	36	+ dA_cs_rev_f = torch.nn.functional.pad(dA_cs_rev_f, (0, pad_len))
	37	+
	38	+ S_padded = ddA_cs_f.shape[2]
	39	+ nchunks = S_padded // chunk_size
	40	+
	41	+ ddA_cs_chunks = ddA_cs_f.reshape(B, H, nchunks, chunk_size)
	42	+ ddA_cs_rev_chunks = ddA_cs_rev_f.reshape(B, H, nchunks, chunk_size)
	43	+ dA_cs_chunks = dA_cs_f.reshape(B, H, nchunks, chunk_size)
	44	+ dA_cs_rev_chunks = dA_cs_rev_f.reshape(B, H, nchunks, chunk_size)
	45	+
	46	+ scaled_ddA_cs = torch.exp(dA_cs_chunks) * ddA_cs_chunks
	47	+ ddt_cs = torch.cumsum(scaled_ddA_cs.flip(3), dim=3).flip(3)
	48	+
	49	+ scaled_ddA_cs_rev = torch.exp(dA_cs_rev_chunks) * ddA_cs_rev_chunks
	50	+ ddt_cs_rev_inclusive = torch.cumsum(scaled_ddA_cs_rev, dim=3)
	51	+ ddt_cs_rev_exclusive = ddt_cs_rev_inclusive - scaled_ddA_cs_rev
	52	+
	53	+ ddt_total = ddt_cs + ddt_cs_rev_exclusive
	54	+ ddt_total = ddt_total.reshape(B, H, S_padded)
	55	+ return ddt_total[:, :, :S]
	56	+
	57	+
	58	+def _get_inputs(input_data: InputDataset):
	59	+ keys = ("ddA_cs", "ddA_cs_rev", "dA_cs", "dA_cs_rev", "ddt_out")
	60	+ if input_data.kwargs:
	61	+ return tuple(input_data.kwargs[key] for key in keys)
	62	+ return tuple(input_data.args[: len(keys)])
	63	+
	64	+
	65	+@register("torch_bwd_dadt_cumsum_fused")
	66	+class TorchBwdDadtCumsumFusedApi(BaseApi):
	67	+ def __call__(self, input_data: InputDataset, with_output: bool = False):
	68	+ ddA_cs, ddA_cs_rev, dA_cs, dA_cs_rev, ddt_out = _get_inputs(input_data)
	69	+ chunk_size = 64
	70	+ output = _bwd_dadt_cumsum_fused_reference(ddA_cs, ddA_cs_rev, dA_cs, dA_cs_rev, chunk_size)
	71	+ return output
	72	+
	73	+
	74	+@register("triton_bwd_dadt_cumsum_fused")
	75	+class TritonBwdDadtCumsumFusedApi(TritonBaseApi):
	76	+ def __call__(self, input_data: InputDataset, with_output: bool = False):
	77	+ ddA_cs, ddA_cs_rev, dA_cs, dA_cs_rev, ddt_out = _get_inputs(input_data)
	78	+ chunk_size = 64
	79	+ output = bwd_dadt_cumsum_fused(ddA_cs, ddA_cs_rev, dA_cs, dA_cs_rev, ddt_out.clone(), chunk_size)
	80	+ return output

Atests/unit_tests/triton/test_bwd_dadt_cumsum_fused.py+84-0

@@ -0,0 +1,84 @@
	1	+# Copyright (c) 2023, Tri Dao, Albert Gu
	2	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	3	+# pylint: disable=duplicate-code
	4	+# - duplicate-code: UT and ATK share the same reference implementation and
	5	+# shape list to keep both checks aligned with the operator formula.
	6	+
	7	+import pytest
	8	+import torch
	9	+
	10	+from mindspeed_ops.api.triton.bwd_dadt_cumsum_fused import bwd_dadt_cumsum_fused
	11	+from mindspeed_ops.api.triton.utils import get_available_device
	12	+from tests.utils import assert_close
	13	+
	14	+
	15	+DEVICE = get_available_device()
	16	+DTYPES = [torch.float32]
	17	+CHUNK_SIZES = [16, 64]
	18	+TEST_SHAPES = [
	19	+ (1, 1, 8), # B=1, H=1, S=8, small sequence
	20	+ (1, 1, 16), # B=1, H=1, S=16
	21	+ (1, 2, 32), # B=1, H=2, S=32
	22	+ (2, 4, 64), # B=2, H=4, S=64, larger batch/heads
	23	+ (1, 1, 128), # B=1, H=1, S=128, longer sequence
	24	+ (1, 8, 64), # B=1, H=8, S=64, many heads
	25	+]
	26	+TEST_CASES = [(shape, chunk_size, dtype) for shape in TEST_SHAPES for chunk_size in CHUNK_SIZES for dtype in DTYPES]
	27	+RATIO = 2
	28	+ATOL = 2e-2
	29	+
	30	+
	31	+def cpu_golden(ddA_cs, ddA_cs_rev, dA_cs, dA_cs_rev, chunk_size):
	L L LinShuaJul 25 当前PR中的torch小算子由于是自己实现的，需要证明下torch小算子实现的正确性，建议在GPU上分别输入相同的输入对比开源triton和torch小算子输出结果，提供相关数据来证明torch小算子实现的正确性。
	32	+ """CPU reference implementation with chunked computation."""
	33	+ B, H, S = ddA_cs.shape
	34	+
	35	+ pad_len = (chunk_size - S % chunk_size) % chunk_size
	36	+ if pad_len > 0:
	37	+ ddA_cs = torch.nn.functional.pad(ddA_cs, (0, pad_len))
	38	+ ddA_cs_rev = torch.nn.functional.pad(ddA_cs_rev, (0, pad_len))
	39	+ dA_cs = torch.nn.functional.pad(dA_cs, (0, pad_len))
	40	+ dA_cs_rev = torch.nn.functional.pad(dA_cs_rev, (0, pad_len))
	41	+
	42	+ S_padded = ddA_cs.shape[2]
	43	+ nchunks = S_padded // chunk_size
	44	+
	45	+ ddA_cs_chunks = ddA_cs.reshape(B, H, nchunks, chunk_size)
	46	+ ddA_cs_rev_chunks = ddA_cs_rev.reshape(B, H, nchunks, chunk_size)
	47	+ dA_cs_chunks = dA_cs.reshape(B, H, nchunks, chunk_size)
	48	+ dA_cs_rev_chunks = dA_cs_rev.reshape(B, H, nchunks, chunk_size)
	49	+
	50	+ scaled_ddA_cs = torch.exp(dA_cs_chunks) * ddA_cs_chunks
	51	+ ddt_cs = torch.cumsum(scaled_ddA_cs.flip(3), dim=3).flip(3)
	52	+
	53	+ scaled_ddA_cs_rev = torch.exp(dA_cs_rev_chunks) * ddA_cs_rev_chunks
	54	+ ddt_cs_rev_inclusive = torch.cumsum(scaled_ddA_cs_rev, dim=3)
	55	+ ddt_cs_rev_exclusive = ddt_cs_rev_inclusive - scaled_ddA_cs_rev
	56	+
	57	+ ddt_total = ddt_cs + ddt_cs_rev_exclusive
	58	+ ddt_total = ddt_total.reshape(B, H, S_padded)
	59	+ return ddt_total[:, :, :S]
	60	+
	61	+
	62	+@pytest.mark.parametrize("shape,chunk_size,dtype", TEST_CASES)
	63	+def test_bwd_dadt_cumsum_fused(shape, chunk_size, dtype):
	64	+ B, H, S = shape
	65	+ if chunk_size > S * 2:
	66	+ pytest.skip(f"chunk_size={chunk_size} too large for S={S}")
	67	+
	68	+ torch.manual_seed(42)
	69	+
	70	+ ddA_cs = torch.randn(B, H, S, dtype=dtype, device="cpu")
	71	+ ddA_cs_rev = torch.randn(B, H, S, dtype=dtype, device="cpu")
	72	+ dA_cs = torch.randn(B, H, S, dtype=dtype, device="cpu") * 0.1
	73	+ dA_cs_rev = torch.randn(B, H, S, dtype=dtype, device="cpu") * 0.1
	74	+
	75	+ out_golden = cpu_golden(ddA_cs.float(), ddA_cs_rev.float(), dA_cs.float(), dA_cs_rev.float(), chunk_size)
	76	+
	77	+ ddA_cs_npu = ddA_cs.to(DEVICE)
	78	+ ddA_cs_rev_npu = ddA_cs_rev.to(DEVICE)
	79	+ dA_cs_npu = dA_cs.to(DEVICE)
	80	+ dA_cs_rev_npu = dA_cs_rev.to(DEVICE)
	81	+ ddt_out_npu = torch.zeros(B, H, S, dtype=dtype, device=DEVICE)
	82	+
	83	+ result = bwd_dadt_cumsum_fused(ddA_cs_npu, ddA_cs_rev_npu, dA_cs_npu, dA_cs_rev_npu, ddt_out_npu, chunk_size)
	84	+ assert_close("cpu_and_triton", out_golden, result.cpu(), RATIO, err_atol=ATOL)