Commit 80d2c0c2c for llama.cpp
commit 80d2c0c2c56e4a1bbf74be98f64b99927bf230bc
Author: Georgi Gerganov <ggerganov@gmail.com>
Date: Sat Nov 7 13:04:13 2026 +0200
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)
diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp
index 72133d86d..0d18a5f01 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-device.cpp
@@ -2241,17 +2241,34 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_l2_norm(ggml_met
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_group_norm(ggml_metal_library_t lib, const ggml_tensor * op) {
assert(op->op == GGML_OP_GROUP_NORM);
- GGML_ASSERT(ggml_is_contiguous(op->src[0]));
+ GGML_ASSERT(ggml_is_contiguous_rows(op->src[0]));
char base[256];
char name[256];
snprintf(base, 256, "kernel_group_norm_f32");
- snprintf(name, 256, "%s", base);
+
+ const ggml_tensor * src0 = op->src[0];
+
+ const bool flat = src0->nb[1] == src0->ne[0]*src0->nb[0] &&
+ src0->nb[2] == src0->nb[1]*src0->ne[1] &&
+ op->nb[1] == op->ne[0] *op->nb[0] &&
+ op->nb[2] == op->nb[1] *op->ne[1];
+
+ if (flat) {
+ snprintf(name, 256, "%s", base);
+ } else {
+ snprintf(name, 256, "%s_strided", base);
+ }
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
if (!res.pipeline) {
- res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+ ggml_metal_cv_t cv = ggml_metal_cv_init();
+ ggml_metal_cv_set_bool(cv, flat, FC_GROUP_NORM + 0);
+
+ res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
+
+ ggml_metal_cv_free(cv);
}
res.smem = 32*sizeof(float);
diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m
index 63fbe2734..5e6099ce1 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.m
+++ b/ggml/src/ggml-metal/ggml-metal-device.m
@@ -1681,9 +1681,13 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
case GGML_OP_CUMSUM:
case GGML_OP_MEAN:
case GGML_OP_SOFT_MAX:
- case GGML_OP_GROUP_NORM:
case GGML_OP_L2_NORM:
return has_simdgroup_reduction && ggml_is_contiguous_rows(op->src[0]);
+ case GGML_OP_GROUP_NORM:
+ return has_simdgroup_reduction &&
+ op->src[0]->type == GGML_TYPE_F32 &&
+ op->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous_rows(op->src[0]);
case GGML_OP_COUNT_EQUAL:
return has_simdgroup_reduction &&
op->src[0]->type == GGML_TYPE_I32 &&
diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h
index 5fc49f9bb..7fd9d97cf 100644
--- a/ggml/src/ggml-metal/ggml-metal-impl.h
+++ b/ggml/src/ggml-metal/ggml-metal-impl.h
@@ -126,6 +126,7 @@
#define FC_LIGHTNING_INDEXER 2200
#define FC_MUL_MV_MMA 2300
#define FC_TOP_K_MASK 2400
+#define FC_GROUP_NORM 2500
// op-specific constants
#define OP_FLASH_ATTN_EXT_NQPSG 8
@@ -677,6 +678,10 @@ typedef struct {
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
+ uint64_t nb03;
+ uint64_t nb1;
+ uint64_t nb2;
+ uint64_t nb3;
int32_t ngrp;
float eps;
} ggml_metal_kargs_group_norm;
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index a42f66080..3d317a0a6 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -4610,6 +4610,10 @@ int ggml_metal_op_group_norm(ggml_metal_op_t ctx, int idx) {
/*.nb00 =*/ nb00,
/*.nb01 =*/ nb01,
/*.nb02 =*/ nb02,
+ /*.nb03 =*/ nb03,
+ /*.nb1 =*/ nb1,
+ /*.nb2 =*/ nb2,
+ /*.nb3 =*/ nb3,
/*.ngrp =*/ ngrp,
/*.eps =*/ eps,
};
@@ -4633,7 +4637,7 @@ int ggml_metal_op_group_norm(ggml_metal_op_t ctx, int idx) {
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
- ggml_metal_encoder_dispatch_threadgroups(enc, ngrp, 1, 1, nth, 1, 1);
+ ggml_metal_encoder_dispatch_threadgroups(enc, ngrp, 1, ne03, nth, 1, 1);
return 1;
}
diff --git a/ggml/src/ggml-metal/kernels/norm.metal b/ggml/src/ggml-metal/kernels/norm.metal
index c76d0c7ff..983d9b5cb 100644
--- a/ggml/src/ggml-metal/kernels/norm.metal
+++ b/ggml/src/ggml-metal/kernels/norm.metal
@@ -236,37 +236,59 @@ typedef decltype(kernel_l2_norm_impl<float, float>) kernel_l2_norm_t;
template [[host_name("kernel_l2_norm_f32_f32")]] kernel kernel_l2_norm_t kernel_l2_norm_impl<float, float>;
template [[host_name("kernel_l2_norm_f32_f32_4")]] kernel kernel_l2_norm_t kernel_l2_norm_impl<float4, float4>;
+constant bool FC_group_norm_flat [[function_constant(FC_GROUP_NORM + 0)]];
+
kernel void kernel_group_norm_f32(
constant ggml_metal_kargs_group_norm & args,
- device const float * src0,
- device float * dst,
+ device const char * src0,
+ device char * dst,
threadgroup float * buf [[threadgroup(0)]],
- uint tgpig[[threadgroup_position_in_grid]],
- uint tpitg[[thread_position_in_threadgroup]],
- uint sgitg[[simdgroup_index_in_threadgroup]],
- uint tiisg[[thread_index_in_simdgroup]],
- uint ntg[[threads_per_threadgroup]]) {
- const int64_t ne = args.ne00*args.ne01*args.ne02;
- const int64_t gs = args.ne00*args.ne01*((args.ne02 + args.ngrp - 1) / args.ngrp);
+ uint3 tgpig[[threadgroup_position_in_grid]],
+ ushort3 tpitg[[thread_position_in_threadgroup]],
+ ushort sgitg[[simdgroup_index_in_threadgroup]],
+ ushort tiisg[[thread_index_in_simdgroup]],
+ ushort3 ntg[[threads_per_threadgroup]]) {
+ const int64_t ne00 = args.ne00;
+ const int64_t ne01 = args.ne01;
+ const int64_t ne02 = args.ne02;
- int start = tgpig * gs;
- int end = start + gs;
+ const int64_t ngrp = args.ngrp;
- start += tpitg;
+ const int64_t i03 = tgpig.z; // batch index
+ const int64_t ig = tgpig.x; // group index
- if (end >= ne) {
- end = ne;
- }
+ const int64_t cpg = (ne02 + ngrp - 1)/ngrp;
+
+ const int64_t c0 = ig*cpg;
+ const int64_t c1 = (c0 + cpg < ne02) ? c0 + cpg : ne02;
+
+ const int64_t npf = ne00*ne01; // elements per channel
+ const int64_t gs = npf*(c1 - c0); // elements in this group
+
+ device const char * x03 = src0 + i03*args.nb03 + c0*args.nb02;
+ device char * y03 = dst + i03*args.nb3 + c0*args.nb2;
float tmp = 0.0f; // partial sum for thread in warp
- for (int j = start; j < end; j += ntg) {
- tmp += src0[j];
+ if (FC_group_norm_flat) {
+ device const float * x = (device const float *) x03;
+ for (int64_t e = tpitg.x; e < gs; e += ntg.x) {
+ tmp += x[e];
+ }
+ } else {
+ for (int64_t e = tpitg.x; e < gs; e += ntg.x) {
+ const int64_t i02 = c0 + e/npf;
+ const int64_t ir = e - (i02 - c0)*npf;
+ const int64_t i01 = ir/ne00;
+ const int64_t i00 = ir - i01*ne00;
+
+ tmp += *(device const float *)(x03 + (i02 - c0)*args.nb02 + i01*args.nb01 + i00*sizeof(float));
+ }
}
threadgroup_barrier(mem_flags::mem_threadgroup);
tmp = simd_sum(tmp);
- if (ntg > N_SIMDWIDTH) {
+ if (ntg.x > N_SIMDWIDTH) {
if (sgitg == 0) {
buf[tiisg] = 0.0f;
}
@@ -286,14 +308,31 @@ kernel void kernel_group_norm_f32(
const float mean = tmp / gs;
tmp = 0.0f;
- for (int j = start; j < end; j += ntg) {
- float xi = src0[j] - mean;
- dst[j] = xi;
- tmp += xi * xi;
+ if (FC_group_norm_flat) {
+ device const float * x = (device const float *) x03;
+ device float * y = (device float *) y03;
+ for (int64_t e = tpitg.x; e < gs; e += ntg.x) {
+ const float xi = x[e] - mean;
+ y[e] = xi;
+ tmp += xi*xi;
+ }
+ } else {
+ for (int64_t e = tpitg.x; e < gs; e += ntg.x) {
+ const int64_t i02 = c0 + e/npf;
+ const int64_t ir = e - (i02 - c0)*npf;
+ const int64_t i01 = ir/ne00;
+ const int64_t i00 = ir - i01*ne00;
+
+ const float xi = *(device const float *)(x03 + (i02 - c0)*args.nb02 + i01*args.nb01 + i00*sizeof(float)) - mean;
+
+ *(device float *)(y03 + (i02 - c0)*args.nb2 + i01*args.nb1 + i00*sizeof(float)) = xi;
+
+ tmp += xi*xi;
+ }
}
tmp = simd_sum(tmp);
- if (ntg > N_SIMDWIDTH) {
+ if (ntg.x > N_SIMDWIDTH) {
if (sgitg == 0) {
buf[tiisg] = 0.0f;
}
@@ -312,7 +351,20 @@ kernel void kernel_group_norm_f32(
const float variance = tmp / gs;
const float scale = 1.0f/sqrt(variance + args.eps);
- for (int j = start; j < end; j += ntg) {
- dst[j] *= scale;
+
+ if (FC_group_norm_flat) {
+ device float * y = (device float *) y03;
+ for (int64_t e = tpitg.x; e < gs; e += ntg.x) {
+ y[e] *= scale;
+ }
+ } else {
+ for (int64_t e = tpitg.x; e < gs; e += ntg.x) {
+ const int64_t i02 = c0 + e/npf;
+ const int64_t ir = e - (i02 - c0)*npf;
+ const int64_t i01 = ir/ne00;
+ const int64_t i00 = ir - i01*ne00;
+
+ *(device float *)(y03 + (i02 - c0)*args.nb2 + i01*args.nb1 + i00*sizeof(float)) *= scale;
+ }
}
}