Searched defs:Is_even_K (Results 1 – 2 of 2) sorted by relevance
| /external/pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/ |
| D | flash_fwd_launch_template.h | 35 …out, bool Is_causal, bool Is_local, bool Has_alibi, bool Is_even_MN, bool Is_even_K, bool Return_s… in DEFINE_FLASH_FORWARD_KERNEL()
44 …nel, bool Is_causal, bool Is_local, bool Has_alibi, bool Is_even_MN, bool Is_even_K, bool Split, b… in DEFINE_FLASH_FORWARD_KERNEL()
52 …FORWARD_KERNEL(flash_fwd_splitkv_combine_kernel, int kBlockM, int Log_max_splits, bool Is_even_K) { in DEFINE_FLASH_FORWARD_KERNEL()
|
| D | flash_bwd_launch_template.h | 36 …_dv_loop_kernel, bool Is_dropout, bool Is_causal, bool Has_alibi, bool Is_even_M, bool Is_even_K) { in DEFINE_FLASH_BACKWARD_KERNEL()
44 … bool Is_dropout, bool Is_causal, bool Is_local, bool Has_alibi, bool Is_even_MN, bool Is_even_K) { in DEFINE_FLASH_BACKWARD_KERNEL()
|