model: M3: Move MSA into a new memory implementation - #26338
Conversation
|
This is better, but we still should not add the |
|
It may worth investigating if we can reuse But I suppose it will be quite complicated due to the |
938640f to
f7e27bd
Compare
|
@ggerganov I removed everything from llama_kv_cache that was added in #24908, I moved the entirety of MSA into the new llama_kv_cache_msa. However I did need to add a small accesor to llama_kv_cache. @ngxson I mirrored llama_kv_cache_dsa fairly closely, with two llama_kv_cache instances, one for K/V and one for the indexer keys as a plain K cache, with every seq op and prepare acted on both so the layouts stay the same. I went with a separate class rather than extending dsa's, because the base caches also differ and the graph inputs diverge, so I think it is cleaner this way. I reused the mapping I had in the previous commit, and the new class just moved where they are populated. This way I could keep the msa_strict_slots logic removed, as block selection can happen in position space. As I mentioned, I did have to add a small get_cells() to the main implementation, with no logic in it. The MSA context uses it to fill the input maps. Alternatively I could try duplicating the cell accounting in the wrapper, but I think that would be a lot more fiddly than this. |
* Move MSA logic from llama-kv-cache into llama-kv-cache-msa * cont : minor * cont : ws fix --------- Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
Overview
Follow up to #24908, addressing #24908 (comment) and #26338 (comment). Everything the original MSA pr added to
llama_kv_cachehas been reverted. MSA is now implemented inside a newllama_kv_cache_msawhich mirrorsllama_kv_cache_dsa, with twollama_kv_cacheinstances, one for K/V and one for the indexer keys as a plain K cache(allocates an unused V cache), kept in sync. I believe a separate class is cleaner than extending dsa's, since the base caches and graph inputs differ.Initially MSA block selection was anchored to cache cell indices, which was only correct for append only streams. This PR also moves the block selection into position space.
The only addition to
llama_kv_cacheis a read only get_cells() accesor, used to fill the MSA position converter maps.None of the changes produced any quality regression up to the 70k ctx I tested. I experienced a minor prefill speed regression, which is most likely the cost of the position -> cell translation. Decode speed remains unchanged.
Additional information
Requirements