diff --git a/.github/workflows/ascend-a2-benchmark-ci.yml b/.github/workflows/ascend-a2-benchmark-ci.yml index 8586f346f0..0316707991 100644 --- a/.github/workflows/ascend-a2-benchmark-ci.yml +++ b/.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/ diff --git a/.github/workflows/ascend-a2-ci.yml b/.github/workflows/ascend-a2-ci.yml index 563fa84f72..5efa6782ed 100644 --- a/.github/workflows/ascend-a2-ci.yml +++ b/.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 \ diff --git a/benchmarks/ops/registry.py b/benchmarks/ops/registry.py index 34e267aa16..85b3374700 100644 --- a/benchmarks/ops/registry.py +++ b/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', diff --git a/benchmarks/ops/verify.py b/benchmarks/ops/verify.py index 869054dbc7..68679d78f4 100644 --- a/benchmarks/ops/verify.py +++ b/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 diff --git a/fla/ops/gdn2/backends/__init__.py b/fla/ops/gdn2/backends/__init__.py new file mode 100644 index 0000000000..3ebfa1d090 --- /dev/null +++ b/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'] diff --git a/fla/ops/gdn2/backends/triton_ascend/__init__.py b/fla/ops/gdn2/backends/triton_ascend/__init__.py new file mode 100644 index 0000000000..97c9bb5189 --- /dev/null +++ b/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, + ) diff --git a/fla/ops/gdn2/backends/triton_ascend/chunk_bwd.py b/fla/ops/gdn2/backends/triton_ascend/chunk_bwd.py new file mode 100644 index 0000000000..ba40c778c9 --- /dev/null +++ b/fla/ops/gdn2/backends/triton_ascend/chunk_bwd.py @@ -0,0 +1,545 @@ +# 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 chunk backward kernels for triton-ascend.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl +from triton.runtime import driver + +from fla.ops.kda.backends.triton_ascend.chunk_bwd import chunk_kda_bwd_kernel_wy_k_part_npu +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.op import exp2 +from fla.utils import input_guard +from fla.utils.ascend_ub_manager import compute_row_tile_block_size + +_BC = 16 +_BWD_MEM_MULT = 10.0 +_SAFETY_MARGIN = 0.80 +_FALLBACK_TILE = 16 +_MAX_TILE = 128 + + +def _get_tile(size: int) -> int: + return compute_row_tile_block_size( + _BC, + size, + _BWD_MEM_MULT, + tiling_row=False, + safety_margin=_SAFETY_MARGIN, + fallback=_FALLBACK_TILE, + min_block=16, + max_block=min(_MAX_TILE, triton.next_power_of_2(size)), + ) + + +def _t_contig_arg(x: torch.Tensor, num_heads: int) -> tuple[torch.Tensor, bool]: + if num_heads == 1: + return x, False + return x.transpose(1, 2).contiguous(), True + + +def _get_npu_properties(): + device = torch.npu.current_device() + return driver.active.utils.get_device_properties(device) + + +def _launch_dA_finalize( + kernel, + *, + nt: int, + bh_total: int, + T: int, + BT: int, + is_varlen: bool, + num_core: int, + kernel_kwargs: dict, +) -> None: + kwargs = dict(kernel_kwargs) + kwargs['num_core'] = num_core + if is_varlen: + kwargs['TAIL_MODE'] = 1 + kwargs['NT_OFFSET'] = 0 + kwargs['task_num'] = nt * bh_total + kernel[(num_core,)](**kwargs) + return + + n_bulk = nt if T % BT == 0 else max(nt - 1, 0) + if n_bulk > 0: + kwargs['TAIL_MODE'] = 0 + kwargs['NT_OFFSET'] = 0 + kwargs['task_num'] = n_bulk * bh_total + kernel[(num_core,)](**kwargs) + if T % BT != 0 and nt > 0: + kwargs['TAIL_MODE'] = 1 + kwargs['NT_OFFSET'] = n_bulk + kwargs['task_num'] = bh_total + kernel[(num_core,)](**kwargs) + + +@triton.jit(do_not_specialize=['T', 'task_num', 'num_core', 'BH']) +def chunk_gdn2_bwd_kernel_wy_v_part_npu( + v, + w_gate, + A, + dv, + dv2, + dw, + dA_acc, + cu_seqlens, + chunk_indices, + T, + BH, + task_num, + num_core, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, + T_CONTIG: tl.constexpr, +): + core_id = tl.program_id(0) + T_seq = T + + for task_id in tl.range(core_id, task_num, num_core): + i_t = task_id // BH + i_bh = task_id % BH + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1, + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + else: + bos = tl.cast(i_b, tl.int64) * T + + if T_CONTIG: + if IS_VARLEN: + head_off = tl.cast(i_h, tl.int64) * T_seq + v_ptr = v + (head_off + bos) * V + w_ptr = w_gate + (head_off + bos) * V + dv_ptr = dv + (head_off + bos) * V + A_ptr = A + (head_off + bos) * BT + else: + head_off = (tl.cast(i_b, tl.int64) * H + i_h) * T_seq + v_ptr = v + head_off * V + w_ptr = w_gate + head_off * V + dv_ptr = dv + head_off * V + A_ptr = A + head_off * BT + value_stride_t = V + a_stride_t = BT + else: + v_ptr = v + (bos * H + i_h) * V + w_ptr = w_gate + (bos * H + i_h) * V + dv_ptr = dv + (bos * H + i_h) * V + A_ptr = A + (bos * H + i_h) * BT + value_stride_t = H * V + a_stride_t = H * BT + + dv2_ptr = dv2 + (bos * H + i_h) * V + dw_ptr = dw + (bos * H + i_h) * V + dA_ptr = dA_acc + (bos * H + i_h) * BT + p_A = tl.make_block_ptr(A_ptr, (BT, T), (1, a_stride_t), (0, i_t * BT), (BT, BT), (0, 1)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v_ptr, (T, V), (value_stride_t, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_w = tl.make_block_ptr(w_ptr, (T, V), (value_stride_t, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv_ptr, (T, V), (value_stride_t, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_dv = tl.load(p_dv, boundary_check=(0, 1)) + b_dv_c = b_dv + 0.0 + b_dA += tl.dot(b_dv, tl.trans(b_v * b_w), allow_tf32=False) + b_A_c = b_A + 0.0 + b_dvb = tl.dot(b_A_c, b_dv_c, allow_tf32=False) + + p_dv2 = tl.make_block_ptr(dv2_ptr, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dw = tl.make_block_ptr(dw_ptr, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + tl.store(p_dv2, (b_dvb * b_w).to(p_dv2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dw, (b_dvb * b_v).to(p_dw.dtype.element_ty), boundary_check=(0, 1)) + + p_dA = tl.make_block_ptr(dA_ptr, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T', 'task_num', 'num_core', 'BH']) +def chunk_gdn2_bwd_kernel_wy_gate_part_npu( + k, + g, + b, + A, + h, + dv, + dA_acc, + db, + dg, + dk, + cu_seqlens, + chunk_indices, + chunk_offsets, + T, + BH, + task_num, + num_core, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STATE_V_FIRST: tl.constexpr, + IS_VARLEN: tl.constexpr, + K_T_CONTIG: tl.constexpr, + G_T_CONTIG: tl.constexpr, + K_OFFSET: tl.constexpr, +): + i_k = K_OFFSET + core_id = tl.program_id(0) + T_seq = T + + for task_id in tl.range(core_id, task_num, num_core): + i_t = task_id // BH + i_bh = task_id % BH + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1, + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + i_tg = tl.load(chunk_offsets + i_n).to(tl.int64) + tl.cast(i_t, tl.int64) + else: + i_tg = tl.cast(i_b, tl.int64) * tl.cdiv(T, BT) + i_t + bos = tl.cast(i_b, tl.int64) * T + + if K_T_CONTIG: + if IS_VARLEN: + k_ptr = k + (tl.cast(i_h, tl.int64) * T_seq + bos) * K + else: + k_ptr = k + (tl.cast(i_b, tl.int64) * H + i_h) * T_seq * K + k_stride_t = K + else: + k_ptr = k + (bos * H + i_h) * K + k_stride_t = H * K + + if G_T_CONTIG: + if IS_VARLEN: + head_off = tl.cast(i_h, tl.int64) * T_seq + bos + else: + head_off = (tl.cast(i_b, tl.int64) * H + i_h) * T_seq + g_ptr = g + head_off * K + b_ptr = b + head_off * K + A_ptr = A + head_off * BT + dv_ptr = dv + head_off * V + g_stride_t = K + a_stride_t = BT + dv_stride_t = V + else: + g_ptr = g + (bos * H + i_h) * K + b_ptr = b + (bos * H + i_h) * K + A_ptr = A + (bos * H + i_h) * BT + dv_ptr = dv + (bos * H + i_h) * V + g_stride_t = H * K + a_stride_t = H * BT + dv_stride_t = H * V + + h_ptr = h + (i_tg * H + i_h) * K * V + dA_ptr = dA_acc + (bos * H + i_h) * BT + db_ptr = db + (bos * H + i_h) * K + dg_ptr = dg + (bos * H + i_h) * K + dk_ptr = dk + (bos * H + i_h) * K + + b_dw = tl.zeros([BT, BK], dtype=tl.float32) + for i_v in range(tl.cdiv(V, BV)): + p_dv = tl.make_block_ptr(dv_ptr, (T, V), (dv_stride_t, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + if STATE_V_FIRST: + p_h = tl.make_block_ptr(h_ptr, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + else: + p_h = tl.make_block_ptr(h_ptr, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + b_dv = tl.load(p_dv, boundary_check=(0, 1)) + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dw += tl.dot(b_dv, b_h.to(b_dv.dtype), allow_tf32=False) + + p_k = tl.make_block_ptr(k_ptr, (T, K), (k_stride_t, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_g = tl.make_block_ptr(g_ptr, (T, K), (g_stride_t, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_b = tl.make_block_ptr(b_ptr, (T, K), (g_stride_t, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_A = tl.make_block_ptr(A_ptr, (BT, T), (1, a_stride_t), (0, i_t * BT), (BT, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_b = tl.load(p_b, boundary_check=(0, 1)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_gk_exp = exp2(b_g) + b_kg = b_k * b_gk_exp + b_dw = -b_dw.to(b_A.dtype) + b_dkgb = tl.dot(b_A, b_dw, allow_tf32=False) + + p_dA = tl.make_block_ptr(dA_ptr, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_dA = tl.load(p_dA, boundary_check=(0, 1)).to(tl.float32) + b_dw_c = b_dw + 0.0 + b_dA += tl.dot(b_dw_c, tl.trans((b_kg * b_b).to(b_A.dtype)), allow_tf32=False) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + + p_db = tl.make_block_ptr(db_ptr, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_db, (b_dkgb * b_kg).to(p_db.dtype.element_ty), boundary_check=(0, 1)) + + p_dk = tl.make_block_ptr(dk_ptr, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_dk = tl.load(p_dk, boundary_check=(0, 1)).to(tl.float32) + b_dk += b_dkgb * b_gk_exp * b_b + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + p_dg = tl.make_block_ptr(dg_ptr, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_dg = tl.load(p_dg, boundary_check=(0, 1)).to(tl.float32) + b_dg += b_kg * b_dkgb * b_b + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T', 'task_num', 'num_core', 'BH', 'NT_OFFSET']) +def chunk_gdn2_bwd_kernel_wy_dA_finalize_npu( + A, + dA_acc, + dA, + cu_seqlens, + chunk_indices, + T, + BH, + task_num, + num_core, + NT_OFFSET, + H: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, + A_T_CONTIG: tl.constexpr, + TAIL_MODE: tl.constexpr, +): + core_id = tl.program_id(0) + T_seq = T + + for task_id in tl.range(core_id, task_num, num_core): + i_t = NT_OFFSET + task_id // BH + i_bh = task_id % BH + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1, + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + else: + bos = tl.cast(i_b, tl.int64) * T + + if A_T_CONTIG: + if IS_VARLEN: + A_ptr = A + (tl.cast(i_h, tl.int64) * T_seq + bos) * BT + else: + A_ptr = A + (tl.cast(i_b, tl.int64) * H + i_h) * T_seq * BT + a_stride_t = BT + else: + A_ptr = A + (bos * H + i_h) * BT + a_stride_t = H * BT + + dA_acc_ptr = dA_acc + (bos * H + i_h) * BT + dA_ptr = dA + (bos * H + i_h) * BT + p_A = tl.make_block_ptr(A_ptr, (BT, T), (1, a_stride_t), (0, i_t * BT), (BT, BT), (0, 1)) + p_dA_acc = tl.make_block_ptr(dA_acc_ptr, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_dA = tl.make_block_ptr(dA_ptr, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + + o_t = i_t * BT + tl.arange(0, BT) + if TAIL_MODE == 0: + b_A = tl.load(p_A) + b_dA = tl.load(p_dA_acc).to(tl.float32) + m_A = o_t[:, None] > o_t[None, :] + else: + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_dA = tl.load(p_dA_acc, boundary_check=(0, 1)).to(tl.float32) + m_t = o_t < T + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t[None, :]) + + b_dA = tl.where(m_A, b_dA, 0) + b_mid = tl.dot(b_dA.to(b_A.dtype), b_A, allow_tf32=False) + b_fin = tl.dot(b_A, b_mid.to(b_A.dtype), allow_tf32=False) + b_fin = tl.where(m_A, -b_fin, 0) + if TAIL_MODE == 0: + tl.store(p_dA, b_fin.to(p_dA.dtype.element_ty)) + else: + tl.store(p_dA, b_fin.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + + +@input_guard +def chunk_gdn2_bwd_wy_dqkg_fused_npu( + 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, +): + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + if BT % _BC != 0: + raise ValueError(f'GDN-2 Ascend bwd requires chunk_size % {_BC} == 0, got {BT}') + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dq = g.new_empty(B, T, H, K, dtype=torch.float) + dk = g.new_empty(B, T, H, K, dtype=torch.float) + dv2 = torch.empty_like(v) + dg = torch.empty_like(g, dtype=torch.float) + db = torch.empty_like(b, dtype=torch.float) + dw = torch.empty_like(w_gate, dtype=torch.float) + dA = torch.empty_like(A, dtype=torch.float) + dA_acc = torch.zeros(B, T, H, BT, dtype=torch.float, device=A.device) + + BK = _get_tile(K) + BV = _get_tile(V) + NK = triton.cdiv(K, BK) + is_varlen = cu_seqlens is not None + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) if is_varlen else g.new_zeros(1, dtype=torch.int64) + bh_total = B * H + task_num = NT * bh_total + num_core = _get_npu_properties()['num_vectorcore'] + + v_arg, t_contig = _t_contig_arg(v, H) + w_arg = _t_contig_arg(w_gate, H)[0] + A_arg = _t_contig_arg(A, H)[0] + dv_arg = _t_contig_arg(dv, H)[0] + chunk_gdn2_bwd_kernel_wy_v_part_npu[(num_core,)]( + v=v_arg, + w_gate=w_arg, + A=A_arg, + dv=dv_arg, + dv2=dv2, + dw=dw, + dA_acc=dA_acc, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + BH=bh_total, + task_num=task_num, + num_core=num_core, + H=H, + V=V, + BT=BT, + BV=BV, + IS_VARLEN=is_varlen, + T_CONTIG=t_contig, + ) + + k_part_kwargs = dict( + q=q, + k=k, + v_new=v_new, + g=g, + h=h, + do=do, + dh=dh, + dq=dq, + dk=dk, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + scale=scale, + T=T, + BH=bh_total, + task_num=task_num, + num_core=num_core, + H=H, + HV=H, + K=K, + V=V, + BT=BT, + BC=32 if BT >= 32 else _BC, + BK=BK, + BV=BV, + STATE_V_FIRST=state_v_first, + IS_VARLEN=is_varlen, + ) + for k_off in range(NK): + k_part_kwargs['K_OFFSET'] = k_off + chunk_kda_bwd_kernel_wy_k_part_npu[(num_core,)](**k_part_kwargs) + + k_arg, k_t_contig = _t_contig_arg(k, H) + g_arg, g_t_contig = _t_contig_arg(g, H) + b_arg = _t_contig_arg(b, H)[0] + gate_kwargs = dict( + k=k_arg, + g=g_arg, + b=b_arg, + A=A_arg, + h=h, + dv=dv_arg, + dA_acc=dA_acc, + db=db, + dg=dg, + dk=dk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + T=T, + BH=bh_total, + task_num=task_num, + num_core=num_core, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + STATE_V_FIRST=state_v_first, + IS_VARLEN=is_varlen, + K_T_CONTIG=k_t_contig, + G_T_CONTIG=g_t_contig, + ) + for k_off in range(NK): + gate_kwargs['K_OFFSET'] = k_off + chunk_gdn2_bwd_kernel_wy_gate_part_npu[(num_core,)](**gate_kwargs) + + _launch_dA_finalize( + chunk_gdn2_bwd_kernel_wy_dA_finalize_npu, + nt=NT, + bh_total=bh_total, + T=T, + BT=BT, + is_varlen=is_varlen, + num_core=num_core, + kernel_kwargs=dict( + A=A_arg, + dA_acc=dA_acc, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + BH=bh_total, + H=H, + BT=BT, + IS_VARLEN=is_varlen, + A_T_CONTIG=t_contig, + ), + ) + return dq, dk, dv2, db, dw, dg, dA diff --git a/fla/ops/gdn2/backends/triton_ascend/chunk_intra.py b/fla/ops/gdn2/backends/triton_ascend/chunk_intra.py new file mode 100644 index 0000000000..ac31c14a83 --- /dev/null +++ b/fla/ops/gdn2/backends/triton_ascend/chunk_intra.py @@ -0,0 +1,628 @@ +# 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 forward intra kernels for triton-ascend.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from fla.ops.gdn2.wy_fast import recompute_w_u_fwd_gdn2 +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp2 +from fla.utils import ascend_compile_kwargs, input_guard +from fla.utils.ascend_ub_manager import ASCEND_MAX_GRID_DIM, compute_row_tile_block_size, max_grid_axis_chunks + +_BC = 16 +_TOKEN_GROUP = 8 +_INTER_MEM_MULT = 18.0 +_SAFETY_MARGIN = 0.80 +_FALLBACK_BK = 16 +_MAX_INTER_BK = 64 +_LAUNCH_BLOCK_BUDGET = 4096 +# Disable auto-multi-buffer and AutoBlockify on the inter kernels for CANN 9.1. +_INTER_COMPILE_KWARGS = ascend_compile_kwargs(blacklist_auto_blockify=True) + + +def _get_inter_bk(K: int) -> int: + return compute_row_tile_block_size( + _BC, + K, + _INTER_MEM_MULT, + tiling_row=False, + safety_margin=_SAFETY_MARGIN, + fallback=_FALLBACK_BK, + min_block=16, + max_block=min(_MAX_INTER_BK, triton.next_power_of_2(K)), + ) + + +def _launch_diag_kernel( + kernel, + *, + nt: int, + nc: int, + bh_total: int, + kernel_kwargs: dict, + sync_stream=None, +) -> None: + budget = _LAUNCH_BLOCK_BUDGET + chunk_indices = kernel_kwargs.get('chunk_indices') + cu_seqlens = kernel_kwargs.get('cu_seqlens') + nt_step = nt if nt * nc * bh_total <= budget else max(1, budget // max(nc * bh_total, 1)) + for nt_off in range(0, nt, nt_step): + nt_len = min(nt_step, nt - nt_off) + if cu_seqlens is not None and chunk_indices is not None: + kernel_kwargs['chunk_indices'] = chunk_indices[nt_off:nt_off + nt_len] + kernel_kwargs['NT_OFFSET'] = 0 + else: + kernel_kwargs['NT_OFFSET'] = nt_off + max_nc = max_grid_axis_chunks(nc, nt_len * bh_total, max_grid=ASCEND_MAX_GRID_DIM) + for nc_off in range(0, nc, max_nc): + nc_len = min(max_nc, nc - nc_off) + kernel_kwargs['NC_OFFSET'] = nc_off + max_bh = max_grid_axis_chunks(bh_total, nt_len * nc_len, max_grid=ASCEND_MAX_GRID_DIM) + for bh_off in range(0, bh_total, max_bh): + bh_len = min(max_bh, bh_total - bh_off) + kernel_kwargs['BH_OFFSET'] = bh_off + kernel[(nt_len, nc_len, bh_len)](**kernel_kwargs) + if sync_stream is not None: + sync_stream.synchronize() + + +def _launch_inter_kernel(kernel, *, nt: int, bh_total: int, kernel_kwargs: dict) -> None: + budget = _LAUNCH_BLOCK_BUDGET + chunk_indices = kernel_kwargs.get('chunk_indices') + cu_seqlens = kernel_kwargs.get('cu_seqlens') + nt_step = nt if nt * bh_total <= budget else max(1, budget // max(bh_total, 1)) + for nt_off in range(0, nt, nt_step): + nt_len = min(nt_step, nt - nt_off) + if cu_seqlens is not None and chunk_indices is not None: + kernel_kwargs['chunk_indices'] = chunk_indices[nt_off:nt_off + nt_len] + kernel_kwargs['NT_OFFSET'] = 0 + else: + kernel_kwargs['NT_OFFSET'] = nt_off + max_bh = max_grid_axis_chunks(bh_total, nt_len, max_grid=ASCEND_MAX_GRID_DIM) + for bh_off in range(0, bh_total, max_bh): + bh_len = min(max_bh, bh_total - bh_off) + kernel_kwargs['BH_OFFSET'] = bh_off + kernel[(nt_len, bh_len)](**kernel_kwargs, **_INTER_COMPILE_KWARGS) + + +@triton.jit(do_not_specialize=['T', 'NT_OFFSET', 'NC_OFFSET', 'BH_OFFSET']) +def chunk_gdn2_fwd_kernel_intra_grouped_npu( + q, + k, + g, + b, + Aqk, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + BR: tl.constexpr, + ROW_GROUP: tl.constexpr, + IS_HEAD_MAJOR: tl.constexpr, + IS_VARLEN: tl.constexpr, + NT_OFFSET, + NC_OFFSET, + BH_OFFSET, +): + """Build one group of causal Aqk/Akk rows without unstable gate factoring.""" + i_t = tl.program_id(0) + NT_OFFSET + i_i = tl.program_id(1) + NC_OFFSET + i_bh = tl.program_id(2) + BH_OFFSET + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1, + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + else: + bos = tl.cast(i_b, tl.int64) * T + + i_ts = i_t * BT + i_i * BC + i_ti = i_ts + ROW_GROUP * BR + if i_ti >= T: + return + + o_r = tl.arange(0, BR) + o_k = tl.arange(0, BK) + o_c = i_ti + o_r + m_r = o_c < T + m_k = o_k < K + m_rk = m_r[:, None] & m_k[None, :] + + if IS_HEAD_MAJOR: + input_offset = ((tl.cast(i_b, tl.int64) * H + i_h) * T) * K + input_stride = K + else: + input_offset = (bos * H + i_h) * K + input_stride = H * K + q += input_offset + k += input_offset + g += input_offset + b += input_offset + Aqk += (bos * H + i_h) * BT + Akk += (bos * H + i_h) * BC + + p_q = q + o_c[:, None] * input_stride + o_k[None, :] + p_k = k + o_c[:, None] * input_stride + o_k[None, :] + p_g = g + o_c[:, None] * input_stride + o_k[None, :] + p_b = b + o_c[:, None] * input_stride + o_k[None, :] + b_q = tl.load(p_q, mask=m_rk, other=0.0).to(tl.float32) + b_k = tl.load(p_k, mask=m_rk, other=0.0).to(tl.float32) + b_g = tl.load(p_g, mask=m_rk, other=0.0).to(tl.float32) + b_b = tl.load(p_b, mask=m_rk, other=0.0).to(tl.float32) + b_k *= b_b + + for j in range(0, (ROW_GROUP + 1) * BR): + i_j = i_ts + j + m_j = i_j < T + p_kj = k + i_j * input_stride + o_k + p_gj = g + i_j * input_stride + o_k + b_kj = tl.load(p_kj, mask=m_j & m_k, other=0.0).to(tl.float32) + b_gj = tl.load(p_gj, mask=m_j & m_k, other=0.0).to(tl.float32) + b_kgj = tl.where(m_k[None, :], b_kj[None, :] * exp2(b_g - b_gj[None, :]), 0.0) + b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale + b_Akk = tl.sum(b_k * b_kgj, axis=1) + row_in_subchunk = ROW_GROUP * BR + o_r + tl.store( + Aqk + o_c * (H * BT) + i_i * BC + j, + b_Aqk.to(Aqk.dtype.element_ty), + mask=m_r & m_j & (j <= row_in_subchunk), + ) + tl.store( + Akk + o_c * (H * BC) + j, + b_Akk.to(Akk.dtype.element_ty), + mask=m_r & m_j & (j < row_in_subchunk), + ) + + +@triton.jit(do_not_specialize=['T', 'NT_OFFSET', 'NC_OFFSET', 'BH_OFFSET']) +def chunk_gdn2_fwd_kernel_diag_solve_npu( + Akkd, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + IS_VARLEN: tl.constexpr, + NT_OFFSET, + NC_OFFSET, + BH_OFFSET, +): + i_t = tl.program_id(0) + NT_OFFSET + i_i = tl.program_id(1) + NC_OFFSET + i_bh = tl.program_id(2) + BH_OFFSET + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + else: + bos = tl.cast(i_b, tl.int64) * T + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + Akkd += (bos * H + i_h).to(tl.int64) * BC + o_i = tl.arange(0, BC) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + p_Akk = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_ti, 0), (BC, BC), (1, 0)) + b_Akk = tl.load(p_Akk, boundary_check=(0, 1)).to(tl.float32) + b_Ai = -tl.where(m_A, b_Akk, 0) + for i in range(2, min(BC, T - i_ti)): + b_a = -tl.load(Akkd + (i_ti + i).to(tl.int64) * H * BC + o_i) + b_a = tl.where(o_i < i, b_a, 0.) + b_a += tl.sum(b_a[:, None] * b_Ai, 0) + b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) + b_Ai += m_I + tl.store(p_Akk, b_Ai.to(Akkd.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T', 'NT_OFFSET', 'BH_OFFSET']) +def chunk_gdn2_fwd_kernel_inter_products_npu( + q, + k, + g, + b, + Aqk, + Akkx, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + DST_BLOCK: tl.constexpr, + SRC_BLOCK: tl.constexpr, + IS_VARLEN: tl.constexpr, + NT_OFFSET, + BH_OFFSET, +): + i_t = tl.program_id(0) + NT_OFFSET + i_bh = tl.program_id(1) + BH_OFFSET + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1, + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + else: + bos = tl.cast(i_b, tl.int64) * T + + i_tc_dst = i_t * BT + DST_BLOCK * BC + i_tc_src = i_t * BT + SRC_BLOCK * BC + + base = bos * H + i_h + q += base * K + k += base * K + g += base * K + b += base * K + Aqk += base * BT + Akkx += base * BT + + o_i = tl.arange(0, BC) + m_dst = (i_tc_dst + o_i) < T + b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk = tl.zeros([BC, BC], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_k_src = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc_src, i_k * BK), (BC, BK), (1, 0)) + p_g_src = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc_src, i_k * BK), (BC, BK), (1, 0)) + b_k_src = tl.load(p_k_src, boundary_check=(0, 1)).to(tl.float32) + b_g_src = tl.load(p_g_src, boundary_check=(0, 1)).to(tl.float32) + + p_q_dst = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc_dst, i_k * BK), (BC, BK), (1, 0)) + p_k_dst = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc_dst, i_k * BK), (BC, BK), (1, 0)) + p_g_dst = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc_dst, i_k * BK), (BC, BK), (1, 0)) + p_b_dst = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_tc_dst, i_k * BK), (BC, BK), (1, 0)) + b_q_dst = tl.load(p_q_dst, boundary_check=(0, 1)).to(tl.float32) + b_k_dst = tl.load(p_k_dst, boundary_check=(0, 1)).to(tl.float32) + b_g_dst = tl.load(p_g_dst, boundary_check=(0, 1)).to(tl.float32) + b_b_dst = tl.load(p_b_dst, boundary_check=(0, 1)).to(tl.float32) + b_gn_dst = tl.load( + g + tl.cast(i_tc_dst, tl.int64) * H * K + o_k, + mask=m_k & (i_tc_dst < T), + other=0.0, + ).to(tl.float32) + b_gq = tl.where(m_dst[:, None], exp2(b_g_dst - b_gn_dst[None, :]), 0.0) + b_kgt = tl.trans(b_k_src * exp2(b_gn_dst[None, :] - b_g_src)) + b_Aqk += tl.dot(b_q_dst * b_gq, b_kgt, allow_tf32=False) + b_Akk += tl.dot((b_b_dst * b_k_dst) * b_gq, b_kgt, allow_tf32=False) + + p_Aqk = tl.make_block_ptr( + Aqk, + (T, BT), + (H * BT, 1), + (i_tc_dst, SRC_BLOCK * BC), + (BC, BC), + (1, 0), + ) + p_Akkx = tl.make_block_ptr( + Akkx, + (T, BT), + (H * BT, 1), + (i_tc_dst, SRC_BLOCK * BC), + (BC, BC), + (1, 0), + ) + tl.store(p_Aqk, (b_Aqk * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akkx, b_Akk, boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T', 'NT_OFFSET', 'BH_OFFSET']) +def chunk_gdn2_fwd_kernel_inter_solve_npu( + Akkd, + Akkx, + Akk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, + NT_OFFSET, + BH_OFFSET, +): + i_t = tl.program_id(0) + NT_OFFSET + i_bh = tl.program_id(1) + BH_OFFSET + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1, + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + else: + bos = tl.cast(i_b, tl.int64) * T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_tc0 + BC + i_tc2 = i_tc0 + 2 * BC + i_tc3 = i_tc0 + 3 * BC + + base = bos * H + i_h + Akkd += base * BC + Akkx += base * BT + Akk += base * BT + + p_Akkx10 = tl.make_block_ptr(Akkx, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + b_Akk10 = tl.load(p_Akkx10, boundary_check=(0, 1)).to(tl.float32) + if NC >= 3: + p_Akkx20 = tl.make_block_ptr(Akkx, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Akkx21 = tl.make_block_ptr(Akkx, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + b_Akk20 = tl.load(p_Akkx20, boundary_check=(0, 1)).to(tl.float32) + b_Akk21 = tl.load(p_Akkx21, boundary_check=(0, 1)).to(tl.float32) + if NC >= 4: + p_Akkx30 = tl.make_block_ptr(Akkx, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_Akkx31 = tl.make_block_ptr(Akkx, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + p_Akkx32 = tl.make_block_ptr(Akkx, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)) + b_Akk30 = tl.load(p_Akkx30, boundary_check=(0, 1)).to(tl.float32) + b_Akk31 = tl.load(p_Akkx31, boundary_check=(0, 1)).to(tl.float32) + b_Akk32 = tl.load(p_Akkx32, boundary_check=(0, 1)).to(tl.float32) + + p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) + b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) + if NC >= 3: + p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc2, 0), (BC, BC), (1, 0)) + b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) + if NC >= 4: + p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc3, 0), (BC, BC), (1, 0)) + b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + + # tl.dot may clobber its lhs on Ascend; materialize copies for later dot and store uses. + b_Ai11_c = b_Ai11 + 0.0 + if NC >= 3: + b_Ai22_c = b_Ai22 + 0.0 + b_Ai22_c2 = b_Ai22 + 0.0 + b_Ai22_c3 = b_Ai22 + 0.0 + if NC >= 4: + b_Ai33_c = b_Ai33 + 0.0 + b_Ai33_c2 = b_Ai33 + 0.0 + b_Ai33_c3 = b_Ai33 + 0.0 + b_Akk31_c = b_Akk31 + 0.0 + b_Akk32_c = b_Akk32 + 0.0 + + b_Ai10 = -tl.dot(tl.dot(b_Ai11, b_Akk10, allow_tf32=False), b_Ai00, allow_tf32=False) + if NC >= 3: + b_Ai21 = -tl.dot(tl.dot(b_Ai22, b_Akk21, allow_tf32=False), b_Ai11_c, allow_tf32=False) + b_Ai20 = -tl.dot( + b_Ai22_c2, + tl.dot(b_Akk20, b_Ai00, allow_tf32=False) + tl.dot(b_Akk21, b_Ai10, allow_tf32=False), + allow_tf32=False, + ) + if NC >= 4: + b_Ai32 = -tl.dot(tl.dot(b_Ai33, b_Akk32, allow_tf32=False), b_Ai22_c3, allow_tf32=False) + b_Ai31 = -tl.dot( + b_Ai33_c2, + tl.dot(b_Akk31, b_Ai11_c, allow_tf32=False) + tl.dot(b_Akk32, b_Ai21, allow_tf32=False), + allow_tf32=False, + ) + b_Ai30 = -tl.dot( + b_Ai33_c3, + tl.dot(b_Akk30, b_Ai00, allow_tf32=False) + + tl.dot(b_Akk31_c, b_Ai10, allow_tf32=False) + + tl.dot(b_Akk32_c, b_Ai20, allow_tf32=False), + allow_tf32=False, + ) + + p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk11, b_Ai11_c.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + if NC >= 3: + p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0)) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk22, b_Ai22_c.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + if NC >= 4: + p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)) + p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0)) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk33, b_Ai33_c.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +@input_guard +def chunk_gdn2_fwd_intra_npu( + 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, +): + # gk is already activated and accumulated; pairwise gate differences serve both gate modes. + del safe_gate + + B, T, H, K = k.shape + BT = chunk_size + BC = _BC + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + is_varlen = cu_seqlens is not None + + Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype) + Akk = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype) + Akkd = torch.zeros(B, T, H, BC, device=k.device, dtype=torch.float32) + Akkx = torch.zeros(B, T, H, BT, device=k.device, dtype=torch.float32) + + sync_stream = torch.npu.current_stream(k.device) + use_head_major_intra = not is_varlen + # serialize split dense launches to avoid CANN queue stalls. + dense_sync_stream = sync_stream if use_head_major_intra else None + if use_head_major_intra: + q_intra = q.transpose(1, 2).contiguous() + k_intra = k.transpose(1, 2).contiguous() + g_intra = gk.transpose(1, 2).contiguous() + b_intra = b.transpose(1, 2).contiguous() + else: + q_intra, k_intra, g_intra, b_intra = q, k, gk, b + for row_group in range(BC // _TOKEN_GROUP): + _launch_diag_kernel( + chunk_gdn2_fwd_kernel_intra_grouped_npu, + nt=NT, + nc=NC, + bh_total=B * H, + sync_stream=dense_sync_stream, + kernel_kwargs=dict( + q=q_intra, + k=k_intra, + g=g_intra, + b=b_intra, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=triton.next_power_of_2(K), + BR=_TOKEN_GROUP, + ROW_GROUP=row_group, + IS_HEAD_MAJOR=use_head_major_intra, + IS_VARLEN=is_varlen, + NT_OFFSET=0, + NC_OFFSET=0, + BH_OFFSET=0, + ), + ) + if not use_head_major_intra and row_group == 0: + # CANN can stall when the two grouped row kernels are queued together. + sync_stream.synchronize() + + _launch_diag_kernel( + chunk_gdn2_fwd_kernel_diag_solve_npu, + nt=NT, + nc=NC, + bh_total=B * H, + kernel_kwargs=dict( + Akkd=Akkd, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + BT=BT, + BC=BC, + IS_VARLEN=is_varlen, + NT_OFFSET=0, + NC_OFFSET=0, + BH_OFFSET=0, + ), + ) + product_kwargs = dict( + q=q, + k=k, + g=gk, + b=b, + Aqk=Aqk, + Akkx=Akkx, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=_get_inter_bk(K), + IS_VARLEN=is_varlen, + NT_OFFSET=0, + BH_OFFSET=0, + ) + for dst_block in range(1, NC): + for src_block in range(dst_block): + _launch_inter_kernel( + chunk_gdn2_fwd_kernel_inter_products_npu, + nt=NT, + bh_total=B * H, + kernel_kwargs=dict(product_kwargs, DST_BLOCK=dst_block, SRC_BLOCK=src_block), + ) + _launch_inter_kernel( + chunk_gdn2_fwd_kernel_inter_solve_npu, + nt=NT, + bh_total=B * H, + kernel_kwargs=dict( + Akkd=Akkd, + Akkx=Akkx, + Akk=Akk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + BT=BT, + BC=BC, + NC=NC, + IS_VARLEN=is_varlen, + NT_OFFSET=0, + BH_OFFSET=0, + ), + ) + w, u, qg, kg = recompute_w_u_fwd_gdn2( + k=k, + v=v, + b=b, + w_gate=w_gate, + A=Akk, + q=q if disable_recompute else None, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # Keep intra scratch alive until its last asynchronous consumer completes. + sync_stream.synchronize() + return w, u, qg, kg, Aqk, Akk diff --git a/fla/ops/gdn2/chunk_bwd.py b/fla/ops/gdn2/chunk_bwd.py index 10690db9ce..340e66794f 100644 --- a/fla/ops/gdn2/chunk_bwd.py +++ b/fla/ops/gdn2/chunk_bwd.py @@ -36,6 +36,7 @@ import triton import triton.language as tl +from fla.ops.backends import dispatch from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h from fla.ops.cp.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu_pre_process, expand_h0 from fla.ops.gdn2.chunk_intra import chunk_gdn2_bwd_intra @@ -252,6 +253,7 @@ def chunk_gdn2_bwd_kernel_wy_dqkg_fused( tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), mask=m_dA) +@dispatch('gdn2') def chunk_gdn2_bwd_wy_dqkg_fused( q: torch.Tensor, k: torch.Tensor, diff --git a/fla/ops/gdn2/chunk_intra.py b/fla/ops/gdn2/chunk_intra.py index 441cfac654..afc706dec3 100644 --- a/fla/ops/gdn2/chunk_intra.py +++ b/fla/ops/gdn2/chunk_intra.py @@ -24,6 +24,7 @@ import triton import triton.language as tl +from fla.ops.backends import dispatch from fla.ops.gdn2.chunk_intra_token_parallel import chunk_gdn2_fwd_intra_token_parallel from fla.ops.gdn2.wy_fast import recompute_w_u_fwd_gdn2 from fla.ops.utils import prepare_chunk_indices @@ -690,6 +691,7 @@ def chunk_gdn2_bwd_kernel_intra( tl.store(p_dg2, b_dg2.to(p_dg2.dtype.element_ty), mask=m_kc) +@dispatch('gdn2') def chunk_gdn2_fwd_intra( q: torch.Tensor, k: torch.Tensor, diff --git a/tests/ops/test_gdn2.py b/tests/ops/test_gdn2.py index 2d73483aca..6fc1d9f896 100644 --- a/tests/ops/test_gdn2.py +++ b/tests/ops/test_gdn2.py @@ -22,7 +22,12 @@ from fla.ops.gdn2 import chunk_gdn2, fused_recurrent_gdn2, naive_recurrent_gdn2 from fla.ops.kda.gate import naive_kda_gate, naive_kda_lowerbound_gate -from fla.utils import assert_close, device +from fla.utils import IS_AMD, IS_NPU, IS_NVIDIA, assert_close, device + +_requires_accelerator = pytest.mark.skipif( + not (IS_NVIDIA or IS_AMD or IS_NPU), + reason="CUDA/ROCm or Ascend NPU required", +) def _activate_g(g, A_log, dt_bias, safe_gate, lower_bound): @@ -66,7 +71,7 @@ def _rand_inputs(B, T, H, HV, K, V, dtype, *, gate_in_kernel=False, b_scale=1.0, # ============================================================================= # fused_recurrent # ============================================================================= -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@_requires_accelerator @pytest.mark.parametrize( ("B", "T", "H", "HV", "K", "V", "scale", "use_qk_l2norm_in_kernel", "dtype"), [ @@ -121,7 +126,7 @@ def test_fused_recurrent(B, T, H, HV, K, V, scale, use_qk_l2norm_in_kernel, dtyp assert_close("ht", ref_ht, tri_ht, 0.005) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@_requires_accelerator @pytest.mark.parametrize( ("B", "T", "H", "K", "V", "has_a_log", "has_dt_bias", "safe_gate"), [ @@ -173,7 +178,7 @@ def test_fused_recurrent_gate_in_kernel(B, T, H, K, V, has_a_log, has_dt_bias, s assert_close("ht", ref_ht, tri_ht, 0.005) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@_requires_accelerator def test_fused_recurrent_state_v_first(): """state_v_first stores the state transposed to [V, K]; output must match.""" dtype = torch.float32 @@ -186,7 +191,7 @@ def test_fused_recurrent_state_v_first(): assert_close("ht", ht0, ht1.transpose(-1, -2), 0.005) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@_requires_accelerator def test_fused_recurrent_initial_state(): dtype = torch.float32 B, T, H, K, V = 2, 64, 2, 64, 64 @@ -218,7 +223,7 @@ def test_fused_recurrent_initial_state(): assert_close("ht", ref_ht, tri_ht, 0.005) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@_requires_accelerator @pytest.mark.parametrize( ("cu_seqlens", "H", "K", "V"), [ @@ -281,7 +286,7 @@ def test_chunk_invalid_chunk_size(chunk_size): chunk_gdn2(q=q, k=k, v=v, g=g, b=b, w=w, chunk_size=chunk_size) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@_requires_accelerator @pytest.mark.parametrize( ("B", "T", "H", "K", "V", "scale", "use_qk_l2norm_in_kernel", "use_gate_in_kernel", "safe_gate", "dtype"), [ @@ -290,6 +295,8 @@ def test_chunk_invalid_chunk_size(chunk_size): (1, 64, 2, 32, 32, 1.0, False, False, False, torch.float32), (2, 256, 2, 64, 64, 0.5, True, False, False, torch.float32), (2, 100, 3, 64, 64, 1.0, True, False, False, torch.float16), # non-multiple T, fp16 + (1, 64, 1, 128, 128, 1.0, True, False, False, torch.bfloat16), + (1, 64, 1, 256, 256, 1.0, True, False, False, torch.bfloat16), (2, 256, 2, 64, 64, 1.0, True, True, False, torch.float32), # gate-in-kernel (1, 128, 2, 64, 64, 1.0, True, True, True, torch.float32), # gate-in-kernel + safe_gate ] @@ -369,7 +376,7 @@ def test_chunk(B, T, H, K, V, scale, use_qk_l2norm_in_kernel, use_gate_in_kernel assert_close("dt_bias", ref_grads["dt_bias"], tri_grads["dt_bias"], 0.02) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@_requires_accelerator def test_chunk_state_v_first(): """state_v_first must give the same output and a transposed final state.""" dtype = torch.float32 @@ -406,7 +413,7 @@ def test_chunk_state_v_first(): assert_close("ht", ht0, ht1.transpose(-1, -2), 0.005) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@_requires_accelerator @pytest.mark.parametrize( ("cu_seqlens", "H", "K", "V", "use_gate_in_kernel", "dtype"), [ @@ -484,7 +491,7 @@ def test_chunk_varlen(cu_seqlens, H, K, V, use_gate_in_kernel, dtype): assert_close("dh0", ref_grads["h0"], tri_grads["h0"], 0.012) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@_requires_accelerator @pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) def test_chunk_matches_fused_recurrent(dtype): """The two production kernels must agree with each other.""" @@ -498,7 +505,7 @@ def test_chunk_matches_fused_recurrent(dtype): assert_close("ht", ht_rec, ht_chunk, 0.006) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@_requires_accelerator @torch.inference_mode() def test_chunk_return_intermediate_states(): """return_intermediate_states yields per-chunk pre-states h; the output must @@ -528,7 +535,7 @@ def test_chunk_return_intermediate_states(): # ============================================================================= # layer — GatedDeltaNet2 (GVA + short conv) end to end # ============================================================================= -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@_requires_accelerator @pytest.mark.parametrize( ("num_heads", "num_v_heads", "use_short_conv"), [