jax-ml/jax · error · ValueError

tcgen05_mma only allows arriving on a Barrier that orders te

Error message

tcgen05_mma only allows arriving on a Barrier that orders tensor core

What it means

The interpreter requires that the Barrier a tcgen05_mma arrives on has orders_tensor_core=True, i.e. it must order tensor-core (async MMA) operations, otherwise synchronization semantics would be wrong.

Source

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

        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,
    thread: memory.Thread,
    acc_allocation_key_as_array: jax.Array,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create the barrier with tensor-core ordering enabled (e.g. Barrier(..., orders_tensor_core=True) or the mosaic API that sets it)
2. Allocate a dedicated commit barrier for MMA separate from load barriers
3. Update jax — API names for ordering flags have changed across versions
4. Check mosaic_gpu examples for the canonical tcgen05 pipeline barrier setup

Example fix

# before
barrier = mgpu.Barrier(threads, num_workers=1)  # no tc ordering
# after
barrier = mgpu.Barrier(threads, num_workers=1, orders_tensor_core=True)

bar = mgpu.Barrier(threads, 1, orders_tensor_core=True)
assert bar.orders_tensor_core

def orders_tc(b) -> bool:
    return isinstance(b, memory.Barrier) and b.orders_tensor_core

try:
    kernel(x)
except ValueError as e:
    if 'orders tensor core' in str(e):
        recreate barrier with orders_tensor_core=True

When it happens

Trigger: Attaching a commit barrier to tcgen05_mma where the Barrier was created without tensor-core ordering (orders_tensor_core flag False), in interpret mode.

Common situations: Reusing a generic barrier intended for TMA/loads as an MMA commit barrier; missing the ordering option when allocating the barrier.

Related errors

• tcgen05_mma only allows arriving on a Barrier
• tcgen05_commit_arrive only allows arriving on a Barrier that
• tcgen05_commit_arrive only allows arriving on a Barrier
• MMA rhs swizzle must match lhs swizzle. {lhs_swizzle=} {rhs_
• Expected metadata dtype to be uint2, got: {meta.dtype}

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