fix(attention): honor flash attention disable flag (#118)
Summary
Honor NVTE_FLASH_ATTN=0 before importing FlashAttention v2 modules.
Previously, TransformerEngine imported flash_attn_2_cuda whenever a
supported flash-attn package was installed, even when FlashAttention
was
explicitly disabled.
This caused import failures on platforms where the package metadata
existed
but the CUDA extension was unavailable, such as MUSA and Enflame.
Changes
Change the FlashAttention v2 import condition to require both:

a supported FlashAttention installation;
NVTE_FLASH_ATTN enabled.


Co-authored-by: canghaiX 59075364+canghaiX@users.noreply.github.com
Co-authored-by: wkhylyh-debug wkhylyh@gmail.com