jax-ml/jax · error · NotImplementedError

Transfer is not a multiple of {WARPGROUP_SIZE} bytes

Error message

Transfer is not a multiple of {WARPGROUP_SIZE} bytes

What it means

The async store's total byte count must be a multiple of WARPGROUP_SIZE bytes because arrive_expect_tx and the TMA transaction accounting operate in warpgroup-sized (typically 128-byte) units. Otherwise the lowering raises NotImplementedError.

Source

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

        atomic_type=atomic_type,
        optimized=optimized,
    )
    return ()

  match remaining_ref_transforms:
    case (gpu_core.UnswizzleRef(swizzle), gpu_core.UntilingTransform(tiling)):
      pass
    case _:
      raise NotImplementedError("async_store_smem requires a tiled and swizzled ref")

  total_bits = math.prod(shape) * dtypes.itemsize_bits(dtype)
  if total_bits % 8:
    raise ValueError(
        f"Can only transfer integer bytes (shape={shape}, dtype={dtype})"
    )
  total_bytes = total_bits // 8
  if total_bytes % WARPGROUP_SIZE:
    raise NotImplementedError(f"Transfer is not a multiple of {WARPGROUP_SIZE} bytes")

  peer_barrier = barrier.remap_to_cluster(gpu_cluster_dim, cluster_idx_val)
  peer_barrier.arrive_expect_tx(total_bytes // WARPGROUP_SIZE)

  lowering._ensure_fa(src, dtype).store_tiled_async(
      ref_smem,
      barrier,
      cluster_dim=gpu_cluster_dim,
      cluster_idx=cluster_idx_val,
      swizzle=swizzle,
      optimized=optimized,
      tiling_rank=len(tiling),
      atomic=atomic,
  )
  return ()


def async_store_smem(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Round buffer/block sizes up so the transfer is a multiple of 128 bytes (e.g. 32 float32 or 64 bf16 elements)
Pad the SMEM buffer and value to the warpgroup-aligned size
Use the Warp (non-warpgroup) semantics or a synchronous store if small transfers are unavoidable

Example fix

# before
async_store_smem(smem, x[:48], barrier)  # 192 bytes, not multiple of 128
# after
async_store_smem(smem, x[:64], barrier)  # 256 bytes = 2 * WARPGROUP_SIZE

Defensive patterns

Strategy: validation

Validate before calling

import math
from jax import dtypes
bytes_ = math.prod(shape) * dtypes.itemsize_bits(dtype) // 8
assert bytes_ % 128 == 0, f'{bytes_}B not a multiple of warpgroup size (128B)'

Prevention

Choose block sizes that are multiples of 128 bytes (e.g. 32 fp32 / 64 bf16 elements)
Beware tail blocks in pipelined loops; pad them

When it happens

Trigger: async_store_smem where prod(shape)*itemsize/8 is not divisible by the warpgroup size (128), e.g. storing a 64-byte block or a small non-conforming tail block.

Common situations: Small tail blocks in pipelined loops; choosing block sizes like (48,) float32 that look aligned but aren't multiples of 128 bytes; porting kernels between warp and warpgroup semantics.

Related errors

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