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

@@ -6,16 +6,19 @@ Scores candidate endpoints using `PrefixCacheMatchInfo` prepared earlier in the
 
 ## What it does
 
-For each candidate endpoint, the scorer reads the `PrefixCacheMatchInfo` attribute and computes:
+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.
+
+The score is computed as:
 
 ```text
-score = matchBlocks / totalBlocks
+score = prefixLengthWeight * matchLengthRatio + (1.0 - prefixLengthWeight) * matchRatio
 ```
 
-This produces a normalized score in the range `[0, 1]`:
+Where:
+- `matchRatio = matchBlocks / totalBlocks` (the fraction of the request prefix matched)
+- `matchLengthRatio = min(1.0, (totalBlocks * blockSize / maxModelLen) ^ 2)` (normalized square of the total request prefix length relative to `maxModelLen`)
 
-- higher score: more of the request prefix is expected to be reusable from cache
-- lower score: less prefix cache reuse is expected
+If `prefixLengthWeight` is `0.0` (the default), the score simplifies to just the `matchRatio`.
 
 If the attribute is missing, has the wrong type, or `totalBlocks` is zero, the endpoint receives score `0`.
 
@@ -29,7 +32,22 @@ The attribute is typically produced by the approximate prefix cache data produce
 
 ## Configuration
 
-This plugin does not define any plugin-specific parameters.
+This plugin supports the following configuration parameters:
+
+| Parameter | Type | Description | Default |
+| :--- | :--- | :--- | :--- |
+| `prefixLengthWeight` | float | Weight of the absolute prefix length in the score. Must be between `0.0` and `1.0`. | `0.0` |
+| `maxModelLen` | integer | The maximum context length (in tokens) supported by the model. Required if `prefixLengthWeight` > `0.0`. | `8192` |
+
+Example configuration:
+
+```yaml
+plugins:
+- type: prefix-cache-scorer
+  parameters:
+    prefixLengthWeight: 0.5
+    maxModelLen: 16384
+```
 
 ## Operational notes
 

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

@@ -19,7 +19,9 @@ package prefix
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"math"
 
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
@@ -32,13 +34,17 @@ import (
 // Config defines the configuration for the prefix cache scorer plugin.
 type Config struct {
 	// The name of the data producer that produces PrefixCacheMatchInfo.
-	PrefixMatchInfoProducerName string `json:"prefixMatchInfoProducerName,omitempty"`
+	PrefixMatchInfoProducerName string   `json:"prefixMatchInfoProducerName,omitempty"`
+	PrefixLengthWeight          *float64 `json:"prefixLengthWeight,omitempty"`
+	MaxModelLen                 *int     `json:"maxModelLen,omitempty"`
 }
 
 // Plugin implements the prefix cache aware scoring logic.
 type Plugin struct {
 	typedName          plugin.TypedName
 	prefixMatchDataKey plugin.DataKey
+	prefixLengthWeight float64
+	maxModelLen        int
 }
 
 // compile-time type assertions
@@ -64,6 +70,23 @@ func PrefixCachePluginFactory(name string, decoder *json.Decoder, handle plugin.
 	if err != nil {
 		return nil, err
 	}
+
+	if cfg.PrefixLengthWeight != nil {
+		if *cfg.PrefixLengthWeight < 0.0 || *cfg.PrefixLengthWeight > 1.0 {
+			return nil, fmt.Errorf("prefixLengthWeight must be between 0.0 and 1.0, got %f", *cfg.PrefixLengthWeight)
+		}
+		p.prefixLengthWeight = *cfg.PrefixLengthWeight
+	}
+	if cfg.MaxModelLen != nil {
+		if *cfg.MaxModelLen <= 0 {
+			return nil, fmt.Errorf("maxModelLen must be greater than 0, got %d", *cfg.MaxModelLen)
+		}
+		p.maxModelLen = *cfg.MaxModelLen
+	}
+	if p.prefixLengthWeight > 0.0 && cfg.MaxModelLen == nil {
+		return nil, errors.New("maxModelLen must be specified when prefixLengthWeight is greater than 0")
+	}
+
 	return p, nil
 }
 
@@ -115,8 +138,18 @@ func (p *Plugin) Score(ctx context.Context, _ *fwksched.InferenceRequest, endpoi
 		}
 
 		if prefixMatchInfo, ok := info.(*attrprefix.PrefixCacheMatchInfo); ok {
-			if prefixMatchInfo.TotalBlocks() != 0 {
-				scores[endpoint] = float64(prefixMatchInfo.MatchBlocks()) / float64(prefixMatchInfo.TotalBlocks())
+			totalBlocks := prefixMatchInfo.TotalBlocks()
+			if totalBlocks != 0 {
+				matchRatio := float64(prefixMatchInfo.MatchBlocks()) / float64(totalBlocks)
+				matchLengthRatio := 0.0
+				blockSize := prefixMatchInfo.BlockSizeTokens()
+				if blockSize > 0 && p.maxModelLen > 0 {
+					// (TotalBlocks * BlockSize / maxModelLen) ^ 2
+					// Capped at 1.0 as the normalized score term cannot be greater than 1.
+					matchLengthRatio = float64(totalBlocks) * float64(blockSize) / float64(p.maxModelLen)
+					matchLengthRatio = math.Min(1.0, matchLengthRatio*matchLengthRatio)
+				}
+				scores[endpoint] = p.prefixLengthWeight*matchLengthRatio + (1.0-p.prefixLengthWeight)*matchRatio
 			}
 		} else {
 			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

@@ -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,93 @@ 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"
+	// prefixLengthWeight = 0.5, maxModelLen = 100
+	p, _ := New(context.Background(), PrefixCacheScorerPluginType, producerName)
+	p.prefixLengthWeight = 0.5
+	p.maxModelLen = 100
+
+	key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(producerName).String()
+
+	// Endpoint 1: match 50, total 80, block size 1
+	// matchRatio = 50/80 = 0.625
+	// matchLengthRatio = 80*1/100 = 0.8 -> squared = 0.64
+	// score = 0.5 * 0.64 + 0.5 * 0.625 = 0.6325
+	endpoint1 := fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, fwkdl.NewMetrics(), nil)
+	endpoint1.Put(key, attrprefix.NewPrefixCacheMatchInfo(50, 80, 1))
+
+	// Endpoint 2: match 10, total 200, block size 1 (exceeds maxModelLen)
+	// matchRatio = 10/200 = 0.05
+	// matchLengthRatio = 200*1/100 = 2.0 -> squared = 4.0 -> capped at 1.0
+	// score = 0.5 * 1.0 + 0.5 * 0.05 = 0.525
+	endpoint2 := fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}}, fwkdl.NewMetrics(), nil)
+	endpoint2.Put(key, attrprefix.NewPrefixCacheMatchInfo(10, 200, 1))
+
+	endpoints := []fwksched.Endpoint{endpoint1, endpoint2}
+	scores := p.Score(context.Background(), nil, endpoints)
+
+	assert.InDelta(t, 0.6325, scores[endpoint1], 1e-6)
+	assert.InDelta(t, 0.525, scores[endpoint2], 1e-6)
+}
+
+func TestPrefixPluginFactoryValidation(t *testing.T) {
+	tests := []struct {
+		name      string
+		config    string
+		expectErr bool
+	}{
+		{
+			name:      "valid config with defaults",
+			config:    `{}`,
+			expectErr: false,
+		},
+		{
+			name:      "valid config with custom values",
+			config:    `{"prefixLengthWeight": 0.5, "maxModelLen": 100}`,
+			expectErr: false,
+		},
+		{
+			name:      "invalid prefixLengthWeight < 0",
+			config:    `{"prefixLengthWeight": -0.1, "maxModelLen": 100}`,
+			expectErr: true,
+		},
+		{
+			name:      "invalid prefixLengthWeight > 1",
+			config:    `{"prefixLengthWeight": 1.1, "maxModelLen": 100}`,
+			expectErr: true,
+		},
+		{
+			name:      "invalid maxModelLen <= 0",
+			config:    `{"prefixLengthWeight": 0.5, "maxModelLen": 0}`,
+			expectErr: true,
+		},
+		{
+			name:      "missing maxModelLen when prefixLengthWeight > 0",
+			config:    `{"prefixLengthWeight": 0.5}`,
+			expectErr: true,
+		},
+		{
+			name:      "zero prefixLengthWeight doesn't require maxModelLen",
+			config:    `{"prefixLengthWeight": 0.0}`,
+			expectErr: false,
+		},
+	}
+
+	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))
+			}
+			_, err := PrefixCachePluginFactory("test", decoder, handle)
+			if tt.expectErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}