jax-ml/jax · error · ValueError

No cluster axes found.

Error message

No cluster axes found.

What it means

For collective MMA (multicast across a cluster), the kernel must be launched with more than one cluster axis; the lowering needs a minormost cluster axis to elect a leader block. If ctx.module_ctx.axis_names.cluster has length <= 1 there is nothing to multicast over, so it errors.

Source

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

    collective = False

  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. Declare at least 2 cluster axes in the kernel launch (e.g. set the cluster/collective dimension so axis_names.cluster has length > 1)
2. Only pass collective_axis when the kernel is actually launched as a multi-block cluster

Example fix

// before
kernel = pl.make_interp... # no cluster
tcgen05_mma(..., collective_axis='io')
// after
# launch with cluster axes, e.g. Interpretation with num clusters > 1 along an axis
kernel = pl.pallas_call(..., in_specs=[BlockSpec(..., collective_axis='io')])

assert collective_axis is None or len(kernel_cluster_axes) > 1

When it happens

Trigger: Calling collective MMA (or async copy to TMEM) with a collective_axis while the kernel's cluster (num_clusters/cluster axes) has 1 or fewer axes, e.g. running without ClusterSpec/collective axes.

Common situations: Developing a single-block kernel then enabling collective_axis; forgetting to declare the cluster dimension in the Pallas grid interpretation.

Related errors

• Can only perform collective MMA along minormost cluster axis
• 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/75322030d465b4bf. Report an issue: GitHub.