Add option of disabling cudnn version checking.

PiperOrigin-RevId: 698889971

jax/BUILD

@@ -305,6 +305,7 @@ py_library_providing_imports_info(
         ":xla_bridge",
         ":xla_metadata",
         "//jax/_src/lib",
+        "//third_party/py/absl/flags",
     ] + py_deps("numpy") + py_deps("scipy") + py_deps("opt_einsum") + py_deps("flatbuffers") + jax_extra_deps,
 )
 

jax/_src/cudnn/fused_attention_stablehlo.py

@@ -17,6 +17,7 @@
 import json
 import math
 
+from absl import flags
 import jax
 from jax import core
 from jax import dtypes
@@ -34,6 +35,15 @@
 
 Array = jnp.ndarray
 
+_DISABLE_CUDNN_VERSION_CHECK = flags.DEFINE_bool(
+    "disable_cudnn_version_check",
+    False,
+    "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 +1035,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 +1059,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