jax-ml/jax · error · ValueError

Can only perform collective MMA along minormost cluster…

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.

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…
• Accumulator Ref must be collective if collective_axis is…
• LHS Ref must be collective if collective_axis is set.
• a_scale and b_scale must both be present or absent.

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

Appendix: 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)