llmgw/internal/router/scheduler/prefix.go

package scheduler

import (
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/kvindex"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix"
)

// KVQuery is the precise index's answer for one request, worked out before
// scheduling because it needs a round trip the scorer must not make.
type KVQuery struct {
	// Fallback, when set, is why the index could not be asked; tier 1
	// decides instead.
	Fallback kvindex.Decision
	// Matched is how many of the prompt's leading blocks each replica the
	// index knows holds, and Total how many blocks the prompt has at that
	// replica's block size. A replica absent from both is one the index has
	// not heard from, and is scored as holding nothing — an under-estimate,
	// which costs a missed preference and never a wrong one.
	Matched map[string]int
	Total   map[string]int
}

// PrefixScorer steers a request towards the replica that has already
// prefilled its prefix: from the precise KV index where the model uses it and
// it has an answer (plan §7.5 tier 2), and otherwise from the approximate tree
// (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)

	// observeKV and minBlocks are the precise index's: its decision metric,
	// and the fewest leading blocks worth preferring a replica for.
	observeKV func(kvindex.Decision)
	minBlocks int
}

// NewPrefixScorer builds the scorer. A nil tree turns tier 1 off: requests
// the precise index cannot answer are then left to load, and nothing is
// recorded for them.
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,
		observeKV: func(kvindex.Decision) {}, minBlocks: 1}
}

// WithKV sets what the precise index needs: the fewest leading blocks a
// replica must hold to be preferred (below one means one), and where its
// decisions are counted.
func (p *PrefixScorer) WithKV(minBlocks int, observe func(kvindex.Decision)) *PrefixScorer {
	if observe != nil {
		p.observeKV = observe
	}
	p.minBlocks = max(minBlocks, 1)
	return p
}

// Name identifies the scorer in weights and metrics.
func (*PrefixScorer) Name() string { return "prefix" }

// Score rates each candidate.
//
// From the index: the share of the prompt's blocks the replica holds, so two
// replicas holding a shared system prompt score alike and the one also
// holding this conversation scores higher. From the tree: 1 for the replica
// it names and 0 for the rest, because a guess has no degrees.
func (p *PrefixScorer) Score(c *Ctx, eps []endpoint.Snapshot) []float64 {
	out := make([]float64, len(eps))
	if q := c.Req.KV; q != nil {
		if q.Fallback == "" {
			p.observeKV(p.scoreKV(q, eps, out))
			return out
		}
		p.observeKV(q.Fallback)
	}
	if p.tree == nil {
		return out
	}
	idx, decision := prefix.Choose(p.tree, c.Req.PromptPrefix, eps, p.thresholds)
	p.observe(decision)
	if idx >= 0 {
		out[idx] = 1
	}
	return out
}

// scoreKV fills out from the index's answer and says which branch it took.
func (p *PrefixScorer) scoreKV(q *KVQuery, eps []endpoint.Snapshot, out []float64) kvindex.Decision {
	if len(eps) == 0 {
		return kvindex.MissMinLoad
	}
	// The guard first, as in tier 1: a certain hit on a replica that cannot
	// take the request quickly is still the wrong place for it.
	if prefix.Imbalanced(eps, p.thresholds) {
		return kvindex.LoadOverride
	}
	best := 0
	for i, e := range eps {
		m, t := q.Matched[e.ID], q.Total[e.ID]
		if t > 0 {
			out[i] = float64(m) / float64(t)
		}
		best = max(best, m)
	}
	if best < p.minBlocks {
		clear(out)
		return kvindex.MissMinLoad
	}
	return kvindex.Hit
}

// 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. The precise index needs no such
// record; the engine reports what it cached.
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)
	}
}