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

Amindspeed_ops/api/triton/bwd_segsum_dadt.py+59-0

@@ -0,0 +1,59 @@
	1	+# Copyright (c) 2023, Tri Dao, Albert Gu
	2	+# Copyright (c) 2026, HUAWEI CORPORATION. All rights reserved.
	Y Y yangcheng7月2日反向算子一般情况下不用单独添加api接口适配 Tsuki 8月10日评论： 8月10日评论：
	3	+
	4	+import torch
	5	+
	6	+from mindspeed_ops.arch32.triton.bwd_segsum_dadt import bwd_segsum_dadt_kernel
	7	+from mindspeed_ops.utils import is_arch35
	8	+
	9	+__all__ = ["bwd_segsum_dadt"]
	10	+
	11	+
	12	+def bwd_segsum_dadt(
	13	+ dSSdA: torch.Tensor,
	14	+ SSdA: torch.Tensor,
	15	+ ddt_out: torch.Tensor,
	16	+ chunk_size: int,
	17	+) -> torch.Tensor:
	18	+ if is_arch35():
	19	+ raise NotImplementedError("bwd_segsum_dadt is not supported on arch35")
	20	+
	21	+ dSSdA = dSSdA.contiguous()
	22	+ SSdA = SSdA.contiguous()
	23	+ ddt_out = ddt_out.contiguous()
	24	+
	25	+ B, H, nchunks, C, _ = dSSdA.shape
	26	+ S = nchunks * C
	27	+
	28	+ assert SSdA.shape == (B, H, nchunks, C, C), (
	29	+ f"SSdA shape mismatch: expected ({B}, {H}, {nchunks}, {C}, {C}), got {SSdA.shape}"
	30	+ )
	31	+ assert ddt_out.shape == (B, H, S), f"ddt_out shape mismatch: expected ({B}, {H}, {S}), got {ddt_out.shape}"
	32	+ assert C == chunk_size, f"C ({C}) must equal chunk_size ({chunk_size})"
	33	+
	34	+ grid = (B, H, nchunks)
	35	+
	36	+ bwd_segsum_dadt_kernel[grid](
	37	+ dSSdA,
	38	+ SSdA,
	39	+ ddt_out,
	40	+ dSSdA.stride(0),
	L L liuzhexu7月22日建议输入转连续后给到kernel，昇腾亲和连续数据，需避免大量stride逻辑。
	41	+ dSSdA.stride(1),
	42	+ dSSdA.stride(2),
	43	+ dSSdA.stride(3),
	44	+ dSSdA.stride(4),
	45	+ SSdA.stride(0),
	46	+ SSdA.stride(1),
	47	+ SSdA.stride(2),
	48	+ SSdA.stride(3),
	49	+ SSdA.stride(4),
	50	+ ddt_out.stride(0),
	51	+ ddt_out.stride(1),
	52	+ ddt_out.stride(2),
	53	+ B=B,
	54	+ H=H,
	55	+ nchunks=nchunks,
	56	+ CHUNK_SIZE=chunk_size,
	57	+ )
	58	+
	59	+ return ddt_out

Amindspeed_ops/arch32/triton/bwd_segsum_dadt.py+60-0

