jax-ml/jax · error · ValueError

`block_id` must be zero when allocating a buffer for all thr

Error message

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

What it means

GPU interpretation-mode counterpart of the thread_id check: HostAllocationRequest.block_id must be 0 for buffers allocated for all threads. An all-thread collective allocation must be anchored at block 0; any other block indicates invalid allocation bookkeeping in the interpreted kernel.

Source

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

    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
    else:
      # We keep the buffer ids in sync across all threads. This implies, in
      # particular, that every instance of the assignment to `key` below assigns
      # an `AllocationKey` object with the same attributes.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Anchor collective allocations at block 0 / outside per-block mapped code
2. If constructing HostAllocationRequest manually, set block_id=0 (and thread_id=0) for all-thread allocations
3. Minimize the kernel to isolate the offending allocation and report upstream if compiler-emitted

Example fix

# before
req = HostAllocationRequest(..., thread_id=tid, block_id=bid)

# after
req = HostAllocationRequest(..., thread_id=0, block_id=0)  # collective alloc

assert allocation_request.block_id == 0, 'collective allocation must come from block 0'

try:
    run_interpreted(kernel)
except ValueError as e:
    if 'block_id' in str(e):
        anchor_collective_allocation_at_block_zero(kernel)
        run_interpreted(kernel)
    else:
        raise

When it happens

Trigger: Interpreting a kernel that issues a for-all-threads allocation from a simulated block other than block 0 (non-zero block_id in HostAllocationRequest).

Common situations: Simulating multi-block kernels with collective allocations; custom allocation callbacks constructed manually with wrong block_id; version changes in the interpret backend's block simulation.

Related errors

• `thread_id` must be zero when allocating a buffer for all th
• 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
• run_scoped_p with collective axes is not supported

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