From 030a7999d2dae43e27c5dd10aee87d1e12a378ef Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 22 Oct 2026 17:45:28 -0700 Subject: [PATCH] ggml-patches/0015: Vulkan RMS norm with 128 threads for rows of at most 128 elements (bit-identical) --- ...0015-ggml-vulkan-rms-norm-short-rows.patch | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 scripts/unsloth/ggml-patches/0015-ggml-vulkan-rms-norm-short-rows.patch diff --git a/scripts/unsloth/ggml-patches/0015-ggml-vulkan-rms-norm-short-rows.patch b/scripts/unsloth/ggml-patches/0015-ggml-vulkan-rms-norm-short-rows.patch new file mode 100644 index 000000000..1cae17c8e --- /dev/null +++ b/scripts/unsloth/ggml-patches/0015-ggml-vulkan-rms-norm-short-rows.patch @@ -0,0 +1,64 @@ +diff --git a/src/ggml-vulkan/ggml-vulkan-types.h b/src/ggml-vulkan/ggml-vulkan-types.h +index 3d71a8b9..990c60b4 100644 +--- a/src/ggml-vulkan/ggml-vulkan-types.h ++++ b/src/ggml-vulkan/ggml-vulkan-types.h +@@ -851,6 +851,8 @@ struct vk_device_struct { + vk_pipeline pipeline_group_norm_f32; + vk_pipeline pipeline_rms_norm_f32; + vk_pipeline pipeline_rms_norm_mul_f32; ++ vk_pipeline pipeline_rms_norm_small_f32; ++ vk_pipeline pipeline_rms_norm_mul_small_f32; + vk_pipeline pipeline_rms_norm_mul_add_f32; + vk_pipeline pipeline_rms_norm_mul_add_mul_f32; + vk_pipeline pipeline_rms_norm_mul_add_partials_f32; +diff --git a/src/ggml-vulkan/ggml-vulkan.cpp b/src/ggml-vulkan/ggml-vulkan.cpp +index dc6eeb62..163e97e9 100644 +--- a/src/ggml-vulkan/ggml-vulkan.cpp ++++ b/src/ggml-vulkan/ggml-vulkan.cpp +@@ -3268,6 +3268,8 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { + + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_f32, "rms_norm_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_f32, "rms_norm_mul_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_rms_norm_small_f32, "rms_norm_small_f32", rms_norm_small_f32_len, rms_norm_small_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); ++ ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_small_f32, "rms_norm_mul_small_f32", rms_norm_small_f32_len, rms_norm_small_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_f32, "rms_norm_mul_add_f32", rms_norm_mul_add_f32_len, rms_norm_mul_add_f32_data, "main", 5, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 0}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_mul_f32, "rms_norm_mul_add_mul_f32", rms_norm_mul_add_f32_len, rms_norm_mul_add_f32_data, "main", 5, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_partials_f32, "rms_norm_mul_add_partials_f32", rms_norm_mul_add_partials_f32_len, rms_norm_mul_add_partials_f32_data, "main", 6, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 0}, 1, true); +@@ -9103,6 +9105,9 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const + if (ctx->do_add_rms_partials) { + return ctx->fused_rms_norm_mode == RMS_NORM_MUL ? ctx->device->pipeline_rms_norm_mul_partials_f32 : ctx->device->pipeline_rms_norm_partials_f32; + } ++ if (src0->ne[0] <= 128) { ++ return ctx->fused_rms_norm_mode == RMS_NORM_MUL ? ctx->device->pipeline_rms_norm_mul_small_f32 : ctx->device->pipeline_rms_norm_small_f32; ++ } + return ctx->fused_rms_norm_mode == RMS_NORM_MUL ? ctx->device->pipeline_rms_norm_mul_f32 : ctx->device->pipeline_rms_norm_f32; + } + return nullptr; +diff --git a/src/ggml-vulkan/vulkan-shaders/rms_norm.comp b/src/ggml-vulkan/vulkan-shaders/rms_norm.comp +index 314c8e56..832262eb 100644 +--- a/src/ggml-vulkan/vulkan-shaders/rms_norm.comp ++++ b/src/ggml-vulkan/vulkan-shaders/rms_norm.comp +@@ -39,7 +39,11 @@ layout (binding = 3) readonly buffer I {uvec2 data_i[];}; + #endif + + #extension GL_EXT_control_flow_attributes : enable ++// Rows of at most 128 elements (per-head QK norms) use a 128-thread variant: each thread still holds at most one ++// element and the tree below only drops the steps that added zeros, so the result is bit-identical to 512 threads. ++#ifndef BLOCK_SIZE + #define BLOCK_SIZE 512 ++#endif + + layout (constant_id = 1) const bool do_multiply = false; + #if RMS_NORM_ADD_FUSION +diff --git a/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +index 0d2f5513..4e567f47 100644 +--- a/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp ++++ b/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +@@ -854,6 +854,7 @@ void process_shaders() { + string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("group_norm_f32", "group_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); ++ string_to_spv("rms_norm_small_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"BLOCK_SIZE", "128"}})); + string_to_spv("rms_norm_mul_add_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_ADD_FUSION", "1"}})); + string_to_spv("rms_norm_mul_add_partials_f32", "rms_norm_partials.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_ADD_FUSION", "1"}})); + string_to_spv("rms_norm_set_rows_f32_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_SET_ROWS_FUSION", "1"}}));