Try by only commenting out SW part and keeping causal part

Signed-off-by: Kshitij Lakhani <[email protected]>

Author: KshitijLakhani

transformer_engine/jax/attention.py

@@ -460,13 +460,13 @@ def _segment_ids_pos_to_seqlens_offsets(
         lambda x, y: jnp.equal(x, y) * x,
     )
     attn_mask = segment_mask
-    # if attn_mask_type.is_causal():
-    #     causal_mask = make_attention_mask(
-    #         segment_pos_q,
-    #         segment_pos_kv,
-    #         jnp.greater_equal,
-    #     )
-    #     attn_mask = jnp.logical_and(segment_mask, causal_mask)
+    if attn_mask_type.is_causal():
+        causal_mask = make_attention_mask(
+            segment_pos_q,
+            segment_pos_kv,
+            jnp.greater_equal,
+        )
+        attn_mask = jnp.logical_and(segment_mask, causal_mask)
 
     # swa_mask = make_swa_mask(segment_pos_q, segment_pos_kv, window_size, dtype=jnp.bool)
     # attn_mask = jnp.logical_and(attn_mask, swa_mask)