FlashMLA-2 (CPU): faster and smaller compute buffer size #253
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR improves the CPU implementation of FlashMLA in 3 ways:
To recall, FlashMLA-2 is enabled via
-mla 2 -fa, and is the variant that works on the CPU and on CUDA.The improvement is achieved by adding implementations for
ggml_mul_matwhere the second operand is notfp32ggml_concatwhere the operands are quantizedggml_repeatwhere the operand is notfp32This allows us to avoid conversions to
fp32that can become quite costly when operating on a very large context.Here is a PP performance comparison for DeepSeek-Lite running on a Ryzen-7950X CPU between the main branch and this PR
And here is a comparison between compute buffer sizes along with KV cache size for
fp16cacheI did a quick attempt to also implement on CUDA, but something wasn't working, so left it for a future PR. This also implies that the new way of preparing the compute graph will only be used if the code was built without support for additional back-ends (even if zero layers are uploaded to them, to avoid fighting with the back-end scheduler).