diff --git a/ggml/src/ggml-cuda/fwht.cu b/ggml/src/ggml-cuda/fwht.cu index 184dc254c..67eea594d 100644 --- a/ggml/src/ggml-cuda/fwht.cu +++ b/ggml/src/ggml-cuda/fwht.cu @@ -1,9 +1,10 @@ #include "common.cuh" +#include "convert.cuh" #include "fwht.cuh" -template +template __launch_bounds__(4*ggml_cuda_get_physical_warp_size(), 1) -__global__ void fwht_cuda(const float * src, float * dst, const int64_t n_rows, const float scale) { +__global__ void fwht_cuda(const T * src, float * dst, const int64_t n_rows, const float scale) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); const int64_t r = (int64_t) blockIdx.x * blockDim.y + threadIdx.y; @@ -22,7 +23,7 @@ __global__ void fwht_cuda(const float * src, float * dst, const int64_t n_rows, ggml_cuda_pdl_sync(); #pragma unroll for (int i = 0; i < el_w; ++i) { - reg[i] = src[i * warp_size + lane] * scale; + reg[i] = ggml_cuda_cast(src[i * warp_size + lane]) * scale; } #pragma unroll @@ -58,15 +59,12 @@ __global__ void fwht_cuda(const float * src, float * dst, const int64_t n_rows, } } -bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) { - GGML_ASSERT(ggml_are_same_shape(src, dst)); - if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst)) { - return false; - } +template +static bool ggml_cuda_op_fwht_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) { const int n = src->ne[0]; const int64_t rows = ggml_nrows(src); - const float * src_d = (const float *) src->data; + const T * src_d = (const T *) src->data; float * dst_d = (float *) dst->data; const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; @@ -84,18 +82,47 @@ bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, switch (n) { case 64: - ggml_cuda_kernel_launch(fwht_cuda<64>, launch_params, src_d, dst_d, rows, scale); + ggml_cuda_kernel_launch(fwht_cuda<64, T>, launch_params, src_d, dst_d, rows, scale); return true; case 128: - ggml_cuda_kernel_launch(fwht_cuda<128>, launch_params, src_d, dst_d, rows, scale); + ggml_cuda_kernel_launch(fwht_cuda<128, T>, launch_params, src_d, dst_d, rows, scale); return true; case 256: - ggml_cuda_kernel_launch(fwht_cuda<256>, launch_params, src_d, dst_d, rows, scale); + ggml_cuda_kernel_launch(fwht_cuda<256, T>, launch_params, src_d, dst_d, rows, scale); return true; case 512: - ggml_cuda_kernel_launch(fwht_cuda<512>, launch_params, src_d, dst_d, rows, scale); + ggml_cuda_kernel_launch(fwht_cuda<512, T>, launch_params, src_d, dst_d, rows, scale); return true; default: return false; } } + +bool ggml_cuda_op_mul_mat_use_fwht(const struct ggml_tensor * op) { + const struct ggml_tensor * a = op->src[0]; + const struct ggml_tensor * b = op->src[1]; + + return op->op == GGML_OP_MUL_MAT && ggml_get_op_params_i32(op, 1) == GGML_HINT_SRC0_IS_HADAMARD && + a->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && + (b->type == GGML_TYPE_F32 || b->type == GGML_TYPE_F16) && ggml_is_contiguous(b) && ggml_is_contiguous(op) && + ggml_are_same_shape(b, op); +} + +bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) { + GGML_ASSERT(ggml_are_same_shape(src, dst)); + if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst)) { + return false; + } + if (dst->type != GGML_TYPE_F32) { + return false; + } + + switch (src->type) { + case GGML_TYPE_F32: + return ggml_cuda_op_fwht_impl(ctx, src, dst); + case GGML_TYPE_F16: + return ggml_cuda_op_fwht_impl(ctx, src, dst); + default: + return false; + } +} diff --git a/ggml/src/ggml-cuda/fwht.cuh b/ggml/src/ggml-cuda/fwht.cuh index cf3df94ca..e7ad53475 100644 --- a/ggml/src/ggml-cuda/fwht.cuh +++ b/ggml/src/ggml-cuda/fwht.cuh @@ -1,4 +1,6 @@ #include "common.cuh" +bool ggml_cuda_op_mul_mat_use_fwht(const struct ggml_tensor * op); + // Returns whether the Fast Walsh-Hadamard transform could be used. bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index e76ff3128..cf9a587f7 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1818,8 +1818,7 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { GGML_TENSOR_BINARY_OP_LOCALS - const int32_t hint = ggml_get_op_params_i32(dst, 1); - if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) { + if (ggml_cuda_op_mul_mat_use_fwht(dst) && ggml_cuda_op_fwht(ctx, src1, dst)) { return; } @@ -5212,7 +5211,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g if (a->nb[0] != ggml_element_size(a) || b->nb[0] != ggml_element_size(b)) { return false; // TODO this could in principle be implemented though currently there is no use case. } - if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16) { + if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16 && !ggml_cuda_op_mul_mat_use_fwht(op)) { return false; } if (op->op == GGML_OP_MUL_MAT_ID && ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) {