fix: Make TokenizedPrompt.MultiModalFeatures per-prompt

pkg/epp/flowcontrol/benchmark/benchmark_test.go

@@ -346,7 +346,7 @@ func BenchmarkFlowController_FullPath(b *testing.B) {
 			//    detector reads to compute saturation.
 			infReq := &scheduling.InferenceRequest{
 				RequestID: reqID,
-				Body:      &requesthandling.InferenceRequestBody{TokenizedPrompt: &requesthandling.TokenizedPrompt{PerPromptTokens: [][]uint32{benchTokenIDs}}},
+				Body:      &requesthandling.InferenceRequestBody{TokenizedRequest: &requesthandling.TokenizedRequest{Prompts: []requesthandling.PromptTokens{{TokenIDs: benchTokenIDs}}}},
 			}
 			schedResult := &scheduling.SchedulingResult{ProfileResults: profileResults}
 			h.producer.PreRequest(ctx, infReq, schedResult)

pkg/epp/flowcontrol/integration_test.go

@@ -82,7 +82,7 @@ func TestConcurrentSaturationReads(t *testing.T) {
 			req := &fwksched.InferenceRequest{
 				RequestID: fmt.Sprintf("req-%d", i),
 				Body: &fwkrh.InferenceRequestBody{
-					TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 10)}},
+					TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 10)}}},
 				},
 			}
 			result := &fwksched.SchedulingResult{
@@ -155,7 +155,7 @@ func TestSaturationFullLoop(t *testing.T) {
 		req := &fwksched.InferenceRequest{
 			RequestID: fmt.Sprintf("prefill-%d", i),
 			Body: &fwkrh.InferenceRequestBody{
-				TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 50)}},
+				TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 50)}}},
 			},
 		}
 		result := &fwksched.SchedulingResult{
@@ -444,7 +444,7 @@ func TestUsageLimitThresholdGatesDispatch(t *testing.T) {
 		req := &fwksched.InferenceRequest{
 			RequestID: fmt.Sprintf("inflight-%d", i),
 			Body: &fwkrh.InferenceRequestBody{
-				TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 10)}},
+				TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 10)}}},
 			},
 		}
 		result := &fwksched.SchedulingResult{
@@ -985,7 +985,7 @@ func TestEndpointReregistrationSaturationAccuracy(t *testing.T) {
 	oldReq := &fwksched.InferenceRequest{
 		RequestID: "old-req",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 50)}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 50)}}},
 		},
 	}
 	oldResult := &fwksched.SchedulingResult{
@@ -1041,7 +1041,7 @@ func TestEndpointReregistrationSaturationAccuracy(t *testing.T) {
 	newReq := &fwksched.InferenceRequest{
 		RequestID: "new-req",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 50)}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 50)}}},
 		},
 	}
 	newResult := &fwksched.SchedulingResult{
@@ -1102,7 +1102,7 @@ func TestEndpointIdentityCollisionDuringPodReplacement(t *testing.T) {
 	req := &fwksched.InferenceRequest{
 		RequestID: "new-pod-req",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 50)}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 50)}}},
 		},
 	}
 	result := &fwksched.SchedulingResult{

pkg/epp/framework/interface/requesthandling/types.go

@@ -112,9 +112,9 @@ type InferenceRequestBody struct {
 	// If the payload is unmarshaled, we can perform advanced processing (like prefix cache aware routing).
 	// If it remains as raw bytes, such processing may not be supported.
 	Payload RequestPayload `json:"-"`
-	// TokenizedPrompt contains parser-derived tokenization results when available.
+	// TokenizedRequest contains parser-derived tokenization results when available.
 	// It is nil when the request was not already tokenized.
-	TokenizedPrompt *TokenizedPrompt `json:"-"`
+	TokenizedRequest *TokenizedRequest `json:"-"`
 
 	// Stream indicates whether the request specifies a streaming response (e.g., via a stream field).
 	// This typically implies the model server's response will be streamed.
@@ -169,32 +169,36 @@ func MaxOutputTokensFromPayload(m PayloadMap, keys ...string) *int64 {
 	return nil
 }
 
-// TokenizedPrompt contains the result of tokenizing the request prompt.
+// TokenizedRequest contains the result of tokenizing the request prompt.
 // It is consumed by scheduling and request-control plugins that benefit from
 // actual token data such as prefix-cache awareness.
-type TokenizedPrompt struct {
-	// PerPromptTokens holds the token IDs for each prompt in the request.
-	// Single-prompt requests (chat, generate, single-string completions) use a
-	// length-1 outer slice. Multi-string completions use one inner slice per
-	// prompt string.
-	PerPromptTokens [][]uint32
-	// MultiModalFeatures holds one entry per multimodal item in prompt order.
-	// Nil if the prompt contains no multimodal content. Offsets are relative
-	// to PerPromptTokens[0] (always single-prompt when multimodal content is
-	// present).
-	MultiModalFeatures []MultiModalFeature
+type TokenizedRequest struct {
+	// Prompts holds the per-prompt token data. Single-prompt requests (chat,
+	// generate, single-string completions) use a length-1 slice. Multi-string
+	// completions use one entry per prompt string.
+	Prompts []PromptTokens
 	// CacheSalt isolates prefix caches across requests. Populated by the token-producer.
 	CacheSalt string
 }
 
+// PromptTokens bundles the token IDs and multimodal features for a single
+// prompt in the request.
+type PromptTokens struct {
+	// TokenIDs holds the token IDs for this prompt.
+	TokenIDs []uint32
+	// MultiModalFeatures holds multimodal items for this prompt, ordered by
+	// token position. Nil if the prompt contains no multimodal content.
+	MultiModalFeatures []MultiModalFeature
+}
+
 // TokenCount returns the total number of tokens across all prompts.
-func (tp *TokenizedPrompt) TokenCount() int {
+func (tp *TokenizedRequest) TokenCount() int {
 	if tp == nil {
 		return 0
 	}
 	n := 0
-	for _, pp := range tp.PerPromptTokens {
-		n += len(pp)
+	for _, p := range tp.Prompts {
+		n += len(p.TokenIDs)
 	}
 	return n
 }
@@ -207,9 +211,10 @@ type MultiModalFeature struct {
 	Modality Modality
 	// Hash is the content hash of the item, used for KV-cache reuse across requests.
 	Hash string
-	// Offset is the index of the first placeholder token for this item in TokenIDs.
+	// Offset is the index of the first placeholder token for this item
+	// in the owning PromptTokens.TokenIDs slice.
 	Offset int
-	// Length is the number of placeholder tokens this item occupies in TokenIDs.
+	// Length is the number of placeholder tokens this item occupies.
 	Length int
 }
 

pkg/epp/framework/interface/scheduling/types.go

@@ -32,7 +32,7 @@ type Modality = fwkrh.Modality
 
 const ModalityImage = fwkrh.ModalityImage
 
-type TokenizedPrompt = fwkrh.TokenizedPrompt
+type TokenizedRequest = fwkrh.TokenizedRequest
 
 type MultiModalFeature = fwkrh.MultiModalFeature
 

pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector_test.go

@@ -807,7 +807,7 @@ func makeTokenRequest(requestID string, inputTokens int) *fwksched.InferenceRequ
 	return &fwksched.InferenceRequest{
 		RequestID: requestID,
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, inputTokens)}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, inputTokens)}}},
 		},
 	}
 }

pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go

@@ -126,12 +126,12 @@ func (p *dataProducer) Produces() map[plugin.DataKey]any {
 	return map[plugin.DataKey]any{p.dk: attrprefix.PrefixCacheMatchInfo{}}
 }
 
-// Consumes declares the TokenizedPrompt dependency so the data-layer DAG orders
+// Consumes declares the TokenizedRequest dependency so the data-layer DAG orders
 // the token-producer before this producer runs and auto-creates one when none
 // is configured.
 func (p *dataProducer) Consumes() plugin.DataDependencies {
 	return plugin.DataDependencies{
-		Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedPrompt{}},
+		Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedRequest{}},
 	}
 }
 

pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go

@@ -53,7 +53,7 @@ func disableMinBlockSizeClamp(t *testing.T) {
 // tokenizedBody returns a request body carrying only a tokenized prompt.
 func tokenizedBody(tokenIDs []uint32) *fwkrh.InferenceRequestBody {
 	return &fwkrh.InferenceRequestBody{
-		TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{tokenIDs}},
+		TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: tokenIDs}}},
 	}
 }
 
@@ -663,8 +663,8 @@ func TestProduce_MultiPrompt(t *testing.T) {
 		RequestID:   uuid.NewString(),
 		TargetModel: "test-model",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{
-				PerPromptTokens: [][]uint32{{1, 2, 3}, {4, 5}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{
+				Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3}}, {TokenIDs: []uint32{4, 5}}},
 			},
 		},
 	}
