Hybrid recurrent cache

[gabe-l-hart]

This is a re-opened version of #13904 after #13746 was merged

Description

This PR introduces the llama_kv_cache_hybrid_recurrent cache implementation. It follows the pattern of llama_kv_cache_unified_iswa by holding two child cache instances and implementing the interface logic such that it manages both correctly for the appropriate layers.

Changes

The main change in this PR is the addition of llama_kv_cache_hybrid_recurrent in llama-kv-cache-hybrid-recurrent.*. In addition to this, the PR does the following:

• Add the llama_model_is_hybrid_recurrent public API (akin to llama_model_is_recurrent)
• Add LLM_KV_ATTENTION_LAYER_INDICES as an hparam to hold the indices of the layers that should use attention (versus recurrent)
  • This part is not well aligned with iswa, but that mechanism also isn't particularly extensible. It might be more appropriate to have a generic mechanism for indicating the type of caching to use for each layer, but that would start to approach the generic hybrid implementation that I originally attempted which ended up being too abstract (feat: Hybrid unified/recurrent cache #13276).
• Abstracting utilities in llm_graph_context that need a specific type of cache to use getters (get_state_unified / get_state_recurrent) that will properly handle llama_kv_cache_hybrid_recurrent
• Make n_embd_k_s / n_embd_v_s layer-dependent and use layer indices when calling them in the existing cache implementations
• Add layer filtering to llama_kv_cache_recurrent
• Updates the logic in llama_model::create_memory to use llm_arch_is_recurrent and llm_arch_is_hybrid_recurrent rather than relying on adding models to the switch statement which was redundant with the implementation of these functions

[ggerganov]

These dynamic casts should not be necessary. Instead you need a new llm_graph_context::build_attn_inp_kv_hybrid_recurrent() method, similar to llm_graph_context::build_attn_inp_kv_unified_iswa().

[gabe-l-hart]

I'm working through this now and a couple of questions are coming up:

1. Would it be best to combine build_inp_s_copy with build_attn_inp_kv for hybrid so that models call just one "build inputs" function, or keep them separate for simplicity?
2. For the build_attn methods, each has a corresponding llm_graph_input_attn_* class. The build_inp_s_* methods don't have this pattern which would make this a bit harder to have code reuse. Are there plans to refactor that further @compilade?
3. In the mamba2 branch, s_mask seems to be totally removed. I'd prefer not to do all of the boilerplate for duplicating build_inp_s_mask for the hybrid recurrent case if that's definitely going to be going away. Is there any reason that might stick around past the merge of mamba2?

[compilade]

Answering out of order, but it should still make sense:

2. For the build_attn methods, each has a corresponding llm_graph_input_attn_* class. The build_inp_s_* methods don't have this pattern

They do follow this pattern, see

llama.cpp/src/llama-graph.h

Line 190 in 7e00e60

