jax-ml/jax · error · ValueError

DMA source/destination/semaphore arguments must be Refs.

Error message

DMA source/destination/semaphore arguments must be Refs.

What it means

dma_start operates on state.Ref arguments, not plain arrays, because the DMA engine writes memory in place and effects tracking depends on ref types. The abstract eval requires src_ref, dst_ref and dst_sem to be state.AbstractRef or state.TransformedRef; anything else (e.g. a jnp array or a triton-style pointer) is rejected.

Source

Thrown at jax/_src/pallas/mosaic/primitives.py:356

      priority=priority,
      device_id_type=device_id_type,
      add=add,
  )
  return []
dma_start_p.to_lojax = _dma_start_to_lojax

@dma_start_p.def_effectful_abstract_eval
def _dma_start_abstract_eval(*args, tree, device_id_type, priority, add):
  if priority < 0:
    raise ValueError(f"DMA start priority must be non-negative: {priority}")
  src_ref_aval, dst_ref_aval, dst_sem_aval, src_sem_aval, device_id_aval = (
      _dma_unflatten(tree, args)
  )
  if not all(
      isinstance(x, (state.AbstractRef, state.TransformedRef))
      for x in [src_ref_aval, dst_ref_aval, dst_sem_aval]
  ):
    raise ValueError(
        "DMA source/destination/semaphore arguments must be Refs.")
  dst_sem_shape = dst_sem_aval.shape
  if dst_sem_shape:
    raise ValueError(
        f"Cannot signal on a non-()-shaped semaphore: {dst_sem_shape}"
    )
  if src_sem_aval is not None:
    if not isinstance(src_sem_aval, (state.AbstractRef, state.TransformedRef)):
      raise ValueError("DMA source semaphore must be a Ref.")
    src_sem_shape = src_sem_aval.shape
    if src_sem_shape:
      raise ValueError(
          f"Cannot signal on a non-()-shaped semaphore: {src_sem_shape}"
      )
  return [], _get_dma_effects(
      src_ref_aval,
      dst_ref_aval,
      dst_sem_aval,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the Ref objects themselves, not blocks read from them
2. Ensure semaphores and buffers come from the kernel's ref parameters or state allocation
3. Check for stray get()/read calls between allocation and dma_start

Example fix

# before
block = src_ref[...]
dma_start(block, dst_ref, sem_ref)  # block is an Array

# after
dma_start(src_ref, dst_ref, sem_ref)  # pass refs directly; index via the DMA's own offsets

from jax._src import state
assert all(isinstance(v, (state.AbstractRef, state.TransformedRef)) or hasattr(v, 'shape') is False for v in (src_ref, dst_ref, dst_sem))

def is_ref(v) -> bool:
  from jax._src import state
  return isinstance(getattr(v, 'aval', v), (state.AbstractRef, state.TransformedRef))

When it happens

Trigger: Calling dma_start with a plain jax.Array as source or destination (e.g. passing a block read out of a ref instead of the ref itself), or a non-ref semaphore.

Common situations: Confusing Pallas refs with arrays because they both support [] indexing — reading a block and passing it to dma_start; adapting example code that used different plumbing; forgetting to declare a semaphore via the ref API.

Related errors

• DMA source semaphore must be a Ref.
• Either both or neither `src_sem` and `device_id` can be set.
• Cannot `wait_send` on a local copy.
• dma_start not implemented in LoJAX yet.
• DMA start priority must be non-negative: {priority}

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