Add weightMode toggle for multimodal cache affinity scorer

Author: guygir

deploy/config/epp-weighted-mm-embeddings-cache-config.yaml

@@ -0,0 +1,48 @@
+# Sample EPP configuration with weighted multimodal embeddings cache affinity.
+apiVersion: llm-d.ai/v1alpha1
+kind: EndpointPickerConfig
+plugins:
+  - type: token-producer
+    parameters:
+      modelName: hf-repo/model-name           # set to the model used in the vLLM deployment
+      vllm:
+        url: http://localhost:8000
+  - type: endpoint-notification-source
+  - type: metrics-data-source
+  - type: core-metrics-extractor
+  - type: mm-embeddings-cache-producer
+    parameters:
+      cacheSize: 10000
+      weightMode: weighted
+  - type: mm-embeddings-cache-scorer
+  - type: precise-prefix-cache-scorer
+    parameters:
+      tokenProcessorConfig:
+        blockSize: 64
+      indexerConfig:
+        kvBlockIndexConfig:
+          enableMetrics: true
+  - type: queue-scorer
+  - type: max-score-picker
+  - type: single-profile-handler
+  - type: decode-filter
+dataLayer:
+  sources:
+    - pluginRef: metrics-data-source
+      extractors:
+        - pluginRef: core-metrics-extractor
+    - pluginRef: endpoint-notification-source
+      extractors:
+        - pluginRef: mm-embeddings-cache-producer
+        - pluginRef: precise-prefix-cache-scorer
+schedulingProfiles:
+  - name: default
+    plugins:
+      - pluginRef: decode-filter
+      - pluginRef: precise-prefix-cache-scorer
+        weight: 10
+      - pluginRef: mm-embeddings-cache-scorer
+        weight: 4
+      - pluginRef: queue-scorer
+        weight: 1
+      - pluginRef: max-score-picker

pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/README.md

@@ -8,8 +8,10 @@ Produces multimodal embeddings cache match data for downstream scheduling plugin
 
 For each request, the producer extracts stable multimodal item hashes from:
 
-- `TokenizedPrompt.MultiModalFeatures`, when a `token-producer` is configured
+- `TokenizedPrompt.MultiModalFeatures`, when weighted mode is enabled and
+  `token-producer` is configured
 - typed OpenAI chat-completions structured media blocks, as a lightweight fallback
+- `Generate.Features.MMHashes`, when present
 
 It keeps an in-memory LRU map from multimodal hash to the set of pods that recently
 handled that item. During scheduling, it attaches `EncoderCacheMatchInfo` to each
@@ -18,15 +20,25 @@ same image, video, or audio input.
 
 Repeated references to the same multimodal hash within one request count once.
 
+## Weight Modes
+
+- `unweighted` (default): each unique multimodal hash has size `1`. No
+  `token-producer` dependency is declared.
+- `weighted`: when tokenized multimodal metadata is available, each item size is
+  `MultiModalFeature.Length` (falls back to `1` when length is zero). Requires
+  `token-producer` in the plugin list so placeholder lengths are available before
+  scheduling.
+
 ## Inputs Consumed
 
-This plugin declares:
+Unweighted mode does not declare upstream dependencies.
+
+Weighted mode declares:
 
-- `TokenizedPrompt`
+- `TokenizedPrompt` from `token-producer`
 
-When `token-producer` is present, this orders tokenization before multimodal match
-data production. If tokenized prompt data is absent at runtime, the producer falls
-back to typed structured chat-completions media blocks.
+If tokenized prompt data is absent at runtime, the producer falls back to typed
+structured chat-completions media blocks with unit item sizes.
 
 ## Data Produced
 
@@ -40,8 +52,9 @@ The producer supports the following runtime parameters:
 
 - `cacheSize` (integer, default: `10000`): maximum number of multimodal hash entries
   retained in the best-effort pod-affinity cache.
+- `weightMode` (string, default: `unweighted`): `unweighted` or `weighted`.
 
-**Configuration Examples:**
+**Unweighted configuration example (tokenizer-free chat media):**
 
 ```yaml
 plugins:
@@ -60,16 +73,19 @@ schedulingProfiles:
         weight: 2
 ```
 
+**Weighted configuration example:**
+
 ```yaml
 plugins:
   - type: token-producer
     parameters:
       modelName: Qwen/Qwen2.5-1.5B-Instruct
       vllm:
-        http: http://localhost:8000
+        url: http://localhost:8000
   - type: mm-embeddings-cache-producer
     parameters:
       cacheSize: 10000
+      weightMode: weighted
   - type: mm-embeddings-cache-scorer
 schedulingProfiles:
   - name: decode
@@ -78,9 +94,15 @@ schedulingProfiles:
         weight: 4
 ```
 
+See also `deploy/config/epp-mm-embeddings-cache-config.yaml` and
+`deploy/config/epp-weighted-mm-embeddings-cache-config.yaml`.
+
 ## Operational Notes
 
 - The cache is a best-effort routing signal, not a correctness dependency.
-- The producer remains tokenizer-free for request shapes where typed media blocks are
-  sufficient; `token-producer` is only required when relying on upstream multimodal
-  metadata.
+- `cacheSize` bounds the EPP routing LRU only; it is not a model-server encoder
+  cache capacity knob.
+- Unweighted mode remains lightweight: structured chat media blocks are enough and
+  `token-producer` is not required.
+- Weighted mode opts into `token-producer` so multimodal placeholder lengths can
+  influence affinity scores.

pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer.go

@@ -38,6 +38,7 @@ import (
 	fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling"
 	"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
 	attrmm "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/multimodal"
+	tokenproducer "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer"
 	k8stypes "k8s.io/apimachinery/pkg/types"
 )
 
@@ -47,6 +48,11 @@ const (
 
 	defaultCacheSize   = 10000
 	podCleanupInterval = 2 * time.Minute
+
+	// WeightModeUnweighted counts each unique multimodal hash once with unit size.
+	WeightModeUnweighted = "unweighted"
+	// WeightModeWeighted uses token-producer multimodal placeholder lengths as item sizes.
+	WeightModeWeighted = "weighted"
 )
 
 var (
@@ -56,12 +62,16 @@ var (
 	_ requestcontrol.DataProducer = &Producer{}
 	_ requestcontrol.PreRequest   = &Producer{}
 	_ fwkdl.EndpointExtractor     = &Producer{}
+	_ plugin.ConsumerPlugin       = &Producer{}
 )
 
 // Parameters configures the multimodal encoder-cache data producer.
 type Parameters struct {
 	// CacheSize defines the maximum number of mm_hash -> pod-set entries to track.
 	CacheSize int `json:"cacheSize"`
+	// WeightMode selects unweighted unit sizes (default) or weighted tokenized
+	// placeholder lengths from token-producer metadata.
+	WeightMode string `json:"weightMode,omitempty"`
 }
 
 // Factory creates a multimodal encoder-cache data producer.
@@ -72,6 +82,12 @@ func Factory(name string, rawParameters *json.Decoder, handle plugin.Handle) (pl
 			return nil, fmt.Errorf("failed to parse the parameters of the '%s' plugin - %w", ProducerType, err)
 		}
 	}
+	if parameters.WeightMode == "" {
+		parameters.WeightMode = WeightModeUnweighted
+	}
+	if parameters.WeightMode != WeightModeUnweighted && parameters.WeightMode != WeightModeWeighted {
+		return nil, fmt.Errorf("failed to parse the parameters of the '%s' plugin - invalid weightMode %q", ProducerType, parameters.WeightMode)
+	}
 
 	return New(handle.Context(), name, &parameters, handle.PodList)
 }
@@ -81,6 +97,7 @@ func Factory(name string, rawParameters *json.Decoder, handle plugin.Handle) (pl
 type Producer struct {
 	typedName   plugin.TypedName
 	dk          plugin.DataKey
+	useWeighted bool
 	cache       *lru.Cache[string, map[string]struct{}]
 	pluginState *plugin.PluginState
 	podList     func() []k8stypes.NamespacedName
@@ -111,9 +128,15 @@ func New(ctx context.Context, name string, params *Parameters, podList func() []
 		return nil, fmt.Errorf("failed to create multimodal encoder-cache LRU with size %d: %w", cacheSize, err)
 	}
 
+	useWeighted := false
+	if params != nil && params.WeightMode == WeightModeWeighted {
+		useWeighted = true
+	}
+
 	p := &Producer{
 		typedName:   plugin.TypedName{Type: ProducerType, Name: name},
 		dk:          attrmm.EncoderCacheMatchInfoKey.WithNonEmptyProducerName(name),
+		useWeighted: useWeighted,
 		cache:       cache,
 		pluginState: plugin.NewPluginState(ctx),
 		podList:     podList,
@@ -147,6 +170,15 @@ func (p *Producer) Produces() map[plugin.DataKey]any {
 	return map[plugin.DataKey]any{p.dk: attrmm.EncoderCacheMatchInfo{}}
 }
 
+// Consumes declares the token-producer dependency only for weighted mode so
+// unweighted configs are not forced to pull in tokenization.
+func (p *Producer) Consumes() map[plugin.DataKey]any {
+	if !p.useWeighted {
+		return nil
+	}
+	return map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: scheduling.TokenizedPrompt{}}
+}
+
 // PluginState returns request-scoped state shared between producer extension points.
 func (p *Producer) PluginState() *plugin.PluginState {
 	return p.pluginState
@@ -155,7 +187,7 @@ func (p *Producer) PluginState() *plugin.PluginState {
 // Produce attaches multimodal encoder-cache match data to endpoints.
 func (p *Producer) Produce(ctx context.Context, request *scheduling.InferenceRequest, endpoints []scheduling.Endpoint) error {
 	logger := log.FromContext(ctx).V(logging.DEBUG)
-	requestItems := ExtractMMItems(request)
+	requestItems := p.extractMMItems(request)
 	if len(requestItems) == 0 {
 		logger.Info("No multimodal content found, skipping encoder-cache match data")
 		return nil
@@ -179,16 +211,25 @@ func (p *Producer) Produce(ctx context.Context, request *scheduling.InferenceReq
 	return nil
 }
 
+func (p *Producer) extractMMItems(request *scheduling.InferenceRequest) []attrmm.MatchItem {
+	return extractMMItems(request, p.useWeighted)
+}
+
 // ExtractMMItems returns deterministic, unique multimodal encoder-cache items
-// for a request. Parser-provided multimodal features are preferred; if
-// unavailable, typed structured media blocks are hashed from stable identifiers.
+// for a request using the unweighted path. Parser-provided multimodal features
+// are preferred; if unavailable, typed structured media blocks are hashed from
+// stable identifiers.
 func ExtractMMItems(request *scheduling.InferenceRequest) []attrmm.MatchItem {
+	return extractMMItems(request, false)
+}
+
+func extractMMItems(request *scheduling.InferenceRequest, useWeighted bool) []attrmm.MatchItem {
 	if request == nil || request.Body == nil {
 		return nil
 	}
 
 	if request.Body.TokenizedPrompt != nil && len(request.Body.TokenizedPrompt.MultiModalFeatures) > 0 {
-		return itemsFromTokenizedPrompt(request.Body.TokenizedPrompt.MultiModalFeatures)
+		return itemsFromTokenizedPrompt(request.Body.TokenizedPrompt.MultiModalFeatures, useWeighted)
 	}
 
 	if g := request.Body.Generate; g != nil && g.Features != nil && len(g.Features.MMHashes) > 0 {
@@ -215,17 +256,28 @@ func itemsFromGenerateFeatures(mmHashes map[string][]string) []attrmm.MatchItem
 	return itemSlice(itemsByHash)
 }
 
-func itemsFromTokenizedPrompt(features []fwkrh.MultiModalFeature) []attrmm.MatchItem {
+func itemsFromTokenizedPrompt(features []fwkrh.MultiModalFeature, useWeighted bool) []attrmm.MatchItem {
 	itemsByHash := map[string]attrmm.MatchItem{}
 	for _, feature := range features {
-		if feature.Hash == "" {
-			continue
-		}
-		addItem(itemsByHash, feature.Hash)
+		addTokenizedItem(itemsByHash, feature, useWeighted)
 	}
 	return itemSlice(itemsByHash)
 }
 
+func addTokenizedItem(itemsByHash map[string]attrmm.MatchItem, feature fwkrh.MultiModalFeature, useWeighted bool) {
+	if feature.Hash == "" {
+		return
+	}
+	if _, exists := itemsByHash[feature.Hash]; exists {
+		return
+	}
+	size := 1
+	if useWeighted && feature.Length > 0 {
+		size = feature.Length
+	}
+	itemsByHash[feature.Hash] = attrmm.MatchItem{Hash: feature.Hash, Size: size}
+}
+
 func itemsFromChat(request *fwkrh.ChatCompletionsRequest) []attrmm.MatchItem {
 	itemsByHash := map[string]attrmm.MatchItem{}
 	for _, message := range request.Messages {

pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer_test.go

@@ -45,8 +45,59 @@ func TestFactory(t *testing.T) {
 	assert.Equal(t, "mm-producer", created.TypedName().Name)
 	assert.Equal(t, ProducerType, created.TypedName().Type)
 
+	weightedRaw, err := json.Marshal(map[string]any{"weightMode": WeightModeWeighted})
+	require.NoError(t, err)
+	weighted, err := Factory("weighted-producer", plugin.StrictDecoder(weightedRaw), &testHandle{ctx: context.Background()})
+	require.NoError(t, err)
+	producer, ok := weighted.(*Producer)
+	require.True(t, ok)
+	assert.True(t, producer.useWeighted)
+	assert.NotNil(t, producer.Consumes())
+
+	unweighted, err := Factory("unweighted-producer", plugin.StrictDecoder(json.RawMessage(`{}`)), &testHandle{ctx: context.Background()})
+	require.NoError(t, err)
+	unweightedProducer, ok := unweighted.(*Producer)
+	require.True(t, ok)
+	assert.False(t, unweightedProducer.useWeighted)
+	assert.Nil(t, unweightedProducer.Consumes())
+
 	_, err = Factory("bad", plugin.StrictDecoder(json.RawMessage(`{"cacheSize":"bad"}`)), &testHandle{ctx: context.Background()})
 	require.Error(t, err)
+
+	_, err = Factory("bad-weight", plugin.StrictDecoder(json.RawMessage(`{"weightMode":"invalid"}`)), &testHandle{ctx: context.Background()})
+	require.Error(t, err)
+}
+
+func TestExtractWeightedMMItemsFromTokenizedPrompt(t *testing.T) {
+	items := extractMMItems(&scheduling.InferenceRequest{
+		Body: &fwkrh.InferenceRequestBody{
+			TokenizedPrompt: &fwkrh.TokenizedPrompt{
+				MultiModalFeatures: []fwkrh.MultiModalFeature{
+					{Hash: "image-a", Length: 576},
+					{Hash: "image-b", Length: 0},
+					{Hash: "image-a", Length: 144},
+				},
+			},
+		},
+	}, true)
+
+	assert.ElementsMatch(t, []attrmm.MatchItem{
+		{Hash: "image-a", Size: 576},
+		{Hash: "image-b", Size: 1},
+	}, items)
+}
+
+func TestProduceWeightedUsesPlaceholderLengths(t *testing.T) {
+	producer := newTestProducer(t, &Parameters{WeightMode: WeightModeWeighted}, nil)
+	podA := k8stypes.NamespacedName{Namespace: "default", Name: "pod-a"}
+	endpointA := newEndpoint(podA)
+	request := requestWithHashes("req-weighted", map[string]int{"hash-a": 80, "hash-c": 20})
+
+	require.NoError(t, producer.Produce(context.Background(), request, []scheduling.Endpoint{endpointA}))
+
+	assertMatchInfo(t, producer, endpointA,
+		nil,
+		[]attrmm.MatchItem{{Hash: "hash-a", Size: 80}, {Hash: "hash-c", Size: 20}})
 }
 
 func TestExtractMMItemsFromTokenizedPrompt(t *testing.T) {