{"record":{"id":"317fd9b4b15cc9bc","repo":"jax-ml/jax","slug":"stores-to-tmem-are-asynchronous-operations-and-can","errorCode":null,"errorMessage":"Stores to TMEM are asynchronous operations and cannot be performed using the usual syntax. Please use plgpu.async_store_tmem instead.","messagePattern":"Stores to TMEM are asynchronous operations and cannot be performed using the usual syntax\\. Please use plgpu\\.async_store_tmem instead\\.","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":2257,"sourceCode":"    raise NotImplementedError(\n        \"Transforms are not yet implemented for warpgroup semantics\"\n    )\n\n  assert isinstance(x_ref, ir.Value)\n  shape = ctx.avals_out[0].shape\n  if shape:\n    return mgpu.dialect.vector_load(x_ref, optimized=optimized)\n  else:\n    return memref_dialect.load(x_ref, [])\n\n\n@register_lowering_rule(sp.swap_p, mgpu.LoweringSemantics.Lane)\n@register_lowering_rule(sp.swap_p, *gpu_core.LANExWARP_SEMANTICS)\ndef _swap_lowering_rule(\n    ctx: LoweringRuleContext, x_ref, value, *leaves, tree\n):\n  if isinstance(x_ref, tcgen05.TMEMRef):\n    raise RuntimeError(\n        \"Stores to TMEM are asynchronous operations and cannot be performed\"\n        \" using the usual syntax. Please use plgpu.async_store_tmem instead.\"\n    )\n  barrier = mgpu.warpgroup_barrier\n  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:\n    if ctx.avals_out[0].shape:\n      raise NotImplementedError(\"Can only store scalars in warp-level lowering.\")\n    i32 = ir.IntegerType.get_signless(32)\n    barrier = functools.partial(\n        nvvm_dialect.bar_warp_sync, arith_dialect.constant(i32, -1)\n    )\n  value = _ensure_fa(value, ctx.avals_in[1].dtype)\n\n  if not isinstance(x_ref, ir.Value) and isinstance(x_ref, ir.MemRefType):\n    raise TypeError(f\"Can only store to references (got {x_ref}).\")\n  v_aval = ctx.avals_in[1]\n  transforms = jax.tree.unflatten(tree, leaves)\n  transform_avals = jax.tree.unflatten(tree, ctx.avals_in[2:])","sourceCodeStart":2239,"sourceCodeEnd":2275,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L2239-L2275","documentation":"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.","triggerScenarios":"Doing `old = tmem_ref[...] = value` (swap_p) where the ref is a tcgen05.TMEMRef, i.e. tensor-core accumulator memory on SM100.","commonSituations":"Writing Blackwell tcgen05 MMA kernels where the accumulator lives in TMEM and attempting normal assignment inside the kernel body.","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"],"exampleFix":"# before\nold = tmem_ref[:, :] = acc\n# after\nplgpu.async_store_tmem(tmem_ref, acc)\nplgpu.commit()  # then plgpu.wait()","handlingStrategy":"type-guard","validationCode":"from jax._src.pallas.mosaic_gpu import tcgen05\nassert not isinstance(ref_obj, tcgen05.TMEMRef), 'use async_store_tmem'","typeGuard":"def is_tmem_ref(r) -> bool:\n    return type(r).__name__ == 'TMEMRef'","tryCatchPattern":"try: store... except RuntimeError as e: if 'async_store_tmem' in str(e): plgpu.async_store_tmem(...)","preventionTips":["Treat TMEM stores as async: always use plgpu.async_store_tmem","Never swap through TMEM refs"],"tags":["jax","pallas","mosaic-gpu","tmem","blackwell","tcgen05"],"backgroundTag":"async-gpu-memory-operation-misuse","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}