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:

ggml/include/ggml/ggml.h

Lines 740 to 751 in 4258875

	GGML_API struct ggml_tensor * ggml_map_unary_f32(
	struct ggml_context * ctx,
	struct ggml_tensor * a,
	const ggml_unary_op_f32_t fun);
	GGML_API struct ggml_tensor * ggml_map_binary_f32(
	struct ggml_context * ctx,
	struct ggml_tensor * a,
	struct ggml_tensor * b,
	const ggml_binary_op_f32_t fun);

[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. ggerganov/llama.cpp#1482