Remove typedef FAUtils for FlashAttentionUtils

Use attn_log instead of att_log

Signed-off-by: Kshitij Janardan Lakhani <[email protected]>

Fix lint error

Signed-off-by: Kshitij Janardan Lakhani <[email protected]>

Author: KshitijLakhani

tests/pytorch/fused_attn/test_fused_attn.py

@@ -21,7 +21,7 @@
     _attention_backends,
 )
 from transformer_engine.pytorch.dot_product_attention.utils import (
-    FAUtils,
+    FlashAttentionUtils,
     get_attention_backend,
     check_set_window_size,
     AttentionParams,
@@ -280,12 +280,12 @@ def test_dot_product_attention(
     # mannually pads and unpads the input and output of FlashAttention for testing purposes
     if (
         pad_between_seqs
-        and FAUtils.is_installed
+        and FlashAttentionUtils.is_installed
         and not (
             config.max_seqlen_q != config.max_seqlen_kv
             and config.attn_mask_type in ["causal", "padding_causal"]
         )
-        and (config.window_size[0] == -1 or FAUtils.v2_3_plus)
+        and (config.window_size[0] == -1 or FlashAttentionUtils.v2_3_plus)
     ):
         flash_attn_supported = True
 
@@ -592,7 +592,7 @@ def test_dpa_bias_shapes(dtype, model_configs, model):
 }
 
 
-@pytest.mark.skipif(not FAUtils.v2_3_plus, reason="Flash-attn 2.3+ is required.")
+@pytest.mark.skipif(not FlashAttentionUtils.v2_3_plus, reason="Flash-attn 2.3+ is required.")
 @pytest.mark.parametrize("dtype", param_types_lean)
 @pytest.mark.parametrize("model_configs", [model_configs_swa])
 @pytest.mark.parametrize("model", model_configs_swa.keys())
@@ -614,7 +614,7 @@ def test_dpa_sliding_window(dtype, model_configs, model):
 }
 
 
-@pytest.mark.skipif(not FAUtils.v2_3_plus, reason="Flash-attn 2.3+ is required.")
+@pytest.mark.skipif(not FlashAttentionUtils.v2_3_plus, reason="Flash-attn 2.3+ is required.")
 @pytest.mark.parametrize("dtype", param_types_lean)
 @pytest.mark.parametrize("model_configs", [model_configs_alibi_slopes])
 @pytest.mark.parametrize("model", model_configs_alibi_slopes.keys())
@@ -1456,7 +1456,7 @@ def test_mha_fp8_vs_f16(dtype, model, qkv_format, input_layernorm, fp8_dpa_bwd,
     ):
         pytest.skip("FP8 with padding or head_dim != 128 is not supported for cuDNN < 9.7")
 
-    if FAUtils.v3_is_installed and not is_training and "padding" not in config.attn_mask_type:
+    if FlashAttentionUtils.v3_is_installed and not is_training and "padding" not in config.attn_mask_type:
         os.environ["NVTE_FLASH_ATTN"] = "1"
         os.environ["NVTE_FUSED_ATTN"] = "0"
         _attention_backends["backend_selection_requires_update"] = True
@@ -1482,7 +1482,7 @@ def test_mha_fp8_vs_f16(dtype, model, qkv_format, input_layernorm, fp8_dpa_bwd,
     rtol = 5e-1
     rmse_tol = 0.15
     logging.debug("========== {:^25s} ==========".format("forward output"))
-    if FAUtils.v3_is_installed and not is_training and "padding" not in config.attn_mask_type:
+    if FlashAttentionUtils.v3_is_installed and not is_training and "padding" not in config.attn_mask_type:
         _error(
             flash_attn_fwd_fp8,
             fused_attn_fwd_f16,
@@ -1667,7 +1667,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training):
     os.environ["NVTE_FP8_DPA_BWD"] = "1" if fp8_dpa_bwd else "0"
     os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "1"
 
-    if FAUtils.v3_is_installed and not is_training and "padding" not in config.attn_mask_type:
+    if FlashAttentionUtils.v3_is_installed and not is_training and "padding" not in config.attn_mask_type:
         os.environ["NVTE_FLASH_ATTN"] = "1"
         os.environ["NVTE_FUSED_ATTN"] = "0"
         _attention_backends["backend_selection_requires_update"] = True
@@ -1696,7 +1696,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training):
     rmse_tol = 0.11
     bwd_names = ["dq", "dk", "dv"]
     logging.debug("========== {:^25s} ==========".format("forward output"))
