Granite Four #13550

gabe-l-hart · 2025-05-14T20:13:13Z

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.

This will be necessary to support Jamba (and other recurrent models mixed with Attention). Doesn't compile yet, and finding a slot isn't yet done correctly for recurrent states.

* llama : begin work on support for variable GQA This will also be useful for Jamba if we consider the Mamba layers to have 0 KV heads. * llama : gracefully fail when not finding hybrid slot

* ggml : simplify SSM-related operators * llama : make recurrent state slot allocation contiguous * llama : adapt internal uses of batches to llama_ubatch

This reduces overhead when running hellaswag on thousands of sequences with very small 100k params Mamba models.

This otherwise was a problem when running the HellaSwag benchmark with small batch sizes, making it crash.

This removes the need for ggml_ssm_conv!!! But performance seems slighly worse on my system, especially for prompt processing. Maybe ggml_mul_mat isn't optimized for small row sizes? More performance testing is necessary until GGML_OP_SSM_CONV is removed. * ggml : make ggml_ssm_scan not modify its source tensors * llama : fix shared recurrent tail cell count for small ubatch sizes Otherwise it was impossible to run the 'parallel' example with '-ub 1' with a Mamba or Jamba model.

* ggml : allow GGML_OP_CONCAT to work on non-contiguous tensors The implementation already supported it, and this makes Mamba's conv step slightly faster.

This can be changed back later if the name change is wrong. I was renaming the functions anyway to generalize kv-cache-related functions to hybrid and recurrent model architectures. I think llama_past is a better name than llama_cache for a combined kv cache and recurrent state cache, because the states it contains pretty much always come before the newly-added ones for any particular sequence. Also 'llama_past_clear' sounds more obvious in what it does than 'llama_kv_cache_clear'. The future is what the models generate. (For embeddings, the kv cache isn't really used anyway) Still, I'm open to better suggestions.

Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

Signed-off-by: Gabe Goodhart <[email protected]> Co-authored-by: Sigbjørn Skjæret <[email protected]>

…o GraniteFour * origin/compilade/refactor-kv-cache: memory : avoid referring to KV in recurrent cache logs model : make falcon-h1 use shared mamba2 layer builder Signed-off-by: Gabe Goodhart <[email protected]>

Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

Some of the tensor names are common with Llama4

…o GraniteFour * origin/compilade/refactor-kv-cache: gguf-py : avoid adding duplicate tensor mappings for Jamba

* origin/master: llama : support Jamba hybrid Transformer-Mamba models (ggml-org#7531) ggml : add ggml_scale_bias (ggml-org#14417)

gabe-l-hart · 2025-07-09T19:18:45Z

@compilade @ggerganov @CISC I think this one is ready to go now!

The main outstanding questions I have with the changes here are:

Since GRANITE_MOE_HYBRID is two characters longer than the longest previous enum key, all the vertical alignment changed in constants.py and llama-arch.cpp. Would it be better to not do this and just have one that doesn't vertically align to avoid the bigger diff?
How do you feel about the use of virtual inheritance to implement the diamond pattern with llm_build_granite_hybrid -> (llm_graph_context_mamba, llm_graph_context_granite) -> llm_graph_context?

If we prefer to focus on composition over inheritance, we could move to a hasa relationship and remove llm_graph_context from the inheritance tree of the llm_graph_context_mamba and llm_graph_context_granite classes. Those would then need to be renamed and take a llm_graph_context & (or llm_graph_context *) as a member, then use it in their layer builder methods (eg graph_ctx->build_attn(...)). Then, the top-level llm_build_granite_hybrid (and llm_build_mamba / llm_build_jamba / llm_build_falcon_h1) would have member instances instantiated with this.

gabe-l-hart · 2025-07-09T19:25:33Z

src/llama-arch.h

@@ -52,6 +52,7 @@ enum llm_arch {
    LLM_ARCH_MAMBA2,
    LLM_ARCH_JAMBA,
    LLM_ARCH_FALCON_H1,
+    LLM_ARCH_BAMBA,


Looking over the code again, I think we could probably collapse LLM_ARCH_BAMBA into LLM_ARCH_GRANITE_MOE_HYBRID

This might then suggest we change GRANITE_MOE_HYBRID to simply GRANITE_HYBRID (which would have the nice benefit of removing the extra indentation!)

The only hang up collapsing these is how to determine when to use rope. Right now, the only way to tell is by the architecture name (bamba uses it, granitehybrid doesn't). This could be handled with an hparam, but I'm not clear exactly which one to use. I see rope_finetuned that doesn't appear to actually be used anywhere which seems like it might be a good option, but the default value is false which is reported as "unknown", so it wouldn't be a perfect fit. I also see n_no_rope_layer_step which is used for this in llama4, but from what I can tell, there's no corresponding constant for plumbing this through via conversion.

ggerganov · 2025-07-09T19:39:00Z

How do you feel about the use of virtual inheritance to implement the diamond pattern with llm_build_granite_hybrid -> (llm_graph_context_mamba, llm_graph_context_granite) -> llm_graph_context?

I was just looking at that - we should avoid this. I haven't used this pattern and I am not sure I understand how it works.

In general, code de-duplication is not an objective here. Even with composition. It's completely fine to copy paste the model architectures instead of reusing them with inheritance. Pretty much all the architectures in llama-model.cpp could be deduplicated, but we don't do it because it's much easier to understand the flow of the networks.

What we typically do for deduplication is to extract common blocks into llm_graph_context. But usually this is done when we see a good pattern.

The llm_graph_context_mamba and llm_build_rwkv6_base are temporary - I think we should eventually move the code of these helper classes into llm_graph_context, by passing all the necessary model tensors as arguments. This way all architectures will be consistently implemented as child of llm_graph_context. This is a well-defined pattern and it's easy to follow.

gabe-l-hart · 2025-07-09T19:50:07Z

That makes sense. For now, I think I'll remove the virtual inheritance and inherit only from llm_graph_context_mamba, then duplicate the granite helpers to keep inheritance linear.

I was just looking at that - we should avoid this. I haven't used this pattern and I am not sure I understand how it works.

I also did not understand this at all until I tried to get it working here. I didn't fully trust it until putting together a dummy version: #13550 (comment)

CISC · 2025-07-09T19:58:51Z

1. Since `GRANITE_MOE_HYBRID` is two characters longer than the longest previous enum key, all the vertical alignment changed in `constants.py` and `llama-arch.cpp`. Would it be better to not do this and just have one that doesn't vertically align to avoid the bigger diff?

It's not an issue, it's easy to filter out by hiding whitespace.

The only key difference is the use of rope which is now set via rope_finetuned in the hparams Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

Per PR discussion, it's simpler to keep this with basic inheritance and not introduce the complexity of virtual inheritance and multiple inheritance ggml-org#13550 (comment) Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

* origin/master: llama : remove llm_graph_input_one (ggml-org#14603) Signed-off-by: Gabe Goodhart <[email protected]>

gabe-l-hart · 2025-07-09T21:11:16Z

I've removed the virtual inheritance now and collapsed Bamba and GraniteMoeHybrid into simply GraniteHybrid which covers all permutations of hybrid architectures (w/ and w/out the granite multipliers on top of llama) and dense/MoE/MoE+shared.

src/llama-model.cpp

Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

This matches how recurrent vs attention heads are identified for Jamba Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

ggerganov

Can merge after @compilade approves.

src/llama-model-loader.cpp

CISC · 2025-07-10T08:30:59Z

src/llama-model.cpp

+                ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);
+                hparams.rope_finetuned = rope_finetuned;
+
+                // A layer is recurrent IFF the n_head_kv value is set to 0


Suggested change

// A layer is recurrent IFF the n_head_kv value is set to 0

// A layer is recurrent IF the n_head_kv value is set to 0

I actually meant IFF as in if and only if. Happy to change it if that's too obscure though

Heh, never heard of that abbreviation, one lives and learns... :)

Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

* origin/master: cmake : do not search for curl libraries by ourselves (ggml-org#14613) SYCL: Initial set_rows kernel implementation (ggml-org#14562) llama : minor coding style fix for smollm3 (ggml-org#14605) cmake : bump llguidance version to v1.0.1 (ggml-org#14609) cmake : llguidance build parser library only (ggml-org#14608) cuda : support Falcon-H1 state size for SSM_SCAN (ggml-org#14602) Signed-off-by: Gabe Goodhart <[email protected]>

compilade added 30 commits April 3, 2024 20:47

wip: llama : separate recurrent states from the KV cache

271104c

This will be necessary to support Jamba (and other recurrent models mixed with Attention). Doesn't compile yet, and finding a slot isn't yet done correctly for recurrent states.

llama : use std::find for seq_nodes in llama_rs_cache

8db1e4d

llama : state checkpoints for recurrent models

0028010

llama : correctly handle more edge cases for the rs cache

0c8b3b2

Merge branch 'master' into compilade/refactor-kv-cache

d66849f

llama : rename many llama_kv_cache_* functions

a09db95

Merge branch 'master' into compilade/refactor-kv-cache

c460ff1

llama : remove useless return value for some llama_cache_* functions

b6fafd1

Merge branch 'master' into compilade/refactor-kv-cache

b7ec12e

Merge branch 'master' into compilade/refactor-kv-cache

3b57b55

llama : rethink recurrent state cell counts

7e13f19

* llama : begin work on support for variable GQA This will also be useful for Jamba if we consider the Mamba layers to have 0 KV heads. * llama : gracefully fail when not finding hybrid slot

llama : support Jamba

cbc743e

Merge branch 'master' into compilade/refactor-kv-cache

0fd13e9

llama : fix BERT inference without KV cache

61a88a1

convert-hf : check for unprocessed Jamba experts

ea2e63e

convert-hf : support Mini-Jamba conversion

fc59407

llama : fix Jamba quantization sanity checks

181dadf

llama : sequence-length-aware batch splitting

3a414b0

Merge branch 'master' into compilade/refactor-kv-cache

4e4c41e

llama : use equal-sequence-length sub-batches for recurrent models

3587a94

* ggml : simplify SSM-related operators * llama : make recurrent state slot allocation contiguous * llama : adapt internal uses of batches to llama_ubatch

Merge branch 'master' into compilade/refactor-kv-cache

5d3c7b9

llama : fix batch split output count for embeddings

72eea49

llama : minimize swaps when reordering logits

18d1c14

This reduces overhead when running hellaswag on thousands of sequences with very small 100k params Mamba models.

llama : fix edge case finding batch seq_id of split recurrent cell

61200ef

This otherwise was a problem when running the HellaSwag benchmark with small batch sizes, making it crash.

llama : avoid copies for simple batch splits

eb589d5

llama : fix .base() compilation error on Windows

17f6c1e

llama : allow doing the equivalent of SSM_CONV with SUM_ROWS and MUL

fee3c1d

* ggml : allow GGML_OP_CONCAT to work on non-contiguous tensors The implementation already supported it, and this makes Mamba's conv step slightly faster.

Merge branch 'master' into compilade/refactor-kv-cache

6840ac0

gabe-l-hart and others added 2 commits July 9, 2025 07:57

fix: Use child inputs for Falcon H1 after merge resolution

f8b81c0

Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

fix: Remove unnecessary prefix on tensor constants

0583d95

Signed-off-by: Gabe Goodhart <[email protected]> Co-authored-by: Sigbjørn Skjæret <[email protected]>

gabe-l-hart force-pushed the GraniteFour branch from 4b93a2a to 0583d95 Compare July 9, 2025 14:04

compilade and others added 7 commits July 9, 2025 10:06

model : make falcon-h1 use shared mamba2 layer builder

7f3955a

memory : avoid referring to KV in recurrent cache logs

452207f

Merge remote-tracking branch 'origin/compilade/refactor-kv-cache' int…

fa159cf

…o GraniteFour * origin/compilade/refactor-kv-cache: memory : avoid referring to KV in recurrent cache logs model : make falcon-h1 use shared mamba2 layer builder Signed-off-by: Gabe Goodhart <[email protected]>

fix: Revert order changes for Falcon H1 to stay consistent with upstream

44cda75

Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

gguf-py : avoid adding duplicate tensor mappings for Jamba

4d6a179

Some of the tensor names are common with Llama4

Merge remote-tracking branch 'origin/compilade/refactor-kv-cache' int…

d1d54d8

…o GraniteFour * origin/compilade/refactor-kv-cache: gguf-py : avoid adding duplicate tensor mappings for Jamba

Merge remote-tracking branch 'origin/master' into GraniteFour

fe34d0e

* origin/master: llama : support Jamba hybrid Transformer-Mamba models (ggml-org#7531) ggml : add ggml_scale_bias (ggml-org#14417)

gabe-l-hart marked this pull request as ready for review July 9, 2025 19:09

gabe-l-hart commented Jul 9, 2025

View reviewed changes

gabe-l-hart added 3 commits July 9, 2025 15:01

refactor: Collapse Bamba and GraniteMoeHybrid into GraniteHybrid

6875697

The only key difference is the use of rope which is now set via rope_finetuned in the hparams Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

refactor: Remove use of diamond inheritance

8dd7f97

Per PR discussion, it's simpler to keep this with basic inheritance and not introduce the complexity of virtual inheritance and multiple inheritance ggml-org#13550 (comment) Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

Merge remote-tracking branch 'origin/master' into GraniteFour

2b36420

* origin/master: llama : remove llm_graph_input_one (ggml-org#14603) Signed-off-by: Gabe Goodhart <[email protected]>

compilade reviewed Jul 9, 2025

View reviewed changes

src/llama-model.cpp Outdated Show resolved Hide resolved

src/llama-model.cpp Outdated Show resolved Hide resolved

src/llama-model.cpp Show resolved Hide resolved

gabe-l-hart added 3 commits July 9, 2025 21:06

feat: Log mamba params for Granite Hybrid

dcf51e0

Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

fix: Remove unused ssm_in_b

5b44f4e

Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

refactor: Remove ATTENTION_LAYER_INDICES hparam in favor of n_head_kv

4e9fef1

This matches how recurrent vs attention heads are identified for Jamba Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

ggerganov approved these changes Jul 10, 2025

View reviewed changes

ggerganov requested review from compilade and CISC July 10, 2025 06:14

CISC approved these changes Jul 10, 2025

View reviewed changes

gabe-l-hart added 2 commits July 10, 2025 07:26

fix: Remove unused template expansion for get_arr

d02d3dd

Branch: GraniteFour Signed-off-by: Gabe Goodhart <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Granite Four #13550

Granite Four #13550

gabe-l-hart commented May 14, 2025 •

edited

Loading

Uh oh!

gabe-l-hart commented Jul 9, 2025

Uh oh!

gabe-l-hart Jul 9, 2025

Uh oh!

gabe-l-hart Jul 9, 2025

Uh oh!

gabe-l-hart Jul 9, 2025

Uh oh!

ggerganov commented Jul 9, 2025

Uh oh!

gabe-l-hart commented Jul 9, 2025

Uh oh!

CISC commented Jul 9, 2025

Uh oh!

gabe-l-hart commented Jul 9, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ggerganov left a comment

Uh oh!

Uh oh!

CISC Jul 10, 2025

Uh oh!

gabe-l-hart Jul 10, 2025

Uh oh!

CISC Jul 10, 2025

Uh oh!

Uh oh!

	// A layer is recurrent IFF the n_head_kv value is set to 0
	// A layer is recurrent IF the n_head_kv value is set to 0

Granite Four #13550

Are you sure you want to change the base?

Granite Four #13550

Conversation

gabe-l-hart commented May 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Outstanding Questions

Testing

Details

python side

c++ side

Uh oh!

gabe-l-hart commented Jul 9, 2025

Uh oh!

gabe-l-hart Jul 9, 2025

Choose a reason for hiding this comment

Uh oh!

gabe-l-hart Jul 9, 2025

Choose a reason for hiding this comment

Uh oh!

gabe-l-hart Jul 9, 2025

Choose a reason for hiding this comment

Uh oh!

ggerganov commented Jul 9, 2025

Uh oh!

gabe-l-hart commented Jul 9, 2025

Uh oh!

CISC commented Jul 9, 2025

Uh oh!

gabe-l-hart commented Jul 9, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ggerganov left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

CISC Jul 10, 2025

Choose a reason for hiding this comment

Uh oh!

gabe-l-hart Jul 10, 2025

Choose a reason for hiding this comment

Uh oh!

CISC Jul 10, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

gabe-l-hart commented May 14, 2025 •

edited

Loading