jax-ml/jax · error · ValueError

Async load only supports TMEM refs

Error message

Async load only supports TMEM refs

What it means

The async_load_tmem Pallas primitive only accepts references in Tensor Memory (TMEM). Passing a ref whose memory_space is not gpu_core.MemorySpace.TMEM raises ValueError in the abstract eval.

Source

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

  )
  if reduce is None:
    [result] = results
    if layout is not None:
      result = gpu_core.layout_cast(result, layout)
    return result
  else:
    result, reduced = results
    if layout is not None:
      result = gpu_core.layout_cast(result, layout)
    return result, reduced


@async_load_tmem_p.def_effectful_abstract_eval
def _async_load_tmem_abstract_eval(
    src, *avals_flat, tree, reduce: tcgen05.LoadReduceOp | None = None
):
  if src.memory_space != gpu_core.MemorySpace.TMEM:
    raise ValueError("Async load only supports TMEM refs")
  val_aval, effects = state_primitives._get_abstract_eval(
      src, *avals_flat, tree=tree
  )
  if reduce is None:
    return (val_aval,), effects
  if val_aval.dtype not in map(jnp.dtype, (jnp.float32, jnp.int32, jnp.uint32)):
    raise ValueError(
        f"Unsupported dtype for reduction: {val_aval.dtype}. Only float32, "
        " int32 and uint32 are supported."
    )
  reduced_aval = jax_core.ShapedArray(val_aval.shape[:-1], val_aval.dtype)
  return (val_aval, reduced_aval), effects


@lowering.register_lowering_rule(async_load_tmem_p, mgpu.LoweringSemantics.Lane)
def _async_load_tmem_lowering_rule(
    ctx: lowering.LoweringRuleContext,
    x_ref,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate the source ref in Tensor Memory (tcgen05 TMEM allocation) before async loading
2. Use the regular async_load for SMEM and other memory spaces
3. Check ref.memory_space before calling to fail fast

Example fix

// before
val = async_load_tmem(smem_ref)
// after
val = async_load_tmem(tmem_ref)  # ref allocated in TMEM

from jax._src.pallas import gpu_core
assert src.memory_space == gpu_core.MemorySpace.TMEM

def is_tmem(ref): return getattr(ref, 'memory_space', None) == gpu_core.MemorySpace.TMEM

When it happens

Trigger: Calling pallas' async TMEM load API (e.g. pl_mosaic.async_load_tmem) on an SMEM or normal (HBM/VMEM) reference.

Common situations: Copy-pasting a TMEM kernel snippet onto a regular grid ref; forgetting to allocate the ref in TMEM for Blackwell tcgen05 usage.

Related errors

• Async store only supports TMEM refs
• Stores to TMEM are asynchronous operations and cannot be per
• Accumulator must be a TMEM Ref.
• LHS must be a TMEM/SMEM Ref.
• Memory space {self.memory_space} is not supported by mesh {s

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