From 77fe4fd982776afa3909e53857d502e557f8416a Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Fri, 13 Dec 2024 17:19:23 +0800 Subject: [PATCH] RWKV_WKV6 Vulkan op tests passed Signed-off-by: Molly Sophia --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 9 ++++----- .../ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp | 2 ++ .../vulkan-shaders/{rwkv_wkv6.comp => wkv6.comp} | 0 3 files changed, 6 insertions(+), 5 deletions(-) rename ggml/src/ggml-vulkan/vulkan-shaders/{rwkv_wkv6.comp => wkv6.comp} (100%) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index e103e67f7..da11e88cd 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1961,7 +1961,7 @@ static void ggml_vk_load_shaders(vk_device& device) { "main", 7, sizeof(vk_op_rwkv_wkv6_push_constants), - {64, 1, 1}, // work group + {1, 1, 1}, // work group {device->subgroup_size}, 1 ); @@ -8344,11 +8344,10 @@ static void ggml_vk_check_results_0(ggml_tensor * tensor) { } else if (tensor->op == GGML_OP_LEAKY_RELU) { const float * op_params = (const float *)tensor->op_params; tensor_clone = ggml_leaky_relu(ggml_ctx, src0_clone, op_params[0], false); + } else if (tensor->op == GGML_OP_RWKV_WKV6) { + tensor_clone = ggml_rwkv_wkv6(ggml_ctx, tensor->src[0], tensor->src[1], tensor->src[2], tensor->src[3], + tensor->src[4], tensor->src[5]); } - // else if (tensor->op == GGML_OP_RWKV_WKV6) { - // tensor_clone = ggml_rwkv_wkv6(ggml_ctx, tensor->src[0], tensor->src[1], tensor->src[2], tensor->src[3], - // tensor->src[4], tensor->src[5]); - // } else { std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl; GGML_ABORT("fatal error"); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index c48a228ae..eff60f3c3 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -479,6 +479,8 @@ void process_shaders() { string_to_spv("pool2d_f32", "pool2d.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("rwkv_wkv6_f32", "wkv6.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"C_TYPE", "float"}, {"D_TYPE", "float"}, {"E_TYPE", "float"}, {"F_TYPE", "float"}, {"S_TYPE", "float"}})); + for (auto &c : compiles) { c.wait(); } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rwkv_wkv6.comp b/ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp similarity index 100% rename from ggml/src/ggml-vulkan/vulkan-shaders/rwkv_wkv6.comp rename to ggml/src/ggml-vulkan/vulkan-shaders/wkv6.comp