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.

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.

Appendix: 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)