mirror/LocalAI
1
0
Fork 0
mirror of https://github.com/mudler/LocalAI.git synced 2026-07-30 18:09:05 -04:00
Code Issues Packages Projects Releases Wiki Activity

feat(classifier/VAD): support voice control on low power devices (#10804)

Browse Source 
* feat(llama-cpp): route Score through the slot loop

Score previously bypassed the slot loop with a direct llama_decode: a
conflict guard aborted the whole process if scoring raced generation, the
config validator had to reject score alongside chat/completion/embeddings,
and every candidate re-decoded the full shared prompt.

Add SERVER_TASK_TYPE_SCORE to the (patched) upstream server so score tasks
are scheduled like any other slot work: generation and scoring serialize
naturally, the shared prompt is decoded once per call, and the slot's
prompt cache carries the conversation prefix across calls. Context
checkpoints at the score boundary and at the cache-divergence point keep
SWA/hybrid/recurrent models (e.g. LFM2.5) from re-prefilling the whole
prompt per candidate: warm-turn scoring on a 6-option set drops from ~8s
to ~0.5s on a desktop CPU.

The conflict guard and the validation split are removed; declaring score
with generation usecases on one config is now supported and shares the
slot cache.

Assisted-by: Claude:claude-fable-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* feat(realtime): classifier wire types and pipeline config

Wire types and YAML config for realtime classifier mode: sessions carry a
localai_classifier extension (options with canned replies/tool calls,
softmax threshold, normalization, history trimming, fallback modes, and a
deterministic wake-word address gate), mirrored by pipeline.classifier in
the model YAML and surfaced in the config-meta registry. The
localai.classifier.result server event reports the full score distribution
per turn.

Assisted-by: Claude:claude-fable-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* feat(realtime): classifier response flow

Classifier-mode responses: instead of autoregressive generation, each user
turn is prefill-scored against the option list (router.ScoreClassifier
prompt/candidate shapes over the Score primitive) and the winning option's
canned reply and tool call are emitted through the existing response
machinery. Below-threshold turns take the configured fallback (none /
canned reply / generate); empty transcripts and unaddressed turns (wake
word not mentioned) skip scoring entirely. The scoring probe defaults to
the latest user message only — small scorers echo canned replies from
prior turns back as the top option otherwise.

Built for hardware that can afford prompt processing but not decode: with
slot-based Score the option list stays KV-cached across turns, so a turn
costs roughly one forward pass over the new words.

session_update_error events now carry the validation cause instead of a
generic message.

Assisted-by: Claude:claude-fable-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* fix(realtime): bound the VAD tick's scan window and buffer retention

The VAD tick loop re-scanned the entire input buffer every 300ms and only
trimmed it on zero-segment ticks or commits. Audio that keeps producing
segments without a committing pause (steady noise a mic pipeline lets
through, music, continuous speech) grew the buffer toward the 100MB cap
with each tick rescanning all of it — O(n^2), measured at ~3.3ms of silero
per buffered second: past ~90s retained, ticks run back to back and pin
~4 cores until the stream stops.

Silero's recurrent state only carries a few hundred ms of context, so
rescanning old audio buys nothing. Clip the slice handed to the VAD to the
largest silence the commit test can need to measure (server_vad silence
window or the semantic eagerness fallback) plus a warm-up margin, and
rebase the returned segment times so every downstream consumer keeps
whole-buffer coordinates. An open turn whose clipped window is all silence
now commits (the silence outran the window) instead of being discarded as
no-speech. Independently, retain at most 90s of raw buffer, rebasing the
live-feed and EOU cursors on trim — this also bounds the previously
unbounded VAD-error path. Turn boundaries are otherwise unchanged: no
forced commits, no new coordinator states.

pipeline.turn_detection.vad_window_sec can widen the scan window; values
below the automatic floor are ignored. The tick body is extracted into
vadTick so specs can drive turn detection synchronously (same shape as
classifySoundWindow); the babble reproduction that pinned 4 cores now
plateaus under 10% of one core.

Assisted-by: Claude:claude-fable-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* fix(backend): let per-model threads override the global default

ModelOptions overrode a set per-model threads value with the app-level
--threads whenever the latter was non-zero — and WithThreads defaults it
to the physical core count, so it always was. The YAML threads: knob has
been dead config: a tiny VAD model could never opt down from the global
pool size.

SetDefaults already fills an unset per-model value from the app config,
which is the intended precedence; resolve threads through a helper that
honors it (explicit threads: 0 still means unset).

Assisted-by: Claude:claude-fable-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* chore(gallery): single-thread the silero VAD

Silero is a ~2MB recurrent model with no exploitable graph parallelism:
measured per-call latency is identical at 1 and 10 ORT threads, while
every extra pool thread just spin-waits between the realtime loop's
frequent tiny inferences.

Assisted-by: Claude:claude-fable-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* docs(realtime): classifier mode, VAD scan window, threads precedence

Document the realtime classifier mode (options, threshold guidance,
wake-word address gate, empty-transcript handling), the VAD scan window
and 90s buffer retention (pipeline.turn_detection.vad_window_sec), the
per-model threads precedence, and the M3 classifier note in the realtime
state-machine design doc.

Assisted-by: Claude:claude-fable-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* perf(llama-cpp): score all candidates in one batched decode

One scoring call is now a single SERVER_TASK_TYPE_SCORE task: the slot
decodes the shared prefix (prompt + longest common candidate token
prefix) once, then forks one sequence per candidate off it
(metadata-only for the unified KV cache, copy-on-write for recurrent
state) and decodes every candidate's unique tail in one llama_decode.
Previously each candidate was its own task that restored the boundary
checkpoint and re-decoded its full tail sequentially, paying
per-candidate task and decode overhead.

The context reserves SERVER_SCORE_FORK_SEQS extra sequence ids (and
recurrent-state cells) beyond the parallel slots via the new
common_params::n_seq_score_forks. Forking requires the unified KV cache
(already this backend's default) since per-sequence streams would shrink
n_ctx_seq; an explicit kv_unified:false disables forking and Score calls
that need it fail cleanly. Candidates beyond the fork/output budget
decode in successive chunks.

Wire contract and scores are unchanged: per-token logprobs are stitched
from the shared region and the forked tails. Verified bitwise
deterministic call-to-call and independent of candidate order (no
cross-fork leakage via equal-length candidate swap); ranking matches the
per-candidate implementation on the drone battery (winner softmax
0.99996 vs 0.99997), and >16-candidate chunking, prefix-of-another and
empty candidates all pass.

Measured on a desktop CPU: warm /api/score calls 0.52s -> 0.23s; warm
realtime classifier turns 196-303ms. The 9-candidate drone turn decodes
~17 unique tail tokens in one batch instead of nine sequential ~220ms
checkpoint-restore tasks.

Assisted-by: Claude:claude-fable-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* fix(realtime): gate scoring capacity by model usecase

Reserve llama.cpp scoring slots only for models that explicitly declare the score usecase, while allowing score to coexist with chat and completion. Reject incompatible unified-KV settings and classifier activation on models without scoring capacity.

Propagate application defaults when resolving realtime and preload pipeline stages so unset thread counts are resolved consistently without overriding explicit model settings.

Assisted-by: Codex:gpt-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* fix(ci): honor APT mirrors in the prebuilt llama-cpp compile step

The builder-prebuilt path installs gcc-14 with apt directly and ignored
the APT_MIRROR/APT_PORTS_MIRROR build args the from-source path already
honors, so an ubuntu mirror outage broke every arm64 backend build. Pass
the args into the stage and run apt-mirror.sh (already in the build
context via COPY . /LocalAI) before the apt step.

Assisted-by: Claude:claude-fable-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* feat(realtime): classifier argument slots via constrained completion

Hybrid classify-then-complete: a classifier option's canned tool call can
declare typed argument slots (number | enum | string, with defaults and
prompt hints) referenced as "{{name}}" in the arguments template. When
the option wins, the slots are filled by a short grammar-constrained
completion that continues the exact scoring prompt — rendered by the same
cached ScoreClassifier, so the llama.cpp prompt cache is already warm —
with the chosen route JSON re-opened at the first slot field. A GBNF
grammar pins the field skeleton and frees only the values; temperature 0,
a couple dozen tokens at most (~300ms on a desktop CPU for two slots).

Slot declarations and hints ride the option descriptions in the shared
system prompt, informing scoring and the fill alike at no per-turn token
cost. The localai.classifier.result event carries the final arguments and
a fill_latency_ms. On inference failure the slots' defaults apply; a slot
without a default fails the response (or falls through with
fallback.mode: generate). Slot filling requires completion alongside
score in the scoring model's known_usecases.

Verified end-to-end on the Pi drone demo: "fly forward three meters" in
distance mode classifies forward and infers {"distance": 3, "units":
"meters"} in ~310ms, and the drone flies exactly 3 units.

Assisted-by: Claude:claude-fable-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* feat(realtime): splice filled slot values into classifier replies

A classifier option's spoken reply can now reference its tool's argument
slots ("Going forward {{distance}} {{units}}."): the values inferred by
the slot-fill completion — or the recovery defaults — are spliced into
the reply as plain text before it is emitted, so what the assistant says
confirms what it actually inferred. Placeholders without a value stay
literal, and options without slots are untouched.

FillToolArguments now returns the raw slot values alongside the spliced
arguments JSON to make the reply templating possible.

Assisted-by: Claude:claude-fable-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* fix(realtime): harden classifier slot completion

Reserve context for constrained slot filling, size completions from their encoded output, and encode enum grammar literals as valid JSON. Reject empty enum values and cover the failure modes with regression tests.

Assisted-by: Codex:gpt-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* feat(realtime): prewarm the classifier scoring prompt on registration

Swapping a session's classifier option list (a voice-switched command
mode, for instance) made the next turns pay a full re-prefill of the new
option-list prompt — measured 2.4s vs 0.3s warm on a desktop CPU, and
worse: on hybrid-memory models like LFM2.5, whose state cannot be
partially rewound (llama.cpp can only restore checkpoints), *every*
probe change re-prefilled from scratch whenever the last checkpoint
missed the probe boundary, so even same-list turns intermittently cost
full prefills.

Registering an option list (pipeline seed or session.update) now fires a
best-effort background prewarm: two throwaway scores with distinct
probes. The first prefills the new option-list prompt; the second,
diverging exactly where per-turn probe text starts, plants the backend's
rewind point (KV checkpoint) at the stable-prefix boundary that every
real turn reuses. The prewarm hides behind the canned mode-switch reply
— by the time it finishes speaking, the cache is warm. Idempotent per
option set, detached from the registering request's lifetime.

Measured on the drone demo (LFM2.5-1.2B, desktop CPU): first turn after
a mode switch 2374ms -> 340ms; intermittent same-list full prefills
(1.3-2.1s) all -> under 0.5s. For clients that swap lists frequently,
options: [parallel:2] on the scoring model additionally keeps one slot
per list via prefix-similarity routing (+26MB RSS, unified KV).

Assisted-by: Claude:claude-fable-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* perf(llama-cpp): checkpoint scoring at the caller-declared stable prefix

Hybrid-memory models (LFM2.5 shortconv, Qwen3.5 deltanet — where new
small models are headed) cannot rewind their state, so any prompt-cache
reuse that needs a rewind falls back to a full re-prefill. For classifier
scoring that meant every probe change re-processed the whole option-list
prompt: the server's checkpoints were placed reactively (at wherever the
previous task happened to diverge), so a checkpoint past the next
divergence was erased rather than restored — measured as intermittent
2-10s turns on prompts with a 95%+ common prefix.

