Searched defs:softmax_scale (Results 1 – 5 of 5) sorted by relevance
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/ |
| H A D | flash_api.cpp | 70 float softmax_scale, in set_params_fprop()
167 void set_params_dgrad(Flash_bwd_params ¶ms, in set_params_dgrad()
356 const float softmax_scale, in mha_fwd()
555 const float softmax_scale, in mha_varlen_fwd()
818 const float softmax_scale, in mha_bwd()
1033 const float softmax_scale, in mha_varlen_bwd()
1265 const float softmax_scale, in mha_fwd_kvcache()
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/transformers/cuda/ |
| H A D | attention_backward.cu | 80 const auto softmax_scale = sdp::calculate_scale(query, scale).as_float_unchecked(); in _flash_attention_backward() local
219 const auto softmax_scale = sdp::calculate_scale(query, scale).as_float_unchecked(); in _scaled_dot_product_cudnn_attention_backward_cuda() local
422 const auto softmax_scale = sdp::calculate_scale(query, scale).as_float_unchecked(); in _efficient_attention_backward() local
|
| H A D | attention.cu | 766 const auto softmax_scale = sdp::calculate_scale(query, scale).as_float_unchecked(); in _scaled_dot_product_cudnn_attention_cuda() local
861 const auto softmax_scale = in _flash_attention_forward() local
1103 const auto softmax_scale = sdp::calculate_scale(query, scale).as_float_unchecked(); in _efficient_attention_forward() local
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/transformers/ |
| H A D | sdp_utils_cpp.h | 59 const auto softmax_scale = scale.has_value() in calculate_scale() local
|
| /aosp_15_r20/external/executorch/extension/llm/custom_ops/ |
| H A D | op_sdpa.cpp | 75 const auto softmax_scale = in calculate_scale() local
|