From 69f568e5e9ff951438d0c9bc43a858a8fd2bf900 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Wed, 6 Dec 2023 00:07:05 +0000
Subject: [PATCH 01/16] add sliding window to FA

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 tests/pytorch/test_fused_attn.py          |  62 +++++++++
 transformer_engine/pytorch/attention.py   | 148 ++++++++++++++++++----
 transformer_engine/pytorch/transformer.py |  35 ++++-
 3 files changed, 220 insertions(+), 25 deletions(-)

diff --git a/tests/pytorch/test_fused_attn.py b/tests/pytorch/test_fused_attn.py
index 085c655bdb..bb4942e5f0 100644
--- a/tests/pytorch/test_fused_attn.py
+++ b/tests/pytorch/test_fused_attn.py
@@ -262,6 +262,68 @@ def get_dummy_cuda_rng_tracker():
 
     return op, inp.grad
 
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes_lean)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("ckpt_attn", [False])#True, False])
+@pytest.mark.parametrize("bias_type", ["no_bias"])#, "post_scale_bias"])
+def test_dpa_sliding_window(dtype, bs, model, ckpt_attn, bias_type):
+    """Test DotProductAttention module with sliding window"""
+
+    # Get configs
+    config = model_configs[model]
+    tols = dict(atol=5e-3, rtol=5e-3)
+    if dtype == torch.bfloat16:
+        tols = dict(atol=2.5e-2, rtol=2.5e-2)
+
+    # Skip if only unfused backend is supported
+    fused_attn_supported = _is_fused_attention_supported(
+        config,
+        dtype,
+        bias_type=bias_type,
+    )
+    flash_attn_supported = _is_flash_attention_supported(bias_type=bias_type)
+    if not (fused_attn_supported or flash_attn_supported):
+        pytest.skip(
+            "Neither FusedAttention nor FlashAttention support this model config"
+        )
+
+    # UnfusedDotProductAttention backend
+    unfused_attn_fwd, unfused_attn_bwd = _run_dot_product_attention(
+        dtype,
+        bs,
+        config,
+        "UnfusedDotProductAttention",
+        ckpt_attn,
+        bias_type,
+    )
+
+    # FusedAttention backend
+    if fused_attn_supported:
+        fused_attn_fwd, fused_attn_bwd = _run_dot_product_attention(
+            dtype,
+            bs,
+            config,
+            "FusedAttention",
+            ckpt_attn,
+            bias_type,
+        )
+        torch.testing.assert_close(fused_attn_fwd, unfused_attn_fwd, **tols)
+        torch.testing.assert_close(fused_attn_bwd, unfused_attn_bwd, **tols)
+
+    # FlashAttention backend
+    if flash_attn_supported:
+        flash_attn_fwd, flash_attn_bwd = _run_dot_product_attention(
+            dtype,
+            bs,
+            config,
+            "FlashAttention",
+            ckpt_attn,
+            bias_type,
+        )
+        torch.testing.assert_close(flash_attn_fwd, unfused_attn_fwd, **tols)
+        torch.testing.assert_close(flash_attn_bwd, unfused_attn_bwd, **tols)
+
 qkv_layouts = [
     'sb3hd', 'sbh3d', 'sbhd_sb2hd', 'sbhd_sbh2d', 'sbhd_sbhd_sbhd',
     'bs3hd', 'bsh3d', 'bshd_bs2hd', 'bshd_bsh2d', 'bshd_bshd_bshd',
diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index 52b17098e2..1ec2d5cf46 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -57,6 +57,7 @@
 _flash_attn_version_required = packaging.version.Version("1.0.6")
 _flash_attn_2_available = _flash_attn_version >= packaging.version.Version("2")
 _flash_attn_2_1_plus = _flash_attn_version >= packaging.version.Version("2.1")
+_flash_attn_2_3_plus = _flash_attn_version >= packaging.version.Version("2.3")
 
 if _flash_attn_2_available:
     from flash_attn.flash_attn_interface import flash_attn_varlen_func as flash_attn_forward_func # pylint: disable=no-name-in-module
@@ -1248,12 +1249,26 @@ def forward(
         cu_seqlens_q: Optional[torch.Tensor] = None,
         cu_seqlens_kv: Optional[torch.Tensor] = None,
         attn_mask_type: str = "causal",
+        window_size: Optional[Tuple[int, int]] = None,
         cp_group: Optional[dist_group_type] = None,
         cp_global_ranks: List[int] = None,
         cp_stream: torch.cuda.Stream = None,
+        window_size: Optional[Tuple[int, int]] = (-1, -1),
     ) -> torch.Tensor:
         """flash-attn fprop"""
 
+        if "causal" in attn_mask_type:
+            if window_size is None:
+                window_size = (-1, 0)
+            else:
+                assert (
+                    window_size[1] == 0
+                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
+        else:
+            if window_size is None:
+                window_size = (-1, -1)
+        print("SWA: ",window_size)
+
         assert (
             query_layer.dtype in [torch.float16, torch.bfloat16]
             and key_layer.dtype in [torch.float16, torch.bfloat16]
@@ -1348,6 +1363,9 @@ def forward(
             max_seqlen_kv = seqlens_kv.max().item()
 
         if context_parallel:
+            assert (
+                window_size == (-1, -1) or window_size == (-1, 0)
+                ), "Sliding window attention is not supported with context parallelism."
             with self.attention_dropout_ctx():
                 output = flash_attn_forward_func_with_cp(
                     query_layer, key_layer, value_layer,
@@ -1368,6 +1386,7 @@ def forward(
                     cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv,
                     self.attention_dropout if self.training else 0.0,
                     softmax_scale=1.0/self.norm_factor, causal=attn_mask_type=="causal",
+                    window_size=window_size,
                     **fa_optional_forward_kwargs
                 )
 
@@ -1771,6 +1790,12 @@ class DotProductAttention(torch.nn.Module):
                    mask can also be applied in conjunction with "`padding`" mask by passing
                    in multiple mask type as a comma separated string, for example,
                    `attn_mask_type="causal,padding"`.
+    window_size: Optional[Tuple[int, int]], default = `None`
+                sliding window size for local attention, where query at position i attends to keys
+                in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q
+                + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding
+                window and causal mask specifically. Similar to :attr:`attn_mask_type`, it can
+                be overridden by :attr:`window_size` in `forward` as well.
     attention_type: str, default = `self`
                    type of attention, either "`self`" and "`cross`".
     layer_number: int, default = `None`
@@ -1820,6 +1845,7 @@ def __init__(
         attention_dropout: float = 0.0,
         qkv_format: str = "sbhd",
         attn_mask_type: str = "causal",
+        window_size: Optional[Tuple[int, int]] = None,
         sequence_parallel: bool = False,
         tp_size: int = 1,
         get_rng_state_tracker: Optional[Callable] = None,
@@ -1834,6 +1860,17 @@ def __init__(
 
         self.qkv_format = qkv_format
         self.attn_mask_type = attn_mask_type
+        self.window_size = window_size
+        if "causal" in attn_mask_type:
+            if window_size is None:
+                self.window_size = (-1, 0)
+            else:
+                assert (
+                    window_size[1] == 0
+                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
+        else:
+            if window_size is None:
+                self.window_size = (-1, -1)
         self.tp_size = tp_size if tp_group is None else get_distributed_world_size(tp_group)
         self.tp_group = tp_group
         self.get_rng_state_tracker = get_rng_state_tracker
@@ -1961,6 +1998,7 @@ def forward(
         cu_seqlens_q: Optional[torch.Tensor] = None,
         cu_seqlens_kv: Optional[torch.Tensor] = None,
         attn_mask_type: Optional[str] = None,
+        window_size: Optional[Tuple[int, int]] = None,
         checkpoint_core_attention: bool = False,
         core_attention_bias_type: str = "no_bias",
         core_attention_bias: Optional[torch.Tensor] = None,
@@ -2029,6 +2067,8 @@ def forward(
                    with shape [batch_size + 1] and dtype torch.int32.
         attn_mask_type: {'causal', 'padding', 'no_mask', 'arbitrary'}, default = `None`
                        type of attention mask passed into softmax operation.
+        window_size: Optional[Tuple[int, int]], default = `None`
+                    sliding window size for local attention.
         checkpoint_core_attention : bool, default = `False`
                                    If true, forward activations for attention are recomputed
                                    during the backward pass in order to save memory that would
@@ -2049,8 +2089,17 @@ def forward(
         assert (key_layer.shape == value_layer.shape
             ), "Keys and values must have the same shape!"
 
+        if attn_mask_type is not None and "causal" in attn_mask_type:
+            if window_size is None:
+                window_size = (-1, 0)
+            else:
+                assert (
+                    window_size[1] == 0
+                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
         if attn_mask_type is None:
             attn_mask_type = self.attn_mask_type
+        if window_size is None:
+            window_size = self.window_size
         if qkv_format is None:
             qkv_format = self.qkv_format
         attn_mask_type, causal_mask = _unpack_attn_mask_type(attn_mask_type)
@@ -2103,6 +2152,7 @@ def forward(
         # is: FlashAttention > FusedAttention (cuDNN) > UnfusedDotProductAttention.
         use_flash_attention = self.use_flash_attention
         use_fused_attention = self.use_fused_attention
+        use_unfused_attention = True 
 
         # The following section filters out some backends based on
         # certain asserts before executing the forward pass.
@@ -2132,9 +2182,11 @@ def forward(
                     and self.device_compute_capability not in ((8, 0), (9, 0)))):
                 use_flash_attention = False
 
+        # Filter: MQA/GQA.
         if not _flash_attn_2_available and self.num_gqa_groups != self.num_attention_heads:
             use_flash_attention = False
 
+        # Filter: cross attention + causal mask.
         if (_flash_attn_2_1_plus
             and causal_mask
             and max_seqlen_q != max_seqlen_kv):
@@ -2145,9 +2197,18 @@ def forward(
             )
             use_flash_attention = False
 
+        # Filter: bias.
         if core_attention_bias_type != "no_bias" or core_attention_bias is not None:
             use_flash_attention = False
 
+        # Filter: sliding window attention.
+        if window_size != (-1, -1) and window_size != (-1, 0):
+            use_fused_attention = False
+            use_unfused_attention = False
+            context_parallel = (self.cp_group is not None) and (get_distributed_world_size(self.cp_group) != 1)
+            if (not _flash_attn_2_3_plus) or context_parallel: 
+                use_flash_attention = False
+
         # Filter: ONNX export.
         if is_in_onnx_export_mode():
             use_flash_attention = False
@@ -2157,7 +2218,7 @@ def forward(
         #    attn_mask_type(s)   |     supported backends
         # ------------------------------------------------
         #   causal               |     All
-        #   padding              |     UnfusedDotProductAttention, FlashAttention
+        #   padding              |     FlashAttention
         #   arbitrary            |     UnfusedDotProductAttention
         #   no_mask              |     All
         #   causal + padding     |     FlashAttention
@@ -2169,6 +2230,10 @@ def forward(
             assert use_flash_attention, "No attention backend available for causal + padding masks."
         elif attn_mask_type == "padding":
             use_fused_attention = False
+        if "padding" in attn_mask_type:
+            use_unfused_attention = False
+        if "causal" in attn_mask_type and max_seqlen_q != max_seqlen_kv:
+            use_unfused_attention = False
 
         if use_fused_attention:
             fused_attention_backend = tex.get_fused_attn_backend(
@@ -2203,6 +2268,7 @@ def forward(
                                         cu_seqlens_q=cu_seqlens_q,
                                         cu_seqlens_kv=cu_seqlens_kv,
                                         attn_mask_type=attn_mask_type,
+                                        window_size=window_size,
                                         cp_group=self.cp_group,
                                         cp_global_ranks=self.cp_global_ranks,
                                         cp_stream=self.cp_stream)
@@ -2235,29 +2301,32 @@ def forward(
                               core_attention_bias = core_attention_bias,
                               fast_zero_fill = fast_zero_fill)
 
-        if checkpoint_core_attention:
-            return self._checkpointed_attention_forward(
-                self.unfused_attention,
-                query_layer,
-                key_layer,
-                value_layer,
-                qkv_layout = qkv_layout,
-                cu_seqlens_q = cu_seqlens_q,
-                cu_seqlens_kv = cu_seqlens_kv,
-                attn_mask_type = attn_mask_type,
-                attention_mask = attention_mask,
-                core_attention_bias_type = core_attention_bias_type,
-                core_attention_bias = core_attention_bias)
-        return self.unfused_attention(query_layer,
-                key_layer,
-                value_layer,
-                qkv_layout = qkv_layout,
-                cu_seqlens_q = cu_seqlens_q,
-                cu_seqlens_kv = cu_seqlens_kv,
-                attn_mask_type = attn_mask_type,
-                attention_mask = attention_mask,
-                core_attention_bias_type = core_attention_bias_type,
-                core_attention_bias = core_attention_bias)
+        if use_unfused_attention:
+            if checkpoint_core_attention:
+                return self._checkpointed_attention_forward(
+                    self.unfused_attention,
+                    query_layer,
+                    key_layer,
+                    value_layer,
+                    qkv_layout = qkv_layout,
+                    cu_seqlens_q = cu_seqlens_q,
+                    cu_seqlens_kv = cu_seqlens_kv,
+                    attn_mask_type = attn_mask_type,
+                    attention_mask = attention_mask,
+                    core_attention_bias_type = core_attention_bias_type,
+                    core_attention_bias = core_attention_bias)
+            return self.unfused_attention(query_layer,
+                    key_layer,
+                    value_layer,
+                    qkv_layout = qkv_layout,
+                    cu_seqlens_q = cu_seqlens_q,
+                    cu_seqlens_kv = cu_seqlens_kv,
+                    attn_mask_type = attn_mask_type,
+                    attention_mask = attention_mask,
+                    core_attention_bias_type = core_attention_bias_type,
+                    core_attention_bias = core_attention_bias)
+
+        raise Exception("No dot product attention support for the provided inputs!")
 
 
 class MultiheadAttention(torch.nn.Module):
@@ -2301,6 +2370,12 @@ class MultiheadAttention(torch.nn.Module):
                    arg is useful for dynamically changing mask types, e.g. a different
                    mask for training and inference. The init arg is useful for cases
                    involving compilation/tracing, e.g. ONNX export.
+    window_size: Optional[Tuple[int, int]], default = `None`
+                sliding window size for local attention, where query at position i attends to keys
+                in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q
+                + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding
+                window and causal mask specifically. Similar to :attr:`attn_mask_type`, it can
+                be overridden by :attr:`window_size` in `forward` as well.
     num_gqa_groups : int, default = `None`
                          number of GQA groups in the transformer layer.
                          Grouped Query Attention is described in
@@ -2392,6 +2467,7 @@ def __init__(
         output_layer_init_method: Optional[Callable] = None,
         layer_number: Optional[int] = None,
         attn_mask_type: str = "causal",
+        window_size: Optional[Tuple[int, int]] = None,
         tp_group: Optional[dist_group_type] = None,
         tp_size: int = 1,
         num_gqa_groups: Optional[int] = None,
@@ -2420,6 +2496,17 @@ def __init__(
         super().__init__()
 
         self.attn_mask_type = attn_mask_type
+        self.window_size = window_size
+        if "causal" in attn_mask_type:
+            if window_size is None:
+                self.window_size = (-1, 0)
+            else:
+                assert (
+                    window_size[1] == 0
+                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
+        else:
+            if window_size is None:
+                self.window_size = (-1, -1)
         self.layer_number = layer_number
         self.input_layernorm = input_layernorm
         self.attention_type = attention_type
@@ -2633,6 +2720,7 @@ def forward(
         attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None,
         encoder_output: Optional[torch.Tensor] = None,
         attn_mask_type: Optional[str] = None,
+        window_size: Optional[Tuple[int, int]] = None,
         is_first_microbatch: Optional[bool] = None,
         checkpoint_core_attention: bool = False,
         inference_params: Optional[InferenceParams] = None,
@@ -2657,6 +2745,8 @@ def forward(
              Boolean tensor used to mask out self-attention softmax input.
         attn_mask_type: {'causal', 'padding', 'no_mask', arbitrary}, default = `None`
                        type of attention mask passed into softmax operation.
+        window_size: Optional[Tuple[int, int]], default = `None`
+                    sliding window size for local attention.
         encoder_output : Optional[torch.Tensor], default = `None`
              Output of the encoder block to be fed into the decoder block if using
              `layer_type="decoder"`.
@@ -2690,8 +2780,17 @@ def forward(
         """
         # hidden_states: [sq, b, h]
 
+        if attn_mask_type is not None and "causal" in attn_mask_type:
+            if window_size is None:
+                window_size = (-1, 0)
+            else:
+                assert (
+                    window_size[1] == 0
+                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
         if attn_mask_type is None:
             attn_mask_type = self.attn_mask_type
+        if window_size is None:
+            window_size = self.window_size
 
         if attn_mask_type == "padding" and attention_mask is not None:
             assert (
@@ -2918,6 +3017,7 @@ def forward(
             cu_seqlens_kv=None,
             attention_mask=attention_mask,
             attn_mask_type=attn_mask_type,
+            window_size=window_size,
             checkpoint_core_attention=checkpoint_core_attention,
             core_attention_bias_type=core_attention_bias_type,
             core_attention_bias=core_attention_bias,
diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py
index dd86260f9f..d8ca57355e 100644
--- a/transformer_engine/pytorch/transformer.py
+++ b/transformer_engine/pytorch/transformer.py
@@ -133,6 +133,12 @@ class TransformerLayer(torch.nn.Module):
                         arg is useful for dynamically changing mask types, e.g. a different
                         mask for training and inference. The init arg is useful for cases
                         involving compilation/tracing, e.g. ONNX export.
+    window_size: Optional[Tuple[int, int]], default = `None`
+                sliding window size for local attention, where query at position i attends to keys
+                in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q
+                + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding
+                window and causal mask specifically. Similar to :attr:`self_attn_mask_type`, it can
+                be overridden by :attr:`window_size` in `forward` as well.
     zero_centered_gamma : bool, default = 'False'
                          if set to 'True', gamma parameter in LayerNorm is initialized to 0 and
                          the LayerNorm formula changes to
@@ -219,6 +225,7 @@ def __init__(
         layer_number: Optional[int] = None,
         kv_channels: Optional[int] = None,
         self_attn_mask_type: str = "causal",
+        window_size: Optional[Tuple[int, int]] = None,
         tp_group: Optional[dist_group_type] = None,
         tp_size: int = 1,
         params_dtype: Optional[torch.dtype] = None,
@@ -250,6 +257,17 @@ def __init__(
             ), "Userbuffer communication backend not available."
 
         self.self_attn_mask_type = self_attn_mask_type
+        self.window_size = window_size 
+        if "causal" in self_attn_mask_type:
+            if window_size is None:
+                self.window_size = (-1, 0)
+            else:
+                assert (
+                    window_size[1] == 0
+                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
+        else:
+            if window_size is None:
+                self.window_size = (-1, -1)
         params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype
         ub_tp_comm_overlap = ub_tp_comm_overlap and bool(int(os.getenv("NVTE_UB_OVERLAP", "1")))
         ub_bulk_wgrad = ub_tp_comm_overlap and bool(int(os.getenv("NVTE_UB_BULK_WGRAD", "1")))
@@ -490,6 +508,7 @@ def forward(
         hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
         self_attn_mask_type: Optional[str] = None,
+        window_size: Optional[Tuple[int, int]] = None,
         encoder_output: Optional[torch.Tensor] = None,
         enc_dec_attn_mask: Optional[torch.Tensor] = None,
         is_first_microbatch: Optional[bool] = None,
@@ -515,8 +534,10 @@ def forward(
         attention_mask : Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], default = `None`
                         Boolean tensor used to mask out self-attention softmax input.
                         Can be a tuple of 2 masks for cross attention with padding masks.
-        self_attn_mask_type: {'causal', 'padding', 'no_mask', 'arbitrary'}, default = `causal`
+        self_attn_mask_type: {'causal', 'padding', 'no_mask', 'arbitrary'}, default = `None`
                             type of attention mask passed into softmax operation.
+        window_size: Optional[Tuple[int, int]], default = `None`
+                    sliding window size for local attention.
         encoder_output : Optional[torch.Tensor], default = `None`
              Output of the encoder block to be fed into the decoder block if using
              `layer_type="decoder"`.
@@ -555,8 +576,18 @@ def forward(
                          to efficienly calculate and store the context during inference.
         """
 
+        if self_attn_mask_type is not None and "causal" in self_attn_mask_type:
+            if window_size is None:
+                window_size = (-1, 0)
+            else:
+                assert (
+                    window_size[1] == 0
+                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
+
         if self_attn_mask_type is None:
             self_attn_mask_type = self.self_attn_mask_type
+        if window_size is None:
+            window_size = self.window_size
 
         assert (
             self_attn_mask_type in AttnMaskTypes
@@ -585,6 +616,7 @@ def forward(
             hidden_states,
             attention_mask=attention_mask,
             attn_mask_type=self_attn_mask_type,
+            window_size=window_size,
             inference_params=inference_params,
             is_first_microbatch=is_first_microbatch,
             checkpoint_core_attention=checkpoint_core_attention,
@@ -611,6 +643,7 @@ def forward(
                 hidden_states,
                 attention_mask=enc_dec_attn_mask,
                 attn_mask_type=self_attn_mask_type,
+                window_size=window_size,
                 encoder_output=encoder_output,
                 is_first_microbatch=is_first_microbatch,
                 checkpoint_core_attention=checkpoint_core_attention,

From 8f7121ff3c6a1c7692af84a6a1cc6b07ef945658 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Wed, 6 Dec 2023 00:22:51 +0000
Subject: [PATCH 02/16] fix forward logic

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 transformer_engine/pytorch/attention.py   | 34 +++++++++++++----------
 transformer_engine/pytorch/transformer.py | 16 +++++++----
 2 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index 1ec2d5cf46..c95df0a4a7 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -1253,7 +1253,6 @@ def forward(
         cp_group: Optional[dist_group_type] = None,
         cp_global_ranks: List[int] = None,
         cp_stream: torch.cuda.Stream = None,
-        window_size: Optional[Tuple[int, int]] = (-1, -1),
     ) -> torch.Tensor:
         """flash-attn fprop"""
 
@@ -1267,7 +1266,6 @@ def forward(
         else:
             if window_size is None:
                 window_size = (-1, -1)
-        print("SWA: ",window_size)
 
         assert (
             query_layer.dtype in [torch.float16, torch.bfloat16]
@@ -2089,13 +2087,17 @@ def forward(
         assert (key_layer.shape == value_layer.shape
             ), "Keys and values must have the same shape!"
 
-        if attn_mask_type is not None and "causal" in attn_mask_type:
-            if window_size is None:
-                window_size = (-1, 0)
+        if attn_mask_type is not None:
+            if "causal" in attn_mask_type:
+                if window_size is None:
+                    window_size = (-1, 0)
+                else:
+                    assert (
+                        window_size[1] == 0
+                    ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
             else:
-                assert (
-                    window_size[1] == 0
-                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
+                if window_size is None:
+                    window_size = (-1, -1)
         if attn_mask_type is None:
             attn_mask_type = self.attn_mask_type
         if window_size is None:
@@ -2780,13 +2782,17 @@ def forward(
         """
         # hidden_states: [sq, b, h]
 
-        if attn_mask_type is not None and "causal" in attn_mask_type:
-            if window_size is None:
-                window_size = (-1, 0)
+        if attn_mask_type is not None:
+            if "causal" in attn_mask_type:
+                if window_size is None:
+                    window_size = (-1, 0)
+                else:
+                    assert (
+                        window_size[1] == 0
+                    ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
             else:
-                assert (
-                    window_size[1] == 0
-                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
+                if window_size is None:
+                    window_size = (-1, -1)
         if attn_mask_type is None:
             attn_mask_type = self.attn_mask_type
         if window_size is None:
diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py
index d8ca57355e..ad6d816350 100644
--- a/transformer_engine/pytorch/transformer.py
+++ b/transformer_engine/pytorch/transformer.py
@@ -576,13 +576,17 @@ def forward(
                          to efficienly calculate and store the context during inference.
         """
 
-        if self_attn_mask_type is not None and "causal" in self_attn_mask_type:
-            if window_size is None:
-                window_size = (-1, 0)
+        if self_attn_mask_type is not None:
+            if "causal" in self_attn_mask_type:
+                if window_size is None:
+                    window_size = (-1, 0)
+                else:
+                    assert (
+                        window_size[1] == 0
+                    ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
             else:
-                assert (
-                    window_size[1] == 0
-                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
+                if window_size is None:
+                    window_size = (-1, -1)
 
         if self_attn_mask_type is None:
             self_attn_mask_type = self.self_attn_mask_type

From f37e87209f323dad644f982be5e21b5ef4724126 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Wed, 6 Dec 2023 00:33:59 +0000
Subject: [PATCH 03/16] fix lint

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 transformer_engine/pytorch/attention.py   | 11 ++++++-----
 transformer_engine/pytorch/transformer.py |  2 +-
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index c95df0a4a7..815a043abd 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -1362,7 +1362,7 @@ def forward(
 
         if context_parallel:
             assert (
-                window_size == (-1, -1) or window_size == (-1, 0)
+                window_size in ((-1, -1), (-1, 0))
                 ), "Sliding window attention is not supported with context parallelism."
             with self.attention_dropout_ctx():
                 output = flash_attn_forward_func_with_cp(
@@ -2154,7 +2154,7 @@ def forward(
         # is: FlashAttention > FusedAttention (cuDNN) > UnfusedDotProductAttention.
         use_flash_attention = self.use_flash_attention
         use_fused_attention = self.use_fused_attention
-        use_unfused_attention = True 
+        use_unfused_attention = True
 
         # The following section filters out some backends based on
         # certain asserts before executing the forward pass.
@@ -2204,11 +2204,12 @@ def forward(
             use_flash_attention = False
 
         # Filter: sliding window attention.
-        if window_size != (-1, -1) and window_size != (-1, 0):
+        if window_size not in ((-1, -1), (-1, 0)):
             use_fused_attention = False
             use_unfused_attention = False
-            context_parallel = (self.cp_group is not None) and (get_distributed_world_size(self.cp_group) != 1)
-            if (not _flash_attn_2_3_plus) or context_parallel: 
+            context_parallel = (self.cp_group is not None
+                and get_distributed_world_size(self.cp_group) != 1)
+            if (not _flash_attn_2_3_plus) or context_parallel:
                 use_flash_attention = False
 
         # Filter: ONNX export.
diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py
index ad6d816350..1ab7edfb92 100644
--- a/transformer_engine/pytorch/transformer.py
+++ b/transformer_engine/pytorch/transformer.py
@@ -257,7 +257,7 @@ def __init__(
             ), "Userbuffer communication backend not available."
 
         self.self_attn_mask_type = self_attn_mask_type
-        self.window_size = window_size 
+        self.window_size = window_size
         if "causal" in self_attn_mask_type:
             if window_size is None:
                 self.window_size = (-1, 0)

From 37f14698f8d2a9caa0a93c0728f98cd64838f998 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Wed, 6 Dec 2023 00:59:08 +0000
Subject: [PATCH 04/16] change bert test to causal as unfused does not support
 padding

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 tests/pytorch/test_sanity.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py
index da4714c7ea..7cfc981f17 100644
--- a/tests/pytorch/test_sanity.py
+++ b/tests/pytorch/test_sanity.py
@@ -518,7 +518,7 @@ def test_sanity_bert(dtype, bs, fp8_recipe, model, skip_wgrad, zero_centered_gam
             apply_residual_connection_post_layernorm=True,
             output_layernorm=True,
             zero_centered_gamma=zero_centered_gamma,
-            self_attn_mask_type="padding",
+            self_attn_mask_type="causal",
             normalization=normalization,
         )
         .to(dtype=dtype)

From e278579f427c2bfca76ee5c3e81635bad8a759e3 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Wed, 6 Dec 2023 20:25:27 +0000
Subject: [PATCH 05/16] fix FlashAttention for v2-2.3 versions

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 transformer_engine/pytorch/attention.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index 815a043abd..654587de13 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -1379,12 +1379,13 @@ def forward(
                 fa_optional_forward_kwargs = {}
                 if not _flash_attn_2_available:
                     fa_optional_forward_kwargs["deterministic"] = self.deterministic
+                if _flash_attn_2_3_plus:
+                    fa_optional_forward_kwargs["window_size"] = window_size
                 output = flash_attn_forward_func(
                     query_layer, key_layer, value_layer,
                     cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv,
                     self.attention_dropout if self.training else 0.0,
                     softmax_scale=1.0/self.norm_factor, causal=attn_mask_type=="causal",
-                    window_size=window_size,
                     **fa_optional_forward_kwargs
                 )
 

From 49e147e322ee73a64cb9a338e1d19984f918a87e Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Thu, 7 Dec 2023 23:39:11 +0000
Subject: [PATCH 06/16] verify FA swa works

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 tests/pytorch/test_fused_attn.py        | 68 ++++++++++++-------------
 transformer_engine/pytorch/attention.py |  2 +-
 2 files changed, 33 insertions(+), 37 deletions(-)

diff --git a/tests/pytorch/test_fused_attn.py b/tests/pytorch/test_fused_attn.py
index bb4942e5f0..d4b2c99f32 100644
--- a/tests/pytorch/test_fused_attn.py
+++ b/tests/pytorch/test_fused_attn.py
@@ -226,6 +226,20 @@ def _run_dot_product_attention(dtype, bs, config, backend, ckpt_attn, bias_type)
     else:
         bias = None
 
+    def get_mask(seq_q, seq_kv):
+        w = torch.randint(0, seq_kv, [2], dtype=torch.int32, device="cuda")
+        #w = torch.Tensor([seq_kv,0]).to(dtype=torch.int32)
+        print('w',w)
+        m = torch.ones(seq_q, seq_kv, dtype=torch.bool, device="cuda")
+        mu = torch.triu(m, diagonal=seq_kv-seq_q-w[0])
+        ml = torch.tril(mu, diagonal=seq_kv-seq_q+w[1])
+        #print(ml.to(dtype=torch.int))
+        ml = ~ ml
+        #print(ml.to(dtype=torch.int))
+        return w, ml
+
+    window_size, attention_mask = get_mask(config.seq_len, config.seq_len)
+
     _DUMMY_CUDA_RNG_STATE_TRACKER = CudaRNGStatesTracker()
     _DUMMY_CUDA_RNG_STATE_TRACKER.add("model-parallel-rng", seed)
 
@@ -255,6 +269,8 @@ def get_dummy_cuda_rng_tracker():
         cu_seqlens_q = cu_seqlens,
         cu_seqlens_kv = cu_seqlens,
         attn_mask_type=config.attn_mask_type,
+        window_size=window_size,
+        attention_mask=attention_mask,
         checkpoint_core_attention=ckpt_attn,
         core_attention_bias_type=bias_type,
         core_attention_bias=bias)
@@ -262,32 +278,34 @@ def get_dummy_cuda_rng_tracker():
 
     return op, inp.grad
 
+model_configs_swa = {
+    "test5": ModelConfig(1, 1024, 16, 64, 128, 0.0, "no_mask"),
+}
 @pytest.mark.parametrize("dtype", param_types)
 @pytest.mark.parametrize("bs", batch_sizes_lean)
-@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("model", model_configs_swa.keys())
 @pytest.mark.parametrize("ckpt_attn", [False])#True, False])
 @pytest.mark.parametrize("bias_type", ["no_bias"])#, "post_scale_bias"])
 def test_dpa_sliding_window(dtype, bs, model, ckpt_attn, bias_type):
     """Test DotProductAttention module with sliding window"""
 
     # Get configs
-    config = model_configs[model]
+    config = model_configs_swa[model]
     tols = dict(atol=5e-3, rtol=5e-3)
     if dtype == torch.bfloat16:
         tols = dict(atol=2.5e-2, rtol=2.5e-2)
 
-    # Skip if only unfused backend is supported
-    fused_attn_supported = _is_fused_attention_supported(
-        config,
+    # FlashAttention backend
+    flash_attn_fwd, flash_attn_bwd = _run_dot_product_attention(
         dtype,
-        bias_type=bias_type,
+        bs,
+        config,
+        "FlashAttention",
+        ckpt_attn,
+        bias_type,
     )
-    flash_attn_supported = _is_flash_attention_supported(bias_type=bias_type)
-    if not (fused_attn_supported or flash_attn_supported):
-        pytest.skip(
-            "Neither FusedAttention nor FlashAttention support this model config"
-        )
 
+    config.attn_mask_type = "arbitrary"
     # UnfusedDotProductAttention backend
     unfused_attn_fwd, unfused_attn_bwd = _run_dot_product_attention(
         dtype,
@@ -298,31 +316,9 @@ def test_dpa_sliding_window(dtype, bs, model, ckpt_attn, bias_type):
         bias_type,
     )
 
-    # FusedAttention backend
-    if fused_attn_supported:
-        fused_attn_fwd, fused_attn_bwd = _run_dot_product_attention(
-            dtype,
-            bs,
-            config,
-            "FusedAttention",
-            ckpt_attn,
-            bias_type,
-        )
-        torch.testing.assert_close(fused_attn_fwd, unfused_attn_fwd, **tols)
-        torch.testing.assert_close(fused_attn_bwd, unfused_attn_bwd, **tols)
-
-    # FlashAttention backend
-    if flash_attn_supported:
-        flash_attn_fwd, flash_attn_bwd = _run_dot_product_attention(
-            dtype,
-            bs,
-            config,
-            "FlashAttention",
-            ckpt_attn,
-            bias_type,
-        )
-        torch.testing.assert_close(flash_attn_fwd, unfused_attn_fwd, **tols)
-        torch.testing.assert_close(flash_attn_bwd, unfused_attn_bwd, **tols)
+    torch.testing.assert_close(flash_attn_fwd, unfused_attn_fwd, **tols)
+    torch.testing.assert_close(flash_attn_bwd, unfused_attn_bwd, **tols)
+    print('flash vs unfused')
 
 qkv_layouts = [
     'sb3hd', 'sbh3d', 'sbhd_sb2hd', 'sbhd_sbh2d', 'sbhd_sbhd_sbhd',
diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index 654587de13..878ea87045 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -2207,7 +2207,7 @@ def forward(
         # Filter: sliding window attention.
         if window_size not in ((-1, -1), (-1, 0)):
             use_fused_attention = False
-            use_unfused_attention = False
+            #use_unfused_attention = False
             context_parallel = (self.cp_group is not None
                 and get_distributed_world_size(self.cp_group) != 1)
             if (not _flash_attn_2_3_plus) or context_parallel:

From 97f807fcd6c663057a89b90c09379d6f31b36ebf Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Fri, 8 Dec 2023 00:54:43 +0000
Subject: [PATCH 07/16] fix mask related restrictions and duplicate code after
 merge

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 tests/pytorch/test_fused_attn.py        | 12 ------------
 transformer_engine/pytorch/attention.py | 10 ++--------
 2 files changed, 2 insertions(+), 20 deletions(-)

diff --git a/tests/pytorch/test_fused_attn.py b/tests/pytorch/test_fused_attn.py
index 4c2fc5fd56..e29437fd69 100644
--- a/tests/pytorch/test_fused_attn.py
+++ b/tests/pytorch/test_fused_attn.py
@@ -281,22 +281,18 @@ def test_dot_product_attention(dtype, model_configs, model, ckpt_attn, workspace
         torch.testing.assert_close(fused_attn_fwd, unfused_attn_fwd, **tols)
         for i,_ in enumerate(unfused_attn_bwd):
             torch.testing.assert_close(fused_attn_bwd[i], unfused_attn_bwd[i], **tols)
-        print('unfused vs fused')
     if unfused_attn_supported and flash_attn_supported:
         torch.testing.assert_close(flash_attn_fwd, unfused_attn_fwd, **tols)
         for i,_ in enumerate(flash_attn_bwd):
             torch.testing.assert_close(unfused_attn_bwd[i], flash_attn_bwd[i], **tols)
-        print('unfused vs flash')
     if fused_attn_supported and flash_attn_supported:
         torch.testing.assert_close(fused_attn_fwd, flash_attn_fwd, **tols)
         for i,_ in enumerate(flash_attn_bwd):
             torch.testing.assert_close(fused_attn_bwd[i], flash_attn_bwd[i], **tols)
-        print('fused vs flash')
     if fused_attn_supported and len(fused_attn_backend) == 2:
         torch.testing.assert_close(fused_attn_fwd, fused_attn_fwd_1, **tols)
         for i,_ in enumerate(fused_attn_bwd):
             torch.testing.assert_close(fused_attn_bwd[i], fused_attn_bwd_1[i], **tols)
-        print('fused 0 vs fused 1')
 
 @pytest.mark.skipif(_cudnn_version() < (8,9,1), reason="cuDNN 8.9.1+ is required.")
 @pytest.mark.parametrize("dtype", param_types)
@@ -322,14 +318,6 @@ def test_dpa_checkpoint(dtype, model_configs, model):
     "mask_6_1": ModelConfig(1, 24, 24, 128, 2048, 4096, 0.0, "padding_causal", "no_bias"),
 }
 
-@pytest.mark.skipif(_cudnn_version() < (8,9,1), reason="cuDNN 8.9.1+ is required.")
-@pytest.mark.parametrize("dtype", param_types_lean)
-@pytest.mark.parametrize("model_configs", [model_configs_mask])
-@pytest.mark.parametrize("model", model_configs_mask.keys())
-def test_dpa_mask(dtype, model_configs, model):
-    """Test DotProductAttention module with different mask types"""
-    test_dot_product_attention(dtype, model_configs, model, False, True, None, False)
-
 @pytest.mark.skipif(_cudnn_version() < (8,9,1), reason="cuDNN 8.9.1+ is required.")
 @pytest.mark.parametrize("dtype", param_types_lean)
 @pytest.mark.parametrize("model_configs", [model_configs_mask])
diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index b44bc4198e..913341e2f2 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -2321,9 +2321,9 @@ def forward(
             use_flash_attention = False
 
         # Filter: sliding window attention.
+        # UnfusedDotProductAttention can support SWA via arbitrary attention mask.
         if window_size not in ((-1, -1), (-1, 0)):
             use_fused_attention = False
-            #use_unfused_attention = False
             context_parallel = (self.cp_group is not None
                 and get_distributed_world_size(self.cp_group) != 1)
             if (not _flash_attn_2_3_plus) or context_parallel:
@@ -2338,7 +2338,7 @@ def forward(
         #    attn_mask_type(s)   |     supported backends
         # ------------------------------------------------
         #   no_mask              |     All
-        #   padding              |     FlashAttention, FusedAttention
+        #   padding              |     UnfusedDotProductAttention, FlashAttention, FusedAttention
         #   causal               |     All
         #   padding + causal     |     FlashAttention, FusedAttention
         #   arbitrary            |     UnfusedDotProductAttention
@@ -2346,12 +2346,6 @@ def forward(
         if attn_mask_type == "arbitrary":
             use_flash_attention = False
             use_fused_attention = False
-        elif attn_mask_type == "padding" and causal_mask:
-            assert use_flash_attention, "No attention backend available for causal + padding masks."
-        elif attn_mask_type == "padding":
-            use_fused_attention = False
-        if "padding" in attn_mask_type:
-            use_unfused_attention = False
         if "causal" in attn_mask_type and max_seqlen_q != max_seqlen_kv:
             use_unfused_attention = False
 

From 603cb522c821bbe8ab33454937eaeb9bfc4ead3d Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Fri, 8 Dec 2023 00:58:29 +0000
Subject: [PATCH 08/16] fix swa test

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 tests/pytorch/test_fused_attn.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tests/pytorch/test_fused_attn.py b/tests/pytorch/test_fused_attn.py
index e29437fd69..1652d00915 100644
--- a/tests/pytorch/test_fused_attn.py
+++ b/tests/pytorch/test_fused_attn.py
@@ -248,12 +248,14 @@ def test_dot_product_attention(dtype, model_configs, model, ckpt_attn, workspace
 
     # UnfusedDotProductAttention backend
     if unfused_attn_supported:
-        attn_mask_type = config.attn_mask_type
-        config.attn_mask_type = "arbitrary"
+        if swa:
+            attn_mask_type = config.attn_mask_type
+            config.attn_mask_type = "arbitrary"
         unfused_attn_fwd, unfused_attn_bwd = _run_dot_product_attention(
             dtype, config, "UnfusedDotProductAttention", ckpt_attn, qkv_layout, workspace_opt, swa,
         )
-        config.attn_mask_type = attn_mask_type
+        if swa:
+            config.attn_mask_type = attn_mask_type
 
     # FusedAttention backend
     if fused_attn_supported:

From e2fcf1d2c9fb90e786d3a24569905080dc5b50db Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Fri, 15 Dec 2023 11:51:49 -0800
Subject: [PATCH 09/16] add docstring for get_swa func

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 tests/pytorch/test_fused_attn.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/pytorch/test_fused_attn.py b/tests/pytorch/test_fused_attn.py
index 1652d00915..596ba27c85 100644
--- a/tests/pytorch/test_fused_attn.py
+++ b/tests/pytorch/test_fused_attn.py
@@ -199,6 +199,8 @@ def _is_unfused_attention_supported(config: ModelConfig) -> bool:
 param_types_lean = [torch.bfloat16]
 
 def get_swa(seq_q, seq_kv, w=None):
+    """Generate a random sliding window size (left, right) if w is None,
+    and create its equivalent attention mask in [seq_q, seq_kv] shape"""
     if w is None:
         w = torch.randint(0, seq_kv, [2], dtype=torch.int32, device="cuda")
     m = torch.ones(seq_q, seq_kv, dtype=torch.bool, device="cuda")

From 0f6b23502e5da54bdb5ac927d11c889b11b633a4 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Fri, 15 Dec 2023 12:24:43 -0800
Subject: [PATCH 10/16] move repeated code into a function

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 transformer_engine/pytorch/attention.py   | 74 ++++++++---------------
 transformer_engine/pytorch/transformer.py | 29 +++------
 2 files changed, 31 insertions(+), 72 deletions(-)

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index 913341e2f2..d19bb7391d 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -1249,6 +1249,25 @@ def run_iteratively(q, k, v):
 
     return qkv_layout, q, k, v
 
+def check_set_window_size(
+        attn_mask_type: str,
+        window_size_in: Tuple[int, int] = None,
+        window_size_out: Tuple[int, int] = None,
+    ):
+    """Check if sliding window size is compliant with mask type and if not,
+    assert or set it to the appropriate size
+    """
+    if "causal" in attn_mask_type:
+        if window_size_in is None:
+            window_size_out = (-1, 0)
+        else:
+            assert (
+                window_size_in[1] == 0
+            ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
+    else:
+        if window_size_in is None:
+            window_size_out = (-1, -1)
+    return window_size_out
 
 class FlashAttention(torch.nn.Module):
     """Dot product attention, using HazyResearch flash-attn package:
@@ -1294,16 +1313,7 @@ def forward(
     ) -> torch.Tensor:
         """flash-attn fprop"""
 
-        if "causal" in attn_mask_type:
-            if window_size is None:
-                window_size = (-1, 0)
-            else:
-                assert (
-                    window_size[1] == 0
-                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
-        else:
-            if window_size is None:
-                window_size = (-1, -1)
+        window_size = check_set_window_size(attn_mask_type, window_size, window_size)
 
         assert (
             query_layer.dtype in [torch.float16, torch.bfloat16]
@@ -1960,16 +1970,7 @@ def __init__(
             attn_mask_type = "padding_causal"
         self.attn_mask_type = attn_mask_type
         self.window_size = window_size
-        if "causal" in attn_mask_type:
-            if window_size is None:
-                self.window_size = (-1, 0)
-            else:
-                assert (
-                    window_size[1] == 0
-                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
-        else:
-            if window_size is None:
-                self.window_size = (-1, -1)
+        self.window_size = check_set_window_size(attn_mask_type, window_size, self.window_size)
         self.tp_size = tp_size if tp_group is None else get_distributed_world_size(tp_group)
         self.tp_group = tp_group
         self.get_rng_state_tracker = get_rng_state_tracker
@@ -2197,16 +2198,7 @@ def forward(
             ), "Keys and values must have the same shape!"
 
         if attn_mask_type is not None:
-            if "causal" in attn_mask_type:
-                if window_size is None:
-                    window_size = (-1, 0)
-                else:
-                    assert (
-                        window_size[1] == 0
-                    ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
-            else:
-                if window_size is None:
-                    window_size = (-1, -1)
+            window_size = check_set_window_size(attn_mask_type, window_size, window_size)
         if attn_mask_type is None:
             attn_mask_type = self.attn_mask_type
         else:
@@ -2624,16 +2616,7 @@ def __init__(
 
         self.attn_mask_type = attn_mask_type
         self.window_size = window_size
-        if "causal" in attn_mask_type:
-            if window_size is None:
-                self.window_size = (-1, 0)
-            else:
-                assert (
-                    window_size[1] == 0
-                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
-        else:
-            if window_size is None:
-                self.window_size = (-1, -1)
+        self.window_size = check_set_window_size(attn_mask_type, window_size, self.window_size)
         self.layer_number = layer_number
         self.input_layernorm = input_layernorm
         self.attention_type = attention_type
@@ -2915,16 +2898,7 @@ def forward(
         # hidden_states: [sq, b, h]
 
         if attn_mask_type is not None:
-            if "causal" in attn_mask_type:
-                if window_size is None:
-                    window_size = (-1, 0)
-                else:
-                    assert (
-                        window_size[1] == 0
-                    ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
-            else:
-                if window_size is None:
-                    window_size = (-1, -1)
+            window_size = check_set_window_size(attn_mask_type, window_size, window_size)
         if attn_mask_type is None:
             attn_mask_type = self.attn_mask_type
         if window_size is None:
diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py
index c757ce637a..5f3c35a65c 100644
--- a/transformer_engine/pytorch/transformer.py
+++ b/transformer_engine/pytorch/transformer.py
@@ -12,7 +12,11 @@
 
 import transformer_engine_extensions as tex
 from transformer_engine.pytorch.module import LayerNormMLP, LayerNorm, RMSNorm
-from transformer_engine.pytorch.attention import InferenceParams, MultiheadAttention
+from transformer_engine.pytorch.attention import (
+    InferenceParams,
+    MultiheadAttention,
+    check_set_window_size,
+)
 from transformer_engine.pytorch.jit import (
     set_jit_fusion_options,
     warmup_jit_bias_dropout_add_all_dtypes,
@@ -259,16 +263,7 @@ def __init__(
 
         self.self_attn_mask_type = self_attn_mask_type
         self.window_size = window_size
-        if "causal" in self_attn_mask_type:
-            if window_size is None:
-                self.window_size = (-1, 0)
-            else:
-                assert (
-                    window_size[1] == 0
-                ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
-        else:
-            if window_size is None:
-                self.window_size = (-1, -1)
+        self.window_size = check_set_window_size(self_attn_mask_type, window_size, self.window_size)
         params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype
         ub_tp_comm_overlap = ub_tp_comm_overlap and bool(int(os.getenv("NVTE_UB_OVERLAP", "1")))
         ub_bulk_wgrad = ub_tp_comm_overlap and bool(int(os.getenv("NVTE_UB_BULK_WGRAD", "1")))
@@ -584,17 +579,7 @@ def forward(
         """
 
         if self_attn_mask_type is not None:
-            if "causal" in self_attn_mask_type:
-                if window_size is None:
-                    window_size = (-1, 0)
-                else:
-                    assert (
-                        window_size[1] == 0
-                    ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
-            else:
-                if window_size is None:
-                    window_size = (-1, -1)
-
+            window_size = check_set_window_size(self_attn_mask_type, window_size, window_size)
         if self_attn_mask_type is None:
             self_attn_mask_type = self.self_attn_mask_type
         if window_size is None:

From 67dd2d97223303c3ea666d95a0de7e4e461fc2bf Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Fri, 15 Dec 2023 12:32:34 -0800
Subject: [PATCH 11/16] revert mask change

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 tests/pytorch/test_sanity.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py
index 7cfc981f17..da4714c7ea 100644
--- a/tests/pytorch/test_sanity.py
+++ b/tests/pytorch/test_sanity.py
@@ -518,7 +518,7 @@ def test_sanity_bert(dtype, bs, fp8_recipe, model, skip_wgrad, zero_centered_gam
             apply_residual_connection_post_layernorm=True,
             output_layernorm=True,
             zero_centered_gamma=zero_centered_gamma,
-            self_attn_mask_type="causal",
+            self_attn_mask_type="padding",
             normalization=normalization,
         )
         .to(dtype=dtype)

From 576112da7d1097639d0e090437ce80d37cd8fd23 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Fri, 15 Dec 2023 22:23:30 +0000
Subject: [PATCH 12/16] add determinism filter and fix FA warning message

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 transformer_engine/pytorch/attention.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index 1268666f84..73917f4521 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -2005,8 +2005,8 @@ def __init__(
         if _flash_attn_2_available and self.deterministic:
             self.use_flash_attention = False
             warnings.warn(
-                "Disabling usage of FlashAttention since version 2 does not support deterministic"
-                "execution. In order to use FA with deterministic behavior, please install"
+                "Disabling usage of FlashAttention since version 2 does not support deterministic "
+                "execution. In order to use FA with deterministic behavior, please install "
                 "FlashAttention version 1."
             )
 
@@ -2361,6 +2361,13 @@ def forward(
             use_fused_attention = (use_fused_attention
                                   and is_backend_avail)
 
+        # Filter: determinism.
+        if (use_fused_attention
+            and fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]
+            and self.deterministic
+            and self.device_compute_capability != (9, 0)):
+                use_fused_attention = False
+
         # Select FusedAttention on sm90 and FlashAttention on others for performance
         if (use_flash_attention
             and use_fused_attention

From 9c7de54bf6591e531a80c5d6b325bac872b8b268 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Fri, 15 Dec 2023 22:37:36 +0000
Subject: [PATCH 13/16] add message for determinism filter

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 transformer_engine/pytorch/attention.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index 73917f4521..8ef9a178f4 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -2327,7 +2327,7 @@ def forward(
             use_fused_attention = False
 
         # Filter: Attention mask type.
-        #    attn_mask_type(s)   |     supported backends
+        #   attn_mask_type(s)    |     supported backends
         # ------------------------------------------------
         #   no_mask              |     All
         #   padding              |     UnfusedDotProductAttention, FlashAttention, FusedAttention
@@ -2362,6 +2362,17 @@ def forward(
                                   and is_backend_avail)
 
         # Filter: determinism.
+        # backend                                  | deterministic
+        # ---------------------------------------------------------
+        # flash-attn v1                            | yes
+        # flash-attn v2                            | no
+        # FusedAttnBackend["F16_max512_seqlen"]    | yes
+        # FusedAttnBackend["F16_arbitrary_seqlen"] | workspace optimization path: yes; otherwise: no
+        # UnfusedDotProductAttention               | yes
+        #
+        # Note that FusedAttnBackend["F16_arbitrary_seqlen"] only has workspace optimization path
+        # on sm90 architectures.
+        #
         if (use_fused_attention
             and fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]
             and self.deterministic

From 53bf1f522e1465413c380e719a8b227450cc69de Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Fri, 15 Dec 2023 23:11:38 +0000
Subject: [PATCH 14/16] simplify check_set_window_size()

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 transformer_engine/pytorch/attention.py | 25 ++++++++++++-------------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index 806d7190ac..f0bdf100fe 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -1251,23 +1251,22 @@ def run_iteratively(q, k, v):
 
 def check_set_window_size(
         attn_mask_type: str,
-        window_size_in: Tuple[int, int] = None,
-        window_size_out: Tuple[int, int] = None,
+        window_size: Tuple[int, int] = None,
     ):
     """Check if sliding window size is compliant with mask type and if not,
     assert or set it to the appropriate size
     """
     if "causal" in attn_mask_type:
-        if window_size_in is None:
-            window_size_out = (-1, 0)
+        if window_size is None:
+            window_size = (-1, 0)
         else:
             assert (
-                window_size_in[1] == 0
+                window_size[1] == 0
             ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!"
     else:
-        if window_size_in is None:
-            window_size_out = (-1, -1)
-    return window_size_out
+        if window_size is None:
+            window_size = (-1, -1)
+    return window_size
 
 class FlashAttention(torch.nn.Module):
     """Dot product attention, using HazyResearch flash-attn package:
@@ -1313,7 +1312,7 @@ def forward(
     ) -> torch.Tensor:
         """flash-attn fprop"""
 
-        window_size = check_set_window_size(attn_mask_type, window_size, window_size)
+        window_size = check_set_window_size(attn_mask_type, window_size)
 
         assert (
             query_layer.dtype in [torch.float16, torch.bfloat16]
@@ -1971,7 +1970,7 @@ def __init__(
             attn_mask_type = "padding_causal"
         self.attn_mask_type = attn_mask_type
         self.window_size = window_size
-        self.window_size = check_set_window_size(attn_mask_type, window_size, self.window_size)
+        self.window_size = check_set_window_size(attn_mask_type, self.window_size)
         self.tp_size = tp_size if tp_group is None else get_distributed_world_size(tp_group)
         self.tp_group = tp_group
         self.get_rng_state_tracker = get_rng_state_tracker
@@ -2200,7 +2199,7 @@ def forward(
             ), "Keys and values must have the same shape!"
 
         if attn_mask_type is not None:
-            window_size = check_set_window_size(attn_mask_type, window_size, window_size)
+            window_size = check_set_window_size(attn_mask_type, window_size)
         if attn_mask_type is None:
             attn_mask_type = self.attn_mask_type
         else:
@@ -2636,7 +2635,7 @@ def __init__(
 
         self.attn_mask_type = attn_mask_type
         self.window_size = window_size
-        self.window_size = check_set_window_size(attn_mask_type, window_size, self.window_size)
+        self.window_size = check_set_window_size(attn_mask_type, self.window_size)
         self.layer_number = layer_number
         self.input_layernorm = input_layernorm
         self.attention_type = attention_type
@@ -2918,7 +2917,7 @@ def forward(
         # hidden_states: [sq, b, h]
 
         if attn_mask_type is not None:
-            window_size = check_set_window_size(attn_mask_type, window_size, window_size)
+            window_size = check_set_window_size(attn_mask_type, window_size)
         if attn_mask_type is None:
             attn_mask_type = self.attn_mask_type
         if window_size is None:

From 24252a2df78c6420123eab7c21dbea50facc0c1f Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Fri, 15 Dec 2023 15:53:36 -0800
Subject: [PATCH 15/16] fix check_set_window_size in transformer layers

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 transformer_engine/pytorch/transformer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py
index 5f3c35a65c..2647e78a7b 100644
--- a/transformer_engine/pytorch/transformer.py
+++ b/transformer_engine/pytorch/transformer.py
@@ -263,7 +263,7 @@ def __init__(
 
         self.self_attn_mask_type = self_attn_mask_type
         self.window_size = window_size
-        self.window_size = check_set_window_size(self_attn_mask_type, window_size, self.window_size)
+        self.window_size = check_set_window_size(self_attn_mask_type, self.window_size)
         params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype
         ub_tp_comm_overlap = ub_tp_comm_overlap and bool(int(os.getenv("NVTE_UB_OVERLAP", "1")))
         ub_bulk_wgrad = ub_tp_comm_overlap and bool(int(os.getenv("NVTE_UB_BULK_WGRAD", "1")))
@@ -579,7 +579,7 @@ def forward(
         """
 
         if self_attn_mask_type is not None:
-            window_size = check_set_window_size(self_attn_mask_type, window_size, window_size)
+            window_size = check_set_window_size(self_attn_mask_type, window_size)
         if self_attn_mask_type is None:
             self_attn_mask_type = self.self_attn_mask_type
         if window_size is None:

From 87c9826a8c63df5516536b61b874a533e57a4ee2 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Fri, 15 Dec 2023 15:54:42 -0800
Subject: [PATCH 16/16] fix indent

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 transformer_engine/pytorch/attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index f0bdf100fe..24a90af39c 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -2378,7 +2378,7 @@ def forward(
             and fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]
             and self.deterministic
             and self.device_compute_capability != (9, 0)):
-                use_fused_attention = False
+            use_fused_attention = False
 
         # Select FusedAttention on sm90 and FlashAttention on others for performance
         if (use_flash_attention