| /third_party/mindspore/mindspore-src/source/tests/ut/python/parallel/ |
| D | test_flash_attention_score_info.py | 38 def generate_inputs(B, N, S, D, input_layout, use_mqa=False, with_real_shift=True, sparse_mode=0): argument
102 def __init__(self, head_num, keep_prob=0.9, input_layout="BSH", sparse_mode=0, use_mqa=False, argument
201 def test_self_attention_standalone_with_compressed_mask(input_layout, sparse_mode): argument
242 def test_flash_attention_semi_auto_parallel_with_compressed_mask(input_layout, sparse_mode): argument
322 def test_flash_attention_compressed_mask_with_seq_parallel(input_layout, sparse_mode): argument
378 def test_flash_attention_compressed_mask_with_load_balance(input_layout, sparse_mode): argument
|
| D | test_prompt_flash_attention.py | 29 def generate_inputs(dims, optinal_inputs, input_layout='BSH', sparse_mode=0): argument
72 def generate_strategy(dp, mp, optinal_inputs, input_layout='BSH', sparse_mode=0, sp=1): argument
116 … num_key_value_heads=0, strategy=None, sparse_mode=0, set_atten_mask_as_constant=False): argument
|
| D | test_fused_infer_attention_score.py | 28 def generate_inputs(dims, optinal_inputs, input_layout='BSH', sparse_mode=0, is_ifa=False): argument
92 def generate_strategy(dp, mp, optinal_inputs, input_layout='BSH', sparse_mode=0, sp=1, is_ifa=False… argument
152 num_key_value_heads=0, strategy=None, sparse_mode=0, inner_precise=1, block_size=0, argument
|
| D | test_ring_attention.py | 33 def generate_inputs(B, N, S, D, input_layout, use_mqa=False, with_real_shift=False, sparse_mode=0): argument
81 def __init__(self, head_num, keep_prob=1.0, input_layout="BSH", sparse_mode=0, use_mqa=False, argument
|
| D | test_model_flops_cal.py | 284 def generate_inputs(B, N, S, D, input_layout, use_mqa=False, with_real_shift=True, sparse_mode=0): argument
340 … def __init__(self, head_num, keep_prob=0.9, input_layout="BSH", sparse_mode=0, use_mqa=False, argument
|
| /third_party/mindspore/mindspore-src/source/tests/st/auto_parallel/ |
| D | test_prompt_flash_attention_info.py | 30 def generate_inputs(dims, optinal_inputs, input_layout='BSH', sparse_mode=0): argument
87 def generate_strategy(dp, mp, optinal_inputs, input_layout='BSH', sparse_mode=0, sp=1): argument
123 … num_key_value_heads=0, strategy=None, sparse_mode=0, set_atten_mask_as_constant=False): argument
|
| D | flash_attention_score.py | 29 def __init__(self, head_num, keep_prob=0.9, input_layout="BSH", sparse_mode=0, use_mqa=False, argument
|
| /third_party/mindspore/mindspore-src/source/mindspore/ccsrc/plugin/device/ascend/kernel/opapi/aclnn/ |
| D | fused_infer_attention_score_aclnn_kernel.cc | 65 auto sparse_mode = transform::ConvertKernelTensor<int64_t>(inputs[real_input_idx_[kIndex23]]); in GetWorkSpaceInfo() local
119 auto sparse_mode = transform::ConvertKernelTensor<int64_t>(inputs[real_input_idx_[kIndex23]]); in Launch() local
|
| D | flash_attention_score_grad_aclnn_kernel.h | 94 auto sparse_mode = inputs[kIndex22]; in DEFINE_GET_WORKSPACE_FOR_RESIZE() local
|
| D | flash_attention_score_aclnn_kernel.cc | 65 auto sparse_mode = inputs[kIndex17]; in GetWorkSpaceInfo() local
|
| D | flash_attention_score_grad_aclnn_kernel.cc | 65 auto sparse_mode = inputs[kIndex22]; in GetWorkSpaceInfo() local
|
| D | flash_attention_score_aclnn_kernel.h | 94 auto sparse_mode = inputs[kIndex17]; in DEFINE_GET_WORKSPACE_FOR_RESIZE() local
|
| /third_party/mindspore/mindspore-src/source/mindspore/lite/tools/optimizer/fusion/ |
| D | flash_attention_fusion.h | 33 int sparse_mode = 0; member
|
| /third_party/mindspore/mindspore-src/source/mindspore/core/ops/ops_func_impl/ |
| D | flash_attention_score.cc | 79 … int64_t sparse_mode, int64_t batch_size, int64_t q_head_num, in CheckFlashAttentionScoreAttnMaskShape()
99 …ttentionScorePrefix(const AbstractBasePtr &prefix, const std::string &op_name, int64_t sparse_mode, in CheckFlashAttentionScorePrefix()
292 auto sparse_mode = sparse_mode_opt.value(); in InferShape() local
|
| D | flash_attention_score_grad.cc | 125 … int64_t sparse_mode, int64_t batch_size, int64_t q_head_num, in CheckFlashAttentionScoreGradAttnMaskShape()
143 …tionScoreGradPrefix(const AbstractBasePtr &prefix, const std::string &op_name, int64_t sparse_mode, in CheckFlashAttentionScoreGradPrefix()
297 auto sparse_mode = sparse_mode_opt.value(); in InferShape() local
|
| /third_party/mindspore/mindspore-src/source/mindspore/ccsrc/plugin/device/ascend/kernel/pyboost/customize/ |
| D | flash_attention_score.cc | 51 const Int64ImmPtr sparse_mode, const std::vector<tensor::BaseTensorPtr> &outputs) { in FlashAttentionScoreAscendCall()
105 const Int64ImmPtr inner_precise, const Int64ImmPtr input_layout, const Int64ImmPtr sparse_mode) { in FlashAttentionScoreAscendCustomize()
|
| D | flash_attention_score_grad.cc | 39 const Int64ImmPtr inner_precise, const Int64ImmPtr input_layout, const Int64ImmPtr sparse_mode, in FlashAttentionScoreGradAscendCall()
93 const Int64ImmPtr inner_precise, const Int64ImmPtr input_layout, const Int64ImmPtr sparse_mode) { in FlashAttentionScoreGradAscendCustomize()
|
| D | fused_infer_attention_score.cc | 40 …const Int64ImmPtr &num_key_value_heads, const Int64ImmPtr &sparse_mode, const Int64ImmPtr &inner_p… in FusedInferAttentionScoreAscendCustomize()
|
| /third_party/mindspore/mindspore-src/source/tests/st/ops/ascend/ |
| D | test_prompt_flash_attention.py | 31 … pre_tokens=2147483547, next_tokens=0, input_layout='BSH', num_key_value_heads=0, sparse_mode=0): argument
|
| D | test_fused_infer_attention_score.py | 32 … pre_tokens=2147483547, next_tokens=0, input_layout='BSH', num_key_value_heads=0, sparse_mode=0, argument
63 pre_tokens=2147483547, next_tokens=0, sparse_mode=0, inner_precise=0): argument
|
| /third_party/mindspore/mindspore-src/source/mindspore/core/ops/ |
| D | prompt_flash_attention.cc | 115 auto sparse_mode = GetValue<int64_t>(primitive->GetAttr("sparse_mode")); in CheckOptinalInputShape() local
234 auto sparse_mode = GetValue<int64_t>(primitive->GetAttr("sparse_mode")); in InferShapeBSH() local
280 auto sparse_mode = GetValue<int64_t>(primitive->GetAttr("sparse_mode")); in InferShapeBNSD() local
|
| /third_party/mindspore/mindspore-src/source/tests/st/ops/ |
| D | test_ops_flash_attention_score.py | 31 … next_tokens=65536, scalar_value=1.0, inner_precise=0, sparse_mode=0): argument
95 next_tokens=65536, inner_precise=0, sparse_mode=0): argument
|
| /third_party/mindspore/mindspore-src/source/mindspore/core/ops/auto_generate/ |
| D | gen_lite_ops.cc | 811 void FlashAttentionScoreGrad::set_sparse_mode(const int64_t &sparse_mode) { (void)this-… in set_sparse_mode()
846 void FlashAttentionScore::set_sparse_mode(const int64_t &sparse_mode) { (void)this->Add… in set_sparse_mode()
2196 void FusedInferAttentionScore::set_sparse_mode(const int64_t &sparse_mode) { (void)this… in set_sparse_mode()
|
| /third_party/mindspore/mindspore-src/source/mindspore/python/mindspore/ops/operations/manually_defined/ |
| D | ops_def.py | 1874 input_layout='BSH', sparse_mode=0): argument
|
| /third_party/mindspore/mindspore-src/source/mindspore/python/mindspore/ops/operations/ |
| D | nn_ops.py | 9606 num_key_value_heads=0, sparse_mode=0, inner_precise=1): argument
|