Pallas Changelog

Pallas Changelog#

This is the list of changes specific to jax.experimental.pallas. For the overall JAX change log see here.

Unreleased#

  • Changes

    • 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 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 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 jax.experimental.export.export(), since the corresponding custom call is not considered export-stable and needs to be enabled explicitly.

Mosaic GPU#

Released with JAX 0.10.2#

  • New features

    • Added 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

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

  • Removals

    • Removed the previously deprecated jax.experimental.pallas.tpu.KernelType and jax.experimental.pallas.tpu.repeat().

  • Deprecations

    • Deprecated pltpu.HOST and pltpu.MemorySpace.HOST in favor of pl.HOST.

Mosaic GPU#

Released with JAX 0.10.1#

TPU#

Mosaic GPU#

  • New features

  • 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 jax.experimental.pallas.kernel() with these keyword arguments must be updated.

TPU#

  • Deprecations

  • 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 jax.experimental.pallas.mosaic_gpu.atomic_add(), jax.experimental.pallas.mosaic_gpu.atomic_max(), jax.experimental.pallas.mosaic_gpu.atomic_min(), jax.experimental.pallas.mosaic_gpu.atomic_and(), jax.experimental.pallas.mosaic_gpu.atomic_or(), and jax.experimental.pallas.mosaic_gpu.atomic_xor().

    • Added a leader_tracked argument to 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:

  • Changes

    • Removed the backend argument of 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 jax.experimental.pallas.triton.

    • Renamed 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 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

  • 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:

  • Deprecations

  • Removals

    • Removed the previously deprecated jax.experimental.pallas.tpu.TPUCompilerParams, jax.experimental.pallas.tpu.TPUMemorySpace, 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

Released with jax 0.7.0#

  • New functionality

    • Added a new decorator jax.experimental.pallas.loop() which allows to write stateless loops as functions.

    • Added new multiple buffering and lookahead functionality to 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

Released with jax 0.6.1#

  • Removals

  • Changes

    • 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.

    • jax.experimental.pallas.pallas_call() now requires compiler_params to be a backend-specific dataclass instead of a param to value mapping.

    • 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 jax.experimental.checkify. Note that debug checks are not executed unless jax.experimental.pallas.enable_debug_checks is set.

Released with jax 0.5.0#

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 jax.experimental.pallas.tpu.CostEstimate and jax.experimental.tpu.run_scoped(). Both are now available in jax.experimental.pallas.

  • New functionality

    • Added a cost estimate tool 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

    • 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.

    • jax.experimental.pallas.BlockSpec no longer supports the previously deprecated argument order, where index_map comes before block_shape.

  • Deprecations

  • New functionality

    • 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.

    • 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 (#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

    • 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.

    • 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 (#22552).

    • The method compute_index of 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 (#22593).

    • Fixed the interpret mode to work with BlockSpec that involve padding (#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: Grids and BlockSpecs.

    • Improved error messages for the jax.experimental.pallas.pallas_call() API.

    • Added lowering rules for TPU for lax.shift_right_arithmetic (#22279) and lax.erf_inv (#22310).

    • Added initial support for shape polymorphism for the Pallas TPU custom kernels
      (#22084).

    • Added TPU support for checkify. (#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)#