[Ops] Add Triton-Ascend compatibility for GDN-2

.github/workflows/ascend-a2-benchmark-ci.yml

@@ -160,7 +160,7 @@ jobs:
           print("matplotlib savefig smoke test: ok")
           PY
 
-      - name: Benchmark chunk_gdn, chunk_kda, fused_attnres, and conv1d
+      - name: Benchmark chunk_gdn, chunk_gdn2, chunk_kda, fused_attnres, and conv1d
         shell: bash
         env:
           FLA_BENCH_WARMUP_MS: "25"
@@ -170,9 +170,45 @@ jobs:
           source /usr/local/Ascend/ascend-toolkit/set_env.sh
           set -uo pipefail
 
-          ASCEND_RT_VISIBLE_DEVICES=0 python -m benchmarks.ops.run \
-            --op chunk_gdn --base '' --json benchmark_chunk_gdn.json \
-            > chunk_gdn.log 2>&1 &
+          (
+            set -e
+            ASCEND_RT_VISIBLE_DEVICES=0 python -m benchmarks.ops.run \
+              --op chunk_gdn --base '' --json benchmark_chunk_gdn.json \
+              > chunk_gdn.log 2>&1
+            ASCEND_RT_VISIBLE_DEVICES=0 python -m benchmarks.ops.run \
+              --op chunk_gdn2 --base '' \
+              --json benchmark_chunk_gdn2.json \
+              > chunk_gdn2.log 2>&1
+            ASCEND_RT_VISIBLE_DEVICES=0 python -m benchmarks.ops.run \
+              --op chunk_gdn2 --base '' \
+              --custom-shapes '{"gdn2_32k":{"B":1,"T":32768,"H":1,"D":32}}' \
+              --json benchmark_chunk_gdn2_32k.json \
+              > chunk_gdn2_32k.log 2>&1
+          python - <<'PY'
+          import json
+
+          from benchmarks.ops.registry import SHAPE_CONFIGS
+
+          def verify(path, shapes):
+              with open(path) as f:
+                  results = [r for r in json.load(f)['results'] if r['op'] == 'chunk_gdn2']
+              actual = {(r['mode'], r['B'], r['T'], r['H'], r['D']) for r in results}
+              expected = {
+                  (mode, shape['B'], shape['T'], shape['H'], shape['D'])
+                  for shape in shapes.values()
+                  for mode in ('fwd', 'fwdbwd')
+              }
+              assert len(results) == len(expected) and actual == expected, {
+                  'missing': sorted(expected - actual),
+                  'unexpected': sorted(actual - expected),
+              }
+
+          verify('benchmark_chunk_gdn2.json', SHAPE_CONFIGS)
+          verify('benchmark_chunk_gdn2_32k.json', {
+              'gdn2_32k': {'B': 1, 'T': 32768, 'H': 1, 'D': 32},
+          })
+          PY
+          ) &
           pid_gdn=$!
           ASCEND_RT_VISIBLE_DEVICES=1 python -m benchmarks.ops.run \
             --op chunk_kda --base '' --json benchmark_chunk_kda.json \
@@ -194,6 +230,10 @@ jobs:
 
           echo "========== chunk_gdn =========="
           cat chunk_gdn.log
+          echo "========== chunk_gdn2 (default shape matrix) =========="
+          cat chunk_gdn2.log
+          echo "========== chunk_gdn2 (B1 T32768 H1 K32 V32) =========="
+          cat chunk_gdn2_32k.log
           echo "========== chunk_kda =========="
           cat chunk_kda.log
           echo "========== fused_attnres =========="
@@ -209,6 +249,8 @@ jobs:
           name: ascend-a2-benchmark-results
           path: |
             benchmark_chunk_gdn.json
+            benchmark_chunk_gdn2.json
+            benchmark_chunk_gdn2_32k.json
             benchmark_chunk_kda.json
             benchmark_fused_attnres.json
             conv_benchmark/

.github/workflows/ascend-a2-ci.yml

