{"record":{"id":"60ff3ab994222d83","repo":"jax-ml/jax","slug":"wait-gmem-to-smem-is-only-supported-on-pre-hopper","errorCode":null,"errorMessage":"wait_gmem_to_smem is only supported on pre-Hopper GPUs, which use cp.async for GMEM->SMEM copies.","messagePattern":"wait_gmem_to_smem is only supported on pre-Hopper GPUs, which use cp\\.async for GMEM->SMEM copies\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":1346,"sourceCode":"\n\n@wait_gmem_to_smem_p.def_effectful_abstract_eval\ndef _wait_gmem_to_smem_abstract_eval(n):\n  del n  # Unused.\n  return (), {gpu_core._memory_effect}\n\n\n@lowering.register_lowering_rule(\n    wait_gmem_to_smem_p, mgpu.LoweringSemantics.Lane)\n@lowering.register_lowering_rule(\n    wait_gmem_to_smem_p, *gpu_core.LANExWARP_SEMANTICS)\n@lowering.register_lowering_rule(\n    wait_gmem_to_smem_p, mgpu.LoweringSemantics.Warpgroup)\n@lowering.register_lowering_rule(\n    wait_gmem_to_smem_p, *gpu_core.WGxWARP_SEMANTICS)\ndef _wait_gmem_to_smem_lowering(ctx: lowering.LoweringRuleContext, n):\n  if mgpu.utils.get_arch().major >= 9:\n    raise ValueError(\n        \"wait_gmem_to_smem is only supported on pre-Hopper GPUs, which use\"\n        \" cp.async for GMEM->SMEM copies.\"\n    )\n  ctx.launch_ctx.await_cp_async_copy(allow_groups=n)\n  return ()\n\n\ndef wait_gmem_to_smem(n: int) -> None:\n  \"\"\"Waits until at most ``n`` ``cp.async`` GMEM->SMEM copies are in flight.\n\n  .. note:: This waiting mechanism is only supported on pre-Hopper GPUs, which\n            use the ``cp.async`` implementation of\n            :func:`jax.experimental.pallas.mosaic_gpu.copy_gmem_to_smem`.\n\n  Args:\n    n: The maximum number of copies allowed to remain in flight.\n  \"\"\"\n  wait_gmem_to_smem_p.bind(n)","sourceCodeStart":1328,"sourceCodeEnd":1364,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L1328-L1364","documentation":"wait_gmem_to_smem waits for completion of cp.async copies, which are only used on pre-Hopper (SM < 9) GPUs. On Hopper+ the GMEM->SMEM copies go through TMA with barriers, so there is no cp.async group to wait on and the call is invalid.","triggerScenarios":"Calling wait_gmem_to_smem in a kernel running on a Hopper (SM90) or newer GPU.","commonSituations":"A100-era kernels with explicit cp.async waits (copy without barrier then wait_gmem_to_smem) run on H100/Blackwell after a cluster migration.","solutions":["Replace copy-without-barrier + wait_gmem_to_smem with an explicit barrier: pass barrier=barrier to copy_gmem_to_smem and call barrier.wait().","Use copy_in / higher-level helpers which select the right synchronization per architecture."],"exampleFix":"# before\ncopy_gmem_to_smem(src_ref, dst_ref)\nwait_gmem_to_smem()\n# after\nbarrier = mgpu.SMEM_ALLOCATOR.get_buffer((1,), mgpu.BarrierType, barrier_init=1)\ncopy_gmem_to_smem(src_ref, dst_ref, barrier=barrier)\nbarrier.wait()","handlingStrategy":"fallback","validationCode":"from jax._src.pallas.mosaic_gpu import mgpu\npre_hopper = mgpu.utils.get_arch().major < 9\nif pre_hopper:\n    copy_gmem_to_smem(src, dst); wait_gmem_to_smem()\nelse:\n    barrier = alloc_barrier(); copy_gmem_to_smem(src, dst, barrier=barrier); barrier.wait()","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use barriers unconditionally; they also work on pre-Hopper.","Avoid cp.async-specific APIs in kernels targeted at multiple GPU generations."],"tags":["mosaic-gpu","pallas","gpu","hopper","synchronization"],"backgroundTag":"gpu-barrier-required","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}