jax-ml/jax · error · ValueError

Async copies with {swizzle=} require the last dimension of t

Error message

Async copies with {swizzle=} require the last dimension of the slice to be exactly {swizzle} bytes i.e.  {(swizzle * 8) // element_bitwidth} elements, but got {slice_shape[-1]} elements.

What it means

When a swizzle mode (e.g. 32/64/128-byte swizzling) is requested for an async copy, TMA requires the innermost slice dimension to exactly match the swizzle width in bytes: slice_shape[-1] must equal (swizzle * 8) // element_bitwidth elements. Otherwise the copy would produce incorrectly swizzled shared-memory data.

Source

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

      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,
      *,
      src_ref: ir.Value,
      dst_ref: ir.Value,
      gmem_slice: Any = (),
      gmem_transform: MemRefTransform | tuple[MemRefTransform, ...] = (),
      gmem_peer_id: int | ir.Value | GlobalBroadcast | None = None,
      barrier: utils.BarrierRef | None = None,
      swizzle: int | None = None,
      arrive: bool | None = None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Set the innermost slice size to (swizzle_bytes * 8) // element_bitwidth elements (e.g. 32 f32 elements for 128B swizzle).
Or change the swizzle mode to match the existing inner dimension byte width, or use kNoSwizzle/None.
Keep dtype and swizzle consistent when retuning tiles: halving element width requires doubling the element count.

Example fix

// before
ctx.async_copy(..., swizzle=mgpu.SwizzlingMode.k128B, gmem_slice=(slice(0, 64), slice(0, 64)))  # f32 -> 256B inner
// after
ctx.async_copy(..., swizzle=mgpu.SwizzlingMode.k128B, gmem_slice=(slice(0, 64), slice(0, 32)))  # 32 f32 = 128B

Defensive patterns

Strategy: validation

Validate before calling

if swizzle is not None and swizzle != mgpu.SwizzlingMode.kNoSwizzle:
    need = (int(swizzle) * 8) // utils.bitwidth(element_type)
    assert slice_shape[-1] == need, f'swizzle {swizzle} requires last dim == {need} elements'

Prevention

Derive swizzle mode from the inner tile byte width rather than hardcoding.
Change swizzle and tile size together when switching dtypes.
Write a helper that picks k32B/k64B/k128B from the last-dim byte size.

When it happens

Trigger: Calling async_copy/async_prefetch with swizzle=SwizzlingMode.k128B (or 32B/64B) and a last-dimension slice whose byte width is not exactly the swizzle width, e.g. 64 elements of f32 with 128B swizzle (256 bytes).

Common situations: Configuring SMEM swizzling for tensor-core MMA layouts in Mosaic; changing tile inner size or dtype without updating the swizzle mode to match.

Related errors

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