[JAX] Fix issues when mask/sequence_descriptor is None (#1477)

Fix issues when mask/sequence_descriptor is None
Signed-off-by: Reese Wang <rewang@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>

tests/jax/test_fused_attn.py

@@ -556,6 +556,9 @@ class FusedAttnRunner:
         else:
             match self.seq_desc_format:
                 case SeqDescFormat.Mask:
+                    if self.attn_mask_type == AttnMaskType.NO_MASK:
+                        self.sequence_desciptor = None
+                    else:
                         self.sequence_desciptor = make_mask(
                             self.segment_ids_q,
                             self.segment_ids_kv,

transformer_engine/jax/attention.py

@@ -950,7 +950,7 @@ def fused_attn(
                              AttnBiasType.NO_BIAS, AttnMaskType.PADDING_CAUSAL_MASK,
                              QKVLayout.T3HD, 0.125, 0, True, 3)
     """
-    if isinstance(sequence_descriptor, jnp.ndarray):
+    if sequence_descriptor is None or isinstance(sequence_descriptor, jnp.ndarray):
         warnings.warn(
             "Pass mask to fused_attn is deprecated, please use SequenceDescriptor instead. "
             + "See help(transformer_engine.jax.attention.SequenceDescriptor) for details.",