feat(prefixcache): index reported KV residency

Add a NATS event contract and exact-residency provider for backend KV cache reports. Keep guessed request observations as the default routing source while maintaining the reported index for future producers.

Assisted-by: Codex:gpt-5

Author: localai-org-maint-bot

core/application/distributed.go

@@ -319,6 +319,14 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
 			return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheInvalidate, err)
 		}
 
+		// Keep an exact-residency index current so backend producers can report
+		// their real KV state without coupling to router internals. Routing stays
+		// on the guessed provider until a backend producer is available.
+		reportedIndex := prefixcache.NewReportedIndex()
+		if _, err := messaging.SubscribeJSON(natsClient, messaging.SubjectPrefixCacheResidency, reportedIndex.Apply); err != nil {
+			return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheResidency, err)
+		}
+
 		// Background eviction: sweep idle entries on the app context. Stopped
 		// when the app context is cancelled (mirrors the reconciler loop which
 		// also runs on options.Context). TTL/2 keeps stale entries from

core/services/messaging/subjects.go

@@ -478,8 +478,28 @@ const subjectSyncStatePrefix = "state."
 const (
 	SubjectPrefixCacheObserve    = "prefixcache.observe"
 	SubjectPrefixCacheInvalidate = "prefixcache.invalidate"
+	SubjectPrefixCacheResidency  = "prefixcache.residency"
 )
 
+// PrefixCacheOperation describes a backend-reported KV-cache residency change.
+type PrefixCacheOperation string
+
+const (
+	PrefixCacheStore  PrefixCacheOperation = "store"
+	PrefixCacheRemove PrefixCacheOperation = "remove"
+	PrefixCacheClear  PrefixCacheOperation = "clear"
+)
+
+// PrefixCacheResidencyEvent reports exact backend KV-cache residency. Chain is
+// the compatible shallow-to-deep prefix hash chain used by the router.
+type PrefixCacheResidencyEvent struct {
+	Operation PrefixCacheOperation `json:"operation"`
+	Model     string               `json:"model"`
+	NodeID    string               `json:"node_id"`
+	Replica   int                  `json:"replica"`
+	Chain     []uint64             `json:"chain,omitempty"`
+}
+
 // PrefixCacheObserveEvent announces that the replica (NodeID, Replica) served a
 // request whose prefix chain ends at the given hashes for model. Chain is the
 // full shallow-to-deep hash chain so peers can insert the same path. Affinity is

core/services/messaging/subjects_prefixcache_test.go

