jax-ml/jax · error · RuntimeError

Stores to TMEM are asynchronous operations and cannot be per

Error message

Stores to TMEM are asynchronous operations and cannot be performed using the usual syntax. Please use plgpu.async_store_tmem instead.

What it means

Raised when swap (the x_ref[...] = v / load-old-then-store pattern) targets a TMEM (tcgen05 tensor memory) reference. TMEM stores are asynchronous on Blackwell and must use plgpu.async_store_tmem; the synchronous swap syntax cannot return the old value.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:2257

    raise NotImplementedError(
        "Transforms are not yet implemented for warpgroup semantics"
    )

  assert isinstance(x_ref, ir.Value)
  shape = ctx.avals_out[0].shape
  if shape:
    return mgpu.dialect.vector_load(x_ref, optimized=optimized)
  else:
    return memref_dialect.load(x_ref, [])


@register_lowering_rule(sp.swap_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(sp.swap_p, *gpu_core.LANExWARP_SEMANTICS)
def _swap_lowering_rule(
    ctx: LoweringRuleContext, x_ref, value, *leaves, tree
):
  if isinstance(x_ref, tcgen05.TMEMRef):
    raise RuntimeError(
        "Stores to TMEM are asynchronous operations and cannot be performed"
        " using the usual syntax. Please use plgpu.async_store_tmem instead."
    )
  barrier = mgpu.warpgroup_barrier
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    if ctx.avals_out[0].shape:
      raise NotImplementedError("Can only store scalars in warp-level lowering.")
    i32 = ir.IntegerType.get_signless(32)
    barrier = functools.partial(
        nvvm_dialect.bar_warp_sync, arith_dialect.constant(i32, -1)
    )
  value = _ensure_fa(value, ctx.avals_in[1].dtype)

  if not isinstance(x_ref, ir.Value) and isinstance(x_ref, ir.MemRefType):
    raise TypeError(f"Can only store to references (got {x_ref}).")
  v_aval = ctx.avals_in[1]
  transforms = jax.tree.unflatten(tree, leaves)
  transform_avals = jax.tree.unflatten(tree, ctx.avals_in[2:])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use plgpu.async_store_tmem(...) for the store and commit/wait as needed
Keep accumulators in TMEM and only read them with tmem load; do not swap through TMEM refs

Example fix

# before
old = tmem_ref[:, :] = acc
# after
plgpu.async_store_tmem(tmem_ref, acc)
plgpu.commit()  # then plgpu.wait()

Defensive patterns

Strategy: type-guard

Validate before calling

from jax._src.pallas.mosaic_gpu import tcgen05
assert not isinstance(ref_obj, tcgen05.TMEMRef), 'use async_store_tmem'

Type guard

def is_tmem_ref(r) -> bool:
    return type(r).__name__ == 'TMEMRef'

Try / catch

try: store... except RuntimeError as e: if 'async_store_tmem' in str(e): plgpu.async_store_tmem(...)

Prevention

Treat TMEM stores as async: always use plgpu.async_store_tmem
Never swap through TMEM refs

When it happens

Trigger: Doing `old = tmem_ref[...] = value` (swap_p) where the ref is a tcgen05.TMEMRef, i.e. tensor-core accumulator memory on SM100.

Common situations: Writing Blackwell tcgen05 MMA kernels where the accumulator lives in TMEM and attempting normal assignment inside the kernel body.

Related errors

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