[JAX] Add the missing 1HSS tests (NVIDIA#1052)

Add the missing 1HSS tests Signed-off-by: Reese Wang <[email protected]>
erhoo82 · Aug 6, 2024 · 5bb3a41 · 5bb3a41
1 parent d74e65f
commit 5bb3a41
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
@@ -295,7 +295,10 @@ def _check_configs(self):
         if self.backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend:
             pytest.skip("Unsupported inputs combination or device compute capability.")
 
-        if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS:
+        if (
+            self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS
+            and self.bias_shape != BiasShape.BIAS_1HSS
+        ):
             if self.attn_mask_type not in [AttnMaskType.NO_MASK, AttnMaskType.CAUSAL_MASK]:
                 pytest.skip(
                     "B1SS, BHSS and 11SS bias shapes are only supported for "