Commit 6d194531e for llama.cpp
commit 6d194531e1e5876d5d1a0641ce3ca7dae23f94ad
Author: Aman Gupta <amangupta052@gmail.com>
Date: Thu Oct 29 16:11:29 2026 +0800
metal : faster MXFP4 mat-vec, full threadgroups for unary/binary/GLU (#31075)
* metal : faster MXFP4 mat-vec kernel
Decode the fp4 values with the fp16 bit trick instead of a threadgroup lookup table, and load the packed nibbles as aligned 32-bit words. The kernel no longer uses threadgroup memory.
* metal : launch contiguous unary and binary ops with full threadgroups
The contiguous unary path and the row-broadcast binary path used one thread per threadgroup. Use up to 256 threads per threadgroup and bound-check the tail.
* metal : split GLU rows across threadgroups
A GLU row ran on one threadgroup, so single-row decode used a single threadgroup for the whole op. Split each row over threadgroups of up to 256 threads.
diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp
index 25a8a3ef6..548b9896a 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-device.cpp
@@ -1103,7 +1103,6 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
{
nsg = N_SG_MXFP4;
nr0 = N_R0_MXFP4;
- smem = 32*sizeof(float);
} break;
case GGML_TYPE_Q2_K:
{
@@ -1399,7 +1398,6 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_m
{
nsg = N_SG_MXFP4;
nr0 = N_R0_MXFP4;
- smem = 32*sizeof(float);
} break;
case GGML_TYPE_Q2_K:
{
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index 14712c3e8..35e3dada2 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -899,7 +899,9 @@ int ggml_metal_op_unary(ggml_metal_op_t ctx, int idx) {
if (pipeline.cnt) {
const int n = pipeline.c4 ? ggml_nelements(op)/4 : ggml_nelements(op);
- ggml_metal_encoder_dispatch_threadgroups(enc, n, 1, 1, 1, 1, 1);
+ const int nth = MIN(256, n);
+
+ ggml_metal_encoder_dispatch_threadgroups(enc, (n + nth - 1)/nth, 1, 1, nth, 1, 1);
} else {
const int nth_max = MIN(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
const int nth = MIN(args.ne00, nth_max);
@@ -952,7 +954,9 @@ int ggml_metal_op_glu(ggml_metal_op_t ctx, int idx) {
const int64_t nrows = ggml_nrows(op->src[0]);
- const int32_t nth = std::max(1, std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne00/2));
+ // split each row over several threadgroups: with a single row (decode) one threadgroup would run on one GPU core
+ const int32_t nth = std::max(1, std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), std::min(256, ne0)));
+ const int32_t nk = (ne0 + nth - 1)/nth;
ggml_metal_encoder_set_pipeline(enc, pipeline);
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
@@ -964,7 +968,7 @@ int ggml_metal_op_glu(ggml_metal_op_t ctx, int idx) {
}
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3);
- ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, nth, 1, 1);
+ ggml_metal_encoder_dispatch_threadgroups(enc, nk, nrows, 1, nth, 1, 1);
return 1;
}
@@ -4328,7 +4332,9 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
ggml_metal_encoder_set_buffer (enc, bid_dst, 3);
if (pipeline.cnt) {
- ggml_metal_encoder_dispatch_threadgroups(enc, args.ne0, ggml_nrows(op), 1, 1, 1, 1);
+ const int nth = MIN(256, args.ne0);
+
+ ggml_metal_encoder_dispatch_threadgroups(enc, (args.ne0 + nth - 1)/nth, ggml_nrows(op), 1, nth, 1, 1);
} else {
const int nth_max = MIN(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
diff --git a/ggml/src/ggml-metal/kernels/binbcast.metal b/ggml/src/ggml-metal/kernels/binbcast.metal
index b367b3ea1..e125b7a03 100644
--- a/ggml/src/ggml-metal/kernels/binbcast.metal
+++ b/ggml/src/ggml-metal/kernels/binbcast.metal
@@ -22,8 +22,13 @@ kernel void kernel_bin_fuse_impl(
if (FC_RB) {
// row broadcast
- const uint i0 = tgpig.y*args.ne00 + tgpig.x;
- const uint i1 = FC_CB ? tgpig.x%args.ne10 : tgpig.x;
+ const uint i = tgpig.x*ntg.x + tpitg.x;
+ if (i >= (uint) args.ne0) {
+ return;
+ }
+
+ const uint i0 = tgpig.y*args.ne00 + i;
+ const uint i1 = FC_CB ? i%args.ne10 : i;
device const T0 * src0_row = (device const T0 *) (src0);
device T * dst_row = (device T *) (dst);
diff --git a/ggml/src/ggml-metal/kernels/mul_mv.metal b/ggml/src/ggml-metal/kernels/mul_mv.metal
index 7ad07a29c..5d012e484 100644
--- a/ggml/src/ggml-metal/kernels/mul_mv.metal
+++ b/ggml/src/ggml-metal/kernels/mul_mv.metal
@@ -3030,6 +3030,22 @@ kernel void kernel_mul_mv_iq4_xs_f32(
kernel_mul_mv_iq4_xs_f32_impl<N_R0_IQ4_XS, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
}
+// decode the E2M1 value in the low nibble of n, scaled by 2^-14, without a lookup table: the 3 magnitude bits
+// land in the fp16 exponent/mantissa (e=0 gives the 0.5 subnormal). same decode as MLX's fp4_e2m1
+static inline half mxfp4_dq1(uint32_t n) {
+ const half c = as_type<half>(ushort((n & 7) << 9));
+ return (n & 8) ? -c : c;
+}
+
+// low / high nibbles of the 4 bytes of w, in byte order, scaled by 2^-14
+static inline float4 mxfp4_lo4(uint32_t w) {
+ return float4(mxfp4_dq1(w), mxfp4_dq1(w >> 8), mxfp4_dq1(w >> 16), mxfp4_dq1(w >> 24));
+}
+
+static inline float4 mxfp4_hi4(uint32_t w) {
+ return float4(mxfp4_dq1(w >> 4), mxfp4_dq1(w >> 12), mxfp4_dq1(w >> 20), mxfp4_dq1(w >> 28));
+}
+
template<int NR0, typename args_t>
void kernel_mul_mv_mxfp4_f32_impl(
args_t args,
@@ -3042,8 +3058,6 @@ void kernel_mul_mv_mxfp4_f32_impl(
ushort sgitg) {
const short NSG = FC_mul_mv_nsg;
- threadgroup float * shmem_f32 = (threadgroup float *) shmem;
-
const int r0 = tgpig.x;
const int r1 = tgpig.y;
const int im = tgpig.z;
@@ -3065,9 +3079,6 @@ void kernel_mul_mv_mxfp4_f32_impl(
const short ix = tiisg/2; // 0...15
const short it = tiisg%2; // 0 or 1
- shmem_f32[tiisg] = kvalues_mxfp4_f[tiisg%16];
- threadgroup_barrier(mem_flags::mem_threadgroup);
-
float4 yl[4];
float sumf[NR0]={0.f};
@@ -3087,14 +3098,25 @@ void kernel_mul_mv_mxfp4_f32_impl(
device const block_mxfp4 & xb = x[row*ns01 + ib];
device const uint8_t * q2 = (device const uint8_t *)(xb.qs + 8*it);
- float4 acc1 = yl[0]*float4(shmem_f32[q2[0] & 0x0F], shmem_f32[q2[1] & 0x0F], shmem_f32[q2[2] & 0x0F], shmem_f32[q2[3] & 0x0F]);
- float4 acc2 = yl[1]*float4(shmem_f32[q2[0] >> 4 ], shmem_f32[q2[1] >> 4 ], shmem_f32[q2[2] >> 4 ], shmem_f32[q2[3] >> 4 ]);
- float4 acc3 = yl[2]*float4(shmem_f32[q2[4] & 0x0F], shmem_f32[q2[5] & 0x0F], shmem_f32[q2[6] & 0x0F], shmem_f32[q2[7] & 0x0F]);
- float4 acc4 = yl[3]*float4(shmem_f32[q2[4] >> 4 ], shmem_f32[q2[5] >> 4 ], shmem_f32[q2[6] >> 4 ], shmem_f32[q2[7] >> 4 ]);
+ // blocks are 17 bytes, so the 8 quant bytes are only byte aligned: read the aligned words that cover them
+ // and funnel-shift. the third word is read only when needed so no word past the block end is touched
+ device const uint32_t * pw = (device const uint32_t *)((ulong) q2 & ~3ul);
+ const uint sh = ((uint)(ulong) q2 & 3u)*8;
+ const uint32_t a = pw[0];
+ const uint32_t b = pw[1];
+ const uint32_t c = pw[sh ? 2 : 1];
+ const uint32_t w0 = (a >> sh) | ((b << 1) << (31 - sh));
+ const uint32_t w1 = (b >> sh) | ((c << 1) << (31 - sh));
+
+ float4 acc1 = yl[0]*mxfp4_lo4(w0);
+ float4 acc2 = yl[1]*mxfp4_hi4(w0);
+ float4 acc3 = yl[2]*mxfp4_lo4(w1);
+ float4 acc4 = yl[3]*mxfp4_hi4(w1);
acc1 = (acc1 + acc3) + (acc2 + acc4);
- sumf[row] += e8m0_to_fp32(xb.e) * ((acc1[0] + acc1[1]) + (acc1[2] + acc1[3]));
+ // 16384 undoes the 2^-14 of the decode
+ sumf[row] += (e8m0_to_fp32(xb.e)*16384.0f) * ((acc1[0] + acc1[1]) + (acc1[2] + acc1[3]));
}
yb += 16 * QK_MXFP4;
diff --git a/ggml/src/ggml-metal/kernels/unary.metal b/ggml/src/ggml-metal/kernels/unary.metal
index 9d24f028c..b361da996 100644
--- a/ggml/src/ggml-metal/kernels/unary.metal
+++ b/ggml/src/ggml-metal/kernels/unary.metal
@@ -38,7 +38,7 @@ kernel void kernel_unary_impl(
int i0;
if (FC_CNT) {
- i0 = tgpig.x;
+ i0 = tgpig.x*ntg.x + tpitg.x;
src0_ptr = (device const T0 *) (src0);
dst_ptr = (device T *) (dst);
@@ -61,6 +61,10 @@ kernel void kernel_unary_impl(
if (i0 >= args.ne0) {
return;
}
+ } else {
+ if (i0 >= args.ne00*args.ne01*args.ne02*args.ne03) {
+ return;
+ }
}
const TC x = (TC) src0_ptr[i0];
@@ -231,14 +235,15 @@ kernel void kernel_reglu(
device const char * src0,
device const char * src1,
device char * dst,
- uint tgpig[[threadgroup_position_in_grid]],
- uint tpitg[[thread_position_in_threadgroup]],
- uint ntg[[threads_per_threadgroup]]) {
- device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
- device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
- device T * dst_row = (device T *) ((device char *) dst + tgpig*args.nb1);
-
- for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+ uint2 tgpig[[threadgroup_position_in_grid]],
+ uint2 tpitg[[thread_position_in_threadgroup]],
+ uint2 ntg[[threads_per_threadgroup]],
+ uint2 ntgs[[threadgroups_per_grid]]) {
+ device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+ device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+ device T * dst_row = (device T *) ((device char *) dst + tgpig.y*args.nb1);
+
+ for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
const float x0 = src0_row[i0];
const float x1 = src1_row[i0];
@@ -257,14 +262,15 @@ kernel void kernel_geglu(
device const char * src0,
device const char * src1,
device char * dst,
- uint tgpig[[threadgroup_position_in_grid]],
- uint tpitg[[thread_position_in_threadgroup]],
- uint ntg[[threads_per_threadgroup]]) {
- device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
- device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
- device T * dst_row = (device T *) ((device char *) dst + tgpig*args.nb1);
-
- for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+ uint2 tgpig[[threadgroup_position_in_grid]],
+ uint2 tpitg[[thread_position_in_threadgroup]],
+ uint2 ntg[[threads_per_threadgroup]],
+ uint2 ntgs[[threadgroups_per_grid]]) {
+ device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+ device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+ device T * dst_row = (device T *) ((device char *) dst + tgpig.y*args.nb1);
+
+ for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
const float x0 = src0_row[i0];
const float x1 = src1_row[i0];
@@ -285,14 +291,15 @@ kernel void kernel_swiglu(
device const char * src0,
device const char * src1,
device char * dst,
- uint tgpig[[threadgroup_position_in_grid]],
- uint tpitg[[thread_position_in_threadgroup]],
- uint ntg[[threads_per_threadgroup]]) {
- device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
- device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
- device T * dst_row = (device T *) ((device char *) dst + tgpig*args.nb1);
-
- for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+ uint2 tgpig[[threadgroup_position_in_grid]],
+ uint2 tpitg[[thread_position_in_threadgroup]],
+ uint2 ntg[[threads_per_threadgroup]],
+ uint2 ntgs[[threadgroups_per_grid]]) {
+ device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+ device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+ device T * dst_row = (device T *) ((device char *) dst + tgpig.y*args.nb1);
+
+ for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
const float x0 = src0_row[i0];
const float x1 = src1_row[i0];
@@ -313,14 +320,15 @@ kernel void kernel_swiglu_oai(
device const char * src0,
device const char * src1,
device char * dst,
- uint tgpig[[threadgroup_position_in_grid]],
- uint tpitg[[thread_position_in_threadgroup]],
- uint ntg[[threads_per_threadgroup]]) {
- device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
- device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
- device T * dst_row = (device T *) ((device char *) dst + tgpig*args.nb1);
-
- for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+ uint2 tgpig[[threadgroup_position_in_grid]],
+ uint2 tpitg[[thread_position_in_threadgroup]],
+ uint2 ntg[[threads_per_threadgroup]],
+ uint2 ntgs[[threadgroups_per_grid]]) {
+ device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+ device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+ device T * dst_row = (device T *) ((device char *) dst + tgpig.y*args.nb1);
+
+ for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
float x0 = src0_row[i0];
float x1 = src1_row[i0];
@@ -345,14 +353,15 @@ kernel void kernel_swiglu_clamp(
device const char * src0,
device const char * src1,
device char * dst,
- uint tgpig[[threadgroup_position_in_grid]],
- uint tpitg[[thread_position_in_threadgroup]],
- uint ntg[[threads_per_threadgroup]]) {
- device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
- device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
- device T * dst_row = (device T *) ((device char *) dst + tgpig*args.nb1);
-
- for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+ uint2 tgpig[[threadgroup_position_in_grid]],
+ uint2 tpitg[[thread_position_in_threadgroup]],
+ uint2 ntg[[threads_per_threadgroup]],
+ uint2 ntgs[[threadgroups_per_grid]]) {
+ device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+ device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+ device T * dst_row = (device T *) ((device char *) dst + tgpig.y*args.nb1);
+
+ for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
const float gate = min((float) src0_row[i0], args.limit);
const float up = clamp((float) src1_row[i0], -args.limit, args.limit);
@@ -371,14 +380,15 @@ kernel void kernel_geglu_erf(
device const char * src0,
device const char * src1,
device char * dst,
- uint tgpig[[threadgroup_position_in_grid]],
- uint tpitg[[thread_position_in_threadgroup]],
- uint ntg[[threads_per_threadgroup]]) {
- device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
- device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
- device T * dst_row = (device T *) ((device char *) dst + tgpig*args.nb1);
-
- for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+ uint2 tgpig[[threadgroup_position_in_grid]],
+ uint2 tpitg[[thread_position_in_threadgroup]],
+ uint2 ntg[[threads_per_threadgroup]],
+ uint2 ntgs[[threadgroups_per_grid]]) {
+ device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+ device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+ device T * dst_row = (device T *) ((device char *) dst + tgpig.y*args.nb1);
+
+ for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
const float x0 = src0_row[i0];
const float x1 = src1_row[i0];
@@ -399,14 +409,15 @@ kernel void kernel_geglu_quick(
device const char * src0,
device const char * src1,
device char * dst,
- uint tgpig[[threadgroup_position_in_grid]],
- uint tpitg[[thread_position_in_threadgroup]],
- uint ntg[[threads_per_threadgroup]]) {
- device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
- device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
- device T * dst_row = (device T *) ((device char *) dst + tgpig*args.nb1);
-
- for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+ uint2 tgpig[[threadgroup_position_in_grid]],
+ uint2 tpitg[[thread_position_in_threadgroup]],
+ uint2 ntg[[threads_per_threadgroup]],
+ uint2 ntgs[[threadgroups_per_grid]]) {
+ device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+ device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+ device T * dst_row = (device T *) ((device char *) dst + tgpig.y*args.nb1);
+
+ for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
const float x0 = src0_row[i0];
const float x1 = src1_row[i0];