jax-ml/jax · error · ValueError

The sharding must divide the mask blocks evenly between devi

Error message

The sharding must divide the mask blocks evenly between devices

What it means

When using dynamic masks with manual sharding (shard_map), the dynamic mask block grid must be divisible across the devices in the sharding. sharding.shard_shape(block_mask_shape) fails when the per-device split is not even, and splash attention surfaces it as this error.

Source

Thrown at jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py:2486

    )
    with ctx:
      return _splash_attention(
          self.fwd_mask_info,
          self.dq_mask_info,
          self.dkv_mask_info,
          *args,
          **kwargs,
          **self.kwargs,
      )

  def manual_sharding_spec(self, sharding: jax.sharding.NamedSharding):
    """Returns a value that can be used as a shard_map partition spec for the kernel."""
    if self.fwd_mask_info.data_next is not None:
      block_mask_shape = self.fwd_mask_info.data_next.shape
      try:
        shard_shape = sharding.shard_shape(block_mask_shape)
      except ValueError as exc:
        raise ValueError(
            "The sharding must divide the mask blocks evenly between devices"
        ) from exc
      if block_mask_shape[-1] != shard_shape[-1]:
        raise ValueError("Sharding the kv sequence dimension is not supported")
    spec = sharding.spec
    assert len(spec) == 2
    replicated = jax.sharding.PartitionSpec()
    partial_mask_blocks_spec = (
        spec if self.fwd_mask_info.is_dynamic_mask else replicated
    )
    # Shard q_sequence over the sequence dimension only.
    q_sequence_spec = jax.sharding.PartitionSpec(spec[1])
    mask_info_specs = mask_info_lib.MaskInfo(
        data_next=spec if self.fwd_mask_info.data_next is not None else None,  # pyrefly: ignore[bad-argument-type]
        mask_next=spec if self.fwd_mask_info.mask_next is not None else None,  # pyrefly: ignore[bad-argument-type]
        block_mask=spec if self.fwd_mask_info.block_mask is not None else None,  # pyrefly: ignore[bad-argument-type]
        partial_mask_blocks=partial_mask_blocks_spec  # pyrefly: ignore[bad-argument-type]
        if self.fwd_mask_info.partial_mask_blocks is not None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad the sequence length so the mask block grid divides evenly across the sharding axes
2. Adjust the mesh/sharding so each axis divides the corresponding mask block dimension
3. Check mask block counts: q_blocks = ceil(q_len / block_q), kv_blocks = ceil(kv_len / block_kv), both must shard evenly

Example fix

// before
sharding = P('heads', 'q')  # 8 q blocks over 3 devices -> uneven
// after
q_len = 8 * block_q * q_shards  # ensure divisible
sharding = P('heads', 'q')

q_blocks = -(-q_len // block_q); kv_blocks = -(-kv_len // block_kv)
assert q_blocks % q_shards == 0 and kv_blocks % kv_shards == 0

When it happens

Trigger: Calling manual_sharding_spec() (directly or via make_splash_kernel_with_shard_map) with a dynamic mask whose per-head block grid (e.g. q_blocks x kv_blocks) is not divisible by the mesh axes in the sharding spec.

Common situations: Multi-host or multi-chip TPU runs where sequence length or mask block counts don't divide the device count; changing mesh shapes without re-deriving mask shape; using sequence lengths not multiples of num_devices * block sizes.

Related errors

• Sharding the kv sequence dimension is not supported
• {head_shards=} should divide {head_count=}.
• Partitioned callback not implemented on {platform} backend.
• Custom Partitioning rules must return Sharding.
• Mosaic kernels cannot be automatically partitioned. Please w

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