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
• llama : support Jamba hybrid Transformer-Mamba models #7531 (needed for further hybrid cache refactors)

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]

Given the changes in #7531 that relate to this same architecture and the draft I have with the mixin pattern and virtual inheritance, I'm bumping this one back behind Jamba in merge order and will update this PR to include the changes in #7531.