{"record":{"id":"e9a026a3797480ae","repo":"jax-ml/jax","slug":"the-cp-async-implementation-only-supports-oob-mode","errorCode":null,"errorMessage":"The cp.async implementation only supports oob_mode=OOBFillMode.PROMISE_IN_BOUNDS","messagePattern":"The cp\\.async implementation only supports oob_mode=OOBFillMode\\.PROMISE_IN_BOUNDS","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":1011,"sourceCode":"\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(\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)","sourceCodeStart":993,"sourceCodeEnd":1029,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L993-L1029","documentation":"cp.async (the pre-Hopper fallback) does not predicate out-of-bounds accesses, so the API requires oob_mode=OOBFillMode.PROMISE_IN_BOUNDS; any other OOB mode raises ValueError on the cp.async path.","triggerScenarios":"Calling copy_gmem_to_smem with oob_mode=OOBFillMode.ZEROS (or any mode other than PROMISE_IN_BOUNDS) on a GPU with compute capability < 9.","commonSituations":"Defaulting to ZEROS fill for boundary blocks (common in flash-attention style kernels) and running on Ampere or older; writing portable kernels without arch-dependent OOB handling.","solutions":["On pre-Hopper GPUs, use oob_mode=OOBFillMode.PROMISE_IN_BOUNDS and guarantee in-bounds copies (e.g. pad the tensor or mask block iteration)","Gate the kernel so the ZEROS path only runs on Hopper+ (TMA)","Handle boundary blocks manually with predicated loads on older hardware"],"exampleFix":"# before\ncopy_gmem_to_smem(src, smem, oob_mode=OOBFillMode.ZEROS)\n# after\nmode = OOBFillMode.ZEROS if mgpu.utils.get_arch().major >= 9 else OOBFillMode.PROMISE_IN_BOUNDS\ncopy_gmem_to_smem(src, smem, oob_mode=mode)","handlingStrategy":"fallback","validationCode":"from jax._src.pallas.mosaic_gpu import mgpu\nfrom jax._src.pallas.mosaic_gpu.primitives import OOBFillMode\nif mgpu.utils.get_arch().major < 9:\n    oob_mode = OOBFillMode.PROMISE_IN_BOUNDS\n    # ensure blocks stay in bounds: pad tensors or mask iteration","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pad tensors to block multiples so PROMISE_IN_BOUNDS is safe on older GPUs","Branch OOB mode per architecture"],"tags":["jax","pallas","mosaic-gpu","gpu-architecture","out-of-bounds","cp-async"],"backgroundTag":"unsupported-hardware-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}