Add option of disabling cudnn version checking.

PiperOrigin-RevId: 698889971

jax/_src/cudnn/fused_attention_stablehlo.py

@@ -20,11 +20,12 @@
 import jax
 from jax import core
 from jax import dtypes
+from jax._src import config
 from jax._src import dispatch
+from jax._src import xla_bridge
 from jax._src.custom_partitioning import custom_partitioning
 from jax._src.interpreters import batching
 from jax._src.lib import cuda_versions
-from jax._src import xla_bridge
 from jax.interpreters import mlir
 from jax.interpreters import xla
 from jax.interpreters.mlir import hlo
@@ -34,6 +35,18 @@
 
 Array = jnp.ndarray
 
+_DISABLE_CUDNN_VERSION_CHECK = config.bool_state(
+    name="disable_cudnn_version_check",
+    default=False,
+    help=(
+        "Whether to disable the cuDNN version check. It's safe to disable the"
+        " check if we don't actually need to run the cudnn kernel, e.g., for"
+        " exporting JaxModule to tensorflow. Without disabling the check, it"
+        " requires running the export on a machine with a100, which is"
+        " unnecessary."
+    ),
+)
+
 
 class AttentionLayout(enum.Enum):
   BTNH = 0
@@ -1025,7 +1038,9 @@ def dot_product_attention(query: Array,
   # check if cuDNN is installed
   cudnn_version = check_cudnn_version()
   # only support at least Ampere
-  if not check_compute_capability("8.0"):
+  if not _DISABLE_CUDNN_VERSION_CHECK.value and not check_compute_capability(
+      "8.0"
+  ):
     raise RuntimeError("Require at least Ampere arch to run")
   layout = _normalize_layout(qkv_layout)
   if has_padding(mask_type) and (q_seqlen is None or kv_seqlen is None):
@@ -1047,7 +1062,7 @@ def dot_product_attention(query: Array,
 
   # combine bias and mask
   if bias is None:
-      bias = mask
+    bias = mask
   else:
     if mask is not None:
       # should be broadcast to same shape