kutacc_af2_invariant_point

By utilizing an attention mechanism to maintain feature point invariance, this interface enhances the model's performance in protein structure prediction. Through the weighted aggregation of input features, the module effectively captures inter-feature relationships and precisely adjusts protein atomic coordinates, thereby improving the accuracy and stability of structure predictions.

Interface Definition

void kutacc_af2_invariant_point(kutacc_af2_ipa_s_inputs_t *ipa_s_ptrs, kutacc_af2_ipa_o_inputs_t *ipa_o_ptrs, kutacc_tensor_h z, kutacc_tensor_h rigid_rot_mats, kutacc_tensor_h rigid_trans, kutacc_tensor_h mask, kutacc_af2_ipa_weights_t *ipa_weight_ptrs);

Parameters

Table 1 Input parameters

Parameter	Type	Description	Input/Output
ipa_s_ptrs	kutacc_af2_ipa_s_inputs_t *	Pointer of the kutacc_af2_ipa_s_inputs_t type. For details, see Table 2.	Input
ipa_o_ptrs	kutacc_af2_ipa_o_inputs_t *	Pointer of the kutacc_af2_ipa_o_inputs_t type. For details, see Table 3.	Input
z	kutacc_tensor_h	Intermediate variable	Input
rigid_rot_mats	kutacc_tensor_h	Rotation parameter for rigid transformation	Input
rigid_trans	kutacc_tensor_h	Translation parameter for rigid transformation	Input
mask	kutacc_tensor_h	Mask, used to ignore invalid residues	Input
ipa_weight_ptrs	kutacc_af2_ipa_weights_t *	Pointer of the kutacc_af2_ipa_weights_t type. For details, see Table 4.	Input

Table 2 kutacc_af2_ipa_s_inputs_t structure definition

Parameter	Type	Description	Input/Output
n_res	int64_t	Length of dimension 0 of the input vector s.	Input
a	kutacc_tensor_h	Temporary result of the matrix multiplication of q transpose and k after applying the Softmax operation.	Input
b	kutacc_tensor_h	Temporary result of linear transformation of vector z.	Input
q	kutacc_tensor_h	Matrix Q generated by linear transformation of input s.	Input
k	kutacc_tensor_h	Matrix K	Input
v	kutacc_tensor_h	Matrix V	Input
q_pts	kutacc_tensor_h	3D point of the query matrix	Input
k_pts	kutacc_tensor_h	3D point of the key matrix	Input
v_pts	kutacc_tensor_h	3D point of the value matrix	Input

Table 3 kutacc_af2_ipa_o_inputs_t structure definition

Parameter	Type	Description	Input/Output
o	kutacc_tensor_h	Vector for storing the qkv and attention calculation results	Input
o_pt	kutacc_tensor_h	Vector for storing the attention calculation results of the pts-related vector	Input
o_pt_norm	kutacc_tensor_h	Stores the vectors that require LayerNorm transformation	Input
o_pair	kutacc_tensor_h	Residue pair feature	Input

Table 4 kutacc_af2_ipa_weights_t structure definition

Parameter	Type	Description	Input/Output
c_z	int64_t	Length of a single	Input
c_hidden	int64_t	Length of the hidden layer	Input
no_heads	int64_t	Number of attention heads	Input
no_qk_points	int64_t	Number of q/k points to be generated	Input
no_v_points	int64_t	Number of v points to be generated	Input
head_weights	kutacc_tensor_h	Weight of the attention head	Input
weights_head_weights	kutacc_tensor_h	Weight of the attention head	Input
linear_b_w	kutacc_tensor_h	Weight for linear transformation of b	Input
linear_b_b	kutacc_tensor_h	Bias for linear transformation of b	Input

Constraints on integer parameters: n_res, c_z, c_hidden, no_heads, no_qk_points, no_v_points > 0

no_heads * c_hidden < INT64_MAX; no_heads * no_v_points * 3 <INT64_MAX

no_heads*(c_hidden+no_v_points*3) <INT64_MAX; no_heads*(c_hidden+no_v_points*4) <INT64_MAX; no_heads * c_z < INT64_MAX

