Fix the scoring function to use matchBlocks instead of totalBlocks

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

Author: azamikram

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

@@ -16,7 +16,7 @@ score = prefixLengthWeight * matchLengthRatio + (1.0 - prefixLengthWeight) * mat
 
 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`)
+- `matchLengthRatio = min(1.0, matchedBlocks * blockSize / maxModelLen) ^ 2` (square of normalized matched blocks relative to `maxModelLen`)
 
 If `prefixLengthWeight` is `0.0` (the default), the score simplifies to just the `matchRatio`.
 

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

@@ -138,16 +138,17 @@ func (p *Plugin) Score(ctx context.Context, _ *fwksched.InferenceRequest, endpoi
 		}
 
 		if prefixMatchInfo, ok := info.(*attrprefix.PrefixCacheMatchInfo); ok {
+			matchBlocks := prefixMatchInfo.MatchBlocks()
 			totalBlocks := prefixMatchInfo.TotalBlocks()
 			if totalBlocks != 0 {
-				matchRatio := float64(prefixMatchInfo.MatchBlocks()) / float64(totalBlocks)
+				matchRatio := float64(matchBlocks) / float64(totalBlocks)
 				matchLengthRatio := 0.0
 				blockSize := prefixMatchInfo.BlockSizeTokens()
 				if blockSize > 0 && p.maxModelLen > 0 {
-					// (TotalBlocks * BlockSize / maxModelLen) ^ 2
+					// (MatchBlocks * 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)
+					matchLengthRatio = math.Min(1.0, float64(matchBlocks)*float64(blockSize)/float64(p.maxModelLen))
+					matchLengthRatio *= matchLengthRatio
 				}
 				scores[endpoint] = p.prefixLengthWeight*matchLengthRatio + (1.0-p.prefixLengthWeight)*matchRatio
 			}

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

@@ -59,25 +59,26 @@ func TestPrefixPluginScoreWithWeights(t *testing.T) {
 
 	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
+	// 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(50, 80, 1))
+	endpoint1.Put(key, attrprefix.NewPrefixCacheMatchInfo(5, 10, 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
+	// 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(10, 200, 1))
+	endpoint2.Put(key, attrprefix.NewPrefixCacheMatchInfo(50, 100, 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)
+	// 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) {