[JAX] Refactor fused attention (#711)

* Remove unused headers
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Unify the fused attn workspace size cpp code
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Reduce the skipped cases
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename self/cross attention to qkvpacked/kvpacked
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update attention mask docs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refine the attn mask implementations
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>

tests/jax/test_distributed_fused_attn.py

@@ -16,7 +16,7 @@ from distributed_test_base import compare_ops
 from utils import make_causal_mask, make_self_mask
 from transformer_engine.jax import fp8_autocast
 from transformer_engine.jax.fused_attn import is_fused_attn_kernel_available
-from transformer_engine.jax.fused_attn import self_fused_attn, cross_fused_attn
+from transformer_engine.jax.fused_attn import fused_attn_qkvpacked, fused_attn_kvpacked
 from transformer_engine.jax.fused_attn import AttnBiasType, AttnMaskType, QKVLayout
 
 DTYPES = [jnp.float16, jnp.bfloat16]
@@ -86,15 +86,15 @@ class TestDistributedSelfAttn:
 
         def target_func(qkv, bias, mask):
             return jnp.mean(
-                self_fused_attn(qkv,
-                                bias,
-                                mask,
-                                None,
-                                attn_bias_type=attn_bias_type,
-                                attn_mask_type=attn_mask_type,
-                                scaling_factor=scaling_factor,
-                                dropout_probability=dropout_prob,
-                                is_training=is_training))
+                fused_attn_qkvpacked(qkv,
+                                     bias,
+                                     mask,
+                                     None,
+                                     attn_bias_type=attn_bias_type,
+                                     attn_mask_type=attn_mask_type,
+                                     scaling_factor=scaling_factor,
+                                     dropout_probability=dropout_prob,
+                                     is_training=is_training))
 
         def ref_func(qkv, bias, mask):
             query, key, value = jnp.split(qkv, [1, 2], axis=-3)
@@ -192,16 +192,16 @@ class TestDistributedCrossAttn:
 
         def target_func(q, kv, mask):
             return jnp.mean(
-                cross_fused_attn(q,
-                                 kv,
-                                 None,
-                                 mask,
-                                 None,
-                                 attn_bias_type=attn_bias_type,
-                                 attn_mask_type=attn_mask_type,
-                                 scaling_factor=scaling_factor,
-                                 dropout_probability=dropout_prob,
-                                 is_training=is_training))
+                fused_attn_kvpacked(q,
+                                    kv,
+                                    None,
+                                    mask,
+                                    None,
+                                    attn_bias_type=attn_bias_type,
+                                    attn_mask_type=attn_mask_type,
+                                    scaling_factor=scaling_factor,
+                                    dropout_probability=dropout_prob,
+                                    is_training=is_training))
 
         def ref_func(query, kv, mask):
             key, value = jnp.split(kv, [1], axis=-3)

tests/jax/test_fused_attn.py

@@ -2,8 +2,6 @@
 #
 # See LICENSE for license information.
 """Tests for fused attention"""
-import sys
-
 from enum import Enum
 from dataclasses import dataclass
 from functools import partial
@@ -21,7 +19,7 @@ from jax import value_and_grad, jit
 from jax.typing import ArrayLike, DTypeLike
 
 from transformer_engine.jax.fused_attn import AttnBiasType, AttnMaskType, QKVLayout
-from transformer_engine.jax.fused_attn import self_fused_attn, cross_fused_attn, fused_attn
+from transformer_engine.jax.fused_attn import fused_attn_qkvpacked, fused_attn_kvpacked, fused_attn
 from transformer_engine.jax.cpp_extensions import FusedAttnHelper
 
 from transformer_engine_jax import NVTE_Fused_Attn_Backend
