---
title: kutacc_af2_gating_attention
description: "gating_attention函数通过在标准多头自注意力（multi-self-attention）输出上施加可学习的门控（gating），来动态地筛选和控制信息流，从而提升模型对关键残基或序列位置的敏感性与表达能力。"
url: https://www.hikunpeng.com/document/detail/zh/kunpenghpcs/hpckit/devg/KunpengHPCKit_developer_133.html
sourcePath: /source/zh/kunpenghpcs/hpckit/devg/KunpengHPCKit_developer_133.html
indexId: 7bf8c55bfef5174db7fdcd919943eedbc70a3d74b7c561e8f9320cc985f0108b67
---
# kutacc_af2_gating_attention

gating_attention函数通过在标准多头自注意力（multi-self-attention）输出上施加可学习的门控（gating），来动态地筛选和控制信息流，从而提升模型对关键残基或序列位置的敏感性与表达能力。

#### 接口定义

void kutacc_af2_gating_attention(kutacc_tensor_h input, kutacc_af2_attention_inputs_t *q_based_ptr, kutacc_tensor_h bias, kutacc_tensor_h nonbatched_bias, kutacc_af2_attention_weights_t *weight_ptr, kutacc_tensor_h out, int64_t block_size);


#### 参数

表1 入参定义

| 参数名 | 类型 | 描述 | 输入/输出 |
| --- | --- | --- | --- |
| input | kutacc\_tensor\_h | 输入数据 | 输入 |
| q\_based\_ptr | kutacc\_af2\_attention\_inputs\_t \* | kutacc\_af2\_attention\_inputs\_t类型的指针，它指向多重输入张量数据的集合，具体见下方kutacc\_af2\_attention\_inputs\_t数据结构表 | 输入 |
| bias | kutacc\_tensor\_h | 神经网络中的attention偏移量数据 | 输入 |
| nonbatched\_bias | kutacc\_tensor\_h | 神经网络中的通用偏移量数据 | 输入 |
| weight\_ptr | kutacc\_af2\_attention\_weights\_t \* | kutacc\_af2\_attention\_weights\_t类型的指针，它指向多重权重数据的结合，具体见下方kutacc\_af2\_attention\_weights\_t数据结构表 | 输入 |
| out | kutacc\_tensor\_h | 输出数据 | 输出 |
| block\_size | int64\_t | 并行操作中的数据块切分大小 | 输入 |


表2 kutacc_af2_attention_inputs_t数据结构表定义

| 参数名 | 类型 | 描述 |
| --- | --- | --- |
| gate | kutacc\_tensor\_h | 输入中间层gate数据 |
| k | kutacc\_tensor\_h | 输入中间层k数据 |
| v | kutacc\_tensor\_h | 输入中间层v数据 |
| q | kutacc\_tensor\_h | 输入中间层q数据 |
| avg | kutacc\_tensor\_h | 权重avg数据 |
| batch | int64\_t | 输入批次 |
| seq\_len | int64\_t | 输入数据的序列长度 |


表3 kutacc_af2_attention_weights_t数据结构表

| 参数名 | 类型 | 描述 |
| --- | --- | --- |
| nchannels | int64\_t | 输入权重的总数 |
| nheads | int64\_t | 输入权重的head数 |
| head\_size | int64\_t | 输入权重每个head中的数据数量 |
| query\_w | kutacc\_tensor\_h | 权重query\_w数据 |
| key\_w | kutacc\_tensor\_h | 权重key\_w数据 |
| gating\_w | kutacc\_tensor\_h | 权重gating\_w数据 |
| gating\_b | kutacc\_tensor\_h | 门控偏移量数据 |
| output\_w | kutacc\_tensor\_h | 权重output\_w数据 |
| output\_b | kutacc\_tensor\_h | 输出矩阵乘法偏移量数据 |
| value\_w | kutacc\_tensor\_h | 权重value\_w数据 |


整数参数的约束关系 应满足:

nchannels = nheads * head_size;

batch, seq_len, nchannels, nheads, head_size, block_size > 0,

nchannels < INT64_MAX


在调用kpex接口时构造tensor应满足以下形状约束

表4 kpex gating_attention入参形状约束

| 参数 | 形状 |
| --- | --- |
| query\_w | [nheads, head\_size, nchannels] |
| key\_w | [nheads, head\_size, nchannels] |
| value\_w | [nheads, head\_size, nchannels] |
| gating\_w | [nchannels, nheads, head\_size] |
| gating\_b | [nheads, head\_size] |
| output\_w | [nchannels, nheads, head\_size] |
| output\_b | [nchannels] |
| q\_data | [batch, seq\_len, nchannels] |
| m\_data | [batch, seq\_len, nchannels] |
| bias | [batch, 1, 1, seq\_len] |
| nonbatched\_bias | [nheads, seq\_len, seq\_len] or [0] |


#### 示例

