{"record":{"id":"b8a9e14f52789bfa","repo":"jax-ml/jax","slug":"warp-reduction-group-size-should-be-a-power-of-2","errorCode":null,"errorMessage":"Warp reduction group size should be a power of 2 (got {group_size})","messagePattern":"Warp reduction group size should be a power of 2 \\(got (.+?)\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":1810,"sourceCode":"  for s, t in zip(shape[-tiling_rank:], tiling):\n    if s % t:\n      raise ValueError(\"Non-divisible tiling:\", shape, tiling)\n  return (\n      *shape[:-tiling_rank],\n      *(s // t for s, t in zip(shape[-tiling_rank:], tiling)),\n      *tiling,\n  )\n\n\ndef warp_tree_reduce(value, op, group_size):\n  \"\"\"Reduce a value across the warpgroup.\"\"\"\n  assert bytewidth(value.type) == 4\n  assert 32 % group_size == 0 and group_size <= 32\n  i32 = ir.IntegerType.get_signless(32)\n  result = value\n  iters = np.log2(group_size)\n  if not iters.is_integer():\n    raise ValueError(\n        f\"Warp reduction group size should be a power of 2 (got {group_size})\"\n    )\n  iters = int(iters)\n  for i in range(iters):\n    other_result = nvvm.shfl_sync(\n        c(0xFFFFFFFF, i32),\n        result,\n        c(1 << i, i32),\n        c(0x1F, i32),\n        nvvm.ShflKind.bfly\n    )\n    result = op(result, other_result)\n\n  return result\n\n\n_MEMORY_SPACES = {f\"#gpu.address_space<{str(x)}>\": x for x in gpu.AddressSpace}\n","sourceCodeStart":1792,"sourceCodeEnd":1828,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L1792-L1828","documentation":"Warp reduction (warp_reduce) uses log2(group_size) shuffle iterations, which only works when group_size is a power of two. The preceding assert also requires group_size to divide 32; np.log2 returning a non-integer triggers this ValueError.","triggerScenarios":"Calling the warp reduction helper with group_size values like 3, 6, 12, 24 — any non-power-of-2 value (the function also requires 32 % group_size == 0 and group_size <= 32, so effectively only 1,2,4,8,16,32 are valid).","commonSituations":"Deriving group_size from warp/thread configuration arithmetic (e.g. num_threads // something) that lands on 24 or 12; porting CUDA code that used arbitrary subgroup sizes; changing the number of lanes per worker in a Mosaic kernel.","solutions":["Round group_size to the nearest power of two <= 32 (1, 2, 4, 8, 16, 32)","Check the computation of group_size (often threads_per_warp or an inferred worker count) upstream","If you need non-power-of-2 groups, pad the group to the next power of 2 and mask out inactive lanes"],"exampleFix":"# before\ngroup_size = num_workers // 3  # e.g. 24 -> raises\n# after\ngroup_size = 1 << (num_workers // 3 - 1).bit_length()  # round up to power of 2","handlingStrategy":"validation","validationCode":"def valid_group_size(g):\n    return g in (1, 2, 4, 8, 16, 32)\nassert valid_group_size(group_size), f'group_size {group_size} not a power of 2 <= 32'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Compute group sizes via 1 << k only","Add a startup assertion on warp config derived from thread counts"],"tags":["mosaic-gpu","warp-reduction","power-of-two"],"backgroundTag":"invalid-group-size","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}