Runtime check for the FlashAttention version rejects local version identifiers when it shouldn't

[lupreCSC]

Describe the bug

The runtime check for the FlashAttention version rejects local version identifiers when it shouldn't (i.e., an installation with version 2.8.3+local_version is rejected by the check in

TransformerEngine/transformer_engine/pytorch/attention/dot_product_attention/backends.py

Lines 102 to 106 in 8260f49

	if torch.cuda.is_available() and get_device_compute_capability() >= (10, 0):
	if fa_utils.version_required_blackwell <= fa_utils.version <= fa_utils.max_version:
	fa_utils.is_installed = True
	elif fa_utils.version_required <= fa_utils.version <= fa_utils.max_version:
	fa_utils.is_installed = True

). As local version identifiers leave the API compatible, this is too restrictive.

Expected behavior

An installation of FlashAttention with a compatible public version identifier (e.g. 2.8.3) should be accepted whether it also provides a local version identifier or not.

Suggested fix
Make the max version check use an exclusive upper bound, i.e.

• FlashAttentionUtils.max_version set to the next higher incompatible version (currently FlashAttentionUtils.max_version = PkgVersion("2.8.4"))
• The version check in transformer_engine/pytorch/attention/dot_product_attention/backends.py then would use < rather than <= to compare to fa_utils.max_version: fa_utils.version_required <= fa_utils.version < fa_utils.max_version.

This would admit local version identifiers for compatible releases (as indicated by the public version identifier) while still rejecting all incompatible versions of FlashAttention (e.g. 2.8.4, 2.9, etc)