jax-ml/jax · error · ValueError

Accumulator must be a TMEM Ref.

Error message

Accumulator must be a TMEM Ref.

What it means

Type guard in the tcgen05_mma abstract eval: Blackwell tcgen05 MMA writes its result to Tensor Memory, so the acc argument's memory_space must be gpu_core.TMEM. An accumulator ref allocated in SMEM or GMEM fails this check.

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.
• a_scale must be a TMEM Ref
• Accumulator Ref must be collective if collective_axis is…
• Async load only supports TMEM refs
• Async store only supports TMEM refs

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

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