Searched defs:num_head (Results 1 – 8 of 8) sorted by relevance
25 int64_t num_head) { in cpu_transform_bias_rescale_qkv()90 int64_t num_head) { in transform_bias_rescale_qkv_kernel_impl()
151 int64_t num_head, in reshape_attn_mask_to_4d()206 int64_t num_head = query.size(2); in cpu_flash_attention() local449 int64_t num_head = query.size(2); in cpu_flash_attention_backward() local
25 size_t num_head = 32; member
233 const int64_t num_head) { in transform_bias_rescale_qkv_cpu()266 const int64_t num_head, in native_multi_head_attention_cpu()782 int64_t num_head = query.size(1); in _scaled_dot_product_flash_attention_cpu() local857 const int64_t num_head, in triton_multi_head_attention()
289 auto num_head = params.query.sym_size(1); in check_attn_mask_shape() local
149 int num_head, num_saved, i; in tcp_update_sack_list() local
379 const int64_t num_head) { in transform_bias_rescale_qkv_cuda()484 const int64_t num_head, in native_multi_head_attention_cuda()
261 int64_t num_head = query.size(1); in cpu_flash_attention() local