jax-ml/jax · error · ValueError

Sharding the kv sequence dimension is not supported

Error message

Sharding the kv sequence dimension is not supported

What it means

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.

Source

Thrown at jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py:2490

          self.dq_mask_info,
          self.dkv_mask_info,
          *args,
          **kwargs,
          **self.kwargs,
      )

  def manual_sharding_spec(self, sharding: jax.sharding.NamedSharding):
    """Returns a value that can be used as a shard_map partition spec for the kernel."""
    if self.fwd_mask_info.data_next is not None:
      block_mask_shape = self.fwd_mask_info.data_next.shape
      try:
        shard_shape = sharding.shard_shape(block_mask_shape)
      except ValueError as exc:
        raise ValueError(
            "The sharding must divide the mask blocks evenly between devices"
        ) from exc
      if block_mask_shape[-1] != shard_shape[-1]:
        raise ValueError("Sharding the kv sequence dimension is not supported")
    spec = sharding.spec
    assert len(spec) == 2
    replicated = jax.sharding.PartitionSpec()
    partial_mask_blocks_spec = (
        spec if self.fwd_mask_info.is_dynamic_mask else replicated
    )
    # Shard q_sequence over the sequence dimension only.
    q_sequence_spec = jax.sharding.PartitionSpec(spec[1])
    mask_info_specs = mask_info_lib.MaskInfo(
        data_next=spec if self.fwd_mask_info.data_next is not None else None,  # pyrefly: ignore[bad-argument-type]
        mask_next=spec if self.fwd_mask_info.mask_next is not None else None,  # pyrefly: ignore[bad-argument-type]
        block_mask=spec if self.fwd_mask_info.block_mask is not None else None,  # pyrefly: ignore[bad-argument-type]
        partial_mask_blocks=partial_mask_blocks_spec  # pyrefly: ignore[bad-argument-type]
        if self.fwd_mask_info.partial_mask_blocks is not None
        else None,
        q_sequence=q_sequence_spec  # pyrefly: ignore[bad-argument-type]
        if self.fwd_mask_info.q_sequence is not None
        else None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the kv-sequence axis from the sharding spec; shard only heads and q blocks
2. Use a different attention implementation (ring/pipelined attention) if kv-sequence parallelism is required
3. Verify block_mask_shape[-1] == shard_shape[-1] before calling

Example fix

// before
P('heads', 'kv')
// after
P('heads', 'q')  # kv dim must be replicated

spec = sharding.spec
assert spec[-1] is None, 'kv sequence dim must be replicated'

def is_valid_splash_spec(spec):
    return len(spec) == 2 and spec[-1] is None  # kv dim replicated

When it happens

Trigger: 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.

Common situations: Adapting 2D-parallel attention code to splash attention; assuming kv-sequence (context) parallelism works like in ring attention; wrong PartitionSpec argument order.

Related errors

• The sharding must divide the mask blocks evenly between devi
• {head_shards=} should divide {head_count=}.
• Partitioned callback not implemented on {platform} backend.
• Length of sharding.spec ({len(out_s.spec)}) must be equal to
• Custom Partitioning rules must return Sharding.

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