Proto commits in jax-ml/jax

These 31 commits are when the Protocol Buffers files have changed:

2026-08-12

[triton] Added support for host-side TMA descriptors PiperOrigin-RevId: 963369899

Implement AOT compilation for autotuned JAX-Triton kernels. We now compile all the autotuning candidate configs during the instantiate phase (i.e. compile time in XLA), and later autotune them during instantiate. Some things worth noting: * We always run instantiate, and so we can always expect its state. During JIT compilations this just happens right before initialize. * Instantiate and Execute happen in the same process, so we don't need to serialize/deserialize protos, we can just use the pointers to the objects in the static cache. PiperOrigin-RevId: 963327724

2026-08-11

Added support for host-side TMA descriptors Note that, unfortunately, `TensorDescriptor` in Gluon assumes PyTorch when validating the arguments. So, using it with jax-triton requires patching or removing `TensorDescriptor.__post_init__` as we do in tests. PiperOrigin-RevId: 962805739

2026-08-07

[triton] Added support for device-side TMA descriptors Closes #38912 and #37482. PiperOrigin-RevId: 960937449

2026-08-06

Implement AOT compilation for non-autotuned JAX-Triton kernels. Previously, JAX-Triton kernels depended on a specific PTX version at runtime, which created backwards-compatibility constraints and increased startup latency when loading the binary. During FFI instantiation (compile time), compile non-autotuned kernels down to the CUBIN level so that executable binaries are ready immediately upon load. This removes the runtime PTX dependency, improves backwards compatibility, and reduces load times. PiperOrigin-RevId: 960276682

Add AOT tests for jax-triton kernels in XLA Created a py script that generates the HLO of the kernel, and added a test on the XLA side that compiles, and runs an AOT binary using that kernel. I didn't want to leak out XLA test rules in jaxlib so I kept the jax parts in the jax directory, and the XLA test in the XLA dir. PiperOrigin-RevId: 958348877

Implement AOT compilation for non-autotuned JAX-Triton kernels. Previously, JAX-Triton kernels depended on a specific PTX version at runtime, which created backwards-compatibility constraints and increased startup latency when loading the binary. During FFI instantiation (compile time), compile non-autotuned kernels down to the CUBIN level so that executable binaries are ready immediately upon load. This removes the runtime PTX dependency, improves backwards compatibility, and reduces load times. PiperOrigin-RevId: 958260793

Implement AOT compilation for autotuned JAX-Triton kernels. We now compile all the autotuning candidate configs during the instantiate phase (i.e. compile time in XLA), and later autotune them during instantiate. PiperOrigin-RevId: 959692915

2026-08-05

Add AOT tests for jax-triton kernels in XLA Created a py script that generates the HLO of the kernel, and added a test on the XLA side that compiles, and runs an AOT binary using that kernel. I didn't want to leak out XLA test rules in jaxlib so I kept the jax parts in the jax directory, and the XLA test in the XLA dir. PiperOrigin-RevId: 958348877

Implement AOT compilation for non-autotuned JAX-Triton kernels. Previously, JAX-Triton kernels depended on a specific PTX version at runtime, which created backwards-compatibility constraints and increased startup latency when loading the binary. During FFI instantiation (compile time), compile non-autotuned kernels down to the CUBIN level so that executable binaries are ready immediately upon load. This removes the runtime PTX dependency, improves backwards compatibility, and reduces load times. PiperOrigin-RevId: 958260793

2026-08-04

Implement AOT compilation for non-autotuned JAX-Triton kernels. Previously, JAX-Triton kernels depended on a specific PTX version at runtime, which created backwards-compatibility constraints and increased startup latency when loading the binary. During FFI instantiation (compile time), compile non-autotuned kernels down to the CUBIN level so that executable binaries are ready immediately upon load. This removes the runtime PTX dependency, improves backwards compatibility, and reduces load times. PiperOrigin-RevId: 958260793

2026-08-03

Implement JAX-Triton Ahead-of-Time (AOT) Compilation & Serialization PiperOrigin-RevId: 958260793

Implement JAX-Triton Ahead-of-Time (AOT) Compilation & Serialization PiperOrigin-RevId: 957098418

2026-07-29

