Assert in fused attn bwd pass for sm100

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

Add check for sm100

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

Author: KshitijLakhani

transformer_engine/jax/cpp_extensions/attention.py

@@ -17,7 +17,10 @@
 from jax.experimental.custom_partitioning import SdyShardingRule
 
 import transformer_engine_jax
-from transformer_engine_jax import NVTE_Fused_Attn_Backend
+from transformer_engine_jax import (
+    NVTE_Fused_Attn_Backend,
+    get_device_compute_capability,
+)
 from transformer_engine.jax.attention import (
     AttnBiasType,
     AttnMaskType,
@@ -2745,6 +2748,11 @@ def fused_attn_bwd(
         assert bias is None
         bias = jnp.zeros(0, dtype=qkv[0].dtype)
 
+    if get_device_compute_capability==100:
+        assert (not(attn_bias_type != "no_bias" and dropout_probability != 0)
+        ),"For sm100, bprop kernel support for dropout + determinism (bias) is not supported"
+
+
     fused_config = _FusedAttnConfig(
         attn_bias_type=attn_bias_type,
         attn_mask_type=attn_mask_type,