class llm_graph_input_s_copy : public llm_graph_input_i {

(this is on the current master)

I think you might mean the build_attn_* methods also return instances of llm_graph_input_attn_*?
That seems to be directly related to llm_graph_context::build_attn() which has multiple implementations which differ by the type of the first argument (e.g. for llm_graph_input_attn_kv_unified, llm_graph_input_attn_no_cache, etc.)

Are there plans to refactor that further @compilade?

Not really, outside of removing s_mask (and related functions and classes) as part of #13834.

1. Would it be best to combine build_inp_s_copy with build_attn_inp_kv for hybrid so that models call just one "build inputs" function, or keep them separate for simplicity?

Personally, I think it would be simpler to keep them separate, because they are fundamentally different (one is intended to be used by build_copy_mask_state (renamed to build_recurrent_state in #13834), while the other is used by build_attn), and they are pretty much independent, even in hybrid models (at least for Jamba, the recurrent and self-attention layers are mostly independent on that front).

I don't see how build_attn would ever need s_copy.

build_inp_s_copy and build_inp_attn_kv_* are called once at the beginning of the graph, while build_attn and build_recurrent_state are called once per layer (where applicable, and so usually different layers for both).

3. Is there any reason [s_mask] might stick around past the merge of mamba2?

No reason to keep it, s_mask will be removed. Its functionality is redundant with s_copy, and otherwise prevents minimizing unnecessary state copies. It was used to clear the states, but the same can be done through inp_s_copy and clearing by copying a zero-ed state (which is the rs_z'th state in the mamba2 branch (and #13834)).

[gabe-l-hart]

Thanks, that's super helpful! I was missing the distinction between build_attn_inp and build_attn which makes perfect sense.

Personally, I think it would be simpler to keep them separate

I agree on my personal gut feeling, so I'll go with this and see how it feels once complete.

[gabe-l-hart]

Ok, I think this feels a lot cleaner now. For build_inp_s_copy, I opted to add an optional parameter so that the caller can take ownership of casting the cache state rather than duplicating the function into build_inp_s_copy_hybrid. That felt a little cleaner w.r.t. code reuse, but I'm happy to do a separate method if that's preferred.

[gabe-l-hart]

It looks like there's one more place that will need changing in build_copy_mask_state (renamed to build_recurrent_state on mamba2). Similar to build_inp_s_copy, I think the cleanest way to do this for code reuse is to add an optional parameter that, if unset, will use the current logic of casting mstate.

[compilade]

Yes, but that will be fixed in #13834

(Noting here in case this gets merged first so that I don't forget to update the comment)

[gabe-l-hart]

Should I strip out this TODO at this point?

[gabe-l-hart]

@ggerganov I've noticed that the Server tests are consistently failing on this branch, but I see them passing on other PRs (eg #14081). The failures seem to be around retried connections to /slots/0?action=restore (here). I've noticed a couple of things when trying to repro locally:

1. For all of the server tests, despite the health check ping loop, I need to manually inject a sleep after server.start() to get any of the tests to pass on my M3 Max 64GB, otherwise the first endpoint call will almost always return an unexpected result (invalid json, or 5xx error). To me, this speaks to the server reporting /health success before it's logically "ready." This is probably correct in that the server is healthy enough to take requests, but the tests are treating this as a readiness check which would indicate that the server is fully initialized which seems incorrect (similar to liveness vs readiness in kubernetes). I wasn't able to find any open issues related to this, so I'd be curious if there's a discussion somewhere of readiness vs liveness?

2. Once I do insert the manual sleep in test_slot_save_restore, the test passes locally for me. This speaks to some difference between my local system and the GH Action runners. I'm not familiar enough with the server code to know whether this would be caused by some kind of memory constraint, disk speed constraint, etc, but my guess would be something like slot1.bin not being fully written somehow. This wouldn't explain why it passes on other branches though. Any insight on the nature of these tests would be much appreciated as I try to debug!

[gabe-l-hart]

I've tried adding retry logic for all requests in 39a93b3 to work around the race between /health, but I'm not sure if this is just going to mask an underling issue.

[ggerganov]

The changes to the server tests should not be needed. Let's revert the commit for now and I'll investigate.

[gabe-l-hart]

@ggerganov Thanks, it looks like those changes didn't fix the failures anyway, so definitely not the right fix. I've reset them out and can open an issue with details of what I see locally

[gabe-l-hart]

Issue for follow up on /health race condition with tests: #14092

[gabe-l-hart]

I've rebased on #13834. Drafting for now until it's merged

[gabe-l-hart]

That was quick! Undrafting now that #13834 is merged

[gabe-l-hart]

I think I've tracked the mysterious break in the Granite 4 rebase chain to something in this method. What I've found is that for the attention layers, the kq_mask is different between my (old) working branch and my rebased broken one. That, in turn led me to llama_kv_cache_unified::set_input_kq_mask which was last modified in #14130 which also had some changes in init_batch. I'm still fumbling around understanding how all the batching works, but it feels like something I'm doing here is probably incorrect for how the unified cache expects its ubatches to show up.

@compilade @ggerganov any chance anything jumps out at you? I think it might be related to #13746 (comment) from several PRs ago?

[gabe-l-hart]

No idea if this is relevant, but if I'm reading this line correctly, I think the index into data changed in #14130 because you end up multiplying s*n_seq_tokens by n_kv whereas it wasn't multiplied before the change.

(changing this back did not fix my issue, so it's not a root cause)

[gabe-l-hart]

I also see #14119 touching the same areas, so it's also possibly related

[gabe-l-hart]

I've tried running a non-hybrid model with the unified cache and attempted to refactor the unified cache to use split equal. The model still works, so I don't think it's a split_simple vs split_equal problem.

[gabe-l-hart]

I've debugged the bad output to this line. The previous op (here) has matching inputs before/after the breakage, but the output of the ggml_soft_max_ext op is different.

I've stepped through the op in cpu mode and I'm pretty confident I've verified that the tensor data in kq_mask is different. In the not-broken run, the only cells with 0 instead of -Inf are in the first position of each row, whereas in the broken run, there are 0s in other positions (eg both 4096 and 4097 whereas 4097 is -Inf in the not-broken run).

This leaves me pretty sure that something in how the kq_mask is initialized changed. The "broken" version is https://github.com/gabe-l-hart/llama.cpp/blob/broken/GraniteFour/src/llama-graph.cpp (based off of ed52f36 on master) while the "working" version is https://github.com/gabe-l-hart/llama.cpp/blob/GraniteFour/src/llama-graph.cpp (based off of f470bc3 on master).

[compilade]

@gabe-l-hart There might be a problem here regarding the attention state: n_kv and head are not really set here (while in llama_kv_cache_unified_state::apply(), they normally are, but that method isn't called here).

But it's not as simple as calling the apply() methods of the child states here because of how they are initialized by the hybrid cache (minimally, which makes the update() step of the unified cache not work). Maybe that initialization needs to be revisited.

[ggerganov]

That's correct. At least the llama_kv_cache_unified_state state has to be initialized in a similar way as it is done in llama_kv_cache_unified_iswa_state:

llama.cpp/src/llama-kv-cache-unified-iswa.cpp

Lines 216 to 232 in fb85a28

	llama_kv_cache_unified_iswa_state::llama_kv_cache_unified_iswa_state(
	llama_kv_cache_unified_iswa * kv,
	llama_sbatch sbatch,
	std::vector<uint32_t> heads_base,
	std::vector<uint32_t> heads_swa,
	std::vector<llama_ubatch> ubatches)
	: status(LLAMA_MEMORY_STATUS_SUCCESS),
	sbatch(std::move(sbatch)),
	ubatches(std::move(ubatches)) {
	// note: here we copy the ubatches. not sure if this is ideal
	state_base.reset(new llama_kv_cache_unified_state(kv->get_base(), {}, std::move(heads_base), this->ubatches));
	state_swa .reset(new llama_kv_cache_unified_state(kv->get_swa (), {}, std::move(heads_swa), this->ubatches));
	status = llama_memory_status_combine(state_base->get_status(), state_swa->get_status());
	}