jax-ml/jax · error · ValueError

Async copies only support striding up to 5 dimensions

Error message

Async copies only support striding up to 5 dimensions

What it means

TMA descriptors can describe at most 5 strided dimensions. After Mosaic squeezes/collapses singleton and leading index dimensions, if the slice still has more than 5 dimensions this ValueError is raised in _prepare_tma. The library mirrors the hardware limit of the TMA tensor map.

Source

Thrown at jax/experimental/mosaic/gpu/launch_context.py:1125

      slice_shape = list(drop.transform_shape(slice_shape))
      # After _prepare_async_copy, squeezed dims have been permuted to the
      # front (via a TransposeTransform in gmem_transform). So in the
      # transformed shape, they occupy the first `num_squeezed_dims`
      # positions.
      # TODO(bchetioui): move the creation of the `TransposeTransform`
      # here instead of in _prepare_async_copy.
      squeezed_dims = tuple(d for i, d in enumerate(squeezed_dims) if i not in unit_dims)
      num_squeezed_dims = len(squeezed_dims)
      if len(slice_shape) > 5 and squeezed_dims:
        # We can try to collapse all squeezed dims into one.
        squeezed_dim_strides = tuple(gmem_strides[d] for d in squeezed_dims)
        collapse = CollapseLeadingIndicesTransform(squeezed_dim_strides)
        gmem_transform = (*gmem_transform, collapse)
        dyn_base_indices = collapse.transform_index(dyn_base_indices)
        slice_shape = list(collapse.transform_shape(tuple(slice_shape)))
        num_squeezed_dims = 1
      if len(slice_shape) > 5:
        raise ValueError("Async copies only support striding up to 5 dimensions")
    del squeezed_dims

    # pyrefly: ignore[redefinition]
    dyn_base_indices: list[ir.Value] = list(dyn_base_indices)
    slice_shape = list(slice_shape)
    assert all(d == 1 for d in slice_shape[:num_squeezed_dims])

    # Partitioned loads have already been processed (before transforms).
    # We process non-partitioned collective loads here, because only here are we
    # able to know in what order the data will be written to SMEM. Transposes
    # and tiling change that order and if we picked a partition based on the
    # untransformed slice shape, we might have ended up with a non-contiguous
    # SMEM window, which would no longer be realizable in a single TMA transfer.
    collective_size = math.prod(self.cluster_size[d] for d in collective)
    if collective_size > 1 and not isinstance(leader_tracked, _Partitioned):
      assert gather_indices is None  # Checked above.
      def partition_dim(dim: int, idx: ir.Value, num_chunks: int):
        # No need to partition squeezed dims. They don't even exist in smem_ref.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Reshape the tensor or slice so the copy spans at most 5 dimensions before calling async_copy.
Ensure leading dimensions can be collapsed (make them contiguous/unit-stride) so the CollapseLeadingIndicesTransform reduces the rank below 5.
Split the copy into multiple async_copy calls over lower-rank slices.

Example fix

// before
ctx.async_copy(gmem_ref, smem_ref, gmem_slice=tuple_of_6d_slices)
// after
x = x.reshape(x.shape[:2] + (-1,) + x.shape[4:])  # fold dims to <=5D
ctx.async_copy(gmem_ref, smem_ref, gmem_slice=fewer_dim_slices)

Defensive patterns

Strategy: validation

Validate before calling

effective_rank = len([d for d in slice_shape if d > 1])  # after collapsing unit dims
assert effective_rank <= 5 or len(slice_shape) <= 5, 'TMA supports at most 5 strided dims'

Prevention

Reshape operands to <=5D before entering Mosaic kernels.
Keep layouts contiguous so leading dims collapse.
Split high-rank copies into multiple lower-rank async_copy calls.

When it happens

Trigger: Calling async_copy/async_prefetch with the TMA implementation on a slice whose shape, after collapsing leading unit-stride dimensions, still has rank > 5, e.g. an effective 6D tile.

Common situations: Operating on high-rank tensors (e.g. batched attention with heads/batch/seq/heads dims) in Mosaic GPU kernels; adding dimensions via transforms instead of reshaping first.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/1d30969bf0d31727. Report an issue: GitHub.