jax-ml/jax · error · ValueError

tcgen05_mma only allows arriving on a Barrier

Error message

tcgen05_mma only allows arriving on a Barrier

What it means

When a tcgen05_mma (tensor-core matmul on Blackwell) op specifies a commit barrier, the interpreter requires it to be a standard memory.Barrier. Any other barrier type is rejected before arriving.

Source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py:1646

        self.thread,
        increment_clock=False,
        logging_info=logging_info,
    )

    if shared_memory.detect_races:
      assert clock is not None
      get_races().check_write(
          self.thread,
          clock.generic_clock,
          self.acc_key,
          acc_range,
          source_info=self.source_info,
      )

    if self.barrier_key:
      barrier = shared_memory.get_barrier(self.barrier_key)
      if not isinstance(barrier, memory.Barrier):
        raise ValueError("tcgen05_mma only allows arriving on a Barrier")
      if not barrier.orders_tensor_core:
        raise ValueError(
            "tcgen05_mma only allows arriving on a Barrier that orders tensor"
            " core"
        )
      barrier.arrive(
          thread=self.thread,
          clock=clock,
          logging_info=logging_info,
      )

    return clock.copy() if clock is not None else None


def tcgen05_mma(
    *,
    token: jax.Array,
    mesh_location: memory.MeshLocation,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Use a plain Barrier created via standard mosaic APIs for the MMA commit
  2. Ensure the barrier_key references a barrier allocated as memory.Barrier
  3. Update jax for newer tcgen05 interpret support
  4. Run on device if the exotic barrier is required
Defensive patterns

Strategy: type-guard

Type guard

def is_plain_barrier(b) -> bool:
    from jax._src.pallas.mosaic_gpu.interpret import memory
    return isinstance(b, memory.Barrier)

Try / catch

try:
    kernel(x)
except ValueError as e:
    if 'tcgen05_mma only allows arriving on a Barrier' in str(e):
        use standard Barrier for the MMA commit

Prevention

When it happens

Trigger: Using a tcgen05 MMA with a barrier_key whose resolved barrier in shared memory is not a memory.Barrier instance (e.g. an mbarrier variant or unsupported type), under interpret mode.

Common situations: Experimenting with Blackwell tensor-core pipelines and mixed barrier types; interpreter not covering newer barrier classes.

Related errors


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