llama.cpp
80d2c0c2 - metal : fix GROUP_NORM for ne03 > 1 and non-contiguous inputs (#31852)

Commit
1 day ago
metal : fix GROUP_NORM for ne03 > 1 and non-contiguous inputs (#31852) The kernel_group_norm_f32 kernel and its host dispatch ignored the batch dimension (ne03) of the source tensor: only ngrp threadgroups were dispatched and each normalized the flat range [tgpig*gs, (tgpig+1)*gs) with gs = ne00*ne01*ceil(ne02/ngrp). For tensors with ne03 > 1, all but the first ne00*ne01*ne02 elements were left unwritten (exposed by the new GROUP_NORM test with ne=[3,1,1,65537], num_groups=1). Dispatch one threadgroup per (group, i03) pair (ngrp in x, ne03 in z, same work distribution as the Vulkan backend) and compute per-batch statistics, matching the CPU reference. Also clamp the group element count for the last group when ne02 is not divisible by ngrp and use int64_t indexing. The kernel now only requires ggml_is_contiguous_rows(op->src[0]) instead of full contiguity: it addresses src0 and dst with their own strides (nb03 and the dst strides are added to the kernel args - dst can be a strided view of src0 with ggml_group_norm_inplace) and falls back from the flat fast path to a strided path that decodes the linear element index when rows or channels are padded. The flat/strided selection is done on the host in ggml_metal_library_get_pipeline_group_norm and compiled into one of two specialized pipeline variants via a FC_group_norm_flat function constant, so the dead path is eliminated at pipeline creation. GROUP_NORM support is also gated on F32 src and dst in supports_op - the kernel only implements f32 and an f16 input would silently be read as f32. Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (Mac Studio, M5 Ultra, 256GB)
Author
Parents
Loading