{"record":{"id":"4dad5017bb8301c4","repo":"jax-ml/jax","slug":"can-only-transfer-integer-bytes-shape-shape-dt","errorCode":null,"errorMessage":"Can only transfer integer bytes (shape={shape}, dtype={dtype})","messagePattern":"Can only transfer integer bytes \\(shape=(.+?), dtype=(.+?)\\)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":680,"sourceCode":"        src,\n        ref_smem,\n        barrier.as_barrier_memref(),\n        gpu_cluster_dim.value,\n        cluster_idx_i32,\n        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,","sourceCodeStart":662,"sourceCodeEnd":698,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L662-L698","documentation":"Before issuing the async store, the lowering computes total transfer size in bits (prod(shape) * itemsize_bits) and requires it to be byte-aligned (divisible by 8). Sub-byte element types (e.g. 4-bit integers) whose total bit count isn't a multiple of 8 raise ValueError.","triggerScenarios":"async_store_smem with a dtype like int4 or a custom sub-byte type where prod(shape)*bits % 8 != 0, e.g. a single int4 element (4 bits).","commonSituations":"Experimenting with 4-bit quantized weights in Pallas kernels; packing an odd number of sub-byte elements.","solutions":["Pack sub-byte elements so the total bit count is a multiple of 8 (e.g. store int4 values in pairs)","Use a wider dtype (int8) for the transfer and unpack afterwards","Ensure the block shape times itemsize yields whole bytes"],"exampleFix":"# before\nasync_store_smem(smem_int4, x_int4, barrier)  # odd element count\n# after\nasync_store_smem(smem_int4, x_int4.reshape(-1, 2), barrier)  # pairs = whole bytes","handlingStrategy":"validation","validationCode":"import math\nfrom jax import dtypes\nbits = math.prod(shape) * dtypes.itemsize_bits(dtype)\nassert bits % 8 == 0, f'transfer of {bits} bits is not byte-aligned'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid sub-byte dtypes in async stores unless packed to whole bytes","Prefer int8+ for transfers and unpack in SMEM"],"tags":["jax","pallas","mosaic-gpu","sub-byte-dtype","alignment","async-store"],"backgroundTag":"byte-alignment-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}