diff --git a/pkg/epp/framework/plugins/scheduling/scorer/prefix/README.md b/pkg/epp/framework/plugins/scheduling/scorer/prefix/README.md index 8e8cfb16c2..d19868d80d 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/prefix/README.md +++ b/pkg/epp/framework/plugins/scheduling/scorer/prefix/README.md @@ -2,26 +2,36 @@ **Type:** `prefix-cache-scorer` -Scores candidate endpoints using `PrefixCacheMatchInfo` prepared earlier in the request pipeline. +Scores candidate endpoints based on how much of the request's prefix is already cached. It can consider both the *ratio* of the prefix matched and the *absolute length* of the match. By default it only considers the ratio of prefix matched. ## What it does -For each candidate endpoint, the scorer reads the `PrefixCacheMatchInfo` attribute and computes: +The scorer computes a score between `0.0` and `1.0` for each endpoint using a weighted combination of two metrics: + +1. **Match Ratio**: The proportion of the request prefix that is cached. +2. **Match Length**: The absolute number of cached tokens, normalized against a scale factor. + +### Scoring Formula ```text -score = matchBlocks / totalBlocks +score = matchLengthWeight*matchLengthScore + (1.0-matchLengthWeight)*matchRatioScore ``` -This produces a normalized score in the range `[0, 1]`: +Where: +* **`matchRatioScore`** = `matchBlocks` / `totalBlocks` +* **`matchLengthScore`** = `min(1.0, matchedTokens / matchLengthScaleTokens) ^ 2` + * `matchedTokens` = `matchBlocks` * `blockSize` + * **`matchLengthScaleTokens`** is the scaling factor. -- higher score: more of the request prefix is expected to be reusable from cache -- lower score: less prefix cache reuse is expected +If `matchLengthWeight` is `0.0` (default), the score simplifies to just the `matchRatioScore`. If the attribute is missing, has the wrong type, or `totalBlocks` is zero, the endpoint receives score `0`. -## Inputs consumed +### Why the Quadratic Term? + +The squaring of `matchLengthScore` introduces a non-linearity motivated by the fact that attention computation grows quadratically as a function of prompt length. This ensures that the scoring function assigns a higher priority to longer requests, where the computational and latency savings of KV-cache reuse are significantly more critical. -This scorer consumes: +## Inputs Consumed - `PrefixCacheMatchInfo` @@ -29,10 +39,22 @@ The attribute is typically produced by the approximate prefix cache data produce ## Configuration -This plugin does not define any plugin-specific parameters. +| Parameter | Type | Description | Default | +| :--- | :--- | :--- | :--- | +| `matchLengthWeight` | float | Weight of the absolute match length in the score. Must be between `0.0` and `1.0`. | `0.0` | +| `matchLengthScaleTokens` | integer | The number of tokens used to normalize `matchLengthScore`. | `8192` | + +### Example + +```yaml +plugins: +- type: prefix-cache-scorer + parameters: + matchLengthWeight: 0.5 + matchLengthScaleTokens: 16384 +``` -## Operational notes +## Operational Notes -- The scorer itself does not hash prompts or maintain cache state. -- It only converts previously prepared prefix match data into endpoint scores. -- To be useful, it should be used together with a data producer that populates `PrefixCacheMatchInfo`. +* **Tuning `matchLengthScaleTokens`**: This should be set to reflect your workload. A good rule of thumb is to set it to the **P95 prompt length** of your typical requests. Setting it too high (e.g., to the model's maximum limit when actual requests are short) will compress the score range and make the scorer less effective. +* The scorer is stateless; it does not manage cache state or hash prompts itself. It relies entirely on the data producer. diff --git a/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin.go b/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin.go index d1e1f3a3e6..dbbca58faf 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin.go @@ -20,6 +20,7 @@ import ( "context" "encoding/json" "fmt" + "math" "sigs.k8s.io/controller-runtime/pkg/log" @@ -33,12 +34,19 @@ import ( type Config struct { // The name of the data producer that produces PrefixCacheMatchInfo. PrefixMatchInfoProducerName string `json:"prefixMatchInfoProducerName,omitempty"` + // The weight assigned to match length, between 0.0 and 1.0. + MatchLengthWeight float64 `json:"matchLengthWeight,omitempty"` + // Normalization factor for match length in terms of tokens. + // Used only when MatchLengthWeight > 0. + MatchLengthScaleTokens int `json:"matchLengthScaleTokens,omitempty"` } // Plugin implements the prefix cache aware scoring logic. type Plugin struct { - typedName plugin.TypedName - prefixMatchDataKey plugin.DataKey + typedName plugin.TypedName + prefixMatchDataKey plugin.DataKey + matchLengthWeight float64 + matchLengthScaleTokens int } // compile-time type assertions @@ -49,11 +57,19 @@ var ( const ( // Type is the unique identifier for the prefix cache scorer plugin. PrefixCacheScorerPluginType = "prefix-cache-scorer" + // The default weight of the absolute match length in the score. + // Set to 0 so by default only the match ratio is considered. + defaultMatchLengthWeight = 0.0 + // Default number of tokens used as a scaling factor. + defaultMatchLengthScaleTokens = 8192 ) // PrefixCachePluginFactory defines the factory function for the Prefix plugin. func PrefixCachePluginFactory(name string, decoder *json.Decoder, handle plugin.Handle) (plugin.Plugin, error) { - var cfg Config + cfg := Config{ + MatchLengthWeight: defaultMatchLengthWeight, + MatchLengthScaleTokens: defaultMatchLengthScaleTokens, + } if decoder != nil { if err := decoder.Decode(&cfg); err != nil { return nil, fmt.Errorf("failed to decode prefix cache scorer parameters: %w", err) @@ -64,6 +80,17 @@ func PrefixCachePluginFactory(name string, decoder *json.Decoder, handle plugin. if err != nil { return nil, err } + + if cfg.MatchLengthWeight < 0.0 || cfg.MatchLengthWeight > 1.0 { + return nil, fmt.Errorf("matchLengthWeight must be between 0.0 and 1.0, got %f", cfg.MatchLengthWeight) + } + p.matchLengthWeight = cfg.MatchLengthWeight + + if p.matchLengthWeight > 0.0 && cfg.MatchLengthScaleTokens <= 0 { + return nil, fmt.Errorf("matchLengthScaleTokens must be greater than 0 when matchLengthWeight is greater than 0, got %d", cfg.MatchLengthScaleTokens) + } + p.matchLengthScaleTokens = cfg.MatchLengthScaleTokens + return p, nil } @@ -114,13 +141,29 @@ func (p *Plugin) Score(ctx context.Context, _ *fwksched.InferenceRequest, endpoi continue } - if prefixMatchInfo, ok := info.(*attrprefix.PrefixCacheMatchInfo); ok { - if prefixMatchInfo.TotalBlocks() != 0 { - scores[endpoint] = float64(prefixMatchInfo.MatchBlocks()) / float64(prefixMatchInfo.TotalBlocks()) - } - } else { + prefixMatchInfo, ok := info.(*attrprefix.PrefixCacheMatchInfo) + if !ok { logger.V(logutil.DEFAULT).Error(nil, "PrefixCacheMatchInfo has unexpected type, assigning score 0", "endpoint", endpoint) + continue + } + + matchBlocks := prefixMatchInfo.MatchBlocks() + totalBlocks := prefixMatchInfo.TotalBlocks() + if totalBlocks == 0 { + logger.V(logutil.DEFAULT).Error(nil, "totalBlocks is set to 0, assigning score 0", "endpoint", endpoint) + continue + } + + matchRatioScore := float64(matchBlocks) / float64(totalBlocks) + blockSize := prefixMatchInfo.BlockSizeTokens() + matchLengthScore := 0.0 + // Calculate matchLengthScore when match length is considered + if p.matchLengthWeight > 0.0 && blockSize > 0 { + // (matchBlocks * blockSize / matchLengthScaleTokens) ^ 2 + normalizedMatchLength := math.Min(1.0, float64(matchBlocks)*float64(blockSize)/float64(p.matchLengthScaleTokens)) + matchLengthScore = normalizedMatchLength * normalizedMatchLength } + scores[endpoint] += p.matchLengthWeight*matchLengthScore + (1.0-p.matchLengthWeight)*matchRatioScore } return scores } diff --git a/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin_test.go b/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin_test.go index 2f334ca17b..c7918dacf2 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin_test.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin_test.go @@ -18,12 +18,15 @@ package prefix import ( "context" + "encoding/json" + "strings" "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" + "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" 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" ) @@ -46,3 +49,110 @@ func TestPrefixPluginScore(t *testing.T) { assert.Equal(t, 0.5, scores[endpoint1]) assert.Equal(t, 0.2, scores[endpoint2]) } + +func TestPrefixPluginScoreWithWeights(t *testing.T) { + producerName := "approx-prefix-cache-producer" + // matchLengthWeight = 0.5, matchLengthScaleTokens = 100 + p, _ := New(context.Background(), PrefixCacheScorerPluginType, producerName) + p.matchLengthWeight = 0.5 + p.matchLengthScaleTokens = 100 + + key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(producerName).String() + + // Endpoint 1: match 5, total 10, block size 1 + // matchRatio = 5/10 = 0.5 + // matchLengthRatio = min(1.0, 5*1/100) = 0.05 -> squared = 0.0025 + // score = 0.5 * 0.0025 + 0.5 * 0.5 = 0.25125 + endpoint1 := fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, fwkdl.NewMetrics(), nil) + endpoint1.Put(key, attrprefix.NewPrefixCacheMatchInfo(5, 10, 1)) + + // Endpoint 2: match 50, total 100, block size 1 + // matchRatio = 50/100 = 0.5 + // matchLengthRatio = min(1.0, 50*1/100) = 0.5 -> squared = 0.25 + // score = 0.5 * 0.25 + 0.5 * 0.5 = 0.375 + endpoint2 := fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}}, fwkdl.NewMetrics(), nil) + endpoint2.Put(key, attrprefix.NewPrefixCacheMatchInfo(50, 100, 1)) + + endpoints := []fwksched.Endpoint{endpoint1, endpoint2} + scores := p.Score(context.Background(), nil, endpoints) + + // matchRatio is the same but we still give longer request higher score + assert.InDelta(t, 0.25125, scores[endpoint1], 1e-6) + assert.InDelta(t, 0.375, scores[endpoint2], 1e-6) +} + +func TestPrefixPluginFactoryValidation(t *testing.T) { + tests := []struct { + name string + config string + expectErr bool + wantMatchLengthWeight float64 + wantMatchLengthScale int + }{ + { + name: "valid config with defaults", + config: `{}`, + expectErr: false, + wantMatchLengthWeight: defaultMatchLengthWeight, + wantMatchLengthScale: defaultMatchLengthScaleTokens, + }, + { + name: "valid config with custom values", + config: `{"matchLengthWeight": 0.5, "matchLengthScaleTokens": 100}`, + expectErr: false, + wantMatchLengthWeight: 0.5, + wantMatchLengthScale: 100, + }, + { + name: "invalid matchLengthWeight < 0", + config: `{"matchLengthWeight": -0.1, "matchLengthScaleTokens": 100}`, + expectErr: true, + }, + { + name: "invalid matchLengthWeight > 1", + config: `{"matchLengthWeight": 1.1, "matchLengthScaleTokens": 100}`, + expectErr: true, + }, + { + name: "invalid matchLengthScaleTokens <= 0", + config: `{"matchLengthWeight": 0.5, "matchLengthScaleTokens": 0}`, + expectErr: true, + }, + { + name: "missing matchLengthScaleTokens when matchLengthWeight > 0 uses default", + config: `{"matchLengthWeight": 0.5}`, + expectErr: false, + wantMatchLengthWeight: 0.5, + wantMatchLengthScale: defaultMatchLengthScaleTokens, + }, + { + name: "zero matchLengthWeight doesn't require matchLengthScaleTokens", + config: `{"matchLengthWeight": 0.0}`, + expectErr: false, + wantMatchLengthWeight: 0.0, + wantMatchLengthScale: defaultMatchLengthScaleTokens, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + handle := plugin.NewEppHandle(context.Background(), nil) + var decoder *json.Decoder + if tt.config != "" { + decoder = json.NewDecoder(strings.NewReader(tt.config)) + } + p, err := PrefixCachePluginFactory("test", decoder, handle) + if tt.expectErr { + assert.Error(t, err) + assert.Nil(t, p) + } else { + assert.NoError(t, err) + prefixPlugin, ok := p.(*Plugin) + if assert.True(t, ok, "plugin must be of type *Plugin") { + assert.Equal(t, tt.wantMatchLengthWeight, prefixPlugin.matchLengthWeight) + assert.Equal(t, tt.wantMatchLengthScale, prefixPlugin.matchLengthScaleTokens) + } + } + }) + } +}