Fused SelfAttention query multipy key and value op (#4660)
* op & kernel * python api * mv bgemm to self_att*.cu * test case * fix bugs * update test * add op to amp gray list * mv to white list * rename bgemm to BatchedGemm * refine transpose * oneflow_api is deprecated * format
Showing
- oneflow/core/job_rewriter/auto_mixed_precision_lists.cpp 3 additions, 2 deletionsoneflow/core/job_rewriter/auto_mixed_precision_lists.cpp
- oneflow/python/ops/domain_ops.py 48 additions, 0 deletionsoneflow/python/ops/domain_ops.py
- oneflow/python/test/ops/test_fused_self_attention_query_mul_key_and_value.py 205 additions, 0 deletions.../ops/test_fused_self_attention_query_mul_key_and_value.py
- oneflow/user/kernels/fused_self_attention_query_mul_key_and_value_kernel.cu 300 additions, 0 deletions...ls/fused_self_attention_query_mul_key_and_value_kernel.cu
- oneflow/user/ops/fused_self_attention_query_mul_key_and_value_ops.cpp 142 additions, 0 deletions.../ops/fused_self_attention_query_mul_key_and_value_ops.cpp
Please register or sign in to comment