{"record":{"id":"b86a8816c1b4304d","repo":"jax-ml/jax","slug":"partitioned-loads-only-supported-for-clusters-of-s","errorCode":null,"errorMessage":"Partitioned loads only supported for clusters of size 2. Got cluster size {ctx.launch_ctx.cluster_size}.","messagePattern":"Partitioned loads only supported for clusters of size 2\\. Got cluster size (.+?)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":989,"sourceCode":"  is_leader_tracked_copy = collective and leader_tracked is not None\n  dst_ty = ir.MemRefType(dst.type)\n  bits = math.prod(dst_ty.shape) * mgpu.bitwidth(dst_ty.element_type)\n  if bits % 8:\n    raise ValueError(\n        f\"Can only transfer integer bytes (shape={dst_ty.shape},\"\n        f\" dtype={dst_ty.element_type})\"\n    )\n  bytes = bits // 8\n\n  if is_leader_tracked_copy:\n    # Leader receives the completion messages from both CTAs.\n    bytes *= 2\n    if len(collective) != 1:\n      raise ValueError(\n          f\"Expected exactly one collective axis, got {collective_axes=}\"\n      )\n    if math.prod(ctx.launch_ctx.cluster_size) != 2:\n      raise NotImplementedError(\n          \"Partitioned loads only supported for clusters of size 2. Got\"\n          f\" cluster size {ctx.launch_ctx.cluster_size}.\"\n      )\n\n  # TMA is only available on Hopper and newer. On older architectures we fall\n  # back to the cp.async implementation.\n  if is_cp_async := mgpu.utils.get_arch().major < 9:\n    if barrier is not None:\n      raise ValueError(\n          \"copy_gmem_to_smem with a barrier is only supported Hopper and newer\"\n          \" GPUs, which use the TMA implementation\"\n      )\n    if collective_axes is not None:\n      raise ValueError(\"Only the TMA implementation supports collective copies\")\n    if leader_tracked is not None:\n      raise ValueError(\n          \"Only the TMA implementation supports leader_tracked copies\"\n      )","sourceCodeStart":971,"sourceCodeEnd":1007,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L971-L1007","documentation":"Leader-tracked partitioned loads are currently only implemented for CUDA clusters of exactly 2 CTAs; the code checks math.prod(cluster_size) == 2 and raises NotImplementedError otherwise, since the message-counting logic (bytes *= 2) assumes two participants.","triggerScenarios":"copy_gmem_to_smem with leader_tracked/partitioned collectives on a launch whose cluster dims multiply to something other than 2 (e.g. cluster=(4,) or cluster=(2,2)).","commonSituations":"Tuning cluster sizes for newer GPUs; running kernels originally written for 2-CTA multicast on configs with larger clusters.","solutions":["Set the cluster dimensions so the product is 2 (e.g. cluster=(2,) or (2,1,1))","Remove leader_tracked to use the non-partitioned path","Wait for / upgrade to a JAX version supporting larger clusters"],"exampleFix":"# before\nkernel = pl.pallas_call(..., cluster=(4,1,1))\n# after\nkernel = pl.pallas_call(..., cluster=(2,1,1))","handlingStrategy":"validation","validationCode":"import math\nif leader_tracked is not None:\n    assert math.prod(cluster_size) == 2, 'partitioned loads need a 2-CTA cluster'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pin cluster dims to product 2 for leader-tracked kernels","Skip leader_tracked on configs with larger clusters"],"tags":["jax","pallas","mosaic-gpu","cluster-size","leader-tracked","not-implemented"],"backgroundTag":"unsupported-hardware-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}