Custom kernels, ternary functions?

[cztomsik]

Is there a way how we could easily extend ggml with custom kernels? Or at least with ternary functions?

I've just implemented RWKV model and it's based a lot on this operation:

f(w, x, x_prev) = w * x + x_prev * (1 - w)

As you can imagine, this might be optimized easily into one operation with 3 parameters, and it would only require one tensor as result. Torch does these kinds of things automatically, but to be honest, I wouldn't mind writing such function by hand, it's just that it's not currently possible - or at least I'm not aware of a way to do so.

[ggerganov]

We can extend the API with ggml_map_ternary_f32() similar to the existing ones:

https://github.com/ggerganov/ggml/blob/4258875830cd14e5e464082e9330c7bec5e28186/include/ggml/ggml.h#L740-L751

[cztomsik]

Ok, I think I should be able to do it. Where is the appropriate place to make PR? Here or in llama.cpp?

[ggerganov]

Both places are ok - in llama.cpp it will get extra attention from other devs

[cztomsik]

PR is here, but I have mixed feelings about the result. ggml-org/llama.cpp#1482