Add support for length-based scoring in prefix cache plugin

Signed-off-by: Azam Ikram <azamikram@google.com>

Author: azamikram

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

@@ -2,37 +2,59 @@
 
 **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`
 
 The attribute is typically produced by the approximate prefix cache data producer before scheduling.
 
 ## 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.

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
 }

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)
+				}
+			}
+		})
+	}
+}