Check that a kernel can be launched before adding it to candidate ker…

…nels

PiperOrigin-RevId: 591324734

jax_triton/triton_lib.py

@@ -496,15 +496,16 @@ def prune_configs(configs, named_args):
     if num_ctas != 1:
       raise ValueError("`num_ctas != 1` is not yet supported.")
 
-    kernel_calls.append(
-        triton_kernel_call_lib.TritonKernelCall(
-            kernel,
-            params["grid"][0],
-            params["grid"][1],
-            params["grid"][2],
-            kernel_params,
-        )
-    )
+    if kernel.can_launch():
+      kernel_calls.append(
+          triton_kernel_call_lib.TritonKernelCall(
+              kernel,
+              params["grid"][0],
+              params["grid"][1],
+              params["grid"][2],
+              kernel_params,
+          )
+      )
 
   if len(kernel_calls) > 1:
     named_scalar_args = {fn.arg_names[i]: v for i, _, v in scalar_args}