{"record":{"id":"e12ef432ebea8d13","repo":"jax-ml/jax","slug":"q-block-size-should-divide-q-seq-len-per-shard","errorCode":null,"errorMessage":"{q_block_size=} should divide {q_seq_len_per_shard=}.","messagePattern":"(.+?) should divide (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py","lineNumber":380,"sourceCode":"    raise ValueError(f'Expected a bool mask, instead got: {mask.dtype}.')\n\n  head_count, q_seq_len, kv_seq_len = mask.shape\n  q_block_size, kv_block_size = block_shape\n  q_blocks_count, q_mod = divmod(q_seq_len, q_block_size)\n  kv_blocks_count, kv_mod = divmod(kv_seq_len, kv_block_size)\n\n  if q_mod != 0:\n    raise ValueError(f'{q_block_size=} should divide {q_seq_len=}.')\n  if kv_mod != 0:\n    raise ValueError(f'{kv_block_size=} should divide {kv_seq_len=}.')\n\n  q_seq_len_per_shard, mod = divmod(q_seq_len, q_seq_shards)\n  if mod != 0:\n    raise ValueError(f'{q_seq_shards=} should divide {q_seq_len=}.')\n\n  q_blocks_per_shard, mod = divmod(q_seq_len_per_shard, q_block_size)\n  if mod != 0:\n    raise ValueError(f'{q_block_size=} should divide {q_seq_len_per_shard=}.')\n\n  heads_per_shard, mod = divmod(head_count, head_shards)\n  if mod != 0:\n    raise ValueError(f'{head_shards=} should divide {head_count=}.')\n\n  block_mask_shape = (\n      head_count,\n      q_blocks_count,\n      kv_blocks_count,\n  )\n\n  # Tile the last 2 dimensions of the mask into 2D tiles of size `block_shape`.\n  partial_mask_blocks = (\n      mask.reshape(\n          head_count,\n          q_blocks_count,\n          q_block_size,\n          kv_blocks_count,","sourceCodeStart":362,"sourceCodeEnd":398,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py#L362-L398","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Make q_seq_len divisible by q_seq_shards * q_block_size (e.g. pad to a multiple of 128 * q_seq_shards)","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)","Reduce q_seq_shards so the per-shard length becomes block-aligned"],"exampleFix":"# before\nmask = make_splash_attention_mask(q_seq_len=1000, ...)  # per-shard 250 % 128 != 0\n# after\nmask = make_splash_attention_mask(q_seq_len=1024, ...)  # per-shard 256 % 128 == 0","handlingStrategy":"validation","validationCode":"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}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always pad sequence lengths to multiples of q_seq_shards * q_block_size","Centralize block-size constants next to sequence-length config"],"tags":["jax","tpu","splash-attention","shape-mismatch","pallas"],"backgroundTag":"shape-divisibility-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}