From 877a495245876200038ad55f8f76206c717ddea7 Mon Sep 17 00:00:00 2001
From: MaggotHATE <clay1326@gmail.com>
Date: Fri, 8 Nov 2024 11:55:10 +0500
Subject: [PATCH] Fix to guarantee K-Shift on the first step only

---
 src/llama-sampling.cpp | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/llama-sampling.cpp b/src/llama-sampling.cpp
index 03ef703e9..59786f708 100644
--- a/src/llama-sampling.cpp
+++ b/src/llama-sampling.cpp
@@ -1107,14 +1107,18 @@ static const char * llama_sampler_k_shift_name(const struct llama_sampler * /*sm
 static void llama_sampler_k_shift_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
     auto * ctx = (llama_sampler_k_shift *) smpl->ctx;
 
-    if (ctx->k_set == true
-        || ctx->k <= 0
-        || ctx->k >= (int) cur_p->size) {
+    // ensures that k-shift can happen on the first step only
+    if (ctx->k_set != true) {
+        ctx->k_set = true;
+    } else {
+        return;
+    }
+
+    if (ctx->k <= 0 || ctx->k >= (int) cur_p->size) {
         return;
     }
 
     llama_sampler_top_shift_impl(cur_p, ctx->k);
-    ctx->k_set = true;
 }
 
 static struct llama_sampler * llama_sampler_k_shift_clone(const struct llama_sampler * smpl) {