Error encountered while using flash attention in TensorFlow

[monatis]

Hi,
I was trying to test the flash attention implementation, but it gives an error:

Traceback (most recent call last):
  File "test.py", line 30, in <module>
    output = test_fmha(queries, keys, values)
  File "test.py", line 17, in test_fmha
    return fmha(q, k, v)
  File "/usr/local/lib/python3.8/dist-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/v/tfplus_src/build/lib.linux-x86_64-cpython-38/tfplus/flash_attn/python/ops/flash_attn_ops.py", line 170, in call
    attn_weight = gen_flash_attention_ops.fmha_forward(
AttributeError: Exception encountered when calling layer 'flash_attention_layer' (type FlashAttentionLayer).
module '4356a221e8f1a9ebe954e817a4687d209fe7f94a' has no attribute 'fmha_forward'
Call arguments received by layer 'flash_attention_layer' (type FlashAttentionLayer):
  • query=tf.Tensor(shape=(2, 4, 8, 64), dtype=bfloat16)
  • key=tf.Tensor(shape=(2, 4, 8, 64), dtype=bfloat16)
  • value=tf.Tensor(shape=(2, 4, 8, 64), dtype=bfloat16)
  • mask=None
  • kwargs={'training': 'None'}

Seems like the native library is not properly built or something. Any hint about it?