Skip to content

Commit da55286

Browse files
azamikramazamikram
committed
Add support for length-based scoring in prefix cache plugin
Signed-off-by: Azam Ikram <azamikram@google.com>
1 parent ac53bff commit da55286

3 files changed

Lines changed: 153 additions & 9 deletions

File tree

‎pkg/epp/framework/plugins/scheduling/scorer/prefix/README.md‎

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,16 +6,19 @@ Scores candidate endpoints using `PrefixCacheMatchInfo` prepared earlier in the
66

77
## What it does
88

9-
For each candidate endpoint, the scorer reads the `PrefixCacheMatchInfo` attribute and computes:
9+
For each candidate endpoint, the scorer reads the `PrefixCacheMatchInfo` attribute and computes a score. By default, this is the ratio of matched blocks to total blocks. Optionally, it can also factor in the absolute length of the prefix.
10+
11+
The score is computed as:
1012

1113
```text
12-
score = matchBlocks / totalBlocks
14+
score = prefixLengthWeight * matchLengthRatio + (1.0 - prefixLengthWeight) * matchRatio
1315
```
1416

15-
This produces a normalized score in the range `[0, 1]`:
17+
Where:
18+
- `matchRatio = matchBlocks / totalBlocks` (the fraction of the request prefix matched)
19+
- `matchLengthRatio = min(1.0, (totalBlocks * blockSize / maxModelLen) ^ 2)` (normalized square of the total request prefix length relative to `maxModelLen`)
1620

17-
- higher score: more of the request prefix is expected to be reusable from cache
18-
- lower score: less prefix cache reuse is expected
21+
If `prefixLengthWeight` is `0.0` (the default), the score simplifies to just the `matchRatio`.
1922

2023
If the attribute is missing, has the wrong type, or `totalBlocks` is zero, the endpoint receives score `0`.
2124

@@ -29,7 +32,22 @@ The attribute is typically produced by the approximate prefix cache data produce
2932

3033
## Configuration
3134

32-
This plugin does not define any plugin-specific parameters.
35+
This plugin supports the following configuration parameters:
36+
37+
| Parameter | Type | Description | Default |
38+
| :--- | :--- | :--- | :--- |
39+
| `prefixLengthWeight` | float | Weight of the absolute prefix length in the score. Must be between `0.0` and `1.0`. | `0.0` |
40+
| `maxModelLen` | integer | The maximum context length (in tokens) supported by the model. Required if `prefixLengthWeight` > `0.0`. | `8192` |
41+
42+
Example configuration:
43+
44+
```yaml
45+
plugins:
46+
- type: prefix-cache-scorer
47+
parameters:
48+
prefixLengthWeight: 0.5
49+
maxModelLen: 16384
50+
```
3351
3452
## Operational notes
3553

‎pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin.go‎

