Support CP with query length larger than 1

[minosfuture]

This PR implements the causal mask for interleave context parallelism to allow query length > 1.

The solution follows the discussion between @LucasWilkinson , @youkaichao , and @youzhedian on slack.

key illustration made by @LucasWilkinson :

Normal:

k_toks >   0 1 2 3 4 5
q_toks v  _____________
       2 | 1 1 1
       3 | 1 1 1 1
       4 | 1 1 1 1 1
       5 | 1 1 1 1 1 1


DCP Rank 0:

k_toks >   0 2 4
q_toks v  _______
       2 | 1 1
       3 | 1 1
       4 | 1 1 1
       5 | 1 1 1 


DCP Rank 1:

k_toks >   1 3 5
q_toks v   ______
       2 | 1
       3 | 1 1
       4 | 1 1
       5 | 1 1 1

In the DCP case, the k/v tokens are distributed in an interleaved fashion, see vllm-project/vllm#23734.
Therefore we have 0,2,4 kv on rank0 and 1,3,5 kv on rank1 in the example above. The mask shape is no longer a bottom right triangle.
This requires FA to be aware of cp world size and cp rank, in order to determine the causal mask.
The block tiling implementation also needs to be updated. As illustrated below, we now needs to process block tile (0,1) in CP case, while it can be skipped previously in normal case.

Tests

Added and passed unit tests for CP.

[LucasWilkinson]

do you mind removing the unrelated formatting changes? trying to stay as close to upstream as possible when possible

[minosfuture]

Removed.

[LucasWilkinson]

can we use cp_tot_seqlen_k to skip the mul here? should branch in the non-cp case to save the mul?

[LucasWilkinson]

we could make cp_tot_seqlen_k == seqlen_k in the params.cp_world_size == 1 case

[minosfuture]

updated

[LucasWilkinson]

ref: https://github.com/vllm-project/flash-attention/pull/93/files#r2391986149

[LucasWilkinson]

Awesome work! left a few comments but its looking really good!

[LucasWilkinson]

LGTM