{"record":{"id":"706f9891c1dad171","repo":"jax-ml/jax","slug":"the-sharding-must-divide-the-mask-blocks-evenly-be","errorCode":null,"errorMessage":"The sharding must divide the mask blocks evenly between devices","messagePattern":"The sharding must divide the mask blocks evenly between devices","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":2486,"sourceCode":"    )\n    with ctx:\n      return _splash_attention(\n          self.fwd_mask_info,\n          self.dq_mask_info,\n          self.dkv_mask_info,\n          *args,\n          **kwargs,\n          **self.kwargs,\n      )\n\n  def manual_sharding_spec(self, sharding: jax.sharding.NamedSharding):\n    \"\"\"Returns a value that can be used as a shard_map partition spec for the kernel.\"\"\"\n    if self.fwd_mask_info.data_next is not None:\n      block_mask_shape = self.fwd_mask_info.data_next.shape\n      try:\n        shard_shape = sharding.shard_shape(block_mask_shape)\n      except ValueError as exc:\n        raise ValueError(\n            \"The sharding must divide the mask blocks evenly between devices\"\n        ) from exc\n      if block_mask_shape[-1] != shard_shape[-1]:\n        raise ValueError(\"Sharding the kv sequence dimension is not supported\")\n    spec = sharding.spec\n    assert len(spec) == 2\n    replicated = jax.sharding.PartitionSpec()\n    partial_mask_blocks_spec = (\n        spec if self.fwd_mask_info.is_dynamic_mask else replicated\n    )\n    # Shard q_sequence over the sequence dimension only.\n    q_sequence_spec = jax.sharding.PartitionSpec(spec[1])\n    mask_info_specs = mask_info_lib.MaskInfo(\n        data_next=spec if self.fwd_mask_info.data_next is not None else None,  # pyrefly: ignore[bad-argument-type]\n        mask_next=spec if self.fwd_mask_info.mask_next is not None else None,  # pyrefly: ignore[bad-argument-type]\n        block_mask=spec if self.fwd_mask_info.block_mask is not None else None,  # pyrefly: ignore[bad-argument-type]\n        partial_mask_blocks=partial_mask_blocks_spec  # pyrefly: ignore[bad-argument-type]\n        if self.fwd_mask_info.partial_mask_blocks is not None","sourceCodeStart":2468,"sourceCodeEnd":2504,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L2468-L2504","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Pad the sequence length so the mask block grid divides evenly across the sharding axes","Adjust the mesh/sharding so each axis divides the corresponding mask block dimension","Check mask block counts: q_blocks = ceil(q_len / block_q), kv_blocks = ceil(kv_len / block_kv), both must shard evenly"],"exampleFix":"// before\nsharding = P('heads', 'q')  # 8 q blocks over 3 devices -> uneven\n// after\nq_len = 8 * block_q * q_shards  # ensure divisible\nsharding = P('heads', 'q')","handlingStrategy":"validation","validationCode":"q_blocks = -(-q_len // block_q); kv_blocks = -(-kv_len // block_kv)\nassert q_blocks % q_shards == 0 and kv_blocks % kv_shards == 0","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pad sequence lengths so block counts divide mesh axes","Re-derive mask shape after changing the mesh"],"tags":["jax","tpu","splash-attention","sharding","dynamic-mask"],"backgroundTag":"sharding-not-divisible","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}