Support the latest Triton API changes

This is a backport of the fix proposed by mattjj@ in #333.

I also fixed a few linter warnings to make sure at least our lint CI is green.

Note that the tests will still be red once this PR lands, because Triton
does not have a release incorporating all the recent API changes. I manually
verified that the tests pass with Triton compiled from their main branch.

Author: superbobry

README.md

@@ -76,13 +76,14 @@ ipynb](https://github.com/jax-ml/jax-triton/blob/main/examples/JAX_%2B_Triton_Fl
 $ pip install jax-triton
 ```
 
-You can either use a stable release of `triton` or a nightly release.
-
 Make sure you have a CUDA-compatible `jax` installed. For example you could run:
 ```bash
 $ pip install "jax[cuda12]"
 ```
 
+`jax-triton` currently requires building the latest version of `triton`
+[from source](https://triton-lang.org/main/getting-started/installation.html#from-source).
+
 ## Development
 
 To develop `jax-triton`, you can clone the repo with:

jax_triton/triton_lib.py

@@ -20,7 +20,6 @@
 import copy
 import dataclasses
 import functools
-import inspect
 import os
 import pprint
 import tempfile
@@ -363,12 +362,19 @@ def get_or_create_triton_kernel(
   alignments = [16] * len(arg_dtypes)
   for i, _, value in scalar_args:
     alignments[i] = value
+  specialize_extra = backend.get_arg_specialization
+  if specialize_impl := getattr(triton.runtime.jit, "specialize_impl", None):
+    # TODO(slebedev): Remove this branch once Triton 3.3 is released.
+    specialize_impl = functools.partial(
+        specialize_impl, specialize_extra=specialize_extra
+    )
+  else:
+    specialize_impl = triton.runtime.jit.create_specialize_impl(specialize_extra)
   specialization = [
-      triton.runtime.jit.specialize_impl(
+      specialize_impl(
           types.SimpleNamespace(
               data_ptr=lambda: alignment, dtype=arg_dtype.removeprefix("*")
           ),
-          backend.get_arg_specialization,
       )
       for arg_dtype, alignment in zip(arg_dtypes, alignments)
   ]

tests/triton_call_test.py

@@ -24,7 +24,6 @@
 import jax_triton as jt
 import numpy as np
 import triton
-from triton.compiler import code_generator as code_gen
 import triton.language as tl
 
 config.parse_flags_with_absl()