vulkan: fix rms_norm_mul to handle broadcasting dim0 (#14817)

2025-07-22 10:35:21 -05:00 · 2025-07-22 10:35:21 -05:00 · 84712b6043
commit 84712b6043
parent d4d1522b20
2 changed files with 9 additions and 3 deletions
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@ -10248,7 +10248,7 @@ static bool ggml_vk_can_fuse(const struct ggml_cgraph * cgraph, int node_idx, st
        }
        // if rms_norm is the B operand, then we don't handle broadcast
        if (rms_norm == mul->src[1] &&
-            mul->src[0]->ne[1] != rms_norm->ne[1]) {
+            !ggml_are_same_shape(mul->src[0], rms_norm)) {
            return false;
        }
        // rms_norm shader assumes contiguous rows