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 +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(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); + } else { + kernel_mul_mv_iq2_xxs_f32_impl(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(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); + kernel_mul_mv_iq2_xxs_f32_disp(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); } template @@ -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 +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(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); + } else { + kernel_mul_mv_iq2_xs_f32_impl(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(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); + kernel_mul_mv_iq2_xs_f32_disp(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 +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(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); + } else { + kernel_mul_mv_iq3_s_f32_impl(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(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); + kernel_mul_mv_iq3_s_f32_disp(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); } template @@ -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 +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(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); + } else { + kernel_mul_mv_iq2_s_f32_impl(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(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); + kernel_mul_mv_iq2_s_f32_disp(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); } template @@ -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 +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(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); + } else { + kernel_mul_mv_iq1_s_f32_impl(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(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); + kernel_mul_mv_iq1_s_f32_disp(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); } template @@ -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 +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(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); + } else { + kernel_mul_mv_iq1_m_f32_impl(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(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); + kernel_mul_mv_iq1_m_f32_disp(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); } template @@ -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>>; template [[host_name("kernel_mul_mv_id_q5_K_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_q6_K_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; -template [[host_name("kernel_mul_mv_id_iq1_s_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; -template [[host_name("kernel_mul_mv_id_iq1_m_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; -template [[host_name("kernel_mul_mv_id_iq2_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; -template [[host_name("kernel_mul_mv_id_iq2_xs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; +template [[host_name("kernel_mul_mv_id_iq1_s_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; +template [[host_name("kernel_mul_mv_id_iq1_m_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; +template [[host_name("kernel_mul_mv_id_iq2_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; +template [[host_name("kernel_mul_mv_id_iq2_xs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_iq3_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; -template [[host_name("kernel_mul_mv_id_iq3_s_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; -template [[host_name("kernel_mul_mv_id_iq2_s_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; +template [[host_name("kernel_mul_mv_id_iq3_s_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; +template [[host_name("kernel_mul_mv_id_iq2_s_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_iq4_nl_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_iq4_xs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_tq2_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>;