Skip to content

Changelog

jax-triton 0.4.0 (unreleased)

  • New features
  • Added support for Gluon kernels.
  • triton_call now accepts pytrees as inputs and outputs. Note that input_output_aliases and zeroed_outputs now index into the flattened inputs and outputs, matching jax.experimental.pallas.pallas_call.
  • triton_call now accepts arbitrary backend options via the backend_options argument.
  • Added a has_side_effect argument to triton_call to prevent dead-code elimination (DCE) at the XLA level.
  • Missing metaparams are now filled in from the kernel parameter defaults.
  • triton_call now accepts kernel operands (arrays and scalars) as keyword arguments, matched by kernel parameter name. Keywords naming constexpr parameters, or names that are not kernel parameters, are still treated as metaparams.
  • jax_triton is now importable on hosts without GPU support; the GPU support check is deferred until lowering time.
  • Exposed the custom call name CUSTOM_CALL_TARGET_NAME, which can be used with disabled_checks in jax.export.export without hand-coding this name.
  • triton_call now accepts Ref arguments for in-place buffer updates.
  • Added support for Triton and Gluon kernels that build TMA descriptors on Hopper and Blackwell GPUs. This feature requires jaxlib >0.11.0.

  • The out_shape argument to triton_call has been renamed to out_type to match the modern Pallas APIs.

  • Deprecations

  • The out_shape argument to triton_call is deprecated in favor of out_type.
  • The input_output_aliases and zeroed_outputs arguments to triton_call are deprecated; use Ref arguments for in-place buffer updates instead.

  • Passing np.float32 scalars to triton_call is deprecated. Use plain Python float instead, which now maps to fp32 by default.

  • Breaking changes

  • Python float scalar arguments are now conditionally mapped to "fp32" if representable in the float32 range (matching upstream Triton), falling back to "fp64" otherwise.
  • Removed the jax_triton.utils submodule together with jax_triton.cdiv and jax_triton.next_power_of_2. Use triton.cdiv and triton.next_power_of_2 instead.
  • Removed the enable_fp_fusion argument to triton_call. Pass backend_options=dict(enable_fp_fusion=...) instead.
  • When using input_output_aliases, aliased buffers must no longer be declared as implicit outputs; the aliased input buffer is used directly. A clearer error is now raised for the old calling convention.

  • triton_call keyword arguments that name a non-constexpr kernel parameter are now treated as operands (subject to scalar-static/array separation), not metaparams.

  • Bug fixes

  • Fixed input_output_aliases indexing when positional scalar arguments precede array inputs.
  • Temporary files are no longer leaked when compiling on ROCm.
  • The compute capability is now resolved from gpu_info when a device is unavailable.
  • jax_triton now unsets TRITON_CACHE_DIR on import for compatibility with Triton >= 3.4.0.