Granite Four

[gabe-l-hart]

Description

This PR is the end-point for architecture support for Granite 4.0 (#13269 . It incorporates a number of changes from other in-flight branches that will need to be merged first:

• Mamba2 model support: llama : initial Mamba-2 support #9126
• Hybrid recurrent cache: Hybrid recurrent cache #13979

Additionally, this PR replaces some work done on other PRs / branches:

• Initial Bamba support: Bamba architecture #10810
  • Bamba is fully supported on this branch, so the old PR can be closed in favor of this PR
• Refactored Bamba support: https://github.com/gabe-l-hart/llama.cpp/tree/BambaArchitectureRefactor
  • I've used this branch as an A/B comparison along the way, but will abandon it now
• Draft Granite 4.0 support: https://github.com/gabe-l-hart/llama.cpp/tree/GraniteFourDraft
  • Build off of the previous Bamba work, this will also be abandoned in favor of this PR
• Initial work on Jamba: llama : support Jamba hybrid Transformer-Mamba models #7531
  • This work is quite out-of-date and would be a lot of work to overhaul to the refactors on master.
  • I had planned to include Jamba support in this branch, but on further inspection, it looks like the Jamba architecture has some additional bells-and-whistles (eg sliding-window-attention) that would need further work, so my plan is to leave Jamba off for now and possibly tackle it later (hopefully it's much easier than the original branch!)

Outstanding Questions

Besides the upstream PRs, there are a few questions to answer before this PR is merge ready:

• This PR contains several changes to llama-kv-cache beyond those in feat: Hybrid unified/recurrent cache #13276, but they depend on the addition of hparams.recurrent_layer_arr which is only populated correctly if there is a valid model architecture to check against. Should I move all of these changes to the hybrid cache PR or keep them here where the model architectures become real?
• Is there a more efficient way to implement hparams.recurrent_layer_arr? Using a max-layer-size std::array doesn't feel quite right.
• There are still some numerical differences between the attention outputs when running Bamba and granite-4.0-tiny-shared-preview on this branch vs the respective draft branches, so I need to determine if this is due to changes in the attention implementation (ie "working as expected") or a bug somewhere.
• The use of dymamic_cast to get the right cache type could be expensive (though it's likely negligible relative to the tensor math). Should we do something more clever to handle different cache types in llama-graph?
• The switch statement for determining the type of KV cache to allocate in llama-model.cpp seems redundant with llama_model_is_recurrent and llama_model_is_hybrid. Should we use those functions instead and eliminate the duplicate logic and additional place to tweak for new recurrent / hybrid models?

Testing

To test out this branch, I've been using the following models:

• granite-4.0-tiny-preview: https://huggingface.co/ibm-granite/granite-4.0-tiny-preview
• Bamba-9B-v1: https://huggingface.co/ibm-ai-platform/Bamba-9B-v1
  • NOTE: v2 is out (here), but I already had v1 from previous branches and stuck with that for consistency
• mamba2-370m-hf: https://huggingface.co/AntonV/mamba2-370m-hf

Details

This PR has a lot of changes in it, some of which are isolated in the prereq-PRs above. In addition to the general mamba2 and llama_kv_cache_hybrid changes, this PR does the following:

python side

• Add conversion support for BambaForCausalLM and GraniteMoeHybridForCausalLM
  • This includes one small tweak to gguf_writer.py that allows duplicate key/value pairs through add_key_value if (and only if) they match both value and type with the existing key. This is a convenience for hybrid models so that the converter doesn't need to rewrite the hparam conversion from multiple parents.
  • This also adds the new HybridAttention section under Keys in constants.py to hold attention.layer_indices. OPEN QUESTION: Should this just go under Attention?

c++ side

• Add a new public API function llama_model_is_hybrid akin to llama_model_is_recurrent
  • I also split up both this function and llama_model_is_recurrent into llm_arch_is_* implemented in llama-arch.* and llama_model_is_* implemented in llama-model.*. This was done so that they could be used during model initialization before the model itself can be passed as the argument, specifically to determine how to populate hparams.recurrent_layer_arr (see below).
• Add hparams.recurrent_layer_arr and support parsing it
  • The current implementation pre-allocates it as a fixed-length array which doesn't feel quite right.
• Add an optional layer id to hparams.n_embd_k_s / hparams.n_embd_v_s
  • This is done because for hybrid models, the values may be different by layer.
  • I plumbed through as many usages of these methods as I could find to properly pass the layer index, but there are some places where it's not available which default to layer 0. This should be fine since none of those places interact with the hybrid caching.
• Add hparams.recurrent_layer(uint32_t) to check whether a given layer is recurrent
• Model name/param/arch plumbing for bamba and granitemoeshared in llama-arch.* (the boring part!)
• (possibly breaking) Add hparams as an additional argument to the llama_model.create_memory method
  • This is done so the hparams can be given to the cache construction and used to determine which layers are recurrent for hybrid cache creation
• In llama-graph, anywhere that a specific cache type needs to be fetched, it is grabbed using new methods get_recurrent_cache / get_unified_cache. These methods use dynamic_cast to handle both non-hybrid caches and hybrid caches.
• Add support for instantiating the hybrid cache in llama-model.cpp
• Add model support for bamba and granitemoehybrid in llama-model
  • Most of this is "business as usual," but that breaks down when trying to avoid code duplication for the hybrid architecture
  • To avoid code duplication, I hoisted build_mamba_layer / build_mamba2_layer from llm_build_mamba and build_attention_layer / build_layer_ffn from llm_build_granite into static methods on their respective classes. This makes for some gross function signatures where member data needs to be explicitly passed, but it allows the hybrid model architecture(s) to use these methods without complex inheritance.
  • I tried an alternative route using diamond inheritance, but this would have required some kind of "don't actually initialize the graph" switch in the parent model builders' constructors to avoid trying to build the parent model graphs during initialization of the hybrid class.

[gabe-l-hart]

@compilade @ggerganov @AnmolS1 I'm going to move the conversation about hybrid cache seg faults over here to avoid cluttering the mamba2 branch since I'm pretty convinced it's a hybrid-only problem.

[gabe-l-hart]

I think the issue may be a logic bug that is somehow being triggered due to the parallel prefill when running without -pps. The logic I think is flawed is:

1. llama_memory_hybrid_context's init_update constructor (here) invokes init_update for both child caches, then calls llama_memory_status_combine on the status values for both children
2. In llama_memory_recurrent, init_update always results in LLAMA_MEMORY_STATUS_NO_UPDATE (here), but in llama_kv_cache_unified, init_update will sometimes result in LLAMA_MEMORY_STATUS_SUCCESS (here)
3. In llama_memory_status_combine, the logic will return LLAMA_MEMORY_STATUS_SUCCESS if either of the children are in the success state
4. The constructor to llama_memory_hybrid_context does not later update the child caches' statuses after combining to set its own

So, the question in my mind is how the two statuses could end up different and what (if any) problems this would cause for the hybrid cache.

[gabe-l-hart]

If I put a conditional check on ctx_recr->get_status() here, it seems to fix the issue. Will keep looking further.

[compilade]

2. In llama_memory_recurrent, init_update always results in LLAMA_MEMORY_STATUS_NO_UPDATE (here), but in llama_kv_cache_unified, init_update will sometimes result in LLAMA_MEMORY_STATUS_SUCCESS (here)

The problem is probably caused when llama_context::kv_self_update is called for the KV cache, but there's nothing to apply for the recurrent cache.

If I put a conditional check on ctx_recr->get_status() here, it seems to fix the issue. Will keep looking further.

It definitely sounds like the recurrent cache's apply should not be called when its state is LLAMA_MEMORY_STATUS_NO_UPDATE.

I wonder if it would be simpler with a separate apply_update method which would unambiguously be a no-op when the sub-cache status is LLAMA_MEMORY_STATUS_NO_UPDATE, so that it could be called from llama_context::kv_self_update. But since the current recurrent cache design never requires separate updates, I guess it could be sufficient to make llama_memory_recurrent_context::apply always be a no-op when status == LLAMA_MEMORY_STATUS_NO_UPDATE.

Not sure if there's a situation where the same problem could happen with iswa caches (i.e. one of the sub-caches needs update while the other doesn't).

[gabe-l-hart]

Ah, yeah, so that makes sense that we could solve this in the recurrent cache as well by simply making it a no-op if status is LLAMA_MEMORY_STATUS_NO_UPDATE rather than asserting success.

[gabe-l-hart]

I'm wondering if the mutex fix that @AnmolS1 did was circumventing this situation where both caches were in different update sates somehow. That and the -pps flag are the only parts of this that still feel mysterious.

[gabe-l-hart]

Ok, yeah, the logic in kv_self_update would definitely trigger this if the two caches had different statuses. I think it makes sense to contain this within the hybrid cache since the problem really lies with the fact that the two are out of sync and it's reasonable for the recurrent cache to assume that its apply will never be called if its status is NO_UPDATE based on the logic in kv_self_update.

I'll open a standalone PR to fix this on master.

[gabe-l-hart]

Fix PR: #14428