Add approximate prefix routing, and find out it did nothing
CI / test (push) Successful in 1m19s

Tier 1 of plan §7.5: a radix tree of request prefixes that estimates which
replica has already prefilled one. Bytes, not tokens — needing the engine's
exact tokenizer is what makes tier 2 a separate project.

It went in inert, and the decision metric is what caught it: 144 lookups, zero
hits. The cause was a rule I had written and tested for. A node's label is the
first request that created it, so the second request for the same system
prompt diverges in its last few bytes — a new question after the same
preamble, the single most common shape there is. Requiring the whole label to
match threw away the 700 bytes they shared. What matters is the shared head,
because that is what the engine cached, and a partial label match now counts.
Two unit tests asserted the wrong rule and have been inverted.

With that fixed, measured end to end against engines whose caches are too
small for the working set: 22.5% cache hits without, 81.7% with, and the
breakdown shows the misses are exactly the first round, when the tree is
empty.

The scenario matters as much as the number. With a handful of prompts and a
large cache, scattering is free — every replica ends up holding everything —
and prefix routing can only add constraint. The first version of the test
measured that case and showed a loss, which was a true statement about a
situation the feature is not for.

Two things keep it from being the disaster prefix routing usually is. The
imbalance guard, which needs both an absolute and a relative gap before it
fires, because either alone misfires — absolute on an idle cluster, relative
on a busy one. And it is a scorer rather than a filter: affinity is a vote the
queue and inflight scorers can outvote, not a rule.

The tree records the replica the picker actually chose, not the one the scorer
proposed, and a node prefers the replica it has seen most rather than most
recently — a count is stable under a randomizing picker, where recency just
chases whatever the picker last did.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Ilker ManapandClaude Opus 5 committed 2026-09-16 12:13:19 +02:00
1 parent de21649ad6
commit a075ff4120
13 files changed
+1207 -7

No files matched your search

