diff --git a/cmd/epp/runner/runner.go b/cmd/epp/runner/runner.go index 9f210bf66f..4d6044b95e 100644 --- a/cmd/epp/runner/runner.go +++ b/cmd/epp/runner/runner.go @@ -83,6 +83,7 @@ import ( "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/admitter/latencyslo" "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/admitter/probabilisticadmitter" reqdataprodprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix" + "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix" "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload" mmproducer "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal" preciseproducer "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache" @@ -520,6 +521,7 @@ func (r *Runner) registerInTreePlugins() { fwkplugin.Register(nohitlru.NoHitLRUType, nohitlru.Factory) fwkplugin.Register(activerequest.ActiveRequestType, activerequest.Factory) fwkplugin.Register(preciseprefixcache.PrecisePrefixCachePluginType, preciseprefixcache.PluginFactory) + fwkplugin.Register(burstprefix.PluginType, burstprefix.Factory) fwkplugin.Register(mmcacheaffinity.Type, mmcacheaffinity.Factory) fwkplugin.Register(preciseproducer.PluginType, preciseproducer.PluginFactory) diff --git a/deploy/config/epp-burst-prefix-cache-config.yaml b/deploy/config/epp-burst-prefix-cache-config.yaml new file mode 100644 index 0000000000..68dbdcddca --- /dev/null +++ b/deploy/config/epp-burst-prefix-cache-config.yaml @@ -0,0 +1,35 @@ +# Sample EPP config: burst prefix cache pipeline +# token-producer -> burst-prefix-cache-producer -> prefix-cache-scorer +# +# Co-locates bursts of prompt-sharing requests (e.g. RL rollout group samples) +# onto shared replicas so a shared prefix is prefilled once instead of scattered +# across replicas on a cold cache. +apiVersion: llm-d.ai/v1alpha1 +kind: EndpointPickerConfig +plugins: + - type: token-producer + parameters: + modelName: hf-repo/model-name # set to the model used in the vLLM deployment + vllm: + url: http://..svc.cluster.local:8000 + - type: single-profile-handler + - type: decode-filter + - type: burst-prefix-cache-producer + parameters: + windowDurationMs: 100 # batch window T + maxPerReplica: -1 # k; -1 sends all samples of a group to one replica + minColocateBlocks: 0 # >0 co-locates distinct groups sharing >= N leading blocks; 0 = load-balanced + - type: prefix-cache-scorer + parameters: + prefixMatchInfoProducerName: burst-prefix-cache-producer + - type: load-aware-scorer + - type: max-score-picker +schedulingProfiles: + - name: default + plugins: + - pluginRef: decode-filter + - pluginRef: prefix-cache-scorer + weight: 2.0 + - pluginRef: load-aware-scorer + weight: 1.0 + - pluginRef: max-score-picker diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/README.md b/pkg/epp/framework/plugins/requestcontrol/dataproducer/README.md index 34388bc1c8..9dca539091 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/README.md +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/README.md @@ -17,6 +17,7 @@ Producers may also implement additional lifecycle hooks: | `token-producer` | [`tokenizer`](tokenizer/) | `TokenizedPrompt` | Tokenizes the request prompt via vLLM `/render`; required by precise-prefix-cache-producer and context-length-aware scorers. | | `approx-prefix-cache-producer` | [`approximateprefix`](approximateprefix/) | `PrefixCacheMatchInfo` | Hashes the prompt into blocks and matches against a per-pod LRU index for approximate prefix-cache affinity. | | `precise-prefix-cache-producer` | [`preciseprefixcache`](preciseprefixcache/) | `PrefixCacheMatchInfo` | Maintains a precise KV-block index by subscribing to vLLM KV-events; requires `token-producer` upstream. | +| `burst-prefix-cache-producer` | [`burstprefix`](burstprefix/) | `PrefixCacheMatchInfo` | Batches requests within a time window and co-locates prompt-sharing samples (e.g. RL rollout groups) onto shared replicas; requires `token-producer` upstream. | | `inflight-load-producer` | [`inflightload`](inflightload/) | `InFlightLoad` | Tracks real-time in-flight request and token counts per endpoint across the full request lifecycle. | | `predicted-latency-producer` | [`predictedlatency`](predictedlatency/) | `LatencyPredictionInfo` | Trains XGBoost models via a sidecar and generates per-endpoint TTFT/TPOT predictions. | | `session-id-producer` | [`sessionid`](sessionid/) | `SessionID` | Extracts a session identifier from a request header or cookie and publishes it for affinity-aware plugins. | @@ -27,6 +28,7 @@ Producers may also implement additional lifecycle hooks: The framework resolves a DAG from each plugin's `Produces` and `Consumes` declarations and runs producers in dependency order. Explicit dependencies to be aware of: - `precise-prefix-cache-producer` **requires** `token-producer` upstream (it consumes `TokenizedPrompt`). +- `burst-prefix-cache-producer` **requires** `token-producer` upstream (it consumes `TokenizedPrompt`). - `mm-embeddings-cache-producer` **optionally** consumes `TokenizedPrompt`; configure `token-producer` first when multimodal features need tokenizer-derived hashes. - `inflight-load-producer` **optionally** consumes `PrefixCacheMatchInfo` from an approx or precise prefix producer; prefix-discounting is applied automatically when the attribute is present. - `predicted-latency-producer` **optionally** consumes `PrefixCacheMatchInfo`; set `prefixMatchInfoProducerName` in its config to the name of the prefix producer instance. @@ -35,6 +37,7 @@ The framework resolves a DAG from each plugin's `Produces` and `Consumes` declar - [Approximate Prefix Cache Producer](approximateprefix/README.md) - [Precise Prefix Cache Producer](preciseprefixcache/README.md) +- [Burst Prefix Cache Producer](burstprefix/README.md) - [Token Producer](tokenizer/README.md) - [In-Flight Load Producer](inflightload/README.md) - [Predicted Latency Producer](predictedlatency/README.md) diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go index 5b7f59694c..ac6a719bb7 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go @@ -32,6 +32,7 @@ import ( fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix" approxprefixconstants "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/constants" + "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash" tokenproducer "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer" ) @@ -179,7 +180,7 @@ func (p *dataProducer) Produce(ctx context.Context, request *fwksched.InferenceR if p.config.MaxPrefixTokensToMatch > 0 && blockSize > 0 { maxBlocks = p.config.MaxPrefixTokensToMatch / blockSize } - perPromptHashes := getBlockHashes(ctx, request, blockSize, maxBlocks) + perPromptHashes := prefixhash.GetBlockHashes(ctx, request, blockSize, maxBlocks) prefixCacheServers := make(map[ServerID]int) totalBlocks := 0 diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go index 35b73ad82e..4de817df5f 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go @@ -32,6 +32,7 @@ import ( fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling" fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix" + "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash" ) func testHandle() plugin.Handle { @@ -138,7 +139,7 @@ func TestPreRequest(t *testing.T) { p.wg.Wait() // 4. Verify indexer was updated - perPromptHashes := getBlockHashes(context.Background(), req1, config.BlockSizeTokens, defaultMaxPrefixBlocks) + perPromptHashes := prefixhash.GetBlockHashes(context.Background(), req1, config.BlockSizeTokens, defaultMaxPrefixBlocks) for _, promptHashes := range perPromptHashes { for _, hash := range promptHashes { pods := p.indexer().Get(hash) @@ -182,7 +183,7 @@ func TestPreRequest(t *testing.T) { p.PreRequest(context.Background(), req, res) p.wg.Wait() - perPromptHashes := getBlockHashes(context.Background(), req, config.BlockSizeTokens, defaultMaxPrefixBlocks) + perPromptHashes := prefixhash.GetBlockHashes(context.Background(), req, config.BlockSizeTokens, defaultMaxPrefixBlocks) allHashes = append(allHashes, perPromptHashes[0]) } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/types.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/types.go index a9109bef5e..36a7c93a85 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/types.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/types.go @@ -22,6 +22,7 @@ import ( k8stypes "k8s.io/apimachinery/pkg/types" "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" + "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash" ) // indexerInterface maintains an LRU cache of prompt prefix hashes and the server(s) that might have that @@ -36,8 +37,9 @@ type indexerInterface interface { // podSet holds a set of pods that may have a specific prefix hash. type podSet map[ServerID]struct{} -// blockHash is a hash of a block of request data. -type blockHash uint64 +// blockHash is a hash of a block of request data. It aliases prefixhash.BlockHash +// so this package and other prefix-aware producers share one block-hash type. +type blockHash = prefixhash.BlockHash // server contains information about a specific server/pod and its cache capacity. type server struct { diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/README.md b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/README.md new file mode 100644 index 0000000000..54034f54b2 --- /dev/null +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/README.md @@ -0,0 +1,73 @@ +# Burst Prefix Cache Producer + +**Type:** `burst-prefix-cache-producer` + +A request-level data producer that co-locates bursts of prompt-sharing requests +so a shared prefix is prefilled once instead of scattered across replicas on a +cold cache. + +## Problem + +When many requests that share a prompt arrive at the same instant (for example +the `n` group samples of an RL rollout step), every replica's prefix cache is +still cold, so a cache-state scorer scores them all zero and load balancing +spreads them. The shared prompt is then prefilled redundantly on several +replicas and the prefix-cache benefit is lost. + +## What it does + +Requests arriving within a configurable window are assigned jointly: + +1. Each request's prompt is hashed into prefix blocks (shared `prefixhash`). +2. Requests with an identical prompt prefix are grouped. An identical group of + more than one member is always a placement unit; a single request becomes a + unit only when `minColocateBlocks > 0` and it shares at least that many leading + blocks with some other request in the batch. +3. Units are steered onto a replica (or a bounded set of replicas), filling one + replica up to `maxPerReplica` before spilling to the next least-loaded replica. + Identical groups are placed first and kept whole, so the proven same-prompt + co-location is the firm structure; prefix-sharing units then attach to it. A + unit prefers a replica that already holds a unit sharing at least + `minColocateBlocks` leading blocks and is still under its fair share of the + batch (prefix co-location bounded by balance, so a shared prefix is prefilled + once without stampeding prefix-sharing units onto one replica); otherwise units + are balanced across replicas. Longer-prefix units are placed first so shorter + units match against the richest set of already-placed prefixes. +4. The producer emits `PrefixCacheMatchInfo` with a full match on the assigned + replica and zero elsewhere. + +A request that overlaps no other in the batch, and any prefix-less request, +receives no affinity (scored zero everywhere), leaving it to other scorers. +With `minColocateBlocks == 0` only identical groups are placed. + +## Scoring + +This producer emits `PrefixCacheMatchInfo`; it does not score. Reuse the +`prefix-cache-scorer` and point it at this producer: + +```yaml +- type: prefix-cache-scorer + parameters: + prefixMatchInfoProducerName: burst-prefix-cache-producer +``` + +## Configuration + +| field | default | meaning | +|---|---|---| +| `windowDurationMs` | 100 | batch window T in milliseconds; must be in `1..10000`. Every request waits up to this long, so keep it small relative to other request processing times. | +| `maxPerReplica` | -1 | max samples of one group per replica (k); -1 = unlimited (whole group to one replica) | +| `blockSizeTokens` | 64 | token block size for prefix hashing | +| `maxPrefixTokensToMatch` | 0 | cap on matched prefix tokens; 0 uses the default block cap. A positive value must be >= `blockSizeTokens`, otherwise it yields zero prefix blocks. | +| `minColocateBlocks` | 0 | min shared leading blocks for inter-unit co-location and for a single request to gain an affinity; 0 disables both (only identical groups are placed, purely load-balanced) | +| `maxBatchSize` | 1000 | max requests one window may accumulate; Produce returns an error once reached. -1 = unlimited. | + +## Operational notes + +- The producer adds up to `windowDurationMs` of latency per request while a + window fills; its producer timeout is extended to cover the window. +- Co-location is decided within a single window from the request prompts. Identity + (the samples of one prompt) is matched exactly; partial overlap between distinct + prompts is matched to `minColocateBlocks` leading blocks. Cross-step warm-cache + reuse is left to the persistent (approximate or precise) prefix cache producer, + which can run alongside this one. diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/assign.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/assign.go new file mode 100644 index 0000000000..5f1e1fa925 --- /dev/null +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/assign.go @@ -0,0 +1,349 @@ +/* +Copyright 2026 The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package burstprefix + +import ( + "cmp" + "slices" + "strings" + + fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" + "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash" +) + +// totalBlocks returns the number of prefix blocks across all prompts. +func totalBlocks(hashes [][]prefixhash.BlockHash) int { + total := 0 + for _, ph := range hashes { + total += len(ph) + } + return total +} + +// encodePrefix serializes the first n prefix blocks. Equal results imply an +// equal n-block leading prefix: each block hash chains its predecessor, so an +// equal leading hash sequence means an equal prompt prefix. +func encodePrefix(hashes [][]prefixhash.BlockHash, n int) string { + var b strings.Builder + var buf [8]byte + count := 0 + for _, ph := range hashes { + for _, h := range ph { + if count >= n { + return b.String() + } + prefixhash.PutBlockHash(buf[:], h) + b.Write(buf[:]) + count++ + } + } + return b.String() +} + +// groupKey identifies requests that share an identical prompt prefix. Requests +// with no prefix return "" and are never grouped. +func groupKey(hashes [][]prefixhash.BlockHash) string { + n := totalBlocks(hashes) + if n == 0 { + return "" + } + return encodePrefix(hashes, n) +} + +// sharedPrefixKeys returns the group keys whose requests share at least +// minColocateBlocks leading blocks with some other request in the batch. It is +// the affinity gate for non-identical requests: a lone request that overlaps no +// other keeps no affinity, while overlapping ones (e.g. one prompt contained in +// another) become placeable. Keys with fewer than minColocateBlocks blocks can +// never meet the threshold and are excluded. +func sharedPrefixKeys(groups map[string][]*entry, order []string, minColocateBlocks int) map[string]bool { + if minColocateBlocks <= 0 { + return nil + } + count := map[string]int{} + keySig := map[string]string{} + for _, key := range order { + hashes := groups[key][0].hashes + if totalBlocks(hashes) < minColocateBlocks { + continue + } + s := encodePrefix(hashes, minColocateBlocks) + keySig[key] = s + count[s] += len(groups[key]) + } + shared := map[string]bool{} + for key, s := range keySig { + if count[s] >= 2 { + shared[key] = true + } + } + return shared +} + +// batchIndex records which replicas hold each prefix block, for matching later +// groups against groups already placed in this batch. Because a whole prompt +// prefix is added at once, holding block i implies holding blocks 0..i, so a +// greedy walk that stops at the first unheld block yields the longest contiguous +// shared prefix per replica (the approximateprefix matchLongestPrefix property). +type batchIndex struct { + holders map[prefixhash.BlockHash]map[string]struct{} +} + +func newBatchIndex() *batchIndex { + return &batchIndex{holders: map[prefixhash.BlockHash]map[string]struct{}{}} +} + +// add records every block of hashes as held by replica. +func (b *batchIndex) add(hashes [][]prefixhash.BlockHash, replica string) { + for _, ph := range hashes { + for _, h := range ph { + s := b.holders[h] + if s == nil { + s = map[string]struct{}{} + b.holders[h] = s + } + s[replica] = struct{}{} + } + } +} + +// longestPrefix returns, per replica, the number of leading prefix blocks that +// replica already holds (summed across prompts) - the shared-prefix length. +func (b *batchIndex) longestPrefix(hashes [][]prefixhash.BlockHash) map[string]int { + res := map[string]int{} +outer: + for _, ph := range hashes { + for _, h := range ph { + holders := b.holders[h] + if len(holders) == 0 { + break outer + } + for name := range holders { + res[name]++ + } + } + } + return res +} + +// assign steers each batched request to a replica so prompt-sharing requests +// co-locate. It groups requests by shared prefix, then places the groups jointly +// over the batch (longest-prefix first) so a shared prefix is prefilled once per +// replica rather than scattered. Identical groups are placed before prefix-sharing +// singletons so proven same-prompt co-location anchors the layout. +func assign(entries []*entry, k, minColocateBlocks int) { + groups := map[string][]*entry{} + order := []string{} + var replicas []fwksched.Endpoint + for _, e := range entries { + key := groupKey(e.hashes) + if key == "" { + continue + } + if _, ok := groups[key]; !ok { + order = append(order, key) + } + groups[key] = append(groups[key], e) + if replicas == nil { + replicas = e.pods + } + } + if len(replicas) == 0 { + return + } + + shared := sharedPrefixKeys(groups, order, minColocateBlocks) + var groupUnits, singleUnits []string + counts := map[string]int{} // prefix-block count per placed unit, computed once + total := 0 + for _, key := range order { + switch { + case len(groups[key]) >= 2: + groupUnits = append(groupUnits, key) + counts[key] = totalBlocks(groups[key][0].hashes) + total += len(groups[key]) + case shared[key]: + singleUnits = append(singleUnits, key) + counts[key] = totalBlocks(groups[key][0].hashes) + total += len(groups[key]) + } + } + if len(groupUnits) == 0 && len(singleUnits) == 0 { + return + } + + // Longest-prefix first within each tier: a longer prefix seeds more blocks in + // the index, so shorter units match against the richest set of placed blocks. + sortByPrefixLen(groupUnits, counts) + sortByPrefixLen(singleUnits, counts) + + maxShare := (total + len(replicas) - 1) / len(replicas) // ceil: equal samples per replica + + idx := newBatchIndex() + load := map[string]int{} // batch-wide samples assigned per replica + // Groups first, then prefix-sharing singletons: same order as before, without + // allocating a combined slice. + for _, key := range groupUnits { + placeGroup(groups[key], replicas, k, minColocateBlocks, maxShare, idx, load) + } + for _, key := range singleUnits { + placeGroup(groups[key], replicas, k, minColocateBlocks, maxShare, idx, load) + } +} + +// sortByPrefixLen orders keys by descending prefix-block count (stable). +func sortByPrefixLen(keys []string, counts map[string]int) { + slices.SortStableFunc(keys, func(a, b string) int { + return cmp.Compare(counts[b], counts[a]) + }) +} + +// placeGroup places one unit (an identical group, or a single prefix-sharing +// request). The first (primary) replica is chosen with the inter-unit prefix +// preference bounded by the fair-share cap; remaining samples (when k caps the +// group) spill to least-loaded replicas. The unit's blocks are then recorded so +// later units can match against it. +func placeGroup(members []*entry, replicas []fwksched.Endpoint, k, minColocateBlocks, maxShare int, idx *batchIndex, load map[string]int) { + if len(replicas) == 0 { + return + } + hashes := members[0].hashes // identical group: any member represents it + + var matches map[string]int + preferMatch := minColocateBlocks > 0 + if preferMatch { + matches = idx.longestPrefix(hashes) + } + + perReplica := map[string]int{} // samples of this group already on a replica + i := 0 + for i < len(members) { + target := pickReplica(replicas, perReplica, k, load, matches, minColocateBlocks, maxShare, preferMatch) + name := target.GetMetadata().NamespacedName.String() + + run := len(members) - i + if k != unlimitedPerReplica { + capLeft := k - perReplica[name] + if capLeft < 1 { + capLeft = 1 // overflow: more members than k*replicas, place one and rebalance + } + if capLeft < run { + run = capLeft + } + } + for j := 0; j < run; j++ { + members[i+j].assigned = target + } + perReplica[name] += run + load[name] += run + i += run + preferMatch = false // only the primary replica gets the prefix preference + } + + for _, m := range members { + if m.assigned != nil { + idx.add(hashes, m.assigned.GetMetadata().NamespacedName.String()) + } + } +} + +// pickReplica chooses the target replica for the next run of a group. When +// preferMatch is set it first tries the replica sharing the longest prefix (at +// least minColocateBlocks blocks) that still has per-group capacity and is below +// its fair share of the batch; otherwise, or when no such replica exists, it +// falls back to the least batch-loaded replica. The fair-share bound is what +// keeps many prefix-sharing groups from stampeding onto a single replica. +func pickReplica(replicas []fwksched.Endpoint, perReplica map[string]int, k int, load, matches map[string]int, minColocateBlocks, maxShare int, preferMatch bool) fwksched.Endpoint { + if preferMatch { + var best fwksched.Endpoint + var bestName string + bestMatch := 0 + for _, r := range replicas { + name := r.GetMetadata().NamespacedName.String() + if k != unlimitedPerReplica && perReplica[name] >= k { + continue + } + if load[name] >= maxShare { + continue // at fair share: co-locating here would unbalance the batch + } + m := matches[name] + if m < minColocateBlocks { + continue + } + if best == nil || m > bestMatch || (m == bestMatch && less(r, name, load, best, bestName)) { + best, bestName, bestMatch = r, name, m + } + } + if best != nil { + return best + } + } + return pickLeastLoaded(replicas, perReplica, k, load) +} + +// pickLeastLoaded returns the least batch-loaded replica with capacity for this +// group, falling back to the overall least-loaded replica when all are at the +// cap. +func pickLeastLoaded(replicas []fwksched.Endpoint, perReplica map[string]int, k int, load map[string]int) fwksched.Endpoint { + var best fwksched.Endpoint + var bestName string + for _, r := range replicas { + name := r.GetMetadata().NamespacedName.String() + if k != unlimitedPerReplica && perReplica[name] >= k { + continue + } + if best == nil || less(r, name, load, best, bestName) { + best, bestName = r, name + } + } + if best != nil { + return best + } + for _, r := range replicas { + name := r.GetMetadata().NamespacedName.String() + if best == nil || less(r, name, load, best, bestName) { + best, bestName = r, name + } + } + return best +} + +// less reports whether replica a should be preferred over replica b: fewer +// assigned samples first, then fewer running requests, then lower name. The +// batch-local assigned-sample count (load) is the primary signal; running +// requests only break ties within a batch and deliberately differ from the +// load-aware-scorer's WaitingQueueSize signal, since placement balances work +// already committed in this window rather than queue depth observed downstream. +func less(a fwksched.Endpoint, aName string, load map[string]int, b fwksched.Endpoint, bName string) bool { + if load[aName] != load[bName] { + return load[aName] < load[bName] + } + ra, rb := runningRequests(a), runningRequests(b) + if ra != rb { + return ra < rb + } + return aName < bName +} + +// runningRequests returns the endpoint's running-request count, or 0 when +// metrics are unavailable. +func runningRequests(e fwksched.Endpoint) int { + if m := e.GetMetrics(); m != nil { + return m.RunningRequestsSize + } + return 0 +} diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/assign_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/assign_test.go new file mode 100644 index 0000000000..df6324a6e3 --- /dev/null +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/assign_test.go @@ -0,0 +1,281 @@ +/* +Copyright 2026 The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package burstprefix + +import ( + "testing" + + "github.com/stretchr/testify/assert" + k8stypes "k8s.io/apimachinery/pkg/types" + + fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer" + fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" + "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash" +) + +func testEndpoint(name string) fwksched.Endpoint { + return fwksched.NewEndpoint( + &fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: name}}, + fwkdl.NewMetrics(), fwkdl.NewAttributes()) +} + +func assignedName(e *entry) string { + if e.assigned == nil { + return "" + } + return e.assigned.GetMetadata().NamespacedName.Name +} + +// group builds n entries sharing one prompt prefix over the given replicas. +func group(n int, prefix []prefixhash.BlockHash, replicas []fwksched.Endpoint) []*entry { + entries := make([]*entry, n) + for i := range entries { + entries[i] = &entry{hashes: [][]prefixhash.BlockHash{prefix}, pods: replicas} + } + return entries +} + +// concat flattens groups into one entry slice in order. +func concat(groups ...[]*entry) []*entry { + n := 0 + for _, g := range groups { + n += len(g) + } + all := make([]*entry, 0, n) + for _, g := range groups { + all = append(all, g...) + } + return all +} + +func counts(entries []*entry) map[string]int { + c := map[string]int{} + for _, e := range entries { + c[assignedName(e)]++ + } + return c +} + +func TestAssign_UnlimitedColocatesWholeGroup(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2"), testEndpoint("pod3"), testEndpoint("pod4")} + entries := group(8, []prefixhash.BlockHash{1, 2, 3}, replicas) + + assign(entries, unlimitedPerReplica, 0) + + for _, e := range entries { + assert.Equal(t, "pod1", assignedName(e), "all samples of one group must co-locate when k=-1") + } +} + +func TestAssign_CapSpreadsEvenly(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2"), testEndpoint("pod3"), testEndpoint("pod4")} + entries := group(8, []prefixhash.BlockHash{1, 2, 3}, replicas) + + assign(entries, 2, 0) + + assert.Equal(t, map[string]int{"pod1": 2, "pod2": 2, "pod3": 2, "pod4": 2}, counts(entries), + "k=2 over 8 samples and 4 replicas must place 2 per replica") +} + +func TestAssign_CapFillsBeforeSpilling(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2"), testEndpoint("pod3"), testEndpoint("pod4")} + entries := group(8, []prefixhash.BlockHash{1, 2, 3}, replicas) + + assign(entries, 4, 0) + + // k=4 fills one replica to 4 before using the next; only two replicas used. + assert.Equal(t, map[string]int{"pod1": 4, "pod2": 4}, counts(entries)) +} + +func TestAssign_SingletonGetsNoAffinity(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")} + entries := group(1, []prefixhash.BlockHash{1, 2, 3}, replicas) + + assign(entries, unlimitedPerReplica, 0) + + assert.Nil(t, entries[0].assigned, "a singleton group has no reuse and must not receive an affinity") +} + +func TestAssign_EmptyPrefixGetsNoAffinity(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")} + entries := []*entry{ + {hashes: nil, pods: replicas}, + {hashes: nil, pods: replicas}, + } + + assign(entries, unlimitedPerReplica, 0) + + for _, e := range entries { + assert.Nil(t, e.assigned, "requests with no prefix must not be grouped") + } +} + +func TestAssign_DistinctGroupsSpreadAcrossReplicas(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2"), testEndpoint("pod3"), testEndpoint("pod4")} + groupA := group(8, []prefixhash.BlockHash{1, 2, 3}, replicas) + groupB := group(8, []prefixhash.BlockHash{9, 9, 9}, replicas) + entries := append(append([]*entry{}, groupA...), groupB...) + + assign(entries, unlimitedPerReplica, 0) + + // Each group co-locates, and the second group lands on a different, less + // loaded replica than the first. + assert.Equal(t, "pod1", assignedName(groupA[0])) + assert.Equal(t, "pod2", assignedName(groupB[0])) + for _, e := range groupA { + assert.Equal(t, "pod1", assignedName(e)) + } + for _, e := range groupB { + assert.Equal(t, "pod2", assignedName(e)) + } +} + +func TestAssign_SharedPrefixFamiliesColocateWithinFairShare(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")} + // Two prefix families of two groups each: famX groups share leading {1, 2}, + // famY groups share leading {7, 8}. + x1 := group(2, []prefixhash.BlockHash{1, 2, 3}, replicas) + x2 := group(2, []prefixhash.BlockHash{1, 2, 4}, replicas) + y1 := group(2, []prefixhash.BlockHash{7, 8, 5}, replicas) + y2 := group(2, []prefixhash.BlockHash{7, 8, 6}, replicas) + entries := concat(x1, x2, y1, y2) + + assign(entries, unlimitedPerReplica, 2) + + // Each family co-locates onto one replica, the two families land on + // different replicas, and the batch stays balanced (4 samples per replica). + assert.Equal(t, assignedName(x1[0]), assignedName(x2[0]), "groups sharing >= minColocateBlocks leading blocks co-locate") + assert.Equal(t, assignedName(y1[0]), assignedName(y2[0]), "groups sharing >= minColocateBlocks leading blocks co-locate") + assert.NotEqual(t, assignedName(x1[0]), assignedName(y1[0]), "distinct families spread across replicas") + assert.Equal(t, map[string]int{"pod1": 4, "pod2": 4}, counts(entries)) +} + +func TestAssign_FairShareSpreadsPrefixSharingGroups(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2"), testEndpoint("pod3"), testEndpoint("pod4")} + // Four groups all sharing leading {1, 2}. Pure longest-match would pile them + // onto one replica; the fair-share cap must spread them instead. + g1 := group(2, []prefixhash.BlockHash{1, 2, 3}, replicas) + g2 := group(2, []prefixhash.BlockHash{1, 2, 4}, replicas) + g3 := group(2, []prefixhash.BlockHash{1, 2, 5}, replicas) + g4 := group(2, []prefixhash.BlockHash{1, 2, 6}, replicas) + entries := concat(g1, g2, g3, g4) + + assign(entries, unlimitedPerReplica, 2) + + assert.Equal(t, map[string]int{"pod1": 2, "pod2": 2, "pod3": 2, "pod4": 2}, counts(entries), + "the fair-share cap must spread prefix-sharing groups, not stampede them onto one replica") +} + +func TestAssign_SingletonAttachesToOverlappingGroup(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")} + // An identical group plus a single request that overlaps its prefix but is + // not identical (shares leading {1,2} / {7,8}). + gx := group(2, []prefixhash.BlockHash{1, 2, 3}, replicas) + sx := group(1, []prefixhash.BlockHash{1, 2, 8}, replicas) + gy := group(2, []prefixhash.BlockHash{7, 8, 5}, replicas) + sy := group(1, []prefixhash.BlockHash{7, 8, 9}, replicas) + entries := concat(gx, sx, gy, sy) + + assign(entries, unlimitedPerReplica, 2) + + // The identical group stays whole, and the overlapping singleton attaches to + // the replica holding the group it shares a prefix with. + assert.Equal(t, assignedName(gx[0]), assignedName(gx[1]), "an identical group is never broken") + assert.Equal(t, assignedName(gx[0]), assignedName(sx[0]), "a prefix-sharing singleton attaches to the overlapping group") + assert.Equal(t, assignedName(gy[0]), assignedName(sy[0]), "a prefix-sharing singleton attaches to the overlapping group") + assert.Equal(t, map[string]int{"pod1": 3, "pod2": 3}, counts(entries)) +} + +func TestAssign_LoneSingletonGetsNoAffinity(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")} + g := group(2, []prefixhash.BlockHash{1, 2, 3}, replicas) + lone := group(1, []prefixhash.BlockHash{5, 5, 5}, replicas) // overlaps nothing + entries := concat(g, lone) + + assign(entries, unlimitedPerReplica, 2) + + assert.Nil(t, lone[0].assigned, "a request overlapping no other must keep no affinity") + assert.Equal(t, assignedName(g[0]), assignedName(g[1]), "the identical group still co-locates") +} + +func TestAssign_OverlappingSingletonsColocateWithoutGroups(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")} + // No identical duplicates at all: four distinct requests sharing leading + // blocks {1,2} (the RAG / shared-system-prompt burst). + s1 := group(1, []prefixhash.BlockHash{1, 2, 3}, replicas) + s2 := group(1, []prefixhash.BlockHash{1, 2, 4}, replicas) + s3 := group(1, []prefixhash.BlockHash{1, 2, 5}, replicas) + s4 := group(1, []prefixhash.BlockHash{1, 2, 6}, replicas) + entries := concat(s1, s2, s3, s4) + + assign(entries, unlimitedPerReplica, 2) + + for _, e := range entries { + assert.NotNil(t, e.assigned, "overlapping singletons receive an affinity even without identical groups") + } + assert.Equal(t, map[string]int{"pod1": 2, "pod2": 2}, counts(entries), + "overlapping singletons co-locate but stay balanced under the fair-share cap") +} + +func TestAssign_SharedPrefixBelowThresholdDoesNotColocate(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")} + // Both families share only 2 leading blocks, below minColocateBlocks=3. + x1 := group(2, []prefixhash.BlockHash{1, 2, 3}, replicas) + x2 := group(2, []prefixhash.BlockHash{1, 2, 4}, replicas) + y1 := group(2, []prefixhash.BlockHash{7, 8, 5}, replicas) + y2 := group(2, []prefixhash.BlockHash{7, 8, 6}, replicas) + entries := concat(x1, x2, y1, y2) + + assign(entries, unlimitedPerReplica, 3) + + // The shared prefix falls short of the threshold, so groups are placed purely + // by load and a family is not kept together. + assert.NotEqual(t, assignedName(x1[0]), assignedName(x2[0]), + "a shared prefix below minColocateBlocks must not co-locate") + assert.Equal(t, map[string]int{"pod1": 4, "pod2": 4}, counts(entries)) +} + +func TestAssign_SingleReplicaPlacesEverything(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1")} + // Distinct groups that would normally spread across replicas; with a single + // replica the fair-share and load logic has nowhere else to send them. + groupA := group(4, []prefixhash.BlockHash{1, 2, 3}, replicas) + groupB := group(4, []prefixhash.BlockHash{9, 9, 9}, replicas) + entries := concat(groupA, groupB) + + assign(entries, unlimitedPerReplica, 2) + + for _, e := range entries { + assert.Equal(t, "pod1", assignedName(e), "with a single replica every request must land on it") + } +} + +func TestAssign_OverflowGuardBalancesBeyondCap(t *testing.T) { + replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")} + // One group of 8 with k=2 over 2 replicas: k*replicas = 4 < 8, so the + // per-replica cap cannot hold the whole group. The capLeft overflow guard must + // still place every member and keep the batch balanced. + entries := group(8, []prefixhash.BlockHash{1, 2, 3}, replicas) + + assign(entries, 2, 0) + + for _, e := range entries { + assert.NotNil(t, e.assigned, "the overflow guard must still assign every member") + } + assert.Equal(t, map[string]int{"pod1": 4, "pod2": 4}, counts(entries), + "members beyond k*replicas must rebalance evenly rather than drop") +} diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin.go new file mode 100644 index 0000000000..b0bae6fade --- /dev/null +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin.go @@ -0,0 +1,212 @@ +/* +Copyright 2026 The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package burstprefix provides a request-level data producer that co-locates +// bursts of prompt-sharing requests. Requests arriving within a configurable +// window are assigned jointly: samples that share a prompt are steered to the +// same replica(s) so the shared prefix is prefilled once instead of scattered +// across replicas on a cold cache. It emits PrefixCacheMatchInfo and reuses the +// prefix-cache-scorer (point its prefixMatchInfoProducerName at this producer). +package burstprefix + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync" + "time" + + "sigs.k8s.io/controller-runtime/pkg/log" + + logutil "github.com/llm-d/llm-d-router/pkg/common/observability/logging" + "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" + "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requestcontrol" + fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" + attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix" + "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash" + tokenproducer "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer" +) + +// produceTimeoutMargin is added to the window so the director's per-producer +// timeout does not cancel a request still waiting for its batch to seal. +const produceTimeoutMargin = 300 * time.Millisecond + +var ( + _ requestcontrol.DataProducer = &dataProducer{} + _ requestcontrol.TimeoutAwareProducer = &dataProducer{} +) + +// dataProducer batches requests within a window and assigns each to a replica +// so prompt-sharing samples co-locate. +type dataProducer struct { + typedName plugin.TypedName + config config + window time.Duration + maxBlocks int + dk plugin.DataKey + + mu sync.Mutex + batch *batch +} + +// TypedName returns the type and name of the plugin. +func (p *dataProducer) TypedName() plugin.TypedName { + return p.typedName +} + +// Produces returns the data produced by the plugin. +func (p *dataProducer) Produces() map[plugin.DataKey]any { + return map[plugin.DataKey]any{p.dk: attrprefix.PrefixCacheMatchInfo{}} +} + +// Consumes declares the TokenizedPrompt dependency so the token-producer runs +// before this producer and one is auto-created when none is configured. +func (p *dataProducer) Consumes() plugin.DataDependencies { + return plugin.DataDependencies{ + Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedPrompt{}}, + } +} + +// ProduceTimeout extends the producer timeout to cover the batch window. +func (p *dataProducer) ProduceTimeout() time.Duration { + return p.window + produceTimeoutMargin +} + +// newDataProducer initializes a new burst prefix cache producer. +func newDataProducer(_ context.Context, name string, cfg config) (*dataProducer, error) { + if cfg.WindowDurationMs <= 0 { + return nil, fmt.Errorf("invalid configuration: windowDurationMs must be > 0 (current value: %d)", cfg.WindowDurationMs) + } + if cfg.WindowDurationMs > maxWindowDurationMs { + return nil, fmt.Errorf("invalid configuration: windowDurationMs must be <= %d (current value: %d)", maxWindowDurationMs, cfg.WindowDurationMs) + } + if cfg.MaxPerReplica != unlimitedPerReplica && cfg.MaxPerReplica < 1 { + return nil, fmt.Errorf("invalid configuration: maxPerReplica must be -1 (unlimited) or >= 1 (current value: %d)", cfg.MaxPerReplica) + } + if cfg.BlockSizeTokens <= 0 { + return nil, fmt.Errorf("invalid configuration: blockSizeTokens must be > 0 (current value: %d)", cfg.BlockSizeTokens) + } + if cfg.MaxPrefixTokensToMatch < 0 { + return nil, fmt.Errorf("invalid configuration: maxPrefixTokensToMatch must be >= 0 (current value: %d)", cfg.MaxPrefixTokensToMatch) + } + if cfg.MaxPrefixTokensToMatch > 0 && cfg.MaxPrefixTokensToMatch < cfg.BlockSizeTokens { + return nil, fmt.Errorf("invalid configuration: maxPrefixTokensToMatch (%d) must be 0 or >= blockSizeTokens (%d); a smaller positive value yields zero prefix blocks", cfg.MaxPrefixTokensToMatch, cfg.BlockSizeTokens) + } + if cfg.MinColocateBlocks < 0 { + return nil, fmt.Errorf("invalid configuration: minColocateBlocks must be >= 0 (current value: %d)", cfg.MinColocateBlocks) + } + if cfg.MaxBatchSize != unlimitedBatchSize && cfg.MaxBatchSize < 1 { + return nil, fmt.Errorf("invalid configuration: maxBatchSize must be -1 (unlimited) or >= 1 (current value: %d)", cfg.MaxBatchSize) + } + + maxBlocks := defaultMaxPrefixBlocks + if cfg.MaxPrefixTokensToMatch > 0 { + maxBlocks = cfg.MaxPrefixTokensToMatch / cfg.BlockSizeTokens + } + + return &dataProducer{ + typedName: plugin.TypedName{Type: PluginType, Name: name}, + config: cfg, + window: time.Duration(cfg.WindowDurationMs) * time.Millisecond, + maxBlocks: maxBlocks, + dk: attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(name), + }, nil +} + +// Produce collects the request into the current batch window, waits for the +// window to seal, then attaches PrefixCacheMatchInfo reflecting the joint +// assignment: a full match on the assigned replica and zero elsewhere. +func (p *dataProducer) Produce(ctx context.Context, request *fwksched.InferenceRequest, pods []fwksched.Endpoint) error { + hashes := prefixhash.GetBlockHashes(ctx, request, p.config.BlockSizeTokens, p.maxBlocks) + e := &entry{hashes: hashes, pods: pods, enqueued: time.Now()} + + p.mu.Lock() + if p.batch == nil { + p.batch = &batch{sealed: make(chan struct{})} + b := p.batch + time.AfterFunc(p.window, func() { p.seal(b) }) + } + b := p.batch + if p.config.MaxBatchSize != unlimitedBatchSize && len(b.entries) >= p.config.MaxBatchSize { + p.mu.Unlock() + return fmt.Errorf("burst prefix batch full: maxBatchSize %d reached", p.config.MaxBatchSize) + } + b.entries = append(b.entries, e) + p.mu.Unlock() + + select { + case <-b.sealed: + case <-ctx.Done(): + // Timed out or cancelled before the batch sealed; leave no affinity. + return ctx.Err() + } + + log.FromContext(ctx).V(logutil.DEBUG).Info("burst batch sealed", "wait_ms", time.Since(e.enqueued).Milliseconds()) + + total := totalBlocks(hashes) + matched := false + for _, pod := range pods { + matchLen := 0 + if e.assigned != nil && pod.GetMetadata().NamespacedName == e.assigned.GetMetadata().NamespacedName { + matchLen = total + matched = true + } + pod.Put(p.dk.String(), attrprefix.NewPrefixCacheMatchInfo(matchLen, total, p.config.BlockSizeTokens)) + } + if e.assigned != nil && !matched { + // The endpoint set changed between batching and release (rolling update or + // scale event): the assigned replica is not among this request's pods, so + // no affinity is produced. Surface it rather than degrading silently. + log.FromContext(ctx).V(logutil.DEBUG).Info("assigned replica absent from request pods; producing zero affinity", + "assigned", e.assigned.GetMetadata().NamespacedName.String()) + } + return nil +} + +// seal finalizes a batch: it detaches the batch so later requests start a fresh +// one, computes the joint assignment, and releases all waiting requests. +func (p *dataProducer) seal(b *batch) { + p.mu.Lock() + if b.closed { + p.mu.Unlock() + return + } + b.closed = true + if p.batch == b { + p.batch = nil + } + entries := b.entries + p.mu.Unlock() + + assign(entries, p.config.MaxPerReplica, p.config.MinColocateBlocks) + close(b.sealed) +} + +// Factory is the factory function for the burst prefix cache producer plugin. +func Factory(name string, rawParameters *json.Decoder, handle plugin.Handle) (plugin.Plugin, error) { + parameters := defaultConfig + if rawParameters != nil { + if err := rawParameters.Decode(¶meters); err != nil { + return nil, fmt.Errorf("failed to unmarshal burst prefix cache parameters: %w", err) + } + } + if handle == nil { + return nil, errors.New("plugin handle is required") + } + log.FromContext(handle.Context()).V(logutil.DEFAULT).Info("Burst prefix DataProducer initialized", "config", parameters) + return newDataProducer(handle.Context(), name, parameters) +} diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin_test.go new file mode 100644 index 0000000000..9894a1b7a7 --- /dev/null +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin_test.go @@ -0,0 +1,170 @@ +/* +Copyright 2026 The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package burstprefix + +import ( + "context" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling" + fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" + attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix" +) + +func tokenizedRequest(tokens []uint32) *fwksched.InferenceRequest { + return &fwksched.InferenceRequest{ + Body: &fwkrh.InferenceRequestBody{ + TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{tokens}}, + }, + } +} + +// assignedReplica returns the name of the endpoint this producer steered the +// request to (the one scored with a non-zero match), or "" if none. +func (p *dataProducer) assignedReplica(t *testing.T, endpoints []fwksched.Endpoint) string { + t.Helper() + for _, ep := range endpoints { + v, ok := ep.Get(p.dk.String()) + require.True(t, ok, "PrefixCacheMatchInfo must be attached to every endpoint") + info, ok := v.(*attrprefix.PrefixCacheMatchInfo) + require.True(t, ok) + if info.MatchBlocks() > 0 { + return ep.GetMetadata().NamespacedName.Name + } + } + return "" +} + +func TestNew_RejectsInvalidConfig(t *testing.T) { + _, err := newDataProducer(context.Background(), "burst", config{WindowDurationMs: 0, MaxPerReplica: -1, BlockSizeTokens: 64}) + assert.Error(t, err, "windowDurationMs must be > 0") + + _, err = newDataProducer(context.Background(), "burst", config{WindowDurationMs: 100, MaxPerReplica: 0, BlockSizeTokens: 64}) + assert.Error(t, err, "maxPerReplica 0 is invalid") + + _, err = newDataProducer(context.Background(), "burst", config{WindowDurationMs: 100, MaxPerReplica: -1, BlockSizeTokens: 0}) + assert.Error(t, err, "blockSizeTokens must be > 0") +} + +func TestProduce_ColocatesIdenticalPromptBurst(t *testing.T) { + p, err := newDataProducer(context.Background(), "burst", config{WindowDurationMs: 100, MaxPerReplica: unlimitedPerReplica, BlockSizeTokens: 4, MaxBatchSize: unlimitedBatchSize}) + require.NoError(t, err) + + const samples = 8 + names := make([]string, samples) + var wg sync.WaitGroup + for i := 0; i < samples; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + // Each request sees its own snapshot of the same four replicas. + endpoints := []fwksched.Endpoint{ + testEndpoint("pod1"), testEndpoint("pod2"), testEndpoint("pod3"), testEndpoint("pod4"), + } + err := p.Produce(context.Background(), tokenizedRequest([]uint32{1, 2, 3, 4, 5, 6, 7, 8}), endpoints) + assert.NoError(t, err) + names[i] = p.assignedReplica(t, endpoints) + }(i) + } + wg.Wait() + + for i := 0; i < samples; i++ { + assert.NotEmpty(t, names[i], "every sample of a duplicated prompt must be assigned a replica") + assert.Equal(t, names[0], names[i], "all samples of one prompt must co-locate on the same replica") + } +} + +func TestProduce_SingletonHasNoAffinity(t *testing.T) { + p, err := newDataProducer(context.Background(), "burst", config{WindowDurationMs: 50, MaxPerReplica: unlimitedPerReplica, BlockSizeTokens: 4, MaxBatchSize: unlimitedBatchSize}) + require.NoError(t, err) + + endpoints := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")} + err = p.Produce(context.Background(), tokenizedRequest([]uint32{1, 2, 3, 4}), endpoints) + require.NoError(t, err) + + assert.Empty(t, p.assignedReplica(t, endpoints), + "a lone request in a window shares no prompt and must not be steered") +} + +func TestProduce_CancelledContextBeforeSeal(t *testing.T) { + // A long window guarantees the batch cannot seal on its own before the + // already-cancelled context is observed. + p, err := newDataProducer(context.Background(), "burst", config{WindowDurationMs: maxWindowDurationMs, MaxPerReplica: unlimitedPerReplica, BlockSizeTokens: 4, MaxBatchSize: unlimitedBatchSize}) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + endpoints := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")} + err = p.Produce(ctx, tokenizedRequest([]uint32{1, 2, 3, 4}), endpoints) + require.Error(t, err, "a context cancelled before seal must return an error") + + for _, ep := range endpoints { + v, ok := ep.Get(p.dk.String()) + if !ok { + continue // no affinity attached is expected when Produce returns early + } + info, ok := v.(*attrprefix.PrefixCacheMatchInfo) + require.True(t, ok) + assert.Zero(t, info.MatchBlocks(), "a cancelled request must not produce a non-zero match") + } +} + +func TestProduce_SecondWindowAfterSeal(t *testing.T) { + p, err := newDataProducer(context.Background(), "burst", config{WindowDurationMs: 30, MaxPerReplica: unlimitedPerReplica, BlockSizeTokens: 4, MaxBatchSize: unlimitedBatchSize}) + require.NoError(t, err) + + runWindow := func(prompt []uint32) []string { + const samples = 4 + names := make([]string, samples) + var wg sync.WaitGroup + for i := 0; i < samples; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + endpoints := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2"), testEndpoint("pod3")} + assert.NoError(t, p.Produce(context.Background(), tokenizedRequest(prompt), endpoints)) + names[i] = p.assignedReplica(t, endpoints) + }(i) + } + wg.Wait() + return names + } + + first := runWindow([]uint32{1, 2, 3, 4, 5, 6, 7, 8}) + + // seal must detach the batch so the next request opens a fresh window. + p.mu.Lock() + reset := p.batch == nil + p.mu.Unlock() + assert.True(t, reset, "seal must reset p.batch to nil after the first window") + + second := runWindow([]uint32{9, 10, 11, 12, 13, 14, 15, 16}) + + for i := range first { + assert.NotEmpty(t, first[i], "every first-window sample must be assigned") + assert.Equal(t, first[0], first[i], "first-window samples must co-locate") + } + for i := range second { + assert.NotEmpty(t, second[i], "every second-window sample must be assigned") + assert.Equal(t, second[0], second[i], "second-window samples must co-locate") + } +} diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/types.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/types.go new file mode 100644 index 0000000000..b6fcaf2dcb --- /dev/null +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/types.go @@ -0,0 +1,111 @@ +/* +Copyright 2026 The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package burstprefix + +import ( + "time" + + fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" + "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash" +) + +const ( + // PluginType is the unique identifier for the burst prefix cache producer plugin. + PluginType = "burst-prefix-cache-producer" + + // unlimitedPerReplica is the MaxPerReplica sentinel that disables the + // per-replica cap, sending every sample of a group to a single replica. + unlimitedPerReplica = -1 + + // unlimitedBatchSize is the MaxBatchSize sentinel that disables the + // per-window cap on accumulated requests. + unlimitedBatchSize = -1 + + defaultWindowDurationMs = 100 + defaultMaxPerReplica = unlimitedPerReplica + defaultBlockSizeTokens = 64 + defaultMinColocateBlocks = 0 + defaultMaxBatchSize = 1000 + + // defaultMaxPrefixBlocks caps how many leading blocks form a group key. + // Two prompts identical up to this many blocks are treated as one group. + defaultMaxPrefixBlocks = 2048 + + // maxWindowDurationMs bounds the batch window. Every request waits up to this + // long, so a misconfigured large value is rejected at construction rather than + // stalling every request. + maxWindowDurationMs = 10000 +) + +// config defines the configuration for the burst prefix cache producer. +type config struct { + // WindowDurationMs is the batch window T in milliseconds. Requests arriving + // within one window are assigned jointly so samples sharing a prompt + // co-locate on the same replica(s). + WindowDurationMs int `json:"windowDurationMs"` + // MaxPerReplica caps how many samples of one group are assigned to a single + // replica (k). -1 disables the cap (all samples of a group to one replica). + MaxPerReplica int `json:"maxPerReplica"` + // BlockSizeTokens is the token block size used to compute prefix hashes. + BlockSizeTokens int `json:"blockSizeTokens"` + // MaxPrefixTokensToMatch caps prefix matching in tokens. When > 0 it sets + // maxBlocks = MaxPrefixTokensToMatch / BlockSizeTokens; otherwise + // defaultMaxPrefixBlocks applies. + MaxPrefixTokensToMatch int `json:"maxPrefixTokensToMatch"` + // MinColocateBlocks is the minimum shared leading blocks for inter-unit prefix + // co-location, and the minimum a single (non-duplicated) request must share + // with another request to gain an affinity at all. 0 disables both: only + // identical groups are placed and placement is purely load-balanced. A larger + // value avoids co-locating on a trivial shared preamble. Co-location is always + // bounded by a per-replica fair-share cap, so raising this never lets + // prefix-sharing units stampede onto one replica. + MinColocateBlocks int `json:"minColocateBlocks"` + // MaxBatchSize caps how many requests one window may accumulate. Once the cap + // is reached, Produce returns an error instead of letting a sustained burst or + // a long window grow the batch unbounded. -1 disables the cap. + MaxBatchSize int `json:"maxBatchSize"` +} + +// defaultConfig provides sensible defaults for the burst prefix cache producer. +var defaultConfig = config{ + WindowDurationMs: defaultWindowDurationMs, + MaxPerReplica: defaultMaxPerReplica, + BlockSizeTokens: defaultBlockSizeTokens, + MaxPrefixTokensToMatch: 0, + MinColocateBlocks: defaultMinColocateBlocks, + MaxBatchSize: defaultMaxBatchSize, +} + +// entry is one request collected into a batch window. +type entry struct { + hashes [][]prefixhash.BlockHash + pods []fwksched.Endpoint + // enqueued is when the request joined the batch, used to report the window wait. + enqueued time.Time + // assigned is the replica this request is steered to, filled when the batch + // is sealed. nil means no affinity (singleton group or empty prompt): the + // request is scored 0 on every endpoint so other scorers decide. + assigned fwksched.Endpoint +} + +// batch accumulates requests arriving within one window and releases them +// together once sealed. +type batch struct { + entries []*entry + sealed chan struct{} + closed bool +} diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/hashing.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing.go similarity index 75% rename from pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/hashing.go rename to pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing.go index a95c6cc3ee..d01438b01c 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/hashing.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing.go @@ -14,7 +14,10 @@ See the License for the specific language governing permissions and limitations under the License. */ -package approximateprefix +// Package prefixhash computes stable prefix block hashes for a tokenized +// prompt. It is shared by the prefix-aware data producers so they derive +// identical block hashes for the same request. +package prefixhash import ( "context" @@ -29,6 +32,9 @@ import ( "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" ) +// BlockHash is a hash of a block of request data. +type BlockHash uint64 + // HashBlock wraps a block of token IDs used for calculating prefix hashes. type HashBlock struct { // Tokens are the token IDs covered by this block. @@ -45,13 +51,13 @@ func (b HashBlock) Hash() uint64 { return 0 } -// getBlockHashes divides the tokenized prompt into blocks and calculates a +// GetBlockHashes divides the tokenized prompt into blocks and calculates a // prefix cache hash for each block. Each prompt in PerPromptTokens is hashed // independently so cross-prompt block adjacency is avoided. The first block // hash of every prompt includes the model name and cache salt (if provided). // For subsequent blocks, the hash is calculated as: hash(block i content, hash(i-1)). // It requires request.Body.TokenizedPrompt to be populated by a token-producer backend. -func getBlockHashes(ctx context.Context, request *scheduling.InferenceRequest, blockSizeTokens int, maxPrefixBlocks int) [][]blockHash { +func GetBlockHashes(ctx context.Context, request *scheduling.InferenceRequest, blockSizeTokens int, maxPrefixBlocks int) [][]BlockHash { loggerDebug := log.FromContext(ctx).V(logutil.DEBUG) if request == nil || request.Body == nil { loggerDebug.Info("Request or request data is nil, skipping hashing") @@ -64,7 +70,7 @@ func getBlockHashes(ctx context.Context, request *scheduling.InferenceRequest, b return nil } - var result [][]blockHash + var result [][]BlockHash for _, tokens := range tp.PerPromptTokens { seq := getKVCacheBlocksFromTokens(tokens, blockSizeTokens) hashes := computeBlockHashes(seq, request, maxPrefixBlocks) @@ -80,8 +86,8 @@ func getBlockHashes(ctx context.Context, request *scheduling.InferenceRequest, b } // computeBlockHashes calculates the hash for content blocks. -func computeBlockHashes(seq iter.Seq[HashBlock], request *scheduling.InferenceRequest, maxPrefixBlocks int) []blockHash { - var blockHashes []blockHash +func computeBlockHashes(seq iter.Seq[HashBlock], request *scheduling.InferenceRequest, maxPrefixBlocks int) []BlockHash { + var blockHashes []BlockHash h := xxhash.New() // Different models should have different hashes even with the same body. @@ -90,7 +96,7 @@ func computeBlockHashes(seq iter.Seq[HashBlock], request *scheduling.InferenceRe _, _ = h.Write([]byte(cacheSalt)) } - prevBlockHash := blockHash(h.Sum64()) + prevBlockHash := BlockHash(h.Sum64()) count := 0 for block := range seq { @@ -99,9 +105,9 @@ func computeBlockHashes(seq iter.Seq[HashBlock], request *scheduling.InferenceRe } h.Reset() blockID := block.Hash() - _, _ = h.Write(toBytes(blockHash(blockID))) + _, _ = h.Write(toBytes(BlockHash(blockID))) _, _ = h.Write(toBytes(prevBlockHash)) - blockHashes = append(blockHashes, blockHash(h.Sum64())) + blockHashes = append(blockHashes, BlockHash(h.Sum64())) prevBlockHash = blockHashes[len(blockHashes)-1] count++ @@ -110,12 +116,19 @@ func computeBlockHashes(seq iter.Seq[HashBlock], request *scheduling.InferenceRe return blockHashes } -func toBytes(i blockHash) []byte { +func toBytes(i BlockHash) []byte { bytes := make([]byte, 8) - binary.LittleEndian.PutUint64(bytes, uint64(i)) + PutBlockHash(bytes, i) return bytes } +// PutBlockHash writes h into the first 8 bytes of buf in little-endian order. +// buf must have length at least 8. It lets callers serialize block hashes into +// a reused buffer without allocating per hash. +func PutBlockHash(buf []byte, h BlockHash) { + binary.LittleEndian.PutUint64(buf, uint64(h)) +} + func getKVCacheBlocksFromTokens(ids []uint32, blockSizeTokens int) iter.Seq[HashBlock] { return func(yield func(HashBlock) bool) { if len(ids) == 0 || blockSizeTokens <= 0 { diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/hashing_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing_test.go similarity index 89% rename from pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/hashing_test.go rename to pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing_test.go index 99592c2423..37d8a239f9 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/hashing_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing_test.go @@ -14,7 +14,7 @@ See the License for the specific language governing permissions and limitations under the License. */ -package approximateprefix +package prefixhash import ( "context" @@ -27,6 +27,10 @@ import ( fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" ) +// testMaxPrefixBlocks is a cap large enough not to truncate any prompt in +// these tests. +const testMaxPrefixBlocks = 2048 + func TestGetKVCacheBlocksFromTokens(t *testing.T) { tests := []struct { name string @@ -122,7 +126,7 @@ func TestGetBlockHashes(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - perPromptHashes := getBlockHashes(context.Background(), tt.request, tt.blockSizeTokens, defaultMaxPrefixBlocks) + perPromptHashes := GetBlockHashes(context.Background(), tt.request, tt.blockSizeTokens, testMaxPrefixBlocks) totalBlocks := 0 for _, ph := range perPromptHashes { totalBlocks += len(ph) @@ -142,12 +146,12 @@ func TestGetBlockHashesCacheSalt(t *testing.T) { } } - noSalt := getBlockHashes(context.Background(), &fwksched.InferenceRequest{ + noSalt := GetBlockHashes(context.Background(), &fwksched.InferenceRequest{ TargetModel: "m", Body: body(""), - }, 2, defaultMaxPrefixBlocks) - salted := getBlockHashes(context.Background(), &fwksched.InferenceRequest{ + }, 2, testMaxPrefixBlocks) + salted := GetBlockHashes(context.Background(), &fwksched.InferenceRequest{ TargetModel: "m", Body: body("salt"), - }, 2, defaultMaxPrefixBlocks) + }, 2, testMaxPrefixBlocks) assert.Equal(t, len(noSalt), len(salted)) assert.NotEqual(t, noSalt, salted, "cache salt must change the block hashes") @@ -194,7 +198,7 @@ func TestGetBlockHashes_MultiPrompt(t *testing.T) { }, }, } - result := getBlockHashes(context.Background(), request, tt.blockSizeTokens, defaultMaxPrefixBlocks) + result := GetBlockHashes(context.Background(), request, tt.blockSizeTokens, testMaxPrefixBlocks) assert.Equal(t, tt.expectedPrompts, len(result)) for i, expected := range tt.expectedBlocksPerPrompt { assert.Equal(t, expected, len(result[i]), "prompt %d block count", i) @@ -221,8 +225,8 @@ func TestGetBlockHashes_MultiPromptHashIndependence(t *testing.T) { }, } - multi := getBlockHashes(context.Background(), multiPrompt, 2, defaultMaxPrefixBlocks) - single := getBlockHashes(context.Background(), singlePrompt, 2, defaultMaxPrefixBlocks) + multi := GetBlockHashes(context.Background(), multiPrompt, 2, testMaxPrefixBlocks) + single := GetBlockHashes(context.Background(), singlePrompt, 2, testMaxPrefixBlocks) assert.Equal(t, 2, len(multi), "multi-prompt should produce 2 inner slices") assert.Equal(t, 1, len(single), "single-prompt should produce 1 inner slice") @@ -232,7 +236,7 @@ func TestGetBlockHashes_MultiPromptHashIndependence(t *testing.T) { "first block of first prompt should match single prompt's first block") // Second prompt's first block [3,4] starts a fresh hash chain from the model // seed, while single prompt's second block [3,4] chains from the first block. - // These must differ — this is the cross-prompt adjacency bug the per-prompt + // These must differ - this is the cross-prompt adjacency bug the per-prompt // split prevents. assert.NotEqual(t, multi[1][0], single[0][1], "second prompt must start its own hash chain, not chain from the first prompt")