Clean up checks in unit test.

Signed-off-by: Michael Goldfarb <[email protected]>

tests/jax/test_distributed_fused_attn.py

@@ -423,7 +423,8 @@ def test_contex_parallel_self_attn(
         num_kv_heads = num_head // kv_groups
         scaling_factor = 1.0 / np.sqrt(num_head)
 
-        if not is_fused_attn_kernel_available(
+        def check_has_backend_for_mask(mask_type):
+            return is_fused_attn_kernel_available(
             dtype,
             dtype,
             qkv_layout,
@@ -435,9 +436,15 @@ def test_contex_parallel_self_attn(
             seqlen,
             seqlen,
             hidden,
-            None,  # no window
-        ):
-            pytest.skip(f"No FusedAttn backend found")
+            None) # no SWA for CP
+
+        # For causal masking we depend on having bottom right support also.
+        has_backend = check_has_backend_for_mask(attn_mask_type)
+        if mask == AttnMaskType.CAUSAL_MASK_MASK:
+            has_backend &= check_has_backend_for_mask(AttnMaskType.CAUSAL_BOTTOM_RIGHT_MASK)
+
+        if not has_backend
+            pytest.skip(f"No FusedAttn backend found {cp_size=} {attn_mask_type=}.")
 
         if dp_size > 1 and batch % dp_size != 0:
             pytest.skip(f"Skipping {batch=} not a multiple of {dp_size=}")