Refactor attention.py part 2 (#1704)

* Move MultiHeadAttention into its own file. Modify tests and files in t_e/pytorch to import from the new MHA module

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

* Resolving lost MHA changes from PR 1614 as a result of rebase

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

* Move context parallelism code into it's own file. Modify test and local imports of cp code accordingly

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

* Move softmax.py frm pytorch/ to pytorch/d_p_a

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

* Move Unfused and Fused attention to backends.py and some utils functions to pytorch/utils.py

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

* Resolving lost mark_activation_offload changes from PR 1678 as a result of rebase

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

* Code clean up

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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Refactor attention dir

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

* Refactor dir structure. Make relevant symbols public in __init__ for attention and d_p_a dirs
Move FA package imports to backends.py
Code cleanup

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

* Modify tests to import attention modules correctly

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

* Lint fixes

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

* Code clean up and fix typo

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

* Allowing InferenceParams and RoPE imports from attention module and pytorch module

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

* Allow InferenceParams and RoPE imports via transformer_engine.pytorch and transformer_engine.pytorch.attention modules
Remove unnecessary checks for check_set_window_size in MHA and TL
Reorder backends such that smaller classes at the start and larger ones at the end
Code clean up

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

* Reinstating changes from PR 1478 for rope.py lost during rebase conflict resolution

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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix lint issues

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

* nit: Code clean up

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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Make imports leaner

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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Kshitij Janardan Lakhani <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: KshitijLakhani

docs/examples/attention/attention.ipynb

@@ -458,7 +458,7 @@
     "  </tr>\n",
     "</table>\n",
     "\n",
-    "Some example usage of the different layouts can be found at [test_dpa_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/fused_attn/test_fused_attn.py) and [test_dpa_qkv_layout_thd](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/fused_attn/test_fused_attn.py). Transformer Engine also provides a utility function [transformer_engine.pytorch.dot_product_attention.utils.get_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py) to help determine which layout a set of `q`, `k`, `v` tensors have (PyTorch only).\n",
+    "Some example usage of the different layouts can be found at [test_dpa_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/fused_attn/test_fused_attn.py) and [test_dpa_qkv_layout_thd](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/fused_attn/test_fused_attn.py). Transformer Engine also provides a utility function [transformer_engine.pytorch.attention.dot_product_attention.utils.get_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py) to help determine which layout a set of `q`, `k`, `v` tensors have (PyTorch only).\n",
     "\n",
     "<div class=\"alert alert-info\">\n",
     "<b>Note</b>\n",

docs/examples/te_llama/te_llama.py

@@ -8,11 +8,9 @@
 from contextlib import contextmanager
 
 import torch
-from torch import nn
 
 import transformer_engine as te
-from transformer_engine.pytorch.dot_product_attention.rope import RotaryPositionEmbedding
-from transformer_engine.pytorch.fp8 import fp8_model_init
+from transformer_engine.pytorch.attention import RotaryPositionEmbedding
 
 import transformers
 from transformers.models.llama.modeling_llama import (

tests/pytorch/fused_attn/run_fused_attn_with_cp.py

@@ -2,12 +2,16 @@
 #
 # See LICENSE for license information.
 
-import os, sys, logging
+import os
+import sys
+import logging
 from contextlib import nullcontext
 import torch
 import torch.distributed as dist
 from transformer_engine.pytorch.attention import DotProductAttention
-from transformer_engine.pytorch.attention import get_cu_seqlens_on_cp_rank
+from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
+    get_cu_seqlens_on_cp_rank,
+)
 import transformer_engine_torch as tex
 from test_fused_attn_with_cp import model_configs_flash_attn, model_configs_fused_attn
 from transformer_engine.pytorch.fp8 import fp8_autocast

tests/pytorch/fused_attn/test_fused_attn.py

@@ -1,12 +1,9 @@
 # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
-
-import functools
 import logging
 import math
 import os
-from importlib.metadata import version
 from typing import Any, Dict, List, Tuple, Union, Optional
 from contextlib import contextmanager
 
@@ -15,26 +12,22 @@
 
 from transformer_engine.common import recipe
 from transformer_engine.pytorch import TransformerLayer, fp8_autocast, fp8_model_init
-from transformer_engine.pytorch.attention import (
+from transformer_engine.pytorch.attention.dot_product_attention import (
     DotProductAttention,
-    MultiheadAttention,
     _attention_backends,
 )
-from transformer_engine.pytorch.dot_product_attention.utils import (
+from transformer_engine.pytorch.attention.multi_head_attention import MultiheadAttention
+from transformer_engine.pytorch.attention.dot_product_attention.utils import (
     FlashAttentionUtils,
     get_attention_backend,
     check_set_window_size,
     AttentionParams,
 )
-from transformer_engine.pytorch.dot_product_attention.inference import InferenceParams
-from transformer_engine.pytorch.dot_product_attention.rope import RotaryPositionEmbedding
-from transformer_engine.pytorch.constants import TE_DType
+from transformer_engine.pytorch.attention import InferenceParams
+from transformer_engine.pytorch.attention import RotaryPositionEmbedding
 import transformer_engine.pytorch.cpp_extensions as ext
 from transformer_engine.pytorch.cpp_extensions.fused_attn import (
-    AttnBiasType,
-    AttnMaskType,
     FusedAttnBackend,
-    QKVLayout,
     fused_attn_bwd,
     fused_attn_fwd,
 )
@@ -49,9 +42,7 @@
 )
 from transformer_engine.pytorch.utils import get_cudnn_version
 import transformer_engine_torch as tex
-from transformer_engine_torch import NVTE_Fused_Attn_Backend
 from transformer_engine.pytorch.tensor.quantized_tensor import (
-    QuantizedTensor,
     Quantizer,
     prepare_for_saving,
     restore_from_saved,

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 FlashAttentionUtils
+from transformer_engine.pytorch.attention.dot_product_attention.utils import FlashAttentionUtils
 from test_fused_attn import ModelConfig
 
 model_configs_flash_attn = {

tests/pytorch/fused_attn/test_kv_cache.py

@@ -11,27 +11,28 @@
 import pytest
 import torch
 
+from test_fused_attn import (
+    ModelConfig,
+    reset_rng_states,
+    _get_attention_backends,
+)
+
 from torch.distributions import Exponential
 from transformer_engine.pytorch import make_graphed_callables
 from transformer_engine.common import recipe
 from transformer_engine.pytorch import fp8_autocast, fp8_model_init
 from transformer_engine.pytorch.transformer import (
     TransformerLayer,
 )
-from transformer_engine.pytorch.attention import DotProductAttention
-from transformer_engine.pytorch.dot_product_attention.inference import InferenceParams
-from transformer_engine.pytorch.dot_product_attention.utils import FlashAttentionUtils as fa_utils
+from transformer_engine.pytorch.attention import DotProductAttention, InferenceParams
+from transformer_engine.pytorch.attention.dot_product_attention.utils import (
+    FlashAttentionUtils as fa_utils,
+)
 from transformer_engine.pytorch.utils import (
-    get_device_compute_capability,
     init_method_normal,
     scaled_init_method_normal,
     is_bf16_compatible,
 )
-from test_fused_attn import (
-    ModelConfig,
-    reset_rng_states,
-    _get_attention_backends,
-)
 
 # Initialize RNG state
 seed = 1234

tests/pytorch/test_fused_optimizer.py

@@ -12,7 +12,7 @@
 from torch.testing._internal.common_device_type import largeTensorTest
 import transformer_engine.pytorch as te
 from transformer_engine.common.recipe import DelayedScaling
-from transformer_engine.pytorch.attention import MultiheadAttention
+from transformer_engine.pytorch.attention.multi_head_attention import MultiheadAttention
 from transformer_engine.pytorch import fp8_model_init
 from transformer_engine.pytorch.utils import is_bf16_compatible
 from transformer_engine.pytorch.fp8 import FP8GlobalStateManager

tests/pytorch/test_fused_rope.py

@@ -1,11 +1,11 @@
 # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
+from typing import Callable, Tuple, Union
 import math
-import pytest
 import torch
-from typing import Callable, Tuple, Union
-from transformer_engine.pytorch.dot_product_attention.rope import (
+import pytest
+from transformer_engine.pytorch.attention.rope import (
     RotaryPositionEmbedding,
     apply_rotary_pos_emb,
 )

tests/pytorch/test_numerics.py

@@ -7,7 +7,6 @@
 import os
 from typing import Dict, List, Tuple, Optional
 import pytest
-import copy
 import random
 
 import torch
@@ -38,7 +37,7 @@
     Fp8Padding,
     Fp8Unpadding,
 )
-from transformer_engine.pytorch.dot_product_attention.inference import InferenceParams
+from transformer_engine.pytorch.attention.inference import InferenceParams
 from transformer_engine.pytorch.distributed import checkpoint as te_checkpoint
 from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm
 from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer

transformer_engine/pytorch/__init__.py

@@ -90,7 +90,8 @@ def _load_library():
 from transformer_engine.pytorch.module import destroy_ub
 from transformer_engine.pytorch.attention import DotProductAttention
 from transformer_engine.pytorch.attention import MultiheadAttention
-from transformer_engine.pytorch.dot_product_attention.inference import InferenceParams
+from transformer_engine.pytorch.attention import InferenceParams
+from transformer_engine.pytorch.attention import RotaryPositionEmbedding
 from transformer_engine.pytorch.transformer import TransformerLayer
 from transformer_engine.pytorch.permutation import (
     moe_permute,