-    if FAUtils.v3_is_installed and not is_training and "padding" not in config.attn_mask_type:
+    if FlashAttentionUtils.v3_is_installed and not is_training and "padding" not in config.attn_mask_type:
         _error(
             flash_attn_fwd_fp8,
             fused_attn_fwd_f16,

tests/pytorch/fused_attn/test_fused_attn_with_cp.py

@@ -11,7 +11,7 @@
     get_device_compute_capability,
     get_cudnn_version,
 )
-from transformer_engine.pytorch.dot_product_attention.utils import FAUtils
+from transformer_engine.pytorch.dot_product_attention.utils import FlashAttentionUtils
 from test_fused_attn import ModelConfig
 
 model_configs_flash_attn = {
@@ -50,7 +50,7 @@ def get_bash_arguments(num_gpus_per_node, **kwargs):
     return args
 
 
-@pytest.mark.skipif(not FAUtils.v2_plus, reason="Flash-attn 2.0+ is required.")
+@pytest.mark.skipif(not FlashAttentionUtils.v2_plus, reason="Flash-attn 2.0+ is required.")
 @pytest.mark.skipif(get_device_compute_capability() < (8, 0), reason="CP tests require sm80+.")
 @pytest.mark.parametrize("dtype", ["bf16", "fp16"])
 @pytest.mark.parametrize("model", model_configs_flash_attn.keys())

transformer_engine/pytorch/attention.py

@@ -78,12 +78,12 @@
 import transformer_engine.pytorch.dot_product_attention.utils as dpa_utils
 import transformer_engine.pytorch.dot_product_attention.inference as dpa_infer
 from transformer_engine.pytorch.dot_product_attention.utils import FlashAttentionUtils as fa_utils
-from transformer_engine.pytorch.dot_product_attention.utils import AttentionLogging as att_log
+from transformer_engine.pytorch.dot_product_attention.utils import AttentionLogging as attn_log
 from transformer_engine.pytorch.dot_product_attention.rope import apply_rotary_pos_emb
 
 
 # Setup Attention Logging
-att_log.setup_logging()
+attn_log.setup_logging()
 
 # Global vars for flash attn imports
 flash_attn_cuda_bwd = None
@@ -101,7 +101,7 @@
         and get_device_compute_capability() >= (8, 0)
         and dpa_utils._NVTE_FLASH_ATTN
     ):
-        att_log.fa_logger.debug(
+        attn_log.fa_logger.debug(
             "flash-attn v2 is not installed. To use, please install it by"
             """ "pip3 install flash-attn".""",
         )
@@ -131,7 +131,7 @@
         and get_device_compute_capability() >= (8, 0)
         and dpa_utils._NVTE_FLASH_ATTN
     ):
-        att_log.fa_logger.warning(
+        attn_log.fa_logger.warning(
             "Supported flash-attn versions are %s. Found flash-attn %s.",
             dpa_utils._get_supported_versions(
                 (
@@ -155,7 +155,7 @@
         and get_device_compute_capability() >= (9, 0)
         and dpa_utils._NVTE_FLASH_ATTN
     ):
-        att_log.fa_logger.debug(
+        attn_log.fa_logger.debug(
             "flash-attn v3 is not installed. To use, please install it by \n%s",
             fa_utils.v3_installation_steps,
         )
@@ -3926,9 +3926,9 @@ def __init__(
         self.layer_number = 1 if layer_number is None else layer_number
         self.deterministic = deterministic
         self.logger = logging.getLogger("FlashAttention")
-        self.logger.setLevel(att_log._log_level)
+        self.logger.setLevel(attn_log._log_level)
         if not self.logger.hasHandlers():
-            self.logger.addHandler(att_log._stream_handler)
+            self.logger.addHandler(attn_log._stream_handler)
 
     def forward(
         self,
@@ -5157,9 +5157,9 @@ def __init__(
         super().__init__()
 
         self.logger = logging.getLogger("DotProductAttention")
-        self.logger.setLevel(att_log._log_level)
+        self.logger.setLevel(attn_log._log_level)
         if not self.logger.hasHandlers():
-            self.logger.addHandler(att_log._stream_handler)
+            self.logger.addHandler(attn_log._stream_handler)
         self.qkv_format = qkv_format
         attn_mask_type = attn_mask_type.replace(",", "_")
         if attn_mask_type == "causal_padding":