Commit aac810230 for llama.cpp
commit aac810230f9ef0cf73a47c56e46e87d0988be348
Author: Foad Abo Dahood <32059146+masterFoad@users.noreply.github.com>
Date: Fri Sep 11 12:30:20 2026 +0300
metal : fix idle threads in the remaining iq mul_mv kernels for ne00 < 1024 (#28692)
* metal : fix idle threads in the remaining iq mul_mv kernels for ne00 < 1024
Generalize the row split from #28086 to the six other kernels that use the
same lane-to-block mapping: iq1_s, iq1_m, iq2_xxs, iq2_xs, iq2_s and iq3_s.
Each of them assigns one 32-element chunk per thread, so when a row has
fewer than 32 chunks the rest of the simdgroup is idle. When nb32 < 32 and
nb32 divides 32, 32/nb32 threads now share each chunk and each takes a
slice of the rows, reusing the FC_mul_mv_split function constant and the
dispatch wrapper introduced for iq3_xxs.
The plain path is untouched: wide matrices keep one thread per chunk and
N_R0_<TYPE> = 4. Only the split path uses N_R0_<TYPE>_SPLIT = 8. The
K-quants have the same idle-thread issue but a different lane mapping, so
they are left for a separate change.
* metal : offset the src0 row pointer once in the iq mul_mv kernels
q2, dh, sc, qh and signs are all derived from xr, so the row slice
offset only has to be applied to xr.
* metal : fold iq mul_mv row split into offset0
Compute row0 and row1 before initializing the source pointers and apply
the row slice directly to offset0.
This keeps x and its derived pointers on the existing path while applying
the split row offset once.
diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp
index 1137c5f6d..bf3d07e78 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-device.cpp
@@ -932,12 +932,24 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
nsg = N_SG_IQ2_XXS;
nr0 = N_R0_IQ2_XXS;
smem = 256*8+128;
+
+ const int nb32 = ne00/32;
+ if (nb32 < 32 && (32 % nb32) == 0) {
+ nr0 = N_R0_IQ2_XXS_SPLIT;
+ split = true;
+ }
} break;
case GGML_TYPE_IQ2_XS:
{
nsg = N_SG_IQ2_XS;
nr0 = N_R0_IQ2_XS;
smem = 512*8+128;
+
+ const int nb32 = ne00/32;
+ if (nb32 < 32 && (32 % nb32) == 0) {
+ nr0 = N_R0_IQ2_XS_SPLIT;
+ split = true;
+ }
} break;
case GGML_TYPE_IQ3_XXS:
{
@@ -957,21 +969,45 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
nsg = N_SG_IQ3_S;
nr0 = N_R0_IQ3_S;
smem = 512*4;
+
+ const int nb32 = ne00/32;
+ if (nb32 < 32 && (32 % nb32) == 0) {
+ nr0 = N_R0_IQ3_S_SPLIT;
+ split = true;
+ }
} break;
case GGML_TYPE_IQ2_S:
{
nsg = N_SG_IQ2_S;
nr0 = N_R0_IQ2_S;
+
+ const int nb32 = ne00/32;
+ if (nb32 < 32 && (32 % nb32) == 0) {
+ nr0 = N_R0_IQ2_S_SPLIT;
+ split = true;
+ }
} break;
case GGML_TYPE_IQ1_S:
{
nsg = N_SG_IQ1_S;
nr0 = N_R0_IQ1_S;
+
+ const int nb32 = ne00/32;
+ if (nb32 < 32 && (32 % nb32) == 0) {
+ nr0 = N_R0_IQ1_S_SPLIT;
+ split = true;
+ }
} break;
case GGML_TYPE_IQ1_M:
{
nsg = N_SG_IQ1_M;
nr0 = N_R0_IQ1_M;
+
+ const int nb32 = ne00/32;
+ if (nb32 < 32 && (32 % nb32) == 0) {
+ nr0 = N_R0_IQ1_M_SPLIT;
+ split = true;
+ }
} break;
case GGML_TYPE_IQ4_NL:
{
@@ -1177,12 +1213,24 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_m
nsg = N_SG_IQ2_XXS;
nr0 = N_R0_IQ2_XXS;
smem = 256*8+128;
+
+ const int nb32 = ne00/32;
+ if (nb32 < 32 && (32 % nb32) == 0) {
+ nr0 = N_R0_IQ2_XXS_SPLIT;
+ split = true;
+ }
} break;
case GGML_TYPE_IQ2_XS:
{
nsg = N_SG_IQ2_XS;
nr0 = N_R0_IQ2_XS;
smem = 512*8+128;
+
+ const int nb32 = ne00/32;
+ if (nb32 < 32 && (32 % nb32) == 0) {
+ nr0 = N_R0_IQ2_XS_SPLIT;
+ split = true;
+ }
} break;
case GGML_TYPE_IQ3_XXS:
{
@@ -1202,21 +1250,45 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_m
nsg = N_SG_IQ3_S;
nr0 = N_R0_IQ3_S;
smem = 512*4;
+
+ const int nb32 = ne00/32;
+ if (nb32 < 32 && (32 % nb32) == 0) {
+ nr0 = N_R0_IQ3_S_SPLIT;
+ split = true;
+ }
} break;
case GGML_TYPE_IQ2_S:
{
nsg = N_SG_IQ2_S;
nr0 = N_R0_IQ2_S;
+
+ const int nb32 = ne00/32;
+ if (nb32 < 32 && (32 % nb32) == 0) {
+ nr0 = N_R0_IQ2_S_SPLIT;
+ split = true;
+ }
} break;
case GGML_TYPE_IQ1_S:
{
nsg = N_SG_IQ1_S;
nr0 = N_R0_IQ1_S;
+
+ const int nb32 = ne00/32;
+ if (nb32 < 32 && (32 % nb32) == 0) {
+ nr0 = N_R0_IQ1_S_SPLIT;
+ split = true;
+ }
} break;
case GGML_TYPE_IQ1_M:
{
nsg = N_SG_IQ1_M;
nr0 = N_R0_IQ1_M;
+
+ const int nb32 = ne00/32;
+ if (nb32 < 32 && (32 % nb32) == 0) {
+ nr0 = N_R0_IQ1_M_SPLIT;
+ split = true;
+ }
} break;
case GGML_TYPE_IQ4_NL:
{
diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h
index 1fe947633..28a9ba101 100644
--- a/ggml/src/ggml-metal/ggml-metal-impl.h
+++ b/ggml/src/ggml-metal/ggml-metal-impl.h
@@ -62,18 +62,23 @@
#define N_R0_IQ1_S 4
#define N_SG_IQ1_S 2
+#define N_R0_IQ1_S_SPLIT 8
#define N_R0_IQ1_M 4
#define N_SG_IQ1_M 2
+#define N_R0_IQ1_M_SPLIT 8
#define N_R0_IQ2_XXS 4
#define N_SG_IQ2_XXS 2
+#define N_R0_IQ2_XXS_SPLIT 8
#define N_R0_IQ2_XS 4
#define N_SG_IQ2_XS 2
+#define N_R0_IQ2_XS_SPLIT 8
#define N_R0_IQ2_S 4
#define N_SG_IQ2_S 2
+#define N_R0_IQ2_S_SPLIT 8
#define N_R0_IQ3_XXS 4
#define N_SG_IQ3_XXS 2
@@ -81,6 +86,7 @@
#define N_R0_IQ3_S 4
#define N_SG_IQ3_S 2
+#define N_R0_IQ3_S_SPLIT 8
#define N_R0_IQ4_NL 2
#define N_SG_IQ4_NL 2
diff --git a/ggml/src/ggml-metal/kernels/mul_mv.metal b/ggml/src/ggml-metal/kernels/mul_mv.metal
index fbe8398ea..8e2df2765 100644
--- a/ggml/src/ggml-metal/kernels/mul_mv.metal
+++ b/ggml/src/ggml-metal/kernels/mul_mv.metal
@@ -1889,8 +1889,19 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
const uint i12 = im%FC_mul_mv_ne12;
const uint i13 = im/FC_mul_mv_ne12;
- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
+ const int nb32 = nb * (QK_K / 32);
+
+ const short ntx = FC_mul_mv_split ? nb32 : 32;
+ const short nrep = 32 / ntx;
+
+ const short ix = tiisg % ntx;
+ const short irep = tiisg / ntx;
+
+ const short row0 = (nr0 * irep ) / nrep;
+ const short row1 = (nr0 * (irep + 1)) / nrep;
+
+ const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
device const block_iq2_xxs * x = (device const block_iq2_xxs *) (src0 + offset0);
device const float * y = (device const float *) (src1 + offset1);
@@ -1898,8 +1909,6 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
float yl[32];
float sumf[nr0]={0.f};
- const int nb32 = nb * (QK_K / 32);
-
threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem);
threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 256);
{
@@ -1912,11 +1921,9 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
threadgroup_barrier(mem_flags::mem_threadgroup);
}
- const int ix = tiisg;
-
device const float * y4 = y + 32 * ix;
- for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+ for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
for (short i = 0; i < 32; ++i) {
yl[i] = y4[i];
}
@@ -1928,7 +1935,7 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
device const uint16_t * q2 = xr->qs + 4 * ib;
device const half * dh = &xr->d;
- for (short row = 0; row < nr0; row++) {
+ for (short row = row0; row < row1; row++) {
const float db = dh[0];
device const uint8_t * aux8 = (device const uint8_t *)q2;
const uint32_t aux32 = q2[2] | (q2[3] << 16);
@@ -1948,7 +1955,7 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
q2 += args.nb01/2;
}
- y4 += 32 * 32;
+ y4 += 32 * ntx;
}
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -1961,6 +1968,23 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
}
}
+template<typename args_t>
+void kernel_mul_mv_iq2_xxs_f32_disp(
+ args_t args,
+ device const char * src0,
+ device const char * src1,
+ device char * dst,
+ threadgroup char * shmem,
+ uint3 tgpig,
+ ushort tiisg,
+ ushort sgitg) {
+ if (FC_mul_mv_split) {
+ kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ } else {
+ kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ }
+}
+
[[host_name("kernel_mul_mv_iq2_xxs_f32")]]
kernel void kernel_mul_mv_iq2_xxs_f32(
constant ggml_metal_kargs_mul_mv & args,
@@ -1971,7 +1995,7 @@ kernel void kernel_mul_mv_iq2_xxs_f32(
uint3 tgpig[[threadgroup_position_in_grid]],
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
- kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ kernel_mul_mv_iq2_xxs_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
}
template<int nr0, typename args_t>
@@ -1997,8 +2021,19 @@ void kernel_mul_mv_iq2_xs_f32_impl(
const uint i12 = im%FC_mul_mv_ne12;
const uint i13 = im/FC_mul_mv_ne12;
- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
+ const int nb32 = nb * (QK_K / 32);
+
+ const short ntx = FC_mul_mv_split ? nb32 : 32;
+ const short nrep = 32 / ntx;
+
+ const short ix = tiisg % ntx;
+ const short irep = tiisg / ntx;
+
+ const short row0 = (nr0 * irep ) / nrep;
+ const short row1 = (nr0 * (irep + 1)) / nrep;
+
+ const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
device const block_iq2_xs * x = (device const block_iq2_xs *) (src0 + offset0);
device const float * y = (device const float *) (src1 + offset1);
@@ -2006,8 +2041,6 @@ void kernel_mul_mv_iq2_xs_f32_impl(
float yl[32];
float sumf[nr0]={0.f};
- const int nb32 = nb * (QK_K / 32);
-
threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem);
threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 512);
{
@@ -2020,11 +2053,9 @@ void kernel_mul_mv_iq2_xs_f32_impl(
threadgroup_barrier(mem_flags::mem_threadgroup);
}
- const int ix = tiisg;
-
device const float * y4 = y + 32 * ix;
- for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+ for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
for (short i = 0; i < 32; ++i) {
yl[i] = y4[i];
}
@@ -2037,7 +2068,7 @@ void kernel_mul_mv_iq2_xs_f32_impl(
device const uint8_t * sc = xr->scales + ib;
device const half * dh = &xr->d;
- for (short row = 0; row < nr0; row++) {
+ for (short row = row0; row < row1; row++) {
const float db = dh[0];
const uint8_t ls1 = sc[0] & 0xf;
const uint8_t ls2 = sc[0] >> 4;
@@ -2066,7 +2097,7 @@ void kernel_mul_mv_iq2_xs_f32_impl(
sc += args.nb01;
}
- y4 += 32 * 32;
+ y4 += 32 * ntx;
}
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2079,6 +2110,23 @@ void kernel_mul_mv_iq2_xs_f32_impl(
}
}
+template<typename args_t>
+void kernel_mul_mv_iq2_xs_f32_disp(
+ args_t args,
+ device const char * src0,
+ device const char * src1,
+ device char * dst,
+ threadgroup char * shmem,
+ uint3 tgpig,
+ ushort tiisg,
+ ushort sgitg) {
+ if (FC_mul_mv_split) {
+ kernel_mul_mv_iq2_xs_f32_impl<N_R0_IQ2_XS_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ } else {
+ kernel_mul_mv_iq2_xs_f32_impl<N_R0_IQ2_XS, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ }
+}
+
[[host_name("kernel_mul_mv_iq2_xs_f32")]]
kernel void kernel_mul_mv_iq2_xs_f32(
constant ggml_metal_kargs_mul_mv & args,
@@ -2090,7 +2138,7 @@ kernel void kernel_mul_mv_iq2_xs_f32(
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
- kernel_mul_mv_iq2_xs_f32_impl<N_R0_IQ2_XS, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ kernel_mul_mv_iq2_xs_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
}
// FC_mul_mv_split: for nb32 < 32 (nb32 divides 32), 32/nb32 threads share each chunk and each takes a slice of the rows
@@ -2117,8 +2165,19 @@ void kernel_mul_mv_iq3_xxs_f32_impl(
const uint i12 = im%FC_mul_mv_ne12;
const uint i13 = im/FC_mul_mv_ne12;
- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
+ const int nb32 = nb * (QK_K / 32);
+
+ const short ntx = FC_mul_mv_split ? nb32 : 32;
+ const short nrep = 32 / ntx;
+
+ const short ix = tiisg % ntx;
+ const short irep = tiisg / ntx;
+
+ const short row0 = (nr0 * irep ) / nrep;
+ const short row1 = (nr0 * (irep + 1)) / nrep;
+
+ const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
device const block_iq3_xxs * x = (device const block_iq3_xxs *) (src0 + offset0);
device const float * y = (device const float *) (src1 + offset1);
@@ -2126,8 +2185,6 @@ void kernel_mul_mv_iq3_xxs_f32_impl(
float yl[32];
float sumf[nr0]={0.f};
- const int nb32 = nb * (QK_K / 32);
-
threadgroup uint32_t * svalues = (threadgroup uint32_t *)(shmem);
threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 256);
{
@@ -2140,15 +2197,6 @@ void kernel_mul_mv_iq3_xxs_f32_impl(
threadgroup_barrier(mem_flags::mem_threadgroup);
}
- const short ntx = FC_mul_mv_split ? nb32 : 32;
- const short nrep = 32 / ntx;
-
- const short ix = tiisg % ntx;
- const short irep = tiisg / ntx;
-
- const short row0 = (nr0 * irep ) / nrep;
- const short row1 = (nr0 * (irep + 1)) / nrep;
-
device const float * y4 = y + 32 * ix;
for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
@@ -2160,9 +2208,9 @@ void kernel_mul_mv_iq3_xxs_f32_impl(
const int ib = ib32 % (QK_K / 32);
device const block_iq3_xxs * xr = x + ibl;
- device const uint8_t * q3 = xr->qs + 8 * ib + (uint64_t) row0*args.nb01;
- device const uint16_t * gas = (device const uint16_t *)(xr->qs + QK_K/4) + 2 * ib + (uint64_t) row0*args.nb01/2;
- device const half * dh = &xr->d + (uint64_t) row0*args.nb01/2;
+ device const uint8_t * q3 = xr->qs + 8 * ib;
+ device const uint16_t * gas = (device const uint16_t *)(xr->qs + QK_K/4) + 2 * ib;
+ device const half * dh = &xr->d;
for (short row = row0; row < row1; row++) {
const float db = dh[0];
@@ -2253,8 +2301,19 @@ void kernel_mul_mv_iq3_s_f32_impl(
const uint i12 = im%FC_mul_mv_ne12;
const uint i13 = im/FC_mul_mv_ne12;
- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
+ const int nb32 = nb * (QK_K / 32);
+
+ const short ntx = FC_mul_mv_split ? nb32 : 32;
+ const short nrep = 32 / ntx;
+
+ const short ix = tiisg % ntx;
+ const short irep = tiisg / ntx;
+
+ const short row0 = (nr0 * irep ) / nrep;
+ const short row1 = (nr0 * (irep + 1)) / nrep;
+
+ const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
device const block_iq3_s * x = (device const block_iq3_s *) (src0 + offset0);
device const float * y = (device const float *) (src1 + offset1);
@@ -2262,8 +2321,6 @@ void kernel_mul_mv_iq3_s_f32_impl(
float yl[32];
float sumf[nr0]={0.f};
- const int nb32 = nb * (QK_K / 32);
-
threadgroup uint32_t * svalues = (threadgroup uint32_t *) shmem;
{
int nval = 8;
@@ -2272,11 +2329,9 @@ void kernel_mul_mv_iq3_s_f32_impl(
threadgroup_barrier(mem_flags::mem_threadgroup);
}
- const int ix = tiisg;
-
device const float * y4 = y + 32 * ix;
- for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+ for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
for (short i = 0; i < 32; ++i) {
yl[i] = y4[i];
}
@@ -2291,7 +2346,7 @@ void kernel_mul_mv_iq3_s_f32_impl(
device const uint8_t * signs = xr->signs + 4 * ib;
device const half * dh = &xr->d;
- for (short row = 0; row < nr0; row++) {
+ for (short row = row0; row < row1; row++) {
const float db = dh[0];
const float d = db * (1 + 2*((sc[0] >> 4*(ib%2)) & 0xf));
@@ -2315,7 +2370,7 @@ void kernel_mul_mv_iq3_s_f32_impl(
signs += args.nb01;
}
- y4 += 32 * 32;
+ y4 += 32 * ntx;
}
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2328,6 +2383,23 @@ void kernel_mul_mv_iq3_s_f32_impl(
}
}
+template<typename args_t>
+void kernel_mul_mv_iq3_s_f32_disp(
+ args_t args,
+ device const char * src0,
+ device const char * src1,
+ device char * dst,
+ threadgroup char * shmem,
+ uint3 tgpig,
+ ushort tiisg,
+ ushort sgitg) {
+ if (FC_mul_mv_split) {
+ kernel_mul_mv_iq3_s_f32_impl<N_R0_IQ3_S_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ } else {
+ kernel_mul_mv_iq3_s_f32_impl<N_R0_IQ3_S, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ }
+}
+
[[host_name("kernel_mul_mv_iq3_s_f32")]]
kernel void kernel_mul_mv_iq3_s_f32(
constant ggml_metal_kargs_mul_mv & args,
@@ -2339,7 +2411,7 @@ kernel void kernel_mul_mv_iq3_s_f32(
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
- kernel_mul_mv_iq3_s_f32_impl<N_R0_IQ3_S, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ kernel_mul_mv_iq3_s_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
}
template<int nr0, typename args_t>
@@ -2365,8 +2437,19 @@ void kernel_mul_mv_iq2_s_f32_impl(
const uint i12 = im%FC_mul_mv_ne12;
const uint i13 = im/FC_mul_mv_ne12;
- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
+ const int nb32 = nb * (QK_K / 32);
+
+ const short ntx = FC_mul_mv_split ? nb32 : 32;
+ const short nrep = 32 / ntx;
+
+ const short ix = tiisg % ntx;
+ const short irep = tiisg / ntx;
+
+ const short row0 = (nr0 * irep ) / nrep;
+ const short row1 = (nr0 * (irep + 1)) / nrep;
+
+ const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
device const block_iq2_s * x = (device const block_iq2_s *) (src0 + offset0);
device const float * y = (device const float *) (src1 + offset1);
@@ -2374,8 +2457,6 @@ void kernel_mul_mv_iq2_s_f32_impl(
float yl[32];
float sumf[nr0]={0.f};
- const int nb32 = nb * (QK_K / 32);
-
//threadgroup uint64_t * svalues = (threadgroup uint64_t *) shmem;
//{
// int nval = 32;
@@ -2384,11 +2465,9 @@ void kernel_mul_mv_iq2_s_f32_impl(
// threadgroup_barrier(mem_flags::mem_threadgroup);
//}
- const short ix = tiisg;
-
device const float * y4 = y + 32 * ix;
- for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+ for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
for (short i = 0; i < 32; ++i) {
yl[i] = y4[i];
}
@@ -2403,7 +2482,7 @@ void kernel_mul_mv_iq2_s_f32_impl(
device const uint8_t * signs = qs + QK_K/8;
device const half * dh = &xr->d;
- for (short row = 0; row < nr0; row++) {
+ for (short row = row0; row < row1; row++) {
const float db = dh[0];
const float d1 = db * (0.5f + (sc[0] & 0xf));
const float d2 = db * (0.5f + (sc[0] >> 4));
@@ -2428,7 +2507,7 @@ void kernel_mul_mv_iq2_s_f32_impl(
signs += args.nb01;
}
- y4 += 32 * 32;
+ y4 += 32 * ntx;
}
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2441,6 +2520,23 @@ void kernel_mul_mv_iq2_s_f32_impl(
}
}
+template<typename args_t>
+void kernel_mul_mv_iq2_s_f32_disp(
+ args_t args,
+ device const char * src0,
+ device const char * src1,
+ device char * dst,
+ threadgroup char * shmem,
+ uint3 tgpig,
+ ushort tiisg,
+ ushort sgitg) {
+ if (FC_mul_mv_split) {
+ kernel_mul_mv_iq2_s_f32_impl<N_R0_IQ2_S_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ } else {
+ kernel_mul_mv_iq2_s_f32_impl<N_R0_IQ2_S, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ }
+}
+
[[host_name("kernel_mul_mv_iq2_s_f32")]]
kernel void kernel_mul_mv_iq2_s_f32(
constant ggml_metal_kargs_mul_mv & args,
@@ -2452,7 +2548,7 @@ kernel void kernel_mul_mv_iq2_s_f32(
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
- kernel_mul_mv_iq2_s_f32_impl<N_R0_IQ2_S, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ kernel_mul_mv_iq2_s_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
}
template<int nr0, typename args_t>
@@ -2478,8 +2574,19 @@ void kernel_mul_mv_iq1_s_f32_impl(
const uint i12 = im%FC_mul_mv_ne12;
const uint i13 = im/FC_mul_mv_ne12;
- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
+ const int nb32 = nb * (QK_K / 32);
+
+ const short ntx = FC_mul_mv_split ? nb32 : 32;
+ const short nrep = 32 / ntx;
+
+ const short ix = tiisg % ntx;
+ const short irep = tiisg / ntx;
+
+ const short row0 = (nr0 * irep ) / nrep;
+ const short row1 = (nr0 * (irep + 1)) / nrep;
+
+ const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
device const block_iq1_s * x = (device const block_iq1_s *) (src0 + offset0);
device const float * y = (device const float *) (src1 + offset1);
@@ -2487,13 +2594,9 @@ void kernel_mul_mv_iq1_s_f32_impl(
float yl[32];
float sumf[nr0]={0.f};
- const int nb32 = nb * (QK_K / 32);
-
- const short ix = tiisg;
-
device const float * y4 = y + 32 * ix;
- for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+ for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
float sumy = 0;
for (short i = 0; i < 32; ++i) {
yl[i] = y4[i];
@@ -2508,7 +2611,7 @@ void kernel_mul_mv_iq1_s_f32_impl(
device const uint16_t * qh = xr->qh + ib;
device const half * dh = &xr->d;
- for (short row = 0; row < nr0; row++) {
+ for (short row = row0; row < row1; row++) {
constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700)));
constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((qh[0] << 5) & 0x700)));
constant uint8_t * grid3 = (constant uint8_t *)(iq1s_grid_gpu + (qs[2] | ((qh[0] << 2) & 0x700)));
@@ -2528,7 +2631,7 @@ void kernel_mul_mv_iq1_s_f32_impl(
qh += args.nb01/2;
}
- y4 += 32 * 32;
+ y4 += 32 * ntx;
}
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2541,6 +2644,23 @@ void kernel_mul_mv_iq1_s_f32_impl(
}
}
+template<typename args_t>
+void kernel_mul_mv_iq1_s_f32_disp(
+ args_t args,
+ device const char * src0,
+ device const char * src1,
+ device char * dst,
+ threadgroup char * shmem,
+ uint3 tgpig,
+ ushort tiisg,
+ ushort sgitg) {
+ if (FC_mul_mv_split) {
+ kernel_mul_mv_iq1_s_f32_impl<N_R0_IQ1_S_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ } else {
+ kernel_mul_mv_iq1_s_f32_impl<N_R0_IQ1_S, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ }
+}
+
[[host_name("kernel_mul_mv_iq1_s_f32")]]
kernel void kernel_mul_mv_iq1_s_f32(
constant ggml_metal_kargs_mul_mv & args,
@@ -2551,7 +2671,7 @@ kernel void kernel_mul_mv_iq1_s_f32(
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
- kernel_mul_mv_iq1_s_f32_impl<N_R0_IQ1_S, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
+ kernel_mul_mv_iq1_s_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
}
template<int nr0, typename args_t>
@@ -2577,8 +2697,19 @@ void kernel_mul_mv_iq1_m_f32_impl(
const uint i12 = im%FC_mul_mv_ne12;
const uint i13 = im/FC_mul_mv_ne12;
- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
+ const int nb32 = nb * (QK_K / 32);
+
+ const short ntx = FC_mul_mv_split ? nb32 : 32;
+ const short nrep = 32 / ntx;
+
+ const short ix = tiisg % ntx;
+ const short irep = tiisg / ntx;
+
+ const short row0 = (nr0 * irep ) / nrep;
+ const short row1 = (nr0 * (irep + 1)) / nrep;
+
+ const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
device const block_iq1_m * x = (device const block_iq1_m *) (src0 + offset0);
device const float * y = (device const float *) (src1 + offset1);
@@ -2586,15 +2717,11 @@ void kernel_mul_mv_iq1_m_f32_impl(
float yl[32];
float sumf[nr0]={0.f};
- const int nb32 = nb * (QK_K / 32);
-
- const short ix = tiisg;
-
device const float * y4 = y + 32 * ix;
iq1m_scale_t scale;
- for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+ for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
float4 sumy = {0.f};
for (short i = 0; i < 8; ++i) {
yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0];
@@ -2611,7 +2738,7 @@ void kernel_mul_mv_iq1_m_f32_impl(
device const uint8_t * qh = xr->qh + 2 * ib;
device const uint16_t * sc = (device const uint16_t *)xr->scales;
- for (short row = 0; row < nr0; row++) {
+ for (short row = row0; row < row1; row++) {
scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700)));
@@ -2637,7 +2764,7 @@ void kernel_mul_mv_iq1_m_f32_impl(
qh += args.nb01;
}
- y4 += 32 * 32;
+ y4 += 32 * ntx;
}
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2650,6 +2777,23 @@ void kernel_mul_mv_iq1_m_f32_impl(
}
}
+template<typename args_t>
+void kernel_mul_mv_iq1_m_f32_disp(
+ args_t args,
+ device const char * src0,
+ device const char * src1,
+ device char * dst,
+ threadgroup char * shmem,
+ uint3 tgpig,
+ ushort tiisg,
+ ushort sgitg) {
+ if (FC_mul_mv_split) {
+ kernel_mul_mv_iq1_m_f32_impl<N_R0_IQ1_M_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ } else {
+ kernel_mul_mv_iq1_m_f32_impl<N_R0_IQ1_M, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+ }
+}
+
[[host_name("kernel_mul_mv_iq1_m_f32")]]
kernel void kernel_mul_mv_iq1_m_f32(
constant ggml_metal_kargs_mul_mv & args,
@@ -2660,7 +2804,7 @@ kernel void kernel_mul_mv_iq1_m_f32(
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
- kernel_mul_mv_iq1_m_f32_impl<N_R0_IQ1_M, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
+ kernel_mul_mv_iq1_m_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
}
template<int NR0, typename args_t>
@@ -3239,13 +3383,13 @@ template [[host_name("kernel_mul_mv_id_q3_K_f32")]] kernel kernel_mul_mv_id_t
template [[host_name("kernel_mul_mv_id_q4_K_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q4_K_f32_impl <N_R0_Q4_K>>>;
template [[host_name("kernel_mul_mv_id_q5_K_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q5_K_f32_impl <N_R0_Q5_K>>>;
template [[host_name("kernel_mul_mv_id_q6_K_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q6_K_f32_impl <N_R0_Q6_K>>>;
-template [[host_name("kernel_mul_mv_id_iq1_s_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_s_f32_impl <N_R0_IQ1_S>>>;
-template [[host_name("kernel_mul_mv_id_iq1_m_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_m_f32_impl <N_R0_IQ1_M>>>;
-template [[host_name("kernel_mul_mv_id_iq2_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS>>>;
-template [[host_name("kernel_mul_mv_id_iq2_xs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xs_f32_impl <N_R0_IQ2_XS>>>;
+template [[host_name("kernel_mul_mv_id_iq1_s_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_s_f32_disp<ggml_metal_kargs_mul_mv>>>;
+template [[host_name("kernel_mul_mv_id_iq1_m_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_m_f32_disp<ggml_metal_kargs_mul_mv>>>;
+template [[host_name("kernel_mul_mv_id_iq2_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xxs_f32_disp<ggml_metal_kargs_mul_mv>>>;
+template [[host_name("kernel_mul_mv_id_iq2_xs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xs_f32_disp<ggml_metal_kargs_mul_mv>>>;
template [[host_name("kernel_mul_mv_id_iq3_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq3_xxs_f32_disp<ggml_metal_kargs_mul_mv>>>;
-template [[host_name("kernel_mul_mv_id_iq3_s_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq3_s_f32_impl <N_R0_IQ3_S>>>;
-template [[host_name("kernel_mul_mv_id_iq2_s_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_s_f32_impl <N_R0_IQ2_S>>>;
+template [[host_name("kernel_mul_mv_id_iq3_s_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq3_s_f32_disp<ggml_metal_kargs_mul_mv>>>;
+template [[host_name("kernel_mul_mv_id_iq2_s_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_s_f32_disp<ggml_metal_kargs_mul_mv>>>;
template [[host_name("kernel_mul_mv_id_iq4_nl_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq4_nl_f32_impl <N_R0_IQ4_NL>>>;
template [[host_name("kernel_mul_mv_id_iq4_xs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq4_xs_f32_impl <N_R0_IQ4_XS>>>;
template [[host_name("kernel_mul_mv_id_tq2_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_tq2_0_f32_impl <N_R0_TQ2_0>>>;