{"record":{"id":"4ce4f746c150d04b","repo":"jax-ml/jax","slug":"gmem-refs-with-peer-ids-are-not-supported-in-warpg","errorCode":null,"errorMessage":"GMEM refs with peer ids are not supported in warpgroup lowering.","messagePattern":"GMEM refs with peer ids are not supported in warpgroup lowering\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":1429,"sourceCode":"    ctx.launch_ctx.async_prefetch(\n        gmem_ref=ref,\n        collective=collective,\n        leader_tracked=leader_tracked,\n        **copy_params,\n        **predicate_kwarg,\n    )\n    return ()\n\n  if \"gmem_slice\" not in copy_params:\n    i32 = ir.IntegerType.get_signless(32)\n    slice_lengths = ir.MemRefType(ref.type).shape\n    indices = [mgpu.utils.c(0, i32)] * len(slice_lengths)\n  else:\n    indices, slice_lengths = _split_gmem_slice(copy_params[\"gmem_slice\"])\n  assert copy_params.get(\"swizzle\") is None\n  assert not copy_params.get(\"gmem_transform\")\n  if copy_params.get(\"gmem_peer_id\", None) is not None:\n    raise NotImplementedError(\n        \"GMEM refs with peer ids are not supported in warpgroup lowering.\"\n    )\n  mgpu.dialect.async_prefetch(\n      ref, indices, slice_lengths, collective=ir.ArrayAttr.get([])\n  )\n  return ()\n\n\ndef async_prefetch(\n    ref: _Ref,\n    *,\n    collective_axes: str | tuple[str, ...] | None = None,\n    leader_tracked: CopyPartition | None = None,\n) -> None:\n  \"\"\"Asynchronously prefetches a GMEM reference to the L2 cache.\n\n  If collective_axes is specified, each CUDA block only prefetches a part of\n  the ``ref``, with other parts covered by blocks that share the same index","sourceCodeStart":1411,"sourceCodeEnd":1447,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L1411-L1447","documentation":"The warpgroup lowering of prefetch_ref (async_prefetch) does not know how to attach a peer id to the prefetch op, so prefetching from a remote GPU's GMEM ref is unsupported in this lowering path.","triggerScenarios":"Calling prefetch_ref on a ref that carries a gmem_peer_id while lowering with Warpgroup semantics.","commonSituations":"NVLink multi-GPU kernels trying to prefetch remote-memory tiles to overlap communication with compute.","solutions":["Remove the prefetch of the remote ref; prefetch only local refs.","Perform a full copy_gmem_to_smem from the peer ref instead of a prefetch (peer copies are supported on jaxlib >= 0.11.1)."],"exampleFix":null,"handlingStrategy":"validation","validationCode":"if getattr(ref, 'gmem_peer_id', None) is not None:\n    raise ValueError('prefetch of remote refs unsupported; use a full copy instead')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Only prefetch local GMEM refs; copy peer refs explicitly."],"tags":["mosaic-gpu","pallas","prefetch","multi-gpu","peer-id"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}