jax-ml/jax · error · ValueError

Can only perform collective MMA along minormost cluster axis

Error message

Can only perform collective MMA along minormost cluster axis. Got {collective_axis}, expected {minormost_cluster_axis}.

What it means

Collective MMA can only multicast along the minormost (first) cluster axis because that is the axis the hardware multicast uses. If the requested collective_axis resolves to a non-zero cluster dimension, the lowering rejects it and tells you which axis it expected.

Source

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

  with predicate_ctx:
    mgpu.dialect.tcgen05_commit_arrive(
        barrier_ref.as_barrier_memref(), collective=collective
    )
  return []


def _collective_mma_predicate(ctx: lowering.LoweringRuleContext,
                              collective_axis: str) -> ir.Value:
  """Computes a predicate to run only on the leader block."""
  cluster_axis = lowering._resolve_cluster_axis(
      ctx.module_ctx.axis_names, collective_axis)
  if cluster_axis != gpu_dialect.Dimension(0):
    # Note: resolve_cluster_axis checks if axis_names exists.
    assert ctx.module_ctx.axis_names is not None
    if len(ctx.module_ctx.axis_names.cluster) <= 1:
      raise ValueError("No cluster axes found.")
    minormost_cluster_axis = ctx.module_ctx.axis_names.cluster[0]
    raise ValueError(
        "Can only perform collective MMA along minormost cluster axis. "
        f"Got {collective_axis}, expected {minormost_cluster_axis}.")
  index = ir.IndexType.get()
  is_leader_block = arith_dialect.cmpi(
      arith_dialect.CmpIPredicate.eq,
      mgpu.utils.cluster_idx(cluster_axis), mgpu.c(0, index))
  return is_leader_block


commit_tmem_p = jax_core.Primitive("commit_tmem")
commit_tmem_p.multiple_results = True


@commit_tmem_p.def_effectful_abstract_eval
def _commit_tmem_abstract_eval():
  return (), {gpu_core._memory_effect}

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reorder your axes so the requested collective axis is the minormost cluster axis (first in cluster list)
2. Change collective_axis to name the first cluster axis

Example fix

// before
# cluster axes declared as ('x','y'); requesting 'y'
spec = BlockSpec(..., collective_axis='y')
// after
spec = BlockSpec(..., collective_axis='x')

assert collective_axis == cluster_axis_names[0]

When it happens

Trigger: Passing collective_axis naming a cluster axis that is not axis_names.cluster[0], e.g. declaring cluster axes ('x','y') and requesting collective_axis='y'.

Common situations: Reordering BlockSpec/collective axes; assuming any cluster axis can be multicast along.

Related errors

• No cluster axes found.
• Accumulator and RHS have incompatible shapes. Expected RHS t
• Accumulator Ref must be collective if collective_axis is set
• LHS Ref must be collective if collective_axis is set.
• dims and idxs must have the same length

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