fix: Unreachable backend check after earlier backend skip

[andrewwhitecdw]

This PR addresses the following issue in tests/jax/test_fused_attn.py: Unreachable backend check after earlier backend skip.

Changes

• tests/jax/test_fused_attn.py: Unreachable backend check after earlier backend skip.

Details

--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -1,13 +1,8 @@
-        if (
-            self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS
-            and self.bias_shape != BiasShape._1HSS
-        ):
-            if self.attn_mask_type.is_padding():
-                pytest.skip(
-                    "B1SS, BHSS and 11SS bias shapes are only supported for non-padding mask"
-                )
-            elif self.backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen:
-                pytest.skip(
-                    "B1SS, BHSS and 11SS bias shapes are only supported for "
-                    "the F16_arbitrary_seqlen backend."
-                )
+        if (
+            self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS
+            and self.bias_shape != BiasShape._1HSS
+        ):
+            if self.attn_mask_type.is_padding():
+                pytest.skip(
+                    "B1SS, BHSS and 11SS bias shapes are only supported for non-padding mask"
+                )

Tests

• tests/jax/test_fused_attn.py

--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -1235,4 +1235,35 @@
     return (left_window_size, 0)
 
+
+def test_post_scale_non_1hss_bias_shapes_do_not_hit_unreachable_backend_skip():
+    """
+    Regression test for the removed unreachable backend check.
+
+    A config that reaches this point with a non-1HSS post-scale bias shape must
+    already be using the F16_arbitrary_seqlen backend, so there should be no
+    additional skip for the backend.
+    """
+    runner = FusedAttnRunner(
+        batch_size=2,
+        max_seqlen_q=2048,
+        max_seqlen_kv=2048,
+        num_heads_q=12,
+        num_heads_kv=12,
+        head_dim_qk=64,
+        head_dim_v=64,
+        attn_bias_type=AttnBiasType.POST_SCALE_BIAS,
+        attn_mask_type=AttnMaskType.NO_MASK,
+        softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+        dropout_prob=0.0,
+        dtype=jnp.bfloat16,
+        is_training=True,
+        qkv_layout=QKVLayout.BSHD_BSHD_BSHD,
+        bias_shape=BiasShape._B1SS,
+        window_size=None,
+        seq_desc_format=SeqDescFormat.Mask,
+    )
+    runner._check_configs()
+    assert runner.backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen
+
+
 @pytest.mark.parametrize(
     "attn_mask_type",

[greptile-apps]

Greptile Summary

The PR removes an unreachable backend-specific skip from the JAX fused-attention configuration checks.

• Retains the padding-mask restriction for non-1HSS post-scale bias shapes.
• Relies on the earlier backend-selection check to skip configurations without a supported fused-attention backend.

Confidence Score: 5/5

The PR appears safe to merge because the removed branch was already made unreachable by the preceding backend check.

Backend selection occurs before the changed block, and unsupported post-scale-bias configurations are skipped there; supported non-padding configurations were unaffected by the removed condition.

Important Files Changed

Filename	Overview
tests/jax/test_fused_attn.py	Removes a redundant skip that cannot be reached after the preceding backend-selection guard; no actionable defect was identified.

_{Reviews (1): Last reviewed commit: "fix: Unreachable backend check after ear..." | Re-trigger Greptile}

[cyanguwa]

This is a no-brainer. Thanks.