Integrate cuDNN frontend v1 to fused attention (#497)

* Integrate cuDNN frontend v1 to fused attention and miscellaneous fixes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix lint
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix jax/paddle for unit tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix jax/pytorch lint
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* simplify stride generation
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix and/or logic in get_backend
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix flag_max512 and test_numerics
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove v.contiguous() since get_qkv_layout covers it
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* skip fp8 tests for sm89
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* further fix jax CI
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix jax CI
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* revert mask type to comma-separated list
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix lint
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix last two commits
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* integrate v1/pre-release-5
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* cleanup prerelease5 integration and fix FA2.1 commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* force dropout to 0 if not training
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix Jax CI
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* testing bias/alibi and padding+causal; add alibi to unfused DPA
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* set flag_arb to false when non determinism is not allowed
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* followup on prev commit; remove redundant python env var setting
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* WIP: minor tweaks for tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* prepare for tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix determinism logic for fused attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add bias to bwd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix gpt_checkpointing/dpa_accuracy problem
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix some seg fault issues
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add failure notes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove use of non-deter var for backend selection
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fix for lint and CI
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix workspace size in bwd and uncomment bias test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix get_alibi and remove check_support
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update tests status
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove workspace_opt from FADescriptor_v1
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* disable arbitrary backend + post scale bias in Jax; waiting on PR 525
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up bhsd order
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* swap bias/rng_state order in aux_ctx_tensor and add bias to aux_ctx_tensor in _qkvpacked/_kvpacked API
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove support for padding_causal + cross for max512
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change alibi bias to float32 for bias_1_4/5 tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* further clean up tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix thd fwd output shape for FlashAttention and add backend info for DPA
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix definition of workspace limit when dbias is present
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* further tweak DP_WORKSPACE_LIMIT definition
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* disallow alibi+no_mask for sdpa flash and update alibi tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update jax/paddle after PR525 and fix DP_WORKSPACE_LIMIT for dbias Jax tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* disable dbias for non-hopper archs
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix layernorm lint
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remode unused arg for lint
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove build dir in setup.py
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change selection logic to prefer fused attn on sm90
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix distributed jax test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix h and s order in header
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update to cudnn fe v1 branch
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove manual setting of workopt path due to dbias after v1 update
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix paddle CI
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add post_scale_bias and alibi to sdpa flash support matrix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix support matrix in header files
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* move headers back to .cu and change seed/offset to int64
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update Megatron commit in L1 test and remove all prints in fused attn test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix L1 Megatron test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix fp8 arg in L1 Megatron script
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* print only when debug flag is on
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove checkpointing loading to avoid loading other tests results
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>

cudnn-frontend @ 9f82dda5

-Subproject commit 12f35fa2be5994c1106367cac2fba21457b064f4
+Subproject commit 9f82dda5c029d15a5f371f0fe003dc0c74a0c987

qa/L1_pytorch_distributed_unittest/test.sh

@@ -8,6 +8,7 @@ set -e
 
 git clone https://github.com/NVIDIA/Megatron-LM.git
 cd Megatron-LM
-git checkout f24fac4ed0dcf0522056521a93445d9a82f501a9 
+git checkout bcce6f54e075e3c3374ea67adefe54f3f2da2b07
+sed -i -e '1504,1505d' megatron/model/transformer.py
 pytest -v -s $TE_PATH/tests/pytorch/distributed/test_convergence.py
 python $TE_PATH/tests/pytorch/distributed/print_logs.py

tests/jax/test_distributed_fused_attn.py

@@ -77,10 +77,11 @@ class TestDistributedSelfAttn:
         is_training = True
         scaling_factor = 1.0
 
-        _, seqlen, _, _, hidden = data_shape
+        _, seqlen, _, num_head, hidden = data_shape
 
         if not is_fused_attn_kernel_available(dtype, dtype, QKVLayout.BS3HD, attn_bias_type,
-                                              attn_mask_type, dropout_prob, seqlen, seqlen, hidden):
+                                              attn_mask_type, dropout_prob, num_head, num_head,
+                                              seqlen, seqlen, hidden):
             pytest.skip(f"No FusedAttn backwend found")
 
         def target_func(qkv, bias, mask):
@@ -182,10 +183,11 @@ class TestDistributedCrossAttn:
         is_training = True
         scaling_factor = 1.0
 
-        _, seqlen, _, hidden = data_shape
+        _, seqlen, num_head, hidden = data_shape
 
         if not is_fused_attn_kernel_available(dtype, dtype, QKVLayout.BSHD_BS2HD, attn_bias_type,
-                                              attn_mask_type, dropout_prob, seqlen, seqlen, hidden):
+                                              attn_mask_type, dropout_prob, num_head, num_head,
+                                              seqlen, seqlen, hidden):
             pytest.skip(f"No FusedAttn backwend found")
 
         def target_func(q, kv, mask):

