jax-ml/jax · error · ValueError

Only 32-, 64- and 128-bit stores are supported

Error message

Only 32-, 64- and 128-bit stores are supported

What it means

Raised by multimem_store, which lowers to the multimem.st PTX instruction for distributed shared memory (DSMEM) multicast stores. The hardware instruction only accepts 32-, 64- and 128-bit stores (v1.f32/v2.f32/v4.f32 equivalents), so wider or narrower values (e.g. a 256-bit vector or an i8) are rejected.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:241

@dataclasses.dataclass(frozen=True)
class MultimemRef:
  ref: ir.Value[ir.MemRefType]

  @property
  def type(self) -> ir.Type:
    return ir.MemRefType(self.ref.type)

  def store(self, value: ir.Value, indices: Sequence[ir.Value]):
    ptr = memref_ptr(memref_slice(self.ref, tuple(indices)))
    multimem_store(ptr, value)


def multimem_store(ptr: ir.Value, value: ir.Value):
  i32 = ir.IntegerType.get_signless(32)
  if (bw := bitwidth(value.type)) not in {32, 64, 128}:
    raise ValueError("Only 32-, 64- and 128-bit stores are supported")
  vector_length = bw // 32
  value = bitcast(value, ir.VectorType.get((vector_length,), i32))
  regs = [
      llvm.extractelement(value, arith.constant(i32, i))
      for i in range(vector_length)
  ]
  if vector_length == 1:
    vec_ptx = "$1"
    vec_mod = ""
  else:
    vec_ptx = f"{{{','.join(f'${i}' for i in range(1, vector_length + 1))}}}"
    vec_mod = ".v" + str(vector_length)
  # It's unclear to me why, but at least according to PTX docs, we have to use
  # the floating-point instructions here to be able to store vectors.
  llvm.inline_asm(
      ir.Type.parse("!llvm.void"),
      [ptr, *regs],
      f"multimem.st.relaxed.sys.global{vec_mod}.f32 [$0], {vec_ptx};",

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Chunk the value into <=128-bit pieces: for vectors, store in slices of at most 4xi32 (128 bits) each
2. Pack narrow elements (i8/i16/f8) into i32 or vector<Nxi32> lanes before the store
3. For 256-bit vectors, issue two 128-bit multimem_store calls on consecutive pointers

Example fix

# before
multimem_store(ptr, vec_v8i32)  # 256-bit -> error
# after
for i in range(0, 8, 4):
  chunk = vector.extract_strided_slice(vec_v8i32, offset=i, size=4, stride=1)
  multimem_store(utils.getelementptr(ptr, [i], i32), chunk)

bw = utils.bitwidth(value.type)
if bw > 128:
    raise ValueError('split value into <=128-bit chunks before multimem_store')
if bw < 32:
    raise ValueError('pack narrow elements into 32-bit lanes before multimem_store')

def is_multimem_storable(value):
    return utils.bitwidth(value.type) in (32, 64, 128)

When it happens

Trigger: Calling utils.multimem_store(ptr, value) where bitwidth(value.type) is not 32/64/128 — e.g. storing vector<8xi32> (256 bits), a bare i16, or an unpacked f8 value. Reached via store/store_tiled/store_untiled and the transfer loop in distributed kernels.

Common situations: Widening SMEM transfer tiles for throughput until they exceed 128 bits; storing narrow dtypes without packing them into 32-bit lanes first.

Related errors

• Unsupported memory space: {ref_type.memory_space}
• Multimem refs are not supported in store_tiled_async
• f16/bf16 SMEM/multimem atomics only support add, got {atomic
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp
• {op} has an unsupported layout: {out_layout_attr}

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