question about the multi_query_attention_rag_buffer()

[bianshuai001]

the implementation of the kernel function multi_query_attention_rag_buffer() seems to follow a similar approach to flash decoding, but during decoding, there is also a branch that directly calls flash decoding. What are the differences between the implementation of multi_query_attention_rag_buffer() and flash decoding?

[unix1986]

The multi_query_attention_rag_buffer() is a variable-length multi-query decoding optimization kernel independently developed by our team earlier. It demonstrates impressive performance on certain hardware configurations, particularly excelling in scenarios with small batch sizes and context lengths below 4k tokens, showing significant advantages on specific GPU models. However, on other hardware platforms like the NVIDIA A800, its performance falls short of flash-decoding implementations when handling long sequence lengths (e.g., >4k tokens) and large batch sizes. To address this limitation, we later integrated the flash-decoding technique, which can be activated by setting the environment variable USE_FA_DECODING=1.