Flash Attention Integration Tests Fail

[weiT1993]

Describe the bug

Flash attention integration tests failed

Expected Behavior

Flash attention integration correctness pytest should pass.

Current Behavior

FAILED flash_attention_correctness.py::test_attention[True-dtype0] - TypeError: multiple values for argument 'softmax_scale'
FAILED flash_attention_correctness.py::test_attention[True-dtype1] - TypeError: multiple values for argument 'softmax_scale'
FAILED flash_attention_correctness.py::test_attention[False-dtype0] - TypeError: multiple values for argument 'softmax_scale'
FAILED flash_attention_correctness.py::test_attention[False-dtype1] - TypeError: multiple values for argument 'softmax_scale'

Reproduction Steps

Running

pytest flash_attention_correctness.py

from test/integration/flash_attention.

Regression Issue

• Select this option if this issue appears to be a regression.

Possible Solution

It is currently calling compiler NKI kernel. Upgrade to the one from the repo?

Additional Information/Context

No response

neuronx-cc version used

2.21

Framework(s) and their versions used (JAX, PyTorch, etc..)

No response

[hannanjgaws]

Hi weiT1993:

Thank you for opening this ticket. We are currently working on reproducing this issue, and will share an update when we've made progress debugging this.

[mrkcath-aws]

We were able to reproduce this problem and are working on a fix.

[mrkcath-aws]

Please see NKI Library for the latest!

[github-actions]

This issue is now closed. Comments on closed issues are hard for our team to see.
If you need more assistance, please open a new issue that references this one.