{"record":{"id":"8d7c560505fa171f","repo":"jax-ml/jax","slug":"transfer-is-not-a-multiple-of-warpgroup-size-byt","errorCode":null,"errorMessage":"Transfer is not a multiple of {WARPGROUP_SIZE} bytes","messagePattern":"Transfer is not a multiple of (.+?) bytes","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":685,"sourceCode":"        atomic_type=atomic_type,\n        optimized=optimized,\n    )\n    return ()\n\n  match remaining_ref_transforms:\n    case (gpu_core.UnswizzleRef(swizzle), gpu_core.UntilingTransform(tiling)):\n      pass\n    case _:\n      raise NotImplementedError(\"async_store_smem requires a tiled and swizzled ref\")\n\n  total_bits = math.prod(shape) * dtypes.itemsize_bits(dtype)\n  if total_bits % 8:\n    raise ValueError(\n        f\"Can only transfer integer bytes (shape={shape}, dtype={dtype})\"\n    )\n  total_bytes = total_bits // 8\n  if total_bytes % WARPGROUP_SIZE:\n    raise NotImplementedError(f\"Transfer is not a multiple of {WARPGROUP_SIZE} bytes\")\n\n  peer_barrier = barrier.remap_to_cluster(gpu_cluster_dim, cluster_idx_val)\n  peer_barrier.arrive_expect_tx(total_bytes // WARPGROUP_SIZE)\n\n  lowering._ensure_fa(src, dtype).store_tiled_async(\n      ref_smem,\n      barrier,\n      cluster_dim=gpu_cluster_dim,\n      cluster_idx=cluster_idx_val,\n      swizzle=swizzle,\n      optimized=optimized,\n      tiling_rank=len(tiling),\n      atomic=atomic,\n  )\n  return ()\n\n\ndef async_store_smem(","sourceCodeStart":667,"sourceCodeEnd":703,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L667-L703","documentation":"The async store's total byte count must be a multiple of WARPGROUP_SIZE bytes because arrive_expect_tx and the TMA transaction accounting operate in warpgroup-sized (typically 128-byte) units. Otherwise the lowering raises NotImplementedError.","triggerScenarios":"async_store_smem where prod(shape)*itemsize/8 is not divisible by the warpgroup size (128), e.g. storing a 64-byte block or a small non-conforming tail block.","commonSituations":"Small tail blocks in pipelined loops; choosing block sizes like (48,) float32 that look aligned but aren't multiples of 128 bytes; porting kernels between warp and warpgroup semantics.","solutions":["Round buffer/block sizes up so the transfer is a multiple of 128 bytes (e.g. 32 float32 or 64 bf16 elements)","Pad the SMEM buffer and value to the warpgroup-aligned size","Use the Warp (non-warpgroup) semantics or a synchronous store if small transfers are unavoidable"],"exampleFix":"# before\nasync_store_smem(smem, x[:48], barrier)  # 192 bytes, not multiple of 128\n# after\nasync_store_smem(smem, x[:64], barrier)  # 256 bytes = 2 * WARPGROUP_SIZE","handlingStrategy":"validation","validationCode":"import math\nfrom jax import dtypes\nbytes_ = math.prod(shape) * dtypes.itemsize_bits(dtype) // 8\nassert bytes_ % 128 == 0, f'{bytes_}B not a multiple of warpgroup size (128B)'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Choose block sizes that are multiples of 128 bytes (e.g. 32 fp32 / 64 bf16 elements)","Beware tail blocks in pipelined loops; pad them"],"tags":["jax","pallas","mosaic-gpu","alignment","warpgroup","async-store"],"backgroundTag":"byte-alignment-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}