jax-ml/jax · error · ValueError

{q_block_size=} should divide {q_seq_len_per_shard=}.

Error message

{q_block_size=} should divide {q_seq_len_per_shard=}.

What it means

Splash Attention (TPU Pallas) requires that the query sequence length per shard be evenly divisible by the query block size. During dynamic-mask processing, mask metadata is tiled into blocks of q_block_size, so a non-divisible shard length makes block tiling impossible. The library raises ValueError eagerly before kernel compilation to fail fast on incompatible shapes.

Source

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

    raise ValueError(f'Expected a bool mask, instead got: {mask.dtype}.')

  head_count, q_seq_len, kv_seq_len = mask.shape
  q_block_size, kv_block_size = block_shape
  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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make q_seq_len divisible by q_seq_shards * q_block_size (e.g. pad to a multiple of 128 * q_seq_shards)
2. Set q_block_size to a divisor of q_seq_len_per_shard (valid splash block sizes are typically powers of two: 64, 128, 256)
3. Reduce q_seq_shards so the per-shard length becomes block-aligned

Example fix

# before
mask = make_splash_attention_mask(q_seq_len=1000, ...)  # per-shard 250 % 128 != 0
# after
mask = make_splash_attention_mask(q_seq_len=1024, ...)  # per-shard 256 % 128 == 0

assert q_seq_len % (q_seq_shards * q_block_size) == 0, f'q_seq_len {q_seq_len} not divisible by {q_seq_shards}*{q_block_size}'

When it happens

Trigger: Calling make_splash_attention_mask / splash_attention_kernel (or mha on TPU) with a dynamic mask whose q_seq_len divided by q_seq_shards is not a multiple of q_block_size. Typical mismatch: q_seq_len=1000, q_seq_shards=4 gives q_seq_len_per_shard=250, not divisible by q_block_size=128.

Common situations: Using non-standard sequence lengths (e.g. 1000, 3000) with default block sizes (128/256); changing the number of TPU shards or mesh q-sharding without adjusting seq len or block size; packing ragged sequences into a padded length that is not shard-and-block aligned.

Related errors

• Can only load scalars from SMEM
• Expected value and mask to have the same shape, but got valu
• Can only store scalars to SMEM
• The number of sources must match the packing factor ({packin
• The shape of the accumulator {acc_shape} is not 2-dimensiona

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