Merge branch 'ikawrakow:main' into main

Author: Thireus

examples/server/server.cpp

@@ -1968,7 +1968,9 @@ struct server_context {
                 slot.generated_text.erase(
                     slot.generated_text.begin() + pos + stop_pos,
                     slot.generated_text.end());
-                pos = std::min(slot.n_sent_text, slot.generated_text.size());
+                // Update n_sent_text to not exceed the new generated_text size
+                slot.n_sent_text = std::min(slot.n_sent_text, slot.generated_text.size());
+                pos = slot.n_sent_text;
             } else {
                 is_stop_full = false;
                 stop_pos = slot.find_stopping_strings(str_test, token_str.size(), false);
@@ -1980,6 +1982,9 @@ struct server_context {
                 result.text_to_send = slot.generated_text.substr(pos, std::string::npos);
                 slot.n_sent_text += result.text_to_send.size();
                 // add the token to slot queue and cache
+            } else if (stop_pos != std::string::npos) {
+                // Handle partial stop - update n_sent_text to the end of the current text
+                slot.n_sent_text = slot.generated_text.size();
             }
 
             slot.add_token_string(result);

ggml/include/ggml.h

@@ -2235,7 +2235,7 @@ extern "C" {
             int                   min_entries,
             float                 thresh);
 
-#define GGML_KQ_MASK_PAD 64
+#define GGML_KQ_MASK_PAD 16
 
     // q:    [n_embd, n_batch,     n_head,    1]
     // k:    [n_embd, n_kv,        n_head_kv, 1]

src/llama-build-context.cpp

@@ -276,6 +276,12 @@ ggml_tensor * llm_build_context::build_inp_out_ids() {
 }
 
 ggml_tensor * llm_build_context::build_inp_KQ_mask(bool causal) {
+    if (causal && flash_attn) {
+        lctx.inp_KQ_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F16, n_kv, GGML_PAD(n_tokens, GGML_KQ_MASK_PAD));
+        cb(lctx.inp_KQ_mask, "KQ_mask", -1);
+        ggml_set_input(lctx.inp_KQ_mask);
+        return lctx.inp_KQ_mask;
+    }
     lctx.inp_KQ_mask = causal
         ? ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv,     GGML_PAD(n_tokens, GGML_KQ_MASK_PAD))
         : ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_tokens, GGML_PAD(n_tokens, GGML_KQ_MASK_PAD));
@@ -287,6 +293,12 @@ ggml_tensor * llm_build_context::build_inp_KQ_mask(bool causal) {
 
 ggml_tensor * llm_build_context::build_inp_KQ_mask_swa(bool causal) {
     GGML_ASSERT(hparams.n_swa > 0);
+    if (causal && flash_attn) {
+        lctx.inp_KQ_mask_swa = ggml_new_tensor_2d(ctx0, GGML_TYPE_F16, n_kv, GGML_PAD(n_tokens, GGML_KQ_MASK_PAD));
+        cb(lctx.inp_KQ_mask_swa, "KQ_mask_swa", -1);
+        ggml_set_input(lctx.inp_KQ_mask_swa);
+        return lctx.inp_KQ_mask_swa;
+    }
 
     lctx.inp_KQ_mask_swa = causal
         ? ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv,     GGML_PAD(n_tokens, GGML_KQ_MASK_PAD))