diff --git a/ggml/src/ggml-metal/ggml-metal-common.cpp b/ggml/src/ggml-metal/ggml-metal-common.cpp index 9c0b9474c..e43023cc6 100644 --- a/ggml/src/ggml-metal/ggml-metal-common.cpp +++ b/ggml/src/ggml-metal/ggml-metal-common.cpp @@ -7,22 +7,29 @@ #include -// must stay in sync with the kernel_fwht__ templates in misc.metal -static bool ggml_metal_fwht_supported_size(int64_t n) { - return n == 64 || n == 128 || n == 256 || n == 512; +// must stay in sync with the kernel_fwht__ templates in misc.metal. Widths up to +// 512 run on the simdgroup kernel and need no threadgroup memory. The wider ones allocate +// float[N] per threadgroup, so they are only available where that fits. +static bool ggml_metal_fwht_supported_size(int64_t n, size_t max_tg_mem) { + if (n == 64 || n == 128 || n == 256 || n == 512) { + return true; + } + + if (n == 1024 || n == 2048 || n == 4096 || n == 8192) { + return (size_t) n * sizeof(float) <= max_tg_mem; + } + + return false; } // the FWHT kernels handle a Hadamard-hinted MUL_MAT only under these conditions. supports_op // and the dispatch must ask the same question: an F16 src1 that is admitted but then falls // through reaches the generic path, which has no F32 src0 by F16 src1 kernel. -bool ggml_metal_op_mul_mat_use_fwht(const struct ggml_tensor * op) { - return ggml_get_op_params_i32(op, 1) == GGML_HINT_SRC0_IS_HADAMARD && - op->type == GGML_TYPE_F32 && - (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && - ggml_is_contiguous(op->src[1]) && - ggml_is_contiguous(op) && - ggml_are_same_shape(op->src[1], op) && - ggml_metal_fwht_supported_size(op->src[1]->ne[0]); +bool ggml_metal_op_mul_mat_use_fwht(const struct ggml_tensor * op, size_t max_tg_mem) { + return ggml_get_op_params_i32(op, 1) == GGML_HINT_SRC0_IS_HADAMARD && op->type == GGML_TYPE_F32 && + (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && ggml_is_contiguous(op->src[1]) && + ggml_is_contiguous(op) && ggml_are_same_shape(op->src[1], op) && + ggml_metal_fwht_supported_size(op->src[1]->ne[0], max_tg_mem); } bool ggml_metal_op_mul_mat_use_mm(const struct ggml_tensor * op, bool has_simdgroup_mm) { diff --git a/ggml/src/ggml-metal/ggml-metal-common.h b/ggml/src/ggml-metal/ggml-metal-common.h index e6a28d032..6b5a1883f 100644 --- a/ggml/src/ggml-metal/ggml-metal-common.h +++ b/ggml/src/ggml-metal/ggml-metal-common.h @@ -3,6 +3,7 @@ #pragma once #include +#include #ifdef __cplusplus extern "C" { @@ -48,7 +49,7 @@ bool ggml_mem_ranges_check(ggml_mem_ranges_t mrs, const struct ggml_tensor * ten void ggml_graph_optimize(struct ggml_cgraph * gf); // mat-mat vs mat-vec dispatch; used by both supports_op and ggml_metal_op_mul_mat* -bool ggml_metal_op_mul_mat_use_fwht (const struct ggml_tensor * op); +bool ggml_metal_op_mul_mat_use_fwht(const struct ggml_tensor * op, size_t max_tg_mem); bool ggml_metal_op_mul_mat_use_mm (const struct ggml_tensor * op, bool has_simdgroup_mm); bool ggml_metal_op_mul_mat_id_use_mm(const struct ggml_tensor * op, bool has_simdgroup_mm); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index c4b86e80d..fa58b8965 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1854,7 +1854,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te // the FWHT kernels read an F16 source directly; every other F16 src1 path // still goes through ggml_metal_supports_mul_mat_op if (op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F16 && - ggml_metal_op_mul_mat_use_fwht(op)) { + ggml_metal_op_mul_mat_use_fwht(op, dev->props.max_theadgroup_memory_size)) { return has_simdgroup_reduction; } return ggml_metal_supports_mul_mat_op( diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 490dd83a1..eed85f283 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -1235,6 +1235,11 @@ typedef struct { int32_t top_k; // k } ggml_metal_kargs_top_k; +// widths at or above this use the threadgroup FWHT kernel, one row per threadgroup +// with GGML_METAL_FWHT_TG_NT threads, instead of one row per simdgroup +#define GGML_METAL_FWHT_TG_MIN_N 1024 +#define GGML_METAL_FWHT_TG_NT 256 + typedef struct { int32_t ne01; // n_tokens uint64_t nb01; // logits row stride diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 54d5b9808..8a46ec66a 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2343,6 +2343,13 @@ int ggml_metal_op_fwht(ggml_metal_op_t ctx, int idx) { const int th_max = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline); const int simd_size = 32; + if (n >= GGML_METAL_FWHT_TG_MIN_N) { + GGML_ASSERT(th_max >= GGML_METAL_FWHT_TG_NT); + ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, GGML_METAL_FWHT_TG_NT, 1, 1); + + return 1; + } + int sg_per_tg = 2; sg_per_tg = std::min(sg_per_tg, th_max/simd_size); sg_per_tg = std::max(sg_per_tg, 1); @@ -2419,10 +2426,11 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) { ggml_metal_library_t lib = ctx->lib; ggml_metal_encoder_t enc = ctx->enc; - if (ggml_metal_op_mul_mat_use_fwht(op)) { + const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); + + if (ggml_metal_op_mul_mat_use_fwht(op, props_dev->max_theadgroup_memory_size)) { return ggml_metal_op_fwht(ctx, idx); } - const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); diff --git a/ggml/src/ggml-metal/kernels/misc.metal b/ggml/src/ggml-metal/kernels/misc.metal index 279d69f8f..d3b01978f 100644 --- a/ggml/src/ggml-metal/kernels/misc.metal +++ b/ggml/src/ggml-metal/kernels/misc.metal @@ -429,6 +429,81 @@ kernel void kernel_fwht( } } +// Wide blocks: one row per threadgroup instead of per simdgroup, so each thread keeps +// N/NT values rather than N/32. Butterflies below the simdgroup width still shuffle; +// those up to NT go through threadgroup memory; the rest stay in registers. +// TODO: try avoiding branch https://github.com/ggml-org/llama.cpp/pull/29094#discussion_r4049563223 +// TODO: try to unroll loops +template +kernel void kernel_fwht_tg( + constant ggml_metal_kargs_fwht & args, + device const src_t * src, + device float * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + + constexpr int NW = N_SIMDWIDTH; + constexpr int NE = N / NT; + + threadgroup float shmem[N]; + + const float scale = 1.0f / sqrt((float) N); + + const int64_t r = tgpig.x; + if (r >= args.nrows) { + return; + } + + src += r * N; + dst += r * N; + + const int tid = sgitg * NW + tiisg; + + float reg[NE]; + for (int i = 0; i < NE; i++) { + reg[i] = float(src[i*NT + tid])*scale; + } + + for (int i = 1; i < NW; i *= 2) { + for (int j = 0; j < NE; j++) { + const float val = reg[j]; + const float val2 = simd_shuffle_xor(val, i); + reg[j] = (tid & i) == 0 ? val2 + val : val2 - val; + } + } + + for (int i = NW; i < NT; i *= 2) { + for (int j = 0; j < NE; j++) { + shmem[j*NT + tid] = reg[j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int j = 0; j < NE; j++) { + const float val = reg[j]; + const float val2 = shmem[j*NT + (tid ^ i)]; + reg[j] = (tid & i) == 0 ? val2 + val : val2 - val; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + for (int i = NT; i < N; i *= 2) { + const int step = i / NT; + for (int j = 0; j < NE; j += (2 * step)) { + for (int k = 0; k < step; k++) { + const float x = reg[j + k ]; + const float y = reg[j + k + step]; + reg[j + k] = x + y; + reg[j + k + step] = x - y; + } + } + } + + for (int i = 0; i < NE; i++) { + dst[i*NT + tid] = reg[i]; + } +} + typedef decltype(kernel_fwht<64, float>) kernel_fwht_f32_t; typedef decltype(kernel_fwht<64, half>) kernel_fwht_f16_t; @@ -442,6 +517,16 @@ template [[host_name("kernel_fwht_f16_128")]] kernel kernel_fwht_f16_t kernel_fw template [[host_name("kernel_fwht_f16_256")]] kernel kernel_fwht_f16_t kernel_fwht<256, half>; template [[host_name("kernel_fwht_f16_512")]] kernel kernel_fwht_f16_t kernel_fwht<512, half>; +template [[host_name("kernel_fwht_f32_1024")]] kernel kernel_fwht_f32_t kernel_fwht_tg<1024, GGML_METAL_FWHT_TG_NT, float>; +template [[host_name("kernel_fwht_f32_2048")]] kernel kernel_fwht_f32_t kernel_fwht_tg<2048, GGML_METAL_FWHT_TG_NT, float>; +template [[host_name("kernel_fwht_f32_4096")]] kernel kernel_fwht_f32_t kernel_fwht_tg<4096, GGML_METAL_FWHT_TG_NT, float>; +template [[host_name("kernel_fwht_f32_8192")]] kernel kernel_fwht_f32_t kernel_fwht_tg<8192, GGML_METAL_FWHT_TG_NT, float>; + +template [[host_name("kernel_fwht_f16_1024")]] kernel kernel_fwht_f16_t kernel_fwht_tg<1024, GGML_METAL_FWHT_TG_NT, half>; +template [[host_name("kernel_fwht_f16_2048")]] kernel kernel_fwht_f16_t kernel_fwht_tg<2048, GGML_METAL_FWHT_TG_NT, half>; +template [[host_name("kernel_fwht_f16_4096")]] kernel kernel_fwht_f16_t kernel_fwht_tg<4096, GGML_METAL_FWHT_TG_NT, half>; +template [[host_name("kernel_fwht_f16_8192")]] kernel kernel_fwht_f16_t kernel_fwht_tg<8192, GGML_METAL_FWHT_TG_NT, half>; + constant int FC_dsv4_hc_n_hc [[function_constant(FC_DSV4_HC + 0)]]; kernel void kernel_dsv4_hc_comb_f32(