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.

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
• `broadcast_to` is a Triton-specific primitive. Please…
• Cannot signal on a non-()-shaped semaphore

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

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