Skip to content

Commit aff5af4

Browse files
committed
feat(routing): Score gRPC primitive, score classifier, L2 embedding cache
Replace the prior feature/knn/llm router classifiers with a single score-based classifier that asks an Arch-Router-style model to rank every policy label as a continuation of the routing prompt and reads off the softmax distribution. Multi-label routing falls out of this naturally: the middleware activates every label whose probability crosses a softmax threshold and picks the first candidate whose labels are a superset of the active set. Wiring summary: - backend.proto adds Score(ScoreRequest) → ScoreResponse. The llama-cpp C++ backend implements Score on top of force-decoded candidates against a freshly-cleared KV cache (prompt-KV sharing optimisation is on the perf TODO list); vLLM uses prompt_logprobs. Other backends return UNIMPLEMENTED. - core/services/routing/router/score.go is the classifier. It builds the ChatML routing prompt once at construction, scores every policy label as a continuation, and applies an activation threshold (default 0.15; 0.40 is a better empirical default on Arch-Router-1.5B per the eval in features/middleware.md). - RouterConfig grows Policies, ActivationThreshold, and an optional EmbeddingCache nested struct. RouterCandidate collapses to {Model, Labels[]} — labels are the matching contract, descriptions live on the policy. - The dead feature/knn/llm/routing_data files are removed. L2 embedding cache: - core/services/routing/router/embedding_cache.go wraps a Classifier decorator that embeds each probe, KNN-searches the per-router local-store collection, returns a cached decision if the cosine similarity passes a threshold (default 0.80, lowered from 0.92 after the eval against nomic-embed-text-v1.5 paraphrases). Low- confidence decisions are deliberately not cached so they can't poison future paraphrases. - Stats include hits, misses, near_misses, low_confidence, and a 10-bin similarity histogram so admins can see where the cosine distribution sits relative to the configured threshold. - Registry tracks built classifiers by fingerprint of the RouterConfig YAML, so config edits invalidate the cache wrapper automatically while the on-disk vectors persist. UI: - The model-editor schema is rewritten: dead KNN/LLM fields gone, policies/activation_threshold/embedding_cache.* added with proper descriptions, sliders, and component bindings. - RouterCandidatesEditor is rewritten for {model, labels[]} with multi-select label chips populated from router.policies via a new FormContext. - RouterPoliciesEditor is the structured editor for the label vocabulary, with duplicate-label detection via a memoised set. - The Routing tab on /app/middleware renders the embedding-cache histogram inline with a threshold marker. Verification: - Unit tests cover the score classifier (multi-label activation, fallback, depth-1) and the embedding cache (hit, near-miss, low-confidence skip, embedder/store error fallthrough, histogram population). - Refreshed e2e specs (router-template.spec.js, middleware-page.spec.js) pass under make test-ui-e2e-docker: 133/135 passing with the two failures unrelated to this slice. - End-to-end eval against the LocalAGI stack with a 30-prompt corpus + 3 paraphrases each produced 35% steady-state hit rate at 0.80 threshold (53% of caching-eligible decisions), 15ms p50 cache-hit latency vs 246ms classifier round-trip — a ~16× speedup on hits. Assisted-by: Claude:claude-opus-4-7 Signed-off-by: Richard Palethorpe <io@richiejp.com>
1 parent 52ccf19 commit aff5af4

51 files changed

Lines changed: 3661 additions & 2754 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

.dockerignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
.devcontainer
55
models
66
backends
7+
volumes
78
examples/chatbot-ui/models
89
backend/go/image/stablediffusion-ggml/build/
910
backend/go/*/build

backend/backend.proto

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,15 @@ service Backend {
4444
// other unformatted PII that regex misses.
4545
rpc TokenClassify(TokenClassifyRequest) returns (TokenClassifyResponse) {}
4646

47+
// Score evaluates the model's joint log-probability of each
48+
// supplied candidate continuation given a shared prompt. The
49+
// prompt's KV cache is computed once and reused across candidates.
50+
// Used for routing-policy multi-label classification, reranking,
51+
// calibrated confidence, and reward-model scoring — any task where
52+
// the consumer wants the model's confidence in a pre-specified
53+
// continuation rather than a generated one.
54+
rpc Score(ScoreRequest) returns (ScoreResponse) {}
55+
4756
rpc GetMetrics(MetricsRequest) returns (MetricsResponse);
4857

4958
rpc VAD(VADRequest) returns (VADResponse) {}
@@ -111,6 +120,48 @@ message TokenClassifyResponse {
111120
repeated TokenClassifyEntity entities = 1;
112121
}
113122

123+
// ScoreRequest carries one shared prompt and one or more continuations
124+
// to score against it. The backend tokenises the prompt once and reuses
125+
// the resulting KV cache across all candidates in this request.
126+
message ScoreRequest {
127+
string prompt = 1;
128+
repeated string candidates = 2;
129+
// Return per-token logprobs for each candidate when true. Default
130+
// false to keep the wire response small; the joint log_prob field
131+
// covers the common ranking case.
132+
bool include_token_logprobs = 3;
133+
// When true, the response also populates length_normalized_log_prob
134+
// (joint log-prob divided by candidate token count). Useful when
135+
// candidates differ in length and the consumer wants a per-token
136+
// measure comparable across them (PMI-style scoring).
137+
bool length_normalize = 4;
138+
}
139+
140+
// CandidateScore is one row in the ScoreResponse, matching by index
141+
// the candidate in ScoreRequest.candidates.
142+
message CandidateScore {
143+
// Sum of log P(token_i | prompt, candidate_token_<i) across the
144+
// candidate's tokens. The primary ranking signal.
145+
double log_prob = 1;
146+
// log_prob / num_tokens — populated when length_normalize=true on
147+
// the request.
148+
double length_normalized_log_prob = 2;
149+
// Per-token detail — populated when include_token_logprobs=true.
150+
repeated TokenLogProb tokens = 3;
151+
// Number of tokens the backend tokenised this candidate into, after
152+
// any backend-specific normalisation (e.g. leading-space handling).
153+
int32 num_tokens = 4;
154+
}
155+
156+
message TokenLogProb {
157+
string token = 1;
158+
double log_prob = 2;
159+
}
160+
161+
message ScoreResponse {
162+
repeated CandidateScore candidates = 1;
163+
}
164+
114165
message RerankRequest {
115166
string query = 1;
116167
repeated string documents = 2;

backend/cpp/llama-cpp/grpc-server.cpp

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,9 @@
3232
#include <grpcpp/health_check_service_interface.h>
3333
#include <grpcpp/security/server_credentials.h>
3434
#include <regex>
35+
#include <algorithm>
3536
#include <atomic>
37+
#include <cmath>
3638
#include <cstdlib>
3739
#include <fstream>
3840
#include <iterator>
@@ -2949,6 +2951,207 @@ class BackendServiceImpl final : public backend::Backend::Service {
29492951
return grpc::Status::OK;
29502952
}
29512953

2954+
// Score returns the model's joint log-probability of each candidate
2955+
// continuation given a shared prompt.
2956+
//
2957+
// WHY bypass the slot/task queue: upstream server_context exposes
2958+
// get_llama_context as "main thread only" and the slot loop's
2959+
// update_slots() owns the context whenever a task is in flight.
2960+
// No public synchronization primitive is available — so Score is
2961+
// unsafe to call concurrently with active generation through this
2962+
// backend. In practice routing-classifier calls happen before the
2963+
// request is routed to a generation backend, so the model used
2964+
// for Score is typically idle. Concurrent Score calls are
2965+
// serialised by a local mutex; KV-cache state is isolated behind
2966+
// a dedicated sequence ID cleared between candidates.
2967+
//
2968+
// A patch to server-context.cpp that adds SERVER_TASK_TYPE_SCORE
2969+
// and routes scoring through the slot loop would be the correct
2970+
// long-term fix; tracked as a follow-up.
2971+
//
2972+
// Perf TODO (measured: ~450 ms warm for 3 candidates on Arch-
2973+
// Router-1.5B Q4_K_M + Intel SYCL): the current loop re-decodes
2974+
// `prompt + candidate` from scratch for every candidate, throwing
2975+
// away the prompt's KV cache between iterations. A smarter
2976+
// version would:
2977+
// 1. Decode just the prompt once into score_seq_id.
2978+
// 2. Snapshot/cp that sequence (llama_memory_seq_cp) into a
2979+
// per-candidate sequence id.
2980+
// 3. For each candidate, decode only its tokens onto the copy
2981+
// (continuing from the saved prompt state), read logits.
2982+
// 4. llama_memory_seq_rm the copy.
2983+
// Estimated speedup: 3-candidate calls 450 ms -> ~150-200 ms,
2984+
// 6-candidate calls 630 ms -> ~220 ms. Single source-file change,
2985+
// no proto / Go-side changes needed. Worth doing once routing is
2986+
// wired into the middleware and Score is on the hot path of every
2987+
// chat request.
2988+
grpc::Status Score(ServerContext* context, const backend::ScoreRequest* request, backend::ScoreResponse* response) override {
2989+
auto auth = checkAuth(context);
2990+
if (!auth.ok()) return auth;
2991+
if (params_base.model.path.empty()) {
2992+
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
2993+
}
2994+
if (request->candidates_size() == 0) {
2995+
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "candidates must be non-empty");
2996+
}
2997+
2998+
// Serialise concurrent Score calls. The slot loop is still
2999+
// free to race with us — see the class comment above.
3000+
static std::mutex score_mutex;
3001+
std::lock_guard<std::mutex> score_lock(score_mutex);
3002+
3003+
llama_context * lctx = ctx_server.get_llama_context();
3004+
if (lctx == nullptr) {
3005+
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "llama context unavailable (sleeping?)");
3006+
}
3007+
const llama_vocab * vocab = ctx_server.impl->vocab;
3008+
const int32_t n_vocab = llama_vocab_n_tokens(vocab);
3009+
const int32_t n_ctx = llama_n_ctx(lctx);
3010+
llama_memory_t mem = llama_get_memory(lctx);
3011+
3012+
// The KV-cache is sized to seq_to_stream.size() at load
3013+
// (typically equal to n_slots, often 1). Sequence IDs must
3014+
// be in [0, n_seq_max), so we can't pick a high-value
3015+
// "private" ID — we have to share with the slot. We clear
3016+
// the cache before AND after each candidate to keep
3017+
// scoring isolated from whatever state the slot held, and
3018+
// the static mutex above guarantees no other Score call is
3019+
// racing in the meantime. The slot loop is still free to
3020+
// race (see comment on this method) — Score must not run
3021+
// concurrently with generation through this backend.
3022+
const llama_seq_id score_seq_id = 0;
3023+
llama_memory_seq_rm(mem, score_seq_id, -1, -1);
3024+
3025+
// Tokenize the shared prompt once with add_special=true so
3026+
// BOS is prepended when the model requires it. parse_special
3027+
// keeps chat-template markers in the prompt intact.
3028+
const std::string prompt = request->prompt();
3029+
std::vector<llama_token> prompt_tokens = common_tokenize(vocab, prompt, /*add_special=*/true, /*parse_special=*/true);
3030+
const int32_t prompt_len = (int32_t) prompt_tokens.size();
3031+
3032+
for (int ci = 0; ci < request->candidates_size(); ci++) {
3033+
const std::string & candidate_text = request->candidates(ci);
3034+
3035+
// Re-tokenize prompt + candidate as a single string. BPE
3036+
// merges across the boundary can shift the tokenization
3037+
// versus tokenize(prompt) ++ tokenize(candidate), so we
3038+
// find the divergence point against prompt_tokens.
3039+
std::vector<llama_token> full_tokens = common_tokenize(vocab, prompt + candidate_text, /*add_special=*/true, /*parse_special=*/true);
3040+
int32_t divergence = prompt_len;
3041+
const int32_t min_len = std::min<int32_t>(prompt_len, (int32_t) full_tokens.size());
3042+
for (int32_t i = 0; i < min_len; i++) {
3043+
if (prompt_tokens[i] != full_tokens[i]) {
3044+
divergence = i;
3045+
break;
3046+
}
3047+
}
3048+
const int32_t cand_len = (int32_t) full_tokens.size() - divergence;
3049+
backend::CandidateScore * cs = response->add_candidates();
3050+
cs->set_num_tokens(cand_len);
3051+
if (cand_len <= 0) {
3052+
cs->set_log_prob(0.0);
3053+
if (request->length_normalize()) {
3054+
cs->set_length_normalized_log_prob(0.0);
3055+
}
3056+
continue;
3057+
}
3058+
if (divergence < 1) {
3059+
// Need at least one prior token (typically BOS) to
3060+
// predict the first candidate token's logit. Tokeniser
3061+
// models without BOS + an empty prompt fall in here.
3062+
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
3063+
"Score: prompt produced no leading tokens; need at least one (e.g. BOS) to predict candidate");
3064+
}
3065+
if ((int32_t) full_tokens.size() > n_ctx) {
3066+
return grpc::Status(grpc::StatusCode::OUT_OF_RANGE,
3067+
"Score: prompt+candidate exceeds context size (got " +
3068+
std::to_string(full_tokens.size()) + ", n_ctx=" + std::to_string(n_ctx) + ")");
3069+
}
3070+
3071+
// Build a batch covering the entire prompt+candidate. We
3072+
// need logits at (divergence-1) onward — those are the
3073+
// predictions for each candidate token.
3074+
llama_batch batch = llama_batch_init((int32_t) full_tokens.size(), 0, 1);
3075+
for (int32_t i = 0; i < (int32_t) full_tokens.size(); i++) {
3076+
batch.token[i] = full_tokens[i];
3077+
batch.pos[i] = i;
3078+
batch.n_seq_id[i] = 1;
3079+
batch.seq_id[i][0] = score_seq_id;
3080+
// logits[i] is "do we want the prediction *for the
3081+
// next token*, computed from this position?"
3082+
// We want predictions for candidate tokens at
3083+
// positions divergence .. full_tokens.size()-1, which
3084+
// come from logits at positions (divergence-1) ..
3085+
// (full_tokens.size()-2).
3086+
bool need_logit = (i >= divergence - 1) && (i < (int32_t) full_tokens.size() - 1);
3087+
batch.logits[i] = need_logit ? 1 : 0;
3088+
}
3089+
batch.n_tokens = (int32_t) full_tokens.size();
3090+
3091+
// Decode the batch. If decode fails (e.g. KV slot
3092+
// exhaustion), surface as INTERNAL — the caller will
3093+
// typically fall back to a sampling-based classifier.
3094+
int decode_err = llama_decode(lctx, batch);
3095+
if (decode_err != 0) {
3096+
llama_batch_free(batch);
3097+
llama_memory_seq_rm(mem, score_seq_id, -1, -1);
3098+
return grpc::Status(grpc::StatusCode::INTERNAL,
3099+
"llama_decode failed during Score: " + std::to_string(decode_err));
3100+
}
3101+
3102+
// Sum log-probabilities of the actual candidate tokens.
3103+
double total_log_prob = 0.0;
3104+
for (int32_t k = 0; k < cand_len; k++) {
3105+
// The k-th candidate token sits at full_tokens index
3106+
// (divergence + k). Its predicting logit is at batch
3107+
// position (divergence + k - 1).
3108+
int32_t logit_pos = divergence + k - 1;
3109+
const float * logits = llama_get_logits_ith(lctx, logit_pos);
3110+
if (logits == nullptr) {
3111+
llama_batch_free(batch);
3112+
llama_memory_seq_rm(mem, score_seq_id, -1, -1);
3113+
return grpc::Status(grpc::StatusCode::INTERNAL,
3114+
"llama_get_logits_ith returned null at position " + std::to_string(logit_pos));
3115+
}
3116+
llama_token target_token = full_tokens[divergence + k];
3117+
3118+
// Compute log_softmax(logits)[target_token] with the
3119+
// max-subtraction stability trick.
3120+
float max_logit = logits[0];
3121+
for (int32_t v = 1; v < n_vocab; v++) {
3122+
if (logits[v] > max_logit) max_logit = logits[v];
3123+
}
3124+
double sum_exp = 0.0;
3125+
for (int32_t v = 0; v < n_vocab; v++) {
3126+
sum_exp += std::exp((double)(logits[v] - max_logit));
3127+
}
3128+
double token_log_prob = (double)(logits[target_token] - max_logit) - std::log(sum_exp);
3129+
total_log_prob += token_log_prob;
3130+
3131+
if (request->include_token_logprobs()) {
3132+
backend::TokenLogProb * tlp = cs->add_tokens();
3133+
std::string piece = common_token_to_piece(lctx, target_token);
3134+
tlp->set_token(piece);
3135+
tlp->set_log_prob(token_log_prob);
3136+
}
3137+
}
3138+
3139+
cs->set_log_prob(total_log_prob);
3140+
if (request->length_normalize() && cand_len > 0) {
3141+
cs->set_length_normalized_log_prob(total_log_prob / (double) cand_len);
3142+
}
3143+
3144+
llama_batch_free(batch);
3145+
// Drop this candidate's KV-cache contribution so the next
3146+
// candidate starts from a clean state. Without this, the
3147+
// next decode would conflict at positions 0..N-1 for our
3148+
// sequence ID.
3149+
llama_memory_seq_rm(mem, score_seq_id, -1, -1);
3150+
}
3151+
3152+
return grpc::Status::OK;
3153+
}
3154+
29523155
grpc::Status TokenizeString(ServerContext* context, const backend::PredictOptions* request, backend::TokenizationResponse* response) override {
29533156
auto auth = checkAuth(context);
29543157
if (!auth.ok()) return auth;

0 commit comments

Comments
 (0)