{"record":{"id":"80778dfffd5c7a39","repo":"jax-ml/jax","slug":"thread-id-must-be-zero-when-allocating-a-buffer","errorCode":null,"errorMessage":"`thread_id` must be zero when allocating a buffer for all threads","messagePattern":"`thread_id` must be zero when allocating a buffer for all threads","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py","lineNumber":271,"sourceCode":"\n  Args:\n    allocation_request_as_array: Array that converts into an\n      `HostAllocationRequest` with `thread_id`/`block_id` set to zero.\n    value: Array of values to initialize the allocated buffer with.\n\n  Returns:\n    `AllocationKey` to refer to the allocated buffer.\n\n  Raises:\n    ValueError: If `thread_id`/`block_id` in `allocation_request` is not zero.\n  \"\"\"\n  allocation_request = HostAllocationRequest.from_array(\n      allocation_request_as_array\n  )\n  del allocation_request_as_array\n\n  if allocation_request.thread_id != 0:\n    raise ValueError(\n        \"`thread_id` must be zero when allocating a buffer for all threads\"\n    )\n  if allocation_request.block_id != 0:\n    raise ValueError(\n        \"`block_id` must be zero when allocating a buffer for all threads\"\n    )\n  assert allocation_request.memory_space_id != memory.get_memory_space_idx(\n      mosaic_gpu_core.MemorySpace.REGS\n  )\n\n  shared_memory = _get_shared_memory()\n\n  key: HostAllocationKey | None = None\n  buffer_id: int | None = None\n  for thread in shared_memory.concurrent_threads(device):\n    buffer_id_for_thread_id = shared_memory.get_next_buffer_id(thread)\n    if not buffer_id:\n      buffer_id = buffer_id_for_thread_id","sourceCodeStart":253,"sourceCodeEnd":289,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py#L253-L289","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Ensure collective (all-thread) allocations are requested only once per block (thread 0), e.g. allocate outside the per-thread mapped function","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","Try lowering num_threads/avoiding warp-specialized allocations in interpret mode"],"exampleFix":"# before (inside per-thread code):\nbuf = allocate_for_all_threads(...)  # called by every thread\n\n# after (hoisted outside the per-thread body):\nwith collective_scope():\n  buf = allocate_for_all_threads(...)  # requested once by thread 0","handlingStrategy":"validation","validationCode":"# in custom interpret-mode code: only thread 0 requests collective allocations\nassert allocation_request.thread_id == 0, 'collective allocation must come from thread 0'","typeGuard":null,"tryCatchPattern":"try:\n    run_interpreted(kernel)\nexcept ValueError as e:\n    if 'thread_id' in str(e):\n        # hoist the allocation out of the per-thread vmap; then retry\n        hoist_collective_allocations(kernel)\n        run_interpreted(kernel)\n    else:\n        raise","preventionTips":["Issue all-thread allocations once, outside per-thread mapped code","Run interpret mode early on new kernels to catch allocation-scope mistakes","Keep num_threads paths simple in interpret mode"],"tags":["jax","pallas","mosaic-gpu","interpret-mode","allocation","warp-specialization"],"backgroundTag":"invalid-thread-allocation-scope","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}