(pallas-changelog)= # Pallas Changelog This is the list of changes specific to {class}`jax.experimental.pallas`. For the overall JAX change log see [here](https://docs.jax.dev/en/latest/changelog.html). ## Unreleased * Deprecations * Pallas Triton ops in {mod}`jax.experimental.pallas.ops.gpu` are deprecated. Please use [`tokamax`](https://github.com/openxla/tokamax) for equivalent implementations if available, e.g., `tokamax.layer_norm`, `tokamax.dot_product_attention`. * {func}`jax.experimental.pallas.core_map` is deprecated. Please migrate to {func}`jax.experimental.pallas.kernel`. ### Mosaic GPU * New features * The ``barrier`` argument of {func}`jax.experimental.pallas.mosaic_gpu.copy_gmem_to_smem` must now be omitted on pre-Hopper GPUs (which use the ``cp.async`` implementation). When omitted, the completion of the copy must be awaited via {func}`jax.experimental.pallas.mosaic_gpu.wait_gmem_to_smem`. ## Released with JAX 0.11.0 (July 16, 2026) * Changes * {func}`jax.experimental.pallas.kernel` now only aliases closed over or passed in refs if the kernel writes to them. If you need aliasing without an explicit write, use {func}`jax.experimental.pallas.tpu.touch`. ### Triton The Triton backend is deprecated and will be removed in a future version of JAX. To keep using Pallas on GPU, please migrate to the Mosaic GPU backend. To keep using Triton, switch to the official Triton bindings and `jax_triton`. * New features * Added {data}`jax.experimental.pallas.triton.CUSTOM_CALL_TARGET_NAME`, the name of the custom call targeted by Pallas Triton. The exact value is not guaranteed to be stable, but it is safe to assume that it is always aligned with what the lowering emits. This is useful, for example, when exporting a computation with Pallas Triton kernels via {func}`jax.experimental.export.export`, since the corresponding custom call is not considered export-stable and needs to be enabled explicitly. ### Mosaic GPU * New features * Added support for Ampere matrix multiply-accumulate instructions via {func}`jax.experimental.pallas.mosaic_gpu.mma`. * Added support for `cp.async` to {func}`jax.experimental.pallas.mosaic_gpu.copy_gmem_to_smem`. * {func}`jax.experimental.pallas.mosaic_gpu.wgmma` now supports mixing the e4m3 and e5m2 FP8 types across its two operands. * Added support for manual "ready" barriers to {func}`jax.experimental.pallas.mosaic_gpu.emit_pipeline_warp_specialized`. * Added ``oob_fill_mode`` to {class}`jax.experimental.pallas.mosaic_gpu.BlockSpec`, which controls the behavior for out-of-bounds accesses during pipelined GMEM-to-SMEM copies. * Changes * Turned on Warpgroup lowering by default. If this causes regression or breakages, switch back to {data}`jax.experimental.pallas.mosaic_gpu.LoweringSemantics.Lane` in your kernel's {class}`jax.experimental.pallas.mosaic_gpu.CompilerParams` and file a bug. * Deprecations * The `idx` parameter of {func}`jax.experimental.pallas.mosaic_gpu.load` is deprecated. Index the ref explicitly via `ref.at[idx]` prior to loading from it. ## Released with JAX 0.10.2 * New features * Added {data}`jax.experimental.pallas.enable_poison_buffers` (config flag `jax_pallas_poison_buffers`) to poison (initialize with NaNs or lowest possible integers) any scratch buffers allocated by Pallas for debugging. * Deprecations * `pl.debug_checks_enabled` is deprecated. Use `pl.enable_debug_checks.value`. * `pl.dot` was moved into {mod}`jax.experimental.pallas.triton`. Accessing it via {mod}`jax.experimental.pallas` is deprecated. You can use {func}`jax.numpy.dot`, {func}`jax.numpy.einsum` or the `@` operator instead in a TPU or MGPU kernel. ### TPU * New features * Added `jax_pallas_auto_assign_collective_ids` config flag to allow two new custom semaphore barrier collective IDs modes: ('yes') assigning missing collective IDs automatically or ('override') overridding all collective IDs and assigning them automatically, both based on the serialized kernel hash. * Changes * {class}`jax.experimental.pallas.tpu.CompilerParams` now defaults ``needs_layout_passes`` to True. The layout passes are still a work in progress. Please file a bug if you encounter a compilation error with them enabled. * {class}`jax.experimental.pallas.tpu.CompilerParams` now defaults ``use_tc_tiling_on_sc`` to True for SparseCore kernels. * Removals * Removed the previously deprecated {class}`jax.experimental.pallas.tpu.KernelType` and {func}`jax.experimental.pallas.tpu.repeat`. * Deprecations * Deprecated `pltpu.HOST` and `pltpu.MemorySpace.HOST` in favor of `pl.HOST`. ### Mosaic GPU * New features * Support using {func}`jax.experimental.pallas.multiple_of` to specify divisibility requirements on dynamic indices. * Support allocating multidimensional `plgpu.Barrier`s and `plgpu.ClusterBarrier`s, by providing a nD shape as the `num_barriers` parameter. * Support for TMEM references with batch dimensions. This includes indexing over batch dimensions and `plgpu.async_store_tmem` support for TMEM references with batch dimensions. * Support TMA scatter. * Rename `TMA_GATHER_INDICES_LAYOUT` to `TMA_INDICES_LAYOUT`. * Changes * Breaking change: {func}`jax.experimental.pallas.program_id` and {func}`jax.experimental.pallas.num_programs` no longer work inside kernels defined via {func}`jax.experimental.pallas.mosaic_gpu.kernel`. Use {func}`jax.lax.axis_index` and {func}`jax.lax.axis_size` instead. * {func}`jax.experimental.pallas.mosaic_gpu.kernel` now has the same API as {func}`jax.experimental.pallas.kernel`, meaning that it can be used as a decorator and also uses ``out_type`` and ``scratch_types`` instead of ``out_shape`` and ``scratch_shapes``. * Removals * Deleted `plgpu.unswizzle_ref` and `plgpu.untile_ref`. * Deprecations * Using {func}`jax.experimental.pallas.pallas_call` for Mosaic GPU kernels is deprecated. Please migrate to {func}`jax.experimental.pallas.mosaic_gpu.kernel` and {func}`jax.experimental.pallas.mosaic_gpu.emit_pipeline`. ## Released with JAX 0.10.1 * Changes * Added {func}`jax.experimental.pallas.align_to`, a utility that rounds a value up to the nearest multiple of a given alignment. * {func}`jax.experimental.pallas.pallas_call` no longer supports checkify. We expect this change to affect few users, as in our experience most kernels either perform no checking or use {func}`jax.experimental.pallas.debug_check` for conditionally-enabled runtime checks. * {func}`jax.experimental.pallas.kernel` now always aliases Refs that are passed in or closed-over. ### TPU * Removals * Removed the `kernel_type` field from {class}`jax.experimental.pallas.tpu.CompilerParams`. It was only used for writing SparseCore kernels via {func}`jax.experimental.pallas.pallas_call`, which is now unsupported. The recommended API for SparseCore kernels is {func}`jax.experimental.pallas.kernel`. ### Mosaic GPU * New features * Added {func}`jax.experimental.pallas.mosaic_gpu.barrier_test` function; a non-blocking equivalent of {func}`jax.experimental.pallas.mosaic_gpu.barrier_wait` only supported in a warp context. * Changes * Breaking change: removed `plgpu.TransposeTransform`. ## Released with JAX 0.10.0 * Changes * Breaking change: refactored `pl.kernel` to use `out_type` instead of `out_shape` and `scratch_types` instead of `scratch_shapes`. Existing usages calling {func}`jax.experimental.pallas.kernel` with these keyword arguments must be updated. ### TPU * Deprecations * `pltpu.semaphore`, `pltpu.DeviceIdType`, `pltpu.semaphore_signal`, `pltpu.semaphore_wait`, and `pltpu.semaphore_read` are now available in {mod}`jax.experimental.pallas`. Accessing them via {mod}`jax.experimental.pallas.tpu` is deprecated. * Removals * Removed the previously deprecated `pltpu.ANY` and `pltpu.MemorySpace.ANY`. Use `pl.ANY` instead. * Removed the deprecated `pltpu.delay`, which is now available as `pl.delay`. ### Mosaic GPU * New features * Added the `leader_tracked` argument to `ClusterBarrier`, which allows tracking barrier completions solely from the leader block along a specific axis in a cluster. ## Released with JAX 0.9.2 * New features: * Added support for atomics to the Mosaic GPU backend, through {func}`jax.experimental.pallas.mosaic_gpu.atomic_add`, {func}`jax.experimental.pallas.mosaic_gpu.atomic_max`, {func}`jax.experimental.pallas.mosaic_gpu.atomic_min`, {func}`jax.experimental.pallas.mosaic_gpu.atomic_and`, {func}`jax.experimental.pallas.mosaic_gpu.atomic_or`, and {func}`jax.experimental.pallas.mosaic_gpu.atomic_xor`. * Added a `leader_tracked` argument to {func}`jax.experimental.pallas.mosaic_gpu.copy_gmem_to_smem`. This adds support for cluster-collective partitioned/replicated GMEM to SMEM copies whose completion is tracked only by the leader block. Removed the `partitioned_axis` argument (which provided limited access to this feature) from the API. ## Released with JAX 0.9.1 * New features: * Added a {func}`jax.experimental.pallas.with_scoped` decorator that provides the function with scope-allocated scratch buffers. * Changes * Removed the `backend` argument of {func}`jax.experimental.pallas.pallas_call` in favor of `compiler_params`. For example, to force the use of the Triton backend you have to now write `compiler_params=pltriton.CompilerParams()`, where `pltriton` refers to {mod}`jax.experimental.pallas.triton`. * Renamed {class}`jax.experimental.pallas.tpu.KernelType` to `CoreType`. The old name is deprecated and will be removed in a future release. ## Released with JAX 0.9.0 * New features: * Added a `reduction_scratch_bytes` field to {class}`jax.experimental.pallas.mosaic_gpu.CompilerParams`. This gives user control over how much shared memory Pallas is allowed to reserve for cross-warp reductions on GPU. Increasing this value typically allows for faster reductions. * Changes * The default lowering path on GPU now goes through Mosaic GPU. To keep using Triton, call {func}`jax.experimental.pallas.pallas_call` with the `backend` argument set to `'triton'`. * Removals * Removed the previously deprecated `pl.atomic_*`, `pl.load`, `pl.store`, `pl.swap` and `pl.max_contiguous`. ## Released with jax 0.8.1 * New features: * Added {func}`jax.experimental.pallas.tpu.get_tpu_info` to get TPU hardware information. * Deprecations * `pl.max_contiguous` has been moved to {mod}`jax.experimental.pallas.triton`. Accessing it via {mod}`jax.experimental.pallas` is deprecated. * `pl.swap` is deprecated and will be removed in a future release. Use indexing or backend-specific loading/storing APIs instead. * Removals * Removed the previously deprecated {class}`jax.experimental.pallas.tpu.TPUCompilerParams`, {class}`jax.experimental.pallas.tpu.TPUMemorySpace`, {class}`jax.experimental.pallas.tpu.TritonCompilerParams`. ## Released with jax 0.7.1 * New features: * `pltpu.make_async_remote_copy` and `pltpu.semaphore_signal`'s `device_id` argument now allows user to pass in a dictionary that only specifies the device index along the communication axis, instead of the full coordinates. It also supports TPU core id index. * `jax.debug.print` now works in Pallas kernels and is the recommended way to print. * Deprecations * `pl.atomic_*` APIs have been moved to {mod}`jax.experimental.pallas.triton`. Accessing them via {mod}`jax.experimental.pallas` is deprecated. * `pl.load` and `pl.store` are deprecated. Use indexing or backend-specific loading/storing APIs instead. ## Released with jax 0.7.0 * New functionality * Added a new decorator {func}`jax.experimental.pallas.loop` which allows to write stateless loops as functions. * Added new multiple buffering and lookahead functionality to {func}`jax.experimental.pallas.tpu.emit_pipeline`. Input buffers can now be multiple-buffered with more than 2 buffers and support a lookahead option to fetch blocks that are an arbitrary number of grid iterations ahead rather than the immediate next iterations. Additionally, pipeline state can now be held in registers to reduce scalar memory usage. * Deprecations * {class}`jax.experimental.pallas.triton.TritonCompilerParams` has been renamed to {class}`jax.experimental.pallas.triton.CompilerParams`. The old name is deprecated and will be removed in a future release. * {class}`jax.experimental.pallas.tpu.TPUCompilerParams` and {class}`jax.experimental.pallas.tpu.TPUMemorySpace` have been renamed to {class}`jax.experimental.pallas.tpu.CompilerParams` and {class}`jax.experimental.pallas.tpu.MemorySpace`. The old names are deprecated and will be removed in a future release. ## Released with jax 0.6.1 * Removals * Removed previously deprecated {mod}`jax.experimental.pallas.gpu`. To use the Triton backend import {mod}`jax.experimental.pallas.triton`. * Changes * {func}`jax.experimental.pallas.BlockSpec` now takes in special types in addition to ints/None in the `block_shape`. `indexing_mode` has been removed. To achieve "Unblocked", pass a `pl.Element(size)` into `block_shape` for each entry that needs unblocked indexing. * {func}`jax.experimental.pallas.pallas_call` now requires `compiler_params` to be a backend-specific dataclass instead of a param to value mapping. * {func}`jax.experimental.pallas.debug_check` is now supported both on TPU and Mosaic GPU. Previously, this functionality was only supported on TPU and required using the APIs from {mod}`jax.experimental.checkify`. Note that debug checks are not executed unless {data}`jax.experimental.pallas.enable_debug_checks` is set. ## Released with jax 0.5.0 * New functionality * Added vector support for {func}`jax.experimental.pallas.debug_print` on TPU. ## Released with jax 0.4.37 * New functionality * Added support for `DotAlgorithmPreset` precision arguments for `dot` lowering on Triton backend. ## Released with jax 0.4.36 (December 6, 2024) ## Released with jax 0.4.35 (October 22, 2024) * Removals * Removed previously deprecated aliases {class}`jax.experimental.pallas.tpu.CostEstimate` and {func}`jax.experimental.tpu.run_scoped`. Both are now available in {mod}`jax.experimental.pallas`. * New functionality * Added a cost estimate tool {func}`pl.estimate_cost` for automatically constructing a kernel cost estimate from a JAX reference function. ## Released with jax 0.4.34 (October 4, 2024) * Changes * {func}`jax.experimental.pallas.debug_print` no longer requires all arguments to be scalars. The restrictions on the arguments are backend-specific: Non-scalar arguments are currently only supported on GPU, when using Triton. * {class}`jax.experimental.pallas.BlockSpec` no longer supports the previously deprecated argument order, where `index_map` comes before `block_shape`. * Deprecations * The {mod}`jax.experimental.pallas.gpu` submodule is deprecated to avoid ambiguite with {mod}`jax.experimental.pallas.mosaic_gpu`. To use the Triton backend import {mod}`jax.experimental.pallas.triton`. * New functionality * {func}`jax.experimental.pallas.pallas_call` now accepts `scratch_shapes`, a PyTree specifying backend-specific temporary objects needed by the kernel, for example, buffers, synchronization primitives etc. * {func}`checkify.check` can now be used to insert runtime asserts when pallas_call is called with the `pltpu.enable_runtime_assert(True)` context manager. ## Released with jax 0.4.33 (September 16, 2024) ## Released with jax 0.4.32 (September 11, 2024) * Changes * The kernel function is not allowed to close over constants. Instead, all the needed arrays must be passed as inputs, with proper block specs ({jax-issue}`#22746`). * New functionality * Improved error messages for mistakes in the signature of the index map functions, to include the name and source location of the index map. ## Released with jax 0.4.31 (July 29, 2024) * Changes * {class}`jax.experimental.pallas.BlockSpec` now expects `block_shape` to be passed *before* `index_map`. The old argument order is deprecated and will be removed in a future release. * {class}`jax.experimental.pallas.GridSpec` does not have anymore the `in_specs_tree`, and the `out_specs_tree` fields, and the `in_specs` and `out_specs` tree now store the values as pytrees of BlockSpec. Previously, `in_specs` and `out_specs` were flattened ({jax-issue}`#22552`). * The method `compute_index` of {class}`jax.experimental.pallas.GridSpec` has been removed because it is private. Similarly, the `get_grid_mapping` and `unzip_dynamic_bounds` have been removed from `BlockSpec` ({jax-issue}`#22593`). * Fixed the interpret mode to work with BlockSpec that involve padding ({jax-issue}`#22275`). Padding in interpret mode will be with NaN, to help debug out-of-bounds errors, but this behavior is not present when running in custom kernel mode, and should not be depended on. * Previously it was possible to import many APIs that are meant to be private, as `jax.experimental.pallas.pallas`. This is not possible anymore. * New Functionality * Added documentation for BlockSpec: {ref}`pallas_grids_and_blockspecs`. * Improved error messages for the {func}`jax.experimental.pallas.pallas_call` API. * Added lowering rules for TPU for `lax.shift_right_arithmetic` ({jax-issue}`#22279`) and `lax.erf_inv` ({jax-issue}`#22310`). * Added initial support for shape polymorphism for the Pallas TPU custom kernels\ ({jax-issue}`#22084`). * Added TPU support for checkify. ({jax-issue}`#22480`) * Added clearer error messages when the block sizes do not match the TPU requirements. Previously, the errors were coming from the Mosaic backend and did not have useful Python stack traces. * Added support for TPU lowering with 1D blocks, and relaxed the requirements for the block sizes with at least 2 dimensions: the last 2 dimensions must be divisible by 8 and 128 respectively, unless they span the entire corresponding array dimension. Previously, block dimensions that spanned the entire array were allowed only if the block dimensions in the last two dimensions were smaller than 8 and 128 respectively. ## Released with JAX 0.4.30 (June 18, 2024) * New Functionality * Added checkify support for {func}`jax.experimental.pallas.pallas_call` in interpret mode ({jax-issue}`#21862`). * Improved support for PRNG keys for TPU kernels ({jax-issue}`#21773`).