kutacc_af2_global_attention
In AF2, global_attention acts as a global attention mechanism. By integrating target residue information along the column direction during the multiple sequence alignment process, it transfers 3D structural information from the first sequence to the other sequences. This enables a better understanding of the residue's role across different evolutionary contexts, as well as the co-evolutionary relationships between residues, ultimately providing richer informational representations for protein structure prediction.
Interface Definition
void kutacc_af2_global_attention(kutacc_af2_attention_inputs_t *q_based_ptr, kutacc_tensor_h q_data, kutacc_tensor_h q_mask, kutacc_af2_attention_weights_t *weight_ptr, kutacc_tensor_h out);
Parameters
Table 1 Input parameter definition
Parameter |
Type |
Description |
Input/Output |
|---|---|---|---|
q_based_ptr |
kutacc_af2_attention_inputs_t * |
Pointer of the kutacc_af2_attention_inputs_t type. For details about the data structure, see Table 2. |
Input |
q_data |
kutacc_tensor_h |
Q matrix data |
Input |
q_mask |
kutacc_tensor_h |
Q mask matrix |
Input |
weight_ptr |
kutacc_af2_attention_weights_t * |
Pointer of the kutacc_af2_attention_weights_t type. For details about the data structure, see Table 3. |
Input |
out |
kutacc_tensor_h |
Output data |
Output |
Table 2 kutacc_af2_attention_inputs_t structure definition
Parameter |
Type |
Description |
gate |
kutacc_tensor_h |
Input intermediate gating tensor |
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 of global_attention:
nchannels = nheads * head_size
batch, seq_len, nchannels, nheads, head_size > 0
Inchannels = 64
The table below describes operator shape constraints when a test case is built or when KPEX is called.
Table 4 Shape constraints on global_attention
Tensor |
Shape |
|---|---|
query_w |
[nheads, head_size, nchannels] |
key_w |
[head_size, nchannels] |
value_w |
[head_size, nchannels] |
gating_w |
[nheads, head_size, nchannels] |
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] |
q_mask |
[batch, seq_len, 1] |
Examples
C++ interface:
//test_global_attention.h
#ifndef KPEX_TPP_ALPHAFOLD_TEST_GLOBAL_ATTENTION_H
#define KPEX_TPP_ALPHAFOLD_TEST_GLOBAL_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);
}
#endif
//bind.h
#include <torch/extension.h>
#include "test_global_attention.h"
namespace alphafold {
inline void bind(pybind11::module &m)
{
auto submodule = m.def_submodule("alphafold");
submodule.def("test_global_attention", &test_global_attention, py::arg("batch"), py::arg("seq_len"), py::arg("nchannels"), py::arg("nheads"), py::arg("head_size"));
}
//test_global_attention.cpp
#include "kutacc.h"
#include <utils/memory.h>
#include <utils/bf16.h>
#include <utils/TensorWrapper.h>
#include "test_global_attention.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) {
float a = 0.75f;
float b = 0.25f;
float c = 0.5f;
at::Tensor q_data = at::ones({batch, seq_len, nchannels}, at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16));
at::Tensor out = at::empty(q_data.sizes(), q_data.options());
at::Tensor q_mask = at::ones({batch, seq_len, 1}, q_data.options());
q_mask = q_mask.contiguous();
at::Tensor q_avg = q_data.new_empty({batch, nchannels});
at::Tensor q = q_data.new_empty({batch, nheads, head_size});
at::Tensor k = q_data.new_empty({batch, seq_len, head_size});
at::Tensor v = q_data.new_empty({head_size, batch, seq_len});
at::Tensor gate = q_data.new_empty({batch, seq_len, nheads, head_size});
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({head_size, nchannels}, b, at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16));
at::Tensor value_w = at::full({head_size, nchannels}, a, at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16));
at::Tensor gating_w = at::full({nheads, head_size, nchannels}, a, at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16));
at::Tensor gating_b = at::full({nheads, head_size}, c, at::TensorOptions().device(kpex::device()).dtype(at::kFloat));
at::Tensor output_w = at::full({nchannels, nheads, head_size}, b, at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16));
at::Tensor output_b = at::full({nchannels}, c, at::TensorOptions().device(kpex::device()).dtype(at::kFloat));
query_w = query_w.to(bf16_opt).contiguous().view({nchannels, nchannels});
key_w = key_w.to(bf16_opt).contiguous().view({head_size, nchannels});
value_w = value_w.to(bf16_opt).contiguous().view({head_size, 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);
kutacc::TensorWrapper q_avg_tw = convert_to_tensor_wrapper(q_avg);
kutacc::TensorWrapper q_tw = convert_to_tensor_wrapper(q);
kutacc::TensorWrapper k_tw = convert_to_tensor_wrapper(k);
kutacc::TensorWrapper v_tw = convert_to_tensor_wrapper(v);
kutacc::TensorWrapper gate_tw = convert_to_tensor_wrapper(gate);
kutacc::TensorWrapper q_data_tw = convert_to_tensor_wrapper(q_data);
kutacc::TensorWrapper q_mask_tw = convert_to_tensor_wrapper(q_mask);
kutacc::TensorWrapper out_tw = convert_to_tensor_wrapper(out);
kutacc::TensorWrapper query_w_tw = convert_to_tensor_wrapper(query_w_res);
kutacc::TensorWrapper key_w_tw = convert_to_tensor_wrapper(key_w_res);
kutacc::TensorWrapper value_w_tw = convert_to_tensor_wrapper(value_w_res);
kutacc::TensorWrapper gating_w_tw = convert_to_tensor_wrapper(gating_w_res);
kutacc::TensorWrapper gating_b_tw = convert_to_tensor_wrapper(gating_b);
kutacc::TensorWrapper output_w_tw = convert_to_tensor_wrapper(output_w_res);
kutacc::TensorWrapper output_b_tw = convert_to_tensor_wrapper(output_b);
kutacc_af2_attention_weights_t_wrapper *global_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 *global_attention_q_ptr = new kutacc_af2_attention_inputs_t_wrapper(q_tw, k_tw, v_tw, gate_tw, q_avg_tw, batch, seq_len);
kutacc_af2_global_attention(global_attention_q_ptr, q_data_tw.get_tensor(), q_mask_tw.get_tensor(), global_attention_weight_ptr, out_tw.get_tensor());
delete global_attention_weight_ptr;
delete global_attention_q_ptr;
return out;
}
}
//test.py
import torch
import kpex._C as kernel
def test_global_attention(batch, seq_len, nchannels, nheads, head_size):
out = kernel.alphafold.test_global_attention(batch, seq_len, nchannels, nheads, head_size)
return out
//output:
tensor([[[768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768.],
[768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768.]],
[[768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768.],
[768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768., 768., 768.,
768., 768., 768., 768., 768., 768., 768., 768., 768.]]],
dtype=torch.bfloat16)