bool use_neox_style = false) {

int elem_nums =

token_num * (q_head_num + kv_head_num) * dim_head; // just q and k

constexpr int PackSize = 16 / sizeof(T);

const int pack_num = elem_nums / PackSize;

const int blocksize = 128;

bool use_neox_style = false) {

int elem_nums =

token_num * (q_head_num + 2 * kv_head_num) * dim_head; // for all q k v

constexpr int PackSize = 16 / sizeof(T);

const int pack_num = elem_nums / PackSize;

const int blocksize = 128;

bool use_neox_style = false) {

int elem_nums =

token_num * (q_head_num + 2 * kv_head_num) * dim_head; // for all q k v

constexpr int PackSize = 16 / sizeof(T);

const int pack_num = elem_nums / PackSize;

const int blocksize = 128;