@@ -0,0 +1,60 @@
	1	+# Copyright (c) 2023, Tri Dao, Albert Gu
	2	+# Copyright (c) 2026, HUAWEI CORPORATION. All rights reserved.
	3	+# pylint: disable=import-error
	4	+import triton
	ascend-robot7月2日此条代码评论区间+2至+4 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。 ascend-robot7月2日此条代码评论区间+2至+4 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。 ascend-robot7月10日此条代码评论区间+2至+4 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。 ascend-robot7月10日此条代码评论区间+2至+4 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。
	5	+import triton.language as tl # pylint: disable=E0611
	6	+
	7	+
	8	+@triton.jit # pylint: disable=E1101
	ascend-robot7月2日此条代码评论区间+4至+8 已过期【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。 ascend-robot7月2日此条代码评论区间+4至+8 已过期【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。 ascend-robot7月10日此条代码评论区间+4至+8 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。 ascend-robot7月10日此条代码评论区间+4至+8 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。
	9	+def bwd_segsum_dadt_kernel(
	ascend-robot7月2日此条代码评论区间+5至+9 已过期【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。 ascend-robot7月2日此条代码评论区间+5至+9 已过期【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。 ascend-robot7月10日此条代码评论区间+5至+9 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。 ascend-robot7月10日此条代码评论区间+5至+9 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。
	10	+ dSSdA_ptr,
	11	+ SSdA_cs_ptr,
	12	+ ddt_out_ptr,
	13	+ stride_dSSdA_batch,
	14	+ stride_dSSdA_head,
	15	+ stride_dSSdA_chunk,
	16	+ stride_dSSdA_row,
	17	+ stride_dSSdA_col,
	18	+ stride_SSdA_batch,
	19	+ stride_SSdA_head,
	20	+ stride_SSdA_chunk,
	21	+ stride_SSdA_row,
	22	+ stride_SSdA_col,
	23	+ stride_ddt_batch,
	24	+ stride_ddt_head,
	25	+ stride_ddt_seq,
	26	+ B: tl.constexpr,
	27	+ H: tl.constexpr,
	28	+ nchunks: tl.constexpr,
	29	+ CHUNK_SIZE: tl.constexpr,
	30	+):
	31	+ pid_batch = tl.program_id(0)
	32	+ pid_head = tl.program_id(1)
	33	+ pid_chunk = tl.program_id(2)
	34	+
	35	+ chunk_start = pid_chunk * CHUNK_SIZE
	36	+ offs_c = tl.arange(0, CHUNK_SIZE)
	37	+ offs_seq = chunk_start + offs_c
	38	+
	39	+ dSSdA_offset = dSSdA_ptr + (
	40	+ pid_batch * stride_dSSdA_batch + pid_head * stride_dSSdA_head + pid_chunk * stride_dSSdA_chunk
	41	+ )
	42	+ SSdA_offset = SSdA_cs_ptr + (
	43	+ pid_batch * stride_SSdA_batch + pid_head * stride_SSdA_head + pid_chunk * stride_SSdA_chunk
	44	+ )
	45	+ ddt_ptrs = ddt_out_ptr + (pid_batch * stride_ddt_batch + pid_head * stride_ddt_head + offs_seq * stride_ddt_seq)
	46	+
	47	+ dSSdA_block = tl.load(dSSdA_offset + offs_c[:, None] * stride_dSSdA_col + offs_c[None, :] * stride_dSSdA_row)
	48	+ SSdA_block = tl.load(SSdA_offset + offs_c[:, None] * stride_SSdA_row + offs_c[None, :] * stride_SSdA_col)
	49	+
	50	+ dSSdA_block = dSSdA_block * tl.exp(SSdA_block)
	51	+ dSSdA_block = tl.cumsum(dSSdA_block, axis=0, reverse=True)
	52	+
	53	+ offs_i = tl.arange(0, CHUNK_SIZE)[:, None]
	54	+ offs_j = tl.arange(0, CHUNK_SIZE)[None, :]
	55	+ SS_mask = offs_i > offs_j
	56	+ dSSdA = tl.where(SS_mask, dSSdA_block, 0.0)
	57	+
	58	+ ddt_chunk = tl.load(ddt_ptrs)
	59	+ ddt_chunk += tl.sum(dSSdA, axis=1)
	60	+ tl.store(ddt_ptrs, ddt_chunk)

Atests/atk_tests/triton/bwd_segsum_dadt/bwd_segsum_dadt.yaml+72-0

