feat: Hybrid unified/recurrent cache

[gabe-l-hart]

Description

This implementation covers both llama_memory_i and llama_kv_cache interfaces, but they could very well not be correct.

Discussion

I'm putting this up for discussion even though it doesn't have much value as standalone. My ultimate goal is support for the just-released granite 4 which is a combination of mamba2 and granitemoeshared layers. I opened #13275 to track the full scope of model architecture changes.

[compilade]

Awesome to see this progress!

[compilade]

Yes it will cause problems if this breaks the coherency between caches. (e.g. part of a sequence is removed in one cache but not the other).

This is what I was referring to in #12799 (comment) when I wrote:

The hardest part will be handling errors and properly keeping coherency between the different types of caches (because they don't necessarily roll-back states in the same way).

I think the seq_rm API might fundamentally be too specific to self-attention KV cache. Recurrent models can't rollback their state, because intermediate states are not kept since keeping them for all tokens would take too much space. (when seq_rm returns false, it means the states have to be re-calculated from scratch for the affected sequence (at least that was the intention in #5328))

Ideally, if there was some API to create snapshots and rollback to them, the implementation would be simpler for recurrent models (and for hybrid models by extension). (technically, sequences (with seq_id) already kind of do this (and are copy-on-write), but snapshots within sequences might be more convenient to manage in user code, since managing which state is the latest per sequence could be done transparently)

But that would also mean having to manage the lifetime of explicit state snapshots (in examples/server/server.cpp among others) instead of directly dealing with ranges of token positions (and might make things like largest-common-prefix context caching harder to handle). I've previously shared some ideas about state snapshots/checkpoints in #7531 (comment) (although the first half of the comment is about session restore as in state_read).

[gabe-l-hart]

Ok, interesting. I'm definitely still learning on-the-fly here, but based on this description and the logic here in server.cpp, it seems like the most correct implementation would be to leak implementation details of the child caches or introduce a new member API for can_seq_rm that is const but returns the same logic. I think I'll give that a shot and see how far I can get.

[gabe-l-hart]

Ok, I've pushed an attempt at doing this safely. One thing I noticed is that these mutating methods don't seem to have any sort of locking mechanism, so the way I have it implemented could certainly be prone to thread safety problems if concurrent threads tried to call seq_rm. I don't think this is any different than for the current cache implementations since those would also be sensitive to the same races where the validated condition changes after validation but before the members get mutated, but I wanted to double check if this kind of thread safety is guarded against elsewhere (or just assumed to be handled in the client layer).

[compilade]

Simple split should not be used with recurrent models, they expect equal split.

See #7531 (comment) which illustrates the splits

Suggested change

	// If any of the caches are recurrent, require simple split
	return llama_sbatch(batch, m_hparams.n_embd, m_has_recurrent, logits_all);
	// If any of the caches are recurrent, require non-simple split
	return llama_sbatch(batch, m_hparams.n_embd, !m_has_recurrent, logits_all);

[gabe-l-hart]

Thanks for the comment pointer, this is super helpful for understanding what the consequences of these actually are!

[compilade]

This will not work, recurrent models expect split_equal to be used.

[gabe-l-hart]

Ok, I'm following now. I had them backwards in my head

[compilade]

Maybe this should be if all children can shift, then return true.

But as you've noticed elsewhere, can_shift should technically always be true for all currently-implemented cache types, so I don't know if that part of the API will stay anyway.

[gabe-l-hart]

@compilade I now have a proof-point (#13550) that this works to some extend, though I haven't tested it robustly for edge cases. There are a few additional changes I needed to make on that branch that should maybe come over to this branch, but it gets a little hairy because they interact with adding the actual model architectures. Some possible paths forward:

1. Review this one as-is and take the updates from the Granite 4 branch when they come
2. Try to move the changes to hparams to determine per-layer recurrence over to this branch, and then just have unreachable code when initializing the hybrid cache in llama-model.cpp
3. Scrap this PR and just review all hybrid cache changes with the incoming models in the Granite 4 branch

Thoughts / Preferences?

[gabe-l-hart]

I thought about it further and decided that the cleanest separation between this and the Granite 4 branch is to pull over the key parts of hparam parsing and kv cache instantiation and isolate the mismatch to the llm_arch_is_hybrid function which will always return false for now until there are real hybrid models.

[gabe-l-hart]

I just realized that this is entirely unnecessary since llama_model already has hparams as a member variable 🤦

[gabe-l-hart]

@compilade @ggerganov I'm going to keep working on some unit tests for the different cache types and try to sort through the CI failure. In the mean time, would it be possible to get an initial review to make sure I'm on the right track?

[gabe-l-hart]

It looks like the failing tests were all on Windows and therefore fall into this note about not running tests using internal APIs on Windows due to the lack of LLAMA_API exports.

[ggerganov]

I think the hybrid implementation can become similar to the new iSWA implementation - mainly for consistency.

I would like to now simplify the llama_kv_cache interface, namely:

• Try to eliminate set_full() completely. We can instead clear() + init batches with appropriate sizes
• Rework the sbatch_init(), ubatch_next() and find_slot(). It seems that this should be possible to be replaced with a single init(llama_batch) call that would prepare all the ubatches beforehand

Will open a draft PR and we can discuss details there.

Adding tests would be really useful - we can merge these from a separate PR without too much delays.

[gabe-l-hart]

@ggerganov Thank you for pointing that out! I'll work to rebase this change on the latest iSWA changes. I'll also try to pull out the unit tests into their own PR.

[gabe-l-hart]

@ggerganov Ported PR with minimal unit tests: #13669

I didn't try to expand them to incorporate the iSWA cache yet, but will look to do that as I go through the changes there and refactor this hybrid cache branch.

[gabe-l-hart]

@ggerganov This branch is now rebuilt based on the iSWA cache. I kept most of the implementation of llama_kv_cache_hybrid the way it was where the sub-types of the child caches are not specified to hopefully make it easier to extend this to other per-layer hybrid styles in the future. Like before, I've plumbed through the instantiation of the hybrid cache to the point where any model that answers true to llm_arch_is_hybrid_recurrent will instantiate a hybrid cache in create_memory, but currently there are no models that fit that, so it will not be exercised until the next PR with the concrete hybrid architectures.

[gabe-l-hart]

This is buggy because push_back is used below, so later when we index directly into the given per-layer tensor vectors, the final layers will be out-of-bounds reads.

[gabe-l-hart]

I've updated the Granite 4 / Bamba branch (#13550) to test this out and found a couple of bugs that should now be fixed and functional.