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.
This commit is contained in:
Foad Abo Dahood
2026-09-11 12:30:20 +03:00
committed by GitHub
parent 5cdd3d1dad
commit aac810230f
3 changed files with 304 additions and 82 deletions
+72
View File
@@ -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:
{
+6
View File
@@ -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
+226 -82
View File
@@ -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>>>;