{"record":{"id":"cc9e1246a5f71f6e","repo":"jax-ml/jax","slug":"the-cp-async-implementation-does-not-support-user","errorCode":null,"errorMessage":"The cp.async implementation does not support user-defined predicates","messagePattern":"The cp\\.async implementation does not support user-defined predicates","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":1016,"sourceCode":"      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(\n          \"The cp.async implementation does not support user-defined predicates\"\n      )\n  else:\n    if oob_mode is None:\n      oob_mode = OOBFillMode.ZEROS\n\n    if barrier is None:\n      raise ValueError(\n          \"copy_gmem_to_smem without a barrier is only supported on pre-Hopper\"\n          \" GPUs, which use the cp.async implementation\"\n      )\n\n  i32 = ir.IntegerType.get_signless(32)\n  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:\n    if (\n        ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warpgroup\n        and ctx.module_ctx.auto_barriers\n    ):","sourceCodeStart":998,"sourceCodeEnd":1034,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L998-L1034","documentation":"The cp.async fallback does not support user-supplied predicates (has_user_predicate); on pre-Hopper GPUs passing a predicate to copy_gmem_to_smem raises NotImplementedError. Predicated copies are only available via TMA on Hopper+.","triggerScenarios":"Calling copy_gmem_to_smem(..., predicate=...) on a GPU with compute capability < 9.","commonSituations":"Using predicates to skip boundary blocks in pipelined loops; kernels written for TMA being run on A100 or older cards.","solutions":["Remove the predicate on pre-Hopper GPUs and instead mask/skip the copy at the kernel level (e.g. conditionally execute the copy op)","Use predicated element-wise loads (regular pl.load with a mask) on older hardware","Gate the predicated fast path on get_arch().major >= 9"],"exampleFix":"# before\ncopy_gmem_to_smem(src, smem, predicate=pred)\n# after\nif mgpu.utils.get_arch().major >= 9:\n  copy_gmem_to_smem(src, smem, predicate=pred)\nelif pred:\n  copy_gmem_to_smem(src, smem)","handlingStrategy":"fallback","validationCode":"from jax._src.pallas.mosaic_gpu import mgpu\nif mgpu.utils.get_arch().major < 9:\n    if not predicate:  # emulate predicated copy\n        skip_copy = True  # caller must skip the copy op entirely","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Wrap predicated copies in an `if pred:` block on pre-Hopper GPUs","Use masked pl.load as the portable alternative"],"tags":["jax","pallas","mosaic-gpu","gpu-architecture","predicate","cp-async","not-implemented"],"backgroundTag":"unsupported-hardware-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}