Lines changed: 36 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,9 @@ package prefix
1919
import (
2020
"context"
2121
"encoding/json"
22+
"errors"
2223
"fmt"
24+
"math"
2325

2426
"sigs.k8s.io/controller-runtime/pkg/log"
2527

@@ -32,13 +34,17 @@ import (
3234
// Config defines the configuration for the prefix cache scorer plugin.
3335
type Config struct {
3436
// The name of the data producer that produces PrefixCacheMatchInfo.
35-
PrefixMatchInfoProducerName string `json:"prefixMatchInfoProducerName,omitempty"`
37+
PrefixMatchInfoProducerName string `json:"prefixMatchInfoProducerName,omitempty"`
38+
PrefixLengthWeight *float64 `json:"prefixLengthWeight,omitempty"`
39+
MaxModelLen *int `json:"maxModelLen,omitempty"`
3640
}
3741

3842
// Plugin implements the prefix cache aware scoring logic.
3943
type Plugin struct {
4044
typedName plugin.TypedName
4145
prefixMatchDataKey plugin.DataKey
46+
prefixLengthWeight float64
47+
maxModelLen int
4248
}
4349

4450
// compile-time type assertions
@@ -64,6 +70,23 @@ func PrefixCachePluginFactory(name string, decoder *json.Decoder, handle plugin.
6470
if err != nil {
6571
return nil, err
6672
}
73+
74+
if cfg.PrefixLengthWeight != nil {
75+
if *cfg.PrefixLengthWeight < 0.0 || *cfg.PrefixLengthWeight > 1.0 {
76+
return nil, fmt.Errorf("prefixLengthWeight must be between 0.0 and 1.0, got %f", *cfg.PrefixLengthWeight)
77+
}
78+
p.prefixLengthWeight = *cfg.PrefixLengthWeight
79+
}
80+
if cfg.MaxModelLen != nil {
81+
if *cfg.MaxModelLen <= 0 {
82+
return nil, fmt.Errorf("maxModelLen must be greater than 0, got %d", *cfg.MaxModelLen)
83+
}
84+
p.maxModelLen = *cfg.MaxModelLen
85+
}
86+
if p.prefixLengthWeight > 0.0 && cfg.MaxModelLen == nil {
87+
return nil, errors.New("maxModelLen must be specified when prefixLengthWeight is greater than 0")
88+
}
89+
6790
return p, nil
6891
}
6992

@@ -115,8 +138,18 @@ func (p *Plugin) Score(ctx context.Context, _ *fwksched.InferenceRequest, endpoi
115138
}
116139

117140
if prefixMatchInfo, ok := info.(*attrprefix.PrefixCacheMatchInfo); ok {
118-
if prefixMatchInfo.TotalBlocks() != 0 {
119-
scores[endpoint] = float64(prefixMatchInfo.MatchBlocks()) / float64(prefixMatchInfo.TotalBlocks())
141+
totalBlocks := prefixMatchInfo.TotalBlocks()
142+
if totalBlocks != 0 {
143+
matchRatio := float64(prefixMatchInfo.MatchBlocks()) / float64(totalBlocks)
144+
matchLengthRatio := 0.0
145+
blockSize := prefixMatchInfo.BlockSizeTokens()
146+
if blockSize > 0 && p.maxModelLen > 0 {
147+
// (TotalBlocks * BlockSize / maxModelLen) ^ 2
148+
// Capped at 1.0 as the normalized score term cannot be greater than 1.
149+
matchLengthRatio = float64(totalBlocks) * float64(blockSize) / float64(p.maxModelLen)
150+
matchLengthRatio = math.Min(1.0, matchLengthRatio*matchLengthRatio)
151+
}
152+
scores[endpoint] = p.prefixLengthWeight*matchLengthRatio + (1.0-p.prefixLengthWeight)*matchRatio
120153
}
121154
} else {
122155
logger.V(logutil.DEFAULT).Error(nil, "PrefixCacheMatchInfo has unexpected type, assigning score 0", "endpoint", endpoint)

‎pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin_test.go‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,15 @@ package prefix
1818

1919
import (
2020
"context"
21+
"encoding/json"
22+
"strings"
2123
"testing"
2224

2325
"github.com/stretchr/testify/assert"
2426
k8stypes "k8s.io/apimachinery/pkg/types"
2527

2628
fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
29+
"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
2730
fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
2831
attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix"
2932
)
@@ -46,3 +49,93 @@ func TestPrefixPluginScore(t *testing.T) {
4649
assert.Equal(t, 0.5, scores[endpoint1])
4750
assert.Equal(t, 0.2, scores[endpoint2])
4851
}
52+
53+
func TestPrefixPluginScoreWithWeights(t *testing.T) {
54+
producerName := "approx-prefix-cache-producer"
55+
// prefixLengthWeight = 0.5, maxModelLen = 100
56+
p, _ := New(context.Background(), PrefixCacheScorerPluginType, producerName)
57+
p.prefixLengthWeight = 0.5
58+
p.maxModelLen = 100
59+
60+
key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(producerName).String()
61+
62+
// Endpoint 1: match 50, total 80, block size 1
63+
// matchRatio = 50/80 = 0.625
64+
// matchLengthRatio = 80*1/100 = 0.8 -> squared = 0.64
65+
// score = 0.5 * 0.64 + 0.5 * 0.625 = 0.6325
66+
endpoint1 := fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, fwkdl.NewMetrics(), nil)
67+
endpoint1.Put(key, attrprefix.NewPrefixCacheMatchInfo(50, 80, 1))
68+
69+
// Endpoint 2: match 10, total 200, block size 1 (exceeds maxModelLen)
70+
// matchRatio = 10/200 = 0.05
71+
// matchLengthRatio = 200*1/100 = 2.0 -> squared = 4.0 -> capped at 1.0
72+
// score = 0.5 * 1.0 + 0.5 * 0.05 = 0.525
73+
endpoint2 := fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}}, fwkdl.NewMetrics(), nil)
74+
endpoint2.Put(key, attrprefix.NewPrefixCacheMatchInfo(10, 200, 1))
75+
76+
endpoints := []fwksched.Endpoint{endpoint1, endpoint2}
77+
scores := p.Score(context.Background(), nil, endpoints)
78+
79+
assert.InDelta(t, 0.6325, scores[endpoint1], 1e-6)
80+
assert.InDelta(t, 0.525, scores[endpoint2], 1e-6)
81+
}
82+
83+
func TestPrefixPluginFactoryValidation(t *testing.T) {
84+
tests := []struct {
85+
name string
86+
config string
87+
expectErr bool
88+
}{
89+
{
90+
name: "valid config with defaults",
91+
config: `{}`,
92+
expectErr: false,
93+
},
94+
{
95+
name: "valid config with custom values",
96+
config: `{"prefixLengthWeight": 0.5, "maxModelLen": 100}`,
97+
expectErr: false,
98+
},
99+
{
100+
name: "invalid prefixLengthWeight < 0",
101+
config: `{"prefixLengthWeight": -0.1, "maxModelLen": 100}`,
102+
expectErr: true,
103+
},
104+
{
105+
name: "invalid prefixLengthWeight > 1",
106+
config: `{"prefixLengthWeight": 1.1, "maxModelLen": 100}`,
107+
expectErr: true,
108+
},
109+
{
110+
name: "invalid maxModelLen <= 0",
111+
config: `{"prefixLengthWeight": 0.5, "maxModelLen": 0}`,
112+
expectErr: true,
113+
},
114+
{
115+
name: "missing maxModelLen when prefixLengthWeight > 0",
116+
config: `{"prefixLengthWeight": 0.5}`,
117+
expectErr: true,
118+
},
119+
{
120+
name: "zero prefixLengthWeight doesn't require maxModelLen",
121+
config: `{"prefixLengthWeight": 0.0}`,
122+
expectErr: false,
123+
},
124+
}
125+
126+
for _, tt := range tests {
127+
t.Run(tt.name, func(t *testing.T) {
128+
handle := plugin.NewEppHandle(context.Background(), nil)
129+
var decoder *json.Decoder
130+
if tt.config != "" {
131+
decoder = json.NewDecoder(strings.NewReader(tt.config))
132+
}
133+
_, err := PrefixCachePluginFactory("test", decoder, handle)
134+
if tt.expectErr {
135+
assert.Error(t, err)
136+
} else {
137+
assert.NoError(t, err)
138+
}
139+
})
140+
}
141+
}

0 commit comments

Comments
 (0)