@@ -0,0 +1,72 @@
	1	+api: pytorch
	2	+version: v2.1
	3	+name: torch_bwd_segsum_dadt
	4	+api_type: torch_bwd_segsum_dadt
	5	+triton_name: triton_bwd_segsum_dadt.TritonBwdSegsumDadtApi
	6	+triton_api_type: triton_bwd_segsum_dadt
	7	+generate: generator_bwd_segsum_dadt
	8	+dtype_numbers: 15
	9	+standard:
	10	+ acc:
	11	+ single_bm:
	12	+ type: high_performance
	13	+ fp16_error: 0.02
	14	+ fp16_eb: 0.002
	15	+ bf16_error: 0.01
	16	+ bf16_eb: 0.001
	17	+ fp32_error: 0.01
	18	+ fp32_eb: 0.001
	19	+ perf: not_key
	20	+inputs:
	21	+ - name: dSSdA
	22	+ type: tensor
	23	+ required: true
	24	+ dtypes:
	25	+ values: [ bf16, fp32 ]
	26	+ ranges:
	27	+ valid:
	28	+ values: [ [-1, 1] ]
	29	+ invalid:
	30	+ values: [ [-1, 1] ]
	31	+ shapes:
	32	+ dim_numbers:
	33	+ values: [ 5 ]
	34	+ max_length: 262144
	35	+ - name: SSdA
	36	+ type: tensor
	37	+ required: true
	38	+ dtypes:
	39	+ values: [ bf16, fp32 ]
	40	+ ranges:
	41	+ valid:
	42	+ values: [ [-1, 1] ]
	43	+ invalid:
	44	+ values: [ [-1, 1] ]
	45	+ shapes:
	46	+ dim_numbers:
	47	+ values: [ 5 ]
	48	+ max_length: 262144
	49	+ - name: ddt_out
	50	+ type: tensor
	51	+ required: true
	52	+ dtypes:
	53	+ values: [ fp32 ]
	54	+ ranges:
	55	+ valid:
	56	+ values: [ [-1, 1] ]
	57	+ invalid:
	58	+ values: [ [-1, 1] ]
	59	+ shapes:
	60	+ dim_numbers:
	61	+ values: [ 3 ]
	62	+ max_length: 262144
	63	+ - name: chunk_size
	64	+ type: attr
	65	+ required: true
	66	+ dtypes:
	67	+ values: [ int64_t ]
	68	+ ranges:
	69	+ valid:
	70	+ values: [ 16 ]
	71	+ invalid:
	72	+ values: [ 16 ]

Atests/atk_tests/triton/bwd_segsum_dadt/generate_bwd_segsum_dadt.py+62-0

