jax-ml/jax · error · ValueError

LHS must be a TMEM/SMEM Ref.

Error message

LHS must be a TMEM/SMEM Ref.

What it means

Type guard in the tcgen05_mma abstract eval: the lhs operand must live in shared memory or tensor memory (gpu_core.SMEM or gpu_core.TMEM), since those are the only sources the tensor-core MMA instruction can read. An 'a' ref in any other memory space (e.g., GMEM) is rejected.

Solutions

1. Load the LHS tile into an SMEM buffer with proper swizzle before the MMA
2. Or keep A in TMEM if using a TMEM-resident LHS layout
3. Ensure the ref's memory_space attribute is SMEM or TMEM

Example fix

# before
tcgen05.mma(a_gmem_ref, b_smem, acc, k_dim=k)
# after
a_smem = load_to_smem(a_gmem_ref, swizzle=128)
tcgen05.mma(a_smem, b_smem, acc, k_dim=k)

assert a.memory_space in (gpu_core.SMEM, gpu_core.TMEM)

def lhs_space_ok(ref):
    return getattr(ref, 'memory_space', None) in (gpu_core.SMEM, gpu_core.TMEM)

When it happens

Trigger: Passing a raw kernel input (GMEM) directly as the LHS without first loading/swizzling it into SMEM or TMEM.

Common situations: Forgetting to stage the A tile through an SMEM scratch buffer with the required swizzle/tiling; treating tcgen05.mma like a regular dot that accepts GMEM operands.

Related errors

• Accumulator must be a TMEM Ref.
• RHS must be an SMEM Ref.
• a_scale must be a TMEM Ref
• Accumulator Ref must be collective if collective_axis is…
• All aliased Refs must have the same memory space (SMEM or…

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

Appendix: source

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

@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:
    barrier, *scales_and_transforms_leaves = barrier_scales_and_transforms_leaves
    orders_tensor_core = getattr(

View on GitHub (pinned to 1e1c6a8fc0)