+33
View File
@@ -110,6 +110,39 @@ router:
failure_threshold: 3
default_profile: chat # chat | batch | long-context | debug
# Approximate prefix routing (plan §7.5 tier 1): send a request to the replica
# that has probably already prefilled its prefix.
#
# It earns its keep when the working set does not fit in one replica's cache.
# With a handful of prompts and a large cache, scattering is free — every
# replica ends up holding everything — and this only adds constraint. Measured
# with sim-engines whose caches were too small for the working set: cache hits
# went from 22.5% to 81.7%.
#
# Watch llmgw_prefix_decisions_total to tune it. A high
# load_balance_override count means the guard is doing its job and affinity is
# not getting a chance; a high cache_miss_min_load means match_ratio or
# min_match is too strict for your traffic.
prefix:
enabled: false
# The shortest match worth acting on, in bytes. Below this the overlap is
# coincidence — two English prompts share "The".
min_match: 64
# The share of the prompt that must match. A short shared header on a very
# long prompt saves almost no prefill.
match_ratio: 0.3
# The imbalance guard, and both must trip before affinity is abandoned.
# Absolute alone fires on an idle cluster where one replica has two requests
# and another none; relative alone fires on a busy one where 200 against 180
# is a rounding error. Without the guard the cache hit rate rises while p99
# collapses.
balance_abs: 8
balance_rel: 1.5
# An engine evicts on its own schedule, so an old record may describe a
# cache that is gone.
ttl: 10m
max_nodes: 100000
session_header: x-session-id # affinity key; falls back to the `user` field
# Ask engines for usage even when the client did not, so token accounting
# works regardless of client behaviour. The extra chunk is dropped on the way
+23
View File
@@ -149,6 +149,11 @@ type Router struct {
// DefaultProfile names the routing profile used when a model picks none.
DefaultProfile string `koanf:"default_profile"`
// Prefix turns on the approximate prefix tree (plan §7.5 tier 1). Off, the
// router uses session affinity alone, which is tier 0 and already most of
// the win for multi-turn traffic.
Prefix Prefix `koanf:"prefix"`
// SessionHeader carries the session id for affinity (plan §7.5 tier 0).
SessionHeader string `koanf:"session_header"`
// InjectUsage asks engines for usage even when the client did not, so token
@@ -181,6 +186,24 @@ type Admission struct {
Interval time.Duration `koanf:"interval"`
}
// Prefix configures approximate prefix routing.
type Prefix struct {
Enabled bool `koanf:"enabled"`
// MaxNodes bounds the tree, MinMatch is the shortest match worth acting on
// in bytes, and TTL is how long an entry is believed — an engine evicts on
// its own schedule, so an old record may describe a cache that is gone.
MaxNodes int `koanf:"max_nodes"`
MinMatch int `koanf:"min_match"`
TTL time.Duration `koanf:"ttl"`
// MatchRatio is the share of the prompt that must match to be worth acting
// on. BalanceAbs and BalanceRel are the imbalance guard, and both must trip
// before affinity is abandoned.
MatchRatio float64 `koanf:"match_ratio"`
BalanceAbs float64 `koanf:"balance_abs"`
BalanceRel float64 `koanf:"balance_rel"`
}
// Band is one priority class.
type Band struct {
Name string `koanf:"name"`
+17 -1
View File
@@ -6,6 +6,7 @@ import (
"github.com/prometheus/client_golang/prometheus"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler"
)
@@ -26,10 +27,17 @@ type metrics struct {
admissionWait *prometheus.HistogramVec
admissionRejected *prometheus.CounterVec
prefixDecisions *prometheus.CounterVec
}
func newMetrics() *metrics {
return &metrics{
prefixDecisions: prometheus.NewCounterVec(prometheus.CounterOpts{
Name: "llmgw_prefix_decisions_total",
Help: "Prefix routing decisions by branch: cache_hit, cache_miss_min_load, " +
"load_balance_override, no_tree, no_healthy_worker.",
}, []string{"decision"}),
requests: prometheus.NewCounterVec(prometheus.CounterOpts{
Name: "llmgw_requests_total",
Help: "Requests handled, by model, pool and response status.",
@@ -86,10 +94,18 @@ func (m *metrics) collectors() []prometheus.Collector {
return []prometheus.Collector{
m.requests, m.duration, m.ttft, m.tokens,
m.scorerValue, m.filtered, m.considered,
m.admissionWait, m.admissionRejected,
m.admissionWait, m.admissionRejected, m.prefixDecisions,
}
}
// observePrefix records which branch a prefix decision took.
//
// Without it, "affinity is not helping" and "affinity never applied" look
// identical from the outside, and the thresholds can only be tuned by guessing.
func (m *metrics) observePrefix(d prefix.Decision) {
m.prefixDecisions.WithLabelValues(string(d)).Inc()
}
func (m *metrics) observeAdmission(band, model string, wait time.Duration) {
if band == "" {
band = "none"
+126
View File
@@ -0,0 +1,126 @@
package prefix
import (
"math"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
)
// Decision is which branch a prefix lookup took.
//
// Recorded as a metric because tuning prefix routing without it is guesswork:
// a cache hit rate says nothing about how often affinity was overridden, and
// "affinity is not helping" and "affinity never applied" look identical from
// the outside.
type Decision string
// The branches.
const (
Hit Decision = "cache_hit"
MissMinLoad Decision = "cache_miss_min_load"
LoadOverride Decision = "load_balance_override"
NoTree Decision = "no_tree"
NoCandidates Decision = "no_healthy_worker"
)
// Thresholds decide when the prefix is worth following.
type Thresholds struct {
// MatchRatio is the share of the prompt that must match. Below it, the
// overlap is not enough to save a meaningful prefill.
MatchRatio float64
// BalanceAbs and BalanceRel are the imbalance guard, and both must trip
// before it does. Without this the cache hit rate rises while p99 latency
// collapses: affinity keeps sending a hot prefix to one replica long after
// that replica has stopped being able to serve it quickly.
//
// Absolute alone would fire on an idle cluster where one replica has two
// requests and another none. Relative alone would fire on a busy one where
// 200 against 180 is a rounding error. Together they mean "meaningfully
// more loaded, and by enough to matter".
BalanceAbs float64
BalanceRel float64
}
// DefaultThresholds returns values that work without measurement.
func DefaultThresholds() Thresholds {
return Thresholds{MatchRatio: 0.3, BalanceAbs: 8, BalanceRel: 1.5}
}
func (t *Thresholds) normalize() {
d := DefaultThresholds()
if t.MatchRatio <= 0 {
t.MatchRatio = d.MatchRatio
}
if t.BalanceAbs <= 0 {
t.BalanceAbs = d.BalanceAbs
}
if t.BalanceRel <= 1 {
t.BalanceRel = d.BalanceRel
}
}
// Choose picks the replica a prompt should go to, and says why.
//
// It returns an index into eps, or -1 when the caller should decide on load
// alone — which is the common case and not a failure.
func Choose(t *Tree, prompt string, eps []endpoint.Snapshot, th Thresholds) (int, Decision) {
th.normalize()
if len(eps) == 0 {
return -1, NoCandidates
}
if t == nil || prompt == "" {
return -1, NoTree
}
// The guard runs before the lookup, not after. Checking afterwards would
// mean paying for a match only to discard it, and — worse — would invite
// the shape where a cache hit quietly wins because somebody moved the
// check.
if imbalanced(eps, th) {
return -1, LoadOverride
}
replica, matched := t.Match(prompt)
if replica == "" || float64(matched)/float64(len(prompt)) < th.MatchRatio {
return -1, MissMinLoad
}
for i, e := range eps {
if e.ID == replica {
return i, Hit
}
}
// The match names a replica that is no longer a candidate: gone, draining,
// or its breaker is open. That is a miss, not an error.
return -1, MissMinLoad
}
// imbalanced reports whether load is lopsided enough to ignore affinity.
func imbalanced(eps []endpoint.Snapshot, th Thresholds) bool {
if len(eps) < 2 {
return false
}
maxL, minL := math.Inf(-1), math.Inf(1)
for _, e := range eps {
l := load(e)
maxL = math.Max(maxL, l)
minL = math.Min(minL, l)
}
if maxL-minL <= th.BalanceAbs {
return false
}
// A ratio against zero is infinite, and an idle replica is exactly where
// the next request should go, so treat it as tripping the guard.
if minL <= 0 {
return true
}
return maxL/minL > th.BalanceRel
}
// load is what the guard compares.
//
// Queue depth plus what the router has dispatched but not seen finish: the
// metric alone lags by a scrape interval, which is long enough for affinity to
// pile a burst onto one replica before anything notices.
func load(e endpoint.Snapshot) float64 {
return float64(e.Stats.Waiting + e.Stats.Running + e.Inflight)
}
+151
View File
@@ -0,0 +1,151 @@
package prefix
import (
"strings"
"testing"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
)
func ep(id string, waiting, running, inflight int) endpoint.Snapshot {
return endpoint.Snapshot{
Info: endpoint.Info{ID: id},
Stats: endpoint.Stats{Waiting: waiting, Running: running},
Inflight: inflight,
}
}
const prompt = "You are a helpful assistant answering questions about distributed systems. "
func TestAMatchIsFollowedWhenLoadIsEven(t *testing.T) {
tr := New(Config{MinMatch: 16})
tr.Observe(prompt, "b")
eps := []endpoint.Snapshot{ep("a", 1, 1, 0), ep("b", 1, 1, 0)}
i, d := Choose(tr, prompt+"What is a quorum?", eps, DefaultThresholds())
if d != Hit || i != 1 {
t.Errorf("index %d, decision %s; want a hit on b", i, d)
}
}
func TestAffinityIsAbandonedWhenLoadIsLopsided(t *testing.T) {
// The guard the plan insists on. Without it the cache hit rate rises while
// p99 collapses: affinity keeps feeding a hot prefix to one replica long
// after it stopped being able to serve it quickly.
tr := New(Config{MinMatch: 16})
tr.Observe(prompt, "b")
eps := []endpoint.Snapshot{ep("a", 0, 1, 0), ep("b", 40, 8, 4)}
i, d := Choose(tr, prompt+"What is a quorum?", eps, DefaultThresholds())
if d != LoadOverride {
t.Errorf("decision %s; the guard should have fired", d)
}
if i != -1 {
t.Errorf("index %d; the caller decides on load alone", i)
}
}
func TestBothHalvesOfTheGuardMustTrip(t *testing.T) {
tr := New(Config{MinMatch: 16})
tr.Observe(prompt, "b")
full := prompt + "What is a quorum?"
th := Thresholds{MatchRatio: 0.1, BalanceAbs: 8, BalanceRel: 1.5}
cases := []struct {
name string
eps []endpoint.Snapshot
want Decision
}{
{
// Relative is huge, absolute is tiny: an idle cluster where one
// replica has two requests and another none.
"a big ratio over a small difference",
[]endpoint.Snapshot{ep("a", 0, 0, 0), ep("b", 2, 0, 0)},
Hit,
},
{
// Absolute is large, relative is nothing: a busy cluster where 200
// against 180 is a rounding error.
"a big difference at a small ratio",
[]endpoint.Snapshot{ep("a", 180, 0, 0), ep("b", 200, 0, 0)},
Hit,
},
{
"both, which is what the guard is for",
[]endpoint.Snapshot{ep("a", 2, 0, 0), ep("b", 60, 0, 0)},
LoadOverride,
},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
if _, d := Choose(tr, full, c.eps, th); d != c.want {
t.Errorf("decision %s, want %s", d, c.want)
}
})
}
}
func TestAnIdleReplicaAlwaysWinsOverAffinity(t *testing.T) {
// A ratio against zero is infinite, and an idle replica is exactly where
// the next request should go.
tr := New(Config{MinMatch: 16})
tr.Observe(prompt, "b")
eps := []endpoint.Snapshot{ep("a", 0, 0, 0), ep("b", 20, 0, 0)}
if _, d := Choose(tr, prompt+"x", eps, DefaultThresholds()); d != LoadOverride {
t.Errorf("decision %s; an idle replica should win", d)
}
}
func TestASmallMatchRatioIsAMiss(t *testing.T) {
// A short shared header on a very long prompt saves almost no prefill, and
// pretending otherwise trades load balance for nothing.
tr := New(Config{MinMatch: 8})
tr.Observe("SYSTEM: be brief. ", "b")
long := "SYSTEM: be brief. " + strings.Repeat("unique content here. ", 200)
eps := []endpoint.Snapshot{ep("a", 1, 0, 0), ep("b", 1, 0, 0)}
if _, d := Choose(tr, long, eps, Thresholds{MatchRatio: 0.5}); d != MissMinLoad {
t.Errorf("decision %s, want a miss", d)
}
}
func TestAMatchOnAReplicaThatIsGoneIsAMiss(t *testing.T) {
// The tree outlives a replica leaving the candidate set — drained, breaker
// open, removed. That is a miss, not an error.
tr := New(Config{MinMatch: 16})
tr.Observe(prompt, "vanished")
eps := []endpoint.Snapshot{ep("a", 1, 0, 0), ep("b", 1, 0, 0)}
if i, d := Choose(tr, prompt+"x", eps, DefaultThresholds()); d != MissMinLoad || i != -1 {
t.Errorf("index %d decision %s", i, d)
}
}
func TestTheDegenerateCasesAreNamed(t *testing.T) {
// Every branch is a distinct decision, because "affinity is not helping"
// and "affinity never applied" look identical from the outside otherwise.
tr := New(Config{MinMatch: 4})
if _, d := Choose(tr, "p", nil, DefaultThresholds()); d != NoCandidates {
t.Errorf("no endpoints: %s", d)
}
eps := []endpoint.Snapshot{ep("a", 0, 0, 0)}
if _, d := Choose(nil, "p", eps, DefaultThresholds()); d != NoTree {
t.Errorf("no tree: %s", d)
}
if _, d := Choose(tr, "", eps, DefaultThresholds()); d != NoTree {
t.Errorf("no prompt: %s", d)
}
}
func TestOneCandidateNeverTripsTheGuard(t *testing.T) {
// There is nothing to balance against, so overriding affinity would throw
// away a cache hit for no reason.
tr := New(Config{MinMatch: 16})
tr.Observe(prompt, "only")
eps := []endpoint.Snapshot{ep("only", 500, 500, 500)}
if _, d := Choose(tr, prompt+"x", eps, DefaultThresholds()); d != Hit {
t.Errorf("decision %s, want a hit", d)
}
}
+317
View File
@@ -0,0 +1,317 @@
// Package prefix estimates which replica already holds the prefix of a request
// (plan §7.5 tier 1).
//
// It is an estimate and says so. The router does not know an engine's real
// cache; it knows what it sent where, which is a good enough proxy because
// engines cache what they are sent. Nothing here is ever authoritative — a
// wrong guess costs a prefill, not a wrong answer.
//
// Tokens are deliberately absent. Matching on bytes needs no tokenizer, and
// needing the engine's exact tokenizer is what makes tier 2 expensive enough to
// be a separate project.
package prefix
import (
"sync"
"time"
)
// Config tunes the tree.
type Config struct {
// MaxNodes bounds the tree. Past it the least recently used branches go, or
// a long-running router accumulates every prompt it has ever seen.
MaxNodes int
// MinMatch is the shortest match worth acting on, in bytes. Below it the
// overlap is coincidence — two English prompts share "The" — and treating
// coincidence as affinity is how prefix routing drowns one replica.
MinMatch int
// TTL is how long an entry is believed. An engine evicts on its own
// schedule, so an old record describes a cache that may be long gone.
TTL time.Duration
}
// DefaultConfig returns values that work without measurement.
func DefaultConfig() Config {
return Config{MaxNodes: 100_000, MinMatch: 64, TTL: 10 * time.Minute}
}
func (c *Config) normalize() {
d := DefaultConfig()
if c.MaxNodes <= 0 {
c.MaxNodes = d.MaxNodes
}
if c.MinMatch <= 0 {
c.MinMatch = d.MinMatch
}
if c.TTL <= 0 {
c.TTL = d.TTL
}
}
// seen is how often and how recently a replica was sent this prefix.
//
// The count is what makes the estimate stable. Picking the most *recent*
// replica sounds right and is not: with a randomizing picker every replica ends
// up in the node eventually, the freshest one flips with every request, and the
// tree spends its time chasing whatever the picker last did. Measured: prefix
// routing made cache hits worse than no prefix routing at all, 28.7% against
// 32.8%.
type seen struct {
count int
last time.Time
}
// Tree is a radix tree of request prefixes, per replica.
type Tree struct {
cfg Config
now func() time.Time
mu sync.Mutex
root *node
nodes int
}
type node struct {
// label is this node's slice of the prefix, shared by everything below it.
label string
children map[byte]*node
// replicas is which endpoints were sent a request passing through this
// node — not only ending at it. Recorded along the whole path, because a
// prompt that diverges partway through a label still shares everything
// before the divergence, and shared leading bytes are exactly what a prefix
// cache holds.
//
// Per replica rather than a set, so one replica's entry can expire without
// taking the others with it.
replicas map[string]*seen
// touched drives eviction.
touched time.Time
}
// New builds an empty tree.
func New(cfg Config) *Tree {
cfg.normalize()
return &Tree{cfg: cfg, now: time.Now, root: &node{children: map[byte]*node{}}, nodes: 1}
}
// Match reports the replica with the longest usable prefix in common, and how
// many bytes matched.
//
// Empty when nothing matched far enough to be worth acting on — which is a
// normal answer, not a failure, and the caller must be able to proceed on load
// alone when it comes back.
func (t *Tree) Match(prompt string) (replica string, matched int) {
if prompt == "" {
return "", 0
}
now := t.now()
t.mu.Lock()
defer t.mu.Unlock()
cur := t.root
consumed := 0
best, bestLen := "", 0
for consumed < len(prompt) {
child, ok := cur.children[prompt[consumed]]
if !ok {
break
}
shared := commonPrefix(child.label, prompt[consumed:])
if shared == 0 {
break
}
consumed += shared
cur = child
child.touched = now
// A partial label counts. Everything recorded at or below this node
// began with the bytes just matched, so a replica here holds that much
// of this prompt however the rest diverges — and a cache hit is about
// shared leading bytes, not about where two prompts stop agreeing.
//
// Requiring the whole label to match made this tree inert: the first
// request for a prompt becomes one long label, the second diverges in
// its last few bytes, and every lookup after that returned nothing.
// Measured as 144 misses and zero hits before this was fixed.
if r, ok := cur.owner(now, t.cfg.TTL); ok && consumed > bestLen {
best, bestLen = r, consumed
}
if shared < len(child.label) {
break
}
}
if bestLen < t.cfg.MinMatch {
return "", 0
}
return best, bestLen
}
// Observe records that a replica was sent this prompt.
func (t *Tree) Observe(prompt, replica string) {
if prompt == "" || replica == "" {
return
}
now := t.now()
t.mu.Lock()
defer t.mu.Unlock()
cur := t.root
rest := prompt
for len(rest) > 0 {
child, ok := cur.children[rest[0]]
if !ok {
leaf := &node{label: rest, children: map[byte]*node{}, touched: now}
cur.children[rest[0]] = leaf
t.nodes++
cur = leaf
break
}
shared := commonPrefix(child.label, rest)
if shared < len(child.label) {
// The existing label diverges partway, so it splits into a shared
// head and what used to hang below it.
tail := &node{
label: child.label[shared:], children: child.children,
replicas: child.replicas, touched: child.touched,
}
child.label = child.label[:shared]
child.children = map[byte]*node{tail.label[0]: tail}
child.replicas = nil
t.nodes++
}
child.touched = now
cur = child
// Marked on the way through, not only at the end: a later prompt that
// diverges here still shares everything above, and that is what makes
// a partial match meaningful.
cur.record(replica, now)
rest = rest[shared:]
}
cur.record(replica, now)
if t.nodes > t.cfg.MaxNodes {
t.evictLocked(now)
}
}
// Forget drops a replica from the tree, for one that has gone away.
//
// Leaving it would have the router steering traffic towards a cache that no
// longer exists, and a replica id is reused by nothing, so the entries would
// never be corrected — only expire.
func (t *Tree) Forget(replica string) {
t.mu.Lock()
defer t.mu.Unlock()
forget(t.root, replica)
}
func forget(n *node, replica string) {
delete(n.replicas, replica)
for _, c := range n.children {
forget(c, replica)
}
}
// Size reports how many nodes the tree holds, for metrics and tests.
func (t *Tree) Size() int {
t.mu.Lock()
defer t.mu.Unlock()
return t.nodes
}
// record marks that a replica was sent something passing through this node.
func (n *node) record(replica string, now time.Time) {
if n.replicas == nil {
n.replicas = map[string]*seen{}
}
if e, ok := n.replicas[replica]; ok {
e.count++
e.last = now
} else {
n.replicas[replica] = &seen{count: 1, last: now}
}
n.touched = now
}
// owner returns the replica most likely to hold this prefix.
//
// Most often sent, not most recently: a count is stable under a randomizing
// picker, and it still moves when traffic genuinely does. Ties go to the more
// recent one, which is the only thing recency is good for here.
func (n *node) owner(now time.Time, ttl time.Duration) (string, bool) {
var best string
var bestSeen *seen
for r, s := range n.replicas {
if now.Sub(s.last) > ttl {
continue
}
switch {
case bestSeen == nil,
s.count > bestSeen.count,
s.count == bestSeen.count && s.last.After(bestSeen.last):
best, bestSeen = r, s
}
}
return best, best != ""
}
// evictLocked drops the coldest half of the tree's branches.
//
// Half rather than just enough: trimming to exactly the limit means evicting
// again on the very next insert, and the walk is the expensive part.
func (t *Tree) evictLocked(now time.Time) {
cutoff := t.coldestHalf(now)
t.nodes = 1 + prune(t.root, cutoff)
}
// coldestHalf finds the touch time that splits the tree in two.
func (t *Tree) coldestHalf(now time.Time) time.Time {
var times []time.Time
collect(t.root, &times)
if len(times) == 0 {
return now
}
// A partial sort would do, but the tree is evicted rarely and this is
// simpler to be sure of.
for i := 1; i < len(times); i++ {
for j := i; j > 0 && times[j].Before(times[j-1]); j-- {
times[j], times[j-1] = times[j-1], times[j]
}
}
return times[len(times)/2]
}
func collect(n *node, out *[]time.Time) {
for _, c := range n.children {
*out = append(*out, c.touched)
collect(c, out)
}
}
// prune removes branches untouched since the cutoff, returning what remains.
func prune(n *node, cutoff time.Time) int {
kept := 0
for b, c := range n.children {
if c.touched.Before(cutoff) {
delete(n.children, b)
continue
}
kept += 1 + prune(c, cutoff)
}
return kept
}
func commonPrefix(a, b string) int {
n := min(len(a), len(b))
i := 0
for i < n && a[i] == b[i] {
i++
}
return i
}
+165
View File
@@ -0,0 +1,165 @@
package prefix
import (
"fmt"
"strings"
"sync"
"testing"
"time"
)
func tree(t *testing.T, cfg Config) *Tree {
t.Helper()
return New(cfg)
}
func TestTheLongestMatchWins(t *testing.T) {
// Two replicas have seen overlapping prefixes; the one that has seen more
// of this prompt is the one that has more of it cached.
tr := tree(t, Config{MinMatch: 4})
tr.Observe("You are a helpful assistant. ", "a")
tr.Observe("You are a helpful assistant. Answer in French. ", "b")
replica, n := tr.Match("You are a helpful assistant. Answer in French. What is 2+2?")
if replica != "b" {
t.Errorf("matched %q with %d bytes; the longer prefix belongs to b", replica, n)
}
if n < len("You are a helpful assistant. Answer in French. ") {
t.Errorf("matched only %d bytes", n)
}
}
func TestAShortOverlapIsNotAffinity(t *testing.T) {
// Two English prompts share "The". Treating coincidence as affinity is how
// prefix routing drowns one replica.
tr := tree(t, Config{MinMatch: 64})
tr.Observe("The quick brown fox", "a")
if replica, n := tr.Match("The slow green turtle"); replica != "" {
t.Errorf("a %d-byte coincidence was treated as a match to %q", n, replica)
}
}
func TestSplittingAnExistingBranch(t *testing.T) {
// The case a radix tree exists for: a new prompt diverges partway through
// a label that was one node.
tr := tree(t, Config{MinMatch: 4})
tr.Observe("abcdefgh", "a")
tr.Observe("abcdXYZW", "b")
if r, _ := tr.Match("abcdefgh and more"); r != "a" {
t.Errorf("the original branch matched %q", r)
}
if r, _ := tr.Match("abcdXYZW and more"); r != "b" {
t.Errorf("the new branch matched %q", r)
}
// The shared head belongs to both, because both were sent something
// beginning with it — and a prefix cache holds leading bytes, not whole
// prompts. An earlier version treated this as no match, which made the
// whole tree inert in practice.
if r, n := tr.Match("abcd"); r == "" || n != 4 {
t.Errorf("the shared head matched %q with %d bytes; both replicas hold it", r, n)
}
}
func TestAPartialLabelIsStillAMatch(t *testing.T) {
// The correction that made this feature work at all.
//
// The first request for a prompt becomes one long label. The second differs
// in its last few bytes — a new question after the same system prompt,
// which is the single most common shape there is. Requiring the whole label
// to match meant every lookup after the first returned nothing: measured as
// 144 misses and zero hits.
//
// What matters is the shared head, because that is what the engine cached.
tr := tree(t, Config{MinMatch: 2})
tr.Observe("abcdefghij", "a")
r, n := tr.Match("abcdeZZZZ")
if r != "a" {
t.Errorf("a partial match returned %q", r)
}
if n != 5 {
t.Errorf("matched %d bytes, want the 5 they share", n)
}
}
func TestEntriesExpire(t *testing.T) {
// An engine evicts on its own schedule, so an old record describes a cache
// that may be long gone. Believing it forever sends traffic towards
// something that is not there.
now := time.Now()
tr := New(Config{MinMatch: 4, TTL: time.Minute})
tr.now = func() time.Time { return now }
tr.Observe("shared prefix here", "a")
if r, _ := tr.Match("shared prefix here and more"); r != "a" {
t.Fatalf("fresh entry matched %q", r)
}
now = now.Add(2 * time.Minute)
if r, _ := tr.Match("shared prefix here and more"); r != "" {
t.Errorf("a stale entry still matched %q", r)
}
}
func TestForgettingAReplica(t *testing.T) {
// A replica that has gone away must stop attracting traffic. Nothing reuses
// a replica id, so the entries would otherwise only ever expire.
tr := tree(t, Config{MinMatch: 4})
tr.Observe("a long enough shared prefix", "a")
tr.Observe("a long enough shared prefix", "b")
tr.Forget("a")
r, _ := tr.Match("a long enough shared prefix and more")
if r == "a" {
t.Error("a forgotten replica still matches")
}
if r != "b" {
t.Errorf("forgetting one replica lost the other: %q", r)
}
}
func TestTheTreeIsBounded(t *testing.T) {
// A router that runs for months would otherwise accumulate every prompt it
// has ever seen.
tr := tree(t, Config{MinMatch: 4, MaxNodes: 200})
for i := range 2000 {
tr.Observe(fmt.Sprintf("prompt number %d with some trailing text", i), "a")
}
if got := tr.Size(); got > 400 {
t.Errorf("the tree grew to %d nodes against a limit of 200", got)
}
// And it still works afterwards.
tr.Observe("a fresh prompt after all that eviction", "b")
if r, _ := tr.Match("a fresh prompt after all that eviction, continued"); r != "b" {
t.Errorf("the tree stopped working after eviction: %q", r)
}
}
func TestConcurrentUse(t *testing.T) {
tr := tree(t, Config{MinMatch: 4, MaxNodes: 500})
var wg sync.WaitGroup
for i := range 32 {
wg.Add(1)
go func() {
defer wg.Done()
p := strings.Repeat("x", 40) + fmt.Sprint(i%4)
for range 50 {
tr.Observe(p, fmt.Sprintf("r%d", i%4))
tr.Match(p)
tr.Forget("r0")
}
}()
}
wg.Wait()
}
func TestEmptyInputsAreHarmless(t *testing.T) {
tr := tree(t, Config{})
tr.Observe("", "a")
tr.Observe("something", "")
if r, n := tr.Match(""); r != "" || n != 0 {
t.Errorf("an empty prompt matched %q/%d", r, n)
}
}
+52 -1
View File
@@ -23,6 +23,7 @@ import (
"gitea.mynodes.xyz/ilker/llmgw/internal/router/admission"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/scrape"
@@ -59,6 +60,9 @@ type Router struct {
acct *accounting.Accountant
recorder *accounting.Recorder
// prefix is nil unless approximate prefix routing is turned on.
prefix *scheduler.PrefixScorer
// clusterView is nil until the control plane hands it over, and stays nil
// in a router-only deployment.
clusterView web.Source
@@ -102,6 +106,27 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) {
KVCeiling: cfg.Router.Scoring.KVCeiling,
SessionLoadFactor: cfg.Router.Scoring.SessionLoadFactor,
}
// The metrics come first because the prefix scorer reports into them, and
// the decision metric is the whole difference between tuning this and
// guessing at it.
metrics := newMetrics()
var prefixScorer *scheduler.PrefixScorer
if cfg.Router.Prefix.Enabled {
prefixScorer = scheduler.NewPrefixScorer(
prefix.New(prefix.Config{
MaxNodes: cfg.Router.Prefix.MaxNodes,
MinMatch: cfg.Router.Prefix.MinMatch,
TTL: cfg.Router.Prefix.TTL,
}),
prefix.Thresholds{
MatchRatio: cfg.Router.Prefix.MatchRatio,
BalanceAbs: cfg.Router.Prefix.BalanceAbs,
BalanceRel: cfg.Router.Prefix.BalanceRel,
},
metrics.observePrefix)
sc.Prefix = prefixScorer
}
profiles := map[string]*scheduler.Profile{}
for _, name := range []string{"chat", "batch", "long-context", "debug"} {
profiles[name] = scheduler.Builtin(name, sc)
@@ -125,7 +150,8 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) {
reg: reg,
profiles: profiles,
models: models,
m: newMetrics(),
prefix: prefixScorer,
m: metrics,
px: proxy.New(proxy.Config{
DialTimeout: cfg.Router.Timeouts.Dial,
ResponseHeaderTimeout: cfg.Router.Timeouts.ResponseHeader,
@@ -297,6 +323,23 @@ func (r *Router) Registry() *endpoint.Registry { return r.reg }
// — so a replica that was already serving does not start from cold just
// because another one moved.
func (r *Router) SetEndpoints(ctx context.Context, infos []endpoint.Info) {
if r.prefix != nil {
// A replica that has gone takes its cache with it. Leaving its entries
// in the tree would steer traffic towards something that is not there,
// and since nothing reuses a replica id the entries would never be
// corrected — only expire, minutes later.
before := map[string]struct{}{}
for _, s := range r.reg.Snapshots("") {
before[s.ID] = struct{}{}
}
for _, i := range infos {
delete(before, i.ID)
}
for id := range before {
r.prefix.Forget(id)
}
}
r.reg.Replace(infos)
r.poller.Sync(ctx)
}
@@ -425,6 +468,14 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa
res := dispatchResult{endpointID: dec.Endpoint.ID, pool: dec.Endpoint.Pool, outcome: out, err: perr}
r.m.observeScores(dec)
// Recorded after the decision, not during it: the scorer proposes and
// the picker disposes, so recording the proposal would fill the tree
// with replicas that were never sent anything — the one thing that
// would make the estimate worse than no estimate at all.
if r.prefix != nil {
r.prefix.Observe(req.PromptPrefix, dec.Endpoint.ID)
}
if perr == nil || out.Wrote || !out.Retryable {
return res
}
+65
View File
@@ -0,0 +1,65 @@
package scheduler
import (
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix"
)
// PrefixScorer steers a request towards the replica that has probably already
// prefilled its prefix (plan §7.5 tier 1).
//
// A scorer rather than a filter, deliberately. A filter would make affinity a
// rule, and a rule that sends every request carrying a popular system prompt to
// one replica is how prefix routing produces a beautiful cache hit rate and an
// unusable p99. As a scorer it is one voice among several, and the queue and
// inflight scorers can outvote it — on top of the imbalance guard inside the
// lookup itself, which stops it speaking at all when load has gone lopsided.
type PrefixScorer struct {
tree *prefix.Tree
thresholds prefix.Thresholds
// observe reports which branch the decision took, for the metric without
// which tuning this is guesswork.
observe func(prefix.Decision)
}
// NewPrefixScorer builds the scorer. A nil tree makes it inert, which is what
// a model with prefix routing turned off should get.
func NewPrefixScorer(t *prefix.Tree, th prefix.Thresholds, observe func(prefix.Decision)) *PrefixScorer {
if observe == nil {
observe = func(prefix.Decision) {}
}
return &PrefixScorer{tree: t, thresholds: th, observe: observe}
}
// Name identifies the scorer in weights and metrics.
func (*PrefixScorer) Name() string { return "prefix" }
// Score gives the matched replica 1 and everything else 0.
func (p *PrefixScorer) Score(c *Ctx, eps []endpoint.Snapshot) []float64 {
out := make([]float64, len(eps))
idx, decision := prefix.Choose(p.tree, c.Req.PromptPrefix, eps, p.thresholds)
p.observe(decision)
if idx >= 0 {
out[idx] = 1
}
return out
}
// Observe records where a request was actually sent.
//
// After the decision rather than during it: the scorer proposes, the picker
// disposes, and recording the proposal would fill the tree with replicas that
// were never sent anything — which is the one thing that would make the
// estimate worse than no estimate at all.
func (p *PrefixScorer) Observe(promptPrefix, replicaID string) {
if p.tree != nil {
p.tree.Observe(promptPrefix, replicaID)
}
}
// Forget drops a replica that has gone away.
func (p *PrefixScorer) Forget(replicaID string) {
if p.tree != nil {
p.tree.Forget(replicaID)
}
}
+24 -4
View File
@@ -10,6 +10,11 @@ type ProfileConfig struct {
KVCeiling float64
// SessionLoadFactor is the load ceiling multiplier for session affinity.
SessionLoadFactor float64
// Prefix, when set, adds the approximate prefix tree to the profiles that
// benefit from it. Nil leaves them exactly as they were — session affinity
// alone, which is tier 0 and already most of the win.
Prefix *PrefixScorer
}
// DefaultProfileConfig returns sensible starting values.
@@ -46,17 +51,30 @@ func Chat(c ProfileConfig) *Profile {
return &Profile{
Name: "chat",
Filters: []Filter{BreakerOpen(), PoolPriority()},
Scorers: []WeightedScorer{
Scorers: c.withPrefix([]WeightedScorer{
{Session(c.SessionLoadFactor), 3.0},
{Queue(c.TargetWaiting), 1.0},
{KV(c.KVCeiling), 1.0},
{Inflight(c.TargetInflight), 1.0},
{Weight(), 0.5},
},
}, 2.0),
Picker: P2C(),
}
}
// withPrefix appends the prefix scorer when one is configured.
//
// Weighted below session affinity on purpose. A session id is a statement from
// the client that these requests belong together; a prefix match is the
// router's guess that they might. When the two disagree the client is more
// likely to be right.
func (c ProfileConfig) withPrefix(scorers []WeightedScorer, weight float64) []WeightedScorer {
if c.Prefix == nil {
return scorers
}
return append(scorers, WeightedScorer{c.Prefix, weight})
}
// Batch is the profile for bulk work: throughput matters, TTFT does not, so
// finding the shortest queue outranks stickiness.
func Batch(c ProfileConfig) *Profile {
@@ -82,12 +100,14 @@ func LongContext(c ProfileConfig) *Profile {
return &Profile{
Name: "long-context",
Filters: []Filter{BreakerOpen(), PoolPriority()},
Scorers: []WeightedScorer{
// Prefix matching earns more here than anywhere else: a long prompt is
// mostly prefill, and prefill is what a cache hit skips.
Scorers: c.withPrefix([]WeightedScorer{
{KV(c.KVCeiling), 3.0},
{Session(c.SessionLoadFactor), 2.0},
{Queue(c.TargetWaiting), 1.0},
{Inflight(c.TargetInflight), 1.0},
},
}, 3.0),
Picker: Softmax(0.2),
}
}
+5 -1
View File
@@ -1424,7 +1424,11 @@ The launch path is now built too:
no identity to consult, which let any client promote itself.
- [ ] Budget fallback
- [ ] Routing profiles from configuration plus the `x-llmgw-profile` header
- [ ] Approximate radix tree (§7.5 tier 1), with its thresholds and decision metrics
- [x] Approximate radix tree (§7.5 tier 1), with its thresholds and the
`llmgw_prefix_decisions_total` breakdown. Measured against sim-engines
whose caches were too small for the working set: 22.5% → 81.7% cache
hits. It does nothing when the working set fits in one replica's cache,
which is worth knowing rather than assuming.
- [ ] Fallback models and pool spillover
- [x] `ctl bench`: a load ramp that reports capacity as the last concurrency
level meeting the SLO, plus the queue depth and KV usage observed there,
+18
View File
@@ -68,12 +68,27 @@ func withLoadedEngines() func(*config.Server) {
return nil // marker; see start()
}
// tightPrefixCache shrinks the engines' prefix cache so the working set does
// not fit in any one of them.
//
// That condition is the whole reason tier 1 exists. With a cache large enough
// for every prompt, scattering is free — every replica ends up holding
// everything — and prefix routing can only add constraint. It is worth knowing
// the feature does nothing there rather than assuming it helps everywhere.
var tightPrefixCache = func(c *sim.Config) { c.PrefixCacheBlocks = 96 }
// start brings up n sim-engines behind a router.
//
// A nil tune function is the "small engines" marker, which keeps the call sites
// readable without a second variadic parameter.
func start(t *testing.T, n int, tune ...func(*config.Server)) *harness {
t.Helper()
return startTuned(t, n, nil, tune...)
}
// startTuned is start with the sim-engines themselves adjusted.
func startTuned(t *testing.T, n int, engineOpts []engineOpt, tune ...func(*config.Server)) *harness {
t.Helper()
small := false
kept := tune[:0]
@@ -99,6 +114,9 @@ func start(t *testing.T, n int, tune ...func(*config.Server)) *harness {
sc.DecodeTokensPerSec = 800
sc.BatchContention = 0.005
}
for _, opt := range engineOpts {
opt(&sc)
}
e := sim.New(sc)
srv := httptest.NewServer(sim.NewServer(e, slog.New(slog.DiscardHandler)))
h.engines = append(h.engines, engine{eng: e, srv: srv})
+211
View File
@@ -0,0 +1,211 @@
package e2e
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"sync"
"testing"
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
)
// tenants is how many distinct system prompts the load uses.
//
// Enough that the working set does not fit in one replica's cache, which is the
// condition tier 1 exists for. With a handful of prompts and a large cache,
// scattering is free — every replica ends up holding everything — and prefix
// routing can only add constraint. Measured that way it made cache hits worse,
// which is a true statement about a scenario the feature is not for.
const tenants = 24
func systemPrompt(i int) string {
return fmt.Sprintf("SYSTEM PROMPT %d. %s ", i,
strings.Repeat("You are a careful assistant with detailed instructions. ", 12))
}
// askAll sends one request per tenant, repeated, and reports how many prompt
// tokens the engines served from cache.
func askAll(t *testing.T, h *harness, rounds int) (cached, prompt int64) {
t.Helper()
var mu sync.Mutex
var wg sync.WaitGroup
for round := range rounds {
for tenant := range tenants {
wg.Add(1)
go func() {
defer wg.Done()
body, _ := json.Marshal(map[string]any{
"model": testModel,
"max_tokens": 4,
"messages": []map[string]string{
{"role": "system", "content": systemPrompt(tenant)},
{"role": "user", "content": fmt.Sprintf("question %d", round)},
},
})
req, err := http.NewRequestWithContext(context.Background(), http.MethodPost,
h.router.URL+"/v1/chat/completions", bytes.NewReader(body))
if err != nil {
return
}
req.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(req)
if err != nil {
return
}
defer resp.Body.Close()
raw, _ := io.ReadAll(resp.Body)
var out struct {
Usage struct {
PromptTokens int `json:"prompt_tokens"`
PromptTokensDetails struct {
CachedTokens int `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
} `json:"usage"`
}
if json.Unmarshal(raw, &out) != nil {
return
}
mu.Lock()
cached += int64(out.Usage.PromptTokensDetails.CachedTokens)
prompt += int64(out.Usage.PromptTokens)
mu.Unlock()
}()
}
wg.Wait()
}
return cached, prompt
}
// TestPrefixRoutingRaisesCacheHits is the measurement the tier exists for.
//
// Many tenants with distinct system prompts against four replicas whose caches
// are too small to hold all of them. Without prefix awareness a tenant's
// requests scatter, each replica re-prefills the same prompt, and they evict
// each other's; with it a tenant converges on the replica that already has it.
func TestPrefixRoutingRaisesCacheHits(t *testing.T) {
tight := []engineOpt{tightPrefixCache}
withoutPrefix := startTuned(t, 4, tight, func(c *config.Server) {
c.Router.DefaultProfile = "long-context"
})
baseCached, basePrompt := askAll(t, withoutPrefix, 6)
withPrefix := startTuned(t, 4, tight, func(c *config.Server) {
c.Router.DefaultProfile = "long-context"
c.Router.Prefix = config.Prefix{
Enabled: true, MinMatch: 32, MatchRatio: 0.2,
BalanceAbs: 8, BalanceRel: 1.5,
}
})
prefixCached, prefixPrompt := askAll(t, withPrefix, 6)
if basePrompt == 0 || prefixPrompt == 0 {
t.Fatal("no requests completed")
}
baseRate := float64(baseCached) / float64(basePrompt)
prefixRate := float64(prefixCached) / float64(prefixPrompt)
t.Logf("cache hit rate: %.1f%% without prefix routing, %.1f%% with",
baseRate*100, prefixRate*100)
t.Logf("prefix decisions: %v", prefixDecisions(t, withPrefix))
if prefixRate <= baseRate {
t.Errorf("prefix routing did not improve cache hits: %.3f vs %.3f", prefixRate, baseRate)
}
}
// TestPrefixRoutingDoesNotDrownOneReplica is the other half, and the reason the
// guard exists. A cache hit rate bought by sending everything to one replica is
// not a win.
func TestPrefixRoutingDoesNotDrownOneReplica(t *testing.T) {
h := start(t, 4, func(c *config.Server) {
c.Router.DefaultProfile = "long-context"
c.Router.Prefix = config.Prefix{
Enabled: true, MinMatch: 32, MatchRatio: 0.2,
BalanceAbs: 8, BalanceRel: 1.5,
}
})
// One tenant, hammered: the shape that makes affinity dangerous.
var wg sync.WaitGroup
for range 60 {
wg.Add(1)
go func() {
defer wg.Done()
body, _ := json.Marshal(map[string]any{
"model": testModel,
"max_tokens": 8,
"messages": []map[string]string{
{"role": "system", "content": systemPrompt(0)},
{"role": "user", "content": "go"},
},
})
req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost,
h.router.URL+"/v1/chat/completions", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(req)
if err != nil {
return
}
_, _ = io.Copy(io.Discard, resp.Body)
resp.Body.Close()
}()
}
wg.Wait()
// Every replica should have done some of it. PrefixQueries counts requests
// the engine looked up, which is one per request it was given.
var served int
for _, e := range h.engines {
if e.eng.Stats().PrefixQueries > 0 {
served++
}
}
if served < 2 {
t.Errorf("only %d of %d replicas served anything; affinity drowned the rest",
served, len(h.engines))
}
}
// prefixDecisions reads the router's own count of which branch each prefix
// lookup took. This is the metric that exists so tuning is not guesswork.
func prefixDecisions(t *testing.T, h *harness) map[string]int {
t.Helper()
mux, err := h.rt.AdminHandler()
if err != nil {
t.Fatal(err)
}
srv := httptest.NewServer(mux)
defer srv.Close()
req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, srv.URL+"/metrics", nil)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
body, _ := io.ReadAll(resp.Body)
out := map[string]int{}
for _, line := range strings.Split(string(body), "\n") {
if !strings.HasPrefix(line, "llmgw_prefix_decisions_total{") {
continue
}
name := line[strings.Index(line, `"`)+1 : strings.LastIndex(line, `"`)]
fields := strings.Fields(line)
n, _ := strconv.Atoi(strings.TrimSuffix(fields[len(fields)-1], ".0"))
if f, err := strconv.ParseFloat(fields[len(fields)-1], 64); err == nil {
n = int(f)
}
out[name] = n
}
return out
}