From fae6572757b1db6d367937ec441f5d2b51ef2821 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Fri, 26 Apr 2024 22:37:26 +0000
Subject: [PATCH 01/21] fix inconsistency for attn mask; now True means
 participating in attn

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 tests/paddle/test_layers.py                    |  4 ++--
 tests/pytorch/fused_attn/test_fused_attn.py    |  4 +---
 tests/pytorch/test_numerics.py                 |  5 +++--
 .../fused_softmax/scaled_masked_softmax.cu     |  2 +-
 transformer_engine/paddle/csrc/custom_ops.cu   |  4 ++--
 transformer_engine/paddle/layer/softmax.py     |  4 ++--
 transformer_engine/paddle/utils.py             |  2 +-
 transformer_engine/pytorch/softmax.py          | 18 +++++++++---------
 transformer_engine/pytorch/utils.py            |  4 ++++
 9 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/tests/paddle/test_layers.py b/tests/paddle/test_layers.py
index 1aeafe030a..25b040e37e 100644
--- a/tests/paddle/test_layers.py
+++ b/tests/paddle/test_layers.py
@@ -858,7 +858,7 @@ def test_dot_product_attention(bs, hidden_size, num_heads, q_seqlen, kv_seqlen,
     q_actual_seqlen = paddle.randint(low=20, high=q_seqlen, shape=(bs,), dtype='int32')
     kv_actual_seqlen = paddle.randint(low=20, high=kv_seqlen, shape=(bs,),
                                       dtype='int32') if attn_type == 'cross' else q_actual_seqlen
-    attn_mask = paddle.ones(shape=(bs, 1, q_seqlen, kv_seqlen), dtype='bool')
+    attn_mask = paddle.zeros(shape=(bs, 1, q_seqlen, kv_seqlen), dtype='bool')
 
     grad_out = paddle.normal(mean=0.0, std=0.02,
                              shape=(bs, q_seqlen, num_heads, head_size)).astype('float32')
@@ -867,7 +867,7 @@ def test_dot_product_attention(bs, hidden_size, num_heads, q_seqlen, kv_seqlen,
     grad_out = grad_out.astype(math_dtype)
 
     for i in range(0, bs):
-        attn_mask[i, 0, 0:q_actual_seqlen[i], 0:kv_actual_seqlen[i]] = False
+        attn_mask[i, 0, 0:q_actual_seqlen[i], 0:kv_actual_seqlen[i]] = True
 
     head_size = hidden_size // num_heads
     layer_te = te.DotProductAttention(num_heads,
diff --git a/tests/pytorch/fused_attn/test_fused_attn.py b/tests/pytorch/fused_attn/test_fused_attn.py
index 40cfdd34b7..5df4de36af 100644
--- a/tests/pytorch/fused_attn/test_fused_attn.py
+++ b/tests/pytorch/fused_attn/test_fused_attn.py
@@ -226,7 +226,6 @@ def get_swa(seq_q, seq_kv, w=None):
     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])
-    ml = ~ ml
     return w, ml
 
 
@@ -558,12 +557,11 @@ def _run_dot_product_attention(
                     .to(dtype=torch.bool).unsqueeze(0).unsqueeze(0).unsqueeze(0)], dim=0)
             attention_mask = (
                     attention_mask_q.to(device="cuda"), attention_mask_kv.to(device="cuda"))
+    window_size = None
     if swa:
         window_size, attention_mask = get_swa(config.max_seqlen_q, config.max_seqlen_kv)
     elif "causal" in config.attn_mask_type:
         window_size, attention_mask = (-1, 0), None
-    else:
-        window_size, attention_mask = None, None
 
     alibi_slopes = None
     if config.attn_bias_type == "alibi" and config.alibi_type == "custom":
diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index 90cfce8a6f..f066e75841 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -76,7 +76,7 @@ def __init__(self, hidden_size, eps, num_attention_heads, embed, num_layers, seq
 
 
 def get_causal_attn_mask(sq: int) -> torch.Tensor:
-    return torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool()
+    return torch.tril(torch.ones(sq, sq, device="cuda"), diagonal=0).bool()
 
 
 def dtype_tols(dtype: torch.dtype) -> Dict[str, float]:
@@ -324,6 +324,7 @@ def __init__(self, hidden_size: int, num_attention_heads: int):
         )
 
     def forward(self, x, attention_mask=None):
+        attention_mask = attention_mask.logical_not() if attention_mask is not None else None
         output = self.mhsa(x, x, x, attn_mask=attention_mask, need_weights=False)
         if isinstance(output, tuple):
             output = output[0]
@@ -914,7 +915,7 @@ def _test_granular_accuracy(block, bs, dtype, config):
 def _test_dpa_accuracy(block, bs, dtype, config):
     reset_rng_states()
 
