{"record":{"id":"95501ca8bf18f1d1","repo":"jax-ml/jax","slug":"loading-from-a-remote-ref-is-only-supported-in-jax","errorCode":null,"errorMessage":"Loading from a remote ref is only supported in jaxlib version 0.11.1 or higher under Warpgroup lowering semantics","messagePattern":"Loading from a remote ref is only supported in jaxlib version 0\\.11\\.1 or higher under Warpgroup lowering semantics","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":1188,"sourceCode":"    bytes = mgpu.c(bytes, ir.IntegerType.get_signless(32))\n\n  if predicate is not None:\n    # We can not enter this branch with bytes as int\n    # because NotImplementedError is raised earlier for\n    # jaxlib<0.11.1 and predicate is not None\n    assert isinstance(bytes, ir.Value)\n    bytes = arith_dialect.select(predicate, bytes, mgpu.c(0, i32))\n\n  with arrive_ctx:\n    mgpu.dialect.arrive_expect_tx(barrier_ref, bytes)\n\n  peer_id = copy_params.get(\"gmem_peer_id\")\n  # TODO(bchetioui): Remove once 0.11.1 is the minimum jaxlib version.\n  if \"gmem_peer_id\" in inspect.signature(mgpu.dialect.async_load).parameters:\n    peer_kwarg = dict(gmem_peer_id=peer_id)\n  else:\n    if peer_id is not None:\n      raise NotImplementedError(\n          \"Loading from a remote ref is only supported in jaxlib version \"\n          \"0.11.1 or higher under Warpgroup lowering semantics\"\n      )\n    peer_kwarg = {}\n\n  mgpu.dialect.async_load(\n      src,\n      dst,\n      indices,  # pyrefly: ignore[bad-argument-type]\n      slice_lengths,  # pyrefly: ignore[bad-argument-type]\n      barrier=barrier_ref,  # pyrefly: ignore[bad-keyword-argument]\n      predicate=predicate,\n      collective=ir.ArrayAttr.get(  # pyrefly: ignore[bad-keyword-argument]\n          [ir.IntegerAttr.get(i32, axis) for axis in collective or []]\n      ),\n      leader_tracked=leader_tracked_attr,\n      oob_fill_mode=ir.IntegerAttr.get(i32, oob_mode.value),\n      **peer_kwarg,  # pyrefly: ignore[bad-argument-type]","sourceCodeStart":1170,"sourceCodeEnd":1206,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L1170-L1206","documentation":"Loading from a remote (peer) GPU's memory via gmem_peer_id requires the gmem_peer_id parameter on mgpu.dialect.async_load, which was added in jaxlib 0.11.1. On older jaxlib the peer id cannot be forwarded, so remote loads under Warpgroup lowering are rejected.","triggerScenarios":"Passing a ref with a peer id (e.g. created for a remote GPU in a NVLink/multi-GPU pallas kernel) into a warpgroup copy on jaxlib < 0.11.1.","commonSituations":"Multi-GPU (NVLink) Mosaic kernels developed against 0.11.1 run on a cluster image with an older jaxlib; CI with version drift between jax and jaxlib.","solutions":["Upgrade jaxlib to >= 0.11.1.","Remove the peer id / use local memory only on older jaxlib."],"exampleFix":null,"handlingStrategy":"validation","validationCode":"from jax._src.pallas.mosaic_gpu import mgpu\nimport inspect\nsupports_peer = 'gmem_peer_id' in inspect.signature(mgpu.dialect.async_load).parameters\nassert peer_id is None or supports_peer, 'remote loads need jaxlib >= 0.11.1'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Check jaxlib version at env setup for multi-GPU kernels.","Gate peer-memory kernels behind a version check with a clear error message."],"tags":["mosaic-gpu","pallas","jaxlib-version","multi-gpu","peer-id"],"backgroundTag":"library-version-too-old","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}