jax-ml/jax · error · ValueError

Can only transfer integer bytes (shape={shape}, dtype={dtype

Error message

Can only transfer integer bytes (shape={shape}, dtype={dtype})

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:680

        src,
        ref_smem,
        barrier.as_barrier_memref(),
        gpu_cluster_dim.value,
        cluster_idx_i32,
        atomic_type=atomic_type,
        optimized=optimized,
    )
    return ()

  match remaining_ref_transforms:
    case (gpu_core.UnswizzleRef(swizzle), gpu_core.UntilingTransform(tiling)):
      pass
    case _:
      raise NotImplementedError("async_store_smem requires a tiled and swizzled ref")

  total_bits = math.prod(shape) * dtypes.itemsize_bits(dtype)
  if total_bits % 8:
    raise ValueError(
        f"Can only transfer integer bytes (shape={shape}, dtype={dtype})"
    )
  total_bytes = total_bits // 8
  if total_bytes % WARPGROUP_SIZE:
    raise NotImplementedError(f"Transfer is not a multiple of {WARPGROUP_SIZE} bytes")

  peer_barrier = barrier.remap_to_cluster(gpu_cluster_dim, cluster_idx_val)
  peer_barrier.arrive_expect_tx(total_bytes // WARPGROUP_SIZE)

  lowering._ensure_fa(src, dtype).store_tiled_async(
      ref_smem,
      barrier,
      cluster_dim=gpu_cluster_dim,
      cluster_idx=cluster_idx_val,
      swizzle=swizzle,
      optimized=optimized,
      tiling_rank=len(tiling),
      atomic=atomic,

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

# before
async_store_smem(smem_int4, x_int4, barrier)  # odd element count
# after
async_store_smem(smem_int4, x_int4.reshape(-1, 2), barrier)  # pairs = whole bytes

Defensive patterns

Strategy: validation

Validate before calling

import math
from jax import dtypes
bits = math.prod(shape) * dtypes.itemsize_bits(dtype)
assert bits % 8 == 0, f'transfer of {bits} bits is not byte-aligned'

Prevention

Avoid sub-byte dtypes in async stores unless packed to whole bytes
Prefer int8+ for transfers and unpack in SMEM

When it happens

Trigger: 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).

Common situations: Experimenting with 4-bit quantized weights in Pallas kernels; packing an odd number of sub-byte elements.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/4dad5017bb8301c4. Report an issue: GitHub.