From 04662748aadc589be83a8dfdaedce2bc4d91fae0 Mon Sep 17 00:00:00 2001
From: Jeff Bolz <jbolz@nvidia.com>
Date: Wed, 20 Nov 2024 13:47:36 -0600
Subject: [PATCH] vulkan: predicate max operation in soft_max shaders/soft_max
 (llama/10437)

Fixes #10434
---
 ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp b/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp
index f9727679..6e20b641 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp
@@ -73,7 +73,9 @@ void soft_max(uint num_iters) {
 
         FLOAT_TYPE v = a * p.scale + slope * b;
 
-        max_val = max(max_val, v);
+        if (col < p.KX) {
+            max_val = max(max_val, v);
+        }
 
         if (idx < DATA_CACHE_SIZE) {
             data_cache[idx] = v;