Skip to content

Commit 02c66f2

Browse files
author
Chris Warren-Smith
committed
LLAMA: nitro fixes
1 parent 287645b commit 02c66f2

3 files changed

Lines changed: 59 additions & 53 deletions

File tree

‎llama/llama-sb.cpp‎

Lines changed: 24 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -274,7 +274,7 @@ bool Llama::add_message(LlamaIter &iter, const string &role, const string &conte
274274
_n_system_tokens = prompt_tokens.size();
275275
}
276276

277-
if (!make_space_for_tokens(prompt_tokens.size())) {
277+
if (!make_space_for_tokens((prompt_tokens.size() * 3) / 2)) {
278278
return false;
279279
}
280280

@@ -485,17 +485,8 @@ bool Llama::batch_decode_tokens(vector<llama_token> &tokens) {
485485
size_t batch_size = std::min((size_t)n_batch, tokens.size() - i);
486486
llama_batch batch = llama_batch_get_one(tokens.data() + i, batch_size);
487487
int result = llama_decode(_ctx, batch);
488-
if (result == 1) {
489-
// KV full mid-batch - try to evict and retry once
490-
if (!make_space_for_tokens(n_batch)) {
491-
_last_error = "KV cache exhausted, cannot evict enough space";
492-
return false;
493-
}
494-
result = llama_decode(_ctx, batch);
495-
}
496488
if (result != 0) {
497-
_last_error = std::format("Failed to decode batch. position:{} error:{} [size:{}]",
498-
i, result, tokens.size());
489+
set_decode_error(result, i, tokens.size());
499490
return false;
500491
}
501492
}
@@ -589,11 +580,13 @@ bool Llama::make_space_for_tokens(int n_tokens) {
589580
return false;
590581
}
591582

583+
llama_pos remove_start = pos_min + _n_system_tokens;
584+
592585
// Remove oldest tokens (from pos_min to pos_min + tokens_to_remove)
593-
llama_memory_seq_rm(mem, 0, pos_min, pos_min + tokens_to_remove);
586+
llama_memory_seq_rm(mem, 0, remove_start, remove_start + tokens_to_remove);
594587

595588
// Shift remaining tokens down
596-
llama_memory_seq_add(mem, 0, pos_min + tokens_to_remove, -1, -tokens_to_remove);
589+
llama_memory_seq_add(mem, 0, remove_start + tokens_to_remove, -1, -tokens_to_remove);
597590

598591
return true;
599592
}
@@ -668,3 +661,21 @@ void Llama::set_last_error(const char *message) {
668661
_last_error = std::format("{} failed", message);
669662
}
670663
}
664+
665+
void Llama::set_decode_error(int32_t error, int index, int num_tokens) {
666+
if (error == 1) {
667+
llama_memory_t mem = llama_get_memory(_ctx);
668+
llama_pos pos_min = llama_memory_seq_pos_min(mem, 0);
669+
llama_pos pos_max = llama_memory_seq_pos_max(mem, 0);
670+
int n_ctx = llama_n_ctx(_ctx);
671+
int current_used = pos_max - pos_min + 1;
672+
int space_needed = num_tokens;
673+
int space_available = n_ctx - current_used;
674+
_n_system_tokens;
675+
_last_error = std::format("KV exhausted. Reduce batch or context sizes. batchNo:{} requested:{} available:{}",
676+
index, space_needed, space_available);
677+
} else {
678+
auto message = error == 2 ? "abort" : error == -1 ? "invalid" : "fatal";
679+
_last_error = std::format("Failed to decode batch. batchNo:{} error:'{}'", index, message);
680+
}
681+
}

‎llama/llama-sb.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,7 @@ struct Llama {
123123
vector<llama_token> tokenize(const string &prompt);
124124
string token_to_string(LlamaIter &iter, llama_token tok);
125125
void set_last_error(const char *message);
126+
void set_decode_error(int32_t error, int index, int num_tokens);
126127

127128
llama_model *_model;
128129
llama_context *_ctx;

‎llama/nitro.cpp‎

Lines changed: 34 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,6 @@
22
// A standalone agentic LLM shell with notcurses TUI.
33
// Uses llama-sb.h as the sole llama.cpp integration layer.
44
//
5-
// Build (example):
6-
// g++ -std=c++20 -O2 nitro.cpp llama-sb.cpp \
7-
// -I/path/to/llama.cpp/include \
8-
// -L/path/to/llama.cpp/build/src \
9-
// -lllama -lggml -lnotcurses-core -lnotcurses -lcurl \
10-
// -o nitro
11-
//
125
// Usage:
136
// ./nitro [options] [project_dir]
147
//
@@ -2247,6 +2240,7 @@ bool AgentState::run_turn(const std::string &user_message, const NitroConfig &cf
22472240
}
22482241
tui.append_line(ICON_ERR + content);
22492242
}
2243+
tui.tokens_per_sec = tokens_per_sec();
22502244
if (!llama->add_message(*iter, "tool_result", content)) {
22512245
tui.append_line(ICON_ERR + "tool result inject: " + llama->last_error());
22522246
}
@@ -2289,11 +2283,15 @@ bool AgentState::run_turn(const std::string &user_message, const NitroConfig &cf
22892283
return ni.id == NCKEY_ESC;
22902284
};
22912285

