jax-ml/jax · error · ValueError

Async copies require the number of bits copied along the las

Error message

Async copies require the number of bits copied along the last dimension to be divisible by 128, but got {zeroth_bw}

What it means

TMA requires the innermost copied dimension to be at least 16 bytes and a multiple of 16 bytes; Mosaic enforces this by checking slice_shape[-1] * element_bitwidth % 128 == 0. Violating this raises the error with the offending bitwidth.

Source

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

        elif rem_collective_size % slice_size == 0:
          # This is an optimization and it lets us skip squeezed dims.
          if slice_size > 1:
            dim_idx = arith.remui(idx, c(slice_size, index))
            partition_dim(dim, dim_idx, slice_size)
            idx = arith.divui(idx, c(slice_size, index))
            rem_collective_size //= slice_size
        else:
          break  # We failed to partition the leading dimensions.
      del idx  # We overwrote the block index in the loop.
      if rem_collective_size > 1:
        raise ValueError(
            "None of the leading dimensions in the transformed slice shape"
            f" {slice_shape} is divisible by the collective size"
            f" {collective_size}"
        )

    if (zeroth_bw := slice_shape[-1] * element_bitwidth) % 128 != 0:
      raise ValueError(
          "Async copies require the number of bits copied along the last"
          f" dimension to be divisible by 128, but got {zeroth_bw}"
      )
    if (
        swizzle is not None
        and swizzle != mgpu_dialect.SwizzlingMode.kNoSwizzle
        and slice_shape[-1] != (swizzle * 8) // element_bitwidth
    ):
      raise ValueError(
          f"Async copies with {swizzle=} require the last dimension of the"
          f" slice to be exactly {swizzle} bytes i.e. "
          f" {(swizzle * 8) // element_bitwidth} elements, but got"
          f" {slice_shape[-1]} elements."
      )
    return (smem_ref, slice_shape, dyn_base_indices, gmem_transform)

  def async_copy(
      self,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Round the last-dimension slice up (pad) so last_dim * element_bytes is a multiple of 16.
2. Choose tile sizes that are multiples of 16 bytes: e.g. multiples of 4 for f32, 8 for f16/bf16, 16 for int8.
3. If padding is unacceptable, avoid the TMA implementation for this copy.

Example fix

// before
ctx.async_copy(..., gmem_slice=(slice(0, 64), slice(0, 3)))  # f32: 12 bytes
// after
ctx.async_copy(..., gmem_slice=(slice(0, 64), slice(0, 4)))  # f32: 16 bytes

last_bw = slice_shape[-1] * utils.bitwidth(element_type)
assert last_bw % 128 == 0, f'last dim must be a multiple of 16 bytes, got {last_bw} bits'

try:
    ctx.async_copy(...)
except ValueError as e:
    if 'divisible by 128' in str(e):
        slice_shape[-1] = _round_up_to_bytes(slice_shape[-1], 16, element_type)
    else:
        raise

When it happens

Trigger: Calling async_copy/async_prefetch with the TMA implementation where the last dimension extent times element size is not a multiple of 16 bytes, e.g. copying a last dim of 3 float32 elements (96 bits).

Common situations: Small or odd innermost tile sizes (e.g. seq length 12 with f32); using narrow element types (int8/f16) with dims not scaled to reach 16-byte multiples.

Related errors

• async_copy requires all GMEM strides except the last one to
• Async copies only support striding up to 5 dimensions
• None of the leading dimensions in the transformed slice shap
• Only byte-aligned shapes are supported. Got shape: {ref.dtyp
• copy_gmem_to_smem with a barrier is only supported Hopper an

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