@@ -0,0 +1,62 @@
	1	+# Copyright (c) 2023, Tri Dao, Albert Gu
	2	+# Copyright (c) 2026, HUAWEI CORPORATION. All rights reserved.
	3	+
	4	+import random
	5	+
	6	+from atk.case_generator.generator.base_generator import CaseGenerator
	7	+from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
	8	+from atk.configs.case_config import CaseConfig
	9	+
	10	+
	11	+@GENERATOR_REGISTRY.register("generator_bwd_segsum_dadt")
	12	+class BwdSegsumDadtGenerator(CaseGenerator):
	13	+ CASES = [
	14	+ (1, 1, 1, 8, 8),
	15	+ (1, 1, 1, 16, 16),
	16	+ (1, 2, 2, 16, 16),
	17	+ (2, 4, 4, 32, 32),
	18	+ (1, 8, 1, 32, 32),
	19	+ (2, 2, 4, 16, 16),
	20	+ (1, 4, 2, 32, 32),
	21	+ (1, 16, 1, 16, 16),
	22	+ ]
	23	+
	24	+ _case_pool = None
	25	+ _case_index = 0
	26	+
	27	+ @classmethod
	28	+ def _build_pool(cls):
	29	+ fixed = list(cls.CASES) * 3
	30	+ random.shuffle(fixed)
	31	+ return fixed
	32	+
	33	+ def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
	34	+ if BwdSegsumDadtGenerator._case_pool is None:
	35	+ BwdSegsumDadtGenerator._case_pool = self._build_pool()
	36	+
	37	+ if BwdSegsumDadtGenerator._case_index < len(BwdSegsumDadtGenerator._case_pool):
	38	+ # pylint: disable-next=unsubscriptable-object
	39	+ B, H, nchunks, C, _ = BwdSegsumDadtGenerator._case_pool[BwdSegsumDadtGenerator._case_index]
	40	+ else:
	41	+ B = random.choice([1, 2])
	42	+ H = random.choice([1, 2, 4, 8])
	43	+ nchunks = random.choice([1, 2, 4])
	44	+ C = random.choice([16, 32])
	45	+
	46	+ BwdSegsumDadtGenerator._case_index += 1
	47	+
	48	+ S = nchunks * C
	49	+
	50	+ dSSdA = case_config.inputs[0]
	51	+ SSdA = case_config.inputs[1]
	52	+ ddt_out = case_config.inputs[2]
	53	+ chunk_size = case_config.inputs[3]
	54	+
	55	+ dSSdA.shape = [B, H, nchunks, C, C]
	56	+ SSdA.shape = [B, H, nchunks, C, C]
	57	+ ddt_out.shape = [B, H, S]
	58	+ chunk_size.range_values = C
	59	+
	60	+ SSdA.dtype = dSSdA.dtype
	61	+
	62	+ return case_config

Atests/atk_tests/triton/bwd_segsum_dadt/triton_bwd_segsum_dadt.py+74-0

@@ -0,0 +1,74 @@
	1	+# Copyright (c) 2023, Tri Dao, Albert Gu
	2	+# Copyright (c) 2026, HUAWEI CORPORATION. All rights reserved.
	3	+# pylint: disable=duplicate-code
	4	+# - duplicate-code: the CPU reference deliberately mirrors the operator formula
	5	+
	6	+import torch
	7	+
	8	+from atk.configs.dataset_config import InputDataset
	9	+from atk.tasks.api_execute import register
	10	+from atk.tasks.api_execute.base_api import BaseApi
	11	+from atk.tasks.api_execute.triton_base_api import TritonBaseApi
	12	+
	13	+from mindspeed_ops.api.triton.bwd_segsum_dadt import bwd_segsum_dadt
	14	+from mindspeed_ops.utils import is_arch35
	15	+
	16	+
	17	+def _bwd_segsum_dadt_reference(
	18	+ dSSdA: torch.Tensor,
	19	+ SSdA: torch.Tensor,
	20	+ ddt_out: torch.Tensor,
	21	+ chunk_size: int,
	22	+) -> torch.Tensor:
	23	+ """CPU golden reference implementation for bwd_segsum_dadt."""
	24	+ B, H, nchunks, C, _ = dSSdA.shape
	25	+
	26	+ ddt_out_ref = ddt_out.clone()
	27	+
	28	+ for b in range(B):
	29	+ for h in range(H):
	30	+ for c in range(nchunks):
	31	+ chunk_start = c * chunk_size
	32	+ chunk_end = chunk_start + chunk_size
	33	+
	34	+ dSSdA_chunk = dSSdA[b, h, c].T
	35	+ SSdA_chunk = SSdA[b, h, c]
	36	+
	37	+ dSSdA_block = dSSdA_chunk * torch.exp(SSdA_chunk)
	38	+ dSSdA_block = torch.cumsum(dSSdA_block.flip(dims=[0]), dim=0).flip(dims=[0])
	39	+
	40	+ offs_i = torch.arange(C)[:, None]
	41	+ offs_j = torch.arange(C)[None, :]
	42	+ SS_mask = offs_i > offs_j
	43	+ dSSdA_block = torch.where(SS_mask, dSSdA_block, torch.tensor(0.0, dtype=dSSdA.dtype))
	44	+
	45	+ ddt_out_ref[b, h, chunk_start:chunk_end] += dSSdA_block.sum(dim=1)
	46	+
	47	+ return ddt_out_ref
	48	+
	49	+
	50	+def _get_inputs(input_data: InputDataset) -> tuple[torch.Tensor, ...]:
	51	+ keys = ("dSSdA", "SSdA", "ddt_out", "chunk_size")
	52	+ if input_data.kwargs:
	53	+ return tuple(input_data.kwargs[key] for key in keys)
	54	+ return tuple(input_data.args[: len(keys)])
	55	+
	56	+
	57	+@register("torch_bwd_segsum_dadt")
	58	+class TorchBwdSegsumDadtApi(BaseApi):
	59	+ def __call__(self, input_data: InputDataset, with_output: bool = False):
	60	+ dSSdA, SSdA, ddt_out, chunk_size = _get_inputs(input_data)
	61	+ chunk_size = int(chunk_size)
	62	+ output = _bwd_segsum_dadt_reference(dSSdA.float(), SSdA.float(), ddt_out.float(), chunk_size)
	63	+ return output.to(ddt_out.dtype)
	64	+
	65	+
	66	+@register("triton_bwd_segsum_dadt")
	67	+class TritonBwdSegsumDadtApi(TritonBaseApi):
	68	+ def __call__(self, input_data: InputDataset, with_output: bool = False):
	69	+ if is_arch35():
	70	+ raise NotImplementedError("bwd_segsum_dadt is not supported on arch35")
	71	+ dSSdA, SSdA, ddt_out, chunk_size = _get_inputs(input_data)
	72	+ chunk_size = int(chunk_size)
	73	+ output = bwd_segsum_dadt(dSSdA, SSdA, ddt_out, chunk_size)
	74	+ return output

