encountered "Batching rule for 'triton_kernel_call' not implemented"

[tlogn]

Hello, I am facing an issue with implementing triton flash attention using jax-triton. I am receiving the error message "Batching rule for 'triton_kernel_call' not implemented". Can anyone provide assistance? Thank you!

[dfm]

I don't think jax-triton supports any of the usual JAX transformations (notably vmap in this case, but also grad, for example) out of the box. I expect that you'll need to implement custom transformation rules using jax.batching.custom_vmap and probably also jax.custom_vjp for grad support. Hope this helps!

[tlogn]

I am also curious about how to customize the batching rules for the Triton kernel. It appears that the 'triton_kernel_call' function applies to all Triton kernels. Is there a possibility of conflicts in batching rules when implementing multiple Triton kernels?

[dfm]

You definitely don't want to register an explicit batching rule for triton_kernel_call_p primitive. Instead, you'll wrap your specific kernel using custom_vmap. For example:

@jax.custom_batching.custom_vmap
def my_function(x, y):
  ...
  return jt.triton_call(...)

@my_function.def_vmap
def my_function_vmap(axis_size, batched_in, x, y):
  ...

Then, only your instance of triton_call will have the custom batching behavior.

[tlogn]

Got it, thank you !