{"record":{"id":"5f228f2b3248b44d","repo":"jax-ml/jax","slug":"copy-gmem-to-smem-with-a-barrier-is-only-supported","errorCode":null,"errorMessage":"copy_gmem_to_smem with a barrier is only supported Hopper and newer GPUs, which use the TMA implementation","messagePattern":"copy_gmem_to_smem with a barrier is only supported Hopper and newer GPUs, which use the TMA implementation","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":998,"sourceCode":"\n  if is_leader_tracked_copy:\n    # Leader receives the completion messages from both CTAs.\n    bytes *= 2\n    if len(collective) != 1:\n      raise ValueError(\n          f\"Expected exactly one collective axis, got {collective_axes=}\"\n      )\n    if math.prod(ctx.launch_ctx.cluster_size) != 2:\n      raise NotImplementedError(\n          \"Partitioned loads only supported for clusters of size 2. Got\"\n          f\" cluster size {ctx.launch_ctx.cluster_size}.\"\n      )\n\n  # TMA is only available on Hopper and newer. On older architectures we fall\n  # back to the cp.async implementation.\n  if is_cp_async := mgpu.utils.get_arch().major < 9:\n    if barrier is not None:\n      raise ValueError(\n          \"copy_gmem_to_smem with a barrier is only supported Hopper and newer\"\n          \" GPUs, which use the TMA implementation\"\n      )\n    if collective_axes is not None:\n      raise ValueError(\"Only the TMA implementation supports collective copies\")\n    if leader_tracked is not None:\n      raise ValueError(\n          \"Only the TMA implementation supports leader_tracked copies\"\n      )\n    # cp.async does not predicate out-of-bounds accesses, so the caller has to\n    # guarantee that the copy stays in bounds.\n    if oob_mode != OOBFillMode.PROMISE_IN_BOUNDS:\n      raise ValueError(\n          \"The cp.async implementation only supports \"\n          \"oob_mode=OOBFillMode.PROMISE_IN_BOUNDS\"\n      )\n    if has_user_predicate:\n      raise NotImplementedError(","sourceCodeStart":980,"sourceCodeEnd":1016,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L980-L1016","documentation":"On GPUs older than Hopper (compute capability < 9, is_cp_async=True), the copy falls back to cp.async, which has no mbarrier support. Passing a barrier to copy_gmem_to_smem on such hardware raises ValueError because only the Hopper+ TMA implementation supports barriers.","triggerScenarios":"Running copy_gmem_to_smem(..., barrier=...) on Ampere (SM80) or older GPUs where the TMA path is unavailable.","commonSituations":"Developing on/for older GPUs (A100, V100) or multi-GPU fleets with mixed generations; code written against Hopper TMA tutorials being run elsewhere.","solutions":["Guard barrier usage on compute capability: only pass a barrier when mgpu.utils.get_arch().major >= 9","Use a non-barrier copy plus explicit synchronization on pre-Hopper GPUs","Target Hopper (H100/H200/B100) hardware for barrier-based pipelines"],"exampleFix":"# before\ncopy_gmem_to_smem(src, smem, barrier=bar)\n# after\nuse_tma = mgpu.utils.get_arch().major >= 9\ncopy_gmem_to_smem(src, smem, barrier=bar if use_tma else None)","handlingStrategy":"validation","validationCode":"from jax._src.pallas.mosaic_gpu import mgpu\nuse_tma = mgpu.utils.get_arch().major >= 9\nassert use_tma or barrier is None, 'barrier copies need Hopper+'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Query compute capability once per kernel and branch on it","Make mbarrier-based pipelines Hopper-gated"],"tags":["jax","pallas","mosaic-gpu","gpu-architecture","tma","barrier","cp-async"],"backgroundTag":"unsupported-hardware-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}