How could I use a query to calculate the attention with multiple k-v

[DongyuXu77]

In this situation:

Query shape is [batch_size, num_heads_q, seq_length, hidden_size]
Key/Value shape is [batch_size, num_heads_kv, seq_length, top_k, hidden_size]

which means each query vector in a specific position select top_k k-v pairs, how could I use flash-attn to calculate it?

[tridao]

You can expand the query top_k times and combine into the batch dimension, so that QKV have batch dimension batch_size * top_k