jax-ml/jax · error · NotImplementedError

async_store_smem requires a tiled and swizzled ref

Error message

async_store_smem requires a tiled and swizzled ref

What it means

In the Warp lowering path, a match statement requires the remaining ref transforms to be exactly (UnswizzleRef(swizzle), UntilingTransform(tiling)) — i.e. the target must be a tiled and swizzled SMEM ref. Any other combination raises NotImplementedError because the async TMA store instruction only works with that layout.

Source

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

    atomic_type = None
    if atomic is not None:
      atomic_type = _atomic_op_type_to_int(AtomicOpType(atomic))
    mgpu.dialect.async_store_smem(
        src,
        ref_smem,
        barrier.as_barrier_memref(),
        gpu_cluster_dim.value,
        cluster_idx_i32,
        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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Construct the destination with the standard tiling+swizzle transforms expected by the pipeline (e.g. use the layouts produced by mgpu/pl.tiled helpers)
2. Copy to a properly tiled+swizzled intermediate SMEM buffer and store from there
3. Fall back to a synchronous store if layout flexibility is more important than async overlap

Example fix

# before
plain_smem = pl.SMEM((128, 128), dtype)
async_store_smem(plain_smem, x, barrier)
# after
tiled_smem = make_tiled_swizzled_smem((128, 128), dtype)  # tiled + swizzled
async_store_smem(tiled_smem, x, barrier)

from jax._src.pallas import gpu_core
def is_tiled_and_swizzled(transforms):
  return (len(transforms) == 2 and
          isinstance(transforms[0], gpu_core.UnswizzleRef) and
          isinstance(transforms[1], gpu_core.UntilingTransform))

When it happens

Trigger: Calling async_store_smem on an SMEM ref that is not both tiled (untiling transform present) and swizzled, e.g. a plain contiguous SMEM block, or one that is tiled but not swizzled.

Common situations: Writing a custom pipeline (e.g. GEMM epilogues, attention) and forgetting to construct the SMEM ref via the tiled/swizzled helpers; using MMA output layouts that skip the swizzle step.

Related errors

• Swizzle {self.swizzle} is not supported. Only 32, 64 and 128
• Swizzle {self.swizzle} requires the trailing dimension to be
• Can't transpose the swizzled dimension.
• Reshape shape {shape} is not divisible by swizzle elements {
• No valid out swizzle{what}: minor dimension has {minor_dim_b

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