2292-
auto fetch_all = [&]() -> void {
2286+
auto fetch_tool = [&]() -> void {
22932287
while (iter->_has_next && !is_escape()) {
22942288
std::string tok = llama->next(*iter);
22952289
buffer += tok;
22962290
tui.tick_spinner();
2291+
auto pos = buffer.find("</think>");
2292+
if (pos != std::string::npos) {
2293+
break;
2294+
}
22972295
}
22982296
};
22992297

@@ -2328,46 +2326,42 @@ bool AgentState::run_turn(const std::string &user_message, const NitroConfig &cf
23282326
end_think("<think|>");
23292327
end_think("<channel|>");
23302328
}
2331-
if (think_mode == t_thunk) {
2332-
auto tool_start = buffer.find("TOOL:");
2333-
if (tool_start == 0) {
2334-
fetch_all();
2335-
invoke_tool(trim(buffer), "TOOL_RESULT: {}");
2336-
buffer.clear();
2337-
think_mode = t_init;
2338-
continue;
2339-
}
2340-
tool_start = buffer.find("<|tool_call>call:");
2341-
if (tool_start != std::string::npos) {
2342-
// see https://ai.google.dev/gemma/docs/core/prompt-formatting-gemma4
2343-
fetch_all();
2344-
auto pos = buffer.find_last_not_of("}<tool_call|>");
2345-
if (pos != std::string::npos) {
2346-
buffer = buffer.substr(0, pos);
2347-
}
2348-
pos = buffer.find_first_not_of('{');
2349-
if (pos != std::string::npos) {
2350-
buffer = buffer.substr(0, pos) + buffer.substr(pos + 1);
2351-
}
2352-
invoke_tool(trim(buffer), "<|tool_response>{}<tool_response|>");
2353-
buffer.clear();
2354-
think_mode = t_init;
2355-
continue;
2356-
}
2357-
auto pos = buffer.find('\n');
2329+
auto tool_start = buffer.find("TOOL:");
2330+
if (tool_start == 0) {
2331+
fetch_tool();
2332+
invoke_tool(trim(buffer), "TOOL_RESULT: {}");
2333+
buffer.clear();
2334+
think_mode = t_init;
2335+
continue;
2336+
}
2337+
tool_start = buffer.find("<|tool_call>call:");
2338+
if (tool_start != std::string::npos) {
2339+
// see https://ai.google.dev/gemma/docs/core/prompt-formatting-gemma4
2340+
fetch_tool();
2341+
auto pos = buffer.find_last_not_of("}<tool_call|>");
23582342
if (pos != std::string::npos) {
2359-
tui.append_token(buffer.substr(0, pos + 1));
2360-
buffer = buffer.substr(pos + 1);
2343+
buffer = buffer.substr(0, pos);
23612344
}
2362-
} else {
2363-
auto pos = buffer.find('\n');
2345+
pos = buffer.find_first_not_of('{');
23642346
if (pos != std::string::npos) {
2347+
buffer = buffer.substr(0, pos) + buffer.substr(pos + 1);
2348+
}
2349+
invoke_tool(trim(buffer), "<|tool_response>{}<tool_response|>");
2350+
buffer.clear();
2351+
think_mode = t_init;
2352+
continue;
2353+
}
2354+
auto pos = buffer.find('\n');
2355+
if (pos != std::string::npos) {
2356+
if (think_mode == t_think) {
23652357
auto thought = buffer.substr(0, pos + 1);
23662358
if (thought.length() > 1) {
23672359
tui.append_token(ICON_THINK + thought);
23682360
}
2369-
buffer = buffer.substr(pos + 1);
2361+
} else {
2362+
tui.append_token(buffer.substr(0, pos + 1));
23702363
}
2364+
buffer = buffer.substr(pos + 1);
23712365
}
23722366
}
23732367

0 commit comments

Comments
 (0)