jax-ml/jax · error · NotImplementedError

dma_start not implemented in LoJAX yet.

Error message

dma_start not implemented in LoJAX yet.

What it means

When a dma_start is lowered into LoJAX (single-device/multi-core) form, both the source and destination refs must be 'high' (per-core global) refs. If either the src or dst ref is not high, the LoJAX lowering does not know how to route the copy and raises NotImplementedError.

Source

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

        for k_ in k:
          effs.add(jax_core.NamedAxisEffect(k_))
  return effs


dma_start_p = jax_core.Primitive('dma_start')
dma_start_p.multiple_results = True

def _dma_is_high(*avals, **params):
  return any(aval.is_high for aval in avals)

dma_start_p.is_high = _dma_is_high

def _dma_start_to_lojax(*args, tree, device_id_type, priority, add):
  src_ref, dst_ref, dst_sem, src_sem, device_id = _dma_unflatten(tree, args)
  src_ref_aval = jax_core.typeof(_get_ref(src_ref))
  dst_ref_aval = jax_core.typeof(_get_ref(dst_ref))
  if not (src_ref_aval.is_high and dst_ref_aval.is_high):
    raise NotImplementedError("dma_start not implemented in LoJAX yet.")
  dst_sem_aval = jax_core.typeof(_get_ref(dst_sem))
  if dst_sem_aval.is_high:
    raise NotImplementedError("dma_start not implemented in LoJAX yet.")
  if _get_ref(src_sem) is not None:
    if jax_core.typeof(_get_ref(src_sem)).is_high:
      raise NotImplementedError("dma_start not implemented in LoJAX yet.")
  src_ref_aval.inner_aval.dma_start(
      src_ref,
      dst_ref,
      src_sem,
      dst_sem,
      device_id=device_id,
      priority=priority,
      device_id_type=device_id_type,
      add=add,
  )
  return []
dma_start_p.to_lojax = _dma_start_to_lojax

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate/rebind the src and dst buffers as high (global) refs so is_high is true for both
2. Restructure so the DMA only references globally visible buffers, with local processing done via ordinary loads/stores
3. Check the JAX version — LoJAX dma_start coverage may have expanded

Example fix

// no user-level one-liner; pattern:
# before: src_ref is core-local (is_high=False)
dma_start(src_ref, dst_ref, sem, ...)

# after: obtain a high/global view for both endpoints
global_src = get_high_ref(...)
global_dst = get_high_ref(...)
dma_start(global_src, global_dst, sem, ...)

src_h = jax_core.typeof(get_ref(src_ref)).is_high
dst_h = jax_core.typeof(get_ref(dst_ref)).is_high
assert src_h and dst_h, "both DMA endpoints must be high refs under LoJAX"

def is_high_ref(r) -> bool:
  import jax._src.core as c
  return c.typeof(r).is_high

try:
  _dma_start_to_lojax(...)
except NotImplementedError:
  # fall back to ordinary load/store copy within one core

When it happens

Trigger: dma_start with a src or dst ref whose aval has is_high == False while running under the LoJAX lowering path (e.g. refs created inside a per-core context being used for a cross-core copy).

Common situations: Running multi-core (LoJAX) Pallas kernels where buffers were allocated in core-local space; mixing core-local scratch refs into DMA operations.

Related errors

• Batching over dynamic grid values is not supported yet.
• DMA partial discharge add=True not yet implemented.
• Acc ref must be at least 2D, got shape {shape}
• {primitive_name}: Buffers with a memory space of HBM or ANY
• masked load_p

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