ChunkGatedDeltaRule

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	×
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

• 算子功能：完成chunk版的Gated Delta Rule计算。

• 计算公式：

  Gated Delta Rule（门控Delta规则，GDR）是一种应用于循环神经网络的算子，也被应用于一种线性注意力机制中。在每个时间步 $t$，GDR根据当前的输入 $q_t$、$k_t$、$v_t$、上一个隐藏状态 $S_{t-1}$、衰减系数 ${\alpha }_t$ 以及更新强度 ${\beta }_t$，计算当前的注意力输出 $o_t$ 和新的隐藏状态 $S_t$，其计算公式如下：

  $$S_t:=S_{t-1}({\alpha }_t(I-{\beta }_t{k}_t{k}_t^T))+{\beta }_t{v}_t{k}_t^T={\alpha }_t{S}_{t-1}+{\beta }_t(v_t-{\alpha }_t{S}_{t-1}{k}_t)k_t^T$$

  $$o_t:=S_t(q_t\cdot scale)$$

  其中，$S_{t-1},S_t\in {\mathbb{R}}^{D_v\times D_k}$，$q_t,k_t\in {\mathbb{R}}^{D_k}$，$v_t\in {\mathbb{R}}^{D_v}$，${\alpha }_t\in \mathbb{R}$，${\beta }_t\in \mathbb{R}$，$o_t\in {\mathbb{R}}^{D_v}$。

  Chunked Gated Delta Rule是GDR的chunk版实现(参考论文)，它通过将输入序列切块，实现了一定的并行效果，在长上下文场景其计算效率相对Recurrent Gated Delta Rule更高，适用于prefill阶段。输入一个长度为L的序列，该算子可以计算出每一步的输出 $o_t,t\in \{1,2,..,L\}$ 以及最终的状态矩阵 $S_L$。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
query	输入	公式中的q。	BFLOAT16	ND
key	输入	公式中的输入k。	BFLOAT16	ND
value	输入	公式中的输入v。	BFLOAT16	ND
beta	输入	公式中的β。	BFLOAT16	ND
initial_state	输入	初始状态矩阵，公式中的输入S_0。	BFLOAT16	ND
actual_seq_lengths	输入	每个batch的序列长度。	INT32	ND
g	输入	衰减系数，公式中的α=e^g	FLOAT32	ND
out	输出	每一步的attention结果，公式中的o_t。	BFLOAT16	ND
final_state	输出	最终的状态矩阵，公式中的S_L。	BFLOAT16	ND
scale_value	可选属性	query的缩放因子，公式中的scale。默认为1.0	FLOAT	-

约束说明

• 为方便理解后续排布格式（如 BNSD、TND 等），统一说明各缩写维度含义：

  • B：输入样本批量大小（Batch）。
  • T：设 $L_i$ 为第 $i$ 个序列长度，则 $T={\sum }_i^B{L}_i$ 表示累积序列长度。
  • Nk：Query 和 Key 头数。
  • Nv：Value 头数。
  • Dk：Query 和 Key 隐藏层维度。
  • Dv：Value 隐藏层维度。

• 当前仅支持 TND 布局：

  • query、key 形状：$(T,Nk,Dk)$
  • value、out 形状：$(T,Nv,Dv)$
  • beta、g 形状：$(T,Nv)$
  • actual_seq_lengths 形状：$(B,)$
  • initial_state、final_state 形状：$(B,Nv,Dv,Dk)$

  维度需满足以下约束：

  • $0<Nv\le 64，0<Nk\le 64$，且 $Nv\mathrm{m}\mathrm{o}\mathrm{d}Nk=0$
  • $0<Dv\le 128，0<Dk\le 128$
  • $B>0，T>0$

• 受算法数值特性限制，需满足以下取值约束，否则易出现数值溢出、精度异常：

  • 张量元素：$0<\text{query}<1$
  • 张量元素：$0<\text{key}<1$
  • 张量元素：$g<0$
  • 张量元素：$0<\text{beta}<1$

调用说明

调用方式	样例代码	说明
aclnn	test_aclnn_chunk_gated_delta_rul.cpp	通过aclnnChunkGatedDeltaRule调用aclnnChunkGatedDeltaRule算子