jax-ml/jax · error · NotImplementedError

Only copies transferring a number of bytes divisible by the

Error message

Only copies transferring a number of bytes divisible by the warpgroup size are supported. Got {bytes=} but warpgroup size is {WARPGROUP_SIZE}

What it means

On Hopper+ (non-cp.async) copies, Mosaic GPU partitions the copy across the threads of a warpgroup (128 threads), so the total number of bytes moved must be divisible by the warpgroup size (128). If the copy's byte count is not divisible by 128, the lowering raises NotImplementedError because it cannot distribute the work evenly.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:1040

    if barrier is None:
      raise ValueError(
          "copy_gmem_to_smem without a barrier is only supported on pre-Hopper"
          " GPUs, which use the cp.async implementation"
      )

  i32 = ir.IntegerType.get_signless(32)
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    if (
        ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warpgroup
        and ctx.module_ctx.auto_barriers
    ):
      mgpu.warpgroup_barrier()  # Make sure all reads have completed.

    if not is_cp_async:
      assert barrier is not None
      if bytes % WARPGROUP_SIZE:
        raise NotImplementedError(
            "Only copies transferring a number of bytes divisible by the"
            f" warpgroup size are supported. Got {bytes=} but warpgroup size is"
            f" {WARPGROUP_SIZE}"
        )
      if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warpgroup:
        # We arrive uniformly from each thread in the WG, so we need to divide the
        # number of bytes by the number of threads in the WG.
        # TODO: apaszke - Relax this. We can just select the WG leader and have it
        # arrive with the whole transfer size, while everyone else arrives with 0.
        # But we should continue using this scheme as it's likely to be faster.
        bytes //= WARPGROUP_SIZE
        if predicate is not None:
          bytes = arith_dialect.select(predicate, mgpu.c(bytes, i32), mgpu.c(0, i32))
        if is_leader_tracked_copy:
          first_block = arith_dialect.cmpi(
              arith_dialect.CmpIPredicate.eq,
              mgpu.utils.cluster_idx(collective[0]),
              mgpu.c(0, ir.IndexType.get()),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pad the copy so total bytes are a multiple of 128 (e.g. round the block shape up and mask out-of-bounds).
Choose block sizes such that num_elements * dtype_bytewidth % 128 == 0 (e.g. multiples of 128 bytes, or 32 f32 elements).
Split the copy into a bulk aligned part plus a scalar fallback path for the remainder.

Example fix

# before
copy_gmem_to_smem(src_ref.at[:n], smem_ref)  # n*4 bytes not divisible by 128
# after
pad_n = (n + 31) // 32 * 32  # f32: 32 elems = 128 bytes
copy_gmem_to_smem(src_ref.at[:pad_n].pad(0, (0, pad_n - n)), smem_ref)

Defensive patterns

Strategy: validation

Validate before calling

BYTES_PER_WG = 128
nbytes = int(np.prod(block_shape)) * np.dtype(dtype).itemsize
assert nbytes % BYTES_PER_WG == 0, f'{nbytes=} not divisible by {BYTES_PER_WG}'

Prevention

Parameterize block shapes so total bytes are always a multiple of 128.
Add a host-side assert on bytes before launching the kernel.

When it happens

Trigger: A warpgroup-level copy_gmem_to_smem (TMA path, barrier given) where bytes = number_of_elements * dtype_size % 128 != 0, e.g. copying 4 float32 elements (16 bytes) or any odd-sized block.

Common situations: Small tail blocks in a tiled kernel (e.g. a 60-element remainder tile), fp8/bfloat8 copies with element counts not a multiple of 128, or block sizes tuned for a different dtype.

Related errors

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