jax-ml/jax · error · NotImplementedError

Partitioned loads only supported for clusters of size 2. Got

Error message

Partitioned loads only supported for clusters of size 2. Got cluster size {ctx.launch_ctx.cluster_size}.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:989

  is_leader_tracked_copy = collective and leader_tracked is not None
  dst_ty = ir.MemRefType(dst.type)
  bits = math.prod(dst_ty.shape) * mgpu.bitwidth(dst_ty.element_type)
  if bits % 8:
    raise ValueError(
        f"Can only transfer integer bytes (shape={dst_ty.shape},"
        f" dtype={dst_ty.element_type})"
    )
  bytes = bits // 8

  if is_leader_tracked_copy:
    # Leader receives the completion messages from both CTAs.
    bytes *= 2
    if len(collective) != 1:
      raise ValueError(
          f"Expected exactly one collective axis, got {collective_axes=}"
      )
    if math.prod(ctx.launch_ctx.cluster_size) != 2:
      raise NotImplementedError(
          "Partitioned loads only supported for clusters of size 2. Got"
          f" cluster size {ctx.launch_ctx.cluster_size}."
      )

  # TMA is only available on Hopper and newer. On older architectures we fall
  # back to the cp.async implementation.
  if is_cp_async := mgpu.utils.get_arch().major < 9:
    if barrier is not None:
      raise ValueError(
          "copy_gmem_to_smem with a barrier is only supported Hopper and newer"
          " GPUs, which use the TMA implementation"
      )
    if collective_axes is not None:
      raise ValueError("Only the TMA implementation supports collective copies")
    if leader_tracked is not None:
      raise ValueError(
          "Only the TMA implementation supports leader_tracked copies"
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set the cluster dimensions so the product is 2 (e.g. cluster=(2,) or (2,1,1))
2. Remove leader_tracked to use the non-partitioned path
3. Wait for / upgrade to a JAX version supporting larger clusters

Example fix

# before
kernel = pl.pallas_call(..., cluster=(4,1,1))
# after
kernel = pl.pallas_call(..., cluster=(2,1,1))

import math
if leader_tracked is not None:
    assert math.prod(cluster_size) == 2, 'partitioned loads need a 2-CTA cluster'

When it happens

Trigger: 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)).

Common situations: Tuning cluster sizes for newer GPUs; running kernels originally written for 2-CTA multicast on configs with larger clusters.

Related errors

• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Sparse metadata format not implemented for {operand_dtype=}

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