jax-ml/jax · error · ValueError

{head_shards=} should divide {head_count=}.

Error message

{head_shards=} should divide {head_count=}.

What it means

Splash Attention shards attention heads across TPU devices; the number of heads must be divisible by the number of head shards. The mask-processing code computes heads_per_shard = head_count // head_shards and requires zero remainder so each shard gets an equal number of heads.

Source

Thrown at jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py:384

  q_blocks_count, q_mod = divmod(q_seq_len, q_block_size)
  kv_blocks_count, kv_mod = divmod(kv_seq_len, kv_block_size)

  if q_mod != 0:
    raise ValueError(f'{q_block_size=} should divide {q_seq_len=}.')
  if kv_mod != 0:
    raise ValueError(f'{kv_block_size=} should divide {kv_seq_len=}.')

  q_seq_len_per_shard, mod = divmod(q_seq_len, q_seq_shards)
  if mod != 0:
    raise ValueError(f'{q_seq_shards=} should divide {q_seq_len=}.')

  q_blocks_per_shard, mod = divmod(q_seq_len_per_shard, q_block_size)
  if mod != 0:
    raise ValueError(f'{q_block_size=} should divide {q_seq_len_per_shard=}.')

  heads_per_shard, mod = divmod(head_count, head_shards)
  if mod != 0:
    raise ValueError(f'{head_shards=} should divide {head_count=}.')

  block_mask_shape = (
      head_count,
      q_blocks_count,
      kv_blocks_count,
  )

  # Tile the last 2 dimensions of the mask into 2D tiles of size `block_shape`.
  partial_mask_blocks = (
      mask.reshape(
          head_count,
          q_blocks_count,
          q_block_size,
          kv_blocks_count,
          kv_block_size,
      )
      .swapaxes(-2, -3)
      .astype(np.bool_)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Choose num_attention_heads divisible by the head-sharding mesh axis (commonly a power of two)
2. Adjust the mesh so the head axis divides num_heads
3. Pad/interpolate head count via tensor parallel replication only if the model permits it

Example fix

# before
num_heads = 12; mesh = Mesh(jax.devices(), ('heads',))  # 12 % 4==ok, 12 % 8 != 0
# after
num_heads = 16; mesh = Mesh(jax.devices(), ('heads',))  # 16 % 4 == 0

assert num_heads % head_shards == 0, f'{num_heads=} not divisible by {head_shards=}'

When it happens

Trigger: Running splash attention on a TPU mesh where the head-sharding axis size does not divide the number of attention heads, e.g. 8 heads sharded over 3 head-shards, or 12 heads with mesh axis 8.

Common situations: Configuring a TPU pod slice or multi-device mesh whose axis length (e.g. 3, 6) doesn't divide the model's num_heads; porting a model with an unusual head count (e.g. 12, 48) to a mesh sized for powers of two.

Related errors

• The sharding must divide the mask blocks evenly between devi
• Sharding the kv sequence dimension is not supported
• {q_block_size=} should divide {q_seq_len_per_shard=}.
• Partitioned callback not implemented on {platform} backend.
• Length of sharding.spec ({len(out_s.spec)}) must be equal to

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