kutacc_af2_gating_attention

The gating_attention function dynamically filters and controls the information flow by applying learnable gating to the standard multi-self-attention output, thereby enhancing the model's sensitivity and expression capability for key residues or sequence positions.

Interface Definition

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);

Parameters

Table 1 Input parameter definition

Parameter	Type	Description	Input/Output
input	kutacc_tensor_h	Input data	Input
q_based_ptr	kutacc_af2_attention_inputs_t *	Pointer of the kutacc_af2_attention_inputs_t type, pointing to the set of multi-input tensor data. For details, see the following table describing the kutacc_af2_attention_inputs_t data structure.	Input
bias	kutacc_tensor_h	Attention bias in the neural network	Input
nonbatched_bias	kutacc_tensor_h	General bias in the neural network	Input
weight_ptr	kutacc_af2_attention_weights_t *	Pointer to the combination of multiple weight data of the kutacc_af2_attention_weights_t type. For details, see the following table describing the kutacc_af2_attention_weights_t structure.	Input
out	kutacc_tensor_h	Output data	Output
block_size	int64_t	Data block size in parallel operations	Input

Table 2 kutacc_af2_attention_inputs_t structure definition

Parameter	Type	Description
gate	kutacc_tensor_h	Input intermediate gate data
k	kutacc_tensor_h	Input intermediate k data
v	kutacc_tensor_h	Input intermediate v data
q	kutacc_tensor_h	Input intermediate q data
avg	kutacc_tensor_h	Weight avg data
batch	int64_t	Input batch
seq_len	int64_t	Sequence length of the input data

Table 3 kutacc_af2_attention_weights_t structure definition

Parameter	Type	Description
nchannels	int64_t	Total number of input weights
nheads	int64_t	Number of heads of input weights
head_size	int64_t	Number of data elements in each head of input weights
query_w	kutacc_tensor_h	Weight query_w data
key_w	kutacc_tensor_h	Weight key_w data
gating_w	kutacc_tensor_h	Weight gating_w data
gating_b	kutacc_tensor_h	Gating bias
output_w	kutacc_tensor_h	Weight output_w data
output_b	kutacc_tensor_h	Output matrix multiplication bias
value_w	kutacc_tensor_h	Weight value_w data

Constraints on integer parameters:

nchannels = nheads * head_size;

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

nchannels < INT64_MAX

The table below describes shape constraints when the kpex API is called to construct tensors.

Table 4 Shape constraints on kpex gating_attention parameters

Parameter	Shape
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]

Examples

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)
 *
 */

Parent topic: AlphaFold2