Skip to content

Granite Four #13550

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 149 commits into
base: master
Choose a base branch
from
Draft

Granite Four #13550

wants to merge 149 commits into from

Conversation

gabe-l-hart
Copy link
Contributor

@gabe-l-hart gabe-l-hart commented May 14, 2025

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:

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

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:

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.

compilade added 30 commits April 3, 2024 20:47
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.
gabe-l-hart and others added 9 commits July 2, 2025 12:49
Branch: GraniteFour

Signed-off-by: Gabe Goodhart <[email protected]>
* origin/master:
gguf-py : add support for chat template jinja files (ggml-org#14508)
* origin/master:
Fix conditional enabling following arch checks for ggml-sycl (ggml-org#14504)
convert : correct gemma 3n conversion (ggml-org#14450)
kv-cache : use ggml_set_rows (ggml-org#14285)
ggml : fix FA mask dim 2 and 3 (ggml-org#14505)
ggml : remove kompute backend (ggml-org#14501)
CUDA: add dynamic shared mem to softmax, refactor general usage (ggml-org#14497)
…o GraniteFourWithJamba

* origin/compilade/refactor-kv-cache: (32 commits)
convert : fix jamba conv1d shape squeezing
llama : partially apply clang-format style
llama : remove implicit recurrent state rollbacks
llama : begin renaming llama_past back to llama_kv_cache
llama : use unused n_embd_k_gqa in k_shift
llama : fix mixed signedness comparison
convert_hf : fix Jamba conversion
llama : session saving and reloading for hybrid models
mamba : fix non-contiguous usage of ggml_silu
examples : replace llama_kv_cache_seq_* with llama_past_seq_*
llama : rename llama_cache to llama_past
llama : allow doing the equivalent of SSM_CONV with SUM_ROWS and MUL
llama : fix .base() compilation error on Windows
llama : use im2col and mul_mat to perform convolution for Mamba
llama : avoid copies for simple batch splits
llama : fix edge case finding batch seq_id of split recurrent cell
llama : minimize swaps when reordering logits
llama : fix batch split output count for embeddings
llama : use equal-sequence-length sub-batches for recurrent models
llama : sequence-length-aware batch splitting
...
…id inputs

Branch: GraniteFourWithJamba

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

Signed-off-by: Gabe Goodhart <[email protected]>
…as mixins

The key is for the mixin classes (llm_graph_context_mamba,
llm_graph_context_granite) to use virtual inheritance from
llm_graph_context. This allows the common members to exist only once in the
class hierarchy. The downside is that llm_graph_context will be
re-initialized once for each parent (ie 2x for single mixin, 3x for two
mixins, etc...).

Branch: GraniteFourWithJamba

Signed-off-by: Gabe Goodhart <[email protected]>
@gabe-l-hart
Copy link
Contributor Author

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.

@gabe-l-hart gabe-l-hart marked this pull request as draft July 3, 2025 19:25
compilade and others added 6 commits July 3, 2025 16:04
But this time it contains the sub-cache graph inputs.
This *should* make it easier to handle updating the inputs
when caching the graph (eventually).
…o GraniteFour

* origin/compilade/refactor-kv-cache:
model : add Jamba to Mamba-specific hparams printing
graph : add back hybrid memory graph input
opencl : broadcast for soft_max (ggml-org#14510)
vulkan: support mixed/deepseekR1 FA head sizes (ggml-org#14509)
ggml: backward pass for split swiglu (ggml-org#14483)
Branch: GraniteFour

Signed-off-by: Gabe Goodhart <[email protected]>
* origin/master:
CUDA: add bf16 and i32 to getrows (ggml-org#14529)
vulkan: increase LOAD_VEC_A to 8 (IQ1/IQ2) or 4 (IQ3) (ggml-org#14485)
vulkan: fix rms_norm+mul fusion (ggml-org#14545)
vulkan: Handle updated FA dim2/3 definition (ggml-org#14518)
server : fix assistant prefilling when content is an array (ggml-org#14360)
opencl: add GELU_ERF (ggml-org#14476)
eval-callback : check for empty input (ggml-org#14539)
test-backend-ops: add support for specifying output format (ggml-org#14368)
metal : disable fast math in all quantize kernels (ggml-org#14528)
batch : add optional for sequential equal split (ggml-org#14511)
graph : prepare for 4D mask (ggml-org#14515)
batch : add n_used count (ggml-org#14512)
CANN: Replace aclrtMemsetSync with aclnnInplaceZero operator (ggml-org#14002)
ggml : implement GEGLU_ERF and GEGLU_QUICK ops (ggml-org#14445)
compilade and others added 7 commits July 7, 2025 14:57
…o GraniteFour

* origin/compilade/refactor-kv-cache:
* origin/master:
model : fix hunyuan moe chat template (ggml-org#14584)
model : add SmolLM3 (ggml-org#14581)
memory : fix broken batch splits for recurrent cache (ggml-org#14575)
vulkan : fix rope with partial rotation and non-cont src (ggml-org#14582)
server: Add ability to mount server at prefix (ggml-org#14544)
model : add hunyuan moe (ggml-org#14425)
vulkan: increase timeout for CI (ggml-org#14574)
cuda : fix rope with partial rotation and non-cont src (ggml-org#14580)
CUDA: add bilinear interpolation for upscale (ggml-org#14563)
musa: fix build warnings (unused variable) (ggml-org#14561)
llama : fix incorrect minicpm3 v_states shape (ggml-org#14571)
llama : remove ggml_cont where possible (ggml-org#14568)
gabe-l-hart and others added 4 commits July 8, 2025 14:50
This was already partially supported via reusing the granite ffn builder,
and there may be models that leverage this architecture going forward. The
naming is a bit odd, but in the transformers version, it reuses the same
model class and simply has zero regular experts and a single shared expert
(which is the same as a single dense FFN).

Branch: GraniteFour

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

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

* origin/compilade/refactor-kv-cache:
model : use ggml_swiglu_split for Mamba
model : remove unnecessary prefix for tensor loading constants
jamba : remove redundant nullptr initializations
vulkan: optimize flash attention split_k_reduce (ggml-org#14554)
Signed-off-by: Gabe Goodhart <[email protected]>

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

Co-authored-by: Sigbjørn Skjæret <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Apple Metal https://en.wikipedia.org/wiki/Metal_(API) examples ggml changes relating to the ggml tensor library for machine learning Nvidia GPU Issues specific to Nvidia GPUs python python script changes server testing Everything test related
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants