jax-ml/jax · error · NotImplementedError

GMEM refs with peer ids are not supported in warpgroup lower

Error message

GMEM refs with peer ids are not supported in warpgroup lowering.

What it means

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.

Source

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

    ctx.launch_ctx.async_prefetch(
        gmem_ref=ref,
        collective=collective,
        leader_tracked=leader_tracked,
        **copy_params,
        **predicate_kwarg,
    )
    return ()

  if "gmem_slice" not in copy_params:
    i32 = ir.IntegerType.get_signless(32)
    slice_lengths = ir.MemRefType(ref.type).shape
    indices = [mgpu.utils.c(0, i32)] * len(slice_lengths)
  else:
    indices, slice_lengths = _split_gmem_slice(copy_params["gmem_slice"])
  assert copy_params.get("swizzle") is None
  assert not copy_params.get("gmem_transform")
  if copy_params.get("gmem_peer_id", None) is not None:
    raise NotImplementedError(
        "GMEM refs with peer ids are not supported in warpgroup lowering."
    )
  mgpu.dialect.async_prefetch(
      ref, indices, slice_lengths, collective=ir.ArrayAttr.get([])
  )
  return ()


def async_prefetch(
    ref: _Ref,
    *,
    collective_axes: str | tuple[str, ...] | None = None,
    leader_tracked: CopyPartition | None = None,
) -> None:
  """Asynchronously prefetches a GMEM reference to the L2 cache.

  If collective_axes is specified, each CUDA block only prefetches a part of
  the ``ref``, with other parts covered by blocks that share the same index

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the prefetch of the remote ref; prefetch only local refs.
2. Perform a full copy_gmem_to_smem from the peer ref instead of a prefetch (peer copies are supported on jaxlib >= 0.11.1).

if getattr(ref, 'gmem_peer_id', None) is not None:
    raise ValueError('prefetch of remote refs unsupported; use a full copy instead')

When it happens

Trigger: Calling prefetch_ref on a ref that carries a gmem_peer_id while lowering with Warpgroup semantics.

Common situations: NVLink multi-GPU kernels trying to prefetch remote-memory tiles to overlap communication with compute.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Loading from a remote ref is only supported in jaxlib versio
• next_fetch is None
• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• Cannot specify both out_shape and out_type. Use out_type.

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