@@ -11,6 +11,15 @@ var _ = Describe("PrefixCache subjects", func() {
 	It("exposes stable subject constants", func() {
 		Expect(messaging.SubjectPrefixCacheObserve).To(Equal("prefixcache.observe"))
 		Expect(messaging.SubjectPrefixCacheInvalidate).To(Equal("prefixcache.invalidate"))
+		Expect(messaging.SubjectPrefixCacheResidency).To(Equal("prefixcache.residency"))
+	})
+
+	It("defines the reported residency event operations", func() {
+		Expect([]messaging.PrefixCacheOperation{
+			messaging.PrefixCacheStore,
+			messaging.PrefixCacheRemove,
+			messaging.PrefixCacheClear,
+		}).To(Equal([]messaging.PrefixCacheOperation{"store", "remove", "clear"}))
 	})
 
 	It("carries a replica index on the observe event", func() {

core/services/nodes/prefixcache/reported.go

@@ -0,0 +1,133 @@
+package prefixcache
+
+import (
+	"sort"
+	"sync"
+	"time"
+
+	"github.com/mudler/LocalAI/core/services/messaging"
+)
+
+// ReportedIndex is an exact-residency Provider populated only by backend
+// events. Request routing observations intentionally do not mutate it.
+type ReportedIndex struct {
+	mu          sync.RWMutex
+	residencies map[string]map[ReplicaKey][][]uint64
+}
+
+func NewReportedIndex() *ReportedIndex {
+	return &ReportedIndex{residencies: map[string]map[ReplicaKey][][]uint64{}}
+}
+
+func (ix *ReportedIndex) Apply(event messaging.PrefixCacheResidencyEvent) {
+	key := ReplicaKey{NodeID: event.NodeID, Replica: event.Replica}
+	if event.Model == "" || key.NodeID == "" {
+		return
+	}
+
+	ix.mu.Lock()
+	defer ix.mu.Unlock()
+	byReplica := ix.residencies[event.Model]
+	switch event.Operation {
+	case messaging.PrefixCacheStore:
+		if len(event.Chain) == 0 {
+			return
+		}
+		if byReplica == nil {
+			byReplica = map[ReplicaKey][][]uint64{}
+			ix.residencies[event.Model] = byReplica
+		}
+		for _, chain := range byReplica[key] {
+			if equalChain(chain, event.Chain) {
+				return
+			}
+		}
+		byReplica[key] = append(byReplica[key], append([]uint64(nil), event.Chain...))
+	case messaging.PrefixCacheRemove:
+		if byReplica == nil || len(event.Chain) == 0 {
+			return
+		}
+		chains := byReplica[key]
+		for i, chain := range chains {
+			if equalChain(chain, event.Chain) {
+				chains = append(chains[:i], chains[i+1:]...)
+				break
+			}
+		}
+		if len(chains) == 0 {
+			delete(byReplica, key)
+		} else {
+			byReplica[key] = chains
+		}
+	case messaging.PrefixCacheClear:
+		delete(byReplica, key)
+	}
+}
+
+func (ix *ReportedIndex) Decide(model string, chain []uint64, candidates []ReplicaKey, _ time.Time) PrefixDecision {
+	order := append([]ReplicaKey(nil), candidates...)
+	sort.Slice(order, func(i, j int) bool { return order[i].less(order[j]) })
+	d := PrefixDecision{ColdOrder: order}
+	if len(chain) == 0 {
+		return d
+	}
+
+	ix.mu.RLock()
+	defer ix.mu.RUnlock()
+	byReplica := ix.residencies[model]
+	bestDepth := 0
+	for _, key := range order {
+		for _, reported := range byReplica[key] {
+			depth := commonDepth(chain, reported)
+			if depth > bestDepth {
+				bestDepth = depth
+				d.Hot = key
+				d.HasHot = true
+			}
+		}
+	}
+	if d.HasHot {
+		d.MatchRatio = float64(bestDepth) / float64(len(chain))
+	}
+	return d
+}
+
+func (ix *ReportedIndex) Observe(string, []uint64, ReplicaKey, time.Time) bool { return false }
+
+func (ix *ReportedIndex) Invalidate(model string, key ReplicaKey) {
+	ix.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheClear, Model: model, NodeID: key.NodeID, Replica: key.Replica})
+}
+
+func (ix *ReportedIndex) InvalidateNode(model, nodeID string) {
+	ix.mu.Lock()
+	defer ix.mu.Unlock()
+	for key := range ix.residencies[model] {
+		if key.NodeID == nodeID {
+			delete(ix.residencies[model], key)
+		}
+	}
+}
+
+func (ix *ReportedIndex) Evict(time.Time) {}
+
+func equalChain(a, b []uint64) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	for i := range a {
+		if a[i] != b[i] {
+			return false
+		}
+	}
+	return true
+}
+
+func commonDepth(a, b []uint64) int {
+	depth := min(len(a), len(b))
+	for i := range depth {
+		if a[i] != b[i] {
+			return i
+		}
+	}
+	return depth
+}

core/services/nodes/prefixcache/reported_test.go

@@ -0,0 +1,53 @@
+package prefixcache_test
+
+import (
+	"time"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/mudler/LocalAI/core/services/messaging"
+	"github.com/mudler/LocalAI/core/services/nodes/prefixcache"
+)
+
+var _ prefixcache.Provider = (*prefixcache.ReportedIndex)(nil)
+
+var _ = Describe("ReportedIndex", func() {
+	var idx *prefixcache.ReportedIndex
+
+	BeforeEach(func() { idx = prefixcache.NewReportedIndex() })
+
+	It("routes to the replica with the longest reported prefix", func() {
+		idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheStore, Model: "m", NodeID: "A", Replica: 0, Chain: []uint64{1, 2}})
+		idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheStore, Model: "m", NodeID: "B", Replica: 0, Chain: []uint64{1, 2, 3, 4}})
+		d := idx.Decide("m", []uint64{1, 2, 3, 9}, []prefixcache.ReplicaKey{rk("A", 0), rk("B", 0)}, t0)
+		Expect(d.HasHot).To(BeTrue())
+		Expect(d.Hot).To(Equal(rk("B", 0)))
+		Expect(d.MatchRatio).To(Equal(0.75))
+	})
+
+	It("removes only the announced residency", func() {
+		for _, chain := range [][]uint64{{1, 2}, {7, 8}} {
+			idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheStore, Model: "m", NodeID: "A", Replica: 0, Chain: chain})
+		}
+		idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheRemove, Model: "m", NodeID: "A", Replica: 0, Chain: []uint64{1, 2}})
+		Expect(idx.Decide("m", []uint64{1, 2}, []prefixcache.ReplicaKey{rk("A", 0)}, t0).HasHot).To(BeFalse())
+		Expect(idx.Decide("m", []uint64{7, 8}, []prefixcache.ReplicaKey{rk("A", 0)}, t0).HasHot).To(BeTrue())
+	})
+
+	It("clears every residency for one replica only", func() {
+		idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheStore, Model: "m", NodeID: "A", Replica: 0, Chain: []uint64{1, 2}})
+		idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheStore, Model: "m", NodeID: "B", Replica: 0, Chain: []uint64{3, 4}})
+		idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheClear, Model: "m", NodeID: "A", Replica: 0})
+		candidates := []prefixcache.ReplicaKey{rk("A", 0), rk("B", 0)}
+		Expect(idx.Decide("m", []uint64{1, 2}, candidates, t0).HasHot).To(BeFalse())
+		Expect(idx.Decide("m", []uint64{3, 4}, candidates, t0).Hot).To(Equal(rk("B", 0)))
+	})
+
+	It("ignores guessed request observations and keeps cold ordering deterministic", func() {
+		Expect(idx.Observe("m", []uint64{1, 2}, rk("A", 0), t0)).To(BeFalse())
+		d := idx.Decide("m", []uint64{1, 2}, []prefixcache.ReplicaKey{rk("B", 1), rk("A", 1), rk("A", 0)}, time.Now())
+		Expect(d.HasHot).To(BeFalse())
+		Expect(d.ColdOrder).To(Equal([]prefixcache.ReplicaKey{rk("A", 0), rk("A", 1), rk("B", 1)}))
+	})
+})

docs/content/features/distributed-mode.md

@@ -948,6 +948,28 @@ Notes:
 
 The scheduling algorithm above is load-based (least in-flight, then least-recently-used). Work is underway to make routing **prefix-cache-aware**: bias each request toward the replica that already holds the relevant KV/prefix cache (multi-turn conversations and shared system prompts), so backends reuse cache instead of recomputing it. The first step is a router-side radix tree of prompt-prefix hashes mapped to nodes, with longest-prefix match, a load guard that preserves round-robin behavior under imbalance, and NATS sync across frontends. It is purely a routing-layer hint (no backend changes) and never routes worse than today's round-robin.
 
+Backends can report exact KV-cache residency on the `prefixcache.residency`
+NATS subject. The JSON event contract is:
+
+```json
+{
+  "operation": "store",
+  "model": "model-name",
+  "node_id": "worker-id",
+  "replica": 0,
+  "chain": [1203053429005847826, 15485907386658061715]
+}
+```
+
+`operation` is `store`, `remove`, or `clear`. `store` adds the announced
+shallow-to-deep chain for one model replica, `remove` removes only that exact
+announced chain, and `clear` removes all reported residency for that model
+replica (and may omit `chain`). Producers must generate the chain with exactly
+the same windowing and hashing algorithm as the router; hashes from a different
+chain algorithm are not compatible and will never match requests correctly.
+Reported events populate the exact-residency provider, but the guessed provider
+remains the routing default until a backend producer is available.
+
 Further enhancements, surfaced from a survey of SGLang, vLLM production-stack, Ray Serve, llm-d, AIBrix, and NVIDIA Dynamo, are tracked under the routing roadmap epic ([#10063](https://github.com/mudler/LocalAI/issues/10063)):
 
 - **Reported/precise KV-event mode** ([#10064](https://github.com/mudler/LocalAI/issues/10064)): subscribe to actual backend KV-cache events for exact residency instead of inferring it from routing history.