{"record":{"id":"ffb598de9e3ff51b","repo":"jax-ml/jax","slug":"only-copies-transferring-a-number-of-bytes-divisib","errorCode":null,"errorMessage":"Only copies transferring a number of bytes divisible by the warpgroup size are supported. Got {bytes=} but warpgroup size is {WARPGROUP_SIZE}","messagePattern":"Only copies transferring a number of bytes divisible by the warpgroup size are supported\\. Got (.+?) but warpgroup size is (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":1040,"sourceCode":"\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    ):\n      mgpu.warpgroup_barrier()  # Make sure all reads have completed.\n\n    if not is_cp_async:\n      assert barrier is not None\n      if bytes % WARPGROUP_SIZE:\n        raise NotImplementedError(\n            \"Only copies transferring a number of bytes divisible by the\"\n            f\" warpgroup size are supported. Got {bytes=} but warpgroup size is\"\n            f\" {WARPGROUP_SIZE}\"\n        )\n      if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warpgroup:\n        # We arrive uniformly from each thread in the WG, so we need to divide the\n        # number of bytes by the number of threads in the WG.\n        # TODO: apaszke - Relax this. We can just select the WG leader and have it\n        # arrive with the whole transfer size, while everyone else arrives with 0.\n        # But we should continue using this scheme as it's likely to be faster.\n        bytes //= WARPGROUP_SIZE\n        if predicate is not None:\n          bytes = arith_dialect.select(predicate, mgpu.c(bytes, i32), mgpu.c(0, i32))\n        if is_leader_tracked_copy:\n          first_block = arith_dialect.cmpi(\n              arith_dialect.CmpIPredicate.eq,\n              mgpu.utils.cluster_idx(collective[0]),\n              mgpu.c(0, ir.IndexType.get()),","sourceCodeStart":1022,"sourceCodeEnd":1058,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L1022-L1058","documentation":"On Hopper+ (non-cp.async) copies, Mosaic GPU partitions the copy across the threads of a warpgroup (128 threads), so the total number of bytes moved must be divisible by the warpgroup size (128). If the copy's byte count is not divisible by 128, the lowering raises NotImplementedError because it cannot distribute the work evenly.","triggerScenarios":"A warpgroup-level copy_gmem_to_smem (TMA path, barrier given) where bytes = number_of_elements * dtype_size % 128 != 0, e.g. copying 4 float32 elements (16 bytes) or any odd-sized block.","commonSituations":"Small tail blocks in a tiled kernel (e.g. a 60-element remainder tile), fp8/bfloat8 copies with element counts not a multiple of 128, or block sizes tuned for a different dtype.","solutions":["Pad the copy so total bytes are a multiple of 128 (e.g. round the block shape up and mask out-of-bounds).","Choose block sizes such that num_elements * dtype_bytewidth % 128 == 0 (e.g. multiples of 128 bytes, or 32 f32 elements).","Split the copy into a bulk aligned part plus a scalar fallback path for the remainder."],"exampleFix":"# before\ncopy_gmem_to_smem(src_ref.at[:n], smem_ref)  # n*4 bytes not divisible by 128\n# after\npad_n = (n + 31) // 32 * 32  # f32: 32 elems = 128 bytes\ncopy_gmem_to_smem(src_ref.at[:pad_n].pad(0, (0, pad_n - n)), smem_ref)","handlingStrategy":"validation","validationCode":"BYTES_PER_WG = 128\nnbytes = int(np.prod(block_shape)) * np.dtype(dtype).itemsize\nassert nbytes % BYTES_PER_WG == 0, f'{nbytes=} not divisible by {BYTES_PER_WG}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Parameterize block shapes so total bytes are always a multiple of 128.","Add a host-side assert on bytes before launching the kernel."],"tags":["mosaic-gpu","pallas","alignment","warpgroup","copy"],"backgroundTag":"memory-alignment-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}