kanshan/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-07-08 20:20:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

wip 3

Browse Source
This commit is contained in:
Xuan Son Nguyen
2026-06-20 18:56:58 +02:00
parent 2b2eed8fd7
commit b786bb2e60
1 changed files with 37 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/server/server-context.cpp
+37 -17
View File
@@ -2621,31 +2621,33 @@ private:
}
llama_batch batch_view;
int32_t i_next = 0;
int32_t off_next = 0;
int32_t n_batch = llama_n_batch(ctx_tgt);
for (int32_t i = 0; i < batch.size(); i = i_next) {
for (int32_t off = 0; off < batch.size(); off = off_next) {
try {
const int32_t n_tokens = std::min(n_batch, batch.size() - i);
// TODO @ngxson : maybe handle n_batch == 1 here instead of inside decode()
bool ok = decode(n_batch, batch_view);
const int32_t n_tokens = std::min(n_batch, batch.size() - off);
batch_view = batch.get_view(off, n_batch);
bool ok = decode(n_batch, off, batch_view);
if (ok) {
// move the head of the batch forward with the number of tokens we just processed
i_next = i + n_tokens;
off_next = off + n_tokens;
// on successful decode, restore the original batch size
n_batch = llama_n_batch(ctx_tgt);
} else {
// get a new batch view and try again
batch_view = batch.get_view(i, n_batch);
continue; // decode again
// try again with the updated n_batch
continue;
}
} catch (const std::exception & e) {
SRV_ERR("decode() failed: %s\n", e.what());
}
try {
post_decode(batch_view);
post_decode(n_batch, off, batch_view);
} catch (const std::exception & e) {
SRV_ERR("post_decode() failed: %s\n", e.what());
}
@@ -3401,8 +3403,8 @@ private:
// returns true = success ; false = retry with smaller batch size
// throw std::runtime_error on fatal error
bool decode(int32_t & n_batch, llama_batch & batch_view) {
SRV_DBG("decoding batch, n_tokens = %d\n", batch.size());
bool decode(int32_t & n_batch, int32_t off, llama_batch & batch_view) {
SRV_DBG("n_batch (effective) = %d, off = %d\n", n_batch, off);
auto & slot_batched = batch.slot_batched;
auto & alora_scale = batch.alora_scale;
@@ -3461,7 +3463,7 @@ private:
// TODO: handle ret == 2 (abort) when we start aborting
if (!err.empty()) {
SRV_ERR("%s n_batch = %d, ret = %d\n", err.c_str(), n_batch, ret);
SRV_ERR("%s off = %d, n_batch = %d, ret = %d\n", err.c_str(), off, n_batch, ret);
for (auto & slot : slots) {
if (slot.is_processing()) {
@@ -3482,10 +3484,9 @@ private:
// retry with half the batch size to try to find a free slot in the KV cache
if (!try_clear_idle_slots()) {
n_batch /= 2;
// TODO @ngxson : handle sub-batching
}
SRV_WRN("failed to find free space in the KV cache, retrying with smaller batch size, n_batch = %d, ret = %d\n", n_batch, ret);
SRV_WRN("failed to find free space in the KV cache, retrying with smaller batch size, off = %d, n_batch = %d, ret = %d\n", off, n_batch, ret);
return false; // retry with the updated n_batch
}
@@ -3526,7 +3527,22 @@ private:
return true;
}
void post_decode(llama_batch & batch_view) {
void post_decode(int32_t n_batch, int32_t off, llama_batch & batch_view) {
// for checking if a given batch index is inside batch_view
auto is_inside_view = [&](int32_t idx) {
return idx >= off && idx < off + n_batch;
};
// TODO @ngxson : it's tricky to make sub-batch compatible with common_sampler_sample_and_accept_n,
// so for now we will throw an error in this case: https://github.com/ggml-org/llama.cpp/issues/24840
for (auto & slot : slots) {
for (auto & i : slot.spec_i_batch) {
if (!is_inside_view(i)) {
throw std::runtime_error(string_format("speculative batch index %d is not inside the current sub-batch [%d, %d)", i, off, off + n_batch));
}
}
}
auto accept_special_token = [&](server_slot & slot, llama_token token) {
return params_base.special ||
slot.task->params.sampling.preserved_tokens.find(token) != slot.task->params.sampling.preserved_tokens.end();
@@ -3540,7 +3556,10 @@ private:
}
}
// TODO @ngxson : bring back slot.i_batch check when sub-batching is implemented
if (is_inside_view(slot.i_batch)) {
// the required token not in this sub-batch, skip
continue; // continue loop of slots
}
if (slot.state == SLOT_STATE_DONE_PROMPT) {
if (slot.task->type == SERVER_TASK_TYPE_EMBEDDING) {
@@ -3574,7 +3593,8 @@ private:
continue; // sample using speculative decoding
}
const int tok_idx = slot.i_batch;
// shifted according to the current sub-batch
const int tok_idx = slot.i_batch - off;
llama_token id = common_sampler_sample(slot.smpl.get(), slot.ctx_tgt, tok_idx);