Atests/unit_tests/triton/test_bwd_segsum_dadt.py+105-0

@@ -0,0 +1,105 @@
	1	+# Copyright (c) 2026, HUAWEI CORPORATION. All rights reserved.
	2	+
	3	+import pytest
	4	+import torch
	5	+
	6	+from mindspeed_ops.api.triton.bwd_segsum_dadt import bwd_segsum_dadt
	7	+from mindspeed_ops.api.triton.utils import get_available_device
	8	+from tests.utils import print_diff, assert_close
	9	+
	10	+
	11	+def _dtype_thresholds(dtype):
	12	+ """Per-dtype (ratio, atol) for assert_close. Empirical from kernel
	13	+ measurements: fp32 near-bitwise, bf16 ~3-4e-3, fp16 ~5e-4.
	14	+ """
	15	+ if dtype == torch.float32:
	16	+ return 1e-4, 1e-5
	17	+ if dtype == torch.bfloat16:
	18	+ return 1e-2, 5e-3
	19	+ if dtype == torch.float16:
	20	+ return 1e-3, 5e-4
	21	+ raise ValueError(f"unsupported dtype {dtype}")
	22	+
	23	+
	24	+TEST_SHAPES = [
	25	+ (1, 1, 1, 8, 8), # B=1, H=1, nchunks=1, C=8, small chunk
	26	+ (1, 1, 1, 16, 16), # B=1, H=1, nchunks=1, C=16
	27	+ (1, 2, 2, 16, 16), # B=1, H=2, nchunks=2, C=16
	28	+ (2, 4, 4, 32, 32), # B=2, H=4, nchunks=4, C=32, larger batch/heads
	29	+ (1, 1, 2, 64, 64), # B=1, H=1, nchunks=2, C=64, multiple chunks
	30	+ (1, 8, 1, 32, 32), # B=1, H=8, nchunks=1, C=32, many heads
	31	+]
	32	+
	33	+DTYPES = [torch.float32, torch.bfloat16, torch.float16]
	34	+
	35	+
	36	+def _case_id(shape, dtype):
	37	+ B, H, nchunks, C, _ = shape
	38	+ dtype_name = {torch.float32: "fp32", torch.bfloat16: "bf16", torch.float16: "fp16"}
	39	+ return f"B{B}_H{H}_N{nchunks}_C{C}_{dtype_name[dtype]}"
	40	+
	41	+
	42	+TEST_CASES = [(shape, dtype) for shape in TEST_SHAPES for dtype in DTYPES]
	43	+
	44	+
	45	+def cpu_golden(dSSdA, SSdA, chunk_size):
	L L LinShua7月25日当前PR中的torch小算子由于是自己实现的，需要证明下torch小算子实现的正确性，建议在GPU上分别输入相同的输入对比开源triton和torch小算子输出结果，提供相关数据来证明torch小算子实现的正确性。 Tsuki 7月27日评论： 7月27日评论：
	46	+ B, H, nchunks, C, _ = dSSdA.shape
	47	+
	48	+ ddt_out = torch.zeros(B, H, nchunks * C, dtype=dSSdA.dtype, device=dSSdA.device)
	49	+
	50	+ for b in range(B):
	51	+ for h in range(H):
	52	+ for c in range(nchunks):
	53	+ chunk_start = c * chunk_size
	54	+ chunk_end = chunk_start + chunk_size
	55	+
	56	+ dSSdA_chunk = dSSdA[b, h, c].T
	57	+ SSdA_chunk = SSdA[b, h, c]
	58	+
	59	+ dSSdA_block = dSSdA_chunk * torch.exp(SSdA_chunk)
	60	+ dSSdA_block = torch.cumsum(dSSdA_block.flip(dims=[0]), dim=0).flip(dims=[0])
	61	+
	62	+ offs_i = torch.arange(C)[:, None]
	63	+ offs_j = torch.arange(C)[None, :]
	64	+ SS_mask = offs_i > offs_j
	65	+ dSSdA_block = torch.where(SS_mask, dSSdA_block, torch.tensor(0.0, dtype=dSSdA.dtype))
	66	+
	67	+ ddt_out[b, h, chunk_start:chunk_end] += dSSdA_block.sum(dim=1)
	68	+
	69	+ return ddt_out
	70	+
	71	+
	72	+class TestBwdSegsumDadtOperator:
	73	+ def setup_method(self):
	74	+ torch.manual_seed(42)
	75	+
	76	+ @pytest.mark.parametrize(
	77	+ ("shape", "dtype"),
	78	+ [pytest.param(test, id=_case_id(test)) for test in TEST_CASES],
	79	+ )
	80	+ def test_bwd_segsum_dadt_npu(
	81	+ self,
	82	+ shape,
	83	+ dtype: torch.dtype,
	84	+ ):
	85	+ B, H, nchunks, C, _ = shape
	86	+ chunk_size = C
	87	+ S = nchunks * C
	88	+
	89	+ dSSdA = torch.randn(shape, dtype=dtype)
	90	+ SSdA = torch.randn(shape, dtype=dtype)
	91	+
	92	+ device = get_available_device()
	93	+
	94	+ out_golden = cpu_golden(dSSdA.float(), SSdA.float(), chunk_size).to(device).to(dtype)
	95	+
	96	+ dSSdA_npu = dSSdA.to(device)
	97	+ SSdA_npu = SSdA.to(device)
	98	+ ddt_out_npu = torch.zeros(B, H, S, dtype=dtype, device=device)
	99	+
	100	+ bwd_segsum_dadt(dSSdA_npu, SSdA_npu, ddt_out_npu, chunk_size)
	101	+
	102	+ ratio, atol = _dtype_thresholds(dtype)
	103	+ print_diff("cpu_and_triton", out_golden, ddt_out_npu, atol)
	104	+
	105	+ assert_close("cpu_and_triton", out_golden, ddt_out_npu, ratio, err_atol=atol)