{"record":{"id":"7b676f607fa8c822","repo":"jax-ml/jax","slug":"multimem-refs-are-not-supported-in-store-tiled-asy","errorCode":null,"errorMessage":"Multimem refs are not supported in store_tiled_async","messagePattern":"Multimem refs are not supported in store_tiled_async","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":3798,"sourceCode":"        is_signed=self.is_signed,\n    )\n    fa.store_untiled(ref)\n\n  def store_tiled_async(\n      self,\n      ref: ir.Value,\n      barrier: utils.BarrierRef,\n      cluster_dim: gpu.Dimension,\n      cluster_idx: ir.Value,\n      swizzle: int | None,\n      optimized: bool = True,\n      tiling_rank: int | None = None,\n      atomic: Literal[\"add\", \"max\", \"min\", \"and\", \"or\", \"xor\"] | None = None,\n  ):\n    i32 = ir.IntegerType.get_signless(32)\n    i64 = ir.IntegerType.get_signless(64)\n    if isinstance(ref, utils.MultimemRef):\n      raise ValueError(\"Multimem refs are not supported in store_tiled_async\")\n    layout, shape = self.layout, self.shape\n    if not isinstance(layout, TiledLayout):\n      raise NotImplementedError(self.layout)\n    if any(\n        isinstance(d, Replicated)\n        for d in itertools.chain(layout.warp_dims, layout.lane_dims)\n    ):\n      raise NotImplementedError(\"Replicated dimensions are not supported\")\n    full_cluster_idx: list[ir.Value] = [\n        gpu.cluster_block_id(d) for d in gpu.Dimension\n    ]\n    full_cluster_idx[cluster_dim] = cluster_idx\n    lin_cluster_idx = arith.index_cast(\n        i32, utils.cluster_idx(tuple(gpu.Dimension), full_cluster_idx)\n    )\n    cluster_barrier_ptr = utils.get_cluster_ptr(\n        barrier.get_ptr(), lin_cluster_idx, generic=False\n    )","sourceCodeStart":3780,"sourceCodeEnd":3816,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L3780-L3816","documentation":"Raised by FragmentedArray.store_tiled_async when the destination reference is a utils.MultimemRef. Async tiled stores lower to PTX st.async which targets shared memory (optionally cluster-wide), and DSMEM/Multimem refs have no async store path in Mosaic, so the code explicitly rejects them.","triggerScenarios":"Calling store_tiled_async(ref, ...) where ref was built with utils.MultimemRef (e.g. a TMA/multicast tensor-memory map descriptor) instead of a plain shared-memory reference.","commonSituations":"Porting a Hopper/Blackwell TMA kernel that used multimem load/stores to the async cluster-store API; passing a reference obtained from a tma map with multimem enabled into the new async barrier-based pipeline.","solutions":["Use store_tiled or a non-async store path for MultimemRef destinations","Construct the destination as a regular shared-memory reference (utils.smem_ref / SharedMemory) and pass that to store_tiled_async","If cluster multicast is needed, rely on the cluster_barrier_ptr/cluster plumbing of store_tiled_async rather than a MultimemRef"],"exampleFix":"// before\nfa.store_tiled_async(multimem_ref, ..., cluster_barrier_ptr=bar)\n// after\nsmem_ref = utils.smem_ref(shape, dtype)\nfa.store_tiled_async(smem_ref, ..., cluster_barrier_ptr=bar)","handlingStrategy":"validation","validationCode":"from jax.experimental.mosaic.gpu import utils\nif isinstance(ref, utils.MultimemRef):\n    raise TypeError('use store_tiled for MultimemRef destinations')\nfa.store_tiled_async(ref, ...)","typeGuard":"def is_async_store_ref(ref) -> bool:\n    return not isinstance(ref, utils.MultimemRef)","tryCatchPattern":"try:\n    fa.store_tiled_async(ref, ...)\nexcept ValueError as e:\n    if 'Multimem refs' in str(e):\n        fa.store_tiled(ref)  # fallback\n    else:\n        raise","preventionTips":["Never pass TMA/multimem descriptors to store_tiled_async","Centralize ref construction so async-store refs are always plain smem refs"],"tags":["jax","mosaic","gpu","multimem","async-store","not-implemented"],"backgroundTag":"unsupported-operation-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}