ggml : add ggml_set_rows

[rgerganov]

Add ggml_set_rows(a, b, c) which copies rows from 'b' into 'a' using indices from 'c'.

ref: #8366

[ggerganov]

So far so good: #14285

I think the ggml_set_rows() alone could be a very useful addition since this mechanism can make the llama_kv_cache_unified::find_slot() to search not just for continuous slots of KV cells, but effectively be able to "scatter" the ubatch. This would be a useful improvement, regardless if the graph reuse works or not, so I think we should proceed to implement this operator.

[ggerganov]

We should aim to reuse the existing cpy/dup code in order to support F32 -> any type, not just F16.

[rgerganov]

Do we need support for quantized data types or just F16,F32,BF16?

[ggerganov]

We should support all possible KV cache quantization types:

llama.cpp/common/arg.cpp

Lines 819 to 830 in 28ee6d2

	const std::vector<ggml_type> kv_cache_types = {
	GGML_TYPE_F32,
	GGML_TYPE_F16,
	GGML_TYPE_BF16,
	GGML_TYPE_Q8_0,
	GGML_TYPE_Q4_0,
	GGML_TYPE_Q4_1,
	GGML_TYPE_IQ4_NL,
	GGML_TYPE_Q5_0,
	GGML_TYPE_Q5_1,
	};

Ideally, it should work for all ggml types that also work with ggml_cpy.

[ggerganov]

We might want to allow broadcasting c into b. It would avoid this ggml_repeat_4d here:

llama.cpp/src/llama-kv-cache-unified.cpp

Lines 795 to 799 in a0c0fb6

	v_cur = ggml_cont_3d(ctx, v_cur, 1, v_cur->ne[0], v_cur->ne[1]);
	kv_idxs = ggml_repeat_4d(ctx, kv_idxs, v_cur->ne[1], v_cur->ne[2], 1, 1);
	return ggml_set_rows(ctx, v_view, v_cur, kv_idxs);