jax-ml/jax · error · ValueError

Multimem refs are not supported in store_tiled_async

Error message

Multimem refs are not supported in store_tiled_async

What it means

Raised by FragmentedArray.store_tiled_async when the destination reference is a utils.MultimemRef. Async tiled stores lower to PTX st.async which targets shared memory (optionally cluster-wide), and DSMEM/Multimem refs have no async store path in Mosaic, so the code explicitly rejects them.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:3798

        is_signed=self.is_signed,
    )
    fa.store_untiled(ref)

  def store_tiled_async(
      self,
      ref: ir.Value,
      barrier: utils.BarrierRef,
      cluster_dim: gpu.Dimension,
      cluster_idx: ir.Value,
      swizzle: int | None,
      optimized: bool = True,
      tiling_rank: int | None = None,
      atomic: Literal["add", "max", "min", "and", "or", "xor"] | None = None,
  ):
    i32 = ir.IntegerType.get_signless(32)
    i64 = ir.IntegerType.get_signless(64)
    if isinstance(ref, utils.MultimemRef):
      raise ValueError("Multimem refs are not supported in store_tiled_async")
    layout, shape = self.layout, self.shape
    if not isinstance(layout, TiledLayout):
      raise NotImplementedError(self.layout)
    if any(
        isinstance(d, Replicated)
        for d in itertools.chain(layout.warp_dims, layout.lane_dims)
    ):
      raise NotImplementedError("Replicated dimensions are not supported")
    full_cluster_idx: list[ir.Value] = [
        gpu.cluster_block_id(d) for d in gpu.Dimension
    ]
    full_cluster_idx[cluster_dim] = cluster_idx
    lin_cluster_idx = arith.index_cast(
        i32, utils.cluster_idx(tuple(gpu.Dimension), full_cluster_idx)
    )
    cluster_barrier_ptr = utils.get_cluster_ptr(
        barrier.get_ptr(), lin_cluster_idx, generic=False
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use store_tiled or a non-async store path for MultimemRef destinations
2. Construct the destination as a regular shared-memory reference (utils.smem_ref / SharedMemory) and pass that to store_tiled_async
3. If cluster multicast is needed, rely on the cluster_barrier_ptr/cluster plumbing of store_tiled_async rather than a MultimemRef

Example fix

// before
fa.store_tiled_async(multimem_ref, ..., cluster_barrier_ptr=bar)
// after
smem_ref = utils.smem_ref(shape, dtype)
fa.store_tiled_async(smem_ref, ..., cluster_barrier_ptr=bar)

from jax.experimental.mosaic.gpu import utils
if isinstance(ref, utils.MultimemRef):
    raise TypeError('use store_tiled for MultimemRef destinations')
fa.store_tiled_async(ref, ...)

def is_async_store_ref(ref) -> bool:
    return not isinstance(ref, utils.MultimemRef)

try:
    fa.store_tiled_async(ref, ...)
except ValueError as e:
    if 'Multimem refs' in str(e):
        fa.store_tiled(ref)  # fallback
    else:
        raise

When it happens

Trigger: Calling store_tiled_async(ref, ...) where ref was built with utils.MultimemRef (e.g. a TMA/multicast tensor-memory map descriptor) instead of a plain shared-memory reference.

Common situations: Porting a Hopper/Blackwell TMA kernel that used multimem load/stores to the async cluster-store API; passing a reference obtained from a tma map with multimem enabled into the new async barrier-based pipeline.

Related errors

• f16/bf16 SMEM/multimem atomics only support add, got {atomic
• Replicated dimensions are not supported
• Unsupported register bitwidth: {reg_bitwidth}
• f32 not supported for async atomics
• f32 only supports add atomics, got {atomic}

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