Public/koboldcpp
1
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-09-19 17:25:07 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'upstream' into concedo_experimental

Browse Source 
# Conflicts:
#	ggml/src/ggml-opencl/ggml-opencl.cpp
#	tests/test-backend-ops.cpp
#	tests/test-quantize-fns.cpp
#	tools/server/README.md
#	tools/ui/src/lib/constants/settings-registry.ts
This commit is contained in:
Concedo
2026-07-11 09:20:40 +08:00
parent fc6e197fdd bbebeec4a8
commit 0f0245161e
103 changed files with 2633 additions and 1170 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/server/README-dev.md
+6 -4
View File
@@ -57,7 +57,7 @@ The core architecture consists of the following components:
- `server_tokens`: Unified representation of token sequences (supports both text and multimodal tokens); used by `server_task` and `server_slot`.
- `server_prompt_checkpoint`: For recurrent (e.g., RWKV) and SWA models, stores snapshots of KV cache state. Enables reuse when subsequent requests share the same prompt prefix, saving redundant computation.
- `server_models`: Standalone component for managing multiple backend instances (used in router mode). It is completely independent of `server_context`.
- `stream_session_manager`: Process wide owner of resumable SSE stream sessions (`g_stream_sessions`), keyed by conversation id. Backs the replay buffer that lets a client reattach to a generation after an HTTP disconnect. See the "Resumable streaming" section below.
- `stream_session_manager`: process wide owner of resumable SSE stream sessions, keyed by conversation id. A file-static singleton inside `server-stream.cpp`, driven through `server_stream_session_manager_start/stop`. Backs the replay buffer that lets a client reattach to a generation after an HTTP disconnect. See the "Resumable streaming" section below.
```mermaid
graph TD
@@ -127,10 +127,12 @@ It is opt in via the `X-Conversation-Id` header on `POST /v1/chat/completions`.
The feature lives entirely in `server-stream.{h,cpp}` and rests on three types:
- `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` stops the producer. One conv maps to at most one live session.
- `stream_session_manager` (`g_stream_sessions`): owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL.
- `stream_session_manager`: a file-static singleton (`g_stream_sessions`) inside `server-stream.cpp`, owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL. Exposed to main only through `server_stream_session_manager_start/stop`.
- `stream_pipe_producer` / `stream_pipe_consumer`: the write and read ends. The producer owns the session lifetime and finalizes it on destruction; the consumer is read only and never finalizes, so a reader detaching cannot kill a running generation.
Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker.
The implementation is hidden in `server-stream.cpp` (pimpl). The header exposes only the route handler factories, `server_stream_session_attach_pipe`, `server_stream_aware_should_stop`, `server_stream_conv_id_from_headers` and the GC lifecycle; the session, manager and consumer types stay in the `.cpp`.
Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `server_stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker.
Lifetime safety: the producer pipe holds a shared `alive` flag also captured by the session cancel hook. `~server_res_generator` calls `cleanup()` to clear that hook while the reader is still alive, so a `cancel` arriving during teardown can never call `stop()` on a freed response. This ordering is the most fragile part of the feature: finalizing or destroying the producer before `cleanup()` runs reintroduces a use after free.
@@ -144,7 +146,7 @@ Routes:
Router mode binds the same paths to proxy handlers. A `conv_id -> child` map (`conv_models`), populated when a POST is routed, resolves the owning child in one lookup with no polling. The lookup groups ids per child; GET and DELETE proxy straight to the owner. This loopback REST hop is expected to move to a websocket IPC later, swapping only the transport.
Lifecycle: `g_stream_sessions.start_gc()` runs in main after common init, `stop_gc()` runs first in `clean_up()` and finalizes every live session so no reader hangs. Reader blocking and the post drop drain both run on httplib worker threads, which block on a condvar rather than spin.
Lifecycle: `server_stream_session_manager_start()` runs in main after common init, `server_stream_session_manager_stop()` runs first in `clean_up()` and finalizes every live session so no reader hangs. Reader blocking and the post drop drain both run on httplib worker threads, which block on a condvar rather than spin.
| Constant | Value | Role |
| --- | --- | --- |
tools/server/server-context.cpp
+56 -111
View File
@@ -897,8 +897,10 @@ private:
server_batch batch;
llama_model_ptr model_dft;
llama_context_ptr ctx_dft;
llama_model * model_dft = nullptr;
llama_context * ctx_dft = nullptr;
common_speculative_init_result_ptr spec_init;
common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO;
common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO;
@@ -939,8 +941,10 @@ private:
void destroy() {
spec.reset();
ctx_dft.reset();
model_dft.reset();
spec_init.reset();
ctx_dft = nullptr;
model_dft = nullptr;
llama_init.reset();
@@ -1084,30 +1088,15 @@ private:
// optionally reserve VRAM for the draft / MTP context before fitting the target model
if (params_base.fit_params) {
if (has_spec) {
common_params params_dft = params_base;
bool measure_model_bytes = true;
// MTP draft context lives on the target model, only context+compute are new
bool measure_model_bytes = has_draft;
if (has_draft) {
const auto & params_spec = params_base.speculative.draft;
params_dft.devices = params_spec.devices;
params_dft.model = params_spec.mparams;
params_dft.n_gpu_layers = params_spec.n_gpu_layers;
params_dft.cache_type_k = params_spec.cache_type_k;
params_dft.cache_type_v = params_spec.cache_type_v;
params_dft.tensor_buft_overrides = params_spec.tensor_buft_overrides;
} else {
// MTP draft context lives on the target model, only context+compute are new
measure_model_bytes = false;
}
params_dft.n_outputs_max = params_base.n_parallel;
common_params params_dft = common_base_params_to_speculative(params_base);
auto mparams_dft = common_model_params_to_llama(params_dft);
auto cparams_dft = common_context_params_to_llama(params_dft);
if (spec_mtp) {
cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
cparams_dft.type_k = params_base.speculative.draft.cache_type_k;
cparams_dft.type_v = params_base.speculative.draft.cache_type_v;
}
cparams_dft.n_rs_seq = 0;
@@ -1175,82 +1164,36 @@ private:
add_bos_token = llama_vocab_get_add_bos(vocab);
if (has_draft) {
// TODO speculative: move to common/speculative.cpp?
const auto & params_spec = params_base.speculative.draft;
SRV_TRC("loading draft model '%s'\n", params_spec.mparams.path.c_str());
auto params_dft = params_base;
params_dft.devices = params_spec.devices;
params_dft.model = params_spec.mparams;
params_dft.n_gpu_layers = params_spec.n_gpu_layers;
params_dft.cache_type_k = params_spec.cache_type_k;
params_dft.cache_type_v = params_spec.cache_type_v;
if (params_spec.cpuparams.n_threads > 0) {
params_dft.cpuparams.n_threads = params_spec.cpuparams.n_threads;
params_dft.cpuparams_batch.n_threads = params_spec.cpuparams_batch.n_threads;
}
params_dft.tensor_buft_overrides = params_spec.tensor_buft_overrides;
auto mparams_dft = common_model_params_to_llama(params_dft);
// progress callback
mparams_dft.progress_callback = load_progress_callback;
mparams_dft.progress_callback_user_data = &load_progress_spec;
model_dft.reset(llama_model_load_from_file(params_dft.model.path.c_str(), mparams_dft));
if (model_dft == nullptr) {
SRV_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str());
return false;
}
auto cparams = common_context_params_to_llama(params_dft);
if (spec_mtp) {
cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
}
// note: for small models maybe we can set this to the maximum possible draft from all speculative types
// the extra memory for small models is likely negligible?
cparams.n_rs_seq = 0;
cparams.ctx_other = ctx_tgt;
ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams));
if (ctx_dft == nullptr) {
SRV_ERR("%s", "failed to create draft context\n");
return false;
}
params_base.speculative.draft.ctx_tgt = ctx_tgt;
params_base.speculative.draft.ctx_dft = ctx_dft.get();
} else if (spec_mtp) {
// no new model load, so we simply report 0.0 and 1.0 progress
if (has_spec) {
// spec_mtp doesn't use load a model internally, so we report 0.0 and 1.0 manually
load_progress_callback(0.0f, &load_progress_spec);
load_progress_spec.t_last_load_progress_ms = 0; // reset so internal cbs aren't delayed
SRV_TRC("creating MTP draft context against the target model '%s'\n",
params_base.model.path.c_str());
{
common_params params_dft = common_base_params_to_speculative(params_base);
auto cparams_mtp = common_context_params_to_llama(params_base);
cparams_mtp.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
cparams_mtp.type_k = params_base.speculative.draft.cache_type_k;
cparams_mtp.type_v = params_base.speculative.draft.cache_type_v;
cparams_mtp.n_rs_seq = 0;
cparams_mtp.n_outputs_max = params_base.n_parallel;
cparams_mtp.ctx_other = ctx_tgt;
// progress callback
params_dft.load_progress_callback = load_progress_callback;
params_dft.load_progress_callback_user_data = &load_progress_spec;
ctx_dft.reset(llama_init_from_model(model_tgt, cparams_mtp));
if (ctx_dft == nullptr) {
SRV_ERR("%s", "failed to create MTP context\n");
return false;
spec_init = common_speculative_init_from_params(params_dft, model_tgt, ctx_tgt);
model_dft = spec_init->model();
ctx_dft = spec_init->context();
if (has_draft && model_dft == nullptr) {
SRV_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str());
return false;
}
if (ctx_dft == nullptr) {
SRV_ERR("%s", "failed to create MTP context\n");
return false;
}
params_base.speculative.draft.ctx_tgt = ctx_tgt;
params_base.speculative.draft.ctx_dft = ctx_dft;
}
params_base.speculative.draft.ctx_tgt = ctx_tgt;
params_base.speculative.draft.ctx_dft = ctx_dft.get();
load_progress_callback(1.0f, &load_progress_spec);
}
@@ -1343,13 +1286,15 @@ private:
}
if (ctx_dft) {
ctx_dft_seq_rm_type = common_context_can_seq_rm(ctx_dft.get());
ctx_dft_seq_rm_type = common_context_can_seq_rm(ctx_dft);
}
if (spec) {
SRV_TRC("%s", "speculative decoding context initialized\n");
} else {
ctx_dft.reset();
spec_init.reset();
ctx_dft = nullptr;
model_dft = nullptr;
}
for (int i = 0; i < params_base.n_parallel; i++) {
@@ -1357,7 +1302,7 @@ private:
slot.id = i;
slot.ctx_tgt = ctx_tgt;
slot.ctx_dft = ctx_dft.get();
slot.ctx_dft = ctx_dft;
slot.spec = spec.get();
slot.n_ctx = n_ctx_slot;
@@ -2362,8 +2307,8 @@ private:
// this is not true for SWA models: https://github.com/ggml-org/llama.cpp/pull/24411#issuecomment-4677983225
cur.update_pos(slot.prompt.n_tokens() - n_tokens_cur, pos_min, pos_max);
cur.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
cur.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
cur.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
cur.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
// stash the draft's speculative state with the checkpoint
common_speculative_get_state(spec.get(), slot.id, cur.data_spec);
@@ -2899,8 +2844,8 @@ private:
common_context_seq_add(ctx_tgt, slot.id, n_keep + n_discard, slot.prompt.n_tokens(), -n_discard);
if (ctx_dft) {
common_context_seq_rm (ctx_dft.get(), slot.id, n_keep , n_keep + n_discard);
common_context_seq_add(ctx_dft.get(), slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard);
common_context_seq_rm (ctx_dft, slot.id, n_keep , n_keep + n_discard);
common_context_seq_add(ctx_dft, slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard);
}
// add generated tokens to cache
@@ -2972,7 +2917,7 @@ private:
llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id));
if (use_ckpt_dft) {
slot.spec_ckpt.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
slot.spec_ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}
slot.spec_prompt = slot.prompt.tokens.get_text_tokens();
@@ -3009,10 +2954,10 @@ private:
if (ctx_dft) {
if (use_ckpt_dft) {
ckpt.load_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}
common_context_seq_rm(ctx_dft.get(), slot.id, ckpt.pos_max + 1, -1);
common_context_seq_rm(ctx_dft, slot.id, ckpt.pos_max + 1, -1);
}
if (!draft.empty()) {
@@ -3021,7 +2966,7 @@ private:
(ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_tgt));
const bool use_ckpt_dft =
(ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_dft.get()));
(ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_dft));
if (use_ckpt_tgt) {
//const int64_t t_start = ggml_time_us();
@@ -3038,7 +2983,7 @@ private:
}
if (use_ckpt_dft) {
ckpt.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}
}
});
@@ -3219,8 +3164,8 @@ private:
common_context_seq_add(ctx_tgt, slot.id, head_c, head_c + n_match, kv_shift);
if (ctx_dft) {
common_context_seq_rm (ctx_dft.get(), slot.id, head_p, head_c);
common_context_seq_add(ctx_dft.get(), slot.id, head_c, head_c + n_match, kv_shift);
common_context_seq_rm (ctx_dft, slot.id, head_p, head_c);
common_context_seq_add(ctx_dft, slot.id, head_c, head_c + n_match, kv_shift);
}
for (size_t i = 0; i < n_match; i++) {
@@ -3320,8 +3265,8 @@ private:
if (!do_reset) {
// restore the context checkpoint
it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
it->load_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
// restore the draft's speculative state
common_speculative_set_state(spec.get(), slot.id, it->data_spec);
@@ -3395,7 +3340,7 @@ private:
common_context_seq_rm(ctx_tgt, slot.id, p0, -1);
if (ctx_dft) {
common_context_seq_rm(ctx_dft.get(), slot.id, p0, -1);
common_context_seq_rm(ctx_dft, slot.id, p0, -1);
}
// If using an alora, there may be uncached tokens that come
@@ -4243,7 +4188,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
}
};
auto effective_should_stop = stream_aware_should_stop(res_this, req.should_stop);
auto effective_should_stop = server_stream_aware_should_stop(res_this, req.should_stop);
try {
if (effective_should_stop()) {
@@ -4339,7 +4284,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
// attach a producer pipe to the response when X-Conversation-Id is present.
// the pipe mirrors SSE chunks into the ring buffer and wires up the cancel hook.
stream_session_attach_pipe(*res, req.headers);
server_stream_session_attach_pipe(*res, req.headers);
return res;
}
tools/server/server-models.cpp
+2 -2
View File
@@ -1681,7 +1681,7 @@ void server_models_routes::init_routes() {
}
// remember which child serves this conversation so the stream routes can route straight
// to it without polling, keyed on the exact conv id from the header
std::string conv_id = stream_conv_id_from_headers(req.headers);
std::string conv_id = server_stream_conv_id_from_headers(req.headers);
if (!conv_id.empty()) {
models.conv_models.remember(conv_id, name);
}
@@ -1896,7 +1896,7 @@ void server_models_routes::init_routes() {
if (!from.empty()) {
child_path += "?from=" + from;
}
SRV_INF("proxying stream resume to model %s on port %d, path=%s\n",
SRV_TRC("proxying stream resume to model %s on port %d, path=%s\n",
owner->name.c_str(), owner->port, child_path.c_str());
auto proxy = std::make_unique<server_http_proxy>(
"GET",
tools/server/server-stream.cpp
+147 -41
View File
@@ -6,6 +6,12 @@
#include <chrono>
#include <memory>
#include <utility>
#include <shared_mutex>
enum class stream_read_status {
OK,
OFFSET_LOST,
};
namespace {
constexpr int64_t STREAM_SESSION_TTL_SECONDS = 300;
@@ -13,7 +19,6 @@ constexpr size_t STREAM_SESSION_MAX_BYTES = 4 * 1024 * 1024;
constexpr int64_t STREAM_SESSION_GC_INTERVAL_SECONDS = 60;
constexpr int64_t STREAM_READ_WAKE_INTERVAL_MS = 200;
// returns unix time in seconds
int64_t now_seconds() {
return std::chrono::duration_cast<std::chrono::seconds>(
std::chrono::system_clock::now().time_since_epoch()
@@ -21,6 +26,91 @@ int64_t now_seconds() {
}
}
// owns all live sessions keyed by conversation_id, one conv = at most one live session.
// a periodic GC evicts expired ones
class stream_session_manager {
public:
stream_session_manager();
~stream_session_manager();
stream_session_manager(const stream_session_manager &) = delete;
stream_session_manager & operator=(const stream_session_manager &) = delete;
// install a new session, evicting and cancelling any previous one. conversation_id must be non empty
stream_session_ptr create_or_replace(const std::string & conversation_id);
stream_session_ptr get(const std::string & conversation_id);
std::vector<stream_session_ptr> list_all() const;
void evict(const std::string & conversation_id);
void evict_and_cancel(const std::string & conversation_id);
void start_gc();
void stop_gc();
private:
void gc_loop();
mutable std::shared_mutex map_mu;
std::unordered_map<std::string, stream_session_ptr> sessions; // key: conversation_id
std::thread gc_thread;
bool running;
std::mutex gc_wake_mu;
std::condition_variable gc_wake_cv;
};
// process wide manager, lifecycle controlled by llama-server main() via start_gc/stop_gc
static stream_session_manager g_stream_sessions;
void server_stream_session_manager_start() {
g_stream_sessions.start_gc();
}
void server_stream_session_manager_stop() {
g_stream_sessions.stop_gc();
}
struct stream_session {
std::string conversation_id;
int64_t started_ts; // unix seconds at construction
stream_session(std::string conversation_id_, size_t max_bytes_);
stream_session(const stream_session &) = delete;
stream_session & operator=(const stream_session &) = delete;
bool append(const char * data, size_t len);
void finalize();
// drain from offset into sink, blocking for more bytes or finalize. OFFSET_LOST if offset
// fell below the dropped prefix
stream_read_status read_from(size_t offset,
const std::function<bool(const char *, size_t)> & sink,
const std::function<bool()> & should_stop);
bool is_done() const;
bool is_cancelled() const;
size_t total_size() const; // bytes that ever entered the session
size_t dropped_prefix() const; // bytes evicted from the front due to cap
int64_t completed_at() const; // 0 while alive, unix seconds after finalize
void set_stop_producer(std::function<void()> fn);
void cancel();
private:
mutable std::mutex mu;
std::condition_variable cv;
std::vector<char> buffer;
size_t prefix_dropped;
size_t cap_bytes;
bool done;
std::atomic<bool> cancelled; // polled lock-free by the should_stop closure, no mu
int64_t completed_ts;
std::function<void()> stop_producer;
};
stream_session::stream_session(std::string conversation_id_, size_t max_bytes_)
: conversation_id(std::move(conversation_id_))
, started_ts(now_seconds())
@@ -38,7 +128,7 @@ bool stream_session::append(const char * data, size_t len) {
}
{
std::lock_guard<std::mutex> lock(mu);
if (done.load(std::memory_order_relaxed)) {
if (done) {
return false;
}
if (len >= cap_bytes) {
@@ -62,11 +152,14 @@ bool stream_session::append(const char * data, size_t len) {
}
void stream_session::finalize() {
bool was_done = done.exchange(true, std::memory_order_acq_rel);
if (was_done) {
return;
{
std::lock_guard<std::mutex> lock(mu);
if (done) {
return;
}
done = true;
completed_ts = now_seconds();
}
completed_ts.store(now_seconds(), std::memory_order_release);
cv.notify_all();
}
@@ -96,7 +189,7 @@ stream_read_status stream_session::read_from(size_t offset,
lock.lock();
continue;
}
if (done.load(std::memory_order_acquire)) {
if (done) {
return stream_read_status::OK;
}
// wait for new bytes, finalize, or a periodic wake to re check should_stop
@@ -105,7 +198,8 @@ stream_read_status stream_session::read_from(size_t offset,
}
bool stream_session::is_done() const {
return done.load(std::memory_order_acquire);
std::lock_guard<std::mutex> lock(mu);
return done;
}
size_t stream_session::total_size() const {
@@ -119,7 +213,8 @@ size_t stream_session::dropped_prefix() const {
}
int64_t stream_session::completed_at() const {
return completed_ts.load(std::memory_order_acquire);
std::lock_guard<std::mutex> lock(mu);
return completed_ts;
}
void stream_session::set_stop_producer(std::function<void()> fn) {
@@ -128,7 +223,7 @@ void stream_session::set_stop_producer(std::function<void()> fn) {
}
void stream_session::cancel() {
// flip cancelled first so the producer-side stream_aware_should_stop can break out of the
// flip cancelled first so the producer-side server_stream_aware_should_stop can break out of the
// recv() wait even if remove_waiting_task_ids does not notify the condvar (the cancel task
// posted by rd.stop() will eventually notify, but we do not want to depend on that timing)
cancelled.store(true, std::memory_order_release);
@@ -237,18 +332,24 @@ void stream_session_manager::evict_and_cancel(const std::string & conversation_i
}
void stream_session_manager::start_gc() {
if (running.exchange(true)) {
return;
{
std::lock_guard<std::mutex> lock(gc_wake_mu);
if (running) {
return;
}
running = true;
}
gc_thread = std::thread([this] { gc_loop(); });
}
void stream_session_manager::stop_gc() {
bool was_running = running.exchange(false);
bool was_running;
{
std::lock_guard<std::mutex> lock(gc_wake_mu);
was_running = running;
running = false;
}
if (was_running) {
{
std::lock_guard<std::mutex> lock(gc_wake_mu);
}
gc_wake_cv.notify_all();
if (gc_thread.joinable()) {
gc_thread.join();
@@ -270,15 +371,15 @@ void stream_session_manager::stop_gc() {
}
void stream_session_manager::gc_loop() {
while (running.load(std::memory_order_acquire)) {
while (true) {
{
std::unique_lock<std::mutex> lock(gc_wake_mu);
gc_wake_cv.wait_for(lock,
std::chrono::seconds(STREAM_SESSION_GC_INTERVAL_SECONDS),
[this] { return !running.load(std::memory_order_acquire); });
}
if (!running.load(std::memory_order_acquire)) {
return;
[this] { return !running; });
if (!running) {
return;
}
}
int64_t cutoff = now_seconds() - STREAM_SESSION_TTL_SECONDS;
std::vector<stream_session_ptr> to_drop;
@@ -301,10 +402,19 @@ void stream_session_manager::gc_loop() {
}
}
// process wide manager, lifecycle controlled by llama-server main() via start_gc/stop_gc
stream_session_manager g_stream_sessions;
// stream_pipe
// stream_pipe ---------------------------------------------------------------------------------
// consumer end: read-only replay of the ring buffer, the destructor does not finalize the session
struct stream_pipe_consumer : stream_pipe {
stream_read_status read(size_t & offset,
const std::function<bool(const char *, size_t)> & sink,
const std::function<bool()> & should_stop);
static std::shared_ptr<stream_pipe_consumer> create(stream_session_ptr session);
private:
explicit stream_pipe_consumer(stream_session_ptr session);
};
stream_pipe::stream_pipe(stream_session_ptr session)
: session_(std::move(session)) {
@@ -408,12 +518,10 @@ static server_http_res_ptr make_error_response(int status, const std::string & m
return res;
}
server_http_context::handler_t make_stream_get_handler() {
server_http_context::handler_t server_stream_make_get_handler() {
return [](const server_http_req & req) -> server_http_res_ptr {
// GET /v1/stream/<conv_id>?from=N replays the SSE bytes already buffered for the
// session, blocks for more bytes when the session is still running, returns when
// the session is finalized. the body is streamed back as text/event-stream so the
// browser EventSource can attach to it like a fresh request
// GET /v1/stream/<conv_id>?from=N replays buffered SSE bytes then blocks for live
// bytes until the session finalizes, streamed as text/event-stream for EventSource
std::string conv_id = req.get_param("conv_id");
if (conv_id.empty()) {
return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST);
@@ -459,11 +567,10 @@ server_http_context::handler_t make_stream_get_handler() {
};
}
server_http_context::handler_t make_streams_lookup_handler() {
server_http_context::handler_t server_stream_make_lookup_handler() {
return [](const server_http_req & req) -> server_http_res_ptr {
// POST /v1/streams/lookup with body {"conversation_ids": ["X", "Y", ...]} returns the
// matching sessions, only for ids the caller already knows. each id matches the exact key
// and any "<id>::<model>" variant, so one lookup covers every per model session for a conv
// POST /v1/streams/lookup returns the matching sessions, only for ids the caller already
// knows. each id matches the exact key and any "<id>::<model>" per model variant
std::vector<std::string> requested;
try {
json body = json::parse(req.body);
@@ -518,11 +625,10 @@ server_http_context::handler_t make_streams_lookup_handler() {
};
}
server_http_context::handler_t make_stream_delete_handler() {
server_http_context::handler_t server_stream_make_delete_handler() {
return [](const server_http_req & req) -> server_http_res_ptr {
// DELETE /v1/stream/<conv_id> is the explicit user Stop, cancels the producer hook
// wired by handle_completions_impl and evicts the buffer. idempotent, a session that
// already finalized or was never created returns 204 either way
// DELETE /v1/stream/<conv_id> is the explicit user Stop, cancels the producer and evicts
// the buffer. idempotent, returns 204 even if the session was already gone
std::string conv_id = req.get_param("conv_id");
if (conv_id.empty()) {
return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST);
@@ -536,7 +642,7 @@ server_http_context::handler_t make_stream_delete_handler() {
};
}
std::string stream_conv_id_from_headers(const std::map<std::string, std::string> & headers) {
std::string server_stream_conv_id_from_headers(const std::map<std::string, std::string> & headers) {
// case-insensitive scan for x-conversation-id
static constexpr char target[] = "x-conversation-id";
static constexpr size_t target_len = sizeof(target) - 1;
@@ -555,8 +661,8 @@ std::string stream_conv_id_from_headers(const std::map<std::string, std::string>
return std::string();
}
void stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers) {
std::string conversation_id = stream_conv_id_from_headers(headers);
void server_stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers) {
std::string conversation_id = server_stream_conv_id_from_headers(headers);
SRV_TRC("conv_id=%s (empty=%d)\n", conversation_id.c_str(), conversation_id.empty() ? 1 : 0);
if (conversation_id.empty()) {
return;
@@ -565,7 +671,7 @@ void stream_session_attach_pipe(server_http_res & res, const std::map<std::strin
res.spipe = stream_pipe_producer::create(session, res);
}
std::function<bool()> stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback) {
std::function<bool()> server_stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback) {
return [res, fallback = std::move(fallback)]() -> bool {
if (res->spipe) {
return res->spipe->is_cancelled();
tools/server/server-stream.h
+15 -136
View File
@@ -3,81 +3,23 @@
#include "server-http.h"
#include <atomic>
#include <condition_variable>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <shared_mutex>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>
enum class stream_read_status {
OK,
OFFSET_LOST,
};
// streaming buffer for one generation, survives HTTP disconnect. the producer appends SSE bytes,
// readers drain from any offset via read_from. keyed by conversation_id, one conv = one live session
// streaming buffer for one generation, survives HTTP disconnect. the producer appends raw SSE
// bytes, readers drain from any offset via read_from and block until more bytes or finalize.
// keyed by conversation_id: one conv = at most one live session
struct stream_session {
std::string conversation_id;
int64_t started_ts; // unix seconds at construction, used by /v1/streams listing
stream_session(std::string conversation_id_, size_t max_bytes_);
stream_session(const stream_session &) = delete;
stream_session & operator=(const stream_session &) = delete;
// append raw bytes, drops from the front if the cap is reached.
// returns false if the session is already finalized
bool append(const char * data, size_t len);
// mark the session as complete, wakes all pending readers
void finalize();
// drain bytes from offset, calling sink for each chunk. blocks until more
// bytes arrive or finalize is called. returns OK on clean exit, OFFSET_LOST
// if offset falls below the dropped prefix
stream_read_status read_from(size_t offset,
const std::function<bool(const char *, size_t)> & sink,
const std::function<bool()> & should_stop);
bool is_done() const;
bool is_cancelled() const;
size_t total_size() const; // bytes that ever entered the session
size_t dropped_prefix() const; // bytes evicted from the front due to cap
int64_t completed_at() const; // 0 while alive, unix seconds after finalize
// attach the producer stop hook used to cancel its reader, pass an empty function to detach
void set_stop_producer(std::function<void()> fn);
// signal the producer to abort its inference asap via the stop hook, idempotent
void cancel();
private:
mutable std::mutex mu;
std::condition_variable cv;
std::vector<char> buffer;
size_t prefix_dropped;
size_t cap_bytes;
std::atomic<bool> done;
std::atomic<bool> cancelled;
std::atomic<int64_t> completed_ts;
std::function<void()> stop_producer; // protected by mu
};
struct stream_session;
using stream_session_ptr = std::shared_ptr<stream_session>;
// one end of a stream_session pipe. the base holds the session and the shared query, the
// producer and consumer ends derive from it. virtual dtor so each end runs its own teardown:
// base of the producer/consumer pipe ends. virtual dtor so each runs its own teardown:
// the producer finalizes the session, the consumer leaves it untouched
struct stream_pipe {
virtual ~stream_pipe() = default;
// true if the session was cancelled (e.g. via DELETE /v1/stream/<conv_id>)
bool is_cancelled() const;
protected:
@@ -95,7 +37,6 @@ protected:
struct stream_pipe_producer : stream_pipe {
~stream_pipe_producer() override;
// append raw bytes to the session's ring buffer, returns false if already finalized
bool write(const char * data, size_t len);
// mark the natural end on the wire so a later close() is a no-op
@@ -121,83 +62,21 @@ private:
server_http_res * res_ = nullptr;
};
// consumer end: read-only replay of the ring buffer, the destructor does not finalize the session
struct stream_pipe_consumer : stream_pipe {
// drain bytes from offset, calling sink for each available chunk. blocks until more data
// arrives or the session finalizes. should_stop is polled, returns OFFSET_LOST if offset
// fell below the dropped prefix
stream_read_status read(size_t & offset,
const std::function<bool(const char *, size_t)> & sink,
const std::function<bool()> & should_stop);
void server_stream_session_manager_start();
void server_stream_session_manager_stop();
static std::shared_ptr<stream_pipe_consumer> create(stream_session_ptr session);
// route handler factories wired under /v1/stream/* by server.cpp
server_http_context::handler_t server_stream_make_get_handler();
server_http_context::handler_t server_stream_make_lookup_handler();
server_http_context::handler_t server_stream_make_delete_handler();
private:
explicit stream_pipe_consumer(stream_session_ptr session);
};
// extract the X-Conversation-Id header value (case-insensitive), empty when absent
std::string server_stream_conv_id_from_headers(const std::map<std::string, std::string> & headers);
// owns all live sessions, runs a periodic GC to evict expired ones.
// the map is keyed by conversation_id, so the invariant "one conv = at most one
// live session" is enforced at the type level
class stream_session_manager {
public:
stream_session_manager();
~stream_session_manager();
stream_session_manager(const stream_session_manager &) = delete;
stream_session_manager & operator=(const stream_session_manager &) = delete;
// install a new session for this conversation, evicting and cancelling any previous one.
// the conversation_id must be non empty, the caller is responsible for that check.
// returns the new session
stream_session_ptr create_or_replace(const std::string & conversation_id);
// lookup, returns null if unknown or already evicted
stream_session_ptr get(const std::string & conversation_id);
// list every live or recently completed session, used by GET /v1/streams without filter
std::vector<stream_session_ptr> list_all() const;
// remove from the map and finalize, wakes any pending readers
void evict(const std::string & conversation_id);
// signal the producer to cancel asap then evict, used by the explicit user Stop path
void evict_and_cancel(const std::string & conversation_id);
void start_gc();
void stop_gc();
private:
void gc_loop();
mutable std::shared_mutex map_mu;
std::unordered_map<std::string, stream_session_ptr> sessions; // key: conversation_id
std::thread gc_thread;
std::atomic<bool> running;
std::mutex gc_wake_mu;
std::condition_variable gc_wake_cv;
};
// process wide manager, linked by both llama-server and llama-cli. llama-server main() drives
// start_gc/stop_gc, llama-cli leaves it idle. the dtor calls stop_gc() unconditionally so exit
// is safe whether or not the GC thread ran
extern stream_session_manager g_stream_sessions;
// route handler factories operating on g_stream_sessions, wired under /v1/stream/* by server.cpp.
// keeps the resumable stream surface confined to server-stream
server_http_context::handler_t make_stream_get_handler();
server_http_context::handler_t make_streams_lookup_handler();
server_http_context::handler_t make_stream_delete_handler();
// extract the X-Conversation-Id header value (case-insensitive), empty when absent. exposed so
// the router can track which child serves a forwarded POST
std::string stream_conv_id_from_headers(const std::map<std::string, std::string> & headers);
// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to
// res. no-op when absent, called from the server_res_generator constructor
void stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers);
// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to res
void server_stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers);
// should_stop closure that ignores peer disconnect when a pipe is attached, so only an explicit
// DELETE stops the producer and generation keeps flowing into the ring buffer. without a pipe it
// delegates to fallback, the legacy non-resumable flow
std::function<bool()> stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback);
std::function<bool()> server_stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback);
tools/server/server-task.cpp
+61 -13
View File
@@ -730,6 +730,10 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() {
}}
});
if (timings.prompt_n >= 0) {
server_sent_events.back().at("data").push_back({"timings", timings.to_json()});
}
return server_sent_events;
}
@@ -1016,6 +1020,7 @@ void server_task_result_cmpl_partial::update(task_result_state & state) {
thinking_block_started = state.thinking_block_started;
text_block_started = state.text_block_started;
oai_resp_created = state.oai_resp_created;
oai_resp_id = state.oai_resp_id;
oai_resp_reasoning_id = state.oai_resp_reasoning_id;
oai_resp_message_id = state.oai_resp_message_id;
@@ -1024,6 +1029,10 @@ void server_task_result_cmpl_partial::update(task_result_state & state) {
// track if the accumulated message has any reasoning content
anthropic_has_reasoning = !state.chat_msg.reasoning_content.empty();
if (res_type == TASK_RESPONSE_TYPE_OAI_RESP && !state.oai_resp_created && (is_progress || n_decoded == 1)) {
state.oai_resp_created = true;
}
// Pre-compute state updates based on diffs (for next chunk)
for (const common_chat_msg_diff & diff : oaicompat_msg_diffs) {
if (!diff.reasoning_content_delta.empty() && !state.thinking_block_started) {
@@ -1181,7 +1190,7 @@ json server_task_result_cmpl_partial::to_json_oaicompat_chat() {
json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
std::vector<json> events;
if (n_decoded == 1) {
if (!oai_resp_created) {
events.push_back(json {
{"event", "response.created"},
{"data", json {
@@ -1204,6 +1213,18 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
}},
}},
});
} else if (is_progress) {
events.push_back(json {
{"event", "response.in_progress"},
{"data", json {
{"type", "response.in_progress"},
{"response", json {
{"id", oai_resp_id},
{"object", "response"},
{"status", "in_progress"},
}},
}},
});
}
for (const common_chat_msg_diff & diff : oaicompat_msg_diffs) {
@@ -1302,6 +1323,17 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
});
}
}
if (!events.empty()) {
json & data = events.back().at("data");
if (timings.prompt_n >= 0) {
data.push_back({"timings", timings.to_json()});
}
if (is_progress) {
data.push_back({"prompt_progress", progress.to_json()});
}
}
return events;
}
@@ -1631,7 +1663,22 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t
}
}
// next, remove any cached prompts that are fully contained in the current prompt
// calculate checkpoints size to see if it will fit with the prompt
size_t checkpoints_size = 0;
for (const auto & ckpt : prompt.checkpoints) {
checkpoints_size += ckpt.size();
}
const size_t state_size_new = state_size_tgt + state_size_dft + checkpoints_size;
// skip over-limit entries to avoid disturbing the cache
if (limit_size > 0 && state_size_new > limit_size) {
SRV_WRN(" - prompt state size %.3f MiB exceeds cache size limit %.3f MiB, skipping\n",
state_size_new / (1024.0 * 1024.0), limit_size / (1024.0 * 1024.0));
return nullptr;
}
// remove any cached prompts that are fully contained in the current prompt
for (auto it = states.begin(); it != states.end();) {
const int len = it->tokens.get_common_prefix(prompt.tokens);
@@ -1644,6 +1691,16 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t
}
}
if (limit_size > 0) {
// make room before allocating the new vectors to avoid breaching the limit
while (!states.empty() && size() + state_size_new > limit_size) {
SRV_WRN(" - making room for prompt cache entry, removing oldest entry (size = %.3f MiB)\n",
states.front().size() / (1024.0 * 1024.0));
states.pop_front();
}
}
std::vector<uint8_t> state_data_tgt;
std::vector<uint8_t> state_data_dft;
@@ -1752,12 +1809,7 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
void server_prompt_cache::update() {
if (limit_size > 0) {
// always keep at least one state, regardless of the limits
while (states.size() > 1 && size() > limit_size) {
if (states.empty()) {
break;
}
while (!states.empty() && size() > limit_size) {
SRV_WRN(" - cache size limit reached, removing oldest entry (size = %.3f MiB)\n", states.front().size() / (1024.0 * 1024.0));
states.pop_front();
@@ -1771,11 +1823,7 @@ void server_prompt_cache::update() {
const size_t limit_tokens_cur = limit_size > 0 ? std::max<size_t>(limit_tokens, limit_size/size_per_token) : limit_tokens;
if (limit_tokens > 0) {
while (states.size() > 1 && n_tokens() > limit_tokens_cur) {
if (states.empty()) {
break;
}
while (!states.empty() && n_tokens() > limit_tokens_cur) {
SRV_WRN(" - cache token limit (%zu, est: %zu) reached, removing oldest entry (size = %.3f MiB)\n",
limit_tokens, limit_tokens_cur, states.front().size() / (1024.0 * 1024.0));
tools/server/server-task.h
+2
View File
@@ -117,6 +117,7 @@ struct task_result_state {
bool text_block_started = false;
// for OpenAI Responses streaming API
bool oai_resp_created = false;
const std::string oai_resp_id;
const std::string oai_resp_reasoning_id;
const std::string oai_resp_message_id;
@@ -440,6 +441,7 @@ struct server_task_result_cmpl_partial : server_task_result {
bool text_block_started = false;
// for OpenAI Responses API
bool oai_resp_created = false;
std::string oai_resp_id;
std::string oai_resp_reasoning_id;
std::string oai_resp_message_id;
tools/server/server.cpp
+8 -8
View File
@@ -85,7 +85,7 @@ int llama_server(int argc, char ** argv) {
// start the stream session manager GC right after common init, before any HTTP route can
// touch it. lifecycle is symmetric, stop_gc() runs in clean_up() before backend free
g_stream_sessions.start_gc();
server_stream_session_manager_start();
if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SERVER)) {
return 1;
@@ -245,8 +245,8 @@ int llama_server(int argc, char ** argv) {
ctx_http.post("/slots/:id_slot", ex_wrapper(routes.post_slots));
// resumable streaming, the conversation_id is the session identity end to end. router and
// child wire different handlers under the same paths: a child binds the local g_stream_sessions
// backed factories, the router binds proxies that resolve the owning child through the
// child wire different handlers under the same paths: a child binds the local session
// factories, the router binds proxies that resolve the owning child through the
// conv_id -> model map
server_http_context::handler_t stream_get_h;
server_http_context::handler_t streams_lookup_h;
@@ -256,9 +256,9 @@ int llama_server(int argc, char ** argv) {
streams_lookup_h = models_routes->router_streams_lookup;
stream_delete_h = models_routes->router_stream_delete;
} else {
stream_get_h = make_stream_get_handler();
streams_lookup_h = make_streams_lookup_handler();
stream_delete_h = make_stream_delete_handler();
stream_get_h = server_stream_make_get_handler();
streams_lookup_h = server_stream_make_lookup_handler();
stream_delete_h = server_stream_make_delete_handler();
}
ctx_http.get ("/v1/stream/:conv_id", ex_wrapper(stream_get_h));
// POST /v1/streams/lookup with body {"conversation_ids": [...]}. you can only ask for ids
@@ -343,7 +343,7 @@ int llama_server(int argc, char ** argv) {
clean_up = [&models_routes]() {
SRV_INF("%s: cleaning up before exit...\n", __func__);
// stop the session GC first, it finalizes live sessions and wakes pending readers
g_stream_sessions.stop_gc();
server_stream_session_manager_stop();
if (models_routes.has_value()) {
models_routes->stopping.store(true); // maybe redundant, but just to be safe
models_routes->models.unload_all();
@@ -371,7 +371,7 @@ int llama_server(int argc, char ** argv) {
clean_up = [&ctx_http, &ctx_server]() {
SRV_INF("%s: cleaning up before exit...\n", __func__);
// stop the session GC first, it finalizes live sessions and wakes pending readers
g_stream_sessions.stop_gc();
server_stream_session_manager_stop();
ctx_http.stop();
ctx_server.terminate();
llama_backend_free();
tools/server/tests/unit/test_compat_oai_responses.py
+41
View File
@@ -71,3 +71,44 @@ def test_responses_stream_with_openai_library():
assert r.response.output[0].id.startswith("msg_")
assert gathered_text == r.response.output_text
assert match_regex("(Suddenly)+", r.response.output_text)
def test_responses_stream_with_llama_telemetry():
global server
server.n_ctx = 256
server.n_batch = 32
server.n_slots = 1
server.start()
saw_progress = False
saw_delta_timings = False
completed = None
res = server.make_stream_request("POST", "/responses", data={
"input": "This is a test" * 10,
"max_output_tokens": 8,
"temperature": 0.8,
"stream": True,
"timings_per_token": True,
"return_progress": True,
})
for data in res:
if "prompt_progress" in data:
assert data["type"] == "response.in_progress"
assert data["prompt_progress"]["total"] > 0
assert data["prompt_progress"]["processed"] >= data["prompt_progress"]["cache"]
saw_progress = True
if "timings" in data:
assert "prompt_per_second" in data["timings"]
assert "predicted_per_second" in data["timings"]
if data["type"] == "response.output_text.delta":
saw_delta_timings = True
if data["type"] == "response.completed":
completed = data
assert saw_progress
assert saw_delta_timings
assert completed is not None
assert "usage" in completed["response"]
assert "timings" in completed