no_heads * c_hidden * 2 * c_s < INT64_MAX; no_heads * 3 *no_qk_points * c_s < INT64_MAX

3 * no_heads * (no_qk_points+ no_v_points) * c_s <INT64_MAX

no_heads * 2 * c_hidden = c_s

no_qk_points * 3 < 16

The table below describes constraints on tensor construction at the KPEX layer.

Tensor	Shape	Description
s	[n_res, c_s]	Input tensor.
z	[n_res, n_res, c_z]	Input tensor. For details, see the description of the input parameter z.
rigid_trans	[n_res, 3]	Input tensor. For details, see the description of the input parameter rigid_trans.
rigid_rot_mats	[n_res, 3, 3]	Input tensor. For details, see the description of the input parameter rigid_rot_mats.
mask	[n_res]	Input tensor. For details, see the description of the input parameter mask.
linear_q_w	[no_heads * c_hidden, c_s]	Weight of q generated by linear transformation of s.
linear_q_b	[no_heads * c_hidden]	Bias of q generated by linear transformation of s.
linear_kv_w	[no_heads * 2 * c_hidden, c_s]	Weight of k and v generated by linear transformation of s.
linear_kv_b	[no_heads * 2 * c_hidden]	Bias of k and v generated by linear transformation of s.
linear_q_points_w	[no_heads * no_qk_points * 3, c_s]	Weight of q_pts generated by linear transformation of s.
linear_q_points_b	[no_heads * no_qk_points *3]	Bias of q_pts generated by linear transformation of s.
linear_kv_points_w	[3 * no_heads * (no_qk_points + no_v_points), c_s]	Weight of k_pts and v_pts generated by linear transformation of s.
linear_kv_points_b	[3 * no_heads * (no_qk_points + no_v_points)]	Bias of k_pts and v_pts generated by linear transformation of s.
linear_b_w	[no_heads, c_z]	See linear_b_w in "Table 4 kutacc_af2_ipa_weights_t structure definition".
linear_b_b	[no_heads]	See linear_b_b in "Table 4 kutacc_af2_ipa_weights_t structure definition".
head_weights	[no_heads]	See head_weights in "Table 4 kutacc_af2_ipa_weights_t structure definition".
linear_out_w	[c_s, no_heads * (c_hidden + no_v_points * 4 + c_z)]	See linear_out_w in "Table 4 kutacc_af2_ipa_weights_t structure definition".
linear_out_b	[c_s]	See linear_out_b in "Table 4 kutacc_af2_ipa_weights_t structure definition".

Examples

C++ interface:

//test_invariant.h
#ifndef KPEX_TPP_ALPHAFOLD_TEST_INVARIAN_H 

#define KPEX_TPP_ALPHAFOLD_TEST_INVARIAN_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>
namespace alphafold { 
at::Tensor test_invariant_point_attention(int64_t n_res, int64_t no_heads, int64_t c_hidden, int64_t no_qk_points, int64_t no_v_points, int64_t c_z, int64_t c_s); 
} 
#endif

//bind.h
#include "test_invariant_point.h"
namespace alphafold { 

inline void bind(pybind11::module &m) 
{ 
    auto submodule = m.def_submodule("alphafold");
    submodule.def("test_invariant_point_attention", &test_invariant_point_attention, py::arg("n_res"), py::arg("no_heads"), py::arg("c_hidden"), py::arg("no_qk_points"), 
        py::arg("no_v_points"), py::arg("c_z"), py::arg("c_s"));
}
}

//test_invariant.cpp
#include "test_invariant_point.h" 

