[Core] Refactor Attention Take 2

[WoosukKwon]

This PR is the second attempt to modularize the attention backends. The main goal of this PR is to hide any backend-specific attention implementation details from the main logic. This refactoring will greatly help introduce new backends, particularly the FlashInfer backend, which requires a different KV cache layout and input data structures from our current attention backends. NOTE: Since this PR just re-organizes the code, it shouldn't affect the functionality or performance of vLLM.

This PR defines three main classes for each attention backend: AttentionBackend, AttentionMetadata, and AttentionImpl. AttentionBackend (which can be queried by get_attn_backend) is a static class that defines the KV cache layout, swapping ops, and also works as a dispatcher for AttentionMetadata and AttentionImpl. AttentionMetadata is same as the current InputMetadata, but it can be different for other backends added in the future. Finally, AttentionImpl is the actual implementation of the attention operator.

While ultimately I'd like to move part of ModelRunner's prepare_inptus to AttentionBackend, I didn't do it to reduce the size of the PR.

[rkooo567]

Really excited to see this PR being merged! Btw, if you have limited bandwidth, I am willing to help writing tests that combine attention backend + prepare_inputs...!

[richardliaw]

Great to see this merged! Thanks for the work here.

[zhyncs]

This PR is the second attempt to modularize the attention backends. The main goal of this PR is to hide any backend-specific attention implementation details from the main logic. This refactoring will greatly help introduce new backends, particularly the FlashInfer backend, which requires a different KV cache layout and input data structures from our current attention backends. NOTE: Since this PR just re-organizes the code, it shouldn't affect the functionality or performance of vLLM.

This PR defines three main classes for each attention backend: AttentionBackend, AttentionMetadata, and AttentionImpl. AttentionBackend (which can be queried by get_attn_backend) is a static class that defines the KV cache layout, swapping ops, and also works as a dispatcher for AttentionMetadata and AttentionImpl. AttentionMetadata is same as the current InputMetadata, but it can be different for other backends added in the future. Finally, AttentionImpl is the actual implementation of the attention operator.

While ultimately I'd like to move part of ModelRunner's prepare_inptus to AttentionBackend, I didn't do it to reduce the size of the PR.

Hi @WoosukKwon After integrating the FlashInfer backend, how much improvement in overall throughput performance is expected? Thanks.

[zhyncs]

For example, in the current Llama2 7b or 13b model, the time taken for attention calculation accounts for about 1/3 of the total time. If FlashInfer is integrated and assuming a 30% improvement on the kernel compared to the current implementation, would there be an overall throughput increase of close to 10%?