Skip to content

Commit 100e271

Browse files
ezrasilveraezrasilvera
authored
burst prefix-cache affinity for prompt-sharing requests (#1727)
* refactor(prefix): extract block-hashing into shared prefixhash package Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * feat(prefix): add burst prefix cache producer for grouped requests Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * feat(prefix): add inter-group longest-prefix placement to burst producer Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * feat(prefix): co-locate prefix-sharing requests, not just identical groups Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * remove duplications Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * fix(burstprefix): stop longest-prefix walk at first unheld block Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * refactor(burstprefix): stabilize prefix key and share hash encoding Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * feat(burstprefix): log when assigned replica is absent from request pods Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * feat(burstprefix): cap batch size with maxBatchSize config Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * feat(burstprefix): log per-request batch window wait time Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * fix(burstprefix): reject window and prefix-token misconfigurations Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * perf(burstprefix): precompute prefix-block counts for unit sorting Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * perf(burstprefix): avoid throwaway slice when placing units Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * docs(burstprefix): explain the less tiebreaker signal choice Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * docs(burstprefix): use cluster-DNS placeholder in deploy sample Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * test(burstprefix): set maxBatchSize in producer construction tests Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * test(burstprefix): cover cancelled-context and second-window paths Signed-off-by: Ezra Silvera <ezra@il.ibm.com> * test(burstprefix): cover single-replica and cap-overflow placement Signed-off-by: Ezra Silvera <ezra@il.ibm.com> --------- Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
1 parent c4ea166 commit 100e271

14 files changed

Lines changed: 1283 additions & 26 deletions

File tree

‎cmd/epp/runner/runner.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ import (
8383
"github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/admitter/latencyslo"
8484
"github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/admitter/probabilisticadmitter"
8585
reqdataprodprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix"
86+
"github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix"
8687
"github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload"
8788
mmproducer "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal"
8889
preciseproducer "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache"
@@ -514,6 +515,7 @@ func (r *Runner) registerInTreePlugins() {
514515
fwkplugin.Register(nohitlru.NoHitLRUType, nohitlru.Factory)
515516
fwkplugin.Register(activerequest.ActiveRequestType, activerequest.Factory)
516517
fwkplugin.Register(preciseprefixcache.PrecisePrefixCachePluginType, preciseprefixcache.PluginFactory)
518+
fwkplugin.Register(burstprefix.PluginType, burstprefix.Factory)
517519
fwkplugin.Register(mmcacheaffinity.Type, mmcacheaffinity.Factory)
518520
fwkplugin.Register(preciseproducer.PluginType, preciseproducer.PluginFactory)
519521

‎deploy/config/epp-burst-prefix-cache-config.yaml‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# Sample EPP config: burst prefix cache pipeline
2+
# token-producer -> burst-prefix-cache-producer -> prefix-cache-scorer
3+
#
4+
# Co-locates bursts of prompt-sharing requests (e.g. RL rollout group samples)
5+
# onto shared replicas so a shared prefix is prefilled once instead of scattered
6+
# across replicas on a cold cache.
7+
apiVersion: llm-d.ai/v1alpha1
8+
kind: EndpointPickerConfig
9+
plugins:
10+
- type: token-producer
11+
parameters:
12+
modelName: hf-repo/model-name # set to the model used in the vLLM deployment
13+
vllm:
14+
url: http://<vllm-service>.<namespace>.svc.cluster.local:8000
15+
- type: single-profile-handler
16+
- type: decode-filter
17+
- type: burst-prefix-cache-producer
18+
parameters:
19+
windowDurationMs: 100 # batch window T
20+
maxPerReplica: -1 # k; -1 sends all samples of a group to one replica
21+
minColocateBlocks: 0 # >0 co-locates distinct groups sharing >= N leading blocks; 0 = load-balanced
22+
- type: prefix-cache-scorer
23+
parameters:
24+
prefixMatchInfoProducerName: burst-prefix-cache-producer
25+
- type: load-aware-scorer
26+
- type: max-score-picker
27+
schedulingProfiles:
28+
- name: default
29+
plugins:
30+
- pluginRef: decode-filter
31+
- pluginRef: prefix-cache-scorer
32+
weight: 2.0
33+
- pluginRef: load-aware-scorer
34+
weight: 1.0
35+
- pluginRef: max-score-picker

‎pkg/epp/framework/plugins/requestcontrol/dataproducer/README.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ Producers may also implement additional lifecycle hooks:
1717
| `token-producer` | [`tokenizer`](tokenizer/) | `TokenizedPrompt` | Tokenizes the request prompt via vLLM `/render`; required by precise-prefix-cache-producer and context-length-aware scorers. |
1818
| `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. |
1919
| `precise-prefix-cache-producer` | [`preciseprefixcache`](preciseprefixcache/) | `PrefixCacheMatchInfo` | Maintains a precise KV-block index by subscribing to vLLM KV-events; requires `token-producer` upstream. |
20+
| `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. |
2021
| `inflight-load-producer` | [`inflightload`](inflightload/) | `InFlightLoad` | Tracks real-time in-flight request and token counts per endpoint across the full request lifecycle. |
2122
| `predicted-latency-producer` | [`predictedlatency`](predictedlatency/) | `LatencyPredictionInfo` | Trains XGBoost models via a sidecar and generates per-endpoint TTFT/TPOT predictions. |
2223
| `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:
2728
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:
2829

2930
- `precise-prefix-cache-producer` **requires** `token-producer` upstream (it consumes `TokenizedPrompt`).
31+
- `burst-prefix-cache-producer` **requires** `token-producer` upstream (it consumes `TokenizedPrompt`).
3032
- `mm-embeddings-cache-producer` **optionally** consumes `TokenizedPrompt`; configure `token-producer` first when multimodal features need tokenizer-derived hashes.
3133
- `inflight-load-producer` **optionally** consumes `PrefixCacheMatchInfo` from an approx or precise prefix producer; prefix-discounting is applied automatically when the attribute is present.
3234
- `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
3537

3638
- [Approximate Prefix Cache Producer](approximateprefix/README.md)
3739
- [Precise Prefix Cache Producer](preciseprefixcache/README.md)
40+
- [Burst Prefix Cache Producer](burstprefix/README.md)
3841
- [Token Producer](tokenizer/README.md)
3942
- [In-Flight Load Producer](inflightload/README.md)
4043
- [Predicted Latency Producer](predictedlatency/README.md)

‎pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ import (
3232
fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
3333
attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix"
3434
approxprefixconstants "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/constants"
35+
"github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash"
3536
tokenproducer "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer"
3637
)
3738

@@ -179,7 +180,7 @@ func (p *dataProducer) Produce(ctx context.Context, request *fwksched.InferenceR
179180
if p.config.MaxPrefixTokensToMatch > 0 && blockSize > 0 {
180181
maxBlocks = p.config.MaxPrefixTokensToMatch / blockSize
181182
}
182-
perPromptHashes := getBlockHashes(ctx, request, blockSize, maxBlocks)
183+
perPromptHashes := prefixhash.GetBlockHashes(ctx, request, blockSize, maxBlocks)
183184

184185
prefixCacheServers := make(map[ServerID]int)
185186
totalBlocks := 0

‎pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ import (
3232
fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling"
3333
fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
3434
attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix"
35+
"github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash"
3536
)
3637

3738
func testHandle() plugin.Handle {
@@ -138,7 +139,7 @@ func TestPreRequest(t *testing.T) {
138139
p.wg.Wait()
139140

140141
// 4. Verify indexer was updated
141-
perPromptHashes := getBlockHashes(context.Background(), req1, config.BlockSizeTokens, defaultMaxPrefixBlocks)
142+
perPromptHashes := prefixhash.GetBlockHashes(context.Background(), req1, config.BlockSizeTokens, defaultMaxPrefixBlocks)
142143
for _, promptHashes := range perPromptHashes {
143144
for _, hash := range promptHashes {
144145
pods := p.indexer().Get(hash)
@@ -182,7 +183,7 @@ func TestPreRequest(t *testing.T) {
182183
p.PreRequest(context.Background(), req, res)
183184
p.wg.Wait()
184185

185-
perPromptHashes := getBlockHashes(context.Background(), req, config.BlockSizeTokens, defaultMaxPrefixBlocks)
186+
perPromptHashes := prefixhash.GetBlockHashes(context.Background(), req, config.BlockSizeTokens, defaultMaxPrefixBlocks)
186187
allHashes = append(allHashes, perPromptHashes[0])
187188
}
188189

‎pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/types.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
k8stypes "k8s.io/apimachinery/pkg/types"
2323

2424
"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
25+
"github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash"
2526
)
2627

2728
// indexerInterface maintains an LRU cache of prompt prefix hashes and the server(s) that might have that
@@ -36,8 +37,9 @@ type indexerInterface interface {
3637
// podSet holds a set of pods that may have a specific prefix hash.
3738
type podSet map[ServerID]struct{}
3839

39-
// blockHash is a hash of a block of request data.
40-
type blockHash uint64
40+
// blockHash is a hash of a block of request data. It aliases prefixhash.BlockHash
41+
// so this package and other prefix-aware producers share one block-hash type.
42+
type blockHash = prefixhash.BlockHash
4143

4244
// server contains information about a specific server/pod and its cache capacity.
4345
type server struct {

‎pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/README.md‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
# Burst Prefix Cache Producer
2+
3+
**Type:** `burst-prefix-cache-producer`
4+
5+
A request-level data producer that co-locates bursts of prompt-sharing requests
6+
so a shared prefix is prefilled once instead of scattered across replicas on a
7+
cold cache.
8+
9+
## Problem
10+
11+
When many requests that share a prompt arrive at the same instant (for example
12+
the `n` group samples of an RL rollout step), every replica's prefix cache is
13+
still cold, so a cache-state scorer scores them all zero and load balancing
14+
spreads them. The shared prompt is then prefilled redundantly on several
15+
replicas and the prefix-cache benefit is lost.
16+
17+
## What it does
18+
19+
Requests arriving within a configurable window are assigned jointly:
20+
21+
1. Each request's prompt is hashed into prefix blocks (shared `prefixhash`).
22+
2. Requests with an identical prompt prefix are grouped. An identical group of
23+
more than one member is always a placement unit; a single request becomes a
24+
unit only when `minColocateBlocks > 0` and it shares at least that many leading
25+
blocks with some other request in the batch.
26+
3. Units are steered onto a replica (or a bounded set of replicas), filling one
27+
replica up to `maxPerReplica` before spilling to the next least-loaded replica.
28+
Identical groups are placed first and kept whole, so the proven same-prompt
29+
co-location is the firm structure; prefix-sharing units then attach to it. A
30+
unit prefers a replica that already holds a unit sharing at least
31+
`minColocateBlocks` leading blocks and is still under its fair share of the
32+
batch (prefix co-location bounded by balance, so a shared prefix is prefilled
33+
once without stampeding prefix-sharing units onto one replica); otherwise units
34+
are balanced across replicas. Longer-prefix units are placed first so shorter
35+
units match against the richest set of already-placed prefixes.
36+
4. The producer emits `PrefixCacheMatchInfo` with a full match on the assigned
37+
replica and zero elsewhere.
38+
39+
A request that overlaps no other in the batch, and any prefix-less request,
40+
receives no affinity (scored zero everywhere), leaving it to other scorers.
41+
With `minColocateBlocks == 0` only identical groups are placed.
42+
43+
## Scoring
44+
45+
This producer emits `PrefixCacheMatchInfo`; it does not score. Reuse the
46+
`prefix-cache-scorer` and point it at this producer:
47+
48+
```yaml
49+
- type: prefix-cache-scorer
50+
parameters:
51+
prefixMatchInfoProducerName: burst-prefix-cache-producer
52+
```
53+
54+
## Configuration
55+
56+
| field | default | meaning |
57+
|---|---|---|
58+
| `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. |
59+
| `maxPerReplica` | -1 | max samples of one group per replica (k); -1 = unlimited (whole group to one replica) |
60+
| `blockSizeTokens` | 64 | token block size for prefix hashing |
61+
| `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. |
62+
| `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) |
63+
| `maxBatchSize` | 1000 | max requests one window may accumulate; Produce returns an error once reached. -1 = unlimited. |
64+
65+
## Operational notes
66+
67+
- The producer adds up to `windowDurationMs` of latency per request while a
68+
window fills; its producer timeout is extended to cover the window.
69+
- Co-location is decided within a single window from the request prompts. Identity
70+
(the samples of one prompt) is matched exactly; partial overlap between distinct
71+
prompts is matched to `minColocateBlocks` leading blocks. Cross-step warm-cache
72+
reuse is left to the persistent (approximate or precise) prefix cache producer,
73+
which can run alongside this one.

0 commit comments

Comments
 (0)