-    mask = torch.triu(torch.ones(config.seq_len, config.seq_len, dtype=torch.bool, device="cuda"), diagonal=1)
+    mask = torch.tril(torch.ones(config.seq_len, config.seq_len, dtype=torch.bool, device="cuda"), diagonal=0)
     query, key, value = [
         torch.randn(
             (config.seq_len, bs, config.num_attention_heads, config.embed),
diff --git a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu
index 7a7194878e..2d02c94d49 100644
--- a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu
+++ b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu
@@ -286,7 +286,7 @@ __global__ void scaled_masked_softmax_warp_forward(
 
 #pragma unroll
                   for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
-                      if (temp_mask[element] != 1) {
+                      if (temp_mask[element] == 1) {
                           elements[i][it + element] = (acc_t)temp_data[element] * scale;
                       } else {
                           elements[i][it + element] = -10000.0;
diff --git a/transformer_engine/paddle/csrc/custom_ops.cu b/transformer_engine/paddle/csrc/custom_ops.cu
index 7dde3d1db2..d467dc6ec7 100644
--- a/transformer_engine/paddle/csrc/custom_ops.cu
+++ b/transformer_engine/paddle/csrc/custom_ops.cu
@@ -1304,12 +1304,12 @@ __global__ __launch_bounds__(BLOCK_SIZE) void mask_to_actual_seqlens_kernel(
     int q = 0, kv = 0;
     for (unsigned int q_idx = tid * kv_seqlen; q_idx < q_seqlen * kv_seqlen;
          q_idx += BLOCK_SIZE * kv_seqlen) {
-        q += (mask[q_idx + batch_offset] ? 0 : 1);
+        q += (mask[q_idx + batch_offset] ? 1 : 0);
     }
 
     if (need_kv) {
         for (unsigned int kv_idx = tid; kv_idx < kv_seqlen; kv_idx += BLOCK_SIZE) {
-            kv += (mask[kv_idx + batch_offset] ? 0 : 1);
+            kv += (mask[kv_idx + batch_offset] ? 1 : 0);
         }
     }
     __syncthreads();
diff --git a/transformer_engine/paddle/layer/softmax.py b/transformer_engine/paddle/layer/softmax.py
index b195f0305f..e310ef7974 100644
--- a/transformer_engine/paddle/layer/softmax.py
+++ b/transformer_engine/paddle/layer/softmax.py
@@ -32,8 +32,8 @@
 def _get_default_causal_mask(seqlen: int) -> paddle.Tensor:
     """Return the causal upper triangular mask for softmax input"""
     if seqlen not in _default_causal_mask:
-        _default_causal_mask[seqlen] = paddle.triu(paddle.ones((seqlen, seqlen)),
-                                                   diagonal=1).cast('bool')
+        _default_causal_mask[seqlen] = paddle.tril(paddle.ones((seqlen, seqlen)),
+                                                   diagonal=0).cast('bool')
     return _default_causal_mask[seqlen]
 
 
diff --git a/transformer_engine/paddle/utils.py b/transformer_engine/paddle/utils.py
index 1ed6e27062..50c0305c81 100644
--- a/transformer_engine/paddle/utils.py
+++ b/transformer_engine/paddle/utils.py
@@ -63,7 +63,7 @@ def attention_mask_func(attention_scores: paddle.Tensor,
 
     def _masked_fill(x, mask, value):
         y = paddle.full(x.shape, value, x.dtype)
-        return paddle.where(mask, y, x)
+        return paddle.where(mask, x, y)
 
     attention_scores = _masked_fill(attention_scores, attention_mask, -10000.0)
     return attention_scores
diff --git a/transformer_engine/pytorch/softmax.py b/transformer_engine/pytorch/softmax.py
index 57fccd80ad..279f680581 100644
--- a/transformer_engine/pytorch/softmax.py
+++ b/transformer_engine/pytorch/softmax.py
@@ -23,12 +23,12 @@
 def _get_default_causal_mask(sq: int, sk: int) -> torch.Tensor:
     """Return the causal upper triangular mask for softmax input"""
     if sq == 1:
-        return torch.zeros((1, sk), dtype=torch.bool, device="cuda")
+        return torch.ones((1, sk), dtype=torch.bool, device="cuda")
 
     matrix_shape = (sq, sk)
     if matrix_shape not in _default_causal_mask:
-        diagonal_offset = sk - sq + 1
-        _default_causal_mask[matrix_shape] = torch.triu(
+        diagonal_offset = sk - sq
+        _default_causal_mask[matrix_shape] = torch.tril(
             torch.ones(sq, sk, dtype=torch.bool, device="cuda"),
             diagonal=diagonal_offset)
     return _default_causal_mask[matrix_shape]
@@ -42,9 +42,9 @@ def _get_onnx_export_causal_mask(
     ONNX does not support dynamic control-flow and requires non-square masks when
     using a KV-cache (seq_k's length len(context)+len(generative) while seq_q's length is 1).
 
-    Argument `onnx_causal_mask` is a square triu (k=1) mask that is sliced to the correct
+    Argument `onnx_causal_mask` is a square tril (k=1) mask that is sliced to the correct
     shape for GPT context and generation phases.
-    In the context phase the derived mask is a square triu of shape (seq_k, seq_k), and in
+    In the context phase the derived mask is a square tril of shape (seq_k, seq_k), and in
     the generation phase the mask is rectangular with shape (1, seq_k).
     """
     assert len(onnx_causal_mask.size()) == 2
@@ -226,8 +226,8 @@ def symbolic(
         # Captures the logic of function scaled_masked_softmax_warp_forward.
         # output = softmax(mask(input*scale)
         # Computed as:
-        #   masked_scaled = (1 - mask)*(input*scale)
-        #   softmax_mask = mask * -10000
+        #   masked_scaled = mask*(input*scale)
+        #   softmax_mask = (1 - mask) * -10000
         #   output = softmax(masked_scaled + softmax_mask)
         scale_input = g.op("Constant", value_t=torch.tensor(scale, dtype=torch.float16))
         scaled = g.op("Mul", inputs, scale_input)
@@ -235,8 +235,8 @@ def symbolic(
         inv_mask = g.op("Sub", one, mask)
         # Note: type is hard coded because softmax uses FP16 or BF16
         neg_tenK = g.op("Constant", value_t=torch.tensor(-10000., dtype=torch.float16))
-        softmax_mask = g.op("Mul", mask, neg_tenK)
-        masked_scaled = g.op("Mul", inv_mask, scaled)
+        softmax_mask = g.op("Mul", inv_mask, neg_tenK)
+        masked_scaled = g.op("Mul", mask, scaled)
         masked = g.op("Add", masked_scaled, softmax_mask)
         out = g.op("Softmax", masked)
         return out
diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py
index f60f8c29c7..2718bac62d 100644
--- a/transformer_engine/pytorch/utils.py
+++ b/transformer_engine/pytorch/utils.py
@@ -44,6 +44,10 @@ def attention_mask_func(
     attention_scores: torch.Tensor, attention_mask: torch.Tensor
 ) -> torch.Tensor:
     """Get attention mask"""
+    #attention_scores.masked_fill_(attention_mask<=0, -10000.0)
+    #attention_mask = attention_mask.logical_not()
+    #attention_mask = ~ attention_mask
+    #attention_scores = attention_scores.where(attention_mask<=0, attention_scores, -10000.0)
     attention_scores.masked_fill_(attention_mask, -10000.0)
     return attention_scores
 

From dc6869d4f04311d81609253fdc96adbfe8acffb9 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Fri, 26 Apr 2024 22:38:47 +0000
Subject: [PATCH 02/21] fix sliding window window_size for decoder+padding
 combination

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

diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py
index 5b6fc1e5c3..0a40dc2cfc 100644
--- a/transformer_engine/pytorch/transformer.py
+++ b/transformer_engine/pytorch/transformer.py
@@ -652,10 +652,11 @@ def forward(
 
         # Cross attention.
         if self.layer_type == "decoder":
+            #window_size = check_set_window_size("padding", None)
             inter_attention_outputs = self.inter_attention(
                 hidden_states,
                 attention_mask=enc_dec_attn_mask,
-                window_size=window_size,
+                #window_size=window_size,
                 encoder_output=encoder_output,
                 is_first_microbatch=is_first_microbatch,
                 checkpoint_core_attention=checkpoint_core_attention,

From 81c4605e1c3c3924177d8e90b118173553a51804 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Wed, 1 May 2024 22:39:41 +0000
Subject: [PATCH 03/21] revert paddle changes regarding mask

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 tests/paddle/test_layers.py                  | 4 ++--
 transformer_engine/paddle/csrc/custom_ops.cu | 4 ++--
 transformer_engine/paddle/layer/softmax.py   | 4 ++--
 transformer_engine/paddle/utils.py           | 2 +-
 4 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/tests/paddle/test_layers.py b/tests/paddle/test_layers.py
index 25b040e37e..1aeafe030a 100644
--- a/tests/paddle/test_layers.py
+++ b/tests/paddle/test_layers.py
@@ -858,7 +858,7 @@ def test_dot_product_attention(bs, hidden_size, num_heads, q_seqlen, kv_seqlen,
     q_actual_seqlen = paddle.randint(low=20, high=q_seqlen, shape=(bs,), dtype='int32')
     kv_actual_seqlen = paddle.randint(low=20, high=kv_seqlen, shape=(bs,),
                                       dtype='int32') if attn_type == 'cross' else q_actual_seqlen
-    attn_mask = paddle.zeros(shape=(bs, 1, q_seqlen, kv_seqlen), dtype='bool')
+    attn_mask = paddle.ones(shape=(bs, 1, q_seqlen, kv_seqlen), dtype='bool')
 
     grad_out = paddle.normal(mean=0.0, std=0.02,
                              shape=(bs, q_seqlen, num_heads, head_size)).astype('float32')
@@ -867,7 +867,7 @@ def test_dot_product_attention(bs, hidden_size, num_heads, q_seqlen, kv_seqlen,
     grad_out = grad_out.astype(math_dtype)
 
     for i in range(0, bs):
-        attn_mask[i, 0, 0:q_actual_seqlen[i], 0:kv_actual_seqlen[i]] = True
+        attn_mask[i, 0, 0:q_actual_seqlen[i], 0:kv_actual_seqlen[i]] = False
 
     head_size = hidden_size // num_heads
     layer_te = te.DotProductAttention(num_heads,
diff --git a/transformer_engine/paddle/csrc/custom_ops.cu b/transformer_engine/paddle/csrc/custom_ops.cu
index d467dc6ec7..7dde3d1db2 100644
--- a/transformer_engine/paddle/csrc/custom_ops.cu
+++ b/transformer_engine/paddle/csrc/custom_ops.cu
@@ -1304,12 +1304,12 @@ __global__ __launch_bounds__(BLOCK_SIZE) void mask_to_actual_seqlens_kernel(
     int q = 0, kv = 0;
     for (unsigned int q_idx = tid * kv_seqlen; q_idx < q_seqlen * kv_seqlen;
          q_idx += BLOCK_SIZE * kv_seqlen) {
-        q += (mask[q_idx + batch_offset] ? 1 : 0);
+        q += (mask[q_idx + batch_offset] ? 0 : 1);
     }
 
     if (need_kv) {
         for (unsigned int kv_idx = tid; kv_idx < kv_seqlen; kv_idx += BLOCK_SIZE) {
-            kv += (mask[kv_idx + batch_offset] ? 1 : 0);
+            kv += (mask[kv_idx + batch_offset] ? 0 : 1);
         }
     }
     __syncthreads();
diff --git a/transformer_engine/paddle/layer/softmax.py b/transformer_engine/paddle/layer/softmax.py
index e310ef7974..b195f0305f 100644
--- a/transformer_engine/paddle/layer/softmax.py
+++ b/transformer_engine/paddle/layer/softmax.py
@@ -32,8 +32,8 @@
 def _get_default_causal_mask(seqlen: int) -> paddle.Tensor:
     """Return the causal upper triangular mask for softmax input"""
     if seqlen not in _default_causal_mask:
-        _default_causal_mask[seqlen] = paddle.tril(paddle.ones((seqlen, seqlen)),
-                                                   diagonal=0).cast('bool')
+        _default_causal_mask[seqlen] = paddle.triu(paddle.ones((seqlen, seqlen)),
+                                                   diagonal=1).cast('bool')
     return _default_causal_mask[seqlen]
 
 
diff --git a/transformer_engine/paddle/utils.py b/transformer_engine/paddle/utils.py
index 50c0305c81..1ed6e27062 100644
--- a/transformer_engine/paddle/utils.py
+++ b/transformer_engine/paddle/utils.py
@@ -63,7 +63,7 @@ def attention_mask_func(attention_scores: paddle.Tensor,
 
     def _masked_fill(x, mask, value):
         y = paddle.full(x.shape, value, x.dtype)
-        return paddle.where(mask, x, y)
+        return paddle.where(mask, y, x)
 
     attention_scores = _masked_fill(attention_scores, attention_mask, -10000.0)
     return attention_scores

From 4ca194cee11bd30b4a9cd8188a9a0bddaaefaefa Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Wed, 1 May 2024 22:57:09 +0000
Subject: [PATCH 04/21]  revert softmax to 1-mask;0-keep

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 .../common/fused_softmax/scaled_masked_softmax.cu               | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu
index 2d02c94d49..7a7194878e 100644
--- a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu
+++ b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu
@@ -286,7 +286,7 @@ __global__ void scaled_masked_softmax_warp_forward(
 
 #pragma unroll
                   for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
-                      if (temp_mask[element] == 1) {
+                      if (temp_mask[element] != 1) {
                           elements[i][it + element] = (acc_t)temp_data[element] * scale;
                       } else {
                           elements[i][it + element] = -10000.0;

From 3a4b3771820ff463917da38f554f8a1abc47f3c1 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Thu, 2 May 2024 00:39:31 +0000
Subject: [PATCH 05/21] enforce 1-mask out; 0-keep rule for jax masks

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 tests/jax/test_fused_attn.py         | 15 ++++++---------
 tests/jax/utils.py                   |  2 +-
 transformer_engine/jax/fused_attn.py | 17 +++++++----------
 3 files changed, 14 insertions(+), 20 deletions(-)

diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index 483f070559..ae792bc25a 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -68,7 +68,7 @@ def general_dot_product_attention(query: ArrayLike, key: ArrayLike, value: Array
     if mask is not None:
         if mask.ndim != logits.ndim:
             mask = jnp.expand_dims(mask, axis=-3)
-        logits = jnp.where(mask, logits, jnp.finfo(dtype).min)
+        logits = jnp.where(mask, jnp.finfo(dtype).min, logits)
 
     softmax_out = jax.nn.softmax(logits).astype(dtype)
 
@@ -96,9 +96,9 @@ def make_decoder_mask(q_tokens: ArrayLike, kv_tokens: ArrayLike) -> Array:
     """
     q_idxs = jnp.broadcast_to(jnp.arange(q_tokens.shape[-1], dtype=jnp.int32), q_tokens.shape)
     kv_idxs = jnp.broadcast_to(jnp.arange(kv_tokens.shape[-1], dtype=jnp.int32), kv_tokens.shape)
-    causal_mask = make_attention_mask(q_idxs, kv_idxs, jnp.greater_equal)
-    padding_mask = make_attention_mask(q_tokens > 0, kv_tokens > 0)
-    return combine_masks(causal_mask, padding_mask)
+    inv_causal_mask = make_attention_mask(q_idxs, kv_idxs, jnp.greater_equal)
+    inv_padding_mask = make_attention_mask(q_tokens > 0, kv_tokens > 0)
+    return jnp.logical_not(combine_masks(inv_causal_mask, inv_padding_mask))
 
 
 def jax_dpa(query, key, value, bias, q_token, kv_token, dropout_rng, **kwargs):
@@ -109,7 +109,7 @@ def jax_dpa(query, key, value, bias, q_token, kv_token, dropout_rng, **kwargs):
     if is_causal_mask(attn_mask_type):
         mask = make_decoder_mask(q_token, kv_token)
     else:
-        mask = make_attention_mask(q_token > 0, kv_token > 0)
+        mask = jnp.logical_not(make_attention_mask(q_token > 0, kv_token > 0))
 
     output = general_dot_product_attention(query,
                                            key,
@@ -132,10 +132,7 @@ def customcall_fused_dpa(query, key, value, bias, q_token, kv_token, dropout_rng
     if is_causal_mask(attn_mask_type):
         mask = make_decoder_mask(q_token, kv_token)
     else:
-        mask = make_attention_mask(q_token > 0, kv_token > 0)
-
-    # mask invert
-    mask = jnp.logical_not(mask)
+        mask = jnp.logical_not(make_attention_mask(q_token > 0, kv_token > 0))
 
     qkv_layout = kwargs.pop('qkv_layout')
     match qkv_layout:
diff --git a/tests/jax/utils.py b/tests/jax/utils.py
index c8e1b1b183..648005fdfa 100644
--- a/tests/jax/utils.py
+++ b/tests/jax/utils.py
@@ -625,7 +625,7 @@ def qkv_init(key, shape, dtype):
                 # position should only attend to those key positions that have already
                 # been generated and cached, not the remaining zero elements.
                 mask = combine_masks(
-                    mask,
+                    jnp.logical_not(mask),
                     jnp.broadcast_to(
                         jnp.arange(length) <= cur_index,
                 # (1, 1, length) represent (head dim, query length, key length)
diff --git a/transformer_engine/jax/fused_attn.py b/transformer_engine/jax/fused_attn.py
index 8b32163811..bd3599aed6 100644
--- a/transformer_engine/jax/fused_attn.py
+++ b/transformer_engine/jax/fused_attn.py
@@ -121,8 +121,7 @@ def _fused_attn_fwd_qkvpacked_rule(qkv: jnp.ndarray, bias: jnp.ndarray | None, m
         actual_seqlen = jnp.full((batch,), seqlen, dtype=jnp.int32)
     else:
         assert mask is not None
-        mask = jnp.logical_not(mask)
-        actual_seqlen = jnp.sum(mask, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+        actual_seqlen = jnp.sum(mask==False, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
     output, softmax_aux, rng_state = fused_attn_fwd_qkvpacked(
         qkv,
         bias,
@@ -208,13 +207,12 @@ def _fused_attn_fwd_kvpacked_rule(q, kv, bias, mask, seed, attn_bias_type, attn_
         kv_actual_seqlen = jnp.full((batch,), s_kv, dtype=jnp.int32)
     else:
         assert mask is not None
-        mask = jnp.logical_not(mask)
-        q_actual_seqlen = jnp.sum(mask, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+        q_actual_seqlen = jnp.sum(mask==False, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
         if attn_mask_type == AttnMaskType.PADDING_MASK:
-            kv_actual_seqlen = jnp.sum(mask, axis=-1, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+            kv_actual_seqlen = jnp.sum(mask==False, axis=-1, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
         else:
             # When mask is causal, the actual seqlen is not the last row, use max to find it
-            kv_actual_seqlen = jnp.max(jnp.sum(mask, axis=-1, dtype=jnp.int32), axis=(-1, -2))
+            kv_actual_seqlen = jnp.max(jnp.sum(mask==False, axis=-1, dtype=jnp.int32), axis=(-1, -2))
 
     output, softmax_aux, rng_state = fused_attn_fwd_kvpacked(
         q,
@@ -304,13 +302,12 @@ def _fused_attn_fwd_rule(q, k, v, bias, mask, seed, attn_bias_type, attn_mask_ty
         kv_actual_seqlen = jnp.full((batch,), s_kv, dtype=jnp.int32)
     else:
         assert mask is not None
-        mask = jnp.logical_not(mask)
-        q_actual_seqlen = jnp.sum(mask, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+        q_actual_seqlen = jnp.sum(mask==False, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
         if attn_mask_type == AttnMaskType.PADDING_MASK:
-            kv_actual_seqlen = jnp.sum(mask, axis=-1, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+            kv_actual_seqlen = jnp.sum(mask==False, axis=-1, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
         else:
             # When mask is causal, the actual seqlen is not the last row, use max to find it
-            kv_actual_seqlen = jnp.max(jnp.sum(mask, axis=-1, dtype=jnp.int32), axis=(-1, -2))
+            kv_actual_seqlen = jnp.max(jnp.sum(mask==False, axis=-1, dtype=jnp.int32), axis=(-1, -2))
 
     output, softmax_aux, rng_state = fused_attn_fwd(q,
                                                     k,

From dfa0ecefb4831b062be04c0da3ed34d0678c61d2 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Thu, 2 May 2024 17:08:35 +0000
Subject: [PATCH 06/21] fix jax lint

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 transformer_engine/jax/fused_attn.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/transformer_engine/jax/fused_attn.py b/transformer_engine/jax/fused_attn.py
index bd3599aed6..77120914f0 100644
--- a/transformer_engine/jax/fused_attn.py
+++ b/transformer_engine/jax/fused_attn.py
@@ -121,7 +121,7 @@ def _fused_attn_fwd_qkvpacked_rule(qkv: jnp.ndarray, bias: jnp.ndarray | None, m
         actual_seqlen = jnp.full((batch,), seqlen, dtype=jnp.int32)
     else:
         assert mask is not None
-        actual_seqlen = jnp.sum(mask==False, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+        actual_seqlen = jnp.sum(not mask, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
     output, softmax_aux, rng_state = fused_attn_fwd_qkvpacked(
         qkv,
         bias,
@@ -207,12 +207,13 @@ def _fused_attn_fwd_kvpacked_rule(q, kv, bias, mask, seed, attn_bias_type, attn_
         kv_actual_seqlen = jnp.full((batch,), s_kv, dtype=jnp.int32)
     else:
         assert mask is not None
-        q_actual_seqlen = jnp.sum(mask==False, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+        q_actual_seqlen = jnp.sum(not mask, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
         if attn_mask_type == AttnMaskType.PADDING_MASK:
-            kv_actual_seqlen = jnp.sum(mask==False, axis=-1, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+            kv_actual_seqlen = jnp.sum(
+                not mask, axis=-1, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
         else:
             # When mask is causal, the actual seqlen is not the last row, use max to find it
-            kv_actual_seqlen = jnp.max(jnp.sum(mask==False, axis=-1, dtype=jnp.int32), axis=(-1, -2))
+            kv_actual_seqlen = jnp.max(jnp.sum(not mask, axis=-1, dtype=jnp.int32), axis=(-1, -2))
 
     output, softmax_aux, rng_state = fused_attn_fwd_kvpacked(
         q,
@@ -302,12 +303,13 @@ def _fused_attn_fwd_rule(q, k, v, bias, mask, seed, attn_bias_type, attn_mask_ty
         kv_actual_seqlen = jnp.full((batch,), s_kv, dtype=jnp.int32)
     else:
         assert mask is not None
-        q_actual_seqlen = jnp.sum(mask==False, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+        q_actual_seqlen = jnp.sum(not mask, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
         if attn_mask_type == AttnMaskType.PADDING_MASK:
-            kv_actual_seqlen = jnp.sum(mask==False, axis=-1, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+            kv_actual_seqlen = jnp.sum(
+                not mask, axis=-1, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
         else:
             # When mask is causal, the actual seqlen is not the last row, use max to find it
-            kv_actual_seqlen = jnp.max(jnp.sum(mask==False, axis=-1, dtype=jnp.int32), axis=(-1, -2))
+            kv_actual_seqlen = jnp.max(jnp.sum(not mask, axis=-1, dtype=jnp.int32), axis=(-1, -2))
 
     output, softmax_aux, rng_state = fused_attn_fwd(q,
                                                     k,

From aa6eaca26e7d63004ffea2f1829bdf03d48eee51 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Thu, 2 May 2024 17:26:10 +0000
Subject: [PATCH 07/21] revert pytorch mask changes; some kept in tests

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 tests/pytorch/fused_attn/test_fused_attn.py |  1 +
 tests/pytorch/test_numerics.py              |  5 ++---
 transformer_engine/pytorch/softmax.py       | 18 +++++++++---------
 transformer_engine/pytorch/utils.py         |  4 ----
 4 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/tests/pytorch/fused_attn/test_fused_attn.py b/tests/pytorch/fused_attn/test_fused_attn.py
index 5df4de36af..1ef22894ea 100644
--- a/tests/pytorch/fused_attn/test_fused_attn.py
+++ b/tests/pytorch/fused_attn/test_fused_attn.py
@@ -226,6 +226,7 @@ def get_swa(seq_q, seq_kv, w=None):
     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])
+    ml = ~ ml
     return w, ml
 
 
diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index f066e75841..90cfce8a6f 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -76,7 +76,7 @@ def __init__(self, hidden_size, eps, num_attention_heads, embed, num_layers, seq
 
 
 def get_causal_attn_mask(sq: int) -> torch.Tensor:
-    return torch.tril(torch.ones(sq, sq, device="cuda"), diagonal=0).bool()
+    return torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool()
 
 
 def dtype_tols(dtype: torch.dtype) -> Dict[str, float]:
@@ -324,7 +324,6 @@ def __init__(self, hidden_size: int, num_attention_heads: int):
         )
 
     def forward(self, x, attention_mask=None):
-        attention_mask = attention_mask.logical_not() if attention_mask is not None else None
         output = self.mhsa(x, x, x, attn_mask=attention_mask, need_weights=False)
         if isinstance(output, tuple):
             output = output[0]
@@ -915,7 +914,7 @@ def _test_granular_accuracy(block, bs, dtype, config):
 def _test_dpa_accuracy(block, bs, dtype, config):
     reset_rng_states()
 
-    mask = torch.tril(torch.ones(config.seq_len, config.seq_len, dtype=torch.bool, device="cuda"), diagonal=0)
+    mask = torch.triu(torch.ones(config.seq_len, config.seq_len, dtype=torch.bool, device="cuda"), diagonal=1)
     query, key, value = [
         torch.randn(
             (config.seq_len, bs, config.num_attention_heads, config.embed),
diff --git a/transformer_engine/pytorch/softmax.py b/transformer_engine/pytorch/softmax.py
index 279f680581..57fccd80ad 100644
--- a/transformer_engine/pytorch/softmax.py
+++ b/transformer_engine/pytorch/softmax.py
@@ -23,12 +23,12 @@
 def _get_default_causal_mask(sq: int, sk: int) -> torch.Tensor:
     """Return the causal upper triangular mask for softmax input"""
     if sq == 1:
-        return torch.ones((1, sk), dtype=torch.bool, device="cuda")
+        return torch.zeros((1, sk), dtype=torch.bool, device="cuda")
 
     matrix_shape = (sq, sk)
     if matrix_shape not in _default_causal_mask:
-        diagonal_offset = sk - sq
-        _default_causal_mask[matrix_shape] = torch.tril(
+        diagonal_offset = sk - sq + 1
+        _default_causal_mask[matrix_shape] = torch.triu(
             torch.ones(sq, sk, dtype=torch.bool, device="cuda"),
             diagonal=diagonal_offset)
     return _default_causal_mask[matrix_shape]
@@ -42,9 +42,9 @@ def _get_onnx_export_causal_mask(
     ONNX does not support dynamic control-flow and requires non-square masks when
     using a KV-cache (seq_k's length len(context)+len(generative) while seq_q's length is 1).
 
-    Argument `onnx_causal_mask` is a square tril (k=1) mask that is sliced to the correct
+    Argument `onnx_causal_mask` is a square triu (k=1) mask that is sliced to the correct
     shape for GPT context and generation phases.
-    In the context phase the derived mask is a square tril of shape (seq_k, seq_k), and in
+    In the context phase the derived mask is a square triu of shape (seq_k, seq_k), and in
     the generation phase the mask is rectangular with shape (1, seq_k).
     """
     assert len(onnx_causal_mask.size()) == 2
@@ -226,8 +226,8 @@ def symbolic(
         # Captures the logic of function scaled_masked_softmax_warp_forward.
         # output = softmax(mask(input*scale)
         # Computed as:
-        #   masked_scaled = mask*(input*scale)
-        #   softmax_mask = (1 - mask) * -10000
+        #   masked_scaled = (1 - mask)*(input*scale)
+        #   softmax_mask = mask * -10000
         #   output = softmax(masked_scaled + softmax_mask)
         scale_input = g.op("Constant", value_t=torch.tensor(scale, dtype=torch.float16))
         scaled = g.op("Mul", inputs, scale_input)
@@ -235,8 +235,8 @@ def symbolic(
         inv_mask = g.op("Sub", one, mask)
         # Note: type is hard coded because softmax uses FP16 or BF16
         neg_tenK = g.op("Constant", value_t=torch.tensor(-10000., dtype=torch.float16))
-        softmax_mask = g.op("Mul", inv_mask, neg_tenK)
-        masked_scaled = g.op("Mul", mask, scaled)
+        softmax_mask = g.op("Mul", mask, neg_tenK)
+        masked_scaled = g.op("Mul", inv_mask, scaled)
         masked = g.op("Add", masked_scaled, softmax_mask)
         out = g.op("Softmax", masked)
         return out
diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py
index 2718bac62d..f60f8c29c7 100644
--- a/transformer_engine/pytorch/utils.py
+++ b/transformer_engine/pytorch/utils.py
@@ -44,10 +44,6 @@ def attention_mask_func(
     attention_scores: torch.Tensor, attention_mask: torch.Tensor
 ) -> torch.Tensor:
     """Get attention mask"""
-    #attention_scores.masked_fill_(attention_mask<=0, -10000.0)
-    #attention_mask = attention_mask.logical_not()
-    #attention_mask = ~ attention_mask
-    #attention_scores = attention_scores.where(attention_mask<=0, attention_scores, -10000.0)
     attention_scores.masked_fill_(attention_mask, -10000.0)
     return attention_scores
 

From 44062a9f3a05cbb8da2696e1920cf0774998ff1d Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Thu, 2 May 2024 19:00:08 +0000
Subject: [PATCH 08/21] revert to jax fused attn on main

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 transformer_engine/jax/fused_attn.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/transformer_engine/jax/fused_attn.py b/transformer_engine/jax/fused_attn.py
index 77120914f0..8b32163811 100644
--- a/transformer_engine/jax/fused_attn.py
+++ b/transformer_engine/jax/fused_attn.py
@@ -121,7 +121,8 @@ def _fused_attn_fwd_qkvpacked_rule(qkv: jnp.ndarray, bias: jnp.ndarray | None, m
         actual_seqlen = jnp.full((batch,), seqlen, dtype=jnp.int32)
     else:
         assert mask is not None
-        actual_seqlen = jnp.sum(not mask, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+        mask = jnp.logical_not(mask)
+        actual_seqlen = jnp.sum(mask, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
     output, softmax_aux, rng_state = fused_attn_fwd_qkvpacked(
         qkv,
         bias,
@@ -207,13 +208,13 @@ def _fused_attn_fwd_kvpacked_rule(q, kv, bias, mask, seed, attn_bias_type, attn_
         kv_actual_seqlen = jnp.full((batch,), s_kv, dtype=jnp.int32)
     else:
         assert mask is not None
-        q_actual_seqlen = jnp.sum(not mask, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+        mask = jnp.logical_not(mask)
+        q_actual_seqlen = jnp.sum(mask, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
         if attn_mask_type == AttnMaskType.PADDING_MASK:
-            kv_actual_seqlen = jnp.sum(
-                not mask, axis=-1, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+            kv_actual_seqlen = jnp.sum(mask, axis=-1, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
         else:
             # When mask is causal, the actual seqlen is not the last row, use max to find it
-            kv_actual_seqlen = jnp.max(jnp.sum(not mask, axis=-1, dtype=jnp.int32), axis=(-1, -2))
+            kv_actual_seqlen = jnp.max(jnp.sum(mask, axis=-1, dtype=jnp.int32), axis=(-1, -2))
 
     output, softmax_aux, rng_state = fused_attn_fwd_kvpacked(
         q,
@@ -303,13 +304,13 @@ def _fused_attn_fwd_rule(q, k, v, bias, mask, seed, attn_bias_type, attn_mask_ty
         kv_actual_seqlen = jnp.full((batch,), s_kv, dtype=jnp.int32)
     else:
         assert mask is not None
-        q_actual_seqlen = jnp.sum(not mask, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+        mask = jnp.logical_not(mask)
+        q_actual_seqlen = jnp.sum(mask, axis=-2, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
         if attn_mask_type == AttnMaskType.PADDING_MASK:
-            kv_actual_seqlen = jnp.sum(
-                not mask, axis=-1, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
+            kv_actual_seqlen = jnp.sum(mask, axis=-1, dtype=jnp.int32)[..., 0, 0]    # shape = (b,)
         else:
             # When mask is causal, the actual seqlen is not the last row, use max to find it
-            kv_actual_seqlen = jnp.max(jnp.sum(not mask, axis=-1, dtype=jnp.int32), axis=(-1, -2))
+            kv_actual_seqlen = jnp.max(jnp.sum(mask, axis=-1, dtype=jnp.int32), axis=(-1, -2))
 
     output, softmax_aux, rng_state = fused_attn_fwd(q,
                                                     k,

From 1c8a07385c220edd05dfae429659b7231772f1f4 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Thu, 2 May 2024 21:48:04 +0000
Subject: [PATCH 09/21] inverse mask logic for get_cu_seqlens/_and_indices in
 PyTorch implementation and mask generation in unit tests

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

diff --git a/tests/pytorch/fused_attn/test_fused_attn.py b/tests/pytorch/fused_attn/test_fused_attn.py
index 1ef22894ea..d80a6e5699 100644
--- a/tests/pytorch/fused_attn/test_fused_attn.py
+++ b/tests/pytorch/fused_attn/test_fused_attn.py
@@ -543,7 +543,7 @@ def _run_dot_product_attention(
             attention_mask_q = torch.Tensor([]).to(dtype=torch.bool)
             for i in range(config.batch_size):
                 attention_mask_q = torch.cat([attention_mask_q,
-                    torch.Tensor([True]*seqlens_q[i] + [False]*(config.max_seqlen_q-seqlens_q[i]))
+                    torch.Tensor([False]*seqlens_q[i] + [True]*(config.max_seqlen_q-seqlens_q[i]))
                     .to(dtype=torch.bool).unsqueeze(0).unsqueeze(0).unsqueeze(0)], dim=0)
             attention_mask = attention_mask_q.to(device="cuda")
         if config.attn_type == 'cross':
@@ -551,10 +551,10 @@ def _run_dot_product_attention(
             attention_mask_kv = torch.Tensor([]).to(dtype=torch.bool)
             for i in range(config.batch_size):
                 attention_mask_q = torch.cat([attention_mask_q,
-                    torch.Tensor([True]*seqlens_q[i] + [False]*(config.max_seqlen_q-seqlens_q[i]))
+                    torch.Tensor([False]*seqlens_q[i] + [True]*(config.max_seqlen_q-seqlens_q[i]))
                     .to(dtype=torch.bool).unsqueeze(0).unsqueeze(0).unsqueeze(0)], dim=0)
                 attention_mask_kv = torch.cat([attention_mask_kv, torch.Tensor(
-                    [True]*seqlens_kv[i] + [False]*(config.max_seqlen_kv-seqlens_kv[i]))
+                    [False]*seqlens_kv[i] + [True]*(config.max_seqlen_kv-seqlens_kv[i]))
                     .to(dtype=torch.bool).unsqueeze(0).unsqueeze(0).unsqueeze(0)], dim=0)
             attention_mask = (
                     attention_mask_q.to(device="cuda"), attention_mask_kv.to(device="cuda"))
@@ -856,7 +856,7 @@ def _run_transformer_layer(
         attention_mask_q = torch.Tensor([]).to(dtype=torch.bool)
         for i in range(config.batch_size):
             attention_mask_q = torch.cat([attention_mask_q,
-                torch.Tensor([True]*seqlens_q[i] + [False]*(config.max_seqlen_q-seqlens_q[i]))
+                torch.Tensor([False]*seqlens_q[i] + [True]*(config.max_seqlen_q-seqlens_q[i]))
                 .to(torch.bool).unsqueeze(0).unsqueeze(0).unsqueeze(0)], dim=0)
         attention_mask = attention_mask_q.to(device="cuda")
 
@@ -942,7 +942,7 @@ def _run_transformer_layer(
 
 model_configs_fp8_vs_f16 = {
     #  test:             b,  h, hg,   d,   sq,  skv,   p,      mask,      bias
-    "fp8_9 ": ModelConfig(2, 24, 24, 128, 2048, 2048, 0.0, "no_mask", "no_bias"),
+    "fp8_9" : ModelConfig(2, 24, 24, 128, 2048, 2048, 0.0, "no_mask", "no_bias"),
     "fp8_10": ModelConfig(2, 24, 24, 128, 2048, 2048, 0.0,  "causal", "no_bias"),
     "fp8_11": ModelConfig(2, 24, 12, 128, 2048, 2048, 0.0, "no_mask", "no_bias"),
     "fp8_12": ModelConfig(2, 24, 12, 128, 2048, 2048, 0.0,  "causal", "no_bias"),
@@ -1141,24 +1141,49 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd):
         dtype, config, False, qkv_layout)
 
     tols = dict(atol=5e-1, rtol=5e-2)
+    rmse_tol = 0.1
+    bwd_names = ['dq', 'dk', 'dv']
+    fwd_rmse = _rmse(fused_attn_fwd_fp8, fused_attn_fwd_f16)
+    fwd_range = max(fused_attn_fwd_fp8.max().item(),
+        fused_attn_fwd_f16.max().item()) - min(fused_attn_fwd_fp8.min().item(),
+        fused_attn_fwd_f16.min().item())
     if _NVTE_DEBUG:
-        print('[test_dpa_fp8_vs_f16]: ', tols)
+        print()
+        print('========== {:^25s} =========='.format('forward output'))
         print('fused_attn_fwd_fp8 min {:.6f} max {:.6f}'.format(
             fused_attn_fwd_fp8.min().item(),fused_attn_fwd_fp8.max().item()))
         print('fused_attn_fwd_f16 min {:.6f} max {:.6f}'.format(
             fused_attn_fwd_f16.min().item(), fused_attn_fwd_f16.max().item()))
-        print('fused_attn_fwd RMSE: {:.6f}'.format(
-            _rmse(fused_attn_fwd_fp8, fused_attn_fwd_f16)))
-    torch.testing.assert_close(fused_attn_fwd_fp8, fused_attn_fwd_f16, **tols)
+        print('fused_attn_fwd RMSE: {:.6f}'.format(fwd_rmse))
+        try:
+            torch.testing.assert_close(fused_attn_fwd_fp8, fused_attn_fwd_f16, **tols)
+        except Exception as e:
+            print(e)
+            print()
+    assert(fwd_rmse < rmse_tol * fwd_range
+        ), "FWD RMSE {:.5f} is over tolerance {:.5f} ({:.5f} * {:.5f})".format(
+        fwd_rmse, rmse_tol * fwd_range, rmse_tol, fwd_range)
     for i,_ in enumerate(fused_attn_bwd_f16):
+        bwd_rmse = _rmse(fused_attn_bwd_fp8[i], fused_attn_bwd_f16[i])
+        bwd_range = max(fused_attn_bwd_fp8[i].max().item(),
+            fused_attn_bwd_f16[i].max().item()) - min(fused_attn_bwd_fp8[i].min().item(),
+            fused_attn_bwd_f16[i].min().item())
         if _NVTE_DEBUG:
-            print('fused_attn_bwd_fp8 min {:.6f} max {:.6f}'.format(
+            print()
+            print('========== {:^25s} =========='.format(bwd_names[i]))
+            print('fused_attn_bwd_fp8[{}] min {:.6f} max {:.6f}'.format(i,
                 fused_attn_bwd_fp8[i].min().item(), fused_attn_bwd_fp8[i].max().item()))
-            print('fused_attn_bwd_f16 min {:.6f} max {:.6f}'.format(
+            print('fused_attn_bwd_f16[{}] min {:.6f} max {:.6f}'.format(i,
                 fused_attn_bwd_f16[i].min().item(), fused_attn_bwd_f16[i].max().item()))
-            print('fused_attn_bwd RMSE: {:.6f}'.format(
-                _rmse(fused_attn_bwd_fp8[i], fused_attn_bwd_f16[i])))
-        torch.testing.assert_close(fused_attn_bwd_fp8[i], fused_attn_bwd_f16[i], **tols)
+            print('fused_attn_bwd RMSE[{}]: {:.6f}'.format(i, bwd_rmse))
+            try:
+                torch.testing.assert_close(fused_attn_bwd_fp8[i], fused_attn_bwd_f16[i], **tols)
+            except Exception as e:
+                print(e)
+                print()
+        assert(bwd_rmse < rmse_tol * bwd_range
+            ), "BWD RMSE {:.5f} is over tolerance {:.5f} ({:.5f} * {:.5f})".format(
+            bwd_rmse, rmse_tol * bwd_range, rmse_tol, bwd_range)
 
 
 def _run_dpa_fp8_vs_f16(dtype, config, fp8_dpa, qkv_layout):
@@ -1229,7 +1254,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker:
             layout = layout.replace('h', 'hg')
             layout = layout.replace('t', 'tg')
         tensor_shape = [dim_to_num[j] for j in layout.split('_')]
-        tensor = 0.1 * torch.randn(tensor_shape, dtype=dtype, device="cuda")
+        tensor = torch.randn(tensor_shape, dtype=dtype, device="cuda")
         tensor_count = 1
         split_dim = 0
         for dim, l in enumerate(layout.split('_')):
@@ -1250,7 +1275,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker:
     qkv_format_kv = qkv_format_kv.replace('s', 'sq')
     out_grad_shape = [dim_to_num[i] for i in qkv_format_kv.split('_')]
     out_grad_shape_new = [*out_grad_shape[:-2], out_grad_shape[-2] * out_grad_shape[-1]]
-    out_grad = 0.1 * torch.randn(out_grad_shape_new, dtype=dtype, device="cuda")
+    out_grad = torch.randn(out_grad_shape_new, dtype=dtype, device="cuda")
 
     with fp8_autocast(enabled=fp8_dpa, fp8_recipe=fp8_recipe):
         out = dpa(inp[0], inp[1], inp[2],
@@ -1357,7 +1382,7 @@ def _run_custom_mha_fp8(dtype, config, backend):
     if backend == "FusedAttention":
         os.environ["NVTE_FUSED_ATTN"] = "1"
 
-    inp = 0.0001 * torch.randint(0, 100,
+    inp = 0.0001 * torch.randint(-100, 100,
             (config.batch_size * config.max_seqlen_q, config.num_heads * config.head_dim),
             dtype=dtype, device="cuda", requires_grad=True)
     seqlens = torch.full([config.batch_size], config.max_seqlen_q,
diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index dbc26d538d..d4fce6bc89 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -223,7 +223,7 @@ def get_cu_seqlens(mask: torch.Tensor) -> torch.Tensor:
     the samples in a batch.
     """
     mask = mask.squeeze(1).squeeze(1)
-    reduced_mask = mask.sum(dim=1)
+    reduced_mask = mask.logical_not().sum(dim=1)
     cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
     zero = torch.zeros(1, dtype=torch.int32, device="cuda")
     cu_seqlens = torch.cat((zero, cu_seqlens))
@@ -241,13 +241,13 @@ def get_cu_seqlens_and_indices(mask: torch.Tensor) -> Tuple[torch.Tensor, torch.
     mask = mask.squeeze(1).squeeze(1)
     bs, seqlen = mask.shape
 
-    reduced_mask = mask.sum(dim=1)
+    reduced_mask = mask.logical_not().sum(dim=1)
     cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
     zero = torch.zeros(1, dtype=torch.int32, device="cuda")
     cu_seqlens = torch.cat((zero, cu_seqlens))
 
     mask = mask.reshape(-1)
-    indices = mask.nonzero()
+    indices = mask.logical_not().nonzero()
     indices = indices.unsqueeze(-1)
 
     num_nonzeros = indices.shape[0]

From 19f271c384f7671e3b52c78fb6ca6911d4274bd1 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Thu, 2 May 2024 22:20:19 +0000
Subject: [PATCH 10/21] temporarily disable update_weight_scale_inv

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

diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py
index 92c7f26f59..f7e7ac810f 100644
--- a/tests/pytorch/test_recipe.py
+++ b/tests/pytorch/test_recipe.py
@@ -93,9 +93,9 @@ def test_amax_and_scale_update(
         ref_scale_forward = (fp8_format.value.max_fwd / ref_amax_forward) / (2 ** margin)
         ref_scale_backward = (fp8_format.value.max_bwd / ref_amax_backward) / (2 ** margin)
         ref_scale_inv_forward = torch.reciprocal(ref_scale_forward)
-        update_weight_scale_inv = is_first_microbatch is None or is_first_microbatch
-        if not update_weight_scale_inv:
-            ref_scale_inv_forward[1].copy_(scale_inv_forward[1])
+        #update_weight_scale_inv = is_first_microbatch is None or is_first_microbatch
+        #if not update_weight_scale_inv:
+        #    ref_scale_inv_forward[1].copy_(scale_inv_forward[1])
         ref_scale_inv_backward = torch.reciprocal(ref_scale_backward)
 
         # Make sure we are not trivially passing tests

From 2a489bfe92580c27d911e5c001bf7319300bfbfc Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Thu, 2 May 2024 22:47:01 +0000
Subject: [PATCH 11/21] enforce window_size for decoder

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 0a40dc2cfc..b920919a4a 100644
--- a/transformer_engine/pytorch/transformer.py
+++ b/transformer_engine/pytorch/transformer.py
@@ -652,11 +652,11 @@ def forward(
 
         # Cross attention.
         if self.layer_type == "decoder":
-            #window_size = check_set_window_size("padding", None)
+            window_size = check_set_window_size("padding", None)
             inter_attention_outputs = self.inter_attention(
                 hidden_states,
                 attention_mask=enc_dec_attn_mask,
-                #window_size=window_size,
+                window_size=window_size,
                 encoder_output=encoder_output,
                 is_first_microbatch=is_first_microbatch,
                 checkpoint_core_attention=checkpoint_core_attention,

From 87d02b666124f0290d7f5e592317e94afa3d57cc Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Thu, 2 May 2024 22:47:52 +0000
Subject: [PATCH 12/21] add docstring for mask definition 1-mask out;0-keep

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

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index d4fce6bc89..e045b0582d 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -3408,7 +3408,9 @@ def forward(
              a single tensor of [batch_size, 1, 1, seqlen_q] for self-attention, and a tuple of
              two tensors in shapes [batch_size, 1, 1, seqlen_q] and [batch_size, 1, 1, seqlen_kv]
              for cross-attention. For the 'arbitrary' mask type, it should be in a shape that is
-             broadcastable to [batch_size, num_heads, max_seqlen_q, max_seqlen_kv].
+             broadcastable to [batch_size, num_heads, max_seqlen_q, max_seqlen_kv]. A `True` value
+             means the corresponding position is masked out and a `False` means that position is
+             allowed to participate in attention.
         qkv_format: str, default = `None`
                    If provided, overrides :attr:`qkv_format` from initialization.
         cu_seqlens_q: Optional[torch.Tensor], default = `None`
@@ -4298,7 +4300,9 @@ def forward(
              a single tensor of [batch_size, 1, 1, seqlen_q] for self-attention, and a tuple of
              two tensors in shapes [batch_size, 1, 1, seqlen_q] and [batch_size, 1, 1, seqlen_kv]
              for cross-attention. For the 'arbitrary' mask type, it should be in a shape that is
-             broadcastable to [batch_size, num_heads, max_seqlen_q, max_seqlen_kv].
+             broadcastable to [batch_size, num_heads, max_seqlen_q, max_seqlen_kv]. A `True` value
+             means the corresponding position is masked out and a `False` means that position is
+             allowed to participate in attention.
         attn_mask_type: {'no_mask', 'padding', 'causal', 'padding_causal', 'arbitrary'},
                        default = `None`
                        type of attention mask passed into softmax operation.
diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py
index b920919a4a..72f3958c52 100644
--- a/transformer_engine/pytorch/transformer.py
+++ b/transformer_engine/pytorch/transformer.py
@@ -542,6 +542,8 @@ def forward(
                         It should be in [batch_size, 1, 1, seqlen_q] for 'padding' mask,
                         and broadcastable to [batch_size, num_heads, max_seqlen_q, max_seqlen_kv]
                         for 'arbitrary'. It should be 'None' for 'causal' and 'no_mask'.
+                        A `True` value means the corresponding position is masked out and
+                        a `False` means that position is allowed to participate in attention.
         self_attn_mask_type: {'no_mask', 'causal', 'padding', 'padding_causal', 'arbitrary'},
                             default = `causal`
                             Type of attention mask passed into softmax operation.
@@ -555,7 +557,9 @@ def forward(
              using `layer_type="decoder"`. It should be a tuple of two masks in
              [batch_size, 1, 1, seqlen_q] and [batch_size, 1, 1, seqlen_kv] for 'padding' mask.
              It should be broadcastable to [batch_size, num_heads, max_seqlen_q, max_seqlen_kv]
-             for 'arbitrary' mask. It should be 'None' for 'causal' and 'no_mask'.
+             for 'arbitrary' mask. It should be 'None' for 'causal' and 'no_mask'. A `True` value
+             means the corresponding position is masked out and a `False` means that position is
+             allowed to participate in attention.
         is_first_microbatch : {True, False, None}, default = None
                              During training using either gradient accumulation or
                              pipeline parallelism a minibatch of data is further split

From d759aeefac801fbab03d6f9561e857cf4caab51e Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Tue, 14 May 2024 01:17:19 +0000
Subject: [PATCH 13/21] add aux_ctx_tensors to save_for_backward

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

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index d32a9b78c8..f67b5f42fb 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -408,7 +408,7 @@ def forward(
         *tensors: Tuple[torch.Tensor, ...]
     ) -> Union[Tuple[torch.Tensor, ...], torch.Tensor]:
         assert 1 <= len(tensors) <= 3, f"Packing {len(tensors)} tensors not supported."
-        ctx.indices = indices
+        ctx.save_for_backward(indices)
         ctx.dim0 = tensors[0].shape[0]
         if len(tensors) == 1:
             return pack_tensor(indices, *tensors)
@@ -418,11 +418,12 @@ def forward(
 
     @staticmethod
     def backward(ctx, *grad_outputs: Tuple[torch.Tensor, ...]):
+        (indices,) = ctx.saved_tensors
         if len(grad_outputs) == 1:
-            return None, unpack_tensor(ctx.indices, ctx.dim0, *grad_outputs)
+            return None, unpack_tensor(indices, ctx.dim0, *grad_outputs)
         if len(grad_outputs) == 2:
-            return None, *unpack_2_tensors(ctx.indices, ctx.dim0, *grad_outputs)
-        return None, *unpack_3_tensors(ctx.indices, ctx.dim0, *grad_outputs)
+            return None, *unpack_2_tensors(indices, ctx.dim0, *grad_outputs)
+        return None, *unpack_3_tensors(indices, ctx.dim0, *grad_outputs)
 
 
 class UnpackTensor(torch.autograd.Function):
@@ -436,12 +437,13 @@ def forward(
         dim0: int,
         tensor: torch.Tensor,
     ) -> torch.Tensor:
-        ctx.indices = indices
+        ctx.save_for_backward(indices)
         return unpack_tensor(indices, dim0, tensor)
 
     @staticmethod
     def backward(ctx, grad_output):
-        return None, None, pack_tensor(ctx.indices, grad_output)
+        (indices,) = ctx.saved_tensors
+        return None, None, pack_tensor(indices, grad_output)
 
 
 def flash_attn_p2p_communicate(rank, send_tensor, send_dst,
@@ -868,8 +870,8 @@ def forward(ctx, is_training, q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q,
         else:
             out = out.view(-1, *out.shape[-2:])
 
-        ctx.save_for_backward(q, kv, out, softmax_lse, cu_seqlens_q, cu_seqlens_k)
-        ctx.rng_states = rng_states
+        ctx.save_for_backward(q, kv, out, softmax_lse,
+            cu_seqlens_q, cu_seqlens_k, rng_states, attn_biases)
         ctx.cp_group = cp_group
         ctx.cp_global_ranks = cp_global_ranks
         ctx.dropout_p = dropout_p
@@ -880,14 +882,14 @@ def forward(ctx, is_training, q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q,
         ctx.qkv_format = qkv_format
         ctx.attn_bias_type = attn_bias_type
         ctx.attn_bias_shape = None if attn_bias is None else attn_bias.shape
-        ctx.attn_biases = attn_biases
         ctx.deterministic = deterministic
         ctx.use_fused_attention = use_fused_attention
         return out
 
     @staticmethod
     def backward(ctx, dout):
-        q, kv, out, softmax_lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors
+        (q, kv, out, softmax_lse,
+            cu_seqlens_q, cu_seqlens_k, rng_states, attn_biases) = ctx.saved_tensors
 
         cp_size = get_distributed_world_size(ctx.cp_group)
         rank = get_distributed_rank(ctx.cp_group)
@@ -897,12 +899,12 @@ def backward(ctx, dout):
 
         qkv_layout = ctx.qkv_format + "_" + ctx.qkv_format + "_" + ctx.qkv_format
 
-        if ctx.attn_biases[0] is not None:
+        if attn_biases[0] is not None:
             # [b, np, sq, 2*cp, sk//(2*cp)]
             attn_dbias = torch.zeros(
                 *ctx.attn_bias_shape,
-                dtype=ctx.attn_biases[0].dtype,
-                device=ctx.attn_biases[0].device
+                dtype=attn_biases[0].dtype,
+                device=attn_biases[0].device
             )
             # [b, np, sq, 2*cp, sk//(2*cp)] -> [b, np, 2, sq//2, 2*cp, sk//(2*cp)]
             attn_dbias_ = attn_dbias.view(
@@ -985,9 +987,9 @@ def backward(ctx, dout):
                             # [2, sq//2, b, np, hn] -> [sq, b, np, hn]
                             out_ = out.view(-1, *out.shape[-3:])
                             dout_ = dout.view(-1, *dout.shape[-3:])
-                        aux_ctx_tensors = [softmax_lse, ctx.rng_states[cp_size-i-1]]
+                        aux_ctx_tensors = [softmax_lse, rng_states[cp_size-i-1]]
                         if attn_dbias is not None:
-                            aux_ctx_tensors += [ctx.attn_biases[cp_size-i-1]]
+                            aux_ctx_tensors += [attn_biases[cp_size-i-1]]
                         dq_, dk_, dv_, dbias_ = fused_attn_bwd(
                             ctx.max_seqlen_q, ctx.max_seqlen_k,
                             cu_seqlens_q, cu_seqlens_k,
@@ -1017,7 +1019,7 @@ def backward(ctx, dout):
                             dq_, dkv_[0], dkv_[1], cu_seqlens_q, cu_seqlens_k,
                             ctx.max_seqlen_q, ctx.max_seqlen_k,
                             ctx.dropout_p, ctx.softmax_scale, True,
-                            rng_state=ctx.rng_states[cp_size-i-1],
+                            rng_state=rng_states[cp_size-i-1],
                             **fa_optional_backward_kwargs
                         )
                 elif i >= (cp_size-rank-1):
@@ -1038,9 +1040,9 @@ def backward(ctx, dout):
                             # [2, sq//2, b, np, hn] -> [sq, b, np, hn]
                             out_ = out.view(-1, *out.shape[-3:])
                             dout_ = dout.view(-1, *dout.shape[-3:])
-                        aux_ctx_tensors = [softmax_lse, ctx.rng_states[cp_size-i-1]]
+                        aux_ctx_tensors = [softmax_lse, rng_states[cp_size-i-1]]
                         if attn_dbias is not None:
-                            aux_ctx_tensors += [ctx.attn_biases[cp_size-i-1]]
+                            aux_ctx_tensors += [attn_biases[cp_size-i-1]]
                         dq_, dk_, dv_, dbias_ = fused_attn_bwd(
                             ctx.max_seqlen_q, ctx.max_seqlen_k//2,
                             cu_seqlens_q, cu_seqlens_k//2,
@@ -1074,7 +1076,7 @@ def backward(ctx, dout):
                             dq_, dkv_[0], dkv_[1], cu_seqlens_q, cu_seqlens_k//2,
                             ctx.max_seqlen_q, ctx.max_seqlen_k//2,
                             ctx.dropout_p, ctx.softmax_scale, False,
-                            rng_state=ctx.rng_states[cp_size-i-1],
+                            rng_state=rng_states[cp_size-i-1],
                             **fa_optional_backward_kwargs
                         )
                 else:
@@ -1095,9 +1097,9 @@ def backward(ctx, dout):
                             # [2, sq//2, b, np, hn] -> [sq//2, b, np, hn]
                             out_ = out[1].contiguous()
                             dout_ = dout[1].contiguous()
-                        aux_ctx_tensors = [softmax_lse_, ctx.rng_states[cp_size-i-1]]
+                        aux_ctx_tensors = [softmax_lse_, rng_states[cp_size-i-1]]
                         if attn_dbias is not None:
-                            aux_ctx_tensors += [ctx.attn_biases[cp_size-i-1]]
+                            aux_ctx_tensors += [attn_biases[cp_size-i-1]]
                         dq_, dk_, dv_, dbias_ = fused_attn_bwd(
                             ctx.max_seqlen_q//2, ctx.max_seqlen_k,
                             cu_seqlens_q//2, cu_seqlens_k,
@@ -1135,14 +1137,14 @@ def backward(ctx, dout):
                             dq_, dkv_[0], dkv_[1], cu_seqlens_q//2, cu_seqlens_k,
                             ctx.max_seqlen_q//2, ctx.max_seqlen_k,
                             ctx.dropout_p, ctx.softmax_scale, False,
-                            rng_state=ctx.rng_states[cp_size-i-1],
+                            rng_state=rng_states[cp_size-i-1],
                             **fa_optional_backward_kwargs
                         )
             else:
                 if ctx.use_fused_attention:
-                    aux_ctx_tensors = [softmax_lse, ctx.rng_states[cp_size-i-1]]
+                    aux_ctx_tensors = [softmax_lse, rng_states[cp_size-i-1]]
                     if attn_dbias is not None:
-                        aux_ctx_tensors += [ctx.attn_biases[cp_size-i-1]]
+                        aux_ctx_tensors += [attn_biases[cp_size-i-1]]
                     dq_, dk_, dv_, dbias_ = fused_attn_bwd(
                         ctx.max_seqlen_q, ctx.max_seqlen_k,
                         cu_seqlens_q, cu_seqlens_k,
@@ -2300,9 +2302,8 @@ def forward(ctx, is_training, max_seqlen, cu_seqlens, qkv, qkv_dtype, attn_bias,
 
         ctx.fp8 = fp8 and int(os.getenv("NVTE_FP8_DPA_BWD", "1"))
         qkvo_tensors = (qkv, out_save) if not ctx.fp8 else (None, None)
-        ctx.save_for_backward(*qkvo_tensors, cu_seqlens, *fp8_tensors)
+        ctx.save_for_backward(*qkvo_tensors, cu_seqlens, *fp8_tensors, *aux_ctx_tensors)
         ctx.fp8_meta = fp8_meta
-        ctx.aux_ctx_tensors = aux_ctx_tensors
         ctx.max_seqlen = max_seqlen
         ctx.qkv_dtype = qkv_dtype
         ctx.attn_scale = attn_scale
@@ -2326,12 +2327,12 @@ def backward(ctx, d_out):
             d_out = d_out._data
 
         d_out = d_out.contiguous()
-        (qkv, out, cu_seqlens,
-            qkv_fp8, out_fp8, fwd_scales, fwd_scale_invs) = ctx.saved_tensors
-        if not ctx.aux_ctx_tensors[0].is_contiguous():
-            ctx.aux_ctx_tensors[0] = ctx.aux_ctx_tensors[0].contiguous()
+        (qkv, out, cu_seqlens, qkv_fp8, out_fp8,
+            fwd_scales, fwd_scale_invs, *aux_ctx_tensors) = ctx.saved_tensors
+        if not aux_ctx_tensors[0].is_contiguous():
+            aux_ctx_tensors[0] = aux_ctx_tensors[0].contiguous()
         if ctx.use_FAv2_bwd:
-            softmax_lse, rng_state = ctx.aux_ctx_tensors
+            softmax_lse, rng_state = aux_ctx_tensors
             dqkv = torch.empty_like(qkv)
             maybe_contiguous = lambda x: x.contiguous() if x.stride(-1) != 1 else x
             d_out, q, k, v, out = [maybe_contiguous(x)
@@ -2363,7 +2364,7 @@ def backward(ctx, d_out):
                     dqkv_fp8, *rest = fused_attn_bwd_qkvpacked(
                         ctx.max_seqlen, cu_seqlens,
                         qkv_fp8, out_fp8, d_out_fp8,
-                        fp8_dtype_forward, fp8_dtype_backward, ctx.aux_ctx_tensors,
+                        fp8_dtype_forward, fp8_dtype_backward, aux_ctx_tensors,
                         ctx.fused_attention_backend,
                         fwd_scale_invs[META_QKV], # d_scale_qkv,
                         fwd_scale_invs[META_S], # d_scale_s,
@@ -2398,7 +2399,7 @@ def backward(ctx, d_out):
                         d_out = d_out_f8tensor.from_float8(qkv.dtype)
                     dqkv, *rest = fused_attn_bwd_qkvpacked(
                         ctx.max_seqlen, cu_seqlens, qkv, out, d_out,
-                        ctx.qkv_dtype, ctx.qkv_dtype, ctx.aux_ctx_tensors,
+                        ctx.qkv_dtype, ctx.qkv_dtype, aux_ctx_tensors,
                         ctx.fused_attention_backend,
                         None, None, None, None, None, None, None, None, None, None,
                         ctx.attn_scale, ctx.dropout_p, ctx.fast_zero_fill,
@@ -2501,9 +2502,9 @@ def forward(ctx, is_training, max_seqlen_q, max_seqlen_kv, cu_seqlens_q, cu_seql
 
         ctx.fp8 = fp8 and int(os.getenv("NVTE_FP8_DPA_BWD", "1"))
         qkvo_tensors = (q, kv, out_save) if not ctx.fp8 else (None, None, None)
-        ctx.save_for_backward(*qkvo_tensors, cu_seqlens_q, cu_seqlens_kv, *fp8_tensors)
+        ctx.save_for_backward(*qkvo_tensors, cu_seqlens_q, cu_seqlens_kv,
+            *fp8_tensors, *aux_ctx_tensors)
         ctx.fp8_meta = fp8_meta
-        ctx.aux_ctx_tensors = aux_ctx_tensors
         ctx.max_seqlen_q = max_seqlen_q
         ctx.max_seqlen_kv = max_seqlen_kv
         ctx.qkv_dtype = qkv_dtype
@@ -2528,12 +2529,12 @@ def backward(ctx, d_out):
             d_out = d_out._data
 
         d_out = d_out.contiguous()
-        (q, kv, out, cu_seqlens_q, cu_seqlens_kv,
-            q_fp8, kv_fp8, out_fp8, fwd_scales, fwd_scale_invs) = ctx.saved_tensors
-        if not ctx.aux_ctx_tensors[0].is_contiguous():
-            ctx.aux_ctx_tensors[0] = ctx.aux_ctx_tensors[0].contiguous()
+        (q, kv, out, cu_seqlens_q, cu_seqlens_kv, q_fp8, kv_fp8, out_fp8,
+            fwd_scales, fwd_scale_invs, *aux_ctx_tensors) = ctx.saved_tensors
+        if not aux_ctx_tensors[0].is_contiguous():
+            aux_ctx_tensors[0] = aux_ctx_tensors[0].contiguous()
         if ctx.use_FAv2_bwd:
-            softmax_lse, rng_state = ctx.aux_ctx_tensors
+            softmax_lse, rng_state = aux_ctx_tensors
             dq = torch.empty_like(q)
             dkv = torch.empty_like(kv)
             maybe_contiguous = lambda x: x.contiguous() if x.stride(-1) != 1 else x
@@ -2567,7 +2568,7 @@ def backward(ctx, d_out):
                     dq_fp8, dkv_fp8, *rest = fused_attn_bwd_kvpacked(
                         ctx.max_seqlen_q, ctx.max_seqlen_kv, cu_seqlens_q, cu_seqlens_kv,
                         q_fp8, kv_fp8, out_fp8, d_out_fp8,
-                        fp8_dtype_forward, fp8_dtype_backward, ctx.aux_ctx_tensors,
+                        fp8_dtype_forward, fp8_dtype_backward, aux_ctx_tensors,
                         ctx.fused_attention_backend,
                         fwd_scale_invs[META_QKV], # d_scale_qkv,
                         fwd_scale_invs[META_S], # d_scale_s,
@@ -2614,7 +2615,7 @@ def backward(ctx, d_out):
                     dq, dkv, *rest = fused_attn_bwd_kvpacked(
                         ctx.max_seqlen_q, ctx.max_seqlen_kv, cu_seqlens_q, cu_seqlens_kv,
                         q, kv, out, d_out,
-                        ctx.qkv_dtype, ctx.qkv_dtype, ctx.aux_ctx_tensors,
+                        ctx.qkv_dtype, ctx.qkv_dtype, aux_ctx_tensors,
                         ctx.fused_attention_backend,
                         None, None, None, None, None, None, None, None, None, None,
                         ctx.attn_scale, ctx.dropout_p, ctx.fast_zero_fill,
@@ -2773,9 +2774,9 @@ def forward(ctx, is_training, max_seqlen_q, max_seqlen_kv, cu_seqlens_q, cu_seql
 
         ctx.fp8 = fp8 and int(os.getenv("NVTE_FP8_DPA_BWD", "1"))
         qkvo_tensors = (q, k, v, out_save) if not ctx.fp8 else (None, None, None, None)
-        ctx.save_for_backward(*qkvo_tensors, cu_seqlens_q, cu_seqlens_kv, *fp8_tensors)
+        ctx.save_for_backward(*qkvo_tensors, cu_seqlens_q, cu_seqlens_kv,
+            *fp8_tensors, *aux_ctx_tensors)
         ctx.fp8_meta = fp8_meta
-        ctx.aux_ctx_tensors = aux_ctx_tensors
         ctx.max_seqlen_q = max_seqlen_q
         ctx.max_seqlen_kv = max_seqlen_kv
         ctx.qkv_dtype = qkv_dtype
@@ -2800,12 +2801,12 @@ def backward(ctx, d_out):
             d_out = d_out._data
 
         d_out = d_out.contiguous()
-        (q, k, v, out, cu_seqlens_q, cu_seqlens_kv,
-            q_fp8, k_fp8, v_fp8, out_fp8, fwd_scales, fwd_scale_invs) = ctx.saved_tensors
-        if not ctx.aux_ctx_tensors[0].is_contiguous():
-            ctx.aux_ctx_tensors[0] = ctx.aux_ctx_tensors[0].contiguous()
+        (q, k, v, out, cu_seqlens_q, cu_seqlens_kv, q_fp8, k_fp8, v_fp8, out_fp8,
+            fwd_scales, fwd_scale_invs, *aux_ctx_tensors) = ctx.saved_tensors
+        if not aux_ctx_tensors[0].is_contiguous():
+            aux_ctx_tensors[0] = aux_ctx_tensors[0].contiguous()
         if ctx.use_FAv2_bwd:
-            softmax_lse, rng_state = ctx.aux_ctx_tensors
+            softmax_lse, rng_state = aux_ctx_tensors
             dq = torch.empty_like(q)
             dk = torch.empty_like(k)
             dv = torch.empty_like(v)
@@ -2840,7 +2841,7 @@ def backward(ctx, d_out):
                     dq_fp8, dk_fp8, dv_fp8, *rest = fused_attn_bwd(
                         ctx.max_seqlen_q, ctx.max_seqlen_kv, cu_seqlens_q, cu_seqlens_kv,
                         q_fp8, k_fp8, v_fp8, out_fp8, d_out_fp8,
-                        fp8_dtype_forward, fp8_dtype_backward, ctx.aux_ctx_tensors,
+                        fp8_dtype_forward, fp8_dtype_backward, aux_ctx_tensors,
                         ctx.fused_attention_backend,
                         fwd_scale_invs[META_QKV], # d_scale_qkv,
                         fwd_scale_invs[META_S], # d_scale_s,
@@ -2923,7 +2924,7 @@ def backward(ctx, d_out):
                     dq, dk, dv, *rest = fused_attn_bwd(
                         ctx.max_seqlen_q, ctx.max_seqlen_kv, cu_seqlens_q, cu_seqlens_kv,
                         q, k, v, out, d_out,
-                        ctx.qkv_dtype, ctx.qkv_dtype, ctx.aux_ctx_tensors,
+                        ctx.qkv_dtype, ctx.qkv_dtype, aux_ctx_tensors,
                         ctx.fused_attention_backend,
                         None, None, None, None, None, None, None, None, None, None,
                         ctx.attn_scale, ctx.dropout_p, ctx.fast_zero_fill,

From 2e84099c2f20ba0f2550bb5d28f7da030c262f6b Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Tue, 14 May 2024 22:17:27 +0000
Subject: [PATCH 14/21] tweak make_decoder_mask and make_mask in jax tests

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

diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index 80f89074fe..2666682c5b 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -90,24 +90,34 @@ def is_causal_mask(mask: AttnMaskType):
 
 def make_decoder_mask(q_tokens: ArrayLike, kv_tokens: ArrayLike) -> Array:
     """
-    Create padded causal mask
+    Create inverse padded causal mask where `True` means allowing the corresponding
+    position to participate in attention and `False` means masking out that position.
     """
     q_idxs = jnp.broadcast_to(jnp.arange(q_tokens.shape[-1], dtype=jnp.int32), q_tokens.shape)
     kv_idxs = jnp.broadcast_to(jnp.arange(kv_tokens.shape[-1], dtype=jnp.int32), kv_tokens.shape)
     inv_causal_mask = make_attention_mask(q_idxs, kv_idxs, jnp.greater_equal)
     inv_padding_mask = make_attention_mask(q_tokens > 0, kv_tokens > 0)
-    return jnp.logical_not(combine_masks(inv_causal_mask, inv_padding_mask))
+    return combine_masks(inv_causal_mask, inv_padding_mask)
 
+def make_mask(q_token: ArrayLike, kv_token: ArrayLike, attn_mask_type: AttnMaskType) -> Array:
+    """
+    Create attention mask based on mask type. A `True` value in the mask means
+    masking out the corresponding position and a `False` value means allowing
+    that position to participate in attention.
+    """
+    if is_causal_mask(attn_mask_type):
+        inv_mask = make_decoder_mask(q_token, kv_token)
+    else:
+        inv_mask = make_attention_mask(q_token > 0, kv_token > 0)
+    mask = jnp.logical_not(inv_mask)
+    return mask
 
 def jax_dpa(query, key, value, bias, q_token, kv_token, dropout_rng, **kwargs):
     """
     JAX native dot product attention implementation
     """
     attn_mask_type = kwargs['attn_mask_type']
-    if is_causal_mask(attn_mask_type):
-        mask = make_decoder_mask(q_token, kv_token)
-    else:
-        mask = jnp.logical_not(make_attention_mask(q_token > 0, kv_token > 0))
+    mask = make_mask(q_token, kv_token, attn_mask_type)
 
     output = general_dot_product_attention(query,
                                            key,
@@ -127,10 +137,7 @@ def customcall_fused_dpa(query, key, value, bias, q_token, kv_token, dropout_rng
     TE customcall dot product attention implementation
     """
     attn_mask_type = kwargs['attn_mask_type']
-    if is_causal_mask(attn_mask_type):
-        mask = make_decoder_mask(q_token, kv_token)
-    else:
-        mask = jnp.logical_not(make_attention_mask(q_token > 0, kv_token > 0))
+    mask = make_mask(q_token, kv_token, attn_mask_type)
 
     qkv_layout = kwargs.pop('qkv_layout')
     match qkv_layout:

From 49a38f08018b48811b5ad76155388bab4a377ea8 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Wed, 15 May 2024 17:30:01 +0000
Subject: [PATCH 15/21] skip dBias for shapes other than 1HSS; otherwise
 dq/dk/dv NaNs

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

diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index 2666682c5b..426d51e906 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -302,6 +302,8 @@ def test_backward(self):
         """
 
         self._setup_inputs()
+        if self.attn_bias_type != AttnBiasType.NO_BIAS and self.bias_shape != BiasShape.BIAS_1HSS:
+            pytest.skip("Bias gradient calculation is only supported for 1HSS bias shape.")
 
         def grad_func(func, *args, **kwargs):
             # Gradient is small, use a gradient multiplier to amplify the gradient

From 5bd7a1a37c532d8e2ef60da3b5f059bafb54f3dc Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Wed, 15 May 2024 17:44:16 +0000
Subject: [PATCH 16/21] expand attn_biases from list to variables in
 save_for_backward

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

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index f67b5f42fb..26ab511754 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -871,7 +871,7 @@ def forward(ctx, is_training, q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q,
             out = out.view(-1, *out.shape[-2:])
 
         ctx.save_for_backward(q, kv, out, softmax_lse,
-            cu_seqlens_q, cu_seqlens_k, rng_states, attn_biases)
+            cu_seqlens_q, cu_seqlens_k, rng_states, *attn_biases)
         ctx.cp_group = cp_group
         ctx.cp_global_ranks = cp_global_ranks
         ctx.dropout_p = dropout_p
@@ -889,7 +889,7 @@ def forward(ctx, is_training, q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q,
     @staticmethod
     def backward(ctx, dout):
         (q, kv, out, softmax_lse,
-            cu_seqlens_q, cu_seqlens_k, rng_states, attn_biases) = ctx.saved_tensors
+            cu_seqlens_q, cu_seqlens_k, rng_states, *attn_biases) = ctx.saved_tensors
 
         cp_size = get_distributed_world_size(ctx.cp_group)
         rank = get_distributed_rank(ctx.cp_group)

From bf9a851f2bcf80f201364bd4f91aaf501b72ea47 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Wed, 15 May 2024 17:55:15 +0000
Subject: [PATCH 17/21] fix use of variable before assignment in jax dact_lu

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

diff --git a/transformer_engine/jax/cpp_extensions.py b/transformer_engine/jax/cpp_extensions.py
index 041b350240..98eb1e4013 100644
--- a/transformer_engine/jax/cpp_extensions.py
+++ b/transformer_engine/jax/cpp_extensions.py
@@ -2661,7 +2661,7 @@ def partition(act_enum, mesh, arg_infos, result_infos):
         """
         act_lu partitioning
         """
-        del result_infos, act_enum
+        del result_infos
         x_spec = get_padded_spec(arg_infos[0])
         arg_shardings = tuple(arg_i.sharding for arg_i in arg_infos)
         out_sharding = NamedSharding(mesh, PartitionSpec(*x_spec[:-2], x_spec[-1]))
@@ -2790,7 +2790,7 @@ def partition(act_enum, mesh, arg_infos, result_infos):
         """
         dact_lu partition
         """
-        del result_infos, act_enum
+        del result_infos
         dx_sharding = NamedSharding(mesh, PartitionSpec(*get_padded_spec(arg_infos[1])))
         arg_shardings = tuple(arg_i.sharding for arg_i in arg_infos)
         out_shardings = dx_sharding

From 84aa282167b8ca06d01ba32b3bf8c2d05f857a71 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Wed, 15 May 2024 21:24:33 +0000
Subject: [PATCH 18/21] remove window size definition for decoder

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

diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py
index 72f3958c52..abb5423850 100644
--- a/transformer_engine/pytorch/transformer.py
+++ b/transformer_engine/pytorch/transformer.py
@@ -656,11 +656,9 @@ def forward(
 
         # Cross attention.
         if self.layer_type == "decoder":
-            window_size = check_set_window_size("padding", None)
             inter_attention_outputs = self.inter_attention(
                 hidden_states,
                 attention_mask=enc_dec_attn_mask,
-                window_size=window_size,
                 encoder_output=encoder_output,
                 is_first_microbatch=is_first_microbatch,
                 checkpoint_core_attention=checkpoint_core_attention,

From 37f25c3cc3edd2dfd63b51ce52882ebee035c67f Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Wed, 15 May 2024 21:47:24 +0000
Subject: [PATCH 19/21] add change notes in README for padding mask in PyTorch

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 README.rst | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/README.rst b/README.rst
index 936dfab077..5baf50153e 100644
--- a/README.rst
+++ b/README.rst
@@ -184,10 +184,17 @@ Compiling with FlashAttention-2
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 Transformer Engine release v0.11.0 adds support for FlashAttention-2 in PyTorch for improved performance. 
 
-It is a known issue that FlashAttention-2 compilation is resource-intensive and requires a large amount of RAM (see `bug <https://github.com/Dao-AILab/flash-attention/issues/358>`_), which may lead to out of memory errors during the installation of Transformer Engine. Please try setting **MAX_JOBS=1** in the environment to circumvent the issue. If the errors persist, install a supported version of FlashAttention-1 (v1.0.6 to v1.0.9).
+It is a known issue that FlashAttention-2 compilation is resource-intensive and requires a large amount of RAM (see `bug <https://github.com/Dao-AILab/flash-attention/issues/358>`_), which may lead to out of memory errors during the installation of Transformer Engine. Please try setting **MAX_JOBS=1** in the environment to circumvent the issue.
 
 Note that NGC PyTorch 23.08+ containers include FlashAttention-2.
 
+Important Changes
+===============
+
+v1.7: Padding mask definition for PyTorch
+^^^^^^^^^^^^^^^^^^^^
+In an effort to unify the definition and usage of the attention mask across all three frameworks in Transformer Engine, the padding mask has changed from `True` meaning inclusion of the corresponding position in attention to exclusion of that position in our PyTorch implementation. Since v1.7, all attention mask types follow the same definition where `True` means masking out the corresponding position and `False` means including that position in attention calculation.
+
 FP8 Convergence
 ===============
 

From 87cea82ec9d706993d887ad32178484bdd1092fd Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Wed, 15 May 2024 22:28:14 +0000
Subject: [PATCH 20/21] tweak padding mask notes in README

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 README.rst | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/README.rst b/README.rst
index 5baf50153e..09eea53b2b 100644
--- a/README.rst
+++ b/README.rst
@@ -188,13 +188,31 @@ It is a known issue that FlashAttention-2 compilation is resource-intensive and
 
 Note that NGC PyTorch 23.08+ containers include FlashAttention-2.
 
-Important Changes
-===============
+Breaking Changes
+================
 
 v1.7: Padding mask definition for PyTorch
-^^^^^^^^^^^^^^^^^^^^
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 In an effort to unify the definition and usage of the attention mask across all three frameworks in Transformer Engine, the padding mask has changed from `True` meaning inclusion of the corresponding position in attention to exclusion of that position in our PyTorch implementation. Since v1.7, all attention mask types follow the same definition where `True` means masking out the corresponding position and `False` means including that position in attention calculation.
 
+An example of this change is,
+
+.. code-block:: bash
+
+    # for a batch of 3 sequences where `a`s, `b`s and `c`s are the useful tokens
+    # and `0`s are the padding tokens,
+    [a, a, a, 0, 0,
+     b, b, 0, 0, 0,
+     c, c, c, c, 0]
+    # the padding mask for this batch before v1.7 is,
+    [ True,  True,  True, False, False,
+      True,  True, False, False, False,
+      True,  True,  True,  True, False]
+    # and for v1.7 onwards it should be,
+    [False, False, False,  True,  True,
+     False, False,  True,  True,  True,
+     False, False, False, False,  True]
+
 FP8 Convergence
 ===============
 

From f07b0bf7e2768a6f52a3887883a6049ff021fa5d Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Thu, 16 May 2024 19:24:03 +0000
Subject: [PATCH 21/21] expand list to tensors for save_for_backwards

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

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index 26ab511754..d4198e688d 100644
--- a/transformer_engine/pytorch/attention.py
+++ b/transformer_engine/pytorch/attention.py
@@ -871,7 +871,7 @@ def forward(ctx, is_training, q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q,
             out = out.view(-1, *out.shape[-2:])
 
         ctx.save_for_backward(q, kv, out, softmax_lse,
-            cu_seqlens_q, cu_seqlens_k, rng_states, *attn_biases)
+            cu_seqlens_q, cu_seqlens_k, *rng_states, *attn_biases)
         ctx.cp_group = cp_group
         ctx.cp_global_ranks = cp_global_ranks
         ctx.dropout_p = dropout_p
@@ -888,10 +888,11 @@ def forward(ctx, is_training, q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q,
 
     @staticmethod
     def backward(ctx, dout):
-        (q, kv, out, softmax_lse,
-            cu_seqlens_q, cu_seqlens_k, rng_states, *attn_biases) = ctx.saved_tensors
-
+        (q, kv, out, softmax_lse, cu_seqlens_q, cu_seqlens_k) = ctx.saved_tensors[:6]
         cp_size = get_distributed_world_size(ctx.cp_group)
+        rng_states = ctx.saved_tensors[6:6+cp_size]
+        attn_biases = ctx.saved_tensors[6+cp_size:6+cp_size*2]
+
         rank = get_distributed_rank(ctx.cp_group)
         send_dst = ctx.cp_global_ranks[(rank - 1) % cp_size]
         recv_src = ctx.cp_global_ranks[(rank + 1) % cp_size]