jax-ml/jax · error · ValueError

RHS must be an SMEM Ref.

Error message

RHS must be an SMEM Ref.

What it means

Type guard in the tcgen05_mma abstract eval: the rhs ('b') operand is read exclusively from shared memory by the tensor-core MMA instruction, so its memory_space must be gpu_core.SMEM. A b ref in TMEM or GMEM fails this check.

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.
• Accumulator must be a TMEM Ref.
• Only workgroup memory is supported but got
• a_scale and b_scale must both be present or absent.
• a_scale must be a TMEM Ref

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

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