{"record":{"id":"4603d7d166e279b1","repo":"jax-ml/jax","slug":"sharding-the-kv-sequence-dimension-is-not-supporte","errorCode":null,"errorMessage":"Sharding the kv sequence dimension is not supported","messagePattern":"Sharding the kv sequence dimension is not supported","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":2490,"sourceCode":"          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\n        else None,\n        q_sequence=q_sequence_spec  # pyrefly: ignore[bad-argument-type]\n        if self.fwd_mask_info.q_sequence is not None\n        else None,","sourceCodeStart":2472,"sourceCodeEnd":2508,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L2472-L2508","documentation":"Splash attention's manual sharding only supports sharding over heads and the query-sequence block axis. If sharding the last (KV sequence) dimension would split kv blocks (shard_shape changes the last dim), the kernel rejects it because kv blocks must stay whole per device.","triggerScenarios":"Passing a sharding spec to manual_sharding_spec whose last positional axis maps to a meshed device axis, changing block_mask_shape[-1] after sharding.","commonSituations":"Adapting 2D-parallel attention code to splash attention; assuming kv-sequence (context) parallelism works like in ring attention; wrong PartitionSpec argument order.","solutions":["Remove the kv-sequence axis from the sharding spec; shard only heads and q blocks","Use a different attention implementation (ring/pipelined attention) if kv-sequence parallelism is required","Verify block_mask_shape[-1] == shard_shape[-1] before calling"],"exampleFix":"// before\nP('heads', 'kv')\n// after\nP('heads', 'q')  # kv dim must be replicated","handlingStrategy":"type-guard","validationCode":"spec = sharding.spec\nassert spec[-1] is None, 'kv sequence dim must be replicated'","typeGuard":"def is_valid_splash_spec(spec):\n    return len(spec) == 2 and spec[-1] is None  # kv dim replicated","tryCatchPattern":null,"preventionTips":["Shard only heads and q blocks with splash attention","Use ring attention for kv-sequence parallelism"],"tags":["jax","tpu","splash-attention","sharding","partition-spec"],"backgroundTag":"unsupported-sharding-axis","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}