Fix for lower versions

Author: ae-foster

folx/experimental/pallas/attention/custom_gradients.py

@@ -6,15 +6,11 @@
 import jax.numpy as jnp
 from jax.experimental import pallas as pl
 
-try:
-    from jax.experimental.pallas import triton as plgpu
-except ImportError:
-    from jax.experimental.pallas import gpu as plgpu
-
 from .mhsa import mhsa_kernel, reference_mhsa_kernel
 from .mhsea import mhsea_kernel, reference_mhsea_kernel
 from .utils import (
     big_number,
+    compiler_params,
     compute_q_and_kv_block_len,
     create_grid,
     get_lse_block_spec,
@@ -58,9 +54,7 @@ def mhsa_forward(
             out_shape=jax.ShapeDtypeStruct(
                 shape=(batch_len, seq_len, num_heads, head_len), dtype=q.dtype
             ),
-            compiler_params=plgpu.TritonCompilerParams(
-                num_warps=num_warps, num_stages=num_stages
-            ),
+            compiler_params=compiler_params(num_warps=num_warps, num_stages=num_stages),
             debug=False,
             interpret=interpret,
             name='mhsa_forward',
@@ -118,9 +112,7 @@ def mhsa_backward(
                     shape=(batch_len, seq_len, num_heads, head_len), dtype=q.dtype
                 ),
             ],
-            compiler_params=plgpu.TritonCompilerParams(
-                num_warps=num_warps, num_stages=num_stages
-            ),
+            compiler_params=compiler_params(num_warps=num_warps, num_stages=num_stages),
             debug=False,
             interpret=interpret,
             name='mhsa_backward',
@@ -273,9 +265,7 @@ def mhsea_forward(
                     shape=(batch_len, seq_len, num_heads), dtype=v.dtype
                 ),
             ],
-            compiler_params=plgpu.TritonCompilerParams(
-                num_warps=num_warps, num_stages=num_stages
-            ),
+            compiler_params=compiler_params(num_warps=num_warps, num_stages=num_stages),
             debug=False,
             interpret=interpret,
             name='mhea_forward',
@@ -377,9 +367,7 @@ def mhsea_backward(
                     shape=(batch_len, seq_len, num_heads, seq_len), dtype=e.dtype
                 ),
             ],
-            compiler_params=plgpu.TritonCompilerParams(
-                num_warps=num_warps, num_stages=num_stages
-            ),
+            compiler_params=compiler_params(num_warps=num_warps, num_stages=num_stages),
             debug=False,
             interpret=interpret,
             name='mhsea_backward_q_vjp',
@@ -438,9 +426,7 @@ def mhsea_backward(
                     shape=(batch_len, seq_len, num_heads, head_len), dtype=v.dtype
                 ),
             ],
-            compiler_params=plgpu.TritonCompilerParams(
-                num_warps=num_warps, num_stages=num_stages
-            ),
+            compiler_params=compiler_params(num_warps=num_warps, num_stages=num_stages),
             debug=False,
             interpret=interpret,
             name='mhsea_backward_kv_vjp',

folx/experimental/pallas/attention/forward_laplacian.py

@@ -6,18 +6,14 @@
 import jax.numpy as jnp
 from jax.experimental import pallas as pl
 
-try:
-    from jax.experimental.pallas import triton as plgpu
-except ImportError:
-    from jax.experimental.pallas import gpu as plgpu
-
 from folx import forward_laplacian
 from folx.api import FwdJacobian, FwdLaplArray
 
 from .mhsa import reference_mhsa_kernel
 from .mhsea import reference_mhsea_kernel
 from .utils import (
     big_number,
+    compiler_params,
     compute_q_and_kv_block_len,
     create_grid,
     get_input_mask_block_spec,
@@ -158,9 +154,7 @@ def mhsa_forward_laplacian(
                     dtype=q.dtype,  # o.laplacian
                 ),
             ],
-            compiler_params=plgpu.TritonCompilerParams(
-                num_warps=num_warps, num_stages=num_stages
-            ),
+            compiler_params=compiler_params(num_warps=num_warps, num_stages=num_stages),
             debug=False,
             interpret=interpret,
             name='mhsa_forward_laplacian',
@@ -593,9 +587,7 @@ def mhsea_forward_laplacian(
                     dtype=v.dtype,  # o.laplacian
                 ),
             ],
-            compiler_params=plgpu.TritonCompilerParams(
-                num_warps=num_warps, num_stages=num_stages
-            ),
+            compiler_params=compiler_params(num_warps=num_warps, num_stages=num_stages),
             debug=False,
             interpret=interpret,
             name='mhsea_forward_laplacian',

folx/experimental/pallas/attention/mhsa.py

@@ -6,13 +6,9 @@
 import jax.numpy as jnp
 from jax.experimental import pallas as pl
 
-try:
-    from jax.experimental.pallas import triton as plgpu
-except ImportError:
-    from jax.experimental.pallas import gpu as plgpu
-
 from .utils import (
     big_number,
+    compiler_params,
     compute_q_and_kv_block_len,
     create_grid,
     get_mask_block_spec,
@@ -63,9 +59,7 @@ def mhsa(
             out_shape=jax.ShapeDtypeStruct(
                 shape=(batch_len, seq_len, num_heads, head_len), dtype=q.dtype
             ),
-            compiler_params=plgpu.TritonCompilerParams(
-                num_warps=num_warps, num_stages=num_stages
-            ),
+            compiler_params=compiler_params(num_warps=num_warps, num_stages=num_stages),
             debug=False,
             interpret=interpret,
             name='mhsa',

folx/experimental/pallas/attention/mhsea.py

@@ -6,13 +6,9 @@
 import jax.numpy as jnp
 from jax.experimental import pallas as pl
 
-try:
-    from jax.experimental.pallas import triton as plgpu
-except ImportError:
-    from jax.experimental.pallas import gpu as plgpu
-
 from .utils import (
     big_number,
+    compiler_params,
     compute_q_and_kv_block_len,
     create_grid,
     get_lse_block_spec,
@@ -63,9 +59,7 @@ def mhsea(
                     shape=(batch_len, seq_len, num_heads), dtype=q.dtype
                 ),  # lse
             ],
-            compiler_params=plgpu.TritonCompilerParams(
-                num_warps=num_warps, num_stages=num_stages
-            ),
+            compiler_params=compiler_params(num_warps=num_warps, num_stages=num_stages),
             debug=False,
             interpret=interpret,
             name='mhsea',

folx/experimental/pallas/attention/utils.py

@@ -4,6 +4,14 @@
 import jax.numpy as jnp
 from jax.experimental import pallas as pl
 
+try:
+    from jax.experimental.pallas import triton as plgpu
+except ImportError:
+    from jax.experimental.pallas import gpu as plgpu
+
+
+from packaging.version import Version
+
 
 def sum_columns(x: jax.Array) -> jax.Array:
     return x.sum(axis=1, keepdims=True)
@@ -210,3 +218,10 @@ def big_number(dtype) -> float:
         return 1e40
     else:
         raise ValueError(f'Unexpected dtype {dtype}')
+
+
+def compiler_params(num_warps, num_stages):
+    if Version(jax.__version__) >= Version('0.4.34'):
+        return plgpu.TritonCompilerParams(num_warps=num_warps, num_stages=num_stages)
+    else:
+        return dict(triton=dict(num_warps=num_warps, num_stages=num_stages))