jax-ml/jax · error · ValueError

LHS Ref must be collective if collective_axis is set.

Error message

LHS Ref must be collective if collective_axis is set.

What it means

When collective_axis is set on tcgen05.mma and the LHS is a TMEM ref, that TMEM ref must be marked collective; otherwise the LHS layout does not match the 2CTA tensor-core expectation.

Source

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

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

  return [], {gpu_core._memory_effect}

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Mark the LHS TMEM ref as collective when using collective_axis
2. Or move the LHS to SMEM (SMEM path is not subject to this check)
3. Keep the collective flag consistent across acc/a allocations

Example fix

# before
a_tmem = allocate(TMEM, (m, k), jnp.float8e4m3)
tcgen05.mma(a_tmem, b_smem, acc, k_dim=k, collective_axis=0)
# after
a_tmem = allocate(TMEM, (m, k), jnp.float8e4m3, collective=True)
tcgen05.mma(a_tmem, b_smem, acc, k_dim=k, collective_axis=0)

if collective_axis is not None and isinstance(a, gpu_core.AbstractTMEMRef):
    assert a.collective, 'LHS TMEM must be collective'

def lhs_collective_ok(a, collective_axis):
    if collective_axis is None:
        return True
    return not isinstance(a, gpu_core.AbstractTMEMRef) or a.collective

When it happens

Trigger: Passing a non-collective TMEM-resident LHS together with collective_axis in tcgen05.mma.

Common situations: Using TMEM for A in a kernel later converted to collective MMA without updating A's allocation; mixing collective accumulator with per-CTA LHS TMEM.

Related errors

• Accumulator 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/ed30b1104f43d7b6. Report an issue: GitHub.