From 820de57d4faa427a3d0bfb14e48057247fae036e Mon Sep 17 00:00:00 2001
From: Xuan-Son Nguyen <son@huggingface.co>
Date: Thu, 24 Jul 2025 13:59:56 +0200
Subject: [PATCH] chat : fix kimi-k2 chat template (#14852)

---
 src/llama-arch.cpp | 12 ++++++------
 src/llama-chat.cpp |  7 +++----
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
index 814ac93a6..062a99776 100644
--- a/src/llama-arch.cpp
+++ b/src/llama-arch.cpp
@@ -1933,12 +1933,6 @@ static const std::map<llm_arch, std::map<llm_tensor, const char *>> LLM_TENSOR_N
             { LLM_TENSOR_TOKEN_EMBD_NORM,   "token_embd_norm" },
         }
     },
-    {
-        LLM_ARCH_UNKNOWN,
-        {
-            { LLM_TENSOR_TOKEN_EMBD,      "token_embd" },
-        },
-    },
     {
         LLM_ARCH_DREAM,
         {
@@ -1956,6 +1950,12 @@ static const std::map<llm_arch, std::map<llm_tensor, const char *>> LLM_TENSOR_N
             { LLM_TENSOR_FFN_UP,          "blk.%d.ffn_up" },
         },
     },
+    {
+        LLM_ARCH_UNKNOWN,
+        {
+            { LLM_TENSOR_TOKEN_EMBD,      "token_embd" },
+        },
+    },
 };
 
 static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
diff --git a/src/llama-chat.cpp b/src/llama-chat.cpp
index 80072ad27..d34bb2687 100644
--- a/src/llama-chat.cpp
+++ b/src/llama-chat.cpp
@@ -718,10 +718,9 @@ int32_t llm_chat_apply_template(
             }
 
             ss << message->content << "<|im_end|>";
-
-            if (add_ass) {
-                ss << "<|im_assistant|>assistant<|im_middle|>";
-            }
+        }
+        if (add_ass) {
+            ss << "<|im_assistant|>assistant<|im_middle|>";
         }
     } else {
         // template not supported