Update bert_layers.py

bert_layers.py

@@ -27,8 +27,9 @@ from bert_padding import (index_first_axis, index_put_first_axis, pad_input, unp
 #from bert_padding_module import (index_first_axis, index_put_first_axis, pad_input, unpad_input, unpad_input_only)
 
 try:
-    from .flash_attn_triton import flash_attn_qkvpacked_func
-    #flash_attn_qkvpacked_func = flash_attn_triton.flash_attn_qkvpacked_func
+    # Force disable triton flash attn due to API incompatibility (trans_b argument) or dtype issues
+    # from .flash_attn_triton import flash_attn_qkvpacked_func
+    flash_attn_qkvpacked_func = None
 except ImportError as e:
     flash_attn_qkvpacked_func = None