C++ interface：
```
//test_gating_attention.h
#include <ATen/core/Tensor.h>
#include <ATen/ops/empty.h>
#include <ATen/ops/ones.h>
#include <ATen/ops/zeros.h>
#include <ATen/ops/full.h>
#include <ATen/native/cpu/utils.h>
#include <c10/core/ScalarType.h>
#include "common_header.h"
namespace alphafold {
at::Tensor test_global_attention(int64_t batch, int64_t seq_len, int64_t nchannels, int64_t nheads, int64_t head_size);
}
//bind.h
#include "test_gating_attention.h"
#include <torch/extension.h>
namespace alphafold {
inline void bind(pybind11::module &m)
{
auto submodule = m.def_submodule("alphafold");
submodule.("test_gating_attention", &test_gating_attention, py::("batch"), py::("seq_len"), py::("nchannels"), py::("nheads"), py::("head_size"), py::("block_size"));
}
}
//test_gating_attention.cpp
#include "kutacc.h"
#include <utils/memory.h>
#include <utils/bf16.h>
#include <utils/TensorWrapper.h>
#include "test_gating_attention.h"
namespace alphafold {
at::Tensor test_gating_attention(int64_t batch, int64_t seq_len, int64_t nchannels, int64_t nheads, int64_t head_size, int64_t block_size) {
float a = 3.0f;
float b = 2.5f;
at::Tensor q_data = at::ones({batch, seq_len, nchannels}, at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16));
auto q = q_data.new_empty({batch, seq_len, nheads, head_size});
auto k = q_data.new_empty({batch, seq_len, nheads, head_size});
auto v = q_data.new_empty({nheads, head_size, batch, seq_len});
auto gate = q_data.new_empty({batch, seq_len, nheads, head_size});
auto weighted_avg = q_data.new_empty({batch, seq_len, nheads, head_size});
at::Tensor input = at::ones({batch * seq_len, nchannels}, q_data.options());
at::Tensor bias =  at::ones({batch, 1, 1, seq_len}, q_data.options());
at::Tensor nonbatched_bias = at::zeros({nheads, seq_len, seq_len}, q_data.options());
at::Tensor query_w = at::full({nheads, head_size, nchannels}, a, q_data.options()).to(at::kBFloat16);
auto bf16_opt = query_w.options().device(kpex::device()).dtype(at::kBFloat16);
auto float_opt = query_w.options().device(kpex::device()).dtype(at::kFloat);
at::Tensor key_w = at::full({nheads, head_size, nchannels}, b, q_data.options()).to(bf16_opt);
at::Tensor value_w = at::full({nheads, head_size, nchannels}, a, q_data.options()).to(bf16_opt);
at::Tensor gating_w = at::full({nheads, head_size, nchannels}, a, q_data.options()).to(bf16_opt);
at::Tensor gating_b = at::full({nheads, head_size}, b, q_data.options()).to(float_opt);
at::Tensor output_w = at::full({nchannels, nheads, head_size}, b, q_data.options()).to(bf16_opt);
at::Tensor output_b = at::full({nchannels}, a, q_data.options()).to(float_opt);
query_w = query_w.to(bf16_opt).contiguous().view({nchannels, nchannels});
key_w = key_w.to(bf16_opt).contiguous().view({nchannels, nchannels});
value_w = value_w.to(bf16_opt).contiguous().view({nchannels, nchannels});
gating_w = gating_w.to(bf16_opt).contiguous().view({nchannels, nchannels});
output_w = output_w.to(bf16_opt).contiguous().view({nchannels, nchannels});
auto query_w_res = linear_weight_prepack(query_w);
auto key_w_res = linear_weight_prepack(key_w);
auto value_w_res = linear_weight_prepack(value_w);
auto gating_w_res = linear_weight_prepack(gating_w);
auto output_w_res = linear_weight_prepack(output_w);
at::Tensor out = at::empty(q_data.sizes(), q_data.options());
auto input_tw = convert_to_tensor_wrapper(input);
auto q_tw = convert_to_tensor_wrapper(q);
auto k_tw = convert_to_tensor_wrapper(k);
auto v_tw = convert_to_tensor_wrapper(v);
auto gate_tw = convert_to_tensor_wrapper(gate);
auto weighted_avg_tw = convert_to_tensor_wrapper(weighted_avg);
auto bias_tw = convert_to_tensor_wrapper(bias);
auto nonbatched_bias_tw = convert_to_tensor_wrapper(nonbatched_bias);
auto query_w_tw = convert_to_tensor_wrapper(query_w_res);
auto key_w_tw = convert_to_tensor_wrapper(key_w_res);
auto value_w_tw = convert_to_tensor_wrapper(value_w_res);
auto gating_w_tw = convert_to_tensor_wrapper(gating_w);
auto gating_b_tw = convert_to_tensor_wrapper(gating_w_res);
auto output_w_tw = convert_to_tensor_wrapper(output_w_res);
auto output_b_tw = convert_to_tensor_wrapper(output_b);
auto out_tw = convert_to_tensor_wrapper(out);
kutacc_af2_attention_weights_t_wrapper *gating_attention_weight_ptr = new kutacc_af2_attention_weights_t_wrapper(query_w_tw, key_w_tw, value_w_tw, gating_w_tw, gating_b_tw,
output_w_tw, output_b_tw, nchannels, nheads, head_size);
kutacc_af2_attention_inputs_t_wrapper *gating_attention_q_ptr = new kutacc_af2_attention_inputs_t_wrapper(q_tw, k_tw, v_tw, gate_tw, weighted_avg_tw, batch, seq_len);
kutacc_af2_gating_attention(input_tw.get_tensor(), gating_attention_q_ptr,bias_tw.get_tensor(), nonbatched_bias_tw.get_tensor(), gating_attention_weight_ptr, out_tw.get_tensor(), block_size);
delete gating_attention_weight_ptr;
delete gating_attention_q_ptr;
return out;
}
}
//test.py
import torch
import kpex._C as kernel
def test_gating_attention(batch, seq_len, nchannels, nheads, head_size, block_size):
out = kernel.alphafold.test_gating_attention(batch, seq_len, nchannels, nheads, head_size, block_size)
print(out)
return out
/**
* input：2, 2, 8, 4, 2, 2
* output:
tensor([[[484., 484., 484., 484., 484., 484., 484., 484.],
[484., 484., 484., 484., 484., 484., 484., 484.]],
[[484., 484., 484., 484., 484., 484., 484., 484.],
[484., 484., 484., 484., 484., 484., 484., 484.]]],
dtype=torch.bfloat16)
*
*/
```