saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-05-22 22:01:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix causual flash attention related kernel run (#14299)

Browse source
This commit is contained in:
Zhang Lei 2023-01-13 21:40:22 -08:00 committed by GitHub
parent 8824f812e0
commit bd39c8f35e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.cu
View file

@ -159,7 +159,7 @@ class FusedMHARunnerFP16v2::mhaImpl {
params.o_ptr = output;
params.cu_seqlens = static_cast<int*>(const_cast<void*>(cu_seqlens));
if (use_flash_attention && flash_attention_kernel != nullptr) {
if (use_flash_attention && flash_attention_kernel != nullptr && !has_causal_mask) {
flash_attention_kernel->run(params, stream);
} else {
xmmaKernel->run(params, stream, use_flash_attention, has_causal_mask);