@@ -706,8 +706,8 @@ func TestMultiPromptMatchAggregation(t *testing.T) {
 		RequestID:   uuid.NewString(),
 		TargetModel: "test-model",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{
-				PerPromptTokens: [][]uint32{{1, 2, 3}, {4, 5}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{
+				Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3}}, {TokenIDs: []uint32{4, 5}}},
 			},
 		},
 	}
@@ -725,8 +725,8 @@ func TestMultiPromptMatchAggregation(t *testing.T) {
 		RequestID:   uuid.NewString(),
 		TargetModel: "test-model",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{
-				PerPromptTokens: [][]uint32{{1, 2, 3}, {4, 5}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{
+				Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3}}, {TokenIDs: []uint32{4, 5}}},
 			},
 		},
 	}
@@ -759,8 +759,8 @@ func TestMultiPromptPartialMatch(t *testing.T) {
 		RequestID:   uuid.NewString(),
 		TargetModel: "test-model",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{
-				PerPromptTokens: [][]uint32{{1, 2}, {3, 4}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{
+				Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2}}, {TokenIDs: []uint32{3, 4}}},
 			},
 		},
 	}
@@ -778,8 +778,8 @@ func TestMultiPromptPartialMatch(t *testing.T) {
 		RequestID:   uuid.NewString(),
 		TargetModel: "test-model",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{
-				PerPromptTokens: [][]uint32{{1, 2}, {5, 6}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{
+				Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2}}, {TokenIDs: []uint32{5, 6}}},
 			},
 		},
 	}

pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin.go

@@ -77,7 +77,7 @@ func (p *dataProducer) Produces() map[plugin.DataKey]any {
 // before this producer and one is auto-created when none is configured.
 func (p *dataProducer) Consumes() plugin.DataDependencies {
 	return plugin.DataDependencies{
-		Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedPrompt{}},
+		Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedRequest{}},
 	}
 }
 

pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin_test.go

@@ -32,7 +32,7 @@ import (
 func tokenizedRequest(tokens []uint32) *fwksched.InferenceRequest {
 	return &fwksched.InferenceRequest{
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{tokens}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: tokens}}},
 		},
 	}
 }

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

@@ -563,7 +563,7 @@ func (p *InFlightLoadProducer) Produces() map[fwkplugin.DataKey]any {
 	}
 }
 
-// Consumes declares TokenizedPrompt as required so the data-layer DAG orders a
+// Consumes declares TokenizedRequest as required so the data-layer DAG orders a
 // token-producer ahead of this producer and auto-creates one when none is
 // configured; without it the input-token estimate silently reads zero.
 // PrefixCacheMatchInfo is optional — used to discount the already-cached prompt