#include "utils/linear.h" 
#include "rigid.h" 
#include "kutacc.h" 
#include "utils/memory.h" 
#include "invariant_point.h" 
namespace alphafold { 
at::Tensor test_invariant_point_attention(int64_t n_res, int64_t no_heads, int64_t c_hidden, int64_t no_qk_points, int64_t no_v_points, int64_t c_z, int64_t c_s) 
{ 
    float a = 0.6f; 
    float b = 0.25f; 
    float c = 0.1f; 
    float d = 0.2f; 
    float e = 0.5f; 
    at::Tensor s = at::full({n_res, c_s}, a, at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16)); 
    at::Tensor out = at::empty(s.sizes(), s.options()); 
    at::Tensor rigid_trans = at::ones({n_res, 1, 1, 3}, s.options()).to(at::kFloat); 
    at::Tensor rigid_rot_mats = at::ones({n_res, 1, 1, 3, 3}, s.options()).to(at::kFloat); 
    at::Tensor linear_q_w = at::full({no_heads, c_hidden, c_s}, b, s.options()); 
    at::Tensor linear_q_b = at::full({no_heads, c_hidden}, c, rigid_rot_mats.options()); 
    at::Tensor linear_k_w = at::ones({no_heads, c_hidden, c_s}, s.options()); 
    at::Tensor linear_k_b = at::zeros({no_heads, c_hidden}, rigid_rot_mats.options()); 
    at::Tensor linear_v_w = at::ones({no_heads, c_hidden, c_s}, s.options()); 
    at::Tensor linear_v_b = at::zeros({no_heads, c_hidden}, rigid_rot_mats.options()); 
    at::Tensor linear_q_points_w = at::full({no_heads, no_qk_points, 3, c_s}, c, s.options()); 
    at::Tensor linear_q_points_b = at::ones({no_heads, no_qk_points, 3}, rigid_rot_mats.options()); 
    at::Tensor linear_k_points_w = at::full({no_heads, no_qk_points, 3, c_s}, d, s.options()); 
    at::Tensor linear_k_points_b = at::full({no_heads, no_qk_points, 3}, d, rigid_rot_mats.options()); 
    at::Tensor linear_v_points_w = at::ones({no_heads, no_v_points, 3, c_s}, s.options()); 
    at::Tensor linear_v_points_b = at::zeros({no_heads, no_v_points, 3}, rigid_rot_mats.options()); 
    at::Tensor linear_b_w = at::ones({no_heads, c_z}, s.options()); 
    at::Tensor linear_b_b = at::ones({no_heads}, rigid_rot_mats.options()); 
    at::Tensor head_weights = at::ones({no_heads}, rigid_rot_mats.options()); 
    at::Tensor linear_out_w = at::ones({c_s, no_heads * (c_hidden + no_v_points * 4 + c_z)}, s.options()); 
    at::Tensor linear_out_b = at::full({c_s}, e, rigid_rot_mats.options()); 
    at::Tensor mask = at::ones({n_res}, s.options()); 
    at::Tensor z = at::ones({n_res, n_res, c_z}, s.options()); 
    at::Tensor q = linear(s, linear_q_w, linear_q_b); 
    at::Tensor k = linear(s, linear_k_w, linear_k_b); 
    at::Tensor v = linear(s, linear_v_w, linear_v_b); 
    at::Tensor q_pts = linear(s, linear_q_points_w, linear_q_points_b); 
    q_pts = rigid_rot_vec_mul(q_pts, rigid_rot_mats, rigid_trans); 
    at::Tensor k_pts = linear(s, linear_k_points_w, linear_k_points_b); 
    k_pts = rigid_rot_vec_mul(k_pts, rigid_rot_mats, rigid_trans); 
    at::Tensor v_pts = linear(s, linear_v_points_w, linear_v_points_b); 
    v_pts = rigid_rot_vec_mul(v_pts, rigid_rot_mats, rigid_trans); 
    v_pts = v_pts.permute({1, 2, 3, 0}).contiguous(); 
    at::Tensor m = at::empty({no_heads, n_res, n_res}, s.options()); // b 
    at::Tensor n = at::empty({no_heads, n_res, n_res}, q.options()); // a 
    at::Tensor head_weights_2 = at::empty(head_weights.sizes(), head_weights.options()); 
    at::Tensor collect = at::empty({n_res, no_heads * (c_hidden + no_v_points * 4 + c_z)}, s.options()); 
    at::Tensor o = collect.narrow(1, 0, no_heads * c_hidden).view({n_res, no_heads, c_hidden}); 
    at::Tensor o_pt = collect.narrow(1, no_heads * c_hidden, no_heads * no_v_points * 3).view({n_res, 3, no_heads, no_v_points}); 
    at::Tensor o_pt_norm = collect.narrow(1, no_heads * (c_hidden + no_v_points * 3), no_heads * no_v_points).view({n_res, no_heads, no_v_points}); 
    at::Tensor o_pair = collect.narrow(1, no_heads * (c_hidden + no_v_points * 4), no_heads * c_z).view({n_res, no_heads, c_z}); 
    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 q_pts_tw = convert_to_tensor_wrapper(q_pts); 
    auto k_pts_tw = convert_to_tensor_wrapper(k_pts); 
    auto v_pts_tw = convert_to_tensor_wrapper(v_pts); 
    auto m_tw = convert_to_tensor_wrapper(m); 
    auto n_tw = convert_to_tensor_wrapper(n); 
    auto head_weights_tw = convert_to_tensor_wrapper(head_weights); 
    auto weights_head_weights_tw = convert_to_tensor_wrapper(head_weights_2); 
    auto o_tw = convert_to_tensor_wrapper(o); 
    auto o_pt_tw = convert_to_tensor_wrapper(o_pt); 
    auto o_pt_norm_tw = convert_to_tensor_wrapper(o_pt_norm); 
    auto o_pair_tw = convert_to_tensor_wrapper(o_pair); 
    auto z_tw = convert_to_tensor_wrapper(z); 
    auto rigid_rot_mats_tw = convert_to_tensor_wrapper(rigid_rot_mats); 
    auto rigid_trans_tw = convert_to_tensor_wrapper(rigid_trans); 
    auto mask_tw = convert_to_tensor_wrapper(mask); 
    auto linear_b_w_tw = convert_to_tensor_wrapper(linear_b_w); 
    auto linear_b_b_tw = convert_to_tensor_wrapper(linear_b_b); 
    kutacc_af2_ipa_weights_t_wrapper *ipa_weight_ptr = new kutacc_af2_ipa_weights_t_wrapper(head_weights_tw, weights_head_weights_tw, linear_b_w_tw, linear_b_b_tw, c_z, c_hidden, no_heads, 
        no_qk_points, no_v_points); 
    kutacc_af2_ipa_s_inputs_t_wrapper *ipa_s_ptrs = new kutacc_af2_ipa_s_inputs_t_wrapper(n_tw, m_tw, q_tw, k_tw, v_tw, q_pts_tw, k_pts_tw, v_pts_tw, n_res); 
    kutacc_af2_ipa_o_inputs_t_wrapper *ipa_o_ptrs = new kutacc_af2_ipa_o_inputs_t_wrapper(o_tw, o_pt_tw, o_pt_norm_tw, o_pair_tw); 
    if (unlikely(ipa_s_ptrs == nullptr || ipa_o_ptrs == nullptr || ipa_weight_ptr == nullptr)) { 
        return out; 
    } 
    kutacc_af2_invariant_point(ipa_s_ptrs, ipa_o_ptrs, z_tw.get_tensor(), rigid_rot_mats_tw.get_tensor(), rigid_trans_tw.get_tensor(), mask_tw.get_tensor(), ipa_weight_ptr); 
    out = linear(collect, linear_out_w, linear_out_b); 
    delete ipa_weight_ptr; 
    delete ipa_s_ptrs; 
    delete ipa_o_ptrs; 
    return out; 
} 
}

//test.py
import torch
import kpex._C as kernel
def test_invariant_point_attention(n_res, no_heads, c_hidden, no_qk_points, no_v_points, c_z, c_s): 
    out = kernel.alphafold.test_invariant_point_attention(n_res, no_heads, c_hidden, no_qk_points, no_v_points, c_z, c_s) 
    return out

//input : 4, 2, 2, 3, 8, 2, 8
//output:
>>> kpex.tpp.alphafold.alphafold.test_invariant_point_attention(4, 2, 2, 3, 8, 2, 8)
tensor([[3296., 3296., 3296., 3296., 3296., 3296., 3296., 3296.],
        [3296., 3296., 3296., 3296., 3296., 3296., 3296., 3296.],
        [3296., 3296., 3296., 3296., 3296., 3296., 3296., 3296.],
        [3296., 3296., 3296., 3296., 3296., 3296., 3296., 3296.]],
       dtype=torch.bfloat16)

Parent topic: AlphaFold2