I think it is not very difficult to provide the indices list (or some other meta information about the mask) without making breaking changes. What we can do is during the construction of the KQ mask, we can create a second tensor that contains non-masked block indices. This tensor could be attached to the flash attention op as an optional src:

cur = ggml_flash_attn_ext(ctx, ...);
ggml_flash_attn_ext_add_idxs(cur, kq_idxs);

The backend can decide wether to use these indices or not. We could also pass a bitmask if that would be more appropriate. Or even both. So from the PoV of ggml and llama.cpp APIs I think we can relatively easy create and provide such additional information to the operator.

On the backend side - if you are seeing significant benefits, I think it would be worth it. I think the existing FA Metal kernel can benefit from a bitmask. Not sure about the list of indices, but once it is provided, we can attempt to write an alternative kernel that uses it.

It turns out that for the new batched inference setup it's faster not to preload any data at all but this is going to hamper the overall performance.

The first test that you did with llama-batched-bench does not benefit at all from the -INF masking optimization because the sequences are always of equal length. But it is performing faster without preloading the mask - that's unexpected. Could this pre-loading not have a significant effect? Can you post the resulting numbers from running the llama-batched-bench on master and this PR to see how the respective numbers change?

What we can do is during the construction of the KQ mask, we can create a second tensor that contains non-masked block indices.

Yes, that is how I would have implemented it as well.

We could also pass a bitmask if that would be more appropriate.

For CUDA a bitmask would not work well because I don't know where in the bit mask I would have to set the start and end indices for evenly distributing data across streaming multiprocessors.

The first test that you did with llama-batched-bench does not benefit at all from the -INF masking optimization because the sequences are always of equal length. But it is performing faster without preloading the mask - that's unexpected. Could this pre-loading not have a significant effect?

For a single sequence the preloading definitely does provide a benefit, the difference is something like 10% end-to-end speedup. Also I just noticed that I forgot that master and the PR are running different kernels for the RTX 4090, master is running the vector kernel, the PR is running the mma kernel. There is also a difference for the RTX 3090, what could be happening there is that due to the lower shared memory use for a single stage pipeline there is some benefit from higher occupancy.

Can you post the resulting numbers from running the llama-batched-bench on master and this PR to see how the respective numbers change?

I'll make a prototype for a kernel using a list of mask slices.

I'm seeing net speedups of ~3%

I think I did the math in my head wrong, sorry.

Try to make the list of indices to be compatible both with unified and split KV cache buffers.

@slaren You mentioned recently that such modification might not be practical (#14363 (comment)). What do you think about the discussed approach above to pass the indices as an optional tensor that the backends can choose to ignore?

Could the backend itself generate the list of indices from the KQ mask in a pre-processing step?

Could the backend itself generate the list of indices from the KQ mask in a pre-processing step?

I considered this but the problem is that I think that that will be inefficient to parallelize. With CUDA I can efficiently generate e.g. a bitmask for which mask slices are all -inf. But a compact list of indices of the active KQ slices needs significantly more communication between CUDA threads.

I think I'll just make a prototype with an extended interface for ggml_flash_attn_ext first. If it turns out that the performance difference is negligible anyways the approach can be scrapped. Otherwise I can implement the generation of the list in the backend and compare vs. generation in llama.cpp.