The classifier now computes the probe-invariant prompt prefix once (the
byte-wise common prefix of two synthetic probe renders) and declares its
length with every Score request; the server maps it to a token boundary
and forces a KV checkpoint exactly there on each score prefill. That
checkpoint sits at or before every future divergence under the same
option list, so it always survives and always restores — repeat scoring
costs probe+candidates regardless of how the probe changes.

Also:
- prewarm reruns on every option-list registration instead of memoizing
  per list: with boundary checkpoints a redundant rewarm costs two
  probe-sized decodes, while skipping one after a slot eviction (three
  lists sharing fewer slots evict in LRU cascades) silently moves a full
  re-prefill onto the user's next turn
- new llama.cpp backend option rs_seq:N exposes bounded recurrent-state
  rollback outside speculative decoding; measured impractical for
  deltanet-scale states (65GB for 64 snapshots on Qwen3.5-4B) but cheap
  insurance for small-state models
- docs: the multi-list recipe (parallel:N + sps:0.5 — the default slot
  similarity threshold funnels distinct lists onto one slot)

Measured on the drone demo (LFM2.5-1.2B scorer, desktop CPU), steady
state: every turn 285-421ms including mode switches, vs 2.4s post-switch
and intermittent 1.3-2.9s re-prefills before.

Assisted-by: Claude:claude-fable-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* fix(realtime): align classifier cache guidance

Document the single-score prewarm behavior and clean the vendored score patch formatting.

Assisted-by: Codex:gpt-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* fix(llama-cpp): guard score task for fork backends

TurboQuant and Bonsai reuse the primary gRPC server against llama.cpp forks that do not carry LocalAI's slot-based Score patches. Compile the Score integration only for the patched primary backend and return UNIMPLEMENTED from fork builds instead of referencing absent task types and common_params fields.

Assisted-by: Codex:gpt-5 [gh]
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* fix(dev): generate gRPC code before commit lint

The coverage phase regenerates ignored protobuf bindings, but lint runs first and can fail against missing or stale output. Generate the pinned bindings before lint so the gate always type-checks the current schema.

Assisted-by: Codex:gpt-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

---------

Signed-off-by: Richard Palethorpe <io@richiejp.com>
This commit is contained in:
Richard Palethorpe
2026-07-29 11:50:22 +01:00
committed by GitHub
parent bc21f832aa
commit 49ef40a187
48 changed files with 4922 additions and 506 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
backend/cpp/llama-cpp/disable-score-task.sh Normal file
View File

