this is built on top of #8333

When sharding spec is provided, we also need to shard the segment ID. The data parallel case is the easiest one.

Q: [B, num_head, Q_S,  head_dim]
K/V: [B, num_kv_head, KV_S,  head_dim]

Q_segment_id: [B, Q_S]
K/V_segment_id: [B, KV_S]

in the data parallel(or fsdp in this manner since we will do a all_gather on all parameters which make parameter full), the mesh is 1D like (num_device, ), name=("data") and the sharding spec we passed to flash_attention will be ("data", None, None, None). We just need to shard the segment_id the same way.

The tricky part is what do we save for the backward. I think we need to save the sharded segment_ids. You can imagine that after the enable_manual_sharding all of the computation becomes based on local shape. segment_ids is not the output of the flash_attnetion hence we don't have to bring it back to full. We saved the full_q/k/v but we also used enable_manual_sharding to shard it again.

Note that another tricky part is that q_segment_id is not what we passed to the pallas kernel, we actually add one dimension to it. check