jax-ml/jax · error · ValueError

Async store only supports TMEM refs

Error message

Async store only supports TMEM refs

What it means

The async_store_tmem primitive only stores into Tensor Memory references. A destination ref whose memory_space is not TMEM raises ValueError in its abstract eval.

Source

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

  Args:
    ref: The TMEM reference to store to.
    value: The value to store.
  """
  ref, ref_transforms = state_primitives.get_ref_and_transforms(
      ref, None, "async_store_tmem"
  )
  flat_ref_transforms, ref_transforms_treedef = tree_util.tree_flatten(
      ref_transforms
  )
  async_store_tmem_p.bind(
      ref, value, *flat_ref_transforms, tree=ref_transforms_treedef
  )

@async_store_tmem_p.def_effectful_abstract_eval
def _async_store_tmem_abstract_eval(ref, val, *avals_flat, tree):
  if ref.memory_space != gpu_core.MemorySpace.TMEM:
    raise ValueError("Async store only supports TMEM refs")
  _, effects = state_primitives._swap_abstract_eval(
      ref, val, *avals_flat, tree=tree
  )
  return (), effects

@lowering.register_lowering_rule(async_store_tmem_p, mgpu.LoweringSemantics.Lane)
def _async_store_tmem_lowering_rule(
    ctx: lowering.LoweringRuleContext, x_ref, value, *leaves, tree
):
  assert isinstance(x_ref, tcgen05.TMEMRef)
  x_aval = ctx.avals_in[0]
  assert isinstance(x_aval, state_types.AbstractRef)
  transforms = jax.tree.unflatten(tree, leaves)
  transform_avals = tree.unflatten(
      ctx.avals_in[2 : 2 + tree.num_leaves]
  )
  x_tmem, _, transforms = lowering._handle_transforms(
      ctx, x_aval, x_ref, transform_avals, transforms, handle_transposes=False,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Store into a TMEM-allocated ref
2. Use regular async_copy or async_store for non-TMEM destinations

Example fix

// before
async_store_tmem(smem_ref, val)
// after
async_store_tmem(tmem_ref, val)

assert ref.memory_space == gpu_core.MemorySpace.TMEM

def is_tmem(ref): return getattr(ref, 'memory_space', None) == gpu_core.MemorySpace.TMEM

When it happens

Trigger: Calling the async TMEM store API with an SMEM or HBM ref as destination.

Common situations: Storing results of a TMEM load back to the wrong buffer; copy-paste errors between load/store kernel templates.

Related errors

• Async load only supports TMEM refs
• Stores to TMEM are asynchronous operations and cannot be per
• Accumulator must be a TMEM Ref.
• LHS must be a TMEM/SMEM Ref.
• Memory space {self.memory_space} is not supported by mesh {s

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