jax-ml/jax · error · NotImplementedError

Transfer of {total_bits} bits is not divisible by {8 * utils

Error message

Transfer of {total_bits} bits is not divisible by {8 * utils.WARPGROUP_SIZE}

What it means

async_store_smem uses arrive_expect_tx to have the cluster barrier count incoming bytes per warp; the total transfer size in bits must be divisible by 8*WARPGROUP_SIZE (i.e. bytes divisible by 128 per warpgroup) or the TMA transaction count cannot be computed.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:619

    raise NotImplementedError(f"Expected TiledLayout, got {type(layout)}")

  ref = op.destination
  transforms_attr = inference_utils.in_transforms(op)[0]
  swizzle = swizzle_from_transforms_attr(transforms_attr)
  unwrapped_ref = unwrap_transformed_memref(ref, transforms_attr)
  tiling_transform, = memref_transforms_from_transforms_attr(transforms_attr)
  assert isinstance(tiling_transform, lc.TileTransform)

  dialect_barrier = utils.DialectBarrierRef.from_barrier_memref(op.barrier)
  barrier_ref = dialect_barrier.barrier_ref

  cluster_dim = gpu.Dimension(op.cluster_dim.value)  # pyrefly: ignore[missing-attribute]
  cluster_idx = arith.index_cast(index, op.cluster_idx)
  cluster_barrier_ref = barrier_ref.remap_to_cluster(cluster_dim, cluster_idx)

  total_bits = math.prod(value.shape) * utils.bitwidth(value.mlir_dtype)
  if total_bits % (8 * utils.WARPGROUP_SIZE):
    raise NotImplementedError(
        f"Transfer of {total_bits} bits is not divisible by "
        f"{8 * utils.WARPGROUP_SIZE}"
    )
  cluster_barrier_ref.arrive_expect_tx(total_bits // 8 // utils.WARPGROUP_SIZE)

  atomic = None
  if op.atomic_type is not None:
    atomic = str(mgpu.AtomicOpType(op.atomic_type.value))  # pyrefly: ignore[missing-attribute]

  def store_tiled_async(optimized: bool):
    value.store_tiled_async(
        unwrapped_ref,
        barrier_ref,
        cluster_dim=cluster_dim,
        cluster_idx=cluster_idx,
        swizzle=swizzle.value if swizzle != mgpu.SwizzlingMode.kNoSwizzle else None,
        optimized=optimized,
        tiling_rank=len(tiling_transform.tiling),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad the stored value's shape so total bits are divisible by 8*WARPGROUP_SIZE
2. Use vector.store / regular store to SMEM instead of async_store_smem for odd sizes
3. Choose a dtype/tile size whose byte size is a multiple of 128 per warpgroup

Example fix

// before
async_store_smem(x, ref)  # x has 30 f32 elements
// after
x = pad_to_multiple(x, 32)  # total bytes divisible by 128
async_store_smem(x, ref)

import math
from jax.experimental.mosaic.gpu import utils
total_bits = math.prod(value.shape) * utils.bitwidth(dtype)
assert total_bits % (8 * utils.WARPGROUP_SIZE) == 0, 'pad store to warpgroup multiple'

When it happens

Trigger: Storing a tensor whose total elements * dtype bitwidth is not divisible by 8 * WARPGROUP_SIZE (1024 bits on 128-thread warpgroups).

Common situations: Small or odd-shaped stores (e.g. f16 tensors with a non-multiple-of-64 element count) to SMEM via async_store_smem with cluster barriers enabled.

Related errors

• Expected TiledLayout, got {type(layout)}
• Offset {i} is not divisible by tile size {t}
• Stride {s} is not divisible by {d} (tile size = {t}). Stride
• Cannot assign layout to async load with gather indices since
• Only byte-aligned bitcasts are supported.

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