jax-ml/jax · error · ValueError

MMA barrier must have orders_tensor_core set to True.

Error message

MMA barrier must have orders_tensor_core set to True.

What it means

When arrive=True in tcgen05.mma, the provided MMA barrier must have orders_tensor_core=True on its dtype so the barrier arrival is ordered by the tensor core, not the SM. This is read via getattr on barrier.inner_aval.dtype.

Source

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

  if collective_axis is not None:
    # TODO(justinfu): If under a core_map, the avals for acc/a
    # become normal MemRefs so we cannot check if they are collective.
    # Figure out a way to fix this.
    if isinstance(acc, gpu_core.AbstractTMEMRef) and not acc.collective:
      raise ValueError(
          "Accumulator Ref must be collective if collective_axis is set.")
    if isinstance(a, gpu_core.AbstractTMEMRef) and not a.collective:
      raise ValueError(
          "LHS Ref must be collective if collective_axis is set.")

  scales_and_transforms_leaves = barrier_scales_and_transforms_leaves
  if arrive:
    barrier, *scales_and_transforms_leaves = barrier_scales_and_transforms_leaves
    orders_tensor_core = getattr(
        barrier.inner_aval.dtype, "orders_tensor_core", False)
    if not orders_tensor_core:
      raise ValueError("MMA barrier must have orders_tensor_core set to True.")
  if scaled:
    a_scale, b_scale = scales_and_transforms_leaves[:2]
    if a_scale.memory_space != gpu_core.TMEM:
      raise ValueError("a_scale must be a TMEM Ref")
    if b_scale.memory_space != gpu_core.TMEM:
      raise ValueError("b_scale must be a TMEM Ref")

  return [], {gpu_core._memory_effect}


@lowering.register_lowering_rule(tcgen05_mma_p, *gpu_core.LANExWG_SEMANTICS)
@lowering.register_lowering_rule(tcgen05_mma_p, *gpu_core.LANExWARP_SEMANTICS)
def _tcgen05_mma_lowering(
    ctx: lowering.LoweringRuleContext,
    acc: tcgen05.TMEMRef,
    a_ref,
    b_ref,
    accumulate: bool | ir.Value,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create the barrier with orders_tensor_core=True (mosaic_gpu barrier allocation option)
2. Or drop arrive=True and manage the arrival manually
3. Check getattr(bar.dtype, 'orders_tensor_core', False) before the call

Example fix

# before
bar = allocate_barrier()
tcgen05.mma(a, b, acc, k_dim=k, barrier=bar, arrive=True)
# after
bar = allocate_barrier(orders_tensor_core=True)
tcgen05.mma(a, b, acc, k_dim=k, barrier=bar, arrive=True)

if arrive:
    assert getattr(barrier.inner_aval.dtype, 'orders_tensor_core', False)

def barrier_orders_tensor_core(barrier):
    return getattr(getattr(barrier, 'inner_aval', barrier).dtype, 'orders_tensor_core', False)

When it happens

Trigger: Passing a regular SemaphoreBarrier (orders_tensor_core defaults False) as the barrier with arrive=True, e.g. tcgen05.mma(..., barrier=bar, arrive=True).

Common situations: Reusing a barrier created for loads/stores in an MMA pipeline; forgetting to construct the barrier with tensor-core ordering when doing software pipelining with tcgen05.

Related errors

• Lookahead is not supported for XLA pipeline emitter lowering
• Only single (1) and double (2) buffering are supported. Got
• DMA partial discharge add=True not yet implemented.
• Num arrivals must be at least 1, but got {n}
• Barriers are not arrays

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