diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 3420ba85116c..260fdb431d9e 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -844,7 +844,17 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_meta const int16_t r2 = (int16_t) (ne12 / op->src[0]->ne[2]); const int16_t r3 = (int16_t) (ne13 / op->src[0]->ne[3]); - snprintf(base, 256, "kernel_mul_mm_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1)); + // Q1_0 products made only of full tiles can use the static-K32 tensor kernel (GGML_METAL_Q1_MM_K32=1) + static const bool q1_k32_env = getenv("GGML_METAL_Q1_MM_K32") && atoi(getenv("GGML_METAL_Q1_MM_K32")) == 1; + const bool q1_k32 = q1_k32_env && has_tensor && tsrc0 == GGML_TYPE_Q1_0 && tsrc1 == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && + !bc_inp && !bc_out && op->ne[0] >= NRA && op->ne[1] >= NRB && + ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op); + + if (q1_k32) { + snprintf(base, 256, "kernel_mul_mm_q1_0_f32_k32"); + } else { + snprintf(base, 256, "kernel_mul_mm_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1)); + } snprintf(name, 256, "%s_bci=%d_bco=%d_ne12=%d_ne13=%d_r2=%d_r3=%d", base, bc_inp, bc_out, ne12, ne13, r2, r3); diff --git a/ggml/src/ggml-metal/kernels/mul_mm.metal b/ggml/src/ggml-metal/kernels/mul_mm.metal index b0029fdc6ad2..106dcea32547 100644 --- a/ggml/src/ggml-metal/kernels/mul_mm.metal +++ b/ggml/src/ggml-metal/kernels/mul_mm.metal @@ -140,6 +140,83 @@ kernel void kernel_mul_mm( cT.store(tD.slice(ra, rb)); } +// Q1_0 products made only of full tiles (M % 64, N % 128, K % 32 all zero): kernel_mul_mm with a static K32 extent and no bounds handling +kernel void kernel_mul_mm_q1_0_f32_k32( + constant ggml_metal_kargs_mul_mm & args, + device const char * srcA, + device const char * srcB, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + (void) sgitg; + + constexpr int NRB = SZ_SIMDGROUP * N_MM_BLOCK_X * N_MM_SIMD_GROUP_X; + constexpr int NRA = SZ_SIMDGROUP * N_MM_BLOCK_Y * N_MM_SIMD_GROUP_Y; + constexpr int NK = N_MM_NK_TOTAL; + constexpr int NUM_THREADS = N_SIMDWIDTH * N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y; + static_assert(NUM_THREADS == NRA * N_MM_NK, "one 16-weight chunk per thread"); + + const int K = args.ne00; + const int M = args.ne0; + const int N = args.ne1; + + const int im = tgpig.z; + const int i12 = im % FC_mul_mm_ne12; + const int i13 = im / FC_mul_mm_ne12; + + const uint64_t offset0 = (i12/FC_mul_mm_r2)*args.nb02 + (i13/FC_mul_mm_r3)*args.nb03; + + const int ra = tgpig.y * NRA; + const int rb = tgpig.x * NRB; + + // same work mapping as kernel_mul_mm: row = tiitg / N_MM_NK, chunk = tiitg % N_MM_NK + const int row = tiitg / N_MM_NK; + const short k_base = (tiitg % N_MM_NK) * 16; + + threadgroup half * sa = (threadgroup half *) shmem; + + device const block_q1_0 * row_ptr = (device const block_q1_0 *)(srcA + args.nb01 * (ra + row) + offset0); + device float * ptrB = (device float *)(srcB + args.nb12*i12 + args.nb13*i13); + const int strideB = args.nb11 / sizeof(float); + + auto tA = tensor(sa, dextents(NK, NRA)); + auto tB = tensor(ptrB + rb * strideB, dextents(NK, NRB), array({1, strideB})); + + mpp::tensor_ops::matmul2d< + mpp::tensor_ops::matmul2d_descriptor( + NRB, NRA, NK, false, true, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate), + execution_simdgroups> mm; + + auto cT = mm.get_destination_cooperative_tensor(); + + for (int loop_k = 0; loop_k < K; loop_k += NK) { + const int k_pos = loop_k + k_base; + + half4x4 temp_a; + dequantize_q1_0(row_ptr + k_pos / QK1_0, (k_pos / 16) % (QK1_0 / 16), temp_a); + + FOR_UNROLL (short i = 0; i < 16; i++) { + sa[row * NK + k_base + i] = temp_a[i/4][i%4]; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + auto tBv = tensor(ptrB + loop_k + rb * strideB, dextents(NK, NRB), array({1, strideB})); + + mm.run(tBv, tA, cT); + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + device float * dstTile = (device float *)dst + (uint64_t) im * N * M + (uint64_t) rb * M + ra; + + auto tD = tensor(dstTile, dextents(NRA, NRB), array({1, M})); + cT.store(tD); +} + #else template< diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 72e580d895c0..9841980e06a5 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9400,6 +9400,10 @@ static std::vector> make_test_cases_eval() { } } + // Q1_0 products made only of full tensor tiles (GGML_METAL_Q1_MM_K32) + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q1_0, GGML_TYPE_F32, 128, 256, 512, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q1_0, GGML_TYPE_F32, 64, 128, 256, {2, 3}, {2, 1})); + // PTQ1_0 small batches, row tails and broadcast dimensions. for (int n : {1, 2, 3, 4, 8}) { for (int k : {128, 384, 5120}) {