jax-ml/jax · error · NotImplementedError

Arriving on a collective barrier is not supported in a warp

Error message

Arriving on a collective barrier is not supported in a warp context

What it means

CollectiveBarrierRef.arrive is a warpgroup-wide collective; it cannot be issued from a warp-scoped context (PrimitiveSemantics.Warp) where only one warp participates. The lowering rejects it to avoid deadlock/incorrect mbarrier state.

Source

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

    transforms_treedef,
):
  transforms = transforms_treedef.unflatten(flat_transforms)
  barrier_aval = ctx.avals_in[0]
  assert isinstance(barrier_aval, state_types.AbstractRef)
  base_index = _get_barrier_base_index(barrier_aval, transforms)
  if base_index is not None:
    barrier = barrier[base_index]
  sem_dtype = barrier_aval.inner_aval.dtype  # pyrefly: ignore[missing-attribute]
  orders_tensor_core = getattr(sem_dtype, "orders_tensor_core", False)

  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    scope = mgpu_utils.ThreadSubset.WARP
  else:
    scope = mgpu_utils.ThreadSubset.WARPGROUP

  if isinstance(barrier, mgpu.CollectiveBarrierRef):
    if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
      raise NotImplementedError(
          "Arriving on a collective barrier is not supported in a warp context"
      )
    barrier.arrive(orders_tensor_core)
  elif ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    barrier.arrive(orders_tensor_core)
  else:
    if scope == mgpu_utils.ThreadSubset.WARP and not orders_tensor_core:
      arrival_count = 4
    else:
      arrival_count = 1

    pred = ctx.module_ctx.single_lane_predicate if orders_tensor_core else None
    barrier.arrive(
        arrival_count=arrival_count,
        orders_tensor_core=orders_tensor_core,
        predicate=pred,
        scope=scope,
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Issue the arrive from warpgroup-level code (outside warp-specialized bodies) so semantics are Warpgroup/WGxWARP.
2. Use a non-collective (per-warpgroup tracked) barrier for warp-scope arrivals.
3. Restructure warp specialization so the collective barrier is only touched at warpgroup scope.

# Only arrive on collective barriers from warpgroup-scope code:
if isinstance(barrier, mgpu.CollectiveBarrierRef):
    assert not in_warp_specialized_body, 'use a non-collective barrier inside warps'

When it happens

Trigger: Calling barrier_arrive on a CollectiveBarrierRef inside code lowered with Warp semantics — e.g. inside warp_specialize warp bodies or when the pipeline lowered to Warp-only semantics.

Common situations: Warp-specialized kernels where a DMA/store warp tries to arrive on a barrier that was allocated as a collective barrier; converting a warpgroup kernel to warp specialization without changing barrier allocation.

Related errors

• num_threads and thread_name must be either both set or both
• `thread_id` must be zero when allocating a buffer for all th
• Unsupported barrier type: {type(barrier)}
• tcgen05_mma only allows arriving on a Barrier
• tcgen05_mma only allows arriving on a Barrier that orders te

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