ggml : support broadcast for ggml_soft_max_ext and ggml_flash_attn_ext

[ggerganov]

Extract broadcast changes from #14363 for ggml_soft_max_ext() and ggml_flash_attn_ext():

llama.cpp/ggml/include/ggml.h

Lines 1435 to 1451 in 236682a

	// a [ne0, ne01, ne02, ne03]
	// mask [ne0, ne11, ne12, ne13] | ne11 >= ne01, F16 or F32, optional
	//
	// broadcast:
	// ne02 % ne12 == 0
	// ne03 % ne13 == 0
	//
	// fused soft_max(a*scale + mask*(ALiBi slope))
	// max_bias = 0.0f for no ALiBi
	GGML_API struct ggml_tensor * ggml_soft_max_ext(
	struct ggml_context * ctx,
	struct ggml_tensor * a,
	struct ggml_tensor * mask,
	float scale,
	float max_bias);

llama.cpp/ggml/include/ggml.h

Lines 1876 to 1896 in 236682a

	// q: [n_embd_k, n_batch, n_head, ne3]
	// k: [n_embd_k, n_kv, n_head_kv, ne3]
	// v: [n_embd_v, n_kv, n_head_kv, ne3] !! not transposed !!
	// mask: [n_kv, n_batch_pad, ne32, 1] !! n_batch_pad = GGML_PAD(n_batch, GGML_KQ_MASK_PAD) !!
	// res: [n_embd_v, n_head, n_batch, ne3] !! permuted !!
	//
	// broadcast:
	// n_head % n_head_kv == 0
	// ne3 % ne32 == 0
	//
	GGML_API struct ggml_tensor * ggml_flash_attn_ext(
	struct ggml_context * ctx,
	struct ggml_tensor * q,
	struct ggml_tensor * k,
	struct ggml_tensor * v,
	struct ggml_tensor * mask,
	float scale,
	float max_bias,
	float logit_softcap);

Both changes should be quite simple. On master we have the assumption that the mask is a 2D matrix and we always broadcast it across the dim 2 (i.e. the heads) and dim 3. With this change, we allow to have separate masks - i.e. generalized broadcast.

Currently, I've added tests and implemented the CPU and Metal to support this. The rest of the backends will fallback to CPU, until this gets implemented:

• CPU
• Metal
• CUDA
• Vulkan
• CANN

Fallback is okay for now since these extensions are not used at the moment by llama.cpp. This support will be needed later for the #14363 PR, although it's better to support this either way.

[jeffbolznv]

I've started working on the Vulkan backend support for this.

[jeffbolznv]

Vulkan support is in #14449, targeted to this branch.