{"record":{"id":"331c735d211871d9","repo":"jax-ml/jax","slug":"async-copies-require-the-number-of-bits-copied-alo","errorCode":null,"errorMessage":"Async copies require the number of bits copied along the last dimension to be divisible by 128, but got {zeroth_bw}","messagePattern":"Async copies require the number of bits copied along the last dimension to be divisible by 128, but got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/launch_context.py","lineNumber":1187,"sourceCode":"        elif rem_collective_size % slice_size == 0:\n          # This is an optimization and it lets us skip squeezed dims.\n          if slice_size > 1:\n            dim_idx = arith.remui(idx, c(slice_size, index))\n            partition_dim(dim, dim_idx, slice_size)\n            idx = arith.divui(idx, c(slice_size, index))\n            rem_collective_size //= slice_size\n        else:\n          break  # We failed to partition the leading dimensions.\n      del idx  # We overwrote the block index in the loop.\n      if rem_collective_size > 1:\n        raise ValueError(\n            \"None of the leading dimensions in the transformed slice shape\"\n            f\" {slice_shape} is divisible by the collective size\"\n            f\" {collective_size}\"\n        )\n\n    if (zeroth_bw := slice_shape[-1] * element_bitwidth) % 128 != 0:\n      raise ValueError(\n          \"Async copies require the number of bits copied along the last\"\n          f\" dimension to be divisible by 128, but got {zeroth_bw}\"\n      )\n    if (\n        swizzle is not None\n        and swizzle != mgpu_dialect.SwizzlingMode.kNoSwizzle\n        and slice_shape[-1] != (swizzle * 8) // element_bitwidth\n    ):\n      raise ValueError(\n          f\"Async copies with {swizzle=} require the last dimension of the\"\n          f\" slice to be exactly {swizzle} bytes i.e. \"\n          f\" {(swizzle * 8) // element_bitwidth} elements, but got\"\n          f\" {slice_shape[-1]} elements.\"\n      )\n    return (smem_ref, slice_shape, dyn_base_indices, gmem_transform)\n\n  def async_copy(\n      self,","sourceCodeStart":1169,"sourceCodeEnd":1205,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/launch_context.py#L1169-L1205","documentation":"TMA requires the innermost copied dimension to be at least 16 bytes and a multiple of 16 bytes; Mosaic enforces this by checking slice_shape[-1] * element_bitwidth % 128 == 0. Violating this raises the error with the offending bitwidth.","triggerScenarios":"Calling async_copy/async_prefetch with the TMA implementation where the last dimension extent times element size is not a multiple of 16 bytes, e.g. copying a last dim of 3 float32 elements (96 bits).","commonSituations":"Small or odd innermost tile sizes (e.g. seq length 12 with f32); using narrow element types (int8/f16) with dims not scaled to reach 16-byte multiples.","solutions":["Round the last-dimension slice up (pad) so last_dim * element_bytes is a multiple of 16.","Choose tile sizes that are multiples of 16 bytes: e.g. multiples of 4 for f32, 8 for f16/bf16, 16 for int8.","If padding is unacceptable, avoid the TMA implementation for this copy."],"exampleFix":"// before\nctx.async_copy(..., gmem_slice=(slice(0, 64), slice(0, 3)))  # f32: 12 bytes\n// after\nctx.async_copy(..., gmem_slice=(slice(0, 64), slice(0, 4)))  # f32: 16 bytes","handlingStrategy":"validation","validationCode":"last_bw = slice_shape[-1] * utils.bitwidth(element_type)\nassert last_bw % 128 == 0, f'last dim must be a multiple of 16 bytes, got {last_bw} bits'","typeGuard":null,"tryCatchPattern":"try:\n    ctx.async_copy(...)\nexcept ValueError as e:\n    if 'divisible by 128' in str(e):\n        slice_shape[-1] = _round_up_to_bytes(slice_shape[-1], 16, element_type)\n    else:\n        raise","preventionTips":["Choose inner tile sizes that are multiples of 16 bytes given the dtype.","For f32 use multiples of 4; f16/int8 multiples of 8/16.","Remember swizzle sizes (32/64/128B) all satisfy this if the last dim matches swizzle width."],"tags":["jax","mosaic-gpu","tma","memory-alignment","shape-validation"],"backgroundTag":"memory-alignment-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}