jax-ml/jax · error · ValueError

async_copy requires all GMEM strides except the last one to

Error message

async_copy requires all GMEM strides except the last one to be a multiple of 16 bytes

What it means

The TMA (Tensor Memory Accelerator) hardware path requires that global-memory tensor strides, except the innermost one, be multiples of 16 bytes. The code checks each stride (in elements) times the element bitwidth is divisible by 128 bits; any outer stride violating this raises the error. This mirrors NVIDIA TMA's global-memory alignment requirement.

Source

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

      swizzle: int | None,
      slice_shape: list[int],
      dyn_base_indices: tuple[ir.Value, ...],
      gather_indices,
      squeezed_dims: tuple[int, ...],
      gmem_transform: tuple[MemRefTransform, ...],
      collective: Sequence[gpu.Dimension],
      leader_tracked: CopyPartition | None = None,
  ):
    """Finalizes setup specific to the TMA implementation of async_copy."""
    index = ir.IndexType.get()
    # The function below is called only to verify the GMEM ref. The output
    # is meant to be ignored.
    _find_kernel_argument_for_gmem_ref(gmem_ref)
    gmem_ref_ty = ir.MemRefType(gmem_ref.type)
    element_bitwidth = utils.bitwidth(gmem_ref_ty.element_type)
    gmem_strides, _ = gmem_ref_ty.get_strides_and_offset()
    if any(s * element_bitwidth % 128 != 0 for s in gmem_strides[:-1]):
      raise ValueError(
          "async_copy requires all GMEM strides except the last one to be a"
          " multiple of 16 bytes"
      )
    # We don't need to do this for gather TMAs, because we'll unroll the
    # transfers ourselves anyway.
    num_squeezed_dims = len(squeezed_dims)
    if gather_indices is None:
      # Drop as many unit-sized dimensions from the transformed shape as we can.
      gmem_shape = tuple(gmem_ref_ty.shape)
      for t in gmem_transform:
        gmem_shape = t.transform_gmem_shape(gmem_shape)
      # The slice shape may pad along 1-sized dimensions. In that case, we do
      # not drop them.
      unit_dims = tuple(
          i for i, (gs, ss) in enumerate(zip(gmem_shape, slice_shape, strict=True))
          if gs == 1 and ss == 1
      )
      # When issuing an `async_prefetch`, there is no SMEM reference to

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pad or reallocate the global tensor so all outer strides are multiples of 16 bytes (e.g. pad leading dims so stride * element_size % 16 == 0).
Check and canonicalize the input layout before the kernel (jnp.reshape/pad to a contiguous 16-byte-aligned layout).
If alignment is impossible, use a non-TMA AsyncCopyImplementation.
Verify utils.bitwidth assumptions: for sub-32-bit types the stride in elements must be proportionally larger.

Example fix

// before
x = jnp.zeros((3, 128), dtype=jnp.float32)  # row stride 3*4=12 bytes
ctx.async_copy(gmem_ref, smem_ref, ..., implementation=mgpu.AsyncCopyImplementation.TMA)
// after
x = jnp.zeros((4, 128), dtype=jnp.float32)  # padded so stride is 16-byte aligned
ctx.async_copy(gmem_ref, smem_ref, ..., implementation=mgpu.AsyncCopyImplementation.TMA)

Defensive patterns

Strategy: validation

Validate before calling

et = gmem_ref.type.element_type
bw = utils.bitwidth(et)
strides, _ = ir.MemRefType(gmem_ref.type).get_strides_and_offset()
assert all(s * bw % 128 == 0 for s in strides[:-1]), 'outer strides must be 16-byte aligned for TMA'

Try / catch

try:
    ctx.async_copy(..., implementation=mgpu.AsyncCopyImplementation.TMA)
except ValueError as e:
    if 'GMEM strides' in str(e):
        x = jnp.ascontiguousarray(_pad_to_16B_align(x))
    else:
        raise

Prevention

Pad input tensors so leading dims give 16-byte-aligned row strides.
Prefer contiguous layouts for tensors passed to TMA kernels.
Add a stride-alignment assert in kernel prologue code.

When it happens

Trigger: Calling async_copy or async_prefetch with implementation=AsyncCopyImplementation.TMA on a GMEM reference whose outer strides * element_bitwidth are not multiples of 128 bits, e.g. an f32 tensor with a row stride of 3 elements (12 bytes).

Common situations: Passing non-contiguous or oddly strided views into a Mosaic kernel; using small leading dimensions (e.g. shape (3, N) f32) in TMA copies; switching a kernel from the non-TMA path to TMA on Hopper+ GPUs.

Related errors

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