@@ -0,0 +1,43 @@
#!/bin/bash
# Mark a copied gRPC server as targeting a llama.cpp fork that does not carry
# LocalAI's slot-based Score patches. The RPC remains present in the shared
# protobuf service, but responds with UNIMPLEMENTED instead of referencing
# server task types and common_params fields absent from those forks.
set -euo pipefail
if [[ $# -ne 1 ]]; then
echo "usage: $0 <grpc-server.cpp>" >&2
exit 2
fi
SRC=$1
if [[ ! -f "$SRC" ]]; then
echo "grpc-server.cpp not found at $SRC" >&2
exit 2
fi
if grep -q '^#define LOCALAI_LLAMA_CPP_NO_SCORE_TASK' "$SRC"; then
echo "==> $SRC already disables the LocalAI score task, skipping"
exit 0
fi
awk '
!done && /^#include/ {
print "#define LOCALAI_LLAMA_CPP_NO_SCORE_TASK 1"
print "// ^ injected by disable-score-task.sh for an unpatched llama.cpp fork"
print ""
done = 1
}
{ print }
END {
if (!done) {
print "disable-score-task.sh: no #include anchor found" > "/dev/stderr"
exit 1
}
}
' "$SRC" > "$SRC.tmp"
mv "$SRC.tmp" "$SRC"
echo "==> LocalAI score task disabled in $SRC"

394
backend/cpp/llama-cpp/grpc-server.cpp
View File

@@ -152,40 +152,6 @@ static std::string base64_encode_bytes(const unsigned char* data, size_t len) {
bool loaded_model; // TODO: add a mutex for this, but happens only once loading the model
// Score bypasses the slot loop (see the comment on Score below) so it
// must not run concurrently with any slot-loop RPC. These counters
// are a defence-in-depth tripwire — ModelConfig.Validate already
// rejects llama-cpp configs that mix score with chat/completion/
// embeddings, so a healthy deployment never trips them. seq_cst is
// load-bearing for the increment-then-check pattern below.
static std::atomic<int> slot_loop_inflight{0};
static std::atomic<int> score_inflight{0};
// Increment-then-check, not check-then-increment: two simultaneous
// racers both observe the other's increment and both abort cleanly.
// Reversed, both could see zero and proceed.
struct conflict_guard {
std::atomic<int>& self;
conflict_guard(const char* rpc, std::atomic<int>& self_, std::atomic<int>& other, const char* other_name)
: self(self_) {
self.fetch_add(1, std::memory_order_seq_cst);
int o = other.load(std::memory_order_seq_cst);
if (o > 0) {
fprintf(stderr,
"FATAL: %s called with %s=%d. The llama-cpp backend cannot "
"service Score and slot-loop RPCs concurrently — Score "
"bypasses the slot loop and races the llama_context. Bind "
"Score-using features to a model dedicated to scoring "
"(known_usecases: [score] with no chat/completion/embeddings).\n",
rpc, other_name, o);
std::abort();
}
}
~conflict_guard() {
self.fetch_sub(1, std::memory_order_seq_cst);
}
};
static std::function<void(int)> shutdown_handler;
static std::atomic_flag is_terminating = ATOMIC_FLAG_INIT;
@@ -732,6 +698,22 @@ static void params_parse(server_context& /*ctx_server*/, const backend::ModelOpt
// If conversion fails, keep default value (0)
}
}
#ifndef LOCALAI_LLAMA_CPP_NO_SCORE_TASK
} else if (!strcmp(optname, "n_rs_seq") || !strcmp(optname, "rs_seq")) {
// Recurrent-state rollback snapshots per sequence. Hybrid models
// (deltanet/conv layers) cannot rewind their state, so without
// snapshots any prompt-cache reuse that needs a rewind — e.g. a
// score task whose probe changed under a stable option-list
// prefix — falls back to a full re-prefill. Costs recurrent-state
// memory x (1 + N) per sequence; unsupported archs clamp to 0.
if (optval != NULL) {
try {
params.n_rs_seq = std::stoi(optval_str);
} catch (const std::exception& e) {
// If conversion fails, keep default value (0)
}
}
#endif
} else if (!strcmp(optname, "slot_prompt_similarity") || !strcmp(optname, "sps")) {
if (optval != NULL) {
try {
@@ -1392,6 +1374,17 @@ static void params_parse(server_context& /*ctx_server*/, const backend::ModelOpt
}
}
#ifndef LOCALAI_LLAMA_CPP_NO_SCORE_TASK
// Score-task suffix forking: reserve seq ids (and recurrent-state cells)
// beyond the slots so one scoring call decodes all candidate tails in a
// single batch (SERVER_TASK_TYPE_SCORE, patches/). Requires the unified
// KV cache — with per-sequence streams the extra ids would shrink every
// sequence's context to n_ctx / n_seq_max. Decided after both option
// passes so an explicit kv_unified:false wins and disables forking.
params.score_enabled = request->enablescore();
params.n_seq_score_forks = params.score_enabled && params.kv_unified ? SERVER_SCORE_FORK_SEQS : 0;
#endif
// Terminate/pad the override vectors only after BOTH the named-option loop
// and the generic passthrough (common_params_parse above) have pushed their
// real entries, so back() is the null sentinel the model loader asserts on.
@@ -1478,6 +1471,16 @@ public:
common_params params;
params_parse(ctx_server, request, params);
#ifndef LOCALAI_LLAMA_CPP_NO_SCORE_TASK
if (params.score_enabled && !params.kv_unified) {
const std::string error_msg =
"Score requires the unified KV cache; remove kv_unified:false or remove score from known_usecases";
result->set_message(error_msg);
result->set_success(false);
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, error_msg);
}
#endif
common_init();
// Ensure debug logs are enabled after common_init() sets up logging
common_log_set_verbosity_thold(params.verbosity);
@@ -1680,7 +1683,6 @@ public:
if (params_base.model.path.empty()) {
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
}
conflict_guard guard("PredictStream", slot_loop_inflight, score_inflight, "score_inflight");
json data = parse_options(true, request, params_base, ctx_server.get_llama_context());
@@ -2249,7 +2251,6 @@ public:
if (params_base.model.path.empty()) {
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
}
conflict_guard guard("Predict", slot_loop_inflight, score_inflight, "score_inflight");
json data = parse_options(true, request, params_base, ctx_server.get_llama_context());
data["stream"] = false;
@@ -2783,7 +2784,6 @@ public:
if (params_base.model.path.empty()) {
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
}
conflict_guard guard("Embedding", slot_loop_inflight, score_inflight, "score_inflight");
json body = parse_options(false, request, params_base, ctx_server.get_llama_context());
body["stream"] = false;
@@ -2893,7 +2893,6 @@ public:
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "\"documents\" must be a non-empty string array");
}
conflict_guard guard("Rerank", slot_loop_inflight, score_inflight, "score_inflight");
// Create and queue the task
auto rd = ctx_server.get_response_reader();
@@ -2970,37 +2969,16 @@ public:
// Score returns the model's joint log-probability of each candidate
// continuation given a shared prompt.
//
// WHY bypass the slot/task queue: upstream server_context exposes
// get_llama_context as "main thread only" and the slot loop's
// update_slots() owns the context whenever a task is in flight.
// No public synchronization primitive is available — so Score is
// unsafe to call concurrently with active generation through this
// backend. In practice routing-classifier calls happen before the
// request is routed to a generation backend, so the model used
// for Score is typically idle. Concurrent Score calls are
// serialised by a local mutex; KV-cache state is isolated behind
// a dedicated sequence ID cleared between candidates.
//
// A patch to server-context.cpp that adds SERVER_TASK_TYPE_SCORE
// and routes scoring through the slot loop would be the correct
// long-term fix; tracked as a follow-up.
//
// Perf TODO (measured: ~450 ms warm for 3 candidates on Arch-
// Router-1.5B Q4_K_M + Intel SYCL): the current loop re-decodes
// `prompt + candidate` from scratch for every candidate, throwing
// away the prompt's KV cache between iterations. A smarter
// version would:
// 1. Decode just the prompt once into score_seq_id.
// 2. Snapshot/cp that sequence (llama_memory_seq_cp) into a
// per-candidate sequence id.
// 3. For each candidate, decode only its tokens onto the copy
// (continuing from the saved prompt state), read logits.
// 4. llama_memory_seq_rm the copy.
// Estimated speedup: 3-candidate calls 450 ms -> ~150-200 ms,
// 6-candidate calls 630 ms -> ~220 ms. Single source-file change,
// no proto / Go-side changes needed. Worth doing once routing is
// wired into the middleware and Score is on the hot path of every
// chat request.
// Scoring runs as a single SERVER_TASK_TYPE_SCORE task through the
// slot loop (added by patches/ on top of upstream server-context), so
// it is safe to interleave with generation on the same process and it
// reuses any KV prefix the slot already holds across turns. The task
// decodes the shared prefix (prompt + longest common candidate token
// prefix) once on the slot's sequence; every candidate's unique tail
// then rides its own forked sequence and all tails are decoded
// together in one batch, so a warm scoring call costs roughly one
// forward pass over the new prompt tokens plus one batched pass over
// the candidate tails.
grpc::Status Score(ServerContext* context, const backend::ScoreRequest* request, backend::ScoreResponse* response) override {
auto auth = checkAuth(context);
if (!auth.ok()) return auth;
@@ -3009,40 +2987,21 @@ public:
if (params_base.model.path.empty()) {
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
}
#ifdef LOCALAI_LLAMA_CPP_NO_SCORE_TASK
(void) request;
(void) response;
return grpc::Status(grpc::StatusCode::UNIMPLEMENTED,
"Score is unavailable in this llama.cpp fork backend");
#else
if (!params_base.score_enabled) {
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION,
"Score was not enabled when the model was loaded; add score to known_usecases");
}
if (request->candidates_size() == 0) {
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "candidates must be non-empty");
}
// Tripwire against the slot loop. Acquired before score_mutex
// so it fires even when this Score is queued behind another.
conflict_guard guard("Score", score_inflight, slot_loop_inflight, "slot_loop_inflight");
// Serialise concurrent Score calls. The slot loop is still
// free to race with us — see the class comment above.
static std::mutex score_mutex;
std::lock_guard<std::mutex> score_lock(score_mutex);
llama_context * lctx = ctx_server.get_llama_context();
if (lctx == nullptr) {
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "llama context unavailable (sleeping?)");
}
const llama_vocab * vocab = ctx_server.impl->vocab;
const int32_t n_vocab = llama_vocab_n_tokens(vocab);
const int32_t n_ctx = llama_n_ctx(lctx);
llama_memory_t mem = llama_get_memory(lctx);
// The KV-cache is sized to seq_to_stream.size() at load
// (typically equal to n_slots, often 1). Sequence IDs must
// be in [0, n_seq_max), so we can't pick a high-value
// "private" ID — we have to share with the slot. We clear
// the cache before AND after each candidate to keep
// scoring isolated from whatever state the slot held, and
// the static mutex above guarantees no other Score call is
// racing in the meantime. The slot loop is still free to
// race (see comment on this method) — Score must not run
// concurrently with generation through this backend.
const llama_seq_id score_seq_id = 0;
llama_memory_seq_rm(mem, score_seq_id, -1, -1);
// Tokenize the shared prompt once with add_special=true so
// BOS is prepended when the model requires it. parse_special
@@ -3051,6 +3010,15 @@ public:
std::vector<llama_token> prompt_tokens = common_tokenize(vocab, prompt, /*add_special=*/true, /*parse_special=*/true);
const int32_t prompt_len = (int32_t) prompt_tokens.size();
// Per candidate: full prompt+candidate token list and the
// divergence point, kept for piece rendering and empty-candidate
// handling after the task comes back.
std::vector<std::vector<llama_token>> cand_tokens(request->candidates_size());
std::vector<int32_t> cand_divergence(request->candidates_size(), 0);
// candidates that actually have tokens to score
std::vector<int32_t> included;
for (int ci = 0; ci < request->candidates_size(); ci++) {
const std::string & candidate_text = request->candidates(ci);
@@ -3067,9 +3035,135 @@ public:
break;
}
}
divergence = std::min<int32_t>(divergence, (int32_t) full_tokens.size());
const int32_t cand_len = (int32_t) full_tokens.size() - divergence;
if (cand_len > 0 && divergence < 1) {
// Need at least one prior token (typically BOS) to
// predict the first candidate token's logit. Tokeniser
// models without BOS + an empty prompt fall in here.
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
"Score: prompt produced no leading tokens; need at least one (e.g. BOS) to predict candidate");
}
if (cand_len > SERVER_SCORE_MAX_CAND_TOKENS) {
// The context reserves logits outputs for at most this many
// candidate tokens per slot (server_n_outputs_max).
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
"Score: candidate " + std::to_string(ci) + " is " + std::to_string(cand_len) +
" tokens; the maximum is " + std::to_string(SERVER_SCORE_MAX_CAND_TOKENS));
}
cand_divergence[ci] = divergence;
cand_tokens[ci] = std::move(full_tokens);
if (cand_len > 0) {
included.push_back(ci);
}
}
auto rd = ctx_server.get_response_reader();
bool posted_task = false;
// Shared prefix bounds, needed again when stitching the results:
// n_shared is the longest common token prefix of the scored
// candidates, n_score_prompt the earliest divergence from the
// bare prompt (scored logprobs start there).
int32_t n_shared = 0;
int32_t n_score_prompt = 0;
if (!included.empty()) {
const auto & first = cand_tokens[included[0]];
// the common prefix of a set is the shortest common prefix
// against any fixed member
n_shared = (int32_t) first.size();
for (int32_t ci : included) {
const auto & ft = cand_tokens[ci];
const int32_t lim = std::min<int32_t>(n_shared, (int32_t) ft.size());
int32_t match = 0;
while (match < lim && ft[match] == first[match]) {
match++;
}
n_shared = match;
}
// below its divergence every candidate equals the prompt
// tokens, so n_score_prompt <= n_shared always holds
n_score_prompt = cand_divergence[included[0]];
for (int32_t ci : included) {
n_score_prompt = std::min(n_score_prompt, cand_divergence[ci]);
}
// Map the caller's stable-prefix byte length onto a token
// index: the last prompt token that ends at or before the
// boundary. A checkpoint forced there survives every future
// probe under the same option list, which is what keeps
// repeat scoring cheap on models that cannot rewind state.
int32_t n_stable_prompt = 0;
if (request->stable_prefix_len() > 0) {
size_t consumed = 0;
for (int32_t ti = 0; ti < n_score_prompt; ti++) {
const size_t piece_len = common_token_to_piece(vocab, prompt_tokens[ti]).size();
// BOS and other zero-length specials consume no prompt bytes
if (consumed + piece_len > (size_t) request->stable_prefix_len()) {
break;
}
consumed += piece_len;
n_stable_prompt = ti + 1;
}
}
server_task task(SERVER_TASK_TYPE_SCORE);
task.id = rd.queue_tasks.get_new_id();
task.index = 0;
task.tokens = server_tokens(llama_tokens(first.begin(), first.begin() + n_shared), false);
task.n_score_prompt = n_score_prompt;
task.n_stable_prompt = n_stable_prompt;
task.score_suffixes.reserve(included.size());
for (int32_t ci : included) {
task.score_suffixes.emplace_back(cand_tokens[ci].begin() + n_shared, cand_tokens[ci].end());
}
std::vector<server_task> tasks;
tasks.push_back(std::move(task));
rd.post_tasks(std::move(tasks));
posted_task = true;
}
// Wait for the shared-prefix and per-candidate logprob vectors.
// Context overflow and decode failures surface here as task errors.
std::vector<float> shared_logprobs;
std::vector<std::vector<float>> cand_logprobs;
if (posted_task) {
auto all_results = rd.wait_for_all([&context]() { return context->IsCancelled(); });
if (all_results.is_terminated) {
return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled by client");
}
if (all_results.error) {
return grpc::Status(grpc::StatusCode::INTERNAL,
all_results.error->to_json().value("message", "Error in receiving score results"));
}
if (all_results.results.size() != 1) {
return grpc::Status(grpc::StatusCode::INTERNAL, "expected a single score result");
}
auto * score_res = dynamic_cast<server_task_result_score*>(all_results.results[0].get());
if (score_res == nullptr) {
return grpc::Status(grpc::StatusCode::INTERNAL, "unexpected result type for score task");
}
shared_logprobs = std::move(score_res->shared_logprobs);
cand_logprobs = std::move(score_res->cand_logprobs);
if (cand_logprobs.size() != included.size()) {
return grpc::Status(grpc::StatusCode::INTERNAL, "score result candidate count mismatch");
}
}
size_t inc = 0; // index into included / cand_logprobs
for (int ci = 0; ci < request->candidates_size(); ci++) {
const int32_t divergence = cand_divergence[ci];
const int32_t cand_len = (int32_t) cand_tokens[ci].size() - divergence;
backend::CandidateScore * cs = response->add_candidates();
cs->set_num_tokens(cand_len);
cs->set_num_tokens(cand_len > 0 ? cand_len : 0);
if (cand_len <= 0) {
cs->set_log_prob(0.0);
if (request->length_normalize()) {
@@ -3077,101 +3171,57 @@ public:
}
continue;
}
if (divergence < 1) {
// Need at least one prior token (typically BOS) to
// predict the first candidate token's logit. Tokeniser
// models without BOS + an empty prompt fall in here.
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
"Score: prompt produced no leading tokens; need at least one (e.g. BOS) to predict candidate");
// Stitch the candidate's scored logprobs back together: the
// stretch inside the shared prefix (identical for every
// candidate) followed by its forked suffix. Suffix entries
// before the candidate's own divergence are prompt tokens
// decoded only as context — not scored.
std::vector<float> lp;
lp.reserve(cand_len);
for (int32_t t = divergence; t < n_shared; t++) {
const int32_t idx = t - n_score_prompt;
if (idx < 0 || idx >= (int32_t) shared_logprobs.size()) {
return grpc::Status(grpc::StatusCode::INTERNAL,
"Score: shared logprob index out of range for candidate " + std::to_string(ci));
}
lp.push_back(shared_logprobs[idx]);
}
if ((int32_t) full_tokens.size() > n_ctx) {
return grpc::Status(grpc::StatusCode::OUT_OF_RANGE,
"Score: prompt+candidate exceeds context size (got " +
std::to_string(full_tokens.size()) + ", n_ctx=" + std::to_string(n_ctx) + ")");
const auto & sfx_lp = cand_logprobs[inc++];
for (int32_t j = std::max(0, divergence - n_shared); j < (int32_t) sfx_lp.size(); j++) {
lp.push_back(sfx_lp[j]);
}
// Build a batch covering the entire prompt+candidate. We
// need logits at (divergence-1) onward — those are the
// predictions for each candidate token.
llama_batch batch = llama_batch_init((int32_t) full_tokens.size(), 0, 1);
for (int32_t i = 0; i < (int32_t) full_tokens.size(); i++) {
batch.token[i] = full_tokens[i];
batch.pos[i] = i;
batch.n_seq_id[i] = 1;
batch.seq_id[i][0] = score_seq_id;
// logits[i] is "do we want the prediction *for the
// next token*, computed from this position?"
// We want predictions for candidate tokens at
// positions divergence .. full_tokens.size()-1, which
// come from logits at positions (divergence-1) ..
// (full_tokens.size()-2).
bool need_logit = (i >= divergence - 1) && (i < (int32_t) full_tokens.size() - 1);
batch.logits[i] = need_logit ? 1 : 0;
}
batch.n_tokens = (int32_t) full_tokens.size();
// Decode the batch. If decode fails (e.g. KV slot
// exhaustion), surface as INTERNAL — the caller will
// typically fall back to a sampling-based classifier.
int decode_err = llama_decode(lctx, batch);
if (decode_err != 0) {
llama_batch_free(batch);
llama_memory_seq_rm(mem, score_seq_id, -1, -1);
if ((int32_t) lp.size() != cand_len) {
return grpc::Status(grpc::StatusCode::INTERNAL,
"llama_decode failed during Score: " + std::to_string(decode_err));
"Score: result for candidate " + std::to_string(ci) + " is missing token logprobs");
}
// Sum log-probabilities of the actual candidate tokens.
double total_log_prob = 0.0;
for (int32_t k = 0; k < cand_len; k++) {
// The k-th candidate token sits at full_tokens index
// (divergence + k). Its predicting logit is at batch
// position (divergence + k - 1).
int32_t logit_pos = divergence + k - 1;
const float * logits = llama_get_logits_ith(lctx, logit_pos);
if (logits == nullptr) {
llama_batch_free(batch);
llama_memory_seq_rm(mem, score_seq_id, -1, -1);
const float token_log_prob = lp[k];
if (std::isnan(token_log_prob)) {
return grpc::Status(grpc::StatusCode::INTERNAL,
"llama_get_logits_ith returned null at position " + std::to_string(logit_pos));
"Score: incomplete result for candidate " + std::to_string(ci) +
" at token " + std::to_string(k));
}
llama_token target_token = full_tokens[divergence + k];
// Compute log_softmax(logits)[target_token] with the
// max-subtraction stability trick.
float max_logit = logits[0];
for (int32_t v = 1; v < n_vocab; v++) {
if (logits[v] > max_logit) max_logit = logits[v];
}
double sum_exp = 0.0;
for (int32_t v = 0; v < n_vocab; v++) {
sum_exp += std::exp((double)(logits[v] - max_logit));
}
double token_log_prob = (double)(logits[target_token] - max_logit) - std::log(sum_exp);
total_log_prob += token_log_prob;
total_log_prob += (double) token_log_prob;
if (request->include_token_logprobs()) {
backend::TokenLogProb * tlp = cs->add_tokens();
std::string piece = common_token_to_piece(lctx, target_token);
tlp->set_token(piece);
tlp->set_token(common_token_to_piece(vocab, cand_tokens[ci][divergence + k]));
tlp->set_log_prob(token_log_prob);
}
}
cs->set_log_prob(total_log_prob);
if (request->length_normalize() && cand_len > 0) {
if (request->length_normalize()) {
cs->set_length_normalized_log_prob(total_log_prob / (double) cand_len);
}
llama_batch_free(batch);
// Drop this candidate's KV-cache contribution so the next
// candidate starts from a clean state. Without this, the
// next decode would conflict at positions 0..N-1 for our
// sequence ID.
llama_memory_seq_rm(mem, score_seq_id, -1, -1);
}
return grpc::Status::OK;
#endif
}
grpc::Status TokenizeString(ServerContext* context, const backend::PredictOptions* request, backend::TokenizationResponse* response) override {
@@ -3182,7 +3232,6 @@ public:
if (params_base.model.path.empty()) {
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
}
conflict_guard guard("TokenizeString", slot_loop_inflight, score_inflight, "score_inflight");
json body = parse_options(false, request, params_base, ctx_server.get_llama_context());
body["stream"] = false;
@@ -3204,7 +3253,6 @@ public:
grpc::Status GetMetrics(ServerContext* /*context*/, const backend::MetricsRequest* /*request*/, backend::MetricsResponse* response) override {
conflict_guard guard("GetMetrics", slot_loop_inflight, score_inflight, "score_inflight");
// request slots data using task queue
auto rd = ctx_server.get_response_reader();

599
backend/cpp/llama-cpp/patches/0001-add-server-task-type-score.patch Normal file
View File

@@ -0,0 +1,599 @@
diff --git a/common/common.cpp b/common/common.cpp
index 8f13217..fc584e1 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1591,8 +1591,10 @@ struct llama_context_params common_context_params_to_llama(const common_params &
auto cparams = llama_context_default_params();
cparams.n_ctx = params.n_ctx;
- cparams.n_seq_max = params.n_parallel;
- cparams.n_rs_seq = params.speculative.need_n_rs_seq();
+ // score-task forks need seq ids (and recurrent-state cells) of their
+ // own beyond the parallel slots
+ cparams.n_seq_max = params.n_parallel + params.n_seq_score_forks;
+ cparams.n_rs_seq = std::max(params.speculative.need_n_rs_seq(), (uint32_t) std::max(0, params.n_rs_seq));
cparams.n_outputs_max = std::max(params.n_outputs_max, 0);
cparams.n_batch = params.n_batch;
cparams.n_ubatch = params.n_ubatch;
diff --git a/common/common.h b/common/common.h
index bffc176..e313bd6 100644
--- a/common/common.h
+++ b/common/common.h
@@ -455,6 +455,9 @@ struct common_params {
int32_t n_keep = 0; // number of tokens to keep from initial prompt
int32_t n_chunks = -1; // max number of chunks to process (-1 = unlimited)
int32_t n_parallel = 1; // number of parallel sequences to decode
+ int32_t n_seq_score_forks = 0; // extra seq ids beyond n_parallel, reserved for server score-task forks
+ int32_t n_rs_seq = 0; // recurrent-state rollback snapshots per seq (hybrid models cannot rewind without them; lets score tasks reuse a cached prompt across probe changes)
+ bool score_enabled = false; // reserve server resources for the Score task type
int32_t n_sequences = 1; // number of sequences to decode
int32_t n_outputs_max = 0; // max outputs in a batch (0 = n_batch)
int32_t grp_attn_n = 1; // group-attention factor
diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt
index 780df32..1d2fe8f 100644
--- a/tools/CMakeLists.txt
+++ b/tools/CMakeLists.txt
@@ -41,3 +41,4 @@ else()
add_subdirectory(fit-params)
add_subdirectory(results)
endif()
+add_subdirectory(grpc-server)
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index 715477e..de5bed8 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -49,7 +49,16 @@ static uint32_t server_n_outputs_max(const common_params & params) {
const uint32_t n_outputs_per_seq = 1 + common_speculative_n_max(&params.speculative);
- const uint64_t n_outputs = (uint64_t) params.n_parallel * n_outputs_per_seq;
+ // score tasks (SERVER_TASK_TYPE_SCORE) output logits for every candidate
+ // token, so reserve room for a bounded candidate tail per parallel slot
+ if (!params.score_enabled) {
+ return std::max<uint32_t>(1, std::min<uint64_t>(n_batch,
+ (uint64_t) params.n_parallel * n_outputs_per_seq));
+ }
+
+ const uint32_t n_outputs_score_seq = 1 + SERVER_SCORE_MAX_CAND_TOKENS;
+
+ const uint64_t n_outputs = (uint64_t) params.n_parallel * std::max(n_outputs_per_seq, n_outputs_score_seq);
return std::max<uint32_t>(1, std::min<uint64_t>(n_batch, n_outputs));
}
@@ -202,6 +211,26 @@ struct server_slot {
std::vector<completion_token_output> generated_token_probs;
+ // SERVER_TASK_TYPE_SCORE: shared-prefix token logprobs harvested
+ // incrementally across batch views (NaN = not yet produced)
+ std::vector<float> score_logprobs;
+
+ // SERVER_TASK_TYPE_SCORE: per-candidate suffix token logprobs; entry
+ // [c][0] comes from the last shared token's logits during prompt
+ // processing, the rest from the forked suffix decode
+ std::vector<std::vector<float>> score_cand_logprobs;
+
+ // SERVER_TASK_TYPE_SCORE: the prompt completed but some candidate has
+ // suffix tokens beyond the first, so a forked decode is still needed
+ bool score_suffix_pending = false;
+
+ // SERVER_TASK_TYPE_SCORE: where the current task's tokens diverged from
+ // the slot's previous cache. When the memory cannot rewind there and a
+ // re-prefill follows, a checkpoint at this position lets the next
+ // scoring call over the same stable prefix (e.g. a classifier's option
+ // list) resume from it instead of re-processing the whole prompt.
+ int32_t score_divergence = -1;
+
bool has_next_token = true;
bool has_new_line = false;
bool truncated = false;
@@ -311,6 +340,10 @@ struct server_slot {
}
generated_tokens.clear();
generated_token_probs.clear();
+ score_logprobs.clear();
+ score_cand_logprobs.clear();
+ score_suffix_pending = false;
+ score_divergence = -1;
json_schema = json();
// clear speculative decoding stats
@@ -2205,6 +2238,229 @@ private:
queue_results.send(std::move(res));
}
+ // log(sum(exp(logits))) with max-subtraction for stability — the
+ // log_softmax denominator shared by every token read from one output
+ static double score_log_denom(const float * logits, int32_t n_vocab) {
+ float max_logit = logits[0];
+ for (int32_t v = 1; v < n_vocab; ++v) {
+ max_logit = std::max(max_logit, logits[v]);
+ }
+ double sum_exp = 0.0;
+ for (int32_t v = 0; v < n_vocab; ++v) {
+ sum_exp += std::exp((double)(logits[v] - max_logit));
+ }
+ return (double) max_logit + std::log(sum_exp);
+ }
+
+ // Harvest logprobs for SCORE tasks from the current batch view: the
+ // shared-prefix scored tokens, and — from the last shared token's
+ // logits — the first suffix token of every candidate. The scored
+ // region can straddle ubatch boundaries for long prompts, so this
+ // accumulates view by view instead of reading everything when the
+ // prompt completes.
+ void collect_score_logprobs(server_slot & slot, const llama_batch & batch) {
+ const int32_t n_prompt = slot.task->n_score_prompt;
+ const int32_t n_total = slot.task->n_tokens();
+ const auto & suffixes = slot.task->score_suffixes;
+
+ const size_t n_shared_scored = (size_t) std::max(0, n_total - n_prompt);
+
+ if (slot.score_logprobs.size() != n_shared_scored) {
+ slot.score_logprobs.assign(n_shared_scored, NAN);
+ }
+ if (slot.score_cand_logprobs.size() != suffixes.size()) {
+ slot.score_cand_logprobs.resize(suffixes.size());
+ for (size_t c = 0; c < suffixes.size(); ++c) {
+ slot.score_cand_logprobs[c].assign(suffixes[c].size(), NAN);
+ }
+ }
+
+ const int32_t n_vocab = llama_vocab_n_tokens(vocab);
+
+ for (int32_t i = 0; i < batch.n_tokens; ++i) {
+ if (!batch.logits[i] || batch.seq_id[i][0] != slot.id) {
+ continue;
+ }
+
+ // the output at position p predicts the task token at index p + 1;
+ // score tasks are text-only, so positions equal token indices
+ const int32_t target = batch.pos[i] + 1;
+ if (target < n_prompt || target > n_total) {
+ continue;
+ }
+
+ const float * logits = llama_get_logits_ith(slot.ctx_tgt, i);
+ if (logits == nullptr) {
+ SLT_ERR(slot, "failed to get logits for score target %d\n", target);
+ continue;
+ }
+
+ const double log_denom = score_log_denom(logits, n_vocab);
+
+ if (target < n_total) {
+ const llama_token tok = slot.task->tokens[target];
+ slot.score_logprobs[target - n_prompt] = (float) ((double) logits[tok] - log_denom);
+ } else {
+ // the last shared token predicts the first suffix token of
+ // every candidate
+ for (size_t c = 0; c < suffixes.size(); ++c) {
+ if (!suffixes[c].empty()) {
+ slot.score_cand_logprobs[c][0] = (float) ((double) logits[suffixes[c][0]] - log_denom);
+ }
+ }
+ }
+ }
+ }
+
+ void send_score(server_slot & slot) {
+ auto res = std::make_unique<server_task_result_score>();
+ res->id = slot.task->id;
+ res->index = slot.task->index;
+ res->shared_logprobs = std::move(slot.score_logprobs);
+ res->cand_logprobs = std::move(slot.score_cand_logprobs);
+
+ slot.score_logprobs.clear();
+ slot.score_cand_logprobs.clear();
+
+ SLT_DBG(slot, "sending score result, n_shared = %zu, n_cand = %zu\n",
+ res->shared_logprobs.size(), res->cand_logprobs.size());
+
+ queue_results.send(std::move(res));
+ }
+
+ // Decode the candidate suffixes of a completed score prompt: fork one
+ // sequence per candidate off the slot's shared prefix (metadata-only
+ // for the unified KV cache, copy-on-write for recurrent state) and
+ // decode all unique suffix tokens in as few llama_decode calls as the
+ // fork/batch/output budgets allow, harvesting a logprob for every
+ // suffix token that predicts a following one.
+ bool decode_score_suffixes(server_slot & slot) {
+ const auto & suffixes = slot.task->score_suffixes;
+
+ auto * mem = llama_get_memory(ctx_tgt);
+
+ // seq ids beyond the slots are reserved for score forks at context
+ // creation (common_params::n_seq_score_forks)
+ const int32_t seq_base = (int32_t) slots.size();
+ const int32_t n_forks_max = std::min<int32_t>(SERVER_SCORE_FORK_SEQS, (int32_t) llama_n_seq_max(ctx_tgt) - seq_base);
+
+ if (n_forks_max < 1) {
+ SLT_ERR(slot, "no fork sequences reserved for score suffixes (n_seq_max = %d, n_slots = %d)\n",
+ (int32_t) llama_n_seq_max(ctx_tgt), seq_base);
+ return false;
+ }
+
+ const int32_t n_batch_max = llama_n_batch(ctx_tgt);
+ const int32_t n_vocab = llama_vocab_n_tokens(vocab);
+ const llama_pos pos0 = slot.prompt.tokens.pos_next();
+
+ std::vector<size_t> pending;
+ for (size_t c = 0; c < suffixes.size(); ++c) {
+ // single-token suffixes were fully scored from the last shared
+ // token's logits during prompt processing
+ if (suffixes[c].size() > 1) {
+ if ((int32_t) suffixes[c].size() > n_batch_max) {
+ SLT_ERR(slot, "score suffix of candidate %zu (%zu tokens) exceeds n_batch (%d)\n",
+ c, suffixes[c].size(), n_batch_max);
+ return false;
+ }
+ pending.push_back(c);
+ }
+ }
+
+ size_t next = 0;
+ while (next < pending.size()) {
+ // greedy-pack candidates into one decode within the fork,
+ // batch and reserved-output budgets
+ std::vector<size_t> chunk;
+ int32_t n_tok = 0;
+ int32_t n_out = 0;
+ while (next < pending.size() && (int32_t) chunk.size() < n_forks_max) {
+ const int32_t m = (int32_t) suffixes[pending[next]].size();
+ if (!chunk.empty() && (n_tok + m > n_batch_max || n_out + m - 1 > SERVER_SCORE_MAX_CAND_TOKENS)) {
+ break;
+ }
+ chunk.push_back(pending[next]);
+ n_tok += m;
+ n_out += m - 1;
+ next++;
+ }
+
+ llama_batch fb = llama_batch_init(n_tok, 0, 1);
+
+ for (size_t k = 0; k < chunk.size(); ++k) {
+ const llama_seq_id seq = seq_base + (llama_seq_id) k;
+ const auto & sfx = suffixes[chunk[k]];
+
+ llama_memory_seq_rm(mem, seq, -1, -1);
+ llama_memory_seq_cp(mem, slot.id, seq, -1, -1);
+
+ for (size_t j = 0; j < sfx.size(); ++j) {
+ common_batch_add(fb, sfx[j], pos0 + (llama_pos) j, { seq }, j + 1 < sfx.size());
+ }
+ }
+
+ const int ret = llama_decode(ctx_tgt, fb);
+
+ if (ret == 0) {
+ int32_t i = 0;
+ for (size_t k = 0; k < chunk.size(); ++k) {
+ const auto & sfx = suffixes[chunk[k]];
+ auto & out = slot.score_cand_logprobs[chunk[k]];
+
+ for (size_t j = 0; j < sfx.size(); ++j, ++i) {
+ if (j + 1 >= sfx.size()) {
+ continue; // last suffix token predicts nothing
+ }
+ const float * logits = llama_get_logits_ith(ctx_tgt, i);
+ if (logits == nullptr) {
+ SLT_ERR(slot, "failed to get logits for suffix token %zu of score candidate %zu\n", j, chunk[k]);
+ continue;
+ }
+ const double log_denom = score_log_denom(logits, n_vocab);
+ out[j + 1] = (float) ((double) logits[sfx[j + 1]] - log_denom);
+ }
+ }
+ }
+
+ for (size_t k = 0; k < chunk.size(); ++k) {
+ llama_memory_seq_rm(mem, seq_base + (llama_seq_id) k, -1, -1);
+ }
+
+ llama_batch_free(fb);
+
+ if (ret != 0) {
+ SLT_ERR(slot, "score suffix decode failed, ret = %d\n", ret);
+ return false;
+ }
+ }
+
+ return true;
+ }
+
+ // score slots whose prompt completed this iteration decode their
+ // candidate suffixes here, after every batch view was consumed — a
+ // mid-view llama_decode would clobber logits other slots still read
+ void update_score_suffixes() {
+ for (auto & slot : slots) {
+ if (!slot.score_suffix_pending) {
+ continue;
+ }
+ slot.score_suffix_pending = false;
+
+ if (!slot.is_processing() || !slot.task || slot.task->type != SERVER_TASK_TYPE_SCORE) {
+ continue; // the task was aborted mid-iteration
+ }
+
+ if (decode_score_suffixes(slot)) {
+ send_score(slot);
+ } else {
+ send_error(slot, "failed to decode score candidate suffixes", ERROR_TYPE_SERVER);
+ }
+ slot.release();
+ }
+ }
+
//
// Functions to process the task
//
@@ -2341,6 +2597,7 @@ private:
case SERVER_TASK_TYPE_INFILL:
case SERVER_TASK_TYPE_EMBEDDING:
case SERVER_TASK_TYPE_RERANK:
+ case SERVER_TASK_TYPE_SCORE:
{
// special case: if input is provided via CLI, tokenize it first
// otherwise, no need to tokenize as it's already done inside the HTTP thread
@@ -2832,6 +3089,13 @@ private:
break; // stop any further processing
}
}
+
+ try {
+ update_score_suffixes();
+ } catch (const std::exception & e) {
+ SRV_ERR("update_score_suffixes() failed: %s\n", e.what());
+ abort_all_slots("update_score_suffixes() failed: " + std::string(e.what()));
+ }
}
void pre_decode() {
@@ -3154,6 +3418,16 @@ private:
n_past = std::min(n_past, slot.alora_invocation_start - 1);
}
+ // score tasks need the logits that predict the first candidate
+ // token, so the last shared-prompt token must be (re-)decoded
+ // even when the cache already covers it
+ if (slot.task->type == SERVER_TASK_TYPE_SCORE) {
+ n_past = std::min(n_past, std::max(0, slot.task->n_score_prompt - 1));
+ // remember the divergence point before the checkpoint
+ // logic below possibly resets n_past to 0
+ slot.score_divergence = n_past;
+ }
+
const auto n_cache_reuse = slot.task->params.n_cache_reuse;
const bool can_cache_reuse =
@@ -3395,8 +3669,12 @@ private:
bool do_checkpoint = params_base.n_ctx_checkpoints > 0;
- // make checkpoints only for completion tasks
- do_checkpoint = do_checkpoint && slot.task->type == SERVER_TASK_TYPE_COMPLETION;
+ // make checkpoints for completion tasks, and for score tasks at the
+ // shared-prompt boundary: models whose memory cannot be partially
+ // rewound (SWA/hybrid/recurrent) would otherwise re-process the whole
+ // prompt for every candidate of a scoring call
+ do_checkpoint = do_checkpoint && (slot.task->type == SERVER_TASK_TYPE_COMPLETION ||
+ slot.task->type == SERVER_TASK_TYPE_SCORE);
// make a checkpoint of the parts of the memory that cannot be rolled back.
// checkpoints are created only if:
@@ -3463,10 +3741,17 @@ private:
// embedding requires all tokens in the batch to be output;
// MTP also wants logits at every prompt position so the
// streaming hook can mirror t_h_nextn into ctx_dft.
+ // score tasks need outputs at the positions that predict
+ // each candidate token (the token at index i predicts the
+ // task token at index i+1).
+ const bool need_score_logit =
+ slot.task->type == SERVER_TASK_TYPE_SCORE &&
+ slot.prompt.n_tokens() + 1 >= slot.task->n_score_prompt &&
+ slot.prompt.n_tokens() + 1 < slot.task->n_tokens();
add_ok &= batch.add(slot.id,
cur_tok,
slot.prompt.tokens.pos_next(),
- slot.need_embd());
+ slot.need_embd() || need_score_logit);
slot.prompt.tokens.push_back(cur_tok);
slot.n_prompt_tokens_processed++;
@@ -3481,6 +3766,32 @@ private:
}
}
+ // score tasks: break at the shared-prompt boundary so the checkpoint
+ // below lands exactly there — the other candidates of the same
+ // scoring call re-process only their own tokens. Also break at the
+ // point where this task diverged from the previous cache: after a
+ // forced re-prefill a checkpoint there serves the next scoring call
+ // over the same stable prefix (e.g. a classifier's option list).
+ // The caller-declared stable-prefix boundary is the strongest of
+ // these: a checkpoint there is at or before every future task's
+ // divergence within the same option list, so it always survives
+ // and always restores.
+ if (do_checkpoint && slot.task->type == SERVER_TASK_TYPE_SCORE &&
+ (slot.prompt.n_tokens() == slot.task->n_score_prompt - 1 ||
+ (slot.task->n_stable_prompt > 0 &&
+ slot.prompt.n_tokens() == slot.task->n_stable_prompt &&
+ slot.prompt.n_tokens() < slot.task->n_score_prompt - 1) ||
+ (slot.prompt.n_tokens() == slot.score_divergence &&
+ slot.prompt.n_tokens() < slot.task->n_score_prompt - 1))) {
+ bool have_ckpt = false;
+ for (const auto & ckpt : slot.prompt.checkpoints) {
+ have_ckpt |= ckpt.n_tokens == slot.prompt.n_tokens();
+ }
+ if (!have_ckpt) {
+ break;
+ }
+ }
+
// process the last few tokens of the prompt separately in order to allow for a checkpoint to be created.
// create checkpoints that many tokens before the end of the prompt:
// - 4 + n_ubatch
@@ -3513,6 +3824,15 @@ private:
const bool is_user_start = spans.is_user_start(n_tokens_start);
const bool is_last_user_message = n_tokens_start == last_user_pos;
+ // a batch starting at the score boundary or divergence point must
+ // always checkpoint — min-step spacing would otherwise suppress it
+ // and every candidate / next scoring call would re-process the prompt
+ const bool is_score_boundary = slot.task->type == SERVER_TASK_TYPE_SCORE &&
+ (n_tokens_start == slot.task->n_score_prompt - 1 ||
+ (slot.task->n_stable_prompt > 0 &&
+ n_tokens_start == slot.task->n_stable_prompt) ||
+ n_tokens_start == slot.score_divergence);
+
// entire prompt has been processed
if (slot.prompt.n_tokens() == slot.task->n_tokens()) {
slot.state = SLOT_STATE_DONE_PROMPT;
@@ -3528,8 +3848,8 @@ private:
slot.init_sampler();
} else {
// skip ordinary mid-prompt checkpoints, unless the batch starts a user
- // message or we are near the end of the prompt
- if (!is_user_start && !near_prompt_end) {
+ // message, the score boundary, or we are near the end of the prompt
+ if (!is_user_start && !is_score_boundary && !near_prompt_end) {
do_checkpoint = false;
}
}
@@ -3546,10 +3866,10 @@ private:
// do not checkpoint after mtmd chunks
do_checkpoint = do_checkpoint && !has_mtmd;
- // no need to create checkpoints that are too close together, unless it's the last user message
+ // no need to create checkpoints that are too close together, unless it's the last user message or the score boundary
do_checkpoint = do_checkpoint && (
slot.prompt.checkpoints.empty() ||
- is_last_user_message || near_prompt_end ||
+ is_last_user_message || near_prompt_end || is_score_boundary ||
n_tokens_start > slot.prompt.checkpoints.back().n_tokens + params_base.checkpoint_min_step);
SLT_DBG(slot, "main/do_checkpoint = %s, pos_min = %d, pos_max = %d\n", do_checkpoint ? "yes" : "no", pos_min, pos_max);
@@ -3703,6 +4023,13 @@ private:
}
}
+ // score slots harvest logprobs from every view that contains
+ // their outputs, not just the one holding the final token
+ if (slot.task && slot.task->type == SERVER_TASK_TYPE_SCORE &&
+ (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_DONE_PROMPT)) {
+ collect_score_logprobs(slot, batch_view);
+ }
+
if (!is_inside_view(slot.i_batch)) {
// the required token not in this sub-batch, skip
return;
@@ -3724,6 +4051,25 @@ private:
return;
}
+ if (slot.task->type == SERVER_TASK_TYPE_SCORE) {
+ // shared-prefix logprobs (and every candidate's first
+ // suffix logprob) were accumulated per view above;
+ // candidates with more suffix tokens still need the
+ // forked decode at the end of update_slots()
+ for (const auto & sfx : slot.task->score_suffixes) {
+ if (sfx.size() > 1) {
+ slot.score_suffix_pending = true;
+ break;
+ }
+ }
+ if (!slot.score_suffix_pending) {
+ send_score(slot);
+ slot.release();
+ }
+ slot.i_batch = -1;
+ return;
+ }
+
GGML_ASSERT(slot.task->need_sampling());
// prompt evaluated for next-token prediction
diff --git a/tools/server/server-task.h b/tools/server/server-task.h
index c3eea2e..fb3c178 100644
--- a/tools/server/server-task.h
+++ b/tools/server/server-task.h
@@ -13,10 +13,25 @@
using json = nlohmann::ordered_json;
+// SERVER_TASK_TYPE_SCORE emits one logits output per candidate token (plus
+// the forced last-token output), and the context's output budget
+// (n_outputs_max) is reserved up front — so candidate length must be
+// bounded. Raising this raises the worst-case compute-buffer reservation
+// by ~n_vocab * 4 bytes per extra output.
+constexpr int32_t SERVER_SCORE_MAX_CAND_TOKENS = 64;
+
+// Maximum sequences forked off the shared prefix in one score suffix
+// decode. The context is created with this many seq ids (and
+// recurrent-state cells) beyond the parallel slots — see
+// common_params::n_seq_score_forks; candidates in excess of the budget
+// are decoded in successive chunks.
+constexpr int32_t SERVER_SCORE_FORK_SEQS = 16;
+
enum server_task_type {
SERVER_TASK_TYPE_COMPLETION,
SERVER_TASK_TYPE_EMBEDDING,
SERVER_TASK_TYPE_RERANK,
+ SERVER_TASK_TYPE_SCORE,
SERVER_TASK_TYPE_INFILL,
SERVER_TASK_TYPE_CANCEL,
SERVER_TASK_TYPE_CONTROL,
@@ -153,6 +168,18 @@ struct server_task {
task_params params;
server_tokens tokens;
+ // used by SERVER_TASK_TYPE_SCORE: `tokens` holds the shared prefix
+ // (prompt + longest common candidate token prefix) and logprobs are
+ // returned for its tokens from n_score_prompt onward. Each candidate's
+ // tokens beyond the shared prefix ride a forked sequence.
+ int32_t n_score_prompt = 0;
+ std::vector<llama_tokens> score_suffixes;
+ // token index where the caller-declared stable prompt prefix ends
+ // (0 = no hint): the option-list system prompt that repeats across
+ // scoring calls. A context checkpoint is forced there so models that
+ // cannot rewind state re-process only the per-call tail next time.
+ int32_t n_stable_prompt = 0;
+
// only used by CLI, this allow tokenizing CLI inputs on server side
// we need this because mtmd_context and vocab are not accessible outside of server_context
bool cli = false;
@@ -197,6 +224,7 @@ struct server_task {
switch (type) {
case SERVER_TASK_TYPE_COMPLETION:
case SERVER_TASK_TYPE_INFILL:
+ case SERVER_TASK_TYPE_SCORE:
return true;
default:
return false;
@@ -494,6 +522,25 @@ struct server_task_result_rerank : server_task_result {
virtual json to_json() override;
};
+struct server_task_result_score : server_task_result {
+ // log P(token | prefix) for the shared-prefix tokens after
+ // n_score_prompt, in order; NaN marks positions the decode never
+ // produced an output for
+ std::vector<float> shared_logprobs;
+
+ // per candidate: logprobs of its suffix tokens, in task order (entry
+ // 0 is the token right after the shared prefix, predicted by the last
+ // shared token's logits)
+ std::vector<std::vector<float>> cand_logprobs;
+
+ virtual json to_json() override {
+ return json {
+ {"shared_logprobs", shared_logprobs},
+ {"cand_logprobs", cand_logprobs},
+ };
+ }
+};
+
struct server_task_result_error : server_task_result {
error_type err_type = ERROR_TYPE_SERVER;
std::string err_msg;