--- title: kutacc_af2_invariant_point description: "利用注意力机制来保持特征点的不变性，从而提高模型在蛋白质结构预测中的性能。通过对输入特征进行加权聚合，模块能够有效地捕获特征之间的关系，并精确地调整蛋白质原子的坐标，从而提高结构预测的准确性和稳定性。" url: https://www.hikunpeng.com/document/detail/zh/kunpenghpcs/hpckit/devg/KunpengHPCKit_developer_141.html sourcePath: /source/zh/kunpenghpcs/hpckit/devg/KunpengHPCKit_developer_141.html indexId: bff9a51d5b05f6a0800caecaa00b364185f35140c9b8614d8d3cb180e2b067db67 --- # kutacc_af2_invariant_point 利用注意力机制来保持特征点的不变性，从而提高模型在蛋白质结构预测中的性能。通过对输入特征进行加权聚合，模块能够有效地捕获特征之间的关系，并精确地调整蛋白质原子的坐标，从而提高结构预测的准确性和稳定性。 #### 接口定义 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); #### 参数 表1 入参说明 | 参数名 | 类型 | 描述 | 输入/输出 | | --- | --- | --- | --- | | ipa\_s\_ptrs | kutacc\_af2\_ipa\_s\_inputs\_t \* | kutacc\_af2\_ipa\_s\_inputs\_t 类型的指针，具体定义见表2 kutacc\_af2\_ipa\_s\_inputs\_t数据结构表定义 | 输入 | | ipa\_o\_ptrs | kutacc\_af2\_ipa\_o\_inputs\_t \* | kutacc\_af2\_ipa\_o\_inputs\_t类型的指针，具体定义见表3 kutacc\_af2\_ipa\_o\_inputs\_t数据结构表定义 | 输入 | | z | kutacc\_tensor\_h | 中间变量 | 输入 | | rigid\_rot\_mats | kutacc\_tensor\_h | 刚体变换：旋转参数 | 输入 | | rigid\_trans | kutacc\_tensor\_h | 刚体变换：平移参数 | 输入 | | mask | kutacc\_tensor\_h | 掩码，用于忽略无效残基 | 输入 | | ipa\_weight\_ptrs | kutacc\_af2\_ipa\_weights\_t \* | kutacc\_af2\_ipa\_weights\_t类型的指针，具体定义见表4 kutacc\_af2\_ipa\_weights\_t数据结构表定义 | 输入 | 表2 kutacc_af2_ipa_s_inputs_t数据结构表定义 | 参数名 | 类型 | 描述 | 输入/输出 | | --- | --- | --- | --- | | n\_res | int64\_t | 输入向量s的第0维长度 | 输入 | | a | kutacc\_tensor\_h | 用于保存q转置与k的矩阵乘，进行softmax后的临时结果 | 输入 | | b | kutacc\_tensor\_h | 用于保存z向量进行线性变换后的临时结果 | 输入 | | q | kutacc\_tensor\_h | 输入s经过线性变化生成的Q矩阵 | 输入 | | k | kutacc\_tensor\_h | K矩阵 | 输入 | | v | kutacc\_tensor\_h | V矩阵 | 输入 | | q\_pts | kutacc\_tensor\_h | 查询矩阵的3D点 | 输入 | | k\_pts | kutacc\_tensor\_h | 键矩阵的3D点 | 输入 | | v\_pts | kutacc\_tensor\_h | 值矩阵的3D点 | 输入 | 表3 kutacc_af2_ipa_o_inputs_t数据结构表定义 | 参数名 | 类型 | 描述 | 输入/输出 | | --- | --- | --- | --- | | o | kutacc\_tensor\_h | 用于保存qkv，attention计算结果的向量 | 输入 | | o\_pt | kutacc\_tensor\_h | 用于保存pts相关向量attention计算结果的向量 | 输入 | | o\_pt\_norm | kutacc\_tensor\_h | 用于保存需要进行layernorm变换的向量 | 输入 | | o\_pair | kutacc\_tensor\_h | 残基对特征 | 输入 | 表4 kutacc_af2_ipa_weights_t数据结构表定义 | 参数名 | 类型 | 描述 | 输入/输出 | | --- | --- | --- | --- | | c\_z | int64\_t | 单一single的长度 | 输入 | | c\_hidden | int64\_t | 隐藏层长度 | 输入 | | no\_heads | int64\_t | attention头数量 | 输入 | | no\_qk\_points | int64\_t | 需要生成的q/k数量 | 输入 | | no\_v\_points | int64\_t | 需要生成的v数量 | 输入 | | head\_weights | kutacc\_tensor\_h | 注意力头的权重 | 输入 | | weights\_head\_weights | kutacc\_tensor\_h | 注意力头的权重 | 输入 | | linear\_b\_w | kutacc\_tensor\_h | 对b进行线性变化时的权重 | 输入 | | linear\_b\_b | kutacc\_tensor\_h | 对b进行线性变化时的偏置 | 输入 | 整数参数应满足的约束关系： 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 #include #include #include #include #include #include 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.6 f; float b = 0.25 f; float c = 0.1 f; float d = 0.2 f; float e = 0.5 f; 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) ```