Integrate Triton up to [@4caedc2](https://github.com/openai/triton/co…

…mmits/4caedc2a51f5c81f7b63212b9bcb76b430102913) PiperOrigin-RevId: 633565503
jax-ml · May 14, 2024 · 893b538 · 893b538
1 parent 4616f13
commit 893b538
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax_triton/triton_lib.py b/jax_triton/triton_lib.py
@@ -296,7 +296,7 @@ def get_or_create_triton_kernel(
   kernel = _COMPILED_KERNEL_CACHE.get(cache_key)
 
   if kernel is None:
-    target = ("cuda", compute_capability)
+    target = cb.GPUTarget('cuda', compute_capability, 32)
     cuda_backend = cb.CUDABackend(target)
     cuda_options = cuda_backend.parse_options(
         dict(