[triton] Pass cluster_dims to TritonKernel and use cuLaunchKernel if …

…size <= 1

PiperOrigin-RevId: 599809560

jax_triton/triton_lib.py

@@ -135,9 +135,6 @@ def get_triton_type(obj: Any) -> str:
   )
 
 
-Grid = Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]]
-GridOrLambda = Union[Grid, Callable[[Dict[str, Any]], Grid]]
-
 triton_kernel_call_p = jax.core.Primitive("triton_kernel_call")
 triton_kernel_call_p.multiple_results = True
 triton_kernel_call_p.def_impl(
@@ -240,7 +237,8 @@ def compile_ttir_to_ptx_inplace(
   if cuda_options.debug:
     print(ptx)
   name = ptx_get_kernel_name(ptx)
-  return ptx, name, shared_mem_bytes, compute_capability
+  cluster_dims = metadata["cluster_dims"]
+  return ptx, name, shared_mem_bytes, compute_capability, cluster_dims
 
 
 _COMPILED_KERNEL_CACHE = {}  # TODO(cjfj): Convert to LRU cache?
@@ -312,7 +310,6 @@ def get_or_create_triton_kernel(
             num_warps=num_warps,
             num_stages=num_stages,
             num_ctas=num_ctas,
-            cluster_dims=(1, 1, 1),
             enable_warp_specialization=enable_warp_specialization,
             enable_persistent=enable_persistent,
             optimize_epilogue=False,
@@ -338,7 +335,7 @@ def get_or_create_triton_kernel(
     )
 
     ttir = str(module)  # `module`` is compiled in-place, so copy TTIR here.
-    ptx, kernel_name, shared_mem_bytes, compute_capability = (
+    ptx, kernel_name, shared_mem_bytes, compute_capability, cluster_dims = (
         compile_ttir_to_ptx_inplace(
             module,
             context,
@@ -350,7 +347,13 @@ def get_or_create_triton_kernel(
     )
 
     kernel = triton_kernel_call_lib.TritonKernel(
-        kernel_name, num_warps, shared_mem_bytes, ptx, ttir, compute_capability
+        kernel_name,
+        num_warps,
+        shared_mem_bytes,
+        ptx,
+        ttir,
+        compute_capability,
+        *cluster_dims,
     )
 
     _COMPILED_KERNEL_CACHE[cache_key] = kernel
@@ -580,7 +583,7 @@ def triton_call(
     custom_call_target_name: str = "triton_kernel_call",
     num_warps: Optional[int] = None,
     num_stages: Optional[int] = None,
-    num_ctas: int = 1,
+    num_ctas: int = 1,  # TODO(giorgioa): Add support for dimensions tuple.
     enable_fp_fusion: bool = True,
     enable_warp_specialization: bool = False,
     enable_persistent: bool = False,
@@ -663,6 +666,8 @@ def add(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
       indices, for outputs that should be zeroed before the kernel is launched.
     num_warps: The number of warps used to execute the Triton kernel.
     num_stages: The number of stages emitted by the Triton compiler.
+    num_ctas: The size of thread blocks per cluster to be used on GPUs with
+      compute capabilities >= 9.0. It must be less or equal to 8.
     debug: Prints out intermediate IRs if True for debugging purposes.
     serialized_metadata: Arbitrary metadata that will be added into the
       serialized kernel call.