jax-ml/jax · error · NotImplementedError

Loading from a remote ref is only supported in jaxlib versio

Error message

Loading from a remote ref is only supported in jaxlib version 0.11.1 or higher under Warpgroup lowering semantics

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:1188

    bytes = mgpu.c(bytes, ir.IntegerType.get_signless(32))

  if predicate is not None:
    # We can not enter this branch with bytes as int
    # because NotImplementedError is raised earlier for
    # jaxlib<0.11.1 and predicate is not None
    assert isinstance(bytes, ir.Value)
    bytes = arith_dialect.select(predicate, bytes, mgpu.c(0, i32))

  with arrive_ctx:
    mgpu.dialect.arrive_expect_tx(barrier_ref, bytes)

  peer_id = copy_params.get("gmem_peer_id")
  # TODO(bchetioui): Remove once 0.11.1 is the minimum jaxlib version.
  if "gmem_peer_id" in inspect.signature(mgpu.dialect.async_load).parameters:
    peer_kwarg = dict(gmem_peer_id=peer_id)
  else:
    if peer_id is not None:
      raise NotImplementedError(
          "Loading from a remote ref is only supported in jaxlib version "
          "0.11.1 or higher under Warpgroup lowering semantics"
      )
    peer_kwarg = {}

  mgpu.dialect.async_load(
      src,
      dst,
      indices,  # pyrefly: ignore[bad-argument-type]
      slice_lengths,  # pyrefly: ignore[bad-argument-type]
      barrier=barrier_ref,  # pyrefly: ignore[bad-keyword-argument]
      predicate=predicate,
      collective=ir.ArrayAttr.get(  # pyrefly: ignore[bad-keyword-argument]
          [ir.IntegerAttr.get(i32, axis) for axis in collective or []]
      ),
      leader_tracked=leader_tracked_attr,
      oob_fill_mode=ir.IntegerAttr.get(i32, oob_mode.value),
      **peer_kwarg,  # pyrefly: ignore[bad-argument-type]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Upgrade jaxlib to >= 0.11.1.
Remove the peer id / use local memory only on older jaxlib.

Defensive patterns

Strategy: validation

Validate before calling

from jax._src.pallas.mosaic_gpu import mgpu
import inspect
supports_peer = 'gmem_peer_id' in inspect.signature(mgpu.dialect.async_load).parameters
assert peer_id is None or supports_peer, 'remote loads need jaxlib >= 0.11.1'

Prevention

Check jaxlib version at env setup for multi-GPU kernels.
Gate peer-memory kernels behind a version check with a clear error message.

When it happens

Trigger: 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.

Common situations: 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.

Related errors

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