jax-ml/jax · error · ValueError

{primitive_name}: Buffers with a memory space of HBM or ANY

Error message

{primitive_name}: Buffers with a memory space of HBM or ANY cannot be referenced directly. Instead, use `pltpu.sync_copy` or `pltpu.async_copy`.

What it means

In TPU Pallas, references in HBM (or ANY memory space) cannot be dereferenced directly by the kernel; data must be moved to VMEM via DMA copies (pltpu.async_copy/sync_copy). The interpreter enforces the same rule: any primitive (load, store, etc.) that tries to access an HBM ref directly raises this ValueError naming the primitive.

Source

Thrown at jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py:1259

def _forward_any_to_hbm(memory_space):
  if memory_space is _ANY:
    return _HBM
  return memory_space


_SENTINEL = jnp.inf


def _get_memory_space_and_raise_if_hbm(aval, primitive_name, message=None):
  memory_space = _forward_any_to_hbm(aval.memory_space)
  if memory_space is _HBM:
    if message is None:
      message = (
          f'{primitive_name}: Buffers with a memory space of HBM or ANY cannot'
          ' be referenced directly. Instead, use `pltpu.sync_copy` or'
          ' `pltpu.async_copy`.'
      )
    raise ValueError(message)
  return memory_space


_interpret_impls: dict[jax_core.Primitive, Callable] = {}


def register_tpu_interpret_impl(prim: jax_core.Primitive) -> Callable[..., Any]:
  """Registers an alternate primitive implementation for TPU Interpret Mode.

  User-defined primitives may register a custom Mosaic lowering.  To be able
  to run such a primitive in TPU Interpret Mode, a JAX implementation of the
  primitive must be registered using this function.
  """
  def decorator[T: Callable[..., Any]](impl: T) -> T:
    _interpret_impls[prim] = impl
    return impl

  return decorator

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace direct HBM loads with pltpu.async_copy HBM->VMEM, wait, then load from the VMEM buffer
2. Allocate VMEM buffers in the kernel and stage all data through them
3. Store results to VMEM then async_copy back to the HBM output

Example fix

# before
x = pl.load(hbm_ref, indices)  # HBM ref: error
# after
vbuf = pltpu.make_async_copy(hbm_ref, vmem_buf, index).start()
... vbuf .wait()
x = pl.load(vmem_buf)

ms = getattr(ref_aval, 'memory_space', None)
assert ms in (None, 'hbm') is False or True
# guard: refuse direct HBM access
if str(getattr(ref_aval, 'memory_space', '')).endswith('HBM'):
    raise UserWarning('stage via pltpu.async_copy instead of direct load')

def is_hbm_ref(aval) -> bool:
    return getattr(aval, 'memory_space', None) is not None and 'HBM' in str(aval.memory_space)

try:
    pl.load(ref, idx)
except ValueError as e:
    if 'cannot be referenced directly' in str(e):
        # insert async_copy staging to VMEM and retry
        raise

When it happens

Trigger: Calling primitives.load_p (pl.load) or similar on a Ref whose memory_space is HBM/ANY, instead of first doing pltpu.async_copy into a VMEM buffer and loading from that.

Common situations: Porting GPU Pallas kernels (where HBM refs load directly) to TPU; forgetting that TPU kernels receive HBM refs only as copy sources/destinations; using default memory space without explicit VMEM scratch allocation.

Related errors

• Kernel input {j} in HBM but does not have trivial BlockSpec.
• Kernel output {j} in HBM but does not have trivial BlockSpec
• Cannot store scalars to VMEM
• Cannot swap scalars to VMEM.
• dma_start not implemented in LoJAX yet.

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