[Mosaic GPU] Move GPU module loading and function resolution out of runtime. This change removes the required compilation of the initialization function for GPU kernels and loads modules directly in custom call handling. * Extract GPU binary and kernel metadata (name, dynamic shared memory size, and cluster size) as module attributes during launch lowering instead of generating a compiled initialization function. * Serialize the GPU binary and metadata directly in MosaicGpuKernelProto. * Invoke CUDA Driver and NVSHMEM APIs directly from custom_call.cc during kernel initialization. * Remove mosaic_gpu_module_load and mosaic_gpu_get_function from runtime.cc. We bump the mosaic_gpu proto version which will trigger XLA recompilation if there is a version mismatch. PiperOrigin-RevId: 955981262

2026-07-22

[Mosaic GPU] Move GPU module loading and function resolution out of JIT runtime. This change removes the required JIT compilation of the initialization function for GPU kernels and loads modules directly in custom call handling. * Extract GPU binary and kernel metadata (name, dynamic shared memory size, and cluster size) as module attributes during launch lowering instead of generating a JIT-compiled initialization function. * Serialize the GPU binary and metadata directly in MosaicGpuKernelProto. * Invoke CUDA Driver and NVSHMEM APIs directly from custom_call.cc during kernel initialization without resolving runtime library symbols. * Remove mosaic_gpu_module_load and mosaic_gpu_get_function from runtime.cc. PiperOrigin-RevId: 952007156

2026-03-04

Mosaic GPU: pre-launch check for multicast support This gives a much better error message from kernels using multimem PTX instructions on incompatible systems.

[Mosaic:GPU] Rename is_comm_used as is_nvshmem_used PiperOrigin-RevId: 878508055

2026-03-02

Add serialization support for Mosaic GPU kernels. This change introduces a proto definition to serialize the state of a compiled Mosaic GPU kernel, including the object file and function names. The `CustomCallResources` struct is updated to stroe the kernel hash and include serialization and deserialization methods, allowing it to be used with XLA FFI's type registry. Also introduces `GetOrCreateKernel` to deduplicate kernels based on their kernel_hash (preserved during de/serialization). PiperOrigin-RevId: 877392204

2026-01-15

Follow up pr after triton integration cl and tokamax pr. Remove obsolete cluster dimensions from TritonKernel proto. PiperOrigin-RevId: 856837148

2025-12-11

Integrate Triton up to 8d445186 https://github.com/openxla/triton/tree/triton_integrate_branch-1.15 PiperOrigin-RevId: 843206657

2025-08-03

[JAX] Migrate most remaining C++ code in jax to the jax:: namespace. PiperOrigin-RevId: 790476584

The documentation is generated from this commit.

2025-08-01

[JAX] Migrate most remaining C++ code in jax to the jax:: namespace. PiperOrigin-RevId: 789738942

The documentation is generated from this commit.

2025-04-23

Move contents of jaxlib/xla into jaxlib/ Having the directory structure of the jaxlib wheel be different to the source tree confuses type checkers such as mypy, since sometimes they find type stubs in the installed jaxlib wheel, and sometimes from the installed source tree. Instead: * don't include type stubs in the jaxlib wheel * don't install the jaxlib wheel as part of pre-commit * make sure that the location of type stubs (and the underlying libraries) is in the same position in the `jaxlib/` directory of the JAX source tree as it would be for the jaxlib wheel when installed. For now, we leave some stubs that forward from the old locations to the new locations for certain headers and modules. These will be removed after migrating some users. PiperOrigin-RevId: 750650528

Move contents of jaxlib/xla into jaxlib/ Having the directory structure of the jaxlib wheel be different to the source tree is confusing type checkers such as mypy. PiperOrigin-RevId: 750398365

2025-03-25

[JAX] [XLA:Python] Migrate more Python modules to JAX. PiperOrigin-RevId: 740507886

[JAX] [XLA:Python] Migrate pytree module to JAX. PiperOrigin-RevId: 740142231

2024-01-19

[triton] Pass cluster_dims to TritonKernel and use cuLaunchKernel if size <= 1 PiperOrigin-RevId: 599809560

2023-10-03

Move jaxlib/ under jax/_src/lib. Prototype only, not ready for review. PiperOrigin-RevId: 570062727

2023-08-16

[jax_triton] Add user-specified `name` field to serialized format. PiperOrigin-RevId: 557415723

2023-07-29

[jax_triton] Add support for float scalar inputs. Python `float`s are inferred as "f64". Values can be passed as "f32" using `np.float32(value)`. PiperOrigin-RevId: 552036612

2023-06-22

Make JAX-Triton calls serializable. PiperOrigin-RevId: 542524794