Home
last modified time | relevance | path

Searched defs:max_seqlen_k (Results 1 – 6 of 6) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/
H A Dmask.h39 const int max_seqlen_k, const int row_idx_offset, in apply_mask_local()
76 const int max_seqlen_k, const int row_idx_offset, in apply_mask_causal()
86 const int col_idx_offset_, const int max_seqlen_k, const int row_idx_offset) in apply_mask_causal_w_idx()
113 const int max_seqlen_k, max_seqlen_q; member
120 : max_seqlen_k(max_seqlen_k) in max_seqlen_k() function
H A Dflash_api.cpp298 const int num_heads, const int head_size, const int max_seqlen_k, const int max_seqlen_q, in set_params_splitkv()
553 const int max_seqlen_k, in mha_varlen_fwd()
1031 const int max_seqlen_k, // max sequence length to choose the kernel in mha_varlen_bwd()
H A Dalibi.h20 const int max_seqlen_k, max_seqlen_q; member
/aosp_15_r20/external/pytorch/torch/utils/
H A Dflop_counter.py333 max_seqlen_k, argument
413 max_seqlen_k, argument
512 max_seqlen_k, argument
/aosp_15_r20/external/pytorch/aten/src/ATen/native/transformers/cuda/
H A Dattention_backward.cu264 int64_t max_seqlen_k, in _efficient_attention_backward()
818 int64_t max_seqlen_k = k_t.size(1); in _scaled_dot_product_efficient_attention_backward_cuda() local
H A Dattention.cu995 int64_t max_seqlen_q = 0, max_seqlen_k = 0; in _efficient_attention_forward() local