tests/jax/test_fused_attn.py

@@ -180,12 +180,14 @@ class TestSelfFusedAttn():
 
     @staticmethod
     def _check_inputs(s, *, attn_bias_type, attn_mask_type, backend, dropout_probability, dtype,
-                      head_dim):
+                      num_heads_q, num_heads_kv, head_dim):
 
         assert isinstance(backend, Backend)
 
         if not is_fused_attn_kernel_available(dtype, dtype, QKVLayout.BS3HD, attn_bias_type,
-                                              attn_mask_type, dropout_probability, s, s, head_dim):
+                                              attn_mask_type, dropout_probability,
+                                              num_heads_q, num_heads_kv,
+                                              s, s, head_dim):
             pytest.skip("Unsupported inputs combination or device compute capability.")
 
     def _set_inputs(self, b, s, h, d, *, attn_bias_type, attn_mask_type, backend,
@@ -197,6 +199,8 @@ class TestSelfFusedAttn():
                                      backend=backend,
                                      dropout_probability=dropout_probability,
                                      dtype=dtype,
+                                     num_heads_q=h,
+                                     num_heads_kv=h,
                                      head_dim=d)
 
         if attn_mask_type in [AttnMaskType.NO_MASK, AttnMaskType.CAUSAL_MASK]:

tests/paddle/parallel_tests/group_sharding.py

@@ -48,11 +48,8 @@ class TestGroupSharding(unittest.TestCase):
     def _get_model_and_optimizer(self, model, stage):
         if stage == 1:
             optimizer = DygraphShardingOptimizer(
-                hcg=fleet.get_hybrid_communicate_group(),
-                user_defined_strategy=self.strategy,
-                params=model.parameters(),
-                inner_optimizer_class=paddle.optimizer.AdamW,
-                learning_rate=0.01,
+                paddle.optimizer.AdamW(learning_rate=0.01, parameters=model.parameters()),
+                fleet.get_hybrid_communicate_group(),
             )
             model = fleet.distributed_model(model)
             optimizer = fleet.distributed_optimizer(optimizer)

tests/paddle/test_layers.py

