jax-ml/jax · error · ValueError

`thread_id` must be zero when allocating a buffer for all th

Error message

`thread_id` must be zero when allocating a buffer for all threads

What it means

In the GPU interpretation (simulation) backend, a buffer allocated for all threads must be requested by thread 0 only; the allocation callback checks HostAllocationRequest.thread_id == 0. A non-zero thread_id means the compiler emitted (or user code simulated) an invalid per-thread collective allocation.

Source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py:271

  Args:
    allocation_request_as_array: Array that converts into an
      `HostAllocationRequest` with `thread_id`/`block_id` set to zero.
    value: Array of values to initialize the allocated buffer with.

  Returns:
    `AllocationKey` to refer to the allocated buffer.

  Raises:
    ValueError: If `thread_id`/`block_id` in `allocation_request` is not zero.
  """
  allocation_request = HostAllocationRequest.from_array(
      allocation_request_as_array
  )
  del allocation_request_as_array

  if allocation_request.thread_id != 0:
    raise ValueError(
        "`thread_id` must be zero when allocating a buffer for all threads"
    )
  if allocation_request.block_id != 0:
    raise ValueError(
        "`block_id` must be zero when allocating a buffer for all threads"
    )
  assert allocation_request.memory_space_id != memory.get_memory_space_idx(
      mosaic_gpu_core.MemorySpace.REGS
  )

  shared_memory = _get_shared_memory()

  key: HostAllocationKey | None = None
  buffer_id: int | None = None
  for thread in shared_memory.concurrent_threads(device):
    buffer_id_for_thread_id = shared_memory.get_next_buffer_id(thread)
    if not buffer_id:
      buffer_id = buffer_id_for_thread_id

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure collective (all-thread) allocations are requested only once per block (thread 0), e.g. allocate outside the per-thread mapped function
2. Reproduce with a simpler kernel to find which allocation is per-thread vs collective; if it looks like a compiler bug, report upstream with a minimal reproducer
3. Try lowering num_threads/avoiding warp-specialized allocations in interpret mode

Example fix

# before (inside per-thread code):
buf = allocate_for_all_threads(...)  # called by every thread

# after (hoisted outside the per-thread body):
with collective_scope():
  buf = allocate_for_all_threads(...)  # requested once by thread 0

# in custom interpret-mode code: only thread 0 requests collective allocations
assert allocation_request.thread_id == 0, 'collective allocation must come from thread 0'

try:
    run_interpreted(kernel)
except ValueError as e:
    if 'thread_id' in str(e):
        # hoist the allocation out of the per-thread vmap; then retry
        hoist_collective_allocations(kernel)
        run_interpreted(kernel)
    else:
        raise

When it happens

Trigger: Running a Pallas Mosaic GPU kernel in interpretation mode where an allocation with collective/all-threads scope is requested with thread_id != 0 in HostAllocationRequest.

Common situations: Writing custom warp-specialized kernels whose vmap-based thread simulation mistakenly issues collective allocations from every 'thread'; changes in how interpret mode maps num_threads to vmapped axes.

Related errors

• `block_id` must be zero when allocating a buffer for all thr
• num_threads and thread_name must be either both set or both
• Out-of-bounds read of {allocation_key}: reading [{read_range
• Out-of-bounds read of ({device_id} {local_core_id} {memory_s
• masked load_p

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