Searched defs:seqlen_k (Results 1 – 4 of 4) sorted by relevance
| /external/pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/ |
| D | flash_api.cpp | 52 const size_t seqlen_k, in set_params_fprop()
171 const size_t seqlen_k, in set_params_dgrad()
392 const int seqlen_k = k.size(1); in mha_fwd() local
869 const int seqlen_k = k.size(1); in mha_bwd() local
1317 const int seqlen_k = !paged_KV ? kcache.size(1) : max_num_blocks_per_seq * page_block_size; in mha_fwd_kvcache() local
|
| D | flash.h | 69 …int b, seqlen_q, seqlen_k, seqlen_knew, d, seqlen_q_rounded, seqlen_k_rounded, d_rounded, rotary_d… member
|
| /external/pytorch/test/ |
| D | test_transformers.py | 2083 …def construct_local_mask(self, seqlen_q, seqlen_k, window_size, query_padding_mask, key_padding_ma… argument
2112 seqlen_k, argument
|
| /external/pytorch/aten/src/ATen/native/transformers/cuda/ |
| D | attention.cu | 970 const std::optional<at::Tensor>& seqlen_k, in _efficient_attention_forward()
|