@@ -634,9 +634,11 @@ def test_dot_product_attention(bs, hidden_size, num_heads, q_seqlen, kv_seqlen,
 
     # Skip if cuDNN fused attention is not supported
     if not is_fused_attention_supported(
-        head_size=head_size,
+        num_heads=num_heads,
+        num_gqa_groups=num_heads,
         q_seqlen=q_seqlen,
         kv_seqlen=kv_seqlen,
+        head_size=head_size,
         dtype=math_dtype,
         dropout=0.0,
         qkv_layout="bs3hd" if attn_type == "self" else "bshd_bs2hd",
@@ -762,9 +764,11 @@ def test_transformer_encoder_layer(bs, hidden_size, num_heads, ffn_hidden_size,
 
     # Skip if cuDNN fused attention is not supported
     if not is_fused_attention_supported(
-        head_size=hidden_size // num_heads,
+        num_heads=num_heads,
+        num_gqa_groups=num_heads,
         q_seqlen=q_seqlen,
         kv_seqlen=kv_seqlen,
+        head_size=hidden_size // num_heads,
         dtype=math_dtype,
         dropout=0.0,
         qkv_layout="bs3hd",
@@ -940,9 +944,11 @@ def test_transformer_decoder_layer(bs, hidden_size, num_heads, ffn_hidden_size,
 
     # Skip if cuDNN fused attention is not supported
     if not is_fused_attention_supported(
-        head_size=hidden_size // num_heads,
+        num_heads=num_heads,
+        num_gqa_groups=num_heads,
         q_seqlen=q_seqlen,
         kv_seqlen=kv_seqlen,
+        head_size=hidden_size // num_heads,
         dtype=math_dtype,
         dropout=0.0,
         qkv_layout="bs3hd",
@@ -952,6 +958,8 @@ def test_transformer_decoder_layer(bs, hidden_size, num_heads, ffn_hidden_size,
         pytest.skip("cuDNN fused attention is not supported")
     if not is_fused_attention_supported(
         head_size=hidden_size // num_heads,
+        num_heads=num_heads,
+        num_gqa_groups=num_heads,
         q_seqlen=q_seqlen,
         kv_seqlen=kv_seqlen,
         dtype=math_dtype,

tests/paddle/test_operators.py

@@ -688,9 +688,11 @@ class TestFusedAttn:
             else "bshd_bs2hd"
         )
         fused_attention_backend = get_fused_attention_backend(
-            head_size=self.head_size,
+            num_heads=self.num_heads,
+            num_gqa_groups=self.num_heads,
             q_seqlen=self.q_seqlen,
             kv_seqlen=self.kv_seqlen,
+            head_size=self.head_size,
             dtype=self.dtype,
             dropout=self.dropout_prob,
             qkv_layout=qkv_layout,
@@ -774,9 +776,11 @@ class TestFusedAttn:
         test self attention forward + backward
         """
         if not is_fused_attention_supported(
-            head_size=d,
+            num_heads=h,
+            num_gqa_groups=h,
             q_seqlen=s,
             kv_seqlen=s,
+            head_size=d,
             dtype=dtype,
             dropout=0.0,
             qkv_layout="bs3hd",
@@ -799,9 +803,11 @@ class TestFusedAttn:
         test cross attention forward + backward
         """
         if not is_fused_attention_supported(
-            head_size=d,
+            num_heads=h,
+            num_gqa_groups=h,
             q_seqlen=s_q,
             kv_seqlen=s_kv,
+            head_size=d,
             dtype=dtype,
             dropout=0.0,
             qkv_layout="bshd_bs2hd",
@@ -825,9 +831,11 @@ class TestFusedAttn:
         test flash attention forward + backward
         """
         if not is_fused_attention_supported(
-            head_size=d,
+            num_heads=h,
+            num_gqa_groups=h,
             q_seqlen=s,
             kv_seqlen=s,
+            head_size=d,
             dtype=dtype,
             dropout=0.0,
             qkv_layout="bs3hd",

tests/paddle/utils.py

@@ -102,9 +102,11 @@ def set_random_seed(seed):
 
 
 def get_fused_attention_backend(
-    head_size: int,
+    num_heads: int,
+    num_gqa_groups: int,
     q_seqlen: int,
     kv_seqlen: int,
+    head_size: int,
     dtype: Union[paddle.dtype, str],
     dropout: float,
     qkv_layout: str = "bs3hd",
@@ -125,6 +127,8 @@ def get_fused_attention_backend(
         AttnBiasType[bias_type],
         AttnMaskType[mask_type],
         dropout,
+        num_heads,
+        num_gqa_groups,
         q_seqlen,
         kv_seqlen,
         head_size,
@@ -132,9 +136,11 @@ def get_fused_attention_backend(
 
 
 def is_fused_attention_supported(
-    head_size: int,
+    num_heads: int,
+    num_gqa_groups: int,
     q_seqlen: int,
     kv_seqlen: int,
+    head_size: int,
     dtype: Union[paddle.dtype, str],
     dropout: float,
     qkv_layout: str = "bs3hd",
@@ -143,9 +149,11 @@ def is_fused_attention_supported(
 ) -> bool:
     """Check if cuDNN fused attention is supported for attention config"""
     backend = get_fused_attention_backend(
-        head_size=head_size,
+        num_heads=num_heads,
+        num_gqa_groups=num_gqa_groups,
         q_seqlen=q_seqlen,
         kv_seqlen=kv_seqlen,
+        head_size=head_size,
         dtype=dtype,
         dropout=dropout,
         qkv_layout=qkv_layout,

tests/pytorch/distributed/run_megatron_lm_gpt.sh

@@ -81,7 +81,6 @@ options=" \
     --merge-file /data/gpt3/pile-cc1-cc2-shuf/bpe/gpt2-merges.txt \
     --save-interval ${SAVE_INTERVAL} \
     --save ${CHECKPOINT_DIR} \
-    --load ${CHECKPOINT_DIR} \
     --split ${SPLIT} \
     --clip-grad ${CLIP_GRAD} \
     --weight-decay ${WEIGHT_DECAY} \
@@ -90,8 +89,6 @@ options=" \
     --init-method-std ${INIT_METHOD_STD} \
     --log-params-norm \
     --log-num-zeros-in-grad \
-    --no-query-key-layer-scaling \
-    --DDP-impl local \
     --transformer-impl ${TRANSFORMER_IMPL} \
     --tensorboard-dir ${TENSORBOARD_DIR} \
     --fp8-margin 0 \
@@ -108,7 +105,7 @@ if [[ "$WGRAD_FUSION" == "False" ]]; then
 fi
 
 if [[ "$FP8" != "False" ]]; then
-        options+=" --fp8-${FP8}"
+        options+=" --fp8-format ${FP8}"
 fi
 
 if [[ "$DTYPE" != "fp32" ]]; then

tests/pytorch/test_numerics.py

@@ -508,6 +508,7 @@ def _test_e2e_checkpointing_get_model(config, dtype):
     sigma = 0.023
     init_method = init_method_normal(sigma)
     output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
     return (
         TransformerLayer(
             config.hidden_size,
@@ -524,7 +525,6 @@ def _test_e2e_checkpointing_get_model(config, dtype):
             params_dtype=dtype,
         )
         .cuda()
-        .eval()
     )
 
 
@@ -559,9 +559,14 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
             if p.requires_grad:
                 param_grads.append(p.grad.clone())
 
+        global _cpu_rng_state, _cuda_rng_state
+        _cpu_rng_state = torch.get_rng_state()
+        _cuda_rng_state = torch.cuda.get_rng_state()
+
         del block
         block = _test_e2e_checkpointing_get_model(config, dtype)
         block.load_state_dict(torch.load(path))
+        reset_rng_states()
 
         for p in block.parameters():
             if p.requires_grad:
@@ -815,21 +820,19 @@ def test_dpa_accuracy(dtype, bs, model):
         DotProductAttention(
             config.num_attention_heads,
             config.embed,
-            attention_dropout=0.1,  # dropout
+            attention_dropout=0.0, # disable dropout, FU uses rng differently
         )
         .to(dtype=dtype)
         .cuda()
-        .eval()
     )
 
     torch_dpa = (
         TorchDotProductAttention(
             config.embed,
-            0.1,  # dropout
+            0.0, # dropout
         )
         .to(dtype=dtype)
         .cuda()
-        .eval()
     )
 
     te_outputs = _test_dpa_accuracy(te_dpa, bs, dtype, config)

transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h

@@ -12,14 +12,13 @@
 #define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_ARBITRARY_SEQLEN_H_
 
 #include "transformer_engine/fused_attn.h"
-
 #include <cudnn.h>
-
 #include "common/common.h"
 
 namespace transformer_engine {
 #if (CUDNN_VERSION >= 8900)
-void fused_attn_arbitrary_seqlen_fwd_qkvpacked(size_t batch, size_t max_seqlen, size_t num_head,
+void fused_attn_arbitrary_seqlen_fwd_qkvpacked(
+                size_t batch, size_t num_attn_heads, size_t max_seqlen,
                 size_t head_size, bool is_training, float attn_scale,
                 float p_dropout, NVTE_QKV_Layout qkv_layout,
                 NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
@@ -28,20 +27,47 @@ void fused_attn_arbitrary_seqlen_fwd_qkvpacked(size_t batch, size_t max_seqlen,
                 const Tensor *cu_seqlens, const Tensor *rng_state,
                 Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle);
 
-void fused_attn_arbitrary_seqlen_bwd_qkvpacked(size_t batch, size_t max_seqlen, size_t num_head,
+void fused_attn_arbitrary_seqlen_bwd_qkvpacked(
+                size_t batch, size_t num_attn_heads, size_t max_seqlen,
                 size_t head_dim, float attn_scale, float p_dropout,
                 NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type,
                 NVTE_Mask_Type mask_type, const Tensor *input_QKV,
-                                      const Tensor *input_O,
-                                      const Tensor *input_dO, Tensor *output_S,
+                const Tensor *input_O, const Tensor *input_dO,
+                const Tensor *input_Bias, Tensor *output_S,
                 Tensor *output_dQKV, Tensor *output_dBias,
                 const Tensor *cu_seqlens, const Tensor *rng_state,
                 Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle);
 
-void fused_attn_arbitrary_seqlen_fwd(size_t batch, size_t max_seqlen_q, size_t max_seqlen_kv,
-                                      size_t num_head, size_t head_size, bool is_training,
+void fused_attn_arbitrary_seqlen_fwd_kvpacked(
+                size_t batch, size_t num_attn_heads, size_t num_gqa_groups,
+                size_t max_seqlen_q, size_t max_seqlen_kv, size_t head_dim,
+                bool is_training, float attn_scale, float p_dropout,
+                NVTE_QKV_Layout qkv_layout,
+                NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
+                const Tensor *input_Q, const Tensor *input_KV, const Tensor *input_Bias,
+                Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors,
+                const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv,
+                const Tensor *rng_state,
+                Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle);
+
+void fused_attn_arbitrary_seqlen_bwd_kvpacked(
+                size_t batch, size_t num_attn_heads, size_t num_gqa_groups,
+                size_t max_seqlen_q, size_t max_seqlen_kv, size_t head_dim,
                 float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout,
                 NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
+                const Tensor *input_Q, const Tensor *input_KV, const Tensor *input_O,
+                const Tensor *input_dO, const Tensor *input_Bias, Tensor *output_S,
+                Tensor *output_dQ, Tensor *output_dKV, Tensor *output_dBias,
+                const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv,
+                const Tensor *rng_state,
+                Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle);
+
+void fused_attn_arbitrary_seqlen_fwd(
+                size_t batch, size_t num_attn_heads, size_t num_gqa_groups,
+                size_t max_seqlen_q, size_t max_seqlen_kv, size_t head_dim,
+                bool is_training, float attn_scale, float p_dropout,
+                NVTE_QKV_Layout qkv_layout,
+                NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
                 const Tensor *input_Q, const Tensor *input_K,
                 const Tensor *input_V, const Tensor *input_Bias,
                 Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors,
@@ -49,13 +75,14 @@ void fused_attn_arbitrary_seqlen_fwd(size_t batch, size_t max_seqlen_q, size_t m
                 const Tensor *rng_state,
                 Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle);
 
-void fused_attn_arbitrary_seqlen_bwd(size_t batch, size_t max_seqlen_q, size_t max_seqlen_kv,
-                                      size_t num_head, size_t head_dim, float attn_scale,
-                                      float p_dropout, NVTE_QKV_Layout qkv_layout,
+void fused_attn_arbitrary_seqlen_bwd(
+                size_t batch, size_t num_attn_heads, size_t num_gqa_groups,
+                size_t max_seqlen_q, size_t max_seqlen_kv, size_t head_dim,
+                float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout,
                 NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
                 const Tensor *input_Q, const Tensor *input_K,
                 const Tensor *input_V, const Tensor *input_O,
-                                      const Tensor *input_dO, Tensor *output_S,
+                const Tensor *input_dO, const Tensor *input_Bias, Tensor *output_S,
                 Tensor *output_dQ, Tensor *output_dK,
                 Tensor *output_dV, Tensor *output_dBias,
                 const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv,

transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu

@@ -217,14 +217,14 @@ static cudnn_frontend::Tensor createMask(int64_t b, int64_t h, int64_t s_q, int6
     int64_t maskOutputTensor_virtual = true;
     cudnnDataType_t maskOutputTensor_dataType = CUDNN_DATA_FLOAT;
     auto maskOutputTensor_reorderType =
-        cudnn_frontend::cudnnBackendTensorReordering_t::CUDNN_TENSOR_REORDERING_NONE;
+        cudnn_frontend::TensorReordering_t::NONE;
 
     if (is_bprop) {
         maskOutputTensor_id = dS_ID;
         maskOutputTensor_virtual = false;
         maskOutputTensor_dataType = tensorType;
         maskOutputTensor_reorderType =
-            cudnn_frontend::cudnnBackendTensorReordering_t::CUDNN_TENSOR_REORDERING_F16x16;
+            cudnn_frontend::TensorReordering_t::F16x16;
     }
 
     auto maskOutputTensor =
@@ -357,7 +357,7 @@ static cudnn_frontend::Tensor createSoftmaxForward(
     // divide (e/ sum(e))
 
     auto reorder_type =
-        cudnn_frontend::cudnnBackendTensorReordering_t::CUDNN_TENSOR_REORDERING_F16x16;
+        cudnn_frontend::TensorReordering_t::F16x16;
 
     auto afterDivisionTensor =
         cudnn_frontend::TensorBuilder()
@@ -448,7 +448,7 @@ static cudnn_frontend::Tensor createDropout(int64_t b, int64_t h, int64_t s_q, i
                                            afterBMM1_stride, true, false);  // is virtual
 
     auto reorder_type =
-        cudnn_frontend::cudnnBackendTensorReordering_t::CUDNN_TENSOR_REORDERING_F16x16;
+        cudnn_frontend::TensorReordering_t::F16x16;
 
     // after dropout tensor
     auto afterDropoutTensor =
@@ -918,7 +918,7 @@ void fused_attn_max_512_bwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
             auto doTensor = tensor_create(tensorType, dO_ID, o_dim, o_stride, false, false);
 
             auto reorder_type =
-                cudnn_frontend::cudnnBackendTensorReordering_t::CUDNN_TENSOR_REORDERING_F16x16;
+                cudnn_frontend::TensorReordering_t::F16x16;
 
             // activation from fprop
             auto pTensor =
@@ -1246,7 +1246,7 @@ void fused_attn_max_512_bwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
 
 using namespace transformer_engine::fused_attn;
 void fused_attn_max_512_fwd_qkvpacked(
-    size_t batch, size_t max_seqlen, size_t num_head, size_t head_dim, bool is_training,
+    size_t batch, size_t num_head, size_t max_seqlen, size_t head_dim, bool is_training,
     float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type,
     NVTE_Mask_Type mask_type, const Tensor *input_QKV, const Tensor *input_Bias, Tensor *output_O,
     NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens, const Tensor *rng_state,
@@ -1312,8 +1312,8 @@ void fused_attn_max_512_fwd_qkvpacked(
     }
 }
 
-void fused_attn_max_512_fwd_kvpacked(size_t batch, size_t q_max_seqlen, size_t kv_max_seqlen,
-                                     size_t num_head, size_t head_dim, bool is_training,
+void fused_attn_max_512_fwd_kvpacked(size_t batch, size_t num_head, size_t q_max_seqlen,
+                                     size_t kv_max_seqlen, size_t head_dim, bool is_training,
                                      float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout,
                                      NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
                                      const Tensor *input_Q, const Tensor *input_KV,
@@ -1389,8 +1389,8 @@ void fused_attn_max_512_fwd_kvpacked(size_t batch, size_t q_max_seqlen, size_t k
         NVTE_ERROR("Unexpected workspace_size.");
     }
 }
-void fused_attn_max_512_fwd(size_t batch, size_t q_max_seqlen, size_t kv_max_seqlen,
-                                     size_t num_head, size_t head_dim, bool is_training,
+void fused_attn_max_512_fwd(size_t batch, size_t num_head, size_t q_max_seqlen,
+                                     size_t kv_max_seqlen, size_t head_dim, bool is_training,
                                      float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout,
                                      NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
                                      const Tensor *input_Q, const Tensor *input_K,
@@ -1460,7 +1460,7 @@ void fused_attn_max_512_fwd(size_t batch, size_t q_max_seqlen, size_t kv_max_seq
     }
 }
 
-void fused_attn_max_512_bwd_qkvpacked(size_t batch, size_t max_seqlen, size_t num_head,
+void fused_attn_max_512_bwd_qkvpacked(size_t batch, size_t num_head, size_t max_seqlen,
                                       size_t head_dim, float attn_scale, float p_dropout,
                                       NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type,
                                       NVTE_Mask_Type mask_type, const Tensor *input_QKV,
@@ -1519,8 +1519,8 @@ void fused_attn_max_512_bwd_qkvpacked(size_t batch, size_t max_seqlen, size_t nu
     }
 }
 
-void fused_attn_max_512_bwd_kvpacked(size_t batch, size_t q_max_seqlen, size_t kv_max_seqlen,
-                                     size_t num_head, size_t head_dim, float attn_scale,
+void fused_attn_max_512_bwd_kvpacked(size_t batch, size_t num_head, size_t q_max_seqlen,
+                                     size_t kv_max_seqlen, size_t head_dim, float attn_scale,
                                      float p_dropout, NVTE_QKV_Layout qkv_layout,
                                      NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
                                      const Tensor *input_Q, const Tensor *input_KV,
@@ -1580,8 +1580,8 @@ void fused_attn_max_512_bwd_kvpacked(size_t batch, size_t q_max_seqlen, size_t k
         NVTE_ERROR("Unexpected workspace_size.");
     }
 }
-void fused_attn_max_512_bwd(size_t batch, size_t q_max_seqlen, size_t kv_max_seqlen,
-                                     size_t num_head, size_t head_dim, float attn_scale,
+void fused_attn_max_512_bwd(size_t batch, size_t num_head, size_t q_max_seqlen,
+                                     size_t kv_max_seqlen, size_t head_dim, float attn_scale,
                                      float p_dropout, NVTE_QKV_Layout qkv_layout,
                                      NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
                                      const Tensor *input_Q, const Tensor *input_K,

transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h

@@ -19,7 +19,7 @@
 
 namespace transformer_engine {
 #if (CUDNN_VERSION >= 8901)
-void fused_attn_max_512_fwd_qkvpacked(size_t batch, size_t max_seqlen, size_t num_head,
+void fused_attn_max_512_fwd_qkvpacked(size_t batch, size_t num_head, size_t max_seqlen,
                                       size_t head_size, bool is_training, float attn_scale,
                                       float p_dropout, NVTE_QKV_Layout qkv_layout,
                                       NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
@@ -28,8 +28,8 @@ void fused_attn_max_512_fwd_qkvpacked(size_t batch, size_t max_seqlen, size_t nu
                                       const Tensor *cu_seqlens, const Tensor *rng_state,
                                       Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle);
 
-void fused_attn_max_512_fwd_kvpacked(size_t batch, size_t q_max_seqlen, size_t kv_max_seqlen,
-                                     size_t num_head, size_t head_dim, bool is_training,
+void fused_attn_max_512_fwd_kvpacked(size_t batch, size_t num_head, size_t q_max_seqlen,
+                                     size_t kv_max_seqlen, size_t head_dim, bool is_training,
                                      float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout,
                                      NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
                                      const Tensor *input_Q, const Tensor *input_KV,
@@ -38,8 +38,8 @@ void fused_attn_max_512_fwd_kvpacked(size_t batch, size_t q_max_seqlen, size_t k
                                      const Tensor *kv_cu_seqlens, const Tensor *rng_state,
                                      Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle);
 
-void fused_attn_max_512_fwd(size_t batch, size_t q_max_seqlen, size_t kv_max_seqlen,
-                                     size_t num_head, size_t head_dim, bool is_training,
+void fused_attn_max_512_fwd(size_t batch, size_t num_head, size_t q_max_seqlen,
+                                     size_t kv_max_seqlen, size_t head_dim, bool is_training,
                                      float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout,
                                      NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
                                      const Tensor *input_Q, const Tensor *input_K,
@@ -49,7 +49,7 @@ void fused_attn_max_512_fwd(size_t batch, size_t q_max_seqlen, size_t kv_max_seq
                                      const Tensor *kv_cu_seqlens, const Tensor *rng_state,
                                      Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle);
 
-void fused_attn_max_512_bwd_qkvpacked(size_t batch, size_t max_seqlen, size_t num_head,
+void fused_attn_max_512_bwd_qkvpacked(size_t batch, size_t num_head, size_t max_seqlen,
                                       size_t head_dim, float attn_scale, float p_dropout,
                                       NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type,
                                       NVTE_Mask_Type mask_type, const Tensor *input_QKV,
@@ -58,8 +58,8 @@ void fused_attn_max_512_bwd_qkvpacked(size_t batch, size_t max_seqlen, size_t nu
                                       const Tensor *cu_seqlens, Tensor *workspace,
                                       cudaStream_t stream, cudnnHandle_t handle);
 
-void fused_attn_max_512_bwd_kvpacked(size_t batch, size_t q_max_seqlen, size_t kv_max_seqlen,
-                                     size_t num_head, size_t head_dim, float attn_scale,
+void fused_attn_max_512_bwd_kvpacked(size_t batch, size_t num_head, size_t q_max_seqlen,
+                                     size_t kv_max_seqlen, size_t head_dim, float attn_scale,
                                      float p_dropout, NVTE_QKV_Layout qkv_layout,
                                      NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
                                      const Tensor *input_Q, const Tensor *input_KV,
@@ -68,8 +68,8 @@ void fused_attn_max_512_bwd_kvpacked(size_t batch, size_t q_max_seqlen, size_t k
                                      const Tensor *q_cu_seqlens, const Tensor *kv_cu_seqlens,
                                      Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle);
 
-void fused_attn_max_512_bwd(size_t batch, size_t q_max_seqlen, size_t kv_max_seqlen,
-                                     size_t num_head, size_t head_dim, float attn_scale,
+void fused_attn_max_512_bwd(size_t batch, size_t num_head, size_t q_max_seqlen,
+                                     size_t kv_max_seqlen, size_t head_dim, float attn_scale,
                                      float p_dropout, NVTE_QKV_Layout qkv_layout,
                                      NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
                                      const Tensor *input_Q, const Tensor *input_K,

transformer_engine/common/fused_attn/fused_attn_fp8.cu

@@ -366,8 +366,7 @@ static cudnn_frontend::Tensor createDropoutForward(
           .setDataType(CUDNN_DATA_FLOAT)
           .setVirtual(true)
           .setByValue(false)
-          .setReorderType(cudnn_frontend::cudnnBackendTensorReordering_t::
-                          CUDNN_TENSOR_REORDERING_F16x16)
+          .setReorderType(cudnn_frontend::TensorReordering_t::F16x16)
           .build();
   // Scale after dropout
   auto scaleDropoutTensor = tensor_create(
@@ -448,8 +447,7 @@ static cudnn_frontend::Tensor createDropoutBackward(
           .setDataType(CUDNN_DATA_FLOAT)
           .setVirtual(true)
           .setByValue(false)
-          .setReorderType(cudnn_frontend::cudnnBackendTensorReordering_t::
-                          CUDNN_TENSOR_REORDERING_F16x16)
+          .setReorderType(cudnn_frontend::TensorReordering_t::F16x16)
           .build();
   // Scale after dropout (1 / (1 - p))
   auto scaleDropoutTensor = tensor_create(
@@ -992,7 +990,7 @@ static cudnn_frontend::Tensor createdSQBMM(
 }
 
 // fused attention FWD FP8
-void fused_attn_fp8_fwd_impl(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
+void fused_attn_fp8_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d,
             bool isTraining, float attnScale,
             float dropoutProbability, NVTE_QKV_Layout layout,
             void* devPtrQ, void* devPtrK, void* devPtrV,
@@ -1305,7 +1303,7 @@ void fused_attn_fp8_fwd_impl(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, in
 }
 
 // fused attention BWD FP8
-void fused_attn_fp8_bwd_impl(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
+void fused_attn_fp8_bwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d,
             float attnScale, float dropoutProbability, NVTE_QKV_Layout layout,
             void* devPtrQ, void* devPtrK, void* devPtrV,
             void* devPtrM, void* devPtrZInv,
@@ -1935,7 +1933,7 @@ void fused_attn_fp8_fwd_qkvpacked(
   size_t workspace_size = 0;
 
   fused_attn::fused_attn_fp8_fwd_impl(
-                  b, max_seqlen, max_seqlen, h, d,
+                  b, h, max_seqlen, max_seqlen, d,
                   is_training, attn_scale, p_dropout, qkv_layout,
                   devPtrQ, devPtrK, devPtrV,
                   devPtrM, devPtrZInv,
@@ -2025,7 +2023,7 @@ void fused_attn_fp8_bwd_qkvpacked(
   size_t workspace_size = 0;
 
   fused_attn::fused_attn_fp8_bwd_impl(
-                  b, max_seqlen, max_seqlen, h, d,
+                  b, h, max_seqlen, max_seqlen, d,
                   attn_scale, p_dropout, qkv_layout,
                   devPtrQ, devPtrK, devPtrV,
                   devPtrM, devPtrZInv,
@@ -2131,7 +2129,7 @@ void fused_attn_fp8_fwd(
   size_t workspace_size = 0;
 
   fused_attn::fused_attn_fp8_fwd_impl(
-                  b, max_seqlen_q, max_seqlen_kv, h, d,
+                  b, h, max_seqlen_q, max_seqlen_kv, d,
                   is_training, attn_scale, p_dropout, qkv_layout,
                   devPtrQ, devPtrK, devPtrV,
                   devPtrM, devPtrZInv,
@@ -2224,7 +2222,7 @@ void fused_attn_fp8_bwd(
   size_t workspace_size = 0;
 
   fused_attn::fused_attn_fp8_bwd_impl(
-                  b, max_seqlen_q, max_seqlen_kv, h, d,
+                  b, h, max_seqlen_q, max_seqlen_kv, d,
                   attn_scale, p_dropout, qkv_layout,
                   devPtrQ, devPtrK, devPtrV,
                   devPtrM, devPtrZInv,

transformer_engine/common/fused_attn/fused_attn_fp8.h

@@ -14,8 +14,7 @@ namespace transformer_engine {
 #if (CUDNN_VERSION >= 8900)
 // fused attention FWD FP8 with packed QKV
 void fused_attn_fp8_fwd_qkvpacked(
-            size_t b, size_t max_seqlen,
-            size_t h, size_t d,
+            size_t b, size_t h, size_t max_seqlen, size_t d,
             bool is_training, float attn_scale,
             float p_dropout, NVTE_QKV_Layout qkv_layout,
             const Tensor *input_QKV,
@@ -30,8 +29,7 @@ void fused_attn_fp8_fwd_qkvpacked(
 
 // fused attention BWD FP8 with packed QKV
 void fused_attn_fp8_bwd_qkvpacked(
-            size_t b, size_t max_seqlen,
-            size_t h, size_t d,
+            size_t b, size_t h, size_t max_seqlen, size_t d,
             float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout,
             const Tensor *input_QKV,
             const Tensor *input_O,
@@ -49,8 +47,7 @@ void fused_attn_fp8_bwd_qkvpacked(
 
 // fused attention FWD FP8 with separate Q, K, V
 void fused_attn_fp8_fwd(
-            size_t b, size_t max_seqlen_q, size_t max_seqlen_kv,
-            size_t h, size_t d,
+            size_t b, size_t h, size_t max_seqlen_q, size_t max_seqlen_kv, size_t d,
             bool is_training, float attn_scale,
             float p_dropout, NVTE_QKV_Layout qkv_layout,
             const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V,
@@ -66,8 +63,7 @@ void fused_attn_fp8_fwd(
 
 // fused attention BWD FP8 with separate Q, K, V
 void fused_attn_fp8_bwd(
-            size_t b, size_t max_seqlen_q, size_t max_seqlen_kv,
-            size_t h, size_t d,
+            size_t b, size_t h, size_t max_seqlen_q, size_t max_seqlen_kv, size_t d,
             float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout,
             const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V,
             const Tensor *input_O,

transformer_engine/common/fused_attn/utils.h

@@ -12,6 +12,7 @@
 
 #include <cudnn.h>
 #include <cudnn_frontend.h>
+#include <cudnn_frontend_utils.h>
 
 #include <cstdint>
 #include <mutex>
@@ -95,6 +96,34 @@ struct FADescriptor {
   }
 };
 
+struct FADescriptor_v1 {
+  std::int64_t b;
+  std::int64_t h;
+  std::int64_t hg;
+  std::int64_t s_q;
+  std::int64_t s_kv;
+  std::int64_t d;
+  float attnScale;
+  bool isTraining;
+  float dropoutProbability;
+  NVTE_QKV_Layout layout;
+  NVTE_Bias_Type bias_type;
+  NVTE_Mask_Type mask_type;
+  cudnn_frontend::DataType_t tensor_type;
+
+  bool operator<(const FADescriptor_v1 &rhs) const {
+    return std::tie(b, h, hg, s_q, s_kv, d,
+                    attnScale, isTraining, dropoutProbability,
+                    layout, mask_type, bias_type, tensor_type)
+                    < std::tie(
+                      rhs.b, rhs.h, rhs.hg, rhs.s_q, rhs.s_kv, rhs.d,
+                      rhs.attnScale, rhs.isTraining,
+                      rhs.dropoutProbability, rhs.layout,
+                      rhs.mask_type, rhs.bias_type,
+                      rhs.tensor_type);
+  }
+};
+
 __global__ void cu_seqlens_to_offsets(size_t b, size_t h, size_t d,
                 int32_t *cu_seqlens_q, int32_t *actual_seqlens_q,
                 int32_t *qkv_ragged_offset, int32_t *o_ragged_offset);
@@ -107,6 +136,7 @@ __global__ void cu_seqlens_to_actual_seqlens(size_t b,
 }  // namespace fused_attn
 
 cudnnDataType_t get_cudnn_dtype(const transformer_engine::DType t);
+cudnn_frontend::DataType_t get_cudnn_fe_dtype(const transformer_engine::DType t);
 
 class cudnnExecutionPlanManager {
  public: