I pulled these into the class so that they can be set differently by derived conversion classes and then used in the common methods below

This has been one of the most annoying changes keeping this branch up to date: The GRANITE_MOE_HYBRID name is two characters longer than the previous longest name, so to keep vertical alignment, it changes the indentation of all values (here and in llama-arch.cpp).

I had originally extracted these as standalone methods so that I could reuse them in the hybrid implementation. Ultimately, any inheritance / static method / mixin approach I tried felt too tangled, so tangled, so I went back to simply duplicating these methods in the hybrid model. I left these separated out here for the symmetry and encapsulation, but I could also revert this set of changes to llm_build_granite to keep the changeset smaller.

This is a copy-paste from llm_build_mamba. Per the other comment, it got too tangled to try to reliably reuse these across model builders. That said, I would still love to find a way to avoid this kind of duplication if there's appetite.

@gabe-l-hart I might have found a way to avoid this kind of duplication, see src/llama-model.cpp and src/llama-graph.cpp in #7531

llm_graph_context_mamba is a child class of llm_graph_context with Mamba-specific layer builders

llm_graph_context_mamba is the parent class of llm_build_mamba and llm_build_jamba. Not sure if that would still be appropriate with multiple-inheritance, though that wasn't necessary for Jamba.

The methods could potentially be moved to llm_graph_context (in src/llama-graph.cpp), but I preferred to keep model-specific graph building methods in src/llama-model.cpp, at least for now.

Note that I've also removed build_inp_mem_hybrid and llm_graph_input_hybrid in favor of directly using the recurrent and self-attention input builders separately. This is relatively clean, I think.

build_inp_rs and build_inp_attn_kv_unified accept an optional mctx override argument.

This makes use of the fact that mctx is stored in inp_rs and inp_attn already, and so build_rs and build_attn were changed to use that instead of trying to cast llm_graph_context::mctx again.

llm_graph_context_mamba is a child class of llm_graph_context with Mamba-specific layer builders

Ah, very nice. This is quite close to something else I tried earlier using the mixin approach. You can see it here:

• llm_build_mamba_mixin: https://github.com/gabe-l-hart/llama.cpp/blob/ed150125d9778ec92c25bd64e258397923d04892/src/llama-model.cpp#L9345-L9359
• llm_build_mamba: https://github.com/gabe-l-hart/llama.cpp/blob/ed150125d9778ec92c25bd64e258397923d04892/src/llama-model.cpp#L9628
• llm_build_granite_mixin: https://github.com/gabe-l-hart/llama.cpp/blob/ed150125d9778ec92c25bd64e258397923d04892/src/llama-model.cpp#L13022
• llm_build_granite: https://github.com/gabe-l-hart/llama.cpp/blob/ed150125d9778ec92c25bd64e258397923d04892/src/llama-model.cpp#L13181
• llm_build_hybrid_mamba (later renamed to llm_build_granite_hybrid): https://github.com/gabe-l-hart/llama.cpp/blob/ed150125d9778ec92c25bd64e258397923d04892/src/llama-model.cpp#L13257

(sidebar, what is the magical GH markdown syntax to get those nice multi-line reference blocks to render inline??)

I think the only real difference here is that the mixins couldn't inherit from llm_graph_context in order to support multiple inheritance and avoid double-initialization. Unfortunately, this meant that mixin classes had to hold a pointer to an llm_graph_context (the eventual child) and the implementations of the shared methods needed to to self->... everywhere instead of just using shared member methods from llm_graph_context.

The other approach I tried was to simply make the layer-builder methods static so that they can simply be called by other classes. I don't have a good point in history to show this because I stopped being super careful about maintaining a working version of this during rebasing, but you can see the gist of it here for the static builder method and here for calling it. On that hash, I think I accidentally dropped the static on the mamba layer builders (it was there at some point), but the idea was the same. Just like in the mixin approach, in order to avoid multiple-inheritance with multi-initialization, the "fix" was to pass this as an argument to the static methods and have them do self->... everywhere in their body.

In both cases, the tradeoff I struggled to make was forcing the implementation of the layer builder to know that it would be used by hybrid models. I'm certainly not against this, but it felt like a bit of an inversion of responsibility (the "child" in the hybrid relationship becoming the "parent" in the implementation hierarchy). I'm very open to revisiting this, especially the mixin approach since I think that is probably the cleanest way to handle this from a design pattern perspective. The only other thought I had was to use a templated generic hybrid builder, but I know templates are generally avoided in this project.

Note that I've also removed build_inp_mem_hybrid and llm_graph_input_hybrid in favor of directly using the recurrent and self-attention input builders separately. This is relatively clean, I think.

build_inp_rs and build_inp_attn_kv_unified accept an optional mctx override argument.

Nice, this feels much better. I had also tried something like this, but @ggerganov had suggested it was best to start by following the build_attn pattern strictly. I like the idea of revisiting this now that we've sorted out some of the other issues with the hybrid cache.

The methods could potentially be moved to llm_graph_context (in src/llama-graph.cpp), but I preferred to keep model-specific graph building methods in src/llama-model.cpp, at least for now.

I thought about this one too, and I agree that putting mamba or granite methods in the upstream base class is not correct. I also tried using virtual inheritance to enable diamond-style multiple inheritance without multiple-initialization of the common base class, but wasn't able to get it working (I can't quite remember why now).

(sidebar, what is the magical GH markdown syntax to get those nice multi-line reference blocks to render inline??)

It's a permalink to ranges of lines on its own Markdown line. I think it has to be from the same repository, though (so only code from https://github.com/ggml-org/llama.cpp (including branches) will have pretty previews here).

Ok, I think I got the mixin pattern while also using diamond inheritance so the mixins can inherit from llm_graph_context. The only change on top of the changes in #7531 is to make llm_graph_context_mamba use virtual inheritance from llm_graph_context (struct llm_graph_context_mamba : public virtual llm_graph_context {). This then means that all classes that inherit from one or more of these mixins need to explicitly initialize llm_graph_context as well as the mixin(s). This does result in some duplicate initialization since the initializer list and constructor of llm_graph_context will still be called once for each parent in the diamond (so twice if single inheritance, three times if diamond inheritance). This should be a negligible runtime hit (especially if we can avoid re-creating res since that's the only heap allocation), but it will still end up setting all of the members multiple times. If we can make this all compatible with #14482, then the hit should be trivial since it will not happen for every token.

I've got a full implementation of this with #7531 merged in https://github.com/gabe-l-hart/llama.cpp/blob/GraniteFourWithJamba. Depending on how we want to order the merges, I'm happy to update this branch to include that refactor and put it behind Jamba for merge.

• Example of mixin with virtual inheritance: https://github.com/gabe-l-hart/llama.cpp/blob/GraniteFourWithJamba/src/llama-model.cpp#L10027-L10028
• Example of multiple inheritance: https://github.com/gabe-l-hart/llama.cpp/blob/GraniteFourWithJamba/src/llama-model.cpp#L14020-L14029

I dug a bit more and I actually think virtual inheritance is exactly what we want and there will be zero duplicate initialization (though there would be duplicate construction if there were a body to the constructor for llm_graph_context). Here's a dummy sample:

// Example program
#include <iostream>
#include <string>

struct Member {
    Member(int i_) : i(i_) {
        std::cout << "Constructing Member with " << i << std::endl;
    }
    
    const int i;
};

struct Base {
    Base(int i) : m(i) {}

    Member m;
};

struct MixinA : public virtual Base {
    MixinA() : Base(1) {}
    
    void do_a() const {
        std::cout << "Doing A: " << m.i << std::endl;
    }
};

struct MixinB : public virtual Base {
    MixinB() : Base(2) {}
    
    void do_b() const {
        std::cout << "Doing B: " << m.i << std::endl;
    }
};

struct Kid : public MixinA, public MixinB {
    Kid() : Base(3), MixinA(), MixinB() {}
    
    void doit() const {
        do_a();
        do_b();
    }
};

int main()
{
    Kid k;
    k.doit();
}

When run, this outputs:

Constructing Member with 3
Doing A: 3
Doing B: 3

This indicates that the explicit initialization of Base in Kid is used, but that the initialization of Base in MixinA and MixinB is ignored which is exactly what we want here.