Description:
- Support flash decoding in
CppFlexAttentionTemplate. We prefer to choose flash decoding instead of flash attention when query length is 1. - For flash decoding, we add a kernel option
PARTITION_SIZEto define the partition size of doing the parallelism on KV length dimension. The default value is 128, which should be multiple of KV cache block size to use flash decoding. - As mentioned in #158617,
flex_attnUTs for the cpu backend are disabled because of the long duration. Here we re-enable them on CPU-only machines. (Already merged in #143553)
Performance:
Here are the E2E results for Llama3.1-8B decoding validated on a GNR machine with 6 NUMA nodes, where we can see E2E improvements from 114% to 121%.
| Data Type | Input/Output tokens | Batch Size | W/O Flash Decoding (tokens/s) | With Flash Decoding (tokens/s) | Speedup |
|---|---|---|---|---|---|
| BF16 | 2016/32 | 25 | 892.196 | 1083.073 | 121.39% |
| FP16 | 2016/32 | 25 | 879.541 | 1015.593 | 115.47% |
| BF16 | 1024/128 | 30 | 1291.349 | 1529.251 | 118.42% |
| FP16 | 1024/128 | 30 | 1294.228 | 1473.049 | 113.82% |
Pull Request resolved: #159835
Approved by: https://github.com/CaoE, https://github.com/mingfeima, https://github.com/jansel