@@ -142,7 +142,7 @@ jobs:
           assert IS_NPU, "Expected NPU backend (IS_NPU=True)"
           PY
 
-      - name: Run tests/modules and tests/ops/utils, gdn, kda, gla, attnres, solve_tril
+      - name: Run tests/modules and tests/ops/utils, gdn, gdn2, kda, gla, attnres, solve_tril
         shell: bash
         env:
           FLA_NPU_XDIST: "1"
@@ -153,6 +153,7 @@ jobs:
             tests/ops/utils \
             tests/ops/test_gdn_kernels.py \
             tests/ops/test_gdn.py \
+            tests/ops/test_gdn2.py \
             tests/ops/test_kda.py \
             tests/ops/test_gla.py \
             tests/ops/test_attnres.py \

benchmarks/ops/registry.py

@@ -318,7 +318,7 @@ def generate_inputs(
     test_file='tests/ops/test_delta.py',
 ))
 
-# --- +gate + beta ---
+# --- Delta-rule variants with decay and update gates ---
 
 register_op(OpConfig(
     name='chunk_gdn',
@@ -334,6 +334,20 @@ def generate_inputs(
     test_file='tests/ops/test_gdn.py',
 ))
 
+register_op(OpConfig(
+    name='chunk_gdn2',
+    import_path='fla.ops.gdn2',
+    inputs={
+        **_simple_qkv,
+        'g': TensorSpec(shape_BTHD, transform=logsigmoid),
+        'b': TensorSpec(shape_BTHD, transform=sigmoid_transform),
+        'w': TensorSpec(shape_BTHD, transform=sigmoid_transform),
+    },
+    extra_kwargs={'use_qk_l2norm_in_kernel': True},
+    category='gate_beta',
+    test_file='tests/ops/test_gdn2.py',
+))
+
 register_op(OpConfig(
     name='chunk_kda',
     import_path='fla.ops.kda',

benchmarks/ops/verify.py

@@ -28,6 +28,9 @@
     # Gate (full pytest) + benchmark vs. main
     python -m benchmarks.ops.verify --op chunk_gla --base main
 
+    # Gate + benchmark for GDN-2
+    python -m benchmarks.ops.verify --op chunk_gdn2 --base main
+
     # Fast signal: gate on a shape subset (pytest -k selection, test unchanged)
     python -m benchmarks.ops.verify --op chunk_gla --gate-k T15 --modes fwd
 

fla/ops/gdn2/backends/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+"""GDN-2 backends."""
+
+from fla.ops.backends import BackendRegistry, dispatch
+from fla.ops.gdn2.backends.triton_ascend import TritonAscendGDN2Backend
+
+gdn2_registry = BackendRegistry("gdn2")
+gdn2_registry.register(TritonAscendGDN2Backend())
+
+
+__all__ = ['dispatch', 'gdn2_registry']

fla/ops/gdn2/backends/triton_ascend/__init__.py

@@ -0,0 +1,159 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+"""Triton-Ascend backend for GDN-2."""
+
+from __future__ import annotations
+
+import torch
+
+from fla.ops.backends import BaseBackend
+
+
+class TritonAscendGDN2Backend(BaseBackend):
+    """Ascend NPU backend for GDN-2 chunk kernels."""
+
+    backend_type = "triton_ascend"
+    package_name = None
+    env_var = None
+    priority = 0
+
+    @classmethod
+    def is_available(cls) -> bool:
+        from fla.utils import IS_NPU
+        return IS_NPU
+
+    def chunk_gdn2_fwd_intra_verifier(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        gk: torch.Tensor,
+        b: torch.Tensor,
+        w_gate: torch.Tensor,
+        scale: float,
+        cu_seqlens: torch.LongTensor | None = None,
+        chunk_size: int = 64,
+        chunk_indices: torch.LongTensor | None = None,
+        safe_gate: bool = False,
+        disable_recompute: bool = False,
+    ) -> tuple[bool, str | None]:
+        del scale, safe_gate, disable_recompute
+        if chunk_size != 64:
+            return False, f"GDN-2 Ascend intra requires chunk_size=64, got {chunk_size}"
+        float_tensors = (q, k, v, gk, b, w_gate)
+        tensors = (*float_tensors, *(t for t in (cu_seqlens, chunk_indices) if t is not None))
+        if any(t.device.type != "npu" for t in tensors):
+            return False, "GDN-2 Ascend intra requires NPU tensors"
+        supported = (torch.float16, torch.bfloat16, torch.float32)
+        if any(t.dtype not in supported for t in float_tensors):
+            return False, "GDN-2 Ascend intra received an unsupported dtype"
+        return True, None
+
+    def chunk_gdn2_fwd_intra(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        gk: torch.Tensor,
+        b: torch.Tensor,
+        w_gate: torch.Tensor,
+        scale: float,
+        cu_seqlens: torch.LongTensor | None = None,
+        chunk_size: int = 64,
+        chunk_indices: torch.LongTensor | None = None,
+        safe_gate: bool = False,
+        disable_recompute: bool = False,
+    ):
+        from fla.ops.gdn2.backends.triton_ascend.chunk_intra import chunk_gdn2_fwd_intra_npu
+        return chunk_gdn2_fwd_intra_npu(
+            q,
+            k,
+            v,
+            gk,
+            b,
+            w_gate,
+            scale,
+            cu_seqlens,
+            chunk_size,
+            chunk_indices,
+            safe_gate,
+            disable_recompute,
+        )
+
+    def chunk_gdn2_bwd_wy_dqkg_fused_verifier(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        v_new: torch.Tensor,
+        g: torch.Tensor,
+        b: torch.Tensor,
+        w_gate: torch.Tensor,
+        A: torch.Tensor,
+        h: torch.Tensor,
+        do: torch.Tensor,
+        dh: torch.Tensor,
+        dv: torch.Tensor,
+        scale: float | None = None,
+        cu_seqlens: torch.LongTensor | None = None,
+        chunk_size: int = 64,
+        chunk_indices: torch.LongTensor | None = None,
+        state_v_first: bool = False,
+    ) -> tuple[bool, str | None]:
+        del scale, state_v_first
+        if chunk_size != 64:
+            return False, f"GDN-2 Ascend backward requires chunk_size=64, got {chunk_size}"
+        float_tensors = (q, k, v, v_new, g, b, w_gate, A, h, do, dh, dv)
+        tensors = (*float_tensors, *(t for t in (cu_seqlens, chunk_indices) if t is not None))
+        if any(t.device.type != "npu" for t in tensors):
+            return False, "GDN-2 Ascend backward requires NPU tensors"
+        supported = (torch.float16, torch.bfloat16, torch.float32)
+        if any(t.dtype not in supported for t in float_tensors):
+            return False, "GDN-2 Ascend backward received an unsupported dtype"
+        return True, None
+
+    def chunk_gdn2_bwd_wy_dqkg_fused(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        v_new: torch.Tensor,
+        g: torch.Tensor,
+        b: torch.Tensor,
+        w_gate: torch.Tensor,
+        A: torch.Tensor,
+        h: torch.Tensor,
+        do: torch.Tensor,
+        dh: torch.Tensor,
+        dv: torch.Tensor,
+        scale: float | None = None,
+        cu_seqlens: torch.LongTensor | None = None,
+        chunk_size: int = 64,
+        chunk_indices: torch.LongTensor | None = None,
+        state_v_first: bool = False,
+    ):
+        from fla.ops.gdn2.backends.triton_ascend.chunk_bwd import chunk_gdn2_bwd_wy_dqkg_fused_npu
+        return chunk_gdn2_bwd_wy_dqkg_fused_npu(
+            q,
+            k,
+            v,
+            v_new,
+            g,
+            b,
+            w_gate,
+            A,
+            h,
+            do,
+            dh,
+            dv,
+            scale,
+            cu_seqlens,
+            chunk_size,
+            chunk_indices,
+            state_v_first,
+        )