jax-ml/jax · error · ValueError

Accumulator must be a TMEM Ref.

Error message

Accumulator must be a TMEM Ref.

What it means

The tcgen05.mma abstract eval requires the accumulator to live in Tensor Memory (TMEM). Passing an accumulator ref in SMEM or GMEM raises this error.

Source

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

@tcgen05_mma_p.def_effectful_abstract_eval
def _tcgen05_mma_abstract_eval(acc, a, b, accumulate,
                               *barrier_scales_and_transforms_leaves,
                               acc_transforms_tree, a_transforms_tree,
                               b_transforms_tree,
                               barrier_transforms_tree,
                               a_scale_transforms_tree,
                               b_scale_transforms_tree,
                               a_sparse_metadata_transforms_tree,
                               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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate the accumulator in TMEM: allocate(TMEM, shape, dtype) or mosaic_gpu.TMEM memory space
2. Keep SMEM only for the A/B operands
3. Check that the accumulator came from a TMEM scratch specification

Example fix

# before
acc = allocate(SMEM, (m, n), jnp.float32)
tcgen05.mma(a, b, acc, k_dim=k)
# after
acc = allocate(TMEM, (m, n), jnp.float32)
tcgen05.mma(a, b, acc, k_dim=k)

from jax._src.pallas.mosaic_gpu import core as gpu_core
assert acc.memory_space == gpu_core.TMEM

def is_tmem(ref):
    return getattr(ref, 'memory_space', None).__class__.__name__ == 'TMEM' or str(getattr(ref, 'memory_space', '')) == 'TMEM'

When it happens

Trigger: Allocating the accumulator with pallas.memory_space=SMEM (or plain allocate without memory_space) and passing it as acc to tcgen05.mma.

Common situations: Copying an SMEM allocation pattern from dot-operand buffers to the accumulator; forgetting that Blackwell tensor-core accumulators must be TMEM refs.

Related errors

• LHS must be a TMEM/SMEM Ref.
• 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
• RHS must be an SMEM Ref.

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