@@ -144,18 +142,22 @@ def customcall_fused_dpa(query, key, value, bias, q_token, kv_token, dropout_rng
         case QKVLayout.BS3HD:
             query, key, value = map(partial(jnp.expand_dims, axis=-3), [query, key, value])
             qkv = jnp.concatenate((query, key, value), axis=-3)
-            return self_fused_attn(qkv, bias, mask, dropout_rng, **kwargs).astype(query.dtype)
+            return fused_attn_qkvpacked(qkv, bias, mask, dropout_rng, **kwargs).astype(query.dtype)
         case QKVLayout.BSHD_BS2HD:
             key, value = map(partial(jnp.expand_dims, axis=-3), [key, value])
             kv = jnp.concatenate((key, value), axis=-3)
-            return cross_fused_attn(query, kv, bias, mask, dropout_rng,
-                                    **kwargs).astype(query.dtype)
+            return fused_attn_kvpacked(query, kv, bias, mask, dropout_rng,
+                                       **kwargs).astype(query.dtype)
         case QKVLayout.BSHD_BSHD_BSHD:
             return fused_attn(query, key, value, bias, mask, dropout_rng,
                               **kwargs).astype(query.dtype)
 
 
 class BiasShape(Enum):
+    """
+    Enum class to represent the different bias shapes used in the fused attention.
+    """
+
     BIAS_1HSS = '1HSS'
     BIAS_B1SS = 'B1SS'
     BIAS_BHSS = 'BHSS'
@@ -188,17 +190,16 @@ class FusedAttnRunner:
         if self.qkv_layout == QKVLayout.BS3HD and self.max_seqlen_q != self.max_seqlen_kv:
             pytest.skip("BS3HD layout requires max_seqlen_q and max_seqlen_kv to be equal.")
 
-        self.backend = FusedAttnHelper(
-            self.dtype, self.dtype, self.qkv_layout.value, self.attn_bias_type.value,
-            self.attn_mask_type.value, self.dropout_prob, self.num_heads_q, self.num_heads_kv,
-            self.max_seqlen_q, self.max_seqlen_kv, self.head_dim).get_fused_attn_backend()
+        self.backend = FusedAttnHelper(self.dtype, self.dtype, self.qkv_layout.value,
+                                       self.attn_bias_type.value, self.attn_mask_type.value,
+                                       self.dropout_prob, self.num_heads_q, self.num_heads_kv,
+                                       self.max_seqlen_q, self.max_seqlen_kv,
+                                       self.head_dim).get_fused_attn_backend()
         if self.backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend:
             pytest.skip("Unsupported inputs combination or device compute capability.")
 
-        if self.bias_shape != BiasShape.BIAS_1HSS:
-            if self.attn_bias_type != AttnBiasType.POST_SCALE_BIAS:
-                pytest.skip("B1SS, BHSS and 11SS bias shapes require POST_SCALE_BIAS.")
-            elif self.attn_mask_type not in [AttnMaskType.NO_MASK, AttnMaskType.CAUSAL_MASK]:
+        if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS:
+            if self.attn_mask_type not in [AttnMaskType.NO_MASK, AttnMaskType.CAUSAL_MASK]:
                 pytest.skip("B1SS, BHSS and 11SS bias shapes are only supported for "
                             "AttnMaskType.NO_MASK and AttnMaskType.CAUSAL_MASK.")
             elif self.backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen:
@@ -213,7 +214,9 @@ class FusedAttnRunner:
         q_shape = (self.batch_size, self.max_seqlen_q, self.num_heads_q, self.head_dim)
         k_shape = v_shape = (self.batch_size, self.max_seqlen_kv, self.num_heads_kv, self.head_dim)
 
-        if self.bias_shape == BiasShape.BIAS_1HSS:
+        if self.attn_bias_type == AttnBiasType.NO_BIAS:
+            bias_shape = None
+        elif self.bias_shape == BiasShape.BIAS_1HSS:
             bias_shape = (1, self.num_heads_q, self.max_seqlen_q, self.max_seqlen_kv)
         elif self.bias_shape == BiasShape.BIAS_B1SS:
             bias_shape = (self.batch_size, 1, self.max_seqlen_q, self.max_seqlen_kv)
@@ -222,7 +225,7 @@ class FusedAttnRunner:
         elif self.bias_shape == BiasShape.BIAS_11SS:
             bias_shape = (1, 1, self.max_seqlen_q, self.max_seqlen_kv)
         else:
-            pytest.xfail("PyTest attempted to use an unrecognized bias layout!")
+            pytest.fail(f"PyTest attempted to use an unrecognized bias_layout = {self.bias_shape}!")
 
         self.q = jax.random.uniform(q_key, q_shape, self.dtype, -1.)
         self.k = jax.random.uniform(k_key, k_shape, self.dtype, -1.)
@@ -237,7 +240,7 @@ class FusedAttnRunner:
                 cudnn_neg_inf = -2.**27. if self.dtype == jnp.bfloat16 else -2.**15.
                 self.bias = jnp.full(bias_shape, cudnn_neg_inf, dtype=self.dtype)
                 max_id = min(self.max_seqlen_q, self.max_seqlen_kv)
-                seq_id_size = max_id * 5 // 128  # 5 ids per interval of 128 sequences
+                seq_id_size = max_id * 5 // 128    # 5 ids per interval of 128 sequences
                 seq_id = jax.random.randint(bias_key, (int(seq_id_size),), 0, max_id).tolist()
                 for i in range(1, len(seq_id)):
                     self.bias = \
@@ -327,8 +330,8 @@ class FusedAttnRunner:
                                                        **kwargs), arg_nums))
         jitted_reference = jit(
             value_and_grad(
-                lambda q, k, v, bias, *args: grad_func(jax_dpa, q, k, v, bias, *args,
-                                                       **kwargs), arg_nums))
+                lambda q, k, v, bias, *args: grad_func(jax_dpa, q, k, v, bias, *args, **kwargs),
+                arg_nums))
 
         primitive_out, primitive_dgrad = jitted_primitive(*args)
         reference_out, reference_dgrad = jitted_reference(*args)
@@ -361,10 +364,10 @@ class FusedAttnRunner:
             primitive_dbias = jnp.float32(primitive_dgrad[3])
             reference_dbias = jnp.float32(reference_dgrad[3])
 
-            assert_allclose(
-                primitive_dbias[..., self.valid_len_q:, self.valid_len_kv:],
-                jnp.zeros_like(primitive_dbias[..., self.valid_len_q:, self.valid_len_kv:]),
-                dtype=self.dtype)
+            assert_allclose(primitive_dbias[..., self.valid_len_q:, self.valid_len_kv:],
+                            jnp.zeros_like(primitive_dbias[..., self.valid_len_q:,
+                                                           self.valid_len_kv:]),
+                            dtype=self.dtype)
 
             # dbias padded part
             assert_allclose(primitive_dbias[..., self.valid_len_q:, self.valid_len_kv:],
@@ -376,15 +379,13 @@ class FusedAttnRunner:
                             reference_dbias[..., :self.valid_len_q, :self.valid_len_kv],
                             dtype=self.dtype)
 
-@pytest.mark.parametrize('bias_shape', [
-    pytest.param(BiasShape.BIAS_1HSS, id='1-H-S-S'),
-    pytest.param(BiasShape.BIAS_B1SS, id='B-1-S-S'),
-    pytest.param(BiasShape.BIAS_BHSS, id='B-H-S-S'),
-    pytest.param(BiasShape.BIAS_11SS, id='1-1-S-S'),
-])
-@pytest.mark.parametrize('attn_bias_type', [
-    pytest.param(AttnBiasType.NO_BIAS, id='NO_BIAS'),
-    pytest.param(AttnBiasType.POST_SCALE_BIAS, id='POST_SCALE_BIAS'),
+
+@pytest.mark.parametrize('attn_bias_type, bias_shape', [
+    pytest.param(AttnBiasType.NO_BIAS, None, id='NO_BIAS'),
+    pytest.param(AttnBiasType.POST_SCALE_BIAS, BiasShape.BIAS_1HSS, id='POST_SCALE_BIAS-1HSS'),
+    pytest.param(AttnBiasType.POST_SCALE_BIAS, BiasShape.BIAS_B1SS, id='POST_SCALE_BIAS-B1SS'),
+    pytest.param(AttnBiasType.POST_SCALE_BIAS, BiasShape.BIAS_BHSS, id='POST_SCALE_BIAS-BHSS'),
+    pytest.param(AttnBiasType.POST_SCALE_BIAS, BiasShape.BIAS_11SS, id='POST_SCALE_BIAS-11SS'),
 ])
 @pytest.mark.parametrize('attn_mask_type', [
     pytest.param(AttnMaskType.NO_MASK, id='NO_MASK'),
@@ -399,31 +400,32 @@ class FusedAttnRunner:
 ])
 @pytest.mark.parametrize('dtype', [
     pytest.param(jnp.bfloat16, id="BF16"),
-    pytest.param(jnp.float16, id="FP16")
+    pytest.param(jnp.float16, id="FP16"),
 ])
-@pytest.mark.parametrize('b, s_q, s_kv, h_q, h_kv, d',[
-    pytest.param(32,  128,  128, 16, 16, 64, id='32-128-128-16-16-64-SELF'),
-    pytest.param( 4, 2048, 2048, 12, 12, 64, id='4-2048-2048-12-12-64-SELF'),
-    pytest.param(32,  512,  128, 16, 16, 64, id='32-512-128-16-16-64-CROSS'),
-    pytest.param( 4, 2048, 1024, 12, 12, 64, id='4-2048-1048-12-12-64-CROSS'),
-    pytest.param(32,  128,  128, 16,  8, 64, id='32-128-128-16-8-64-GQA'),
-    pytest.param( 4, 2048, 2048, 12,  6, 64, id='4-2048-2048-12-6-64-GQA')
+@pytest.mark.parametrize('b, s_q, s_kv, h_q, h_kv, d', [
+    pytest.param(32, 128, 128, 16, 16, 64, id='32-128-128-16-16-64-SELF'),
+    pytest.param(4, 2048, 2048, 12, 12, 64, id='4-2048-2048-12-12-64-SELF'),
+    pytest.param(32, 512, 128, 16, 16, 64, id='32-512-128-16-16-64-CROSS'),
+    pytest.param(4, 2048, 1024, 12, 12, 64, id='4-2048-1048-12-12-64-CROSS'),
+    pytest.param(32, 128, 128, 16, 8, 64, id='32-128-128-16-8-64-GQA'),
+    pytest.param(4, 2048, 2048, 12, 6, 64, id='4-2048-2048-12-6-64-GQA'),
 ])
 @pytest.mark.parametrize('dropout_prob', [
     pytest.param(0.0, id="DROP_0.0"),
-    pytest.param(0.1, id="DROP_0.1")
-])
-@pytest.mark.parametrize('is_training', [
-    pytest.param(True, id='TRAINING'),
-    pytest.param(False, id='INFERENCE'),
+    pytest.param(0.1, id="DROP_0.1"),
 ])
 class TestFusedAttn:
     """
     Fused attention tester
     """
+
     @staticmethod
-    def test_forward(b, s_q, s_kv, h_q, h_kv, d, attn_bias_type, attn_mask_type,
-                     dropout_prob, dtype, is_training, qkv_layout, bias_shape):
+    @pytest.mark.parametrize('is_training', [
+        pytest.param(True, id='TRAINING'),
+        pytest.param(False, id='INFERENCE'),
+    ])
+    def test_forward(b, s_q, s_kv, h_q, h_kv, d, attn_bias_type, attn_mask_type, dropout_prob,
+                     dtype, is_training, qkv_layout, bias_shape):
         """
         Test forward with parameterized configs
         """
@@ -432,13 +434,11 @@ class TestFusedAttn:
         runner.test_forward()
 
     @staticmethod
-    def test_backward(b, s_q, s_kv, h_q, h_kv, d, attn_bias_type, attn_mask_type,
-                      dropout_prob, dtype, is_training, qkv_layout, bias_shape):
+    def test_backward(b, s_q, s_kv, h_q, h_kv, d, attn_bias_type, attn_mask_type, dropout_prob,
+                      dtype, qkv_layout, bias_shape):
         """
         Test backward with parameterized configs
         """
-        if not is_training:
-            pytest.skip("Backward pass does not support inference.")
         runner = FusedAttnRunner(b, s_q, s_kv, h_q, h_kv, d, attn_bias_type, attn_mask_type,
                                  dropout_prob, dtype, True, qkv_layout, bias_shape)
         runner.test_backward()

tests/jax/test_layer.py

@@ -449,6 +449,7 @@ class TestDecoderLayer:
                             hidden_dropout_dims=(sequence_dim,),
                             intermediate_dropout_dims=(sequence_dim,),
                             layer_type=TransformerLayerType.DECODER,
+                            self_attn_mask_type='padding_causal',
                             dtype=dtype,
                             **te_layer_attrs)
         ref_layer, ref_params, ref_others = generate_layer(ref_layer_cls, init_rng, inputs,
@@ -497,6 +498,7 @@ class TestDecoderLayer:
                             hidden_dropout_dims=(sequence_dim,),
                             intermediate_dropout_dims=(sequence_dim,),
                             layer_type=TransformerLayerType.DECODER,
+                            self_attn_mask_type='padding_causal',
                             dtype=dtype,
                             **te_layer_attrs)
         ref_layer, ref_params, ref_others = generate_layer(ref_layer_cls, init_rng, inputs,

tests/jax/test_praxis_layers.py

@@ -730,8 +730,13 @@ class TestDotProductAttn(TestLayer):
     def input_getter(self, shape, dtype):
         key = jax.random.PRNGKey(seed=1234)
         q_key, k_key, v_key = jax.random.split(key, 3)
-        return list(map(partial(jax.random.normal, shape=shape, dtype=dtype),
-                        [q_key, k_key, v_key]))
+        b, s, *_ = shape
+        if self.attrs[DotProductAttnAttr.TRANSPOSE_BS]:
+            b, s = s, b
+        mask = jnp.zeros((b, 1, s, s), dtype=jnp.uint8)
+        return [
+            *map(partial(jax.random.normal, shape=shape, dtype=dtype), [q_key, k_key, v_key]), mask
+        ]
 
     def get_layer_name(self):
         return 'dot_product_attn'
@@ -765,6 +770,7 @@ class TestDotProductAttn(TestLayer):
     @pytest.mark.parametrize('dtype', DTYPE)
     @pytest.mark.parametrize('attrs', DotProductAttnAttr.ATTRS)
     def test_forward_backward(self, data_shape, dtype, attrs, rtol=1e-05, atol=1e-08):
+        self.attrs = attrs
         praxis_p, flax_cls = self.generate_praxis_p_and_flax_cls(dtype, attrs)
         self.forward_backward_runner(data_shape, dtype, praxis_p, flax_cls, rtol, atol)
 
@@ -853,9 +859,11 @@ class MultiHeadAttnAttr:
 class TestMultiHeadAttn(TestLayer):
 
     def input_getter(self, shape, dtype):
-        data_key = jax.random.PRNGKey(seed=1234)
-        return (jax.random.normal(data_key, shape,
-                                  dtype), jax.random.normal(data_key, shape, dtype))
+        key = jax.random.PRNGKey(seed=1234)
+        q_key, kv_key = jax.random.split(key, 2)
+        s, b, *_ = shape
+        mask = jnp.zeros((b, 1, s, s), dtype=jnp.uint8)
+        return [*map(partial(jax.random.normal, shape=shape, dtype=dtype), [q_key, kv_key]), mask]
 
     def get_layer_name(self):
         return 'multi_head_attn'
@@ -1183,9 +1191,15 @@ class TransformerLayerAttr:
 class TestTransformer(TestLayer):
 
     def input_getter(self, shape, dtype):
-        data_key = jax.random.PRNGKey(seed=1234)
-        return (jax.random.normal(data_key, shape,
-                                  dtype), jax.random.normal(data_key, shape, dtype))
+        key = jax.random.PRNGKey(seed=1234)
+        q_key, kv_key = jax.random.split(key, 2)
+        b, s, *_ = shape
+        if self.attrs[TransformerLayerAttr.TRANSPOSE_BS]:
+            b, s = s, b
+        mask = jnp.zeros((b, 1, s, s), dtype=jnp.uint8)
+        return [
+            *map(partial(jax.random.normal, shape=shape, dtype=dtype), [q_key, kv_key]), mask, mask
+        ]
 
     def get_layer_name(self):
         return 'transformerlayer'
@@ -1277,6 +1291,7 @@ class TestTransformer(TestLayer):
     @pytest.mark.parametrize('dtype', DTYPE)
     @pytest.mark.parametrize('attrs', TransformerLayerAttr.ATTRS)
     def test_forward_backward(self, data_shape, dtype, attrs, rtol=1e-05, atol=1e-08):
+        self.attrs = attrs
         praxis_p, flax_cls = self.generate_praxis_p_and_flax_cls(dtype, attrs)
         self.forward_backward_runner(data_shape, dtype, praxis_p, flax_cls, rtol, atol)
 
@@ -1292,7 +1307,7 @@ class TestTransformer(TestLayer):
                                   fp8_format,
                                   rtol=1e-05,
                                   atol=1e-08):
-
+        self.attrs = attrs
         ds = DelayedScaling(fp8_format=fp8_format)
         with fp8_autocast(enabled=True, fp8_recipe=ds):
             praxis_p, flax_cls = self.generate_praxis_p_and_flax_cls(dtype, attrs)

transformer_engine/jax/csrc/extensions.cpp

@@ -49,10 +49,6 @@ pybind11::dict Registrations() {
         EncapsulateFunction(ScaledUpperTriangMaskedSoftmaxForward);
     dict["te_scaled_upper_triang_masked_softmax_backward"] =
         EncapsulateFunction(ScaledUpperTriangMaskedSoftmaxBackward);
-    dict["te_self_fused_attn_forward"] = EncapsulateFunction(SelfFusedAttnForward);
-    dict["te_self_fused_attn_backward"] = EncapsulateFunction(SelfFusedAttnBackward);
-    dict["te_cross_fused_attn_forward"] = EncapsulateFunction(CrossFusedAttnForward);
-    dict["te_cross_fused_attn_backward"] = EncapsulateFunction(CrossFusedAttnBackward);
     dict["te_fused_attn_forward"] = EncapsulateFunction(FusedAttnForward);
     dict["te_fused_attn_backward"] = EncapsulateFunction(FusedAttnBackward);
     return dict;
@@ -72,10 +68,6 @@ PYBIND11_MODULE(transformer_engine_jax, m) {
     m.def("get_dgelu_dbias_ct_workspace_sizes", &GetDGeluDBiasCastTransposeWorkspaceSizes);
     m.def("get_layernorm_fwd_workspace_sizes", &GetLayerNormForwardWorkspaceSizes);
     m.def("get_layernorm_bwd_workspace_sizes", &GetLayerNormBackwardWorkspaceSizes);
-    m.def("get_self_fused_attn_fwd_workspace_sizes", &GetSelfFusedAttnForwardWorkspaceSizes);
-    m.def("get_self_fused_attn_bwd_workspace_sizes", &GetSelfFusedAttnBackwardWorkspaceSizes);
-    m.def("get_cross_fused_attn_fwd_workspace_sizes", &GetCrossFusedAttnForwardWorkspaceSizes);
-    m.def("get_cross_fused_attn_bwd_workspace_sizes", &GetCrossFusedAttnBackwardWorkspaceSizes);
     m.def("get_fused_attn_fwd_workspace_sizes", &GetFusedAttnForwardWorkspaceSizes);
     m.def("get_fused_attn_bwd_workspace_sizes", &GetFusedAttnBackwardWorkspaceSizes);
 

transformer_engine/jax/csrc/modules.h

@@ -118,17 +118,18 @@ struct CustomCallFusedAttnDescriptor {
     float dropout_probability;
     NVTE_Bias_Type bias_type;
     NVTE_Mask_Type mask_type;
+    NVTE_QKV_Layout qkv_layout;
     DType dtype;
     DType wkspace_dtype;
     bool is_training;
 };
 
 pybind11::bytes PackCustomCallFusedAttnDescriptor(
-    size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen,
+    size_t input_batch, size_t batch_size, size_t q_max_seqlen, size_t kv_max_seqlen,
     size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t head_dim,
-    size_t wkspace_size, float scaling_factor, float dropout_probability,
-    NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
-    DType dtype, DType wkspace_dtype, bool is_training);
+    size_t wkspace_size, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type,
+    NVTE_Mask_Type mask_type, NVTE_QKV_Layout qkv_layout, DType dtype, DType wkspace_dtype,
+    bool is_training);
 
 NVTE_Fused_Attn_Backend GetFusedAttnBackend(DType q_dtype, DType kv_dtype,
                                             NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type,
@@ -207,55 +208,19 @@ void ScaledUpperTriangMaskedSoftmaxForward(cudaStream_t stream, void **buffers,
 void ScaledUpperTriangMaskedSoftmaxBackward(cudaStream_t stream, void **buffers, const char *opaque,
                                             std::size_t opaque_len);
 
-pybind11::tuple GetSelfFusedAttnForwardWorkspaceSizes(
-    size_t input_batch, size_t bias_batch, size_t max_seqlen,
-    size_t attn_heads, size_t bias_heads, size_t head_dim,
-    float scaling_factor, float dropout_probability,
-    NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, DType dtype, bool is_training);
-
-void SelfFusedAttnForward(cudaStream_t stream, void **buffers, const char *opaque,
-                          size_t opaque_len);
-
-pybind11::tuple GetSelfFusedAttnBackwardWorkspaceSizes(
-    size_t input_batch, size_t bias_batch, size_t max_seqlen,
-    size_t attn_heads, size_t bias_heads, size_t head_dim,
-    float scaling_factor, float dropout_probability,
-    NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, DType dtype, bool is_training);
-
-void SelfFusedAttnBackward(cudaStream_t stream, void **buffers, const char *opaque,
-                           size_t opaque_len);
-
-pybind11::tuple GetCrossFusedAttnForwardWorkspaceSizes(
-    size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen,
-    size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t head_dim,
-    float scaling_factor, float dropout_probability,
-    NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, DType dtype, bool is_training);
-
-void CrossFusedAttnForward(cudaStream_t stream, void **buffers, const char *opaque,
-                           size_t opaque_len);
-
-pybind11::tuple GetCrossFusedAttnBackwardWorkspaceSizes(
-    size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen,
-    size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t head_dim,
-    float scaling_factor, float dropout_probability,
-    NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, DType dtype, bool is_training);
-
-void CrossFusedAttnBackward(cudaStream_t stream, void **buffers, const char *opaque,
-                            size_t opaque_len);
-
 pybind11::tuple GetFusedAttnForwardWorkspaceSizes(
     size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen,
     size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t head_dim,
-    float scaling_factor, float dropout_probability,
-    NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, DType dtype, bool is_training);
+    float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type,
+    NVTE_Mask_Type mask_type, NVTE_QKV_Layout qkv_layout, DType dtype, bool is_training);
 
 void FusedAttnForward(cudaStream_t stream, void **buffers, const char *opaque, size_t opaque_len);
 
 pybind11::tuple GetFusedAttnBackwardWorkspaceSizes(
     size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen,
     size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t head_dim,
-    float scaling_factor, float dropout_probability,
-    NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, DType dtype, bool is_training);
+    float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type,
+    NVTE_Mask_Type mask_type, NVTE_QKV_Layout qkv_layout, DType dtype, bool is_training);
 
 void FusedAttnBackward(cudaStream_t stream, void **buffers, const char *opaque, size_t opaque_len);
 

transformer_engine/jax/flax/transformer.py

@@ -26,7 +26,7 @@ from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP
 from .module import LayerNorm, Softmax
 from ..fused_attn import AttnBiasType, AttnMaskType, QKVLayout
 from ..fused_attn import is_fused_attn_kernel_available, canonicalize_attn_mask_type
-from ..fused_attn import self_fused_attn, cross_fused_attn, fused_attn
+from ..fused_attn import fused_attn_qkvpacked, fused_attn_kvpacked, fused_attn
 from ..softmax import SoftmaxType
 from ..sharding import num_of_devices
 from ..sharding import get_sharding_map_logic_axis_to_mesh_axis
@@ -190,16 +190,19 @@ class _UnfusedDotProductAttention(nn.Module):    # pylint: disable=too-few-publi
 
         def convert_to_softmax_type(attn_mask_type, mask):
             """Convert the attn_mask_type to SoftmaxType"""
+            # mask is ignored for no_mask and causal_mask
+            if attn_mask_type in [AttnMaskType.NO_MASK, AttnMaskType.CAUSAL_MASK]:
+                mask = None
             if attn_mask_type in [AttnMaskType.CAUSAL_MASK, AttnMaskType.PADDING_CAUSAL_MASK]:
-                return SoftmaxType.SCALED_UPPER_TRIANG_MASKED
+                return SoftmaxType.SCALED_UPPER_TRIANG_MASKED, mask
             if attn_mask_type in [AttnMaskType.NO_MASK, AttnMaskType.PADDING_MASK]:
                 if mask is not None:
-                    return SoftmaxType.SCALED_MASKED
-                return SoftmaxType.SCALED
-            raise ValueError(f"Unsupported {attn_mask_type=}, "
-                             "supported attn_mask_type = {'causal', 'padding'}")
+                    return SoftmaxType.SCALED_MASKED, mask
+                return SoftmaxType.SCALED, mask
+            raise ValueError(f"Unsupported {attn_mask_type=}, supported attn_mask_type="
+                             "{'no_mask', 'padding', 'causal', 'padding_causal', 'causal_padding'}")
 
-        softmax_type = convert_to_softmax_type(self.attn_mask_type, mask)
+        softmax_type, mask = convert_to_softmax_type(self.attn_mask_type, mask)
 
         attn_weights = Softmax(softmax_type=softmax_type,
                                scale_factor=fused_scale_factor)(attn_weights, mask,
@@ -266,15 +269,15 @@ class _FusedDotProductAttention(nn.Module):    # pylint: disable=too-few-public-
             qkv_packed = query
             if self.transpose_batch_sequence:
                 qkv_packed = qkv_packed.transpose([1, 0, 2, 3, 4])
-            x = self_fused_attn(qkv_packed,
-                                bias,
-                                mask,
-                                seed,
-                                attn_mask_type=self.attn_mask_type,
-                                attn_bias_type=self.attn_bias_type,
-                                scaling_factor=scale_factor,
-                                dropout_probability=self.attention_dropout,
-                                is_training=not deterministic)
+            x = fused_attn_qkvpacked(qkv_packed,
+                                     bias,
+                                     mask,
+                                     seed,
+                                     attn_mask_type=self.attn_mask_type,
+                                     attn_bias_type=self.attn_bias_type,
+                                     scaling_factor=scale_factor,
+                                     dropout_probability=self.attention_dropout,
+                                     is_training=not deterministic)
         elif self.qkv_layout == QKVLayout.BSHD_BS2HD:
             """kvpacked format, treat
             query: query tensor, shape = [..., h, d]
@@ -285,16 +288,16 @@ class _FusedDotProductAttention(nn.Module):    # pylint: disable=too-few-public-
             if self.transpose_batch_sequence:
                 query = query.transpose([1, 0, 2, 3])
                 kv_packed = kv_packed.transpose([1, 0, 2, 3, 4])
-            x = cross_fused_attn(query,
-                                 kv_packed,
-                                 bias,
-                                 mask,
-                                 seed,
-                                 attn_mask_type=self.attn_mask_type,
-                                 attn_bias_type=self.attn_bias_type,
-                                 scaling_factor=scale_factor,
-                                 dropout_probability=self.attention_dropout,
-                                 is_training=not deterministic)
+            x = fused_attn_kvpacked(query,
+                                    kv_packed,
+                                    bias,
+                                    mask,
+                                    seed,
+                                    attn_mask_type=self.attn_mask_type,
+                                    attn_bias_type=self.attn_bias_type,
+                                    scaling_factor=scale_factor,
+                                    dropout_probability=self.attention_dropout,
+                                    is_training=not deterministic)
         elif self.qkv_layout == QKVLayout.BSHD_BSHD_BSHD:
             if self.transpose_batch_sequence:
                 query = query.transpose([1, 0, 2, 3])
@@ -358,11 +361,27 @@ class DotProductAttention(nn.Module):    # pylint: disable=too-few-public-method
     attention_dropout: float, default = 0.0
         Dropout probability for the dropout op after the softmax.
     attn_mask_type: str, default = 'causal'
-        Type of the attention mask passed into softmax operation in the self attention.
-        Available options: {'no_mask', 'padding', 'causal', 'causal_padding'}
-        Introduced in v0.10.0.
+        This parameter specifies the type of attention mask to be applied during the softmax
+        operation.
+        Available options are {'no_mask', 'padding', 'causal', 'causal_padding', 'padding_causal'}
+
+        Each described below:
+
+        * no_mask: No attention mask is applied. This means the attention will consider the
+          full sequence without any restrictions.
+        * padding: Indicates the presence of padding at the end of each sequence.
+          Users must provide a mask with the shape [batch, 1, max_seqlen_q, max_seqlen_kv] in the
+          :attr:`__call__` method to specify the padding positions.
+        * causal: An upper triangular mask is applied to the softmax inputs,
+          ensuring that the prediction for a certain position is only dependent on known outputs
+          from positions before it.
+        * causal_padding / padding_causal: A combination of both causal and padding masks.
+          Both 'causal_padding' and 'padding_causal' are acceptable and have the same effect.
+
+        .. note:: :attr:`mask` in :attr:`__call__` is ignored for 'no_mask' and 'causal'.
+
     attn_bias_type: Optional[str], default = None
-        Type of the attention bias passed in the self attention.
+        Type of the attention bias passed in the attention.
         Available options: {'no_bias', 'pre_scale_bias', 'post_scale_bias'}.
         When default is present, the type is automatically decided by the MHA's bias parameter.
         Where it is :attr:`post_scale_bias` if there is bias. Otherwise :attr:`no_bias` is used.
@@ -438,6 +457,7 @@ class DotProductAttention(nn.Module):    # pylint: disable=too-few-public-method
         mask: jax.numpy.ndarray, default = None
             Boolean tensor used to mask out the attention softmax input.
             :attr:`True` means to mask out the corresponding values.
+            Ignored when :attr:`self.attn_mask_type` is either 'no_mask' or 'causal'.
         bias: jax.numpy.ndarray, default = None
             A tensor used to shift attention softmax input.
         *:
@@ -639,9 +659,25 @@ class MultiHeadAttention(nn.Module):    # pylint: disable=too-few-public-methods
     attention_dropout: float, default = 0.0
         Dropout probability for the dropout op after the softmax.
     attn_mask_type: str, default = 'causal'
-        Type of the attention mask passed into softmax operation in the attention.
-        Available options: {'no_mask', 'padding', 'causal', 'causal_padding'}
-        Introduced in v0.10.0.
+        This parameter specifies the type of attention mask to be applied during the softmax
+        operation.
+        Available options are {'no_mask', 'padding', 'causal', 'causal_padding', 'padding_causal'}
+
+        Each described below:
+
+        * no_mask: No attention mask is applied. This means the attention will consider the
+          full sequence without any restrictions.
+        * padding: Indicates the presence of padding at the end of each sequence.
+          Users must provide a mask with the shape [batch, 1, max_seqlen_q, max_seqlen_kv] in the
+          :attr:`__call__` method to specify the padding positions.
+        * causal: An upper triangular mask is applied to the softmax inputs,
+          ensuring that the prediction for a certain position is only dependent on known outputs
+          from positions before it.
+        * causal_padding / padding_causal: A combination of both causal and padding masks.
+          Both 'causal_padding' and 'padding_causal' are acceptable and have the same effect.
+
+        .. note:: :attr:`mask` in :attr:`__call__` is ignored for 'no_mask' and 'causal'.
+
     attn_bias_type: Optional[str], default = None
         Type of the attention bias passed in the attention.
         Available options: {'no_bias', 'pre_scale_bias', 'post_scale_bias'}.
@@ -809,6 +845,7 @@ class MultiHeadAttention(nn.Module):    # pylint: disable=too-few-public-methods
         mask: jax.numpy.ndarray, default = None
             Boolean tensor used to mask out the attention softmax input.
             :attr:`True` means mask out the corresponding values.
+            Ignored when :attr:`self.attn_mask_type` is either 'no_mask' or 'causal'.
         bias: jax.numpy.ndarray, default = None
             A tensor used to shift the attention softmax input.
         *
@@ -1299,9 +1336,25 @@ class TransformerLayer(nn.Module):    # pylint: disable=too-few-public-methods
         is added after self-attention.this can be used for structures like `T5`
         Transformer in conjunction with the TransformerLayerType.ENCODER option.
     self_attn_mask_type: str, default = 'causal'
-        Type of the attention mask passed into softmax operation in the self attention.
-        Available options: {'no_mask', 'padding', 'causal', 'causal_padding'}
-        Introduced in v0.10.0.
+        This parameter specifies the type of attention mask to be applied during the softmax
+        operation in the self attention.
+        Available options are {'no_mask', 'padding', 'causal', 'causal_padding', 'padding_causal'}
+
+        Each described below:
+
+        * no_mask: No attention mask is applied. This means the self attention will consider the
+          full sequence without any restrictions.
+        * padding: Indicates the presence of padding at the end of each sequence.
+          Users must provide a mask with the shape [batch, 1, max_seqlen_q, max_seqlen_kv] in the
+          :attr:`__call__` method to specify the padding positions.
+        * causal: An upper triangular mask is applied to the softmax inputs,
+          ensuring that the prediction for a certain position is only dependent on known outputs
+          from positions before it.
+        * causal_padding / padding_causal: A combination of both causal and padding masks.
+          Both 'causal_padding' and 'padding_causal' are acceptable and have the same effect.
+
+        .. note:: :attr:`attention_mask` in :attr:`__call__` is ignored for 'no_mask' and 'causal'.
+
     self_attn_bias_type: Optional[str], default = None
         Type of the attention bias passed into the self attention.
         Available options: {'no_bias', 'pre_scale_bias', 'post_scale_bias'}.
@@ -1420,9 +1473,12 @@ class TransformerLayer(nn.Module):    # pylint: disable=too-few-public-methods
             :attr:`layer_type=TransformerLayerType.DECODER`.
         attention_mask : jax.numpy.ndarray, default = None
             Boolean tensor used to mask out self-attention softmax input.
+            :attr:`True` means mask out the corresponding values.
+            Ignored when :attr:`self.self_attn_mask_type` is either 'no_mask' or 'causal'.
         encoder_decoder_mask: jax.numpy.ndarray, default = None
             Boolean tensor used to mask out cross-attention softmax input when
             :attr:`layer_type=TransformerLayerType.DECODER`.
+            :attr:`True` means mask out the corresponding values.
         deterministic: bool, default = False
             Disable dropout layers if set to True.
         decode: bool, default = False