Feature Request: Support for Kimi K2 Thinking tool calling

[Lissanro]

Prerequisites

• I am running the latest code. Mention the version if possible as well.
• I carefully followed the README.md.
• I searched using keywords relevant to my issue to make sure that I am creating a new issue that is not already open (or closed).
• I reviewed the Discussions, and have a new and useful enhancement to share.

Feature Description

I noticed that even without thinking tool calls for example in Roo Code often fail and get rendered as XML.

I searching why this could be happening, I found this: "Prevent special token leakage in KimiK2ToolParser streaming mode" vllm-project/vllm#28543 - it suggests the issue may be needed to be handled at the backend level.

I also found a bug report at llama.cpp: "Kimi-K2-Thinking reasoning and tool calling support" ggml-org/llama.cpp#17155 - but I am not sure if it referes to the same issue, because there ggml-org/llama.cpp#16932 is mention as a possible solution, but it does not list Kimi K2 Thinking as a supported model, only Kimi K2, which already was working fine for me.

I would appreciate if someone more knowledgable would look into this, does ik_llama.cpp need an update for tool colling handling? Or is it an issue for agents like Roo Code to solve?

Motivation

Kimi K2 Thinking relies a lot on tool calling (according to the official model card):

Deep Thinking & Tool Orchestration: End-to-end trained to interleave chain-of-thought reasoning with function calls, enabling autonomous research, coding, and writing workflows that last hundreds of steps without drift.

The issue described here happens even with thinking disabled though, and prevents Roo Code and other similar agents from working.

If turns out issues I pointed out are not in ik_llama.cpp but in the agent tool calling handling, please feel free to close this. But seeing how some possibly related issues are being considered to be addressed at the backend level in vllm and llama.cpp, I thought it is worth reporting here and sharing relevant information I found about it.

Possible Implementation

No response

[hksdpc255]

Kimi K2 Thinking and Kimi K2 appear to use the same chat format. If you have time, could you test PR #16932 and share any feedback or bugs you encounter?

[Lissanro]

I tried #958 which is equivalent for #16932 at llama.cpp, maybe it got better for tool calls outside of the think block, but still failing tool calling inside the think block. It just prints XML and then usually fails to generate any actual response, most likely because the model expected tool calling to return something but it did not. It looks like this:

<read_file>
<args>
<file>
<path>...

Since Kimi K2 Thinking likes to call tools while thinking (possibly multiple times in one think block), currently the model does not really work so I could not yet use it for any real work projects. If there is something else I need to test, please let me know.

For reference, this is how I run it:

numactl --cpunodebind=0 --interleave=all ~/pkgs/ik_llama.cpp/build/bin/llama-server \
--model /mnt/neuro/models/Kimi-K2-Thinking-Q8_0-Q4_0.gguf \
--jinja --special \
--ctx-size 131072 --n-gpu-layers 62 --tensor-split 12,26,31,31 -mla 3 -ctk q8_0 -amb 512 -b 4096 -ub 4096 \
-ot "blk\.(3)\.ffn_.*=CUDA0" \
-ot "blk\.(4)\.ffn_.*=CUDA1" \
-ot "blk\.(5)\.ffn_.*=CUDA2" \
-ot "blk\.(6)\.ffn_.*=CUDA3" \
-ot exps=CPU \
--threads 64 --host 0.0.0.0 --port 5000 \
--slot-save-path /var/cache/ik_llama.cpp/k2-thinking

I generate the quant myself using Ubergram Q4_X recipe from https://huggingface.co/ubergarm/Kimi-K2-Thinking-GGUF and also I replaced the chat template with https://huggingface.co/unsloth/Kimi-K2-Thinking/resolve/main/chat_template.jinja in order to include fixes from Unshoth. Still, tool calling in the think block is something is not working yet.

I tried also without --jinja --special, then the model seem to make correct tool calls in Roo Code more often, but also frequently fails to do them, like this for example:

Now let me check the TranslationContext to understand how translations are loaded:functions.read_file:3{"args": [{"path": "src/contexts/TranslationContext.tsx"}]}

Not sure if this is because the model does not work correctly without --jinja --special or if it is another tool call format that fails (it is also possible Roo Code does not display raw strings and strips XML tags outside the thinking block). But until there is support for K2 Thinking tool calling, it is hard to test it while --jinja --special failing tool calls in the think block, since this seems to confuse the model.

I am happy to provide mode debug information if needed, so if there is anything I can do to help, please let me know.

[hksdpc255]

Kimi-K2 uses a tool-call format like this: <|tool_calls_section_begin|><|tool_call_begin|>functions.terminal:0<|tool_call_argument_begin|>{"command": "df -h", "cd": "."}<|tool_call_end|><|tool_calls_section_end|> accroding to its spec.

The snippet:

<read_file>
<args>
<file>
<path>...

appears to be a malformed tool call emitted by the model, rather than a parser bug.

[Lissanro]

This is very consistent and works outside the think block in most cases, without #958 patch similar XML is generated even outside think blocks instead of actual tool calls. With #958 it works outside think block most of the time. So, safe to say this issue related to tool calling, and not caused my malformed tool calls as far as I can tell. How Roo Code shows them in its window may not represent what is happening under the hood, since for example it also transforms many other string including mark up, etc. and does not have an option to display raw tokens model generated.

What I am not sure of, how to verify if current ik_llama.cpp with #958 patch can understand multiple tool calls in the think block. If this supposed to work, then the issue is in Roo Code doing something wrong with thinking block tool calling. But my understanding this feature needs the backend support in the inference engine.

[hksdpc255]

What chat template you're using without the patch in #958?

[Lissanro]

I use built-in chat template. I tried with both with and without --jinja --special, without the #958 K2 Thinking does not work at all, failing most tool calls even without the think block. Patching helps but only if I drop "--jinja --special" to disable thinking capability. Otherwise, all tool calls in the think block still failing. I tested with Roo Code, the latest version. Built-in jinja file in the GGUF in my case is this one: https://huggingface.co/unsloth/Kimi-K2-Thinking/resolve/main/chat_template.jinja

[hksdpc255]

Surprisingly, the built-in chat template does not crash. However, I don’t think Minja supports this snippet:

 {%- set formatted_id = tool_call['id'] -%}
    <|tool_call_begin|>{{ formatted_id }}<|tool_call_argument_begin|>{% if tool_call['function']['arguments'] is string %}{{ tool_call['function']['arguments'] }}{% else %}{{ tool_call['function']['arguments'] | tojson }}{% endif %}<|tool_call_end|>

Could you confirm what tool_call['id'] actually renders to?

[Lissanro]

Please let me know what commands to run, and I will be happy to test.

My apologies if asking something obvious, but until few days ago when I tried Kimi K2 Thinking for the first time, all tool calls where just working with the older models so I am not yet very familiar how to debug them, and only recently started learning details how they work.

[hksdpc255]

Without my patch, Kimi-K2 would fall back to a generic grammar and parser. While this may work, it can result in decreased performance for the model.

The goal of my work is to implement the official, documented tool-call format in order to fully leverage the model’s capabilities. However, it seems that the current chat template still contains some issues and will need further fixes. I don't know where to ask for help.

[hksdpc255]

@Lissanro I generated a further-refined chat template using AI. Hopefully this version is more accurate and can help address the remaining issues:

k2_templates.zip

k2i is for non-reasoning, k2t is for reasoning

[Lissanro]

To debug, I made this simple patch to log exact tokens ik_llama.cpp generates:

diff --git a/examples/server/server.cpp b/examples/server/server.cpp
index 966947c0..d452fdc1 100644
--- a/examples/server/server.cpp
+++ b/examples/server/server.cpp
@@ -2288,6 +2288,10 @@ struct server_context {
         const std::string token_str = result.text_to_send;
         slot.sampled = result.tok;
 
+        const std::string piece = llama_token_to_piece(slot.ctx, result.tok, true);
+        printf("token %6d: '%s'\n", result.tok, piece.c_str());
+        fflush(stdout);
+
         // search stop word and delete it
         slot.generated_text += token_str;
         slot.has_next_token = true;

With #958 applied, I get frequent crashes, like this:

token 163595: '<|tool_calls_section_begin|>'
token 163597: '<|tool_call_begin|>'
token     77: 'n'
token     27: '<'
token   1396: 'read'
terminate called after throwing an instance of 'std::runtime_error'
  what():  Invalid diff: now finding less tool calls!
zsh: IOT instruction (core dumped)  CUDA_VISIBLE_DEVICES="0,1,2,3" numactl --cpunodebind=0 --interleave=all

Or like this:

token 163595: '<|tool_calls_section_begin|>'
token 163597: '<|tool_call_begin|>'
token  41937: 'functions'
token   8827: '.read'
token   6101: '_file'
token     25: ':'
token     15: '0'
token 163598: '<|tool_call_argument_begin|>'
token   8264: '{"'
token   3992: 'args'
token   1289: '":'
token  94369: ' [{"'
token   4953: 'path'
token   1289: '":'
token    414: ' "'
token   6435: 'src'
token  52798: '/providers'
token     14: '/'
token  18482: 'Theme'
token   6977: 'Provider'
token 158213: '.tsx'
token  24438: '"},'
token   9392: ' {"'
token   4953: 'path'
token   1289: '":'
token    414: ' "'
token   6435: 'src'
token  23991: '/components'
token     14: '/'
token   7238: 'Header'
token 158213: '.tsx'
token  24438: '"},'
token   9392: ' {"'
token   4953: 'path'
token   1289: '":'
token    414: ' "'
token   6435: 'src'
token  23991: '/components'
token     14: '/'
token  18482: 'Theme'
token  32336: 'Toggle'
token 158213: '.tsx'
token  24438: '"},'
token   9392: ' {"'
token   4953: 'path'
token   1289: '":'
token    414: ' "'
token   6435: 'src'
token  27938: '/app'
token 155256: '/globals'
token  16184: '.css'
token  24438: '"},'
token   9392: ' {"'
token   4953: 'path'
token   1289: '":'
token    414: ' "'
token   6435: 'src'
token  27938: '/app'
token  92806: '/layout'
token 158213: '.tsx'
token  16934: '"}'
token     60: ']'
token  18727: '}</'
terminate called after throwing an instance of 'std::runtime_error'
  what():  Invalid diff: '{"args":[{"path":"src/providers/ThemeProvider.tsx"},{"path":"src/components/Header.tsx"},{"path":"src/components/ThemeToggle.tsx"},{"path":"src/app/globals.css"},{"path":"src/app/layout.tsx"}]' not found at start of '{"args":"[{\"path\": \"src/providers/ThemeProvider.tsx\"}"'
zsh: IOT instruction (core dumped)  CUDA_VISIBLE_DEVICES="0,1,2,3" numactl --cpunodebind=0 --interleave=all

Another type of crash:

token   9413: 'Let'
token   1019: ' me'
token   5661: ' implement'
token    276: ' the'
token  41291: ' fixes'
token   2215: ' now'
token     25: ':'
token 163595: '<|tool_calls_section_begin|>'
token 163597: '<|tool_call_begin|>'
token   2482: 'user'
token 163606: '<think>'
token   1008: 'The'
token   2742: ' user'
token    387: ' is'
token  51107: ' frustrated'
token   2416: ' because'
token    374: ' I'
token  22381: ' haven't'
token   5305: ' actually'
token   2615: ' made'
token   1178: ' any'
token   3253: ' code'
token   5772: ' changes'
token   5412: ' yet'
...
token   9413: 'Let'
token   1019: ' me'
token   1628: ' make'
token   1946: ' these'
token   5772: ' changes'
token   2389: ' using'
token    276: ' the'
token   5504: ' apply'
token  28715: '_diff'
token   4453: ' tool'
token     13: '.'
token 163607: '</think>'
token     27: '<'
token  10601: 'apply'
terminate called after throwing an instance of 'std::runtime_error'
  what():  Invalid diff: now finding less tool calls!
zsh: IOT instruction (core dumped)  CUDA_VISIBLE_DEVICES="0,1,2,3" numactl --cpunodebind=0 --interleave=all      

(I shortened this log a bit by replacing normal tokens in the middle with "...", since the problem seems to be caused by special token handling)

I tried you k2i, the model still tends to emit think block still (as shown in one of the logs above), but so far it avoided making tool calls inside it, while using the k2i template. Aside from the crashes above, the k2i template seems to work, and while using it I did not see malfolmed tool calls that were happening sometimes without it. But as I mentioned, I get frequent crashes, so I think there are bugs in the #958 patch. I think the k2i chat template made crashes more frequent, making them easier to reproduce.

I will test k2t next, and will report back.

[Lissanro]

(a message was here about k2t template, but it turned out to be Roo Code failure to update todo list if attempted within the think block, and not an actual tool call, so I removed this message for now; I will keep testing the k2t template further and report back when I have more information; issues mentioned in the previous message about testing with #958 patch and k2i template still apply, and I did not yet find any workaround to avoid frequent crashes)