jax-ml/jax · error · ValueError

RHS must be an SMEM Ref.

Error message

RHS must be an SMEM Ref.

What it means

The RHS operand of tcgen05.mma must be an SMEM ref. TMEM or GMEM RHS operands are rejected.

Source

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

                               *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:
    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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Stage the B tile into SMEM with the required swizzle/tiling transforms
2. Do not allocate B in TMEM
3. Verify b.memory_space == SMEM before the call

Example fix

# before
b_tmem = allocate(TMEM, (k, n), jnp.float8_e4m3fnuz)
tcgen05.mma(a_smem, b_tmem, acc, k_dim=k)
# after
b_smem = load_to_smem(b_gmem, swizzle=128)
tcgen05.mma(a_smem, b_smem, acc, k_dim=k)

assert b.memory_space == gpu_core.SMEM

def is_smem(ref):
    return getattr(ref, 'memory_space', None) == gpu_core.SMEM

When it happens

Trigger: Passing b as a TMEM ref or an unstaged GMEM input to tcgen05.mma.

Common situations: Symmetric allocation of A and B in TMEM when only A may live there; skipping the SMEM staging/swizzle for B when porting a kernel.

Related errors

• LHS must be a TMEM/SMEM Ref.
• Only workgroup memory is supported but got {ref}.
• Accumulator must be a TMEM Ref.
• Memory space {self.memory_space} is not supported by mesh {s
• Cannot eagerly run with_memory_space_constraint.

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