{"record":{"id":"631838a2a6a719bd","repo":"jax-ml/jax","slug":"head-shards-should-divide-head-count","errorCode":null,"errorMessage":"{head_shards=} should divide {head_count=}.","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":384,"sourceCode":"  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,\n          kv_block_size,\n      )\n      .swapaxes(-2, -3)\n      .astype(np.bool_)","sourceCodeStart":366,"sourceCodeEnd":402,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py#L366-L402","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Choose num_attention_heads divisible by the head-sharding mesh axis (commonly a power of two)","Adjust the mesh so the head axis divides num_heads","Pad/interpolate head count via tensor parallel replication only if the model permits it"],"exampleFix":"# before\nnum_heads = 12; mesh = Mesh(jax.devices(), ('heads',))  # 12 % 4==ok, 12 % 8 != 0\n# after\nnum_heads = 16; mesh = Mesh(jax.devices(), ('heads',))  # 16 % 4 == 0","handlingStrategy":"validation","validationCode":"assert num_heads % head_shards == 0, f'{num_heads=} not divisible by {head_shards=}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pick head counts divisible by the head-sharding mesh axis","Validate mesh vs model shape in a config test"],"tags":["jax","tpu","splash-attention","sharding","shape-mismatch"],"backgroundTag":"shape-divisibility-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}