jax-ml/jax · error · ValueError

Accumulator Ref must be collective if collective_axis is set

Error message

Accumulator Ref must be collective if collective_axis is set.

What it means

When collective_axis is passed to tcgen05.mma, the accumulator TMEM ref must itself be marked collective. The check is skipped if the aval has degraded to a plain MemRef (e.g. under core_map).

Source

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

                               collective_axis,
                               arrive,
                               scaled,
                               sparse):
  del accumulate, acc_transforms_tree, a_transforms_tree, b_transforms_tree, barrier_transforms_tree

  if acc.memory_space != gpu_core.TMEM:
    raise ValueError("Accumulator must be a TMEM Ref.")
  if a.memory_space not in (gpu_core.SMEM, gpu_core.TMEM):
    raise ValueError("LHS must be a TMEM/SMEM Ref.")
  if b.memory_space != gpu_core.SMEM:
    raise ValueError("RHS must be an SMEM Ref.")

  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")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate the accumulator TMEM ref with the collective flag enabled (see the collective MMA helper/allocate API used by Mosaic GPU examples)
2. Make the collective setting of the accumulator consistent with collective_axis
3. If you did not intend 2CTA mode, remove collective_axis

Example fix

# before
acc = allocate(TMEM, (m, 2*n), jnp.float32)  # not collective
tcgen05.mma(a, b, acc, k_dim=k, collective_axis=0)
# after
acc = allocate(TMEM, (m, 2*n), jnp.float32, collective=True)
tcgen05.mma(a, b, acc, k_dim=k, collective_axis=0)

if collective_axis is not None:
    assert getattr(acc, 'collective', True), 'acc must be collective'

def acc_collective_ok(acc, collective_axis):
    return collective_axis is None or getattr(acc, 'collective', False)

When it happens

Trigger: Calling tcgen05.mma(..., collective_axis=0) with an accumulator that was allocated without the collective TMEM flag.

Common situations: Adding collective_axis to an existing kernel but keeping the old TMEM allocation; using a non-collective accumulator allocation helper when switching to 2CTA mode.

Related errors

• LHS Ref must be collective if collective_axis is set.
• TMEM aliasing only supported for Refs with the same first di
• Unsupported TMEM ref {ref}.
• Stores to TMEM are asynchronous operations and cannot be per
• Accumulator and RHS have incompatible shapes. Expected RHS t

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