{"record":{"id":"77377e09bc8bb14b","repo":"jax-ml/jax","slug":"only-32-64-and-128-bit-stores-are-supported","errorCode":null,"errorMessage":"Only 32-, 64- and 128-bit stores are supported","messagePattern":"Only 32-, 64- and 128-bit stores are supported","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":241,"sourceCode":"\n\n@dataclasses.dataclass(frozen=True)\nclass MultimemRef:\n  ref: ir.Value[ir.MemRefType]\n\n  @property\n  def type(self) -> ir.Type:\n    return ir.MemRefType(self.ref.type)\n\n  def store(self, value: ir.Value, indices: Sequence[ir.Value]):\n    ptr = memref_ptr(memref_slice(self.ref, tuple(indices)))\n    multimem_store(ptr, value)\n\n\ndef multimem_store(ptr: ir.Value, value: ir.Value):\n  i32 = ir.IntegerType.get_signless(32)\n  if (bw := bitwidth(value.type)) not in {32, 64, 128}:\n    raise ValueError(\"Only 32-, 64- and 128-bit stores are supported\")\n  vector_length = bw // 32\n  value = bitcast(value, ir.VectorType.get((vector_length,), i32))\n  regs = [\n      llvm.extractelement(value, arith.constant(i32, i))\n      for i in range(vector_length)\n  ]\n  if vector_length == 1:\n    vec_ptx = \"$1\"\n    vec_mod = \"\"\n  else:\n    vec_ptx = f\"{{{','.join(f'${i}' for i in range(1, vector_length + 1))}}}\"\n    vec_mod = \".v\" + str(vector_length)\n  # It's unclear to me why, but at least according to PTX docs, we have to use\n  # the floating-point instructions here to be able to store vectors.\n  llvm.inline_asm(\n      ir.Type.parse(\"!llvm.void\"),\n      [ptr, *regs],\n      f\"multimem.st.relaxed.sys.global{vec_mod}.f32 [$0], {vec_ptx};\",","sourceCodeStart":223,"sourceCodeEnd":259,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L223-L259","documentation":"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.","triggerScenarios":"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.","commonSituations":"Widening SMEM transfer tiles for throughput until they exceed 128 bits; storing narrow dtypes without packing them into 32-bit lanes first.","solutions":["Chunk the value into <=128-bit pieces: for vectors, store in slices of at most 4xi32 (128 bits) each","Pack narrow elements (i8/i16/f8) into i32 or vector<Nxi32> lanes before the store","For 256-bit vectors, issue two 128-bit multimem_store calls on consecutive pointers"],"exampleFix":"# before\nmultimem_store(ptr, vec_v8i32)  # 256-bit -> error\n# after\nfor i in range(0, 8, 4):\n  chunk = vector.extract_strided_slice(vec_v8i32, offset=i, size=4, stride=1)\n  multimem_store(utils.getelementptr(ptr, [i], i32), chunk)","handlingStrategy":"validation","validationCode":"bw = utils.bitwidth(value.type)\nif bw > 128:\n    raise ValueError('split value into <=128-bit chunks before multimem_store')\nif bw < 32:\n    raise ValueError('pack narrow elements into 32-bit lanes before multimem_store')","typeGuard":"def is_multimem_storable(value):\n    return utils.bitwidth(value.type) in (32, 64, 128)","tryCatchPattern":null,"preventionTips":["Cap DSMEM transfer vectors at 4xi32 (128 bits)","Pack sub-32-bit elements into i32 lanes before multimem stores"],"tags":["gpu","mosaic","dsmem","multimem","store","bit-width"],"backgroundTag":"unsupported-bit-width","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}