Integrate Triton up to [ba7386d](https://github.com/openai/triton/commits/ba7386d942252fcbb336c20cfe5c92d5ee64774f)

https://github.com/openxla/triton/tree/triton_integrate_branch-1.15

PiperOrigin-RevId: 834207274

Author: olegshyshkov

jax_triton/triton_lib.py

@@ -180,7 +180,6 @@ class CompilationResult:
   binary: str
   name: str
   shared_mem_bytes: int
-  cluster_dims: tuple
   ttgir: str | None
   llir: str | None
 
@@ -261,14 +260,12 @@ def compile_ttir_to_ptx_inplace(
   if cuda_options.debug:
     print(ptx)
   name = metadata["name"]
-  cluster_dims = metadata["cluster_dims"]
   ttgir = str(ttgir) if _JAX_TRITON_DUMP_DIR else None
   llir = str(llir) if _JAX_TRITON_DUMP_DIR else None
   return CompilationResult(
       binary=ptx,
       name=name,
       shared_mem_bytes=shared_mem_bytes,
-      cluster_dims=cluster_dims,
       ttgir=ttgir,
       llir=llir,
   )
@@ -306,9 +303,6 @@ def compile_ttir_to_hsaco_inplace(
   name = metadata["name"]
   ttgir = str(ttgir) if _JAX_TRITON_DUMP_DIR else None
   llir = str(llir) if _JAX_TRITON_DUMP_DIR else None
-  # cluster dims are NOT useful on hip backend.
-  # We just fill up with some value for API compatibility
-  cluster_dims = (0, 0, 0)
   # Instead of passing hsaco which are "bytes", we first write
   # to a file and then pass the "string" path. This is needed because
   # nanobind doesn't automatically convert between bytes and string.
@@ -320,7 +314,6 @@ def compile_ttir_to_hsaco_inplace(
       binary=hsaco_path,
       name=name,
       shared_mem_bytes=shared_mem_bytes,
-      cluster_dims=cluster_dims,
       ttgir=ttgir,
       llir=llir,
   )
@@ -469,18 +462,17 @@ def get_or_create_triton_kernel(
       ) as f:
         f.write(
             f"{kernel_name}: shared_mem_bytes:"
-            f" {compilation_result.shared_mem_bytes}, cluster_dims:"
-            f" {compilation_result.cluster_dims}\n"
+            f" {compilation_result.shared_mem_bytes}\n"
         )
 
     kernel = triton_kernel_call_lib.TritonKernel(
         kernel_name,
         num_warps,
+        num_ctas,
         compilation_result.shared_mem_bytes,
         compilation_result.binary,
         ttir,
         compute_capability,
-        *compilation_result.cluster_dims,
     )
 
     _COMPILED_KERNEL_CACHE[cache_key] = kernel

tests/cluster_test.py

@@ -40,29 +40,6 @@ def dummy_kernel(x_ptr, o_ptr):
 
 class ClusterTest(parameterized.TestCase):
 
-  @parameterized.parameters(1, 2, 4, 8)
-  def test_cluster(self, num_ctas):
-    if 'h100' not in jax.devices()[0].device_kind.lower():
-      self.skipTest('Clusters available only on H100s.')
-
-    cluster_dims = []
-    original_compile_ttir_to_ptx_fn = jt.triton_lib.compile_ttir_to_ptx_inplace
-
-    def my_compile_ttir_to_ptx(*args, **kwargs):
-      nonlocal cluster_dims, original_compile_ttir_to_ptx_fn
-      ret_args = original_compile_ttir_to_ptx_fn(*args, **kwargs)
-      cluster_dims = ret_args.cluster_dims
-      return ret_args
-
-    my_triton_call = functools.partial(jt.triton_call, num_ctas=num_ctas)
-
-    with mock.patch.object(jt, 'triton_call', my_triton_call):
-      with mock.patch.object(
-          jt.triton_lib, 'compile_ttir_to_ptx_inplace', my_compile_ttir_to_ptx
-      ):
-        _dummy_fn(jnp.empty((16,)))
-        self.assertEqual(math.prod(cluster_dims), num_ctas)
-
   def test_invalid_cluster_size(self):
     if 'h100' not in jax.devices()[0].device_kind.lower():
       self.skipTest('Clusters available on H100s.')