@@ -572,7 +572,7 @@ func (p *InFlightLoadProducer) Produces() map[fwkplugin.DataKey]any {
 func (p *InFlightLoadProducer) Consumes() fwkplugin.DataDependencies {
 	return fwkplugin.DataDependencies{
 		Required: map[fwkplugin.DataKey]any{
-			tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedPrompt{},
+			tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedRequest{},
 		},
 		Optional: map[fwkplugin.DataKey]any{
 			p.prefixMatchInfoDK: attrprefix.PrefixCacheMatchInfo{},

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

@@ -57,7 +57,7 @@ func TestInFlightLoadProducer_Consumes(t *testing.T) {
 
 	deps := newTestProducer(t).Consumes()
 
-	// TokenizedPrompt is required so the data-layer DAG auto-creates a
+	// TokenizedRequest is required so the data-layer DAG auto-creates a
 	// token-producer and orders it ahead of this producer; without it the input
 	// token estimate silently reads zero.
 	require.Contains(t, deps.Required, tokenproducer.TokenizedPromptDataKey)
@@ -587,8 +587,8 @@ func makeTokenRequest(requestID string, inputTokens int) *fwksched.InferenceRequ
 	return &fwksched.InferenceRequest{
 		RequestID: requestID,
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{
-				PerPromptTokens: [][]uint32{make([]uint32, inputTokens)},
+			TokenizedRequest: &fwkrh.TokenizedRequest{
+				Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, inputTokens)}},
 			},
 		},
 	}

pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/token_estimator.go

@@ -87,10 +87,10 @@ func (e *SimpleTokenEstimator) Estimate(request *fwksched.InferenceRequest) int6
 // EstimateInput returns the input token count read from the tokenized prompt,
 // or 0 when no tokenization is available.
 func (e *SimpleTokenEstimator) EstimateInput(request *fwksched.InferenceRequest) int64 {
-	if request == nil || request.Body == nil || request.Body.TokenizedPrompt == nil {
+	if request == nil || request.Body == nil || request.Body.TokenizedRequest == nil {
 		return 0
 	}
-	return int64(request.Body.TokenizedPrompt.TokenCount())
+	return int64(request.Body.TokenizedRequest.TokenCount())
 }
 
 // EstimateOutput returns the estimated output token count given the input token

pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/token_estimator_test.go

@@ -30,8 +30,8 @@ import (
 func tokenizedRequest(n int) *fwksched.InferenceRequest {
 	return &fwksched.InferenceRequest{
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{
-				PerPromptTokens: [][]uint32{make([]uint32, n)},
+			TokenizedRequest: &fwkrh.TokenizedRequest{
+				Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, n)}},
 			},
 		},
 	}
@@ -73,7 +73,7 @@ func TestSimpleTokenEstimator_Estimate(t *testing.T) {
 			name: "Empty tokenized prompt",
 			request: &fwksched.InferenceRequest{
 				Body: &fwkrh.InferenceRequestBody{
-					TokenizedPrompt: &fwkrh.TokenizedPrompt{},
+					TokenizedRequest: &fwkrh.TokenizedRequest{},
 				},
 			},
 			expected: 0,

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

@@ -250,12 +250,12 @@ func (p *Producer) Produces() map[plugin.DataKey]any {
 	return map[plugin.DataKey]any{p.dk: attrmm.EncoderCacheMatchInfo{}}
 }
 
-// Consumes declares the TokenizedPrompt dependency so the data-layer DAG orders
+// Consumes declares the TokenizedRequest dependency so the data-layer DAG orders
 // the token-producer before this producer runs and auto-creates one when none
 // is configured; multimodal features come from the tokenizer output.
 func (p *Producer) Consumes() plugin.DataDependencies {
 	return plugin.DataDependencies{
-		Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: scheduling.TokenizedPrompt{}},
+		Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: scheduling.TokenizedRequest{}},
 	}
 }
 
@@ -297,16 +297,18 @@ func (p *Producer) Produce(ctx context.Context, request *scheduling.InferenceReq
 // ExtractMMItems returns deterministic, unique multimodal encoder-cache items
 // derived from the tokenized prompt's multimodal features.
 func ExtractMMItems(request *scheduling.InferenceRequest) []attrmm.MatchItem {
-	if request == nil || request.Body == nil || request.Body.TokenizedPrompt == nil {
+	if request == nil || request.Body == nil || request.Body.TokenizedRequest == nil {
 		return nil
 	}
 
 	itemsByHash := map[string]attrmm.MatchItem{}
-	for _, feature := range request.Body.TokenizedPrompt.MultiModalFeatures {
-		if feature.Hash == "" {
-			continue
+	for _, p := range request.Body.TokenizedRequest.Prompts {
+		for _, feature := range p.MultiModalFeatures {
+			if feature.Hash == "" {
+				continue
+			}
+			addItem(itemsByHash, feature.Hash, string(feature.Modality))
 		}
-		addItem(itemsByHash, feature.Hash, string(feature.Modality))
 